This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23888-f248f451bcc02ca5d99502e0bcf1fa3ee20ca448 in repository https://gitbox.apache.org/repos/asf/datafusion.git
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]
