Add support for partitioned by in DDL

mwylde · mwylde · commit 6f82601f1f1c · 2025-10-15T16:46:36.000-07:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -166,7 +166,7 @@ recursive = "0.1.1"
 regex = "1.8"
 rstest = "0.25.0"
 serde_json = "1"
-sqlparser = { git = "https://github.com/ArroyoSystems/sqlparser-rs", branch = "0.55.0/arroyo", features = ["visitor"] }
+sqlparser = { git = "https://github.com/ArroyoSystems/sqlparser-rs", branch = "partitioned_by", features = ["visitor"] }
 tempfile = "3"
 tokio = { version = "1.45", features = ["macros", "rt", "sync"] }
 url = "2.5.4"
diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs
@@ -17,16 +17,16 @@
 
 //! [`DataSource`] and [`DataSourceExec`]
 
-use std::any::Any;
-use std::fmt;
-use std::fmt::{Debug, Formatter};
-use std::sync::Arc;
 use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
 use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet};
 use datafusion_physical_plan::projection::ProjectionExec;
 use datafusion_physical_plan::{
     DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
 };
+use std::any::Any;
+use std::fmt;
+use std::fmt::{Debug, Formatter};
+use std::sync::Arc;
 
 use crate::file_scan_config::FileScanConfig;
 use datafusion_common::config::ConfigOptions;
diff --git a/datafusion/ffi/src/udaf/accumulator.rs b/datafusion/ffi/src/udaf/accumulator.rs
@@ -44,8 +44,7 @@ pub struct FFI_Accumulator {
     ) -> RResult<(), RString>,
 
     // Evaluate and return a ScalarValues as protobuf bytes
-    pub evaluate:
-        unsafe extern "C" fn(accumulator: &Self) -> RResult<RVec<u8>, RString>,
+    pub evaluate: unsafe extern "C" fn(accumulator: &Self) -> RResult<RVec<u8>, RString>,
 
     pub size: unsafe extern "C" fn(accumulator: &Self) -> usize,
 
@@ -256,8 +255,7 @@ impl Accumulator for ForeignAccumulator {
 
     fn state(&self) -> Result<Vec<ScalarValue>> {
         unsafe {
-            let state_protos =
-                df_result!((self.accumulator.state)(&self.accumulator))?;
+            let state_protos = df_result!((self.accumulator.state)(&self.accumulator))?;
 
             state_protos
                 .into_iter()
diff --git a/datafusion/functions-aggregate/src/median.rs b/datafusion/functions-aggregate/src/median.rs
@@ -250,11 +250,9 @@ impl<T: ArrowNumericType> Accumulator for MedianAccumulator<T> {
             OffsetBuffer::new(ScalarBuffer::from(vec![0, self.all_values.len() as i32]));
 
         // Build inner array
-        let values_array = PrimitiveArray::<T>::new(
-            ScalarBuffer::from(self.all_values.clone()),
-            None,
-        )
-            .with_data_type(self.data_type.clone());
+        let values_array =
+            PrimitiveArray::<T>::new(ScalarBuffer::from(self.all_values.clone()), None)
+                .with_data_type(self.data_type.clone());
 
         // Build the result list array
         let list_array = ListArray::new(
diff --git a/datafusion/functions-aggregate/src/sum.rs b/datafusion/functions-aggregate/src/sum.rs
@@ -219,11 +219,14 @@ impl AggregateUDFImpl for Sum {
             .into()];
 
             if args.for_sliding {
-                fields.push(Field::new(
-                    format_state_name(args.name, "count"),
-                    DataType::UInt64,
-                    true,
-                ).into())
+                fields.push(
+                    Field::new(
+                        format_state_name(args.name, "count"),
+                        DataType::UInt64,
+                        true,
+                    )
+                    .into(),
+                )
             }
             Ok(fields)
         }
diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs
@@ -555,8 +555,8 @@ where
         });
         // SAFETY: the offsets were constructed correctly in `insert_if_new` --
         // monotonically increasing, overflows were checked.
-        let offsets = unsafe { OffsetBuffer::new_unchecked(ScalarBuffer::from(offsets.clone())) };
-
+        let offsets =
+            unsafe { OffsetBuffer::new_unchecked(ScalarBuffer::from(offsets.clone())) };
 
         let values = buffer.as_slice().iter().map(|b| *b).collect();
 
@@ -582,7 +582,6 @@ where
         }
     }
 
