test: add range-partitioned sorted time-bin aggregation coverage

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());

datafusion/sqllogictest/src/test_context/range_partitioning.rs

@@ -19,17 +19,19 @@ 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};
 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<Path>,
     schema: SchemaRef,
-    partitions: impl IntoIterator<Item = &'static [(i32, i32, i32)]>,
+    batches: Vec<RecordBatch>,
     output_partitioning: Partitioning,
+    file_sort_order: Option<Vec<Vec<SortExpr>>>,
 ) {
     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<Item = &'static [(i32, i32, i32)]>,
+) -> Vec<RecordBatch> {
+    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")
+}