This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 17694ac0d0 feat(parquet): multi-column lexicographic stats reorder for
TopK sort pushdown (#23888)
17694ac0d0 is described below
commit 17694ac0d09bcecd33cdae574e4dbb523de13e5a
Author: Qi Zhu <[email protected]>
AuthorDate: Tue Aug 4 11:56:20 2026 +0800
feat(parquet): multi-column lexicographic stats reorder for TopK sort
pushdown (#23888)
## Which issue does this PR close?
- Part of #22198 (the **P0 — multi-column lexicographic reorder** item)
- Part of #23036 (Sort Pushdown epic)
## Rationale for this change
`PreparedAccessPlan::reorder_by_statistics` (row-group level) and
`ParquetSource::reorder_files` (file level) key only off the **leading**
sort column's `min` statistic. When the leading key ties across row
groups or files — the `ORDER BY low_cardinality_col, ts LIMIT k` shape —
the reorder degenerates to a no-op, the scan reads data in disk order,
and the TopK dynamic filter's threshold improves monotonically without
ever proving a later row group unwinnable. Result: zero runtime RG
pruning and a full-file decode even though per-column statistics contain
everything needed to read the best row group first.
The end-to-end test in this PR demonstrates the failure concretely on
current main: leading key tied across 5 row groups, secondary key
clustered but stored in adversarial DESC disk order →
`row_groups_pruned_dynamic_filter=0`, all 500 rows decoded. With this
change the scan reads the best RG first and prunes the other 4.
Note the *pruning* machinery itself already handles multi-column
lexicographic dynamic filters (the leading disjunct of `a < x OR (a = x
AND b < y)` prunes on `min(a)` stats — verified by the other new e2e
test, which passes on main unmodified). The gap was purely in the
*reorder* layers feeding it.
## What changes are included in this PR?
Both reorder layers now sort lexicographically over the **longest
plain-`Column` prefix** of the sort order, using per-column `min`
statistics:
- **RG level** (`access_plan.rs`): switch from `sort_to_indices` over
the leading column's mins to `lexsort_to_indices` over per-column min
arrays. The leading key stays normalized ASC (direction still applied by
the downstream `reverse()`, unchanged); secondary keys sort by their
direction *relative to* the leading key, with null placement pre-flipped
when the reverse is coming so the post-reverse order matches the
request. The prefix walk stops at the first non-`Column` /
not-in-file-schema expression (leading-key-only graceful skips
unchanged).
- **File level** (`sort.rs`): `extract_topk_sort_info` extended from the
single leading column to the plain-`Column` prefix, comparator compares
tuple-wise with per-column direction; missing-stats files still sort
last per column.
Behavior for single-column sort orders is unchanged.
## Are these changes tested?
- 4 new unit tests for `reorder_by_statistics`: secondary tie-break,
relative secondary direction (`a ASC, b DESC`), DESC/DESC normalization
through `reverse()`, and prefix-stop on a non-`Column` secondary
expression.
- 1 new unit test for file-level reorder: secondary tie-break with
missing-stats file placement.
- 2 new end-to-end tests in `dynamic_row_group_pruning.rs`:
- multi-column sort with clustered leading key → pruning fires via the
leading disjunct (passes before this PR; guards the existing behavior),
- multi-column sort with tied leading key + adversarial disk order →
pruning fires only with this PR's lexicographic reorder.
- Existing reorder unit tests, `dynamic_row_group_pruning` integration
tests, and `topk.slt` / `sort_pushdown.slt` /
`dynamic_row_group_pruning.slt` / `dynamic_filter_pushdown_config.slt`
all pass.
- `topk_tpch` (sort-tpch `-l 100`) shows no regressions (all 11 queries
within noise; lineitem's secondary keys are not clustered when the
leading key ties, so no wins expected on that suite — the win case is
the `ORDER BY low_card, ts` shape shown in the e2e test).
## Are there any user-facing changes?
No API changes. Queries with multi-column `ORDER BY ... LIMIT` over
parquet whose leading sort key ties across row groups / files can now
skip row groups at runtime instead of decoding the whole file.
---
.../tests/parquet/dynamic_row_group_pruning.rs | 152 +++++++++
datafusion/datasource-parquet/src/access_plan.rs | 345 ++++++++++++++++-----
datafusion/datasource-parquet/src/sort.rs | 95 ++++--
datafusion/datasource-parquet/src/source.rs | 63 ++++
4 files changed, 556 insertions(+), 99 deletions(-)
diff --git a/datafusion/core/tests/parquet/dynamic_row_group_pruning.rs
b/datafusion/core/tests/parquet/dynamic_row_group_pruning.rs
index b72c56ace5..d5d648be9b 100644
--- a/datafusion/core/tests/parquet/dynamic_row_group_pruning.rs
+++ b/datafusion/core/tests/parquet/dynamic_row_group_pruning.rs
@@ -433,3 +433,155 @@ async fn dynamic_rg_pruning_coexists_with_row_filter() {
output.description(),
);
}
+
+/// Build five two-column `RecordBatch`es: `a` is physically clustered
+/// (batch `i` carries `a ∈ [i*100, (i+1)*100)`, disjoint per-RG stats)
+/// and `b` is a per-batch shuffle (identical `[0, 100)` range in every
+/// RG, useless for pruning).
+fn build_two_col_leading_clustered(schema: &Arc<Schema>) -> Vec<RecordBatch> {
+ (0..5i64)
+ .map(|rg| {
+ let base = rg * 100;
+ let a: Vec<i64> = (base..base + 100).collect();
+ // pseudo-shuffled b, same value set in every RG
+ let b: Vec<i64> = (0..100).map(|i| (i * 37) % 100).collect();
+ RecordBatch::try_new(
+ Arc::clone(schema),
+ vec![
+ Arc::new(Int64Array::from(a)) as ArrayRef,
+ Arc::new(Int64Array::from(b)) as ArrayRef,
+ ],
+ )
+ .unwrap()
+ })
+ .collect()
+}
+
+/// Build five two-column `RecordBatch`es where the *leading* sort key
+/// ties everywhere (`a = 1` in every row / RG) and the *secondary* key
+/// is clustered but stored in DESC disk order: batch 0 carries
+/// `b ∈ [400, 500)`, batch 4 carries `b ∈ [0, 100)`.
+///
+/// An `ORDER BY a, b LIMIT k` query wants the rows in batch 4 first;
+/// reading disk order decodes every RG with a monotonically *improving*
+/// threshold that never proves a later RG unwinnable.
+fn build_two_col_leading_tied_desc(schema: &Arc<Schema>) -> Vec<RecordBatch> {
+ (0..5i64)
+ .map(|rg| {
+ let base = (4 - rg) * 100;
+ let a: Vec<i64> = vec![1; 100];
+ let b: Vec<i64> = (base..base + 100).collect();
+ RecordBatch::try_new(
+ Arc::clone(schema),
+ vec![
+ Arc::new(Int64Array::from(a)) as ArrayRef,
+ Arc::new(Int64Array::from(b)) as ArrayRef,
+ ],
+ )
+ .unwrap()
+ })
+ .collect()
+}
+
+fn two_col_schema() -> Arc<Schema> {
+ Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int64, false),
+ Field::new("b", DataType::Int64, false),
+ ]))
+}
+
+/// A multi-column `ORDER BY a, b LIMIT k` must still engage the runtime
+/// RG pruner through the *leading* disjunct of the lexicographic dynamic
+/// filter (`a < x OR (a = x AND b < y)`): once the heap fills from the
+/// first (best) row group, `min(a) > x` alone proves later RGs
+/// unwinnable regardless of `b`.
+#[tokio::test]
+async fn dynamic_rg_pruning_fires_for_multi_column_sort_leading_clustered() {
+ let schema = two_col_schema();
+ let batches = build_two_col_leading_clustered(&schema);
+
+ let mut ctx = ContextWithParquet::with_custom_data(
+ Scenario::Int,
+ RowGroup(100),
+ Arc::clone(&schema),
+ batches,
+ )
+ .await;
+
+ let output = ctx
+ .query("SELECT a, b FROM t ORDER BY a ASC, b ASC LIMIT 5")
+ .await;
+
+ assert_eq!(output.result_rows, 5, "query must return LIMIT rows");
+
+ let pruned = output
+ .row_groups_pruned_dynamic_filter()
+ .expect("`row_groups_pruned_dynamic_filter` metric must be
registered");
+ assert!(
+ pruned >= 1,
+ "multi-column TopK must prune via the leading column's disjunct; \
+ pruned={pruned}\n{}",
+ output.description(),
+ );
+}
+
+/// When the leading sort key ties across all row groups, pruning (and
+/// reading the right RG first) must fall to the *secondary* key: RG
+/// stats give `min(a) = max(a) = 1` everywhere, so the lex dynamic
+/// filter reduces to `a = 1 AND b < y` — prunable via `min(b)`.
+///
+/// The disk order is adversarial (secondary key DESC), so without
+/// multi-column stats reorder the scan reads the worst RG first and the
+/// threshold never proves later RGs unwinnable. With multi-column
+/// reorder the best RG is read first and every other RG is pruned.
+#[tokio::test]
+async fn dynamic_rg_pruning_fires_for_multi_column_sort_leading_tied() {
+ let schema = two_col_schema();
+ let batches = build_two_col_leading_tied_desc(&schema);
+
+ let mut ctx = ContextWithParquet::with_custom_data(
+ Scenario::Int,
+ RowGroup(100),
+ Arc::clone(&schema),
+ batches,
+ )
+ .await;
+
+ let output = ctx
+ .query("SELECT a, b FROM t ORDER BY a ASC, b ASC LIMIT 5")
+ .await;
+
+ assert_eq!(output.result_rows, 5, "query must return LIMIT rows");
+ // The leading key `a = 1` is tied everywhere, so correctness rests
+ // entirely on the secondary key: the five smallest `b` values must come
+ // back, in ascending secondary order. Assert the exact result rows
+ // (full two-column text, in order) rather than just probing for each
+ // `b` — a bare `| {b} ` match would be satisfied by the leading `a = 1`
+ // column even if that `b` were missing or misordered.
+ let formatted = output.pretty_results();
+ let data_rows: Vec<&str> = formatted
+ .lines()
+ .filter(|line| line.starts_with("| 1 |"))
+ .collect();
+ assert_eq!(
+ data_rows,
+ vec![
+ "| 1 | 0 |",
+ "| 1 | 1 |",
+ "| 1 | 2 |",
+ "| 1 | 3 |",
+ "| 1 | 4 |",
+ ],
+ "output must be exactly (a=1, b=0..=4) in ascending secondary order;
got:\n{formatted}",
+ );
+
+ let pruned = output
+ .row_groups_pruned_dynamic_filter()
+ .expect("`row_groups_pruned_dynamic_filter` metric must be
registered");
+ assert!(
+ pruned >= 1,
+ "with the leading key tied everywhere, the secondary key must \
+ drive RG reorder + pruning; pruned={pruned}\n{}",
+ output.description(),
+ );
+}
diff --git a/datafusion/datasource-parquet/src/access_plan.rs
b/datafusion/datasource-parquet/src/access_plan.rs
index 8189c2378c..1e9bae0ff6 100644
--- a/datafusion/datasource-parquet/src/access_plan.rs
+++ b/datafusion/datasource-parquet/src/access_plan.rs
@@ -606,13 +606,24 @@ impl PreparedAccessPlan {
/// Reorder row groups by their min statistics for the given sort order.
///
/// This helps TopK queries find optimal values first. Row groups are
- /// always sorted by min values in ASC order — direction (DESC) is
- /// handled separately by `reverse()` which is applied after reorder.
+ /// lexicographically sorted by per-column min values over the longest
+ /// prefix of the sort order made of plain columns present in the file
+ /// schema. The leading column is always sorted ASC by min — direction
+ /// (DESC) is handled separately by `reverse()` which is applied after
+ /// reorder. Subsequent columns sort by their direction *relative* to
+ /// the leading column (and their null placement is flipped when the
+ /// plan will be reversed), so that the post-`reverse()` order
+ /// approximates the requested lexicographic order.
+ ///
+ /// Secondary sort keys matter when the leading column's min ties
+ /// across row groups (e.g. `ORDER BY low_cardinality_col, ts LIMIT k`)
+ /// — without them the reorder is a no-op on such files and the TopK
+ /// dynamic filter converges only as fast as disk order allows.
///
/// Gracefully skips reordering when:
/// - There is a row_selection (too complex to remap)
/// - 0 or 1 row groups (nothing to reorder)
- /// - Sort expression is not a simple column reference
+ /// - The leading sort expression is not a simple column reference
/// - Statistics are unavailable
pub(crate) fn reorder_by_statistics(
mut self,
@@ -631,88 +642,116 @@ impl PreparedAccessPlan {
return Ok(self);
}
- let first_sort_expr = sort_order.first();
-
- // Extract column name from sort expression
- let column: &Column = match
first_sort_expr.expr.downcast_ref::<Column>() {
- Some(col) => col,
- None => {
- debug!("Skipping RG reorder: sort expr is not a simple
column");
- return Ok(self);
- }
- };
-
- // Expected graceful skip: the sort column lives outside the
- // file schema (e.g. a partition column whose ordering came
- // through `reversed_satisfies` rather than `column_in_file_schema`).
- // Parquet has no per-RG stats for it. Bail out quietly — no
- // `debug_assert!` because this is a normal pushdown shape.
- if arrow_schema.field_with_name(column.name()).is_err() {
- debug!(
- "Skipping RG reorder: column `{}` not in file schema",
- column.name()
- );
- return Ok(self);
- }
-
- // From here, any `StatisticsConverter` / stats read / sort
- // failure is unexpected — the column exists in the file
- // schema, so building the converter and pulling typed mins
- // should succeed on any well-formed parquet file. Trip a
- // `debug_assert!` so CI catches regressions, but stay graceful
- // in release so a single odd file can't take down a scan.
- let converter = match StatisticsConverter::try_new(
- column.name(),
- arrow_schema,
- file_metadata.file_metadata().schema_descr(),
- ) {
- Ok(c) => c,
- Err(e) => {
- debug_assert!(
- false,
- "RG reorder: cannot create stats converter for `{}`: {e}",
- column.name(),
- );
- return Ok(self);
- }
- };
-
- // Always sort ASC by min values — direction is handled by reverse
let rg_metadata: Vec<&RowGroupMetaData> = self
.row_group_indexes
.iter()
.map(|&idx| file_metadata.row_group(idx))
.collect();
- let stat_mins = match
converter.row_group_mins(rg_metadata.iter().copied()) {
- Ok(vals) => vals,
- Err(e) => {
- debug_assert!(
- false,
- "RG reorder: cannot get min values for `{}`: {e}",
- column.name(),
- );
- return Ok(self);
+ let leading_descending = sort_order.first().options.descending;
+
+ // Build one `SortColumn` of per-RG mins for each usable prefix
+ // column of the sort order. The walk stops at the first
+ // expression that isn't a plain `Column` in the file schema —
+ // stats for later columns can't refine the order once an
+ // unresolvable key sits between them and the resolved prefix.
+ let mut sort_columns: Vec<arrow::compute::SortColumn> = Vec::new();
+ for (i, sort_expr) in sort_order.iter().enumerate() {
+ let column: &Column = match
sort_expr.expr.downcast_ref::<Column>() {
+ Some(col) => col,
+ None => {
+ if i == 0 {
+ debug!("Skipping RG reorder: sort expr is not a simple
column");
+ return Ok(self);
+ }
+ break;
+ }
+ };
+
+ // Expected graceful skip: the sort column lives outside the
+ // file schema (e.g. a partition column whose ordering came
+ // through `reversed_satisfies` rather than
+ // `column_in_file_schema`). Parquet has no per-RG stats for
+ // it. Bail out quietly — no `debug_assert!` because this is
+ // a normal pushdown shape.
+ if arrow_schema.field_with_name(column.name()).is_err() {
+ if i == 0 {
+ debug!(
+ "Skipping RG reorder: column `{}` not in file schema",
+ column.name()
+ );
+ return Ok(self);
+ }
+ break;
}
- };
- let sort_options = arrow::compute::SortOptions {
- descending: false,
- nulls_first: first_sort_expr.options.nulls_first,
- };
- let sorted_indices =
- match arrow::compute::sort_to_indices(&stat_mins,
Some(sort_options), None) {
- Ok(indices) => indices,
+ // From here, any `StatisticsConverter` / stats read / sort
+ // failure is unexpected — the column exists in the file
+ // schema, so building the converter and pulling typed mins
+ // should succeed on any well-formed parquet file. Trip a
+ // `debug_assert!` so CI catches regressions, but stay graceful
+ // in release so a single odd file can't take down a scan.
+ let converter = match StatisticsConverter::try_new(
+ column.name(),
+ arrow_schema,
+ file_metadata.file_metadata().schema_descr(),
+ ) {
+ Ok(c) => c,
Err(e) => {
debug_assert!(
false,
- "RG reorder: arrow sort_to_indices failed for `{}`:
{e}",
+ "RG reorder: cannot create stats converter for `{}`:
{e}",
column.name(),
);
- return Ok(self);
+ if i == 0 {
+ return Ok(self);
+ }
+ break;
}
};
+ let stat_mins = match
converter.row_group_mins(rg_metadata.iter().copied()) {
+ Ok(vals) => vals,
+ Err(e) => {
+ debug_assert!(
+ false,
+ "RG reorder: cannot get min values for `{}`: {e}",
+ column.name(),
+ );
+ if i == 0 {
+ return Ok(self);
+ }
+ break;
+ }
+ };
+
+ // The plan is later `reverse()`d iff the leading column is
+ // DESC, which flips both value order and null placement of
+ // every column. Sort each column by its direction relative
+ // to the leading column (leading itself is therefore always
+ // ASC), and pre-flip null placement when the reverse is
+ // coming, so the post-reverse order matches the request.
+ // Nulls here are row groups with *missing stats*, so their
+ // placement is a heuristic, not a correctness matter.
+ let sort_options = arrow::compute::SortOptions {
+ descending: sort_expr.options.descending != leading_descending,
+ nulls_first: sort_expr.options.nulls_first !=
leading_descending,
+ };
+ sort_columns.push(arrow::compute::SortColumn {
+ values: stat_mins,
+ options: Some(sort_options),
+ });
+ }
+
+ let sorted_indices = match
arrow::compute::lexsort_to_indices(&sort_columns, None)
+ {
+ Ok(indices) => indices,
+ Err(e) => {
+ debug_assert!(false, "RG reorder: arrow lexsort_to_indices
failed: {e}");
+ return Ok(self);
+ }
+ };
+
// Apply the reordering
let original_indexes = self.row_group_indexes.clone();
self.row_group_indexes = sorted_indices
@@ -1223,4 +1262,172 @@ mod test {
assert_eq!(result.row_group_indexes, vec![0, 1]);
}
+
+ // ----------------------------------------------------------------
+ // multi-column `reorder_by_statistics` tests
+ // ----------------------------------------------------------------
+
+ /// Two-column int32 schema named "a", "b".
+ fn two_col_schema_descr() -> SchemaDescPtr {
+ use parquet::basic::Type as PhysicalType;
+ use parquet::schema::types::Type as SchemaType;
+ let fields = ["a", "b"]
+ .iter()
+ .map(|name| {
+ Arc::new(
+ SchemaType::primitive_type_builder(name,
PhysicalType::INT32)
+ .build()
+ .unwrap(),
+ )
+ })
+ .collect();
+ let schema = SchemaType::group_type_builder("schema")
+ .with_fields(fields)
+ .build()
+ .unwrap();
+ Arc::new(SchemaDescriptor::new(Arc::new(schema)))
+ }
+
+ /// Build a `ParquetMetaData` with one row group per element of
+ /// `mins`: `(min(a), min(b))` per row group, `min == max`.
+ fn parquet_metadata_with_two_col_mins(mins: &[(i32, i32)]) ->
ParquetMetaData {
+ let schema_descr = two_col_schema_descr();
+ let row_groups: Vec<RowGroupMetaData> = mins
+ .iter()
+ .map(|&(a, b)| {
+ let columns = [(0, a), (1, b)]
+ .iter()
+ .map(|&(col, m)| {
+ let stats = ParquetStatistics::int32(
+ Some(m),
+ Some(m),
+ None,
+ Some(0),
+ false,
+ );
+ ColumnChunkMetaData::builder(schema_descr.column(col))
+ .set_statistics(stats)
+ .set_num_values(100)
+ .build()
+ .unwrap()
+ })
+ .collect();
+ RowGroupMetaData::builder(schema_descr.clone())
+ .set_num_rows(100)
+ .set_column_metadata(columns)
+ .build()
+ .unwrap()
+ })
+ .collect();
+ let file_metadata =
+ FileMetaData::new(0, 0, None, None, schema_descr.clone(), None);
+ ParquetMetaData::new(file_metadata, row_groups)
+ }
+
+ fn arrow_schema_ab_int() -> Schema {
+ Schema::new(vec![
+ Field::new("a", DataType::Int32, true),
+ Field::new("b", DataType::Int32, true),
+ ])
+ }
+
+ fn sort_expr(name: &str, index: usize, descending: bool) ->
PhysicalSortExpr {
+ PhysicalSortExpr {
+ expr: Arc::new(Column::new(name, index)),
+ options: SortOptions {
+ descending,
+ nulls_first: true,
+ },
+ }
+ }
+
+ /// `ORDER BY a ASC, b ASC` with the leading key tied everywhere:
+ /// the secondary key must break the tie, so RGs order by `min(b)`.
+ #[test]
+ fn reorder_by_statistics_breaks_leading_ties_with_secondary_column() {
+ let metadata =
+ parquet_metadata_with_two_col_mins(&[(1, 300), (1, 100), (1,
200)]);
+ let plan = PreparedAccessPlan::new(vec![0, 1, 2], None).unwrap();
+ let order =
+ LexOrdering::new(vec![sort_expr("a", 0, false), sort_expr("b", 1,
false)])
+ .unwrap();
+
+ let result = plan
+ .reorder_by_statistics(&order, &metadata, &arrow_schema_ab_int())
+ .unwrap();
+
+ assert_eq!(result.row_group_indexes, vec![1, 2, 0]);
+ }
+
+ /// `ORDER BY a ASC, b DESC`: the secondary key's direction is
+ /// honored relative to the leading key, so ties on `min(a)` order
+ /// by `min(b)` DESC.
+ #[test]
+ fn reorder_by_statistics_honors_secondary_direction() {
+ let metadata =
+ parquet_metadata_with_two_col_mins(&[(1, 100), (1, 300), (0,
500)]);
+ let plan = PreparedAccessPlan::new(vec![0, 1, 2], None).unwrap();
+ let order =
+ LexOrdering::new(vec![sort_expr("a", 0, false), sort_expr("b", 1,
true)])
+ .unwrap();
+
+ let result = plan
+ .reorder_by_statistics(&order, &metadata, &arrow_schema_ab_int())
+ .unwrap();
+
+ // a=0 first, then the two a=1 groups by b DESC: 300 before 100.
+ assert_eq!(result.row_group_indexes, vec![2, 1, 0]);
+ }
+
+ /// `ORDER BY a DESC, b DESC` is normalized to ASC lexsort here and
+ /// flipped by the later `reverse()`: both keys sort ASC relative to
+ /// the leading direction, so reversing yields `(a DESC, b DESC)`.
+ #[test]
+ fn reorder_by_statistics_normalizes_desc_desc_for_reverse() {
+ let metadata =
+ parquet_metadata_with_two_col_mins(&[(1, 300), (2, 100), (1,
100)]);
+ let plan = PreparedAccessPlan::new(vec![0, 1, 2], None).unwrap();
+ let order =
+ LexOrdering::new(vec![sort_expr("a", 0, true), sort_expr("b", 1,
true)])
+ .unwrap();
+
+ let result = plan
+ .reorder_by_statistics(&order, &metadata, &arrow_schema_ab_int())
+ .unwrap();
+
+ // ASC lexsort of (a, b): (1,100) < (1,300) < (2,100); the later
+ // reverse() produces (2,100), (1,300), (1,100) = (a DESC, b DESC).
+ assert_eq!(result.row_group_indexes, vec![2, 0, 1]);
+ }
+
+ /// A non-`Column` *secondary* expression stops the stats walk but
+ /// keeps the leading column's reorder (prefix semantics).
+ #[test]
+ fn reorder_by_statistics_keeps_leading_prefix_on_non_column_secondary() {
+ let metadata =
+ parquet_metadata_with_two_col_mins(&[(5, 300), (3, 100), (4,
200)]);
+ let plan = PreparedAccessPlan::new(vec![0, 1, 2], None).unwrap();
+ let order = LexOrdering::new(vec![
+ sort_expr("a", 0, false),
+ PhysicalSortExpr {
+ expr: Arc::new(BinaryExpr::new(
+ Arc::new(Column::new("b", 1)),
+ Operator::Plus,
+ lit(1i32),
+ )),
+ options: SortOptions {
+ descending: false,
+ nulls_first: true,
+ },
+ },
+ ])
+ .unwrap();
+
+ let result = plan
+ .reorder_by_statistics(&order, &metadata, &arrow_schema_ab_int())
+ .unwrap();
+
+ // Ordered by min(a) ASC only: 3, 4, 5.
+ assert_eq!(result.row_group_indexes, vec![1, 2, 0]);
+ }
}
diff --git a/datafusion/datasource-parquet/src/sort.rs
b/datafusion/datasource-parquet/src/sort.rs
index c1cf4e8b78..ea33fb0e2e 100644
--- a/datafusion/datasource-parquet/src/sort.rs
+++ b/datafusion/datasource-parquet/src/sort.rs
@@ -124,9 +124,14 @@ pub fn reverse_row_selection(
/// Reorder a file list so the most "promising" files are read first,
/// matching `PreparedAccessPlan::reorder_by_statistics` at the
-/// row-group level: key off the file's `min(col)`, and let the sort
-/// direction follow the request (ASC by `min` for ASC requests, DESC
-/// by `min` for DESC requests).
+/// row-group level: key lexicographically off the file's per-column
+/// `min` for the longest plain-`Column` prefix of the sort order, and
+/// let the leading sort direction follow the request (ASC by `min`
+/// for ASC requests, DESC by `min` for DESC requests).
+///
+/// Secondary sort keys break ties when the leading column's `min` is
+/// equal across files (e.g. `ORDER BY low_cardinality_col, ts LIMIT k`),
+/// mirroring the row-group level lexicographic reorder.
///
/// Keeping both layers consistent matters because they share the same
/// convergence story for TopK's dynamic filter: file `i`'s `min` is a
@@ -147,51 +152,81 @@ pub(crate) fn reorder_files_by_min_statistics(
reverse_row_groups: bool,
table_schema: &Schema,
) -> Vec<PartitionedFile> {
- let Some((col_name, descending)) =
- extract_topk_sort_info(sort_order, reverse_row_groups)
- else {
+ let sort_keys = extract_topk_sort_info(sort_order, reverse_row_groups);
+ if sort_keys.is_empty() {
return files;
- };
+ }
- let Ok(col_idx) = table_schema.index_of(&col_name) else {
- return files;
- };
+ // Resolve names to column indexes; the leading key is required, later
+ // keys are best-effort (stop at the first unresolvable one).
+ let mut keys: Vec<(usize, bool)> = Vec::with_capacity(sort_keys.len());
+ for (col_name, descending) in &sort_keys {
+ match table_schema.index_of(col_name) {
+ Ok(idx) => keys.push((idx, *descending)),
+ Err(_) if keys.is_empty() => return files,
+ Err(_) => break,
+ }
+ }
files.sort_by(|a, b| {
- let key_a = file_min_value(a, col_idx);
- let key_b = file_min_value(b, col_idx);
- match (key_a, key_b) {
- (Some(va), Some(vb)) => {
- let cmp =
va.partial_cmp(&vb).unwrap_or(std::cmp::Ordering::Equal);
- if descending { cmp.reverse() } else { cmp }
+ for &(col_idx, descending) in &keys {
+ let key_a = file_min_value(a, col_idx);
+ let key_b = file_min_value(b, col_idx);
+ let ord = match (key_a, key_b) {
+ (Some(va), Some(vb)) => {
+ let cmp =
va.partial_cmp(&vb).unwrap_or(std::cmp::Ordering::Equal);
+ if descending { cmp.reverse() } else { cmp }
+ }
+ // Missing stats always sort last, regardless of direction.
+ (Some(_), None) => std::cmp::Ordering::Less,
+ (None, Some(_)) => std::cmp::Ordering::Greater,
+ (None, None) => std::cmp::Ordering::Equal,
+ };
+ if ord != std::cmp::Ordering::Equal {
+ return ord;
}
- (Some(_), None) => std::cmp::Ordering::Less,
- (None, Some(_)) => std::cmp::Ordering::Greater,
- (None, None) => std::cmp::Ordering::Equal,
}
+ std::cmp::Ordering::Equal
});
log::debug!(
- "Reordered {} files by min({}) {} for TopK optimization",
+ "Reordered {} files by lexicographic min of {:?} for TopK
optimization",
files.len(),
- col_name,
- if descending { "DESC" } else { "ASC" }
+ sort_keys,
);
files
}
-/// Extract the `(column name, descending)` tuple used by file-level
-/// reordering. Returns `None` when the sort order isn't set or the
-/// leading sort expression isn't a plain `Column`.
+/// Extract the `(column name, descending)` keys used by file-level
+/// reordering: the longest prefix of the sort order made of plain
+/// `Column` expressions. Returns an empty vec when the sort order isn't
+/// set or the leading sort expression isn't a plain `Column`.
+///
+/// The leading key's direction is `reverse_row_groups` (the pushdown's
+/// authoritative flip decision, which may differ from the raw
+/// expression's `descending` in the `reversed_satisfies` case);
+/// subsequent keys apply their direction *relative to the leading
+/// expression* on top of that flag, so a request like
+/// `[a DESC, b ASC]` with `reverse_row_groups=true` sorts by
+/// `(min(a) DESC, min(b) ASC)`.
fn extract_topk_sort_info(
sort_order: Option<&LexOrdering>,
reverse_row_groups: bool,
-) -> Option<(String, bool)> {
- let sort_order = sort_order?;
- let first = sort_order.first();
- let col = first.expr.downcast_ref::<Column>()?;
- Some((col.name().to_string(), reverse_row_groups))
+) -> Vec<(String, bool)> {
+ let Some(sort_order) = sort_order else {
+ return vec![];
+ };
+ let leading_descending = sort_order.first().options.descending;
+ let mut keys = Vec::new();
+ for sort_expr in sort_order.iter() {
+ let Some(col) = sort_expr.expr.downcast_ref::<Column>() else {
+ break;
+ };
+ let relative_desc = sort_expr.options.descending != leading_descending;
+ keys.push((col.name().to_string(), reverse_row_groups !=
relative_desc));
+ }
+ keys
}
/// File's per-column `min` for the reorder key.
diff --git a/datafusion/datasource-parquet/src/source.rs
b/datafusion/datasource-parquet/src/source.rs
index 3443b08475..3e62023767 100644
--- a/datafusion/datasource-parquet/src/source.rs
+++ b/datafusion/datasource-parquet/src/source.rs
@@ -1756,6 +1756,69 @@ mod tests {
assert_eq!(names(&reordered), vec!["has_min", "no_stats"]);
}
+ /// Multi-column TopK: when the leading column's `min` ties across
+ /// files, the secondary sort key breaks the tie (lexicographic,
+ /// mirroring the row-group level reorder).
+ #[test]
+ fn reorder_files_breaks_leading_ties_with_secondary_column() {
+ use datafusion_common::stats::Precision;
+ use datafusion_common::{ColumnStatistics, ScalarValue, Statistics};
+ use datafusion_datasource::PartitionedFile;
+ use pushdown_sort_helpers::*;
+ use reorder_files_helpers::*;
+
+ fn file_with_two_mins(
+ name: &str,
+ min_a: i32,
+ min_b: Option<i32>,
+ ) -> PartitionedFile {
+ let mut pf = PartitionedFile::new(name.to_string(), 0);
+ let col = |min: Option<i32>| ColumnStatistics {
+ null_count: Precision::Absent,
+ max_value: Precision::Absent,
+ min_value: min
+ .map(|v| Precision::Exact(ScalarValue::Int32(Some(v))))
+ .unwrap_or(Precision::Absent),
+ sum_value: Precision::Absent,
+ distinct_count: Precision::Absent,
+ byte_size: Precision::Absent,
+ };
+ pf.statistics = Some(Arc::new(Statistics {
+ num_rows: Precision::Absent,
+ total_byte_size: Precision::Absent,
+ column_statistics: vec![col(Some(min_a)), col(min_b)],
+ }));
+ pf
+ }
+
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int32, true),
+ Field::new("b", DataType::Int32, true),
+ ]));
+ let mut source = ParquetSource::new(Arc::clone(&schema));
+ source.sort_order_for_reorder = Some(
+ LexOrdering::new(vec![
+ sort_expr_on(&schema, "a", false),
+ sort_expr_on(&schema, "b", false),
+ ])
+ .unwrap(),
+ );
+
+ let reordered = source.reorder_files(vec![
+ file_with_two_mins("tie_late", 1, Some(300)),
+ file_with_two_mins("first", 0, Some(999)),
+ file_with_two_mins("tie_early", 1, Some(100)),
+ file_with_two_mins("tie_no_b_stats", 1, None),
+ ]);
+
+ // `first` wins on the leading key; the `a = 1` ties order by
+ // `min(b)` ASC with missing-`b`-stats last.
+ assert_eq!(
+ names(&reordered),
+ vec!["first", "tie_early", "tie_late", "tie_no_b_stats"]
+ );
+ }
+
/// When no sort pushdown has fired (`sort_order_for_reorder` is
/// `None`), `reorder_files` is a no-op and preserves input order.
#[test]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