-
     /// Total number of entries (including null, if present)
     pub fn len(&self) -> usize {
         self.non_null_len() + self.null.map(|_| 1).unwrap_or(0)
diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs
@@ -130,7 +130,9 @@ impl AggregateMode {
             AggregateMode::Partial
             | AggregateMode::Single
             | AggregateMode::SinglePartitioned => true,
-            AggregateMode::Final | AggregateMode::FinalPartitioned | AggregateMode::CombinePartial => false,
+            AggregateMode::Final
+            | AggregateMode::FinalPartitioned
+            | AggregateMode::CombinePartial => false,
         }
     }
 }
@@ -1254,7 +1256,9 @@ pub fn aggregate_expressions(
             })
             .collect()),
         // In this mode, we build the merge expressions of the aggregation.
-        AggregateMode::Final | AggregateMode::FinalPartitioned | AggregateMode::CombinePartial => {
+        AggregateMode::Final
+        | AggregateMode::FinalPartitioned
+        | AggregateMode::CombinePartial => {
             let mut col_idx_base = col_idx_base;
             aggr_expr
                 .iter()
diff --git a/datafusion/physical-plan/src/aggregates/no_grouping.rs b/datafusion/physical-plan/src/aggregates/no_grouping.rs
@@ -234,7 +234,9 @@ fn aggregate_batch(
                 AggregateMode::Partial
                 | AggregateMode::Single
                 | AggregateMode::SinglePartitioned => accum.update_batch(&values),
-                AggregateMode::Final | AggregateMode::FinalPartitioned | AggregateMode::CombinePartial => accum.merge_batch(&values),
+                AggregateMode::Final
+                | AggregateMode::FinalPartitioned
+                | AggregateMode::CombinePartial => accum.merge_batch(&values),
             };
             let size_post = accum.size();
             allocated += size_post.saturating_sub(size_pre);
diff --git a/datafusion/physical-plan/src/aggregates/row_hash.rs b/datafusion/physical-plan/src/aggregates/row_hash.rs
@@ -955,8 +955,9 @@ impl GroupedHashAggregateStream {
         // Next output each aggregate value
         for acc in self.accumulators.iter_mut() {
             match self.mode {
-                AggregateMode::Partial
-                | AggregateMode::CombinePartial => output.extend(acc.state(emit_to)?),
+                AggregateMode::Partial | AggregateMode::CombinePartial => {
+                    output.extend(acc.state(emit_to)?)
+                }
                 _ if spilling => {
                     // If spilling, output partial state because the spilled data will be
                     // merged and re-evaluated later.
diff --git a/datafusion/physical-plan/src/joins/hash_join.rs b/datafusion/physical-plan/src/joins/hash_join.rs
@@ -28,10 +28,7 @@ use super::utils::{
     asymmetric_join_output_partitioning, get_final_indices_from_shared_bitmap,
     reorder_output_after_swap, swap_join_projection,
 };
-use super::{
-    utils::{OnceFut},
-    PartitionMode, SharedBitmapBuilder,
-};
+use super::{utils::OnceFut, PartitionMode, SharedBitmapBuilder};
 use super::{JoinOn, JoinOnRef};
 use crate::execution_plan::{boundedness_from_children, EmissionType};
 use crate::projection::{
@@ -811,7 +808,7 @@ impl ExecutionPlan for HashJoinExec {
                     need_produce_result_in_final(self.join_type),
                     self.right().output_partitioning().partition_count(),
                 ))
-            },
+            }
             PartitionMode::Partitioned => {
                 let left_stream = self.left.execute(partition, Arc::clone(&context))?;
 
diff --git a/datafusion/physical-plan/src/metrics/builder.rs b/datafusion/physical-plan/src/metrics/builder.rs
@@ -17,11 +17,9 @@
 
 //! Builder for creating arbitrary metrics
 
-use std::{borrow::Cow};
+use std::borrow::Cow;
 
-use super::{
-    Count, ExecutionPlanMetricsSet, Gauge, Label, MetricValue, Time, Timestamp,
-};
+use super::{Count, ExecutionPlanMetricsSet, Gauge, Label, MetricValue, Time, Timestamp};
 
 /// Structure for constructing metrics, counters, timers, etc.
 ///
@@ -87,8 +85,7 @@ impl<'a> MetricBuilder<'a> {
 
     /// Consume self and create a metric of the specified value
     /// registered with the MetricsSet
-    pub fn build(self, _value: MetricValue) {
-    }
+    pub fn build(self, _value: MetricValue) {}
 
     /// Consume self and create a new counter for recording output rows
     pub fn output_rows(self, partition: usize) -> Count {
diff --git a/datafusion/physical-plan/src/metrics/value.rs b/datafusion/physical-plan/src/metrics/value.rs
@@ -31,8 +31,7 @@ use datafusion_execution::memory_pool::human_readable_size;
 ///
 /// Note `clone`ing counters update the same underlying metrics
 #[derive(Debug, Clone)]
-pub struct Count {
-}
+pub struct Count {}
 
 impl PartialEq for Count {
     fn eq(&self, other: &Self) -> bool {
@@ -55,13 +54,11 @@ impl Default for Count {
 impl Count {
     /// create a new counter
     pub fn new() -> Self {
-        Self {
-        }
+        Self {}
     }
 
     /// Add `n` to the metric's value
-    pub fn add(&self, _n: usize) {
-    }
+    pub fn add(&self, _n: usize) {}
 
     /// Get the current value
     pub fn value(&self) -> usize {
@@ -74,8 +71,7 @@ impl Count {
 ///
 /// Note `clone`ing gauge update the same underlying metrics
 #[derive(Debug, Clone)]
-pub struct Gauge {
-}
+pub struct Gauge {}
 
 impl PartialEq for Gauge {
     fn eq(&self, other: &Self) -> bool {
@@ -98,21 +94,17 @@ impl Default for Gauge {
 impl Gauge {
     /// create a new gauge
     pub fn new() -> Self {
-        Self {
-        }
+        Self {}
     }
 
     /// Add `n` to the metric's value
-    pub fn add(&self, _n: usize) {
-    }
+    pub fn add(&self, _n: usize) {}
 
     /// Sub `n` from the metric's value
-    pub fn sub(&self, _n: usize) {
-    }
+    pub fn sub(&self, _n: usize) {}
 
     /// Set metric's value to maximum of `n` and current value
-    pub fn set_max(&self, _n: usize) {
-    }
+    pub fn set_max(&self, _n: usize) {}
 
     /// Set the metric's value to `n` and return the previous value
     pub fn set(&self, _n: usize) -> usize {
@@ -127,8 +119,7 @@ impl Gauge {
 
 /// Measure a potentially non contiguous duration of time
 #[derive(Debug, Clone)]
-pub struct Time {
-}
+pub struct Time {}
 
 impl Default for Time {
     fn default() -> Self {
@@ -153,13 +144,11 @@ impl Time {
     /// Create a new [`Time`] wrapper suitable for recording elapsed
     /// times for operations.
     pub fn new() -> Self {
-        Self {
-        }
+        Self {}
     }
 
     /// Add elapsed nanoseconds since `start`to self
-    pub fn add_elapsed(&self, _start: Instant) {
-    }
+    pub fn add_elapsed(&self, _start: Instant) {}
 
     /// Add duration of time to self
     ///
@@ -171,12 +160,10 @@ impl Time {
     /// This is based on the assumption that the timing logic in most cases is likely
     /// to take at least a nanosecond, and so this is reasonable mechanism to avoid
     /// ambiguity, especially on systems with low-resolution monotonic clocks
-    pub fn add_duration(&self, _duration: Duration) {
-    }
+    pub fn add_duration(&self, _duration: Duration) {}
 
     /// Add the number of nanoseconds of other `Time` to self
-    pub fn add(&self, _other: &Time) {
-    }
+    pub fn add(&self, _other: &Time) {}
 
     /// return a scoped guard that adds the amount of time elapsed
     /// between its creation and its drop or call to `stop` to the
@@ -197,8 +184,7 @@ impl Time {
 /// Stores a single timestamp, stored as the number of nanoseconds
 /// elapsed from Jan 1, 1970 UTC
 #[derive(Debug, Clone)]
-pub struct Timestamp {
-}
+pub struct Timestamp {}
 
 impl Default for Timestamp {
     fn default() -> Self {
@@ -209,8 +195,7 @@ impl Default for Timestamp {
 impl Timestamp {
     /// Create a new timestamp and sets its value to 0
     pub fn new() -> Self {
-        Self {
-        }
+        Self {}
     }
 
     /// Sets the timestamps value to the current time
@@ -219,8 +204,7 @@ impl Timestamp {
     }
 
     /// Sets the timestamps value to a specified time
-    pub fn set(&self, _now: DateTime<Utc>) {
-    }
+    pub fn set(&self, _now: DateTime<Utc>) {}
 
     /// return the timestamps value at the last time `record()` was
     /// called.
@@ -231,12 +215,10 @@ impl Timestamp {
     }
 
     /// sets the value of this timestamp to the minimum of this and other
-    pub fn update_to_min(&self, _other: &Timestamp) {
-    }
+    pub fn update_to_min(&self, _other: &Timestamp) {}
 
     /// sets the value of this timestamp to the maximum of this and other
-    pub fn update_to_max(&self, _other: &Timestamp) {
-    }
+    pub fn update_to_max(&self, _other: &Timestamp) {}
 }
 
 impl PartialEq for Timestamp {
diff --git a/datafusion/physical-plan/src/work_table.rs b/datafusion/physical-plan/src/work_table.rs
@@ -240,7 +240,8 @@ impl ExecutionPlan for WorkTableExec {
         self.metrics.reset();
         self.work_table.reset();
         Ok(())
-    }}
+    }
+}
 
 #[cfg(test)]
 mod tests {
diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs
@@ -938,10 +938,8 @@ impl protobuf::PhysicalPlanNode {
             protobuf::AggregateMode::SinglePartitioned => {
                 AggregateMode::SinglePartitioned
             }
-        protobuf::AggregateMode::CombinePartial => {
-                        AggregateMode::CombinePartial
-                    }
-                };
+            protobuf::AggregateMode::CombinePartial => AggregateMode::CombinePartial,
+        };
 
         let num_expr = hash_agg.group_expr.len();
 
@@ -2160,10 +2158,9 @@ impl protobuf::PhysicalPlanNode {
             AggregateMode::Single => protobuf::AggregateMode::Single,
             AggregateMode::SinglePartitioned => {
                 protobuf::AggregateMode::SinglePartitioned
-
-                }
-                AggregateMode::CombinePartial => protobuf::AggregateMode::CombinePartial,
-            };
+            }
+            AggregateMode::CombinePartial => protobuf::AggregateMode::CombinePartial,
+        };
         let input_schema = exec.input_schema();
         let input = protobuf::PhysicalPlanNode::try_from_physical_plan(
             exec.input().to_owned(),
diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs
