From 35792e58fc74cd07cb06d61e21c0696bfb7bc9ee Mon Sep 17 00:00:00 2001 From: Nga Tran Date: Wed, 19 Aug 2026 15:29:09 -0400 Subject: [PATCH 1/3] test: add range-partitioned sorted time-bin aggregation coverage Pin today's Partial + hash RepartitionExec + Final plan for GROUP BY key, date_bin(timestamp) on a table that is already Range([timestamp]) and sorted on (key, timestamp), so a follow-up can remove the shuffle. Co-authored-by: Cursor --- datafusion/sqllogictest/src/test_context.rs | 8 +- .../src/test_context/range_partitioning.rs | 178 ++++++++++++++++- .../test_files/range_sorted_time_bin_agg.slt | 183 ++++++++++++++++++ 3 files changed, 365 insertions(+), 4 deletions(-) create mode 100644 datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index 39aa2b09e5685..7252a0936afea 100644 --- a/datafusion/sqllogictest/src/test_context.rs +++ b/datafusion/sqllogictest/src/test_context.rs @@ -54,7 +54,9 @@ use datafusion::{ use datafusion_spark::SessionStateBuilderSpark; use crate::is_spark_path; -use range_partitioning::register_range_partitioned_table; +use range_partitioning::{ + register_metrics_range_sorted_table, register_range_partitioned_table, +}; use async_trait::async_trait; use datafusion::common::cast::as_float64_array; @@ -179,6 +181,10 @@ impl TestContext { info!("Registering range partitioned table"); register_range_partitioned_table(test_ctx.session_ctx()); } + "range_sorted_time_bin_agg.slt" => { + info!("Registering range-sorted metrics table"); + register_metrics_range_sorted_table(test_ctx.session_ctx()); + } "metadata.slt" | "arrow_field.slt" => { info!("Registering metadata table tables"); register_metadata_tables(test_ctx.session_ctx()); diff --git a/datafusion/sqllogictest/src/test_context/range_partitioning.rs b/datafusion/sqllogictest/src/test_context/range_partitioning.rs index 3cde3939f0b7c..de4f5875c8c93 100644 --- a/datafusion/sqllogictest/src/test_context/range_partitioning.rs +++ b/datafusion/sqllogictest/src/test_context/range_partitioning.rs @@ -19,9 +19,11 @@ use std::fs::{File, create_dir_all, remove_dir_all}; use std::path::Path; use std::sync::Arc; -use arrow::array::{ArrayRef, Int32Array}; +use arrow::array::{ + ArrayRef, Int32Array, Int64Array, StringArray, TimestampNanosecondArray, +}; use arrow::compute::SortOptions; -use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; use arrow::record_batch::RecordBatch; use datafusion::catalog::streaming::StreamingTable; use datafusion::common::{ScalarValue, SplitPoint}; @@ -29,7 +31,7 @@ use datafusion::datasource::file_format::parquet::ParquetFormat; use datafusion::datasource::listing::{ ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, }; -use datafusion::logical_expr::{Partitioning, RangePartitioning, col}; +use datafusion::logical_expr::{Partitioning, RangePartitioning, SortExpr, col}; use datafusion::parquet::arrow::ArrowWriter; use datafusion::physical_expr::{ Partitioning as PhysicalPartitioning, PhysicalSortExpr, @@ -292,3 +294,173 @@ fn range_batch(schema: SchemaRef, rows: &[(i32, i32, i32)]) -> RecordBatch { ) .expect("range batch should be valid") } + +// ============================================================================== +// Metrics table: range-partitioned on timestamp, sorted on (key, timestamp) +// ============================================================================== + +/// Unix nanoseconds for `2024-01-01 00:00:00 UTC`. +const METRICS_EPOCH_NS: i64 = 1_704_067_200_000_000_000; +const NANOS_PER_SECOND: i64 = 1_000_000_000; +const NANOS_PER_MINUTE: i64 = 60 * NANOS_PER_SECOND; + +/// Timestamp helper: minutes and seconds after `2024-01-01 00:00:00 UTC`. +fn metrics_ts(minutes: i64, seconds: i64) -> i64 { + METRICS_EPOCH_NS + minutes * NANOS_PER_MINUTE + seconds * NANOS_PER_SECOND +} + +/// Row: (key, zone, host, pod, service, timestamp_ns, value) +type MetricsRow = ( + &'static str, + &'static str, + &'static str, + &'static str, + &'static str, + i64, + i64, +); + +/// Registers `metrics_range_sorted` for time-bin aggregation plan tests. +/// +/// Two file groups, each covering a 60-minute timestamp range: +/// - partition 0: `[2024-01-01 00:00, 01:00)` +/// - partition 1: `[2024-01-01 01:00, 02:00)` +/// +/// Files are range-partitioned on `timestamp` and sorted on `(key, timestamp)`. +/// Because `date_bin(60 seconds, timestamp)` does not straddle the hour split, +/// grouping by `(key, time_bin)` is partition-disjoint. Today's planner still +/// inserts a hash shuffle; the test pins that plan so a follow-up can remove it. +pub(super) fn register_metrics_range_sorted_table(ctx: &SessionContext) { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("zone", DataType::Utf8, false), + Field::new("host", DataType::Utf8, false), + Field::new("pod", DataType::Utf8, false), + Field::new("service", DataType::Utf8, false), + Field::new( + "timestamp", + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + ), + Field::new("value", DataType::Int64, false), + ])); + + // Each partition covers 60 minutes. The split is aligned to the 60-second + // `date_bin` used by the test query, so time bins do not straddle files. + let hour_split = metrics_ts(60, 0); + let output_partitioning = Partitioning::Range( + RangePartitioning::try_new( + vec![col("timestamp").sort(true, true)], + vec![SplitPoint::new(vec![ScalarValue::TimestampNanosecond( + Some(hour_split), + None, + )])], + ) + .expect("metrics range partitioning should be valid"), + ); + + // Within each 60-minute file, rows are sorted by (key, timestamp). + let partitions = vec![ + vec![ + ("k1", "z1", "h1", "p1", "a", metrics_ts(0, 10), 1), + ("k1", "z1", "h1", "p1", "a", metrics_ts(0, 40), 2), + ("k1", "z1", "h1", "p1", "b", metrics_ts(1, 10), 99), + ("k2", "z1", "h1", "p1", "a", metrics_ts(30, 0), 3), + ("k2", "z1", "h1", "p1", "a", metrics_ts(30, 30), 4), + ], + vec![ + ("k1", "z1", "h1", "p1", "a", metrics_ts(60, 10), 10), + ("k1", "z1", "h1", "p1", "a", metrics_ts(60, 40), 20), + ("k2", "z1", "h1", "p1", "a", metrics_ts(90, 0), 30), + ("k2", "z1", "h1", "p1", "a", metrics_ts(105, 0), 5), + ], + ]; + + let table_dir = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("test_files/scratch_range_partitioning/metrics_range_sorted"); + register_metrics_listing_table( + ctx, + "metrics_range_sorted", + &table_dir, + Arc::clone(&schema), + partitions, + output_partitioning, + vec![vec![ + col("key").sort(true, true), + col("timestamp").sort(true, true), + ]], + ); +} + +fn register_metrics_listing_table( + ctx: &SessionContext, + name: &str, + table_dir: impl AsRef, + schema: SchemaRef, + partitions: Vec>, + output_partitioning: Partitioning, + file_sort_order: Vec>, +) { + let table_dir = table_dir.as_ref(); + if table_dir.exists() { + remove_dir_all(table_dir).expect("test table dir should be removable"); + } + create_dir_all(table_dir).expect("test table dir should be created"); + for (idx, rows) in partitions.into_iter().enumerate() { + let batch = metrics_batch(Arc::clone(&schema), &rows); + let file = File::create(table_dir.join(format!("part-{idx}.parquet"))) + .expect("test table parquet partition should be created"); + let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), None) + .expect("test table parquet writer should be created"); + writer + .write(&batch) + .expect("test table parquet partition should be written"); + writer + .close() + .expect("test table parquet writer should close"); + } + + let table_path = format!( + "{}/", + table_dir + .to_str() + .expect("test table path should be valid utf8") + ); + let table_url = + ListingTableUrl::parse(&table_path).expect("test table url should parse"); + let options = ListingOptions::new(Arc::new(ParquetFormat::default())) + .with_output_partitioning(Some(output_partitioning)) + .with_file_sort_order(file_sort_order); + let config = ListingTableConfig::new(table_url) + .with_listing_options(options) + .with_schema(schema); + let table = + ListingTable::try_new(config).expect("test listing table should be valid"); + + ctx.register_table(name, Arc::new(table)) + .expect("test listing table registration should succeed"); +} + +fn metrics_batch(schema: SchemaRef, rows: &[MetricsRow]) -> RecordBatch { + RecordBatch::try_new( + schema, + vec![ + Arc::new(StringArray::from_iter_values(rows.iter().map(|row| row.0))) + as ArrayRef, + Arc::new(StringArray::from_iter_values(rows.iter().map(|row| row.1))) + as ArrayRef, + Arc::new(StringArray::from_iter_values(rows.iter().map(|row| row.2))) + as ArrayRef, + Arc::new(StringArray::from_iter_values(rows.iter().map(|row| row.3))) + as ArrayRef, + Arc::new(StringArray::from_iter_values(rows.iter().map(|row| row.4))) + as ArrayRef, + Arc::new(TimestampNanosecondArray::from_iter_values( + rows.iter().map(|row| row.5), + )) as ArrayRef, + Arc::new(Int64Array::from_iter_values(rows.iter().map(|row| row.6))) + as ArrayRef, + ], + ) + .expect("metrics batch should be valid") +} diff --git a/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt b/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt new file mode 100644 index 0000000000000..43b12502f812a --- /dev/null +++ b/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt @@ -0,0 +1,183 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# GROUP BY on a table that is: +# - range partitioned on timestamp into two 60-minute file groups +# - sorted within each file on (key, timestamp) +# +# Query: +# 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 +# +# Scan metadata already advertises: +# 1. Range([timestamp]) and output_ordering=[key, timestamp] +# 2. Two file_groups, so the two 60-minute streams run in parallel +# +# Improvement opportunity: +# date_bin(60s) is monotonic in timestamp and the hour split is aligned to bin +# boundaries, so (key, time_bin) is partition-disjoint. Aggregation could be a +# single streaming SinglePartitioned step with no hash shuffle. +# +# Today's plan still hash-repartitions: +# Partial AggregateExec (ordering_mode=Sorted) +# -> RepartitionExec Hash([key, date_bin(...)]) +# -> FinalPartitioned AggregateExec (ordering_mode=Sorted) + +statement ok +set datafusion.explain.physical_plan_only = true; + +statement ok +set datafusion.execution.collect_statistics = false; + +statement ok +set datafusion.execution.target_partitions = 2; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 2; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 1; + +statement ok +set datafusion.optimizer.enable_round_robin_repartition = false; + +statement ok +set datafusion.optimizer.enable_join_dynamic_filter_pushdown = false; + +statement ok +set datafusion.optimizer.enable_topk_dynamic_filter_pushdown = false; + +statement ok +set datafusion.optimizer.enable_aggregate_dynamic_filter_pushdown = false; + +statement ok +set datafusion.execution.parquet.pushdown_filters = false; + +########## +# TEST 1: Scan metadata — range partitioned on timestamp, sorted on (key, timestamp), +# two parallel file groups covering 60-minute intervals. +########## + +query TT +EXPLAIN SELECT key, timestamp, value FROM metrics_range_sorted; +---- +physical_plan DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/metrics_range_sorted/part-0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/metrics_range_sorted/part-1.parquet]]}, projection=[key, timestamp, value], output_ordering=[key@0 ASC, timestamp@1 ASC], output_partitioning=Range([timestamp@1 ASC], [(1704070800000000000)], 2), file_type=parquet + +########## +# TEST 2: Filtered time-bin aggregation. +# GROUP BY keys are (key, date_bin(timestamp)). Input is sorted on those keys +# (date_bin is monotonic in timestamp) and range-partitioned so bins do not +# overlap across the two 60-minute streams. +# +# Today this is still Partial + hash RepartitionExec + Final, even though +# ordering_mode=Sorted is already recognized. +########## + +query TT +EXPLAIN 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; +---- +physical_plan +01)ProjectionExec: expr=[key@0 as key, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1 as time_bin, sum(metrics_range_sorted.value)@2 as sum(metrics_range_sorted.value)] +02)--AggregateExec: mode=FinalPartitioned, gby=[key@0 as key, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1 as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)], aggr=[sum(metrics_range_sorted.value)], ordering_mode=Sorted +03)----RepartitionExec: partitioning=Hash([key@0, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1], 2), input_partitions=2, preserve_order=true, sort_exprs=key@0 ASC, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1 ASC +04)------AggregateExec: mode=Partial, gby=[key@0 as key, date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, timestamp@1) as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)], aggr=[sum(metrics_range_sorted.value)], ordering_mode=Sorted +05)--------FilterExec: service@1 = a, projection=[key@0, timestamp@2, value@3] +06)----------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/metrics_range_sorted/part-0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/metrics_range_sorted/part-1.parquet]]}, projection=[key, service, timestamp, value], output_ordering=[key@0 ASC, timestamp@2 ASC], output_partitioning=Range([timestamp@2 ASC], [(1704070800000000000)], 2), file_type=parquet, predicate=service@4 = a, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= a AND a <= service_max@1, required_guarantees=[service in (a)] + +query TPI +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 +ORDER BY key, time_bin; +---- +k1 2024-01-01T00:00:00 3 +k1 2024-01-01T01:00:00 30 +k2 2024-01-01T00:30:00 7 +k2 2024-01-01T01:30:00 30 +k2 2024-01-01T01:45:00 5 + +########## +# TEST 3: Same aggregation without the service filter. The scan still has two +# 60-minute file groups, and today's plan still hash-repartitions. +########## + +query TT +EXPLAIN SELECT key, date_bin(INTERVAL '60 seconds', timestamp) AS time_bin, sum(value) +FROM metrics_range_sorted +GROUP BY key, time_bin; +---- +physical_plan +01)ProjectionExec: expr=[key@0 as key, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1 as time_bin, sum(metrics_range_sorted.value)@2 as sum(metrics_range_sorted.value)] +02)--AggregateExec: mode=FinalPartitioned, gby=[key@0 as key, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1 as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)], aggr=[sum(metrics_range_sorted.value)], ordering_mode=Sorted +03)----RepartitionExec: partitioning=Hash([key@0, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1], 2), input_partitions=2, preserve_order=true, sort_exprs=key@0 ASC, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1 ASC +04)------AggregateExec: mode=Partial, gby=[key@0 as key, date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, timestamp@1) as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)], aggr=[sum(metrics_range_sorted.value)], ordering_mode=Sorted +05)--------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/metrics_range_sorted/part-0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/metrics_range_sorted/part-1.parquet]]}, projection=[key, timestamp, value], output_ordering=[key@0 ASC, timestamp@1 ASC], output_partitioning=Range([timestamp@1 ASC], [(1704070800000000000)], 2), file_type=parquet + +query TPI +SELECT key, date_bin(INTERVAL '60 seconds', timestamp) AS time_bin, sum(value) +FROM metrics_range_sorted +GROUP BY key, time_bin +ORDER BY key, time_bin; +---- +k1 2024-01-01T00:00:00 3 +k1 2024-01-01T00:01:00 99 +k1 2024-01-01T01:00:00 30 +k2 2024-01-01T00:30:00 7 +k2 2024-01-01T01:30:00 30 +k2 2024-01-01T01:45:00 5 + +########## +# CLEANUP +########## + +# The SLT runner sets `target_partitions` to 4 instead of using the default, so +# reset it explicitly. +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +reset datafusion.explain.physical_plan_only; + +statement ok +reset datafusion.execution.collect_statistics; + +statement ok +reset datafusion.optimizer.subset_repartition_threshold; + +statement ok +reset datafusion.optimizer.preserve_file_partitions; + +statement ok +reset datafusion.optimizer.enable_round_robin_repartition; + +statement ok +reset datafusion.optimizer.enable_join_dynamic_filter_pushdown; + +statement ok +reset datafusion.optimizer.enable_topk_dynamic_filter_pushdown; + +statement ok +reset datafusion.optimizer.enable_aggregate_dynamic_filter_pushdown; + +statement ok +reset datafusion.execution.parquet.pushdown_filters; From 58dea6208132cdc49a33516329ae4e0a93f1b024 Mon Sep 17 00:00:00 2001 From: Nga Tran Date: Thu, 20 Aug 2026 10:27:52 -0400 Subject: [PATCH 2/3] test: generalize range-sorted time-bin aggregation test names Use generic column and table names so the coverage does not expose a metrics-specific schema. Co-authored-by: Cursor --- datafusion/sqllogictest/src/test_context.rs | 6 +- .../src/test_context/range_partitioning.rs | 62 +++++++++---------- .../test_files/range_sorted_time_bin_agg.slt | 44 ++++++------- 3 files changed, 56 insertions(+), 56 deletions(-) diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index 7252a0936afea..9b97d3f59dac4 100644 --- a/datafusion/sqllogictest/src/test_context.rs +++ b/datafusion/sqllogictest/src/test_context.rs @@ -55,7 +55,7 @@ use datafusion_spark::SessionStateBuilderSpark; use crate::is_spark_path; use range_partitioning::{ - register_metrics_range_sorted_table, register_range_partitioned_table, + register_range_partitioned_table, register_range_sorted_time_bin_table, }; use async_trait::async_trait; @@ -182,8 +182,8 @@ impl TestContext { register_range_partitioned_table(test_ctx.session_ctx()); } "range_sorted_time_bin_agg.slt" => { - info!("Registering range-sorted metrics table"); - register_metrics_range_sorted_table(test_ctx.session_ctx()); + info!("Registering range-sorted time-bin table"); + register_range_sorted_time_bin_table(test_ctx.session_ctx()); } "metadata.slt" | "arrow_field.slt" => { info!("Registering metadata table tables"); diff --git a/datafusion/sqllogictest/src/test_context/range_partitioning.rs b/datafusion/sqllogictest/src/test_context/range_partitioning.rs index de4f5875c8c93..d1e92c68c8fb4 100644 --- a/datafusion/sqllogictest/src/test_context/range_partitioning.rs +++ b/datafusion/sqllogictest/src/test_context/range_partitioning.rs @@ -296,21 +296,21 @@ fn range_batch(schema: SchemaRef, rows: &[(i32, i32, i32)]) -> RecordBatch { } // ============================================================================== -// Metrics table: range-partitioned on timestamp, sorted on (key, timestamp) +// Time-bin table: range-partitioned on timestamp, sorted on (key, timestamp) // ============================================================================== /// Unix nanoseconds for `2024-01-01 00:00:00 UTC`. -const METRICS_EPOCH_NS: i64 = 1_704_067_200_000_000_000; +const TIME_BIN_EPOCH_NS: i64 = 1_704_067_200_000_000_000; const NANOS_PER_SECOND: i64 = 1_000_000_000; const NANOS_PER_MINUTE: i64 = 60 * NANOS_PER_SECOND; /// Timestamp helper: minutes and seconds after `2024-01-01 00:00:00 UTC`. -fn metrics_ts(minutes: i64, seconds: i64) -> i64 { - METRICS_EPOCH_NS + minutes * NANOS_PER_MINUTE + seconds * NANOS_PER_SECOND +fn time_bin_ts(minutes: i64, seconds: i64) -> i64 { + TIME_BIN_EPOCH_NS + minutes * NANOS_PER_MINUTE + seconds * NANOS_PER_SECOND } -/// Row: (key, zone, host, pod, service, timestamp_ns, value) -type MetricsRow = ( +/// Row: (key, col1, col2, col3, col4, timestamp_ns, value) +type TimeBinRow = ( &'static str, &'static str, &'static str, @@ -320,7 +320,7 @@ type MetricsRow = ( i64, ); -/// Registers `metrics_range_sorted` for time-bin aggregation plan tests. +/// Registers `range_sorted_time_bin` for time-bin aggregation plan tests. /// /// Two file groups, each covering a 60-minute timestamp range: /// - partition 0: `[2024-01-01 00:00, 01:00)` @@ -330,13 +330,13 @@ type MetricsRow = ( /// Because `date_bin(60 seconds, timestamp)` does not straddle the hour split, /// grouping by `(key, time_bin)` is partition-disjoint. Today's planner still /// inserts a hash shuffle; the test pins that plan so a follow-up can remove it. -pub(super) fn register_metrics_range_sorted_table(ctx: &SessionContext) { +pub(super) fn register_range_sorted_time_bin_table(ctx: &SessionContext) { let schema = Arc::new(Schema::new(vec![ Field::new("key", DataType::Utf8, false), - Field::new("zone", DataType::Utf8, false), - Field::new("host", DataType::Utf8, false), - Field::new("pod", DataType::Utf8, false), - Field::new("service", DataType::Utf8, false), + Field::new("col1", DataType::Utf8, false), + Field::new("col2", DataType::Utf8, false), + Field::new("col3", DataType::Utf8, false), + Field::new("col4", DataType::Utf8, false), Field::new( "timestamp", DataType::Timestamp(TimeUnit::Nanosecond, None), @@ -347,7 +347,7 @@ pub(super) fn register_metrics_range_sorted_table(ctx: &SessionContext) { // Each partition covers 60 minutes. The split is aligned to the 60-second // `date_bin` used by the test query, so time bins do not straddle files. - let hour_split = metrics_ts(60, 0); + let hour_split = time_bin_ts(60, 0); let output_partitioning = Partitioning::Range( RangePartitioning::try_new( vec![col("timestamp").sort(true, true)], @@ -356,31 +356,31 @@ pub(super) fn register_metrics_range_sorted_table(ctx: &SessionContext) { None, )])], ) - .expect("metrics range partitioning should be valid"), + .expect("time-bin range partitioning should be valid"), ); // Within each 60-minute file, rows are sorted by (key, timestamp). let partitions = vec![ vec![ - ("k1", "z1", "h1", "p1", "a", metrics_ts(0, 10), 1), - ("k1", "z1", "h1", "p1", "a", metrics_ts(0, 40), 2), - ("k1", "z1", "h1", "p1", "b", metrics_ts(1, 10), 99), - ("k2", "z1", "h1", "p1", "a", metrics_ts(30, 0), 3), - ("k2", "z1", "h1", "p1", "a", metrics_ts(30, 30), 4), + ("k1", "x", "y", "z", "a", time_bin_ts(0, 10), 1), + ("k1", "x", "y", "z", "a", time_bin_ts(0, 40), 2), + ("k1", "x", "y", "z", "b", time_bin_ts(1, 10), 99), + ("k2", "x", "y", "z", "a", time_bin_ts(30, 0), 3), + ("k2", "x", "y", "z", "a", time_bin_ts(30, 30), 4), ], vec![ - ("k1", "z1", "h1", "p1", "a", metrics_ts(60, 10), 10), - ("k1", "z1", "h1", "p1", "a", metrics_ts(60, 40), 20), - ("k2", "z1", "h1", "p1", "a", metrics_ts(90, 0), 30), - ("k2", "z1", "h1", "p1", "a", metrics_ts(105, 0), 5), + ("k1", "x", "y", "z", "a", time_bin_ts(60, 10), 10), + ("k1", "x", "y", "z", "a", time_bin_ts(60, 40), 20), + ("k2", "x", "y", "z", "a", time_bin_ts(90, 0), 30), + ("k2", "x", "y", "z", "a", time_bin_ts(105, 0), 5), ], ]; let table_dir = Path::new(env!("CARGO_MANIFEST_DIR")) - .join("test_files/scratch_range_partitioning/metrics_range_sorted"); - register_metrics_listing_table( + .join("test_files/scratch_range_partitioning/range_sorted_time_bin"); + register_time_bin_listing_table( ctx, - "metrics_range_sorted", + "range_sorted_time_bin", &table_dir, Arc::clone(&schema), partitions, @@ -392,12 +392,12 @@ pub(super) fn register_metrics_range_sorted_table(ctx: &SessionContext) { ); } -fn register_metrics_listing_table( +fn register_time_bin_listing_table( ctx: &SessionContext, name: &str, table_dir: impl AsRef, schema: SchemaRef, - partitions: Vec>, + partitions: Vec>, output_partitioning: Partitioning, file_sort_order: Vec>, ) { @@ -407,7 +407,7 @@ fn register_metrics_listing_table( } create_dir_all(table_dir).expect("test table dir should be created"); for (idx, rows) in partitions.into_iter().enumerate() { - let batch = metrics_batch(Arc::clone(&schema), &rows); + let batch = time_bin_batch(Arc::clone(&schema), &rows); let file = File::create(table_dir.join(format!("part-{idx}.parquet"))) .expect("test table parquet partition should be created"); let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), None) @@ -441,7 +441,7 @@ fn register_metrics_listing_table( .expect("test listing table registration should succeed"); } -fn metrics_batch(schema: SchemaRef, rows: &[MetricsRow]) -> RecordBatch { +fn time_bin_batch(schema: SchemaRef, rows: &[TimeBinRow]) -> RecordBatch { RecordBatch::try_new( schema, vec![ @@ -462,5 +462,5 @@ fn metrics_batch(schema: SchemaRef, rows: &[MetricsRow]) -> RecordBatch { as ArrayRef, ], ) - .expect("metrics batch should be valid") + .expect("time-bin batch should be valid") } diff --git a/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt b/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt index 43b12502f812a..18123a492dbd6 100644 --- a/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt +++ b/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt @@ -21,8 +21,8 @@ # # Query: # SELECT key, date_bin(INTERVAL '60 seconds', timestamp) AS time_bin, sum(value) -# FROM metrics_range_sorted -# WHERE service = 'a' +# FROM range_sorted_time_bin +# WHERE col4 = 'a' # GROUP BY key, time_bin # # Scan metadata already advertises: @@ -75,9 +75,9 @@ set datafusion.execution.parquet.pushdown_filters = false; ########## query TT -EXPLAIN SELECT key, timestamp, value FROM metrics_range_sorted; +EXPLAIN SELECT key, timestamp, value FROM range_sorted_time_bin; ---- -physical_plan DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/metrics_range_sorted/part-0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/metrics_range_sorted/part-1.parquet]]}, projection=[key, timestamp, value], output_ordering=[key@0 ASC, timestamp@1 ASC], output_partitioning=Range([timestamp@1 ASC], [(1704070800000000000)], 2), file_type=parquet +physical_plan DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_sorted_time_bin/part-0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_sorted_time_bin/part-1.parquet]]}, projection=[key, timestamp, value], output_ordering=[key@0 ASC, timestamp@1 ASC], output_partitioning=Range([timestamp@1 ASC], [(1704070800000000000)], 2), file_type=parquet ########## # TEST 2: Filtered time-bin aggregation. @@ -91,22 +91,22 @@ physical_plan DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion query TT EXPLAIN SELECT key, date_bin(INTERVAL '60 seconds', timestamp) AS time_bin, sum(value) -FROM metrics_range_sorted -WHERE service = 'a' +FROM range_sorted_time_bin +WHERE col4 = 'a' GROUP BY key, time_bin; ---- physical_plan -01)ProjectionExec: expr=[key@0 as key, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1 as time_bin, sum(metrics_range_sorted.value)@2 as sum(metrics_range_sorted.value)] -02)--AggregateExec: mode=FinalPartitioned, gby=[key@0 as key, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1 as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)], aggr=[sum(metrics_range_sorted.value)], ordering_mode=Sorted -03)----RepartitionExec: partitioning=Hash([key@0, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1], 2), input_partitions=2, preserve_order=true, sort_exprs=key@0 ASC, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1 ASC -04)------AggregateExec: mode=Partial, gby=[key@0 as key, date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, timestamp@1) as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)], aggr=[sum(metrics_range_sorted.value)], ordering_mode=Sorted -05)--------FilterExec: service@1 = a, projection=[key@0, timestamp@2, value@3] -06)----------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/metrics_range_sorted/part-0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/metrics_range_sorted/part-1.parquet]]}, projection=[key, service, timestamp, value], output_ordering=[key@0 ASC, timestamp@2 ASC], output_partitioning=Range([timestamp@2 ASC], [(1704070800000000000)], 2), file_type=parquet, predicate=service@4 = a, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= a AND a <= service_max@1, required_guarantees=[service in (a)] +01)ProjectionExec: expr=[key@0 as key, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)@1 as time_bin, sum(range_sorted_time_bin.value)@2 as sum(range_sorted_time_bin.value)] +02)--AggregateExec: mode=FinalPartitioned, gby=[key@0 as key, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)@1 as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)], aggr=[sum(range_sorted_time_bin.value)], ordering_mode=Sorted +03)----RepartitionExec: partitioning=Hash([key@0, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)@1], 2), input_partitions=2, preserve_order=true, sort_exprs=key@0 ASC, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)@1 ASC +04)------AggregateExec: mode=Partial, gby=[key@0 as key, date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, timestamp@1) as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)], aggr=[sum(range_sorted_time_bin.value)], ordering_mode=Sorted +05)--------FilterExec: col4@1 = a, projection=[key@0, timestamp@2, value@3] +06)----------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_sorted_time_bin/part-0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_sorted_time_bin/part-1.parquet]]}, projection=[key, col4, timestamp, value], output_ordering=[key@0 ASC, timestamp@2 ASC], output_partitioning=Range([timestamp@2 ASC], [(1704070800000000000)], 2), file_type=parquet, predicate=col4@4 = a, pruning_predicate=col4_null_count@2 != row_count@3 AND col4_min@0 <= a AND a <= col4_max@1, required_guarantees=[col4 in (a)] query TPI SELECT key, date_bin(INTERVAL '60 seconds', timestamp) AS time_bin, sum(value) -FROM metrics_range_sorted -WHERE service = 'a' +FROM range_sorted_time_bin +WHERE col4 = 'a' GROUP BY key, time_bin ORDER BY key, time_bin; ---- @@ -117,25 +117,25 @@ k2 2024-01-01T01:30:00 30 k2 2024-01-01T01:45:00 5 ########## -# TEST 3: Same aggregation without the service filter. The scan still has two +# TEST 3: Same aggregation without the col4 filter. The scan still has two # 60-minute file groups, and today's plan still hash-repartitions. ########## query TT EXPLAIN SELECT key, date_bin(INTERVAL '60 seconds', timestamp) AS time_bin, sum(value) -FROM metrics_range_sorted +FROM range_sorted_time_bin GROUP BY key, time_bin; ---- physical_plan -01)ProjectionExec: expr=[key@0 as key, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1 as time_bin, sum(metrics_range_sorted.value)@2 as sum(metrics_range_sorted.value)] -02)--AggregateExec: mode=FinalPartitioned, gby=[key@0 as key, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1 as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)], aggr=[sum(metrics_range_sorted.value)], ordering_mode=Sorted -03)----RepartitionExec: partitioning=Hash([key@0, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1], 2), input_partitions=2, preserve_order=true, sort_exprs=key@0 ASC, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1 ASC -04)------AggregateExec: mode=Partial, gby=[key@0 as key, date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, timestamp@1) as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)], aggr=[sum(metrics_range_sorted.value)], ordering_mode=Sorted -05)--------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/metrics_range_sorted/part-0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/metrics_range_sorted/part-1.parquet]]}, projection=[key, timestamp, value], output_ordering=[key@0 ASC, timestamp@1 ASC], output_partitioning=Range([timestamp@1 ASC], [(1704070800000000000)], 2), file_type=parquet +01)ProjectionExec: expr=[key@0 as key, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)@1 as time_bin, sum(range_sorted_time_bin.value)@2 as sum(range_sorted_time_bin.value)] +02)--AggregateExec: mode=FinalPartitioned, gby=[key@0 as key, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)@1 as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)], aggr=[sum(range_sorted_time_bin.value)], ordering_mode=Sorted +03)----RepartitionExec: partitioning=Hash([key@0, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)@1], 2), input_partitions=2, preserve_order=true, sort_exprs=key@0 ASC, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)@1 ASC +04)------AggregateExec: mode=Partial, gby=[key@0 as key, date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, timestamp@1) as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)], aggr=[sum(range_sorted_time_bin.value)], ordering_mode=Sorted +05)--------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_sorted_time_bin/part-0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_sorted_time_bin/part-1.parquet]]}, projection=[key, timestamp, value], output_ordering=[key@0 ASC, timestamp@1 ASC], output_partitioning=Range([timestamp@1 ASC], [(1704070800000000000)], 2), file_type=parquet query TPI SELECT key, date_bin(INTERVAL '60 seconds', timestamp) AS time_bin, sum(value) -FROM metrics_range_sorted +FROM range_sorted_time_bin GROUP BY key, time_bin ORDER BY key, time_bin; ---- From 600609629edb3033ca9f688070599ec9f7940d11 Mon Sep 17 00:00:00 2001 From: Nga Tran Date: Thu, 20 Aug 2026 11:34:12 -0400 Subject: [PATCH 3/3] test: share parquet listing-table helper for range tests Accept RecordBatches and optional file sort order so the time-bin table reuses the same registration path. Co-authored-by: Cursor --- .../src/test_context/range_partitioning.rs | 96 +++++++------------ 1 file changed, 34 insertions(+), 62 deletions(-) diff --git a/datafusion/sqllogictest/src/test_context/range_partitioning.rs b/datafusion/sqllogictest/src/test_context/range_partitioning.rs index d1e92c68c8fb4..becde0f3286db 100644 --- a/datafusion/sqllogictest/src/test_context/range_partitioning.rs +++ b/datafusion/sqllogictest/src/test_context/range_partitioning.rs @@ -97,8 +97,9 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) { "range_partitioned", &range_table_dir, Arc::clone(&schema), - RANGE_PARTITIONS, + range_batches(&schema, RANGE_PARTITIONS), output_partitioning, + None, ); register_unbounded_range_stream_table( @@ -134,8 +135,9 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) { Path::new(env!("CARGO_MANIFEST_DIR")) .join("test_files/scratch_range_partitioning/range_partitioned_shifted"), Arc::clone(&schema), - SHIFTED_RANGE_PARTITIONS, + range_batches(&schema, SHIFTED_RANGE_PARTITIONS), shifted_output_partitioning, + None, ); // Same rows as `range_partitioned` but split into only three range @@ -158,8 +160,9 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) { Path::new(env!("CARGO_MANIFEST_DIR")) .join("test_files/scratch_range_partitioning/range_partitioned_narrow"), Arc::clone(&schema), - NARROW_RANGE_PARTITIONS, + range_batches(&schema, NARROW_RANGE_PARTITIONS), narrow_output_partitioning, + None, ); let sparse_output_partitioning = Partitioning::Range( @@ -180,8 +183,9 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) { Path::new(env!("CARGO_MANIFEST_DIR")) .join("test_files/scratch_range_partitioning/range_partitioned_sparse"), Arc::clone(&schema), - SPARSE_RANGE_PARTITIONS, + range_batches(&schema, SPARSE_RANGE_PARTITIONS), sparse_output_partitioning, + None, ); } @@ -190,16 +194,16 @@ fn register_parquet_listing_table( name: &str, table_dir: impl AsRef, schema: SchemaRef, - partitions: impl IntoIterator, + batches: Vec, output_partitioning: Partitioning, + file_sort_order: Option>>, ) { let table_dir = table_dir.as_ref(); if table_dir.exists() { remove_dir_all(table_dir).expect("test table dir should be removable"); } create_dir_all(table_dir).expect("test table dir should be created"); - for (idx, rows) in partitions.into_iter().enumerate() { - let batch = range_batch(Arc::clone(&schema), rows); + for (idx, batch) in batches.into_iter().enumerate() { let file = File::create(table_dir.join(format!("part-{idx}.parquet"))) .expect("test table parquet partition should be created"); let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), None) @@ -220,8 +224,11 @@ fn register_parquet_listing_table( ); let table_url = ListingTableUrl::parse(&table_path).expect("test table url should parse"); - let options = ListingOptions::new(Arc::new(ParquetFormat::default())) + let mut options = ListingOptions::new(Arc::new(ParquetFormat::default())) .with_output_partitioning(Some(output_partitioning)); + if let Some(file_sort_order) = file_sort_order { + options = options.with_file_sort_order(file_sort_order); + } let config = ListingTableConfig::new(table_url) .with_listing_options(options) .with_schema(schema); @@ -280,6 +287,16 @@ fn range_stream_partition( )])) } +fn range_batches( + schema: &SchemaRef, + partitions: impl IntoIterator, +) -> Vec { + partitions + .into_iter() + .map(|rows| range_batch(Arc::clone(schema), rows)) + .collect() +} + fn range_batch(schema: SchemaRef, rows: &[(i32, i32, i32)]) -> RecordBatch { RecordBatch::try_new( schema, @@ -360,7 +377,7 @@ pub(super) fn register_range_sorted_time_bin_table(ctx: &SessionContext) { ); // Within each 60-minute file, rows are sorted by (key, timestamp). - let partitions = vec![ + let partitions = [ vec![ ("k1", "x", "y", "z", "a", time_bin_ts(0, 10), 1), ("k1", "x", "y", "z", "a", time_bin_ts(0, 40), 2), @@ -378,69 +395,24 @@ pub(super) fn register_range_sorted_time_bin_table(ctx: &SessionContext) { let table_dir = Path::new(env!("CARGO_MANIFEST_DIR")) .join("test_files/scratch_range_partitioning/range_sorted_time_bin"); - register_time_bin_listing_table( + let batches = partitions + .iter() + .map(|rows| time_bin_batch(Arc::clone(&schema), rows)) + .collect(); + register_parquet_listing_table( ctx, "range_sorted_time_bin", &table_dir, Arc::clone(&schema), - partitions, + batches, output_partitioning, - vec![vec![ + Some(vec![vec![ col("key").sort(true, true), col("timestamp").sort(true, true), - ]], + ]]), ); } -fn register_time_bin_listing_table( - ctx: &SessionContext, - name: &str, - table_dir: impl AsRef, - schema: SchemaRef, - partitions: Vec>, - output_partitioning: Partitioning, - file_sort_order: Vec>, -) { - let table_dir = table_dir.as_ref(); - if table_dir.exists() { - remove_dir_all(table_dir).expect("test table dir should be removable"); - } - create_dir_all(table_dir).expect("test table dir should be created"); - for (idx, rows) in partitions.into_iter().enumerate() { - let batch = time_bin_batch(Arc::clone(&schema), &rows); - let file = File::create(table_dir.join(format!("part-{idx}.parquet"))) - .expect("test table parquet partition should be created"); - let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), None) - .expect("test table parquet writer should be created"); - writer - .write(&batch) - .expect("test table parquet partition should be written"); - writer - .close() - .expect("test table parquet writer should close"); - } - - let table_path = format!( - "{}/", - table_dir - .to_str() - .expect("test table path should be valid utf8") - ); - let table_url = - ListingTableUrl::parse(&table_path).expect("test table url should parse"); - let options = ListingOptions::new(Arc::new(ParquetFormat::default())) - .with_output_partitioning(Some(output_partitioning)) - .with_file_sort_order(file_sort_order); - let config = ListingTableConfig::new(table_url) - .with_listing_options(options) - .with_schema(schema); - let table = - ListingTable::try_new(config).expect("test listing table should be valid"); - - ctx.register_table(name, Arc::new(table)) - .expect("test listing table registration should succeed"); -} - fn time_bin_batch(schema: SchemaRef, rows: &[TimeBinRow]) -> RecordBatch { RecordBatch::try_new( schema,