Skip to content

Commit 8332cfa

Browse files
NGA-TRANcursoragent
NGA-TRAN
and
cursoragent
authored
test: add range-partitioned sorted time-bin aggregation coverage (#24500)
## Which issue does this PR close? - Related to #23569 (skip hash shuffle for `GROUP BY key, date_bin(timestamp)` on `Range([timestamp])`). - Related to #24438. This PR does not close those issues; it only pins today's plan. A follow-up will update the expected plan once the shuffle is removed. ## Rationale for this change Adds a table that matches a common metrics layout: - columns `key, zone, host, pod, service, timestamp, value` - range-partitioned on `timestamp` into two 60-minute file groups - sorted within each file on `(key, timestamp)` and pins today's physical plan for: ```sql SELECT key, date_bin(INTERVAL '60 seconds', timestamp) AS time_bin, sum(value) FROM metrics_range_sorted WHERE service = 'a' GROUP BY key, time_bin ``` The scan already advertises `Range([timestamp])`, `output_ordering=[key, timestamp]`, and two parallel file groups. Time bins do not straddle the hour split, so `(key, time_bin)` is partition-disjoint. **Today DataFusion still inserts a hash `RepartitionExec` between Partial and Final**, even though `ordering_mode=Sorted` is already recognized. This PR documents that plan so a follow-up can remove the shuffle and the expected plan will shrink to a single `SinglePartitioned` aggregate. ## What changes are included in this PR? - Register `metrics_range_sorted` for sqllogictests. - EXPLAIN + result coverage for the filtered and unfiltered queries. ## Are these changes tested? This PR is tests only. ## Are there any user-facing changes? No. ## Test plan - [x] `cargo test --test sqllogictests -- range_sorted_time_bin_agg.slt` - [ ] CI sqllogictests --------- Co-authored-by: Cursor <[email protected]>
1 parent eec2dff commit 8332cfa

3 files changed

Lines changed: 345 additions & 12 deletions

File tree

‎datafusion/sqllogictest/src/test_context.rs‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,9 @@ use datafusion::{
5454
use datafusion_spark::SessionStateBuilderSpark;
5555

5656
use crate::is_spark_path;
57-
use range_partitioning::register_range_partitioned_table;
57+
use range_partitioning::{
58+
register_range_partitioned_table, register_range_sorted_time_bin_table,
59+
};
5860

5961
use async_trait::async_trait;
6062
use datafusion::common::cast::as_float64_array;
@@ -179,6 +181,10 @@ impl TestContext {
179181
info!("Registering range partitioned table");
180182
register_range_partitioned_table(test_ctx.session_ctx());
181183
}
184+
"range_sorted_time_bin_agg.slt" => {
185+
info!("Registering range-sorted time-bin table");
186+
register_range_sorted_time_bin_table(test_ctx.session_ctx());
187+
}
182188
"metadata.slt" | "arrow_field.slt" => {
183189
info!("Registering metadata table tables");
184190
register_metadata_tables(test_ctx.session_ctx());

‎datafusion/sqllogictest/src/test_context/range_partitioning.rs‎

Lines changed: 155 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -19,17 +19,19 @@ use std::fs::{File, create_dir_all, remove_dir_all};
1919
use std::path::Path;
2020
use std::sync::Arc;
2121

22-
use arrow::array::{ArrayRef, Int32Array};
22+
use arrow::array::{
23+
ArrayRef, Int32Array, Int64Array, StringArray, TimestampNanosecondArray,
24+
};
2325
use arrow::compute::SortOptions;
24-
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
26+
use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
2527
use arrow::record_batch::RecordBatch;
2628
use datafusion::catalog::streaming::StreamingTable;
2729
use datafusion::common::{ScalarValue, SplitPoint};
2830
use datafusion::datasource::file_format::parquet::ParquetFormat;
2931
use datafusion::datasource::listing::{
3032
ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl,
3133
};
32-
use datafusion::logical_expr::{Partitioning, RangePartitioning, col};
34+
use datafusion::logical_expr::{Partitioning, RangePartitioning, SortExpr, col};
3335
use datafusion::parquet::arrow::ArrowWriter;
3436
use datafusion::physical_expr::{
3537
Partitioning as PhysicalPartitioning, PhysicalSortExpr,
@@ -95,8 +97,9 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) {
9597
"range_partitioned",
9698
&range_table_dir,
9799
Arc::clone(&schema),
98-
RANGE_PARTITIONS,
100+
range_batches(&schema, RANGE_PARTITIONS),
99101
output_partitioning,
102+
None,
100103
);
101104

102105
register_unbounded_range_stream_table(
@@ -132,8 +135,9 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) {
132135
Path::new(env!("CARGO_MANIFEST_DIR"))
133136
.join("test_files/scratch_range_partitioning/range_partitioned_shifted"),
134137
Arc::clone(&schema),
135-
SHIFTED_RANGE_PARTITIONS,
138+
range_batches(&schema, SHIFTED_RANGE_PARTITIONS),
136139
shifted_output_partitioning,
140+
None,
137141
);
138142

139143
// Same rows as `range_partitioned` but split into only three range
@@ -156,8 +160,9 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) {
156160
Path::new(env!("CARGO_MANIFEST_DIR"))
157161
.join("test_files/scratch_range_partitioning/range_partitioned_narrow"),
158162
Arc::clone(&schema),
159-
NARROW_RANGE_PARTITIONS,
163+
range_batches(&schema, NARROW_RANGE_PARTITIONS),
160164
narrow_output_partitioning,
165+
None,
161166
);
162167

163168
let sparse_output_partitioning = Partitioning::Range(
@@ -178,8 +183,9 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) {
178183
Path::new(env!("CARGO_MANIFEST_DIR"))
179184
.join("test_files/scratch_range_partitioning/range_partitioned_sparse"),
180185
Arc::clone(&schema),
181-
SPARSE_RANGE_PARTITIONS,
186+
range_batches(&schema, SPARSE_RANGE_PARTITIONS),
182187
sparse_output_partitioning,
188+
None,
183189
);
184190
}
185191

@@ -188,16 +194,16 @@ fn register_parquet_listing_table(
188194
name: &str,
189195
table_dir: impl AsRef<Path>,
190196
schema: SchemaRef,
191-
partitions: impl IntoIterator<Item = &'static [(i32, i32, i32)]>,
197+
batches: Vec<RecordBatch>,
192198
output_partitioning: Partitioning,
199+
file_sort_order: Option<Vec<Vec<SortExpr>>>,
193200
) {
194201
let table_dir = table_dir.as_ref();
195202
if table_dir.exists() {
196203
remove_dir_all(table_dir).expect("test table dir should be removable");
197204
}
198205
create_dir_all(table_dir).expect("test table dir should be created");
199-
for (idx, rows) in partitions.into_iter().enumerate() {
200-
let batch = range_batch(Arc::clone(&schema), rows);
206+
for (idx, batch) in batches.into_iter().enumerate() {
201207
let file = File::create(table_dir.join(format!("part-{idx}.parquet")))
202208
.expect("test table parquet partition should be created");
203209
let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), None)
@@ -218,8 +224,11 @@ fn register_parquet_listing_table(
218224
);
219225
let table_url =
220226
ListingTableUrl::parse(&table_path).expect("test table url should parse");
221-
let options = ListingOptions::new(Arc::new(ParquetFormat::default()))
227+
let mut options = ListingOptions::new(Arc::new(ParquetFormat::default()))
222228
.with_output_partitioning(Some(output_partitioning));
229+
if let Some(file_sort_order) = file_sort_order {
230+
options = options.with_file_sort_order(file_sort_order);
231+
}
223232
let config = ListingTableConfig::new(table_url)
224233
.with_listing_options(options)
225234
.with_schema(schema);
@@ -278,6 +287,16 @@ fn range_stream_partition(
278287
)]))
279288
}
280289

290+
fn range_batches(
291+
schema: &SchemaRef,
292+
partitions: impl IntoIterator<Item = &'static [(i32, i32, i32)]>,
293+
) -> Vec<RecordBatch> {
294+
partitions
295+
.into_iter()
296+
.map(|rows| range_batch(Arc::clone(schema), rows))
297+
.collect()
298+
}
299+
281300
fn range_batch(schema: SchemaRef, rows: &[(i32, i32, i32)]) -> RecordBatch {
282301
RecordBatch::try_new(
283302
schema,
@@ -292,3 +311,128 @@ fn range_batch(schema: SchemaRef, rows: &[(i32, i32, i32)]) -> RecordBatch {
292311
)
293312
.expect("range batch should be valid")
294313
}
314+
315+
// ==============================================================================
316+
// Time-bin table: range-partitioned on timestamp, sorted on (key, timestamp)
317+
// ==============================================================================
318+
319+
/// Unix nanoseconds for `2024-01-01 00:00:00 UTC`.
320+
const TIME_BIN_EPOCH_NS: i64 = 1_704_067_200_000_000_000;
321+
const NANOS_PER_SECOND: i64 = 1_000_000_000;
322+
const NANOS_PER_MINUTE: i64 = 60 * NANOS_PER_SECOND;
323+
324+
/// Timestamp helper: minutes and seconds after `2024-01-01 00:00:00 UTC`.
325+
fn time_bin_ts(minutes: i64, seconds: i64) -> i64 {
326+
TIME_BIN_EPOCH_NS + minutes * NANOS_PER_MINUTE + seconds * NANOS_PER_SECOND
327+
}
328+
329+
/// Row: (key, col1, col2, col3, col4, timestamp_ns, value)
330+
type TimeBinRow = (
331+
&'static str,
332+
&'static str,
333+
&'static str,
334+
&'static str,
335+
&'static str,
336+
i64,
337+
i64,
338+
);
339+
340+
/// Registers `range_sorted_time_bin` for time-bin aggregation plan tests.
341+
///
342+
/// Two file groups, each covering a 60-minute timestamp range:
343+
/// - partition 0: `[2024-01-01 00:00, 01:00)`
344+
/// - partition 1: `[2024-01-01 01:00, 02:00)`
345+
///
346+
/// Files are range-partitioned on `timestamp` and sorted on `(key, timestamp)`.
347+
/// Because `date_bin(60 seconds, timestamp)` does not straddle the hour split,
348+
/// grouping by `(key, time_bin)` is partition-disjoint. Today's planner still
349+
/// inserts a hash shuffle; the test pins that plan so a follow-up can remove it.
350+
pub(super) fn register_range_sorted_time_bin_table(ctx: &SessionContext) {
351+
let schema = Arc::new(Schema::new(vec![
352+
Field::new("key", DataType::Utf8, false),
353+
Field::new("col1", DataType::Utf8, false),
354+
Field::new("col2", DataType::Utf8, false),
355+
Field::new("col3", DataType::Utf8, false),
356+
Field::new("col4", DataType::Utf8, false),
357+
Field::new(
358+
"timestamp",
359+
DataType::Timestamp(TimeUnit::Nanosecond, None),
360+
false,
361+
),
362+
Field::new("value", DataType::Int64, false),
363+
]));
364+
365+
// Each partition covers 60 minutes. The split is aligned to the 60-second
366+
// `date_bin` used by the test query, so time bins do not straddle files.
367+
let hour_split = time_bin_ts(60, 0);
368+
let output_partitioning = Partitioning::Range(
369+
RangePartitioning::try_new(
370+
vec![col("timestamp").sort(true, true)],
371+
vec![SplitPoint::new(vec![ScalarValue::TimestampNanosecond(
372+
Some(hour_split),
373+
None,
374+
)])],
375+
)
376+
.expect("time-bin range partitioning should be valid"),
377+
);
378+
379+
// Within each 60-minute file, rows are sorted by (key, timestamp).
380+
let partitions = [
381+
vec![
382+
("k1", "x", "y", "z", "a", time_bin_ts(0, 10), 1),
383+
("k1", "x", "y", "z", "a", time_bin_ts(0, 40), 2),
384+
("k1", "x", "y", "z", "b", time_bin_ts(1, 10), 99),
385+
("k2", "x", "y", "z", "a", time_bin_ts(30, 0), 3),
386+
("k2", "x", "y", "z", "a", time_bin_ts(30, 30), 4),
387+
],
388+
vec![
389+
("k1", "x", "y", "z", "a", time_bin_ts(60, 10), 10),
390+
("k1", "x", "y", "z", "a", time_bin_ts(60, 40), 20),
391+
("k2", "x", "y", "z", "a", time_bin_ts(90, 0), 30),
392+
("k2", "x", "y", "z", "a", time_bin_ts(105, 0), 5),
393+
],
394+
];
395+
396+
let table_dir = Path::new(env!("CARGO_MANIFEST_DIR"))
397+
.join("test_files/scratch_range_partitioning/range_sorted_time_bin");
398+
let batches = partitions
399+
.iter()
400+
.map(|rows| time_bin_batch(Arc::clone(&schema), rows))
401+
.collect();
402+
register_parquet_listing_table(
403+
ctx,
404+
"range_sorted_time_bin",
405+
&table_dir,
406+
Arc::clone(&schema),
407+
batches,
408+
output_partitioning,
409+
Some(vec![vec![
410+
col("key").sort(true, true),
411+
col("timestamp").sort(true, true),
412+
]]),
413+
);
414+
}
415+
416+
fn time_bin_batch(schema: SchemaRef, rows: &[TimeBinRow]) -> RecordBatch {
417+
RecordBatch::try_new(
418+
schema,
419+
vec![
420+
Arc::new(StringArray::from_iter_values(rows.iter().map(|row| row.0)))
421+
as ArrayRef,
422+
Arc::new(StringArray::from_iter_values(rows.iter().map(|row| row.1)))
423+
as ArrayRef,
424+
Arc::new(StringArray::from_iter_values(rows.iter().map(|row| row.2)))
425+
as ArrayRef,
426+
Arc::new(StringArray::from_iter_values(rows.iter().map(|row| row.3)))
427+
as ArrayRef,
428+
Arc::new(StringArray::from_iter_values(rows.iter().map(|row| row.4)))
429+
as ArrayRef,
430+
Arc::new(TimestampNanosecondArray::from_iter_values(
431+
rows.iter().map(|row| row.5),
432+
)) as ArrayRef,
433+
Arc::new(Int64Array::from_iter_values(rows.iter().map(|row| row.6)))
434+
as ArrayRef,
435+
],
436+
)
437+
.expect("time-bin batch should be valid")
438+
}

0 commit comments

Comments
 (0)