diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index 39aa2b09e5685..9b97d3f59dac4 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_range_partitioned_table, register_range_sorted_time_bin_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 time-bin table"); + register_range_sorted_time_bin_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..becde0f3286db 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, @@ -95,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( @@ -132,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 @@ -156,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( @@ -178,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, ); } @@ -188,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) @@ -218,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); @@ -278,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, @@ -292,3 +311,128 @@ fn range_batch(schema: SchemaRef, rows: &[(i32, i32, i32)]) -> RecordBatch { ) .expect("range batch should be valid") } + +// ============================================================================== +// Time-bin table: range-partitioned on timestamp, sorted on (key, timestamp) +// ============================================================================== + +/// Unix nanoseconds for `2024-01-01 00:00:00 UTC`. +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 time_bin_ts(minutes: i64, seconds: i64) -> i64 { + TIME_BIN_EPOCH_NS + minutes * NANOS_PER_MINUTE + seconds * NANOS_PER_SECOND +} + +/// Row: (key, col1, col2, col3, col4, timestamp_ns, value) +type TimeBinRow = ( + &'static str, + &'static str, + &'static str, + &'static str, + &'static str, + i64, + i64, +); + +/// 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)` +/// - 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_range_sorted_time_bin_table(ctx: &SessionContext) { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", 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), + 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 = time_bin_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("time-bin range partitioning should be valid"), + ); + + // Within each 60-minute file, rows are sorted by (key, timestamp). + 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), + ("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", "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/range_sorted_time_bin"); + 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), + batches, + output_partitioning, + Some(vec![vec![ + col("key").sort(true, true), + col("timestamp").sort(true, true), + ]]), + ); +} + +fn time_bin_batch(schema: SchemaRef, rows: &[TimeBinRow]) -> 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("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 new file mode 100644 index 0000000000000..18123a492dbd6 --- /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 range_sorted_time_bin +# WHERE col4 = '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 range_sorted_time_bin; +---- +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. +# 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 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 }"),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 range_sorted_time_bin +WHERE col4 = '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 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 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 }"),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 range_sorted_time_bin +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;