Add session config, limit-after-reverse fix, SLT tests, and plan display for exact reverse scan

zhuqi-lucas · zhuqi-lucas · commit fdfa2c773ce4 · 2026-04-15T11:52:24.000+08:00
- Add `enable_exact_reverse_scan` to ParquetOptions (default false)
- Wire config through `with_table_parquet_options` to set `exact_reverse`
- Fix limit correctness: skip passing limit to parquet reader when
  reverse_rows=true; apply limit in ReversedRowGroupStream after reversal
- Display `scan_direction=Reversed` for exact, `reverse_row_groups=true` for inexact
- Add 4 snapshot tests for exact reverse (removes Sort, fetch pushdown, through projection)
- Add 8 SLT tests: EXPLAIN plans + result verification, LIMIT, OFFSET+LIMIT, ASC unchanged, toggle off
diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs
@@ -732,6 +732,13 @@ config_namespace! {
         /// parquet reader setting. 0 means no caching.
         pub max_predicate_cache_size: Option<usize>, default = None
 
+        /// (reading) If true, reverse scans produce exact descending order
+        /// by reversing rows within each row group. This allows the Sort
+        /// operator to be removed entirely and fetch/limit to be pushed
+        /// down to the scan. If false (default), reverse scans only reverse
+        /// row group order (inexact), keeping TopK above for final sorting.
+        pub enable_exact_reverse_scan: bool, default = false
+
         // The following options affect writing to parquet files
         // and map to parquet::file::properties::WriterProperties
 
diff --git a/datafusion/common/src/file_options/parquet_writer.rs b/datafusion/common/src/file_options/parquet_writer.rs
@@ -209,6 +209,7 @@ impl ParquetOptions {
             coerce_int96: _,     // not used for writer props
             skip_arrow_metadata: _,
             max_predicate_cache_size: _,
+            enable_exact_reverse_scan: _, // reads not used for writer props
         } = self;
 
         let mut builder = WriterProperties::builder()
diff --git a/datafusion/core/tests/physical_optimizer/pushdown_sort.rs b/datafusion/core/tests/physical_optimizer/pushdown_sort.rs
@@ -33,9 +33,9 @@ use std::sync::Arc;
 
 use crate::physical_optimizer::test_utils::{
     OptimizationTest, coalesce_batches_exec, coalesce_partitions_exec, parquet_exec,
-    parquet_exec_with_sort, projection_exec, projection_exec_with_alias,
-    repartition_exec, schema, simple_projection_exec, sort_exec, sort_exec_with_fetch,
-    sort_expr, sort_expr_named, test_scan_with_ordering,
+    parquet_exec_with_sort, parquet_exec_with_sort_exact_reverse, projection_exec,
+    projection_exec_with_alias, repartition_exec, schema, simple_projection_exec,
+    sort_exec, sort_exec_with_fetch, sort_expr, sort_expr_named, test_scan_with_ordering,
 };
 
 #[test]
@@ -1038,3 +1038,119 @@ fn test_sort_pushdown_with_test_scan_arbitrary_ordering() {
     "
     );
 }
+
+// ============================================================================
+// EXACT REVERSE SCAN TESTS
+// ============================================================================
+// These tests verify behavior when exact_reverse is enabled on ParquetSource.
+// With exact reverse, the Sort operator is removed entirely and fetch is pushed
+// down to the scan.
+
+#[test]
+fn test_exact_reverse_removes_sort() {
+    // With exact_reverse=true, Sort should be removed entirely
+    let schema = schema();
+    let a = sort_expr("a", &schema);
+    let source_ordering = LexOrdering::new(vec![a.clone()]).unwrap();
+    let source =
+        parquet_exec_with_sort_exact_reverse(schema.clone(), vec![source_ordering]);
+
+    let desc_ordering = LexOrdering::new(vec![a.reverse()]).unwrap();
+    let plan = sort_exec(desc_ordering, source);
+
+    insta::assert_snapshot!(
+        OptimizationTest::new(plan, PushdownSort::new(), true),
+        @r"
+    OptimizationTest:
+      input:
+        - SortExec: expr=[a@0 DESC NULLS LAST], preserve_partitioning=[false]
+        -   DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet
+      output:
+        Ok:
+          - DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet, scan_direction=Reversed
+    "
+    );
+}
+
+#[test]
+fn test_exact_reverse_with_fetch_pushes_limit() {
+    // With exact_reverse=true, Sort with fetch should be removed and fetch
+    // pushed down to the scan
+    let schema = schema();
+    let a = sort_expr("a", &schema);
+    let source_ordering = LexOrdering::new(vec![a.clone()]).unwrap();
+    let source =
+        parquet_exec_with_sort_exact_reverse(schema.clone(), vec![source_ordering]);
+
+    let desc_ordering = LexOrdering::new(vec![a.reverse()]).unwrap();
+    let plan = sort_exec_with_fetch(desc_ordering, Some(10), source);
+
+    insta::assert_snapshot!(
+        OptimizationTest::new(plan, PushdownSort::new(), true),
+        @r"
+    OptimizationTest:
+      input:
+        - SortExec: TopK(fetch=10), expr=[a@0 DESC NULLS LAST], preserve_partitioning=[false]
+        -   DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet
+      output:
+        Ok:
+          - DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], limit=10, output_ordering=[a@0 ASC], file_type=parquet, scan_direction=Reversed
+    "
+    );
+}
+
+#[test]
+fn test_exact_reverse_through_projection_with_fetch() {
+    // Exact reverse with fetch pushes through projection
+    let schema = schema();
+    let a = sort_expr("a", &schema);
+    let source_ordering = LexOrdering::new(vec![a.clone()]).unwrap();
+    let source =
+        parquet_exec_with_sort_exact_reverse(schema.clone(), vec![source_ordering]);
+
+    let projection = simple_projection_exec(source, vec![0, 1]);
+
+    let desc_ordering = LexOrdering::new(vec![a.reverse()]).unwrap();
+    let plan = sort_exec_with_fetch(desc_ordering, Some(5), projection);
+
+    insta::assert_snapshot!(
+        OptimizationTest::new(plan, PushdownSort::new(), true),
+        @r"
+    OptimizationTest:
+      input:
+        - SortExec: TopK(fetch=5), expr=[a@0 DESC NULLS LAST], preserve_partitioning=[false]
+        -   ProjectionExec: expr=[a@0 as a, b@1 as b]
+        -     DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet
+      output:
+        Ok:
+          - ProjectionExec: expr=[a@0 as a, b@1 as b]
+          -   DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], limit=5, output_ordering=[a@0 ASC], file_type=parquet, scan_direction=Reversed
+    "
+    );
+}
+
+#[test]
+fn test_exact_reverse_without_fetch_no_limit() {
+    // Exact reverse without fetch: Sort removed, no limit on scan
+    let schema = schema();
+    let a = sort_expr("a", &schema);
+    let source_ordering = LexOrdering::new(vec![a.clone()]).unwrap();
+    let source =
+        parquet_exec_with_sort_exact_reverse(schema.clone(), vec![source_ordering]);
+
+    let desc_ordering = LexOrdering::new(vec![a.reverse()]).unwrap();
+    let plan = sort_exec(desc_ordering, source); // no fetch
+
+    insta::assert_snapshot!(
+        OptimizationTest::new(plan, PushdownSort::new(), true),
+        @r"
+    OptimizationTest:
+      input:
+        - SortExec: expr=[a@0 DESC NULLS LAST], preserve_partitioning=[false]
+        -   DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet
+      output:
+        Ok:
+          - DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet, scan_direction=Reversed
+    "
+    );
+}
diff --git a/datafusion/core/tests/physical_optimizer/test_utils.rs b/datafusion/core/tests/physical_optimizer/test_utils.rs
@@ -101,6 +101,23 @@ pub(crate) fn parquet_exec_with_sort(
     DataSourceExec::from_data_source(config)
 }
 
+/// Create a single parquet file that is sorted with exact_reverse enabled
+pub(crate) fn parquet_exec_with_sort_exact_reverse(
+    schema: SchemaRef,
+    output_ordering: Vec<LexOrdering>,
+) -> Arc<DataSourceExec> {
+    let source = ParquetSource::new(schema).with_exact_reverse(true);
+    let config = FileScanConfigBuilder::new(
+        ObjectStoreUrl::parse("test:///").unwrap(),
+        Arc::new(source),
+    )
+    .with_file(PartitionedFile::new("x".to_string(), 100))
+    .with_output_ordering(output_ordering)
+    .build();
+
+    DataSourceExec::from_data_source(config)
+}
+
 fn int64_stats() -> ColumnStatistics {
     ColumnStatistics {
         null_count: Precision::Absent,
diff --git a/datafusion/datasource-parquet/src/opener.rs b/datafusion/datasource-parquet/src/opener.rs
@@ -581,8 +581,14 @@ impl FileOpener for ParquetOpener {
             // Apply the prepared plan to the builder
             builder = prepared_plan.apply_to_builder(builder);
 
+            // When reverse_rows is enabled, limit must be applied AFTER row
+            // reversal (in ReversedRowGroupStream), not at the parquet reader
+            // level. Applying limit here would read the first N rows in forward
+            // order and then reverse them, giving wrong results.
             if let Some(limit) = limit {
-                builder = builder.with_limit(limit)
+                if !reverse_rows {
+                    builder = builder.with_limit(limit)
+                }
             }
 
             if let Some(max_predicate_cache_size) = max_predicate_cache_size {
@@ -705,9 +711,11 @@ impl FileOpener for ParquetOpener {
 
             // When exact reverse is enabled, wrap the stream to buffer
             // and reverse rows per row group. Memory cost: O(largest_RG).
+            // The limit is applied here (after reversal) instead of at the
+            // parquet reader level so that we get the correct reversed rows.
             let stream: futures::stream::BoxStream<'static, Result<RecordBatch>> =
                 if reverse_rows {
-                    ReversedRowGroupStream::new(stream, rg_row_counts).boxed()
+                    ReversedRowGroupStream::new(stream, rg_row_counts, limit).boxed()
                 } else {
                     stream.boxed()
                 };
@@ -746,10 +754,12 @@ struct ReversedRowGroupStream<S> {
     output_buffer: VecDeque<RecordBatch>,
     /// Whether the inner stream is exhausted
     done: bool,
+    /// Optional row limit (applied after reversal for correct results)
+    remaining_limit: Option<usize>,
 }
 
 impl<S> ReversedRowGroupStream<S> {
-    fn new(inner: S, rg_row_counts: Vec<usize>) -> Self {
+    fn new(inner: S, rg_row_counts: Vec<usize>, limit: Option<usize>) -> Self {
         let rows_remaining = rg_row_counts.first().copied().unwrap_or(0);
         Self {
             inner,
@@ -759,6 +769,25 @@ impl<S> ReversedRowGroupStream<S> {
             buffer: Vec::new(),
             output_buffer: VecDeque::new(),
             done: false,
+            remaining_limit: limit,
+        }
+    }
+
+    /// Truncate batch to remaining limit and update the counter.
+    /// Returns the (possibly truncated) batch.
+    fn apply_limit(&mut self, batch: RecordBatch) -> RecordBatch {
+        if let Some(remaining) = self.remaining_limit.as_mut() {
+            let rows = batch.num_rows();
+            if rows <= *remaining {
+                *remaining -= rows;
+                batch
+            } else {
+                let truncated = batch.slice(0, *remaining);
+                *remaining = 0;
+                truncated
+            }
+        } else {
+            batch
         }
     }
 
@@ -798,11 +827,14 @@ where
         mut self: Pin<&mut Self>,
         cx: &mut Context<'_>,
     ) -> Poll<Option<Self::Item>> {
-        use Poll;
+        // Check if limit has been reached
+        if self.remaining_limit == Some(0) {
+            return Poll::Ready(None);
+        }
 
         // First, emit any already-reversed batches
         if let Some(batch) = self.output_buffer.pop_front() {
-            return Poll::Ready(Some(Ok(batch)));
+            return Poll::Ready(Some(Ok(self.apply_limit(batch))));
         }
 
         if self.done {
@@ -824,7 +856,7 @@ where
                             return Poll::Ready(Some(Err(e)));
                         }
                         if let Some(batch) = self.output_buffer.pop_front() {
-                            return Poll::Ready(Some(Ok(batch)));
+                            return Poll::Ready(Some(Ok(self.apply_limit(batch))));
                         }
                     }
                 }
@@ -838,7 +870,7 @@ where
                         return Poll::Ready(Some(Err(e)));
                     }
                     if let Some(batch) = self.output_buffer.pop_front() {
-                        return Poll::Ready(Some(Ok(batch)));
+                        return Poll::Ready(Some(Ok(self.apply_limit(batch))));
                     }
                     return Poll::Ready(None);
                 }
diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs
@@ -336,6 +336,7 @@ impl ParquetSource {
         mut self,
         table_parquet_options: TableParquetOptions,
     ) -> Self {
+        self.exact_reverse = table_parquet_options.global.enable_exact_reverse_scan;
         self.table_parquet_options = table_parquet_options;
         self
     }
@@ -651,8 +652,10 @@ impl FileSource for ParquetSource {
 
                 write!(f, "{predicate_string}")?;
 
-                // Add reverse_scan info if enabled
-                if self.reverse_row_groups {
+                // Add reverse scan info if enabled
+                if self.reverse_row_groups && self.reverse_rows {
+                    write!(f, ", scan_direction=Reversed")?;
+                } else if self.reverse_row_groups {
                     write!(f, ", reverse_row_groups=true")?;
                 }
 
diff --git a/datafusion/sqllogictest/test_files/sort_pushdown.slt b/datafusion/sqllogictest/test_files/sort_pushdown.slt
