From 118730667dc54ebe8d7936070197c850beaa8eee Mon Sep 17 00:00:00 2001 From: Qi Zhu <821684824@qq.com> Date: Sat, 25 Jul 2026 23:39:11 +0800 Subject: [PATCH] feat(parquet): multi-column lexicographic stats reorder for TopK sort pushdown MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements #22198 P0: PreparedAccessPlan::reorder_by_statistics and ParquetSource::reorder_files previously keyed only off the leading sort column's min statistic. When the leading key ties across row groups or files (e.g. ORDER BY low_cardinality_col, ts LIMIT k), the reorder was a no-op, the scan read row groups in disk order, and the TopK dynamic filter never tightened enough for the runtime RG pruner to fire. Both layers now sort lexicographically over the longest plain-Column prefix of the sort order, by per-column min with per-column directions expressed relative to the leading key (the leading key stays normalized ASC at the RG level; direction is applied by the later reverse() as before). The included end-to-end test shows the effect: leading key tied across 5 row groups, secondary key clustered but stored in adversarial DESC disk order — before this change row_groups_pruned_dynamic_filter=0 and all 500 rows decode; after it, the scan reads the best RG first and prunes the other 4. --- .../parquet/dynamic_row_group_pruning.rs | 137 +++++++ .../datasource-parquet/src/access_plan.rs | 345 ++++++++++++++---- datafusion/datasource-parquet/src/sort.rs | 95 +++-- datafusion/datasource-parquet/src/source.rs | 63 ++++ 4 files changed, 541 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 b72c56ace5acd..38fa41bb374b5 100644 --- a/datafusion/core/tests/parquet/dynamic_row_group_pruning.rs +++ b/datafusion/core/tests/parquet/dynamic_row_group_pruning.rs @@ -433,3 +433,140 @@ 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) -> Vec { + (0..5i64) + .map(|rg| { + let base = rg * 100; + let a: Vec = (base..base + 100).collect(); + // pseudo-shuffled b, same value set in every RG + let b: Vec = (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) -> Vec { + (0..5i64) + .map(|rg| { + let base = (4 - rg) * 100; + let a: Vec = vec![1; 100]; + let b: Vec = (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 { + 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"); + let formatted = output.pretty_results(); + for b in [0i64, 1, 2, 3, 4] { + assert!( + formatted.contains(&format!("| {b} ")), + "output must contain b={b}; 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 8189c2378cece..1e9bae0ff6ba3 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::() { - 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 = Vec::new(); + for (i, sort_expr) in sort_order.iter().enumerate() { + let column: &Column = match sort_expr.expr.downcast_ref::() { + 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 = 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 c1cf4e8b7824e..ea33fb0e2ecb2 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 { - 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::()?; - 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::() 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 3443b08475e0d..3e620237679fc 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, + ) -> PartitionedFile { + let mut pf = PartitionedFile::new(name.to_string(), 0); + let col = |min: Option| 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]