zhuqi-lucas commented on code in PR #23888:
URL: https://github.com/apache/datafusion/pull/23888#discussion_r3699221521
##########
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<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");
+ let formatted = output.pretty_results();
+ for b in [0i64, 1, 2, 3, 4] {
+ assert!(
+ formatted.contains(&format!("| {b} ")),
Review Comment:
Good catch, thanks @kosiew! You're right that `| {b} ` could be satisfied by
the leading `a = 1` column. Updated it to assert the exact result rows — the
full two-column text (`| 1 | 0 |` … `| 1 | 4 |`) in ascending secondary order —
so it now validates the secondary key values and their ordering specifically.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
