Handle new Arroyo syntax forms

Author: mwylde

Cargo.lock

@@ -166,11 +166,12 @@ recursive = "0.1.1"
 regex = "1.8"
 rstest = "0.25.0"
 serde_json = "1"
-sqlparser = { version = "0.55.0", features = ["visitor"] }
+sqlparser = { git = "https://github.com/ArroyoSystems/sqlparser-rs", branch = "0.55.0/arroyo", features = ["visitor"] }
 tempfile = "3"
 tokio = { version = "1.45", features = ["macros", "rt", "sync"] }
 url = "2.5.4"
 
+
 [profile.release]
 codegen-units = 1
 lto = true

Cargo.toml

@@ -134,7 +134,7 @@ impl Accumulator for GeometricMean {
     // This function serializes our state to `ScalarValue`, which DataFusion uses
     // to pass this state between execution stages.
     // Note that this can be arbitrary data.
-    fn state(&mut self) -> Result<Vec<ScalarValue>> {
+    fn state(&self) -> Result<Vec<ScalarValue>> {
         Ok(vec![
             ScalarValue::from(self.prod),
             ScalarValue::from(self.n),
@@ -143,7 +143,7 @@ impl Accumulator for GeometricMean {
 
     // DataFusion expects this function to return the final value of this aggregator.
     // in this case, this is the formula of the geometric mean
-    fn evaluate(&mut self) -> Result<ScalarValue> {
+    fn evaluate(&self) -> Result<ScalarValue> {
         let value = self.prod.powf(1.0 / self.n as f64);
         Ok(ScalarValue::from(value))
     }

datafusion-examples/examples/advanced_udaf.rs

@@ -71,7 +71,7 @@ impl Accumulator for GeometricMean {
     // This function serializes our state to `ScalarValue`, which DataFusion uses
     // to pass this state between execution stages.
     // Note that this can be arbitrary data.
-    fn state(&mut self) -> Result<Vec<ScalarValue>> {
+    fn state(&self) -> Result<Vec<ScalarValue>> {
         Ok(vec![
             ScalarValue::from(self.prod),
             ScalarValue::from(self.n),
@@ -80,7 +80,7 @@ impl Accumulator for GeometricMean {
 
     // DataFusion expects this function to return the final value of this aggregator.
     // in this case, this is the formula of the geometric mean
-    fn evaluate(&mut self) -> Result<ScalarValue> {
+    fn evaluate(&self) -> Result<ScalarValue> {
         let value = self.prod.powf(1.0 / self.n as f64);
         Ok(ScalarValue::from(value))
     }

datafusion-examples/examples/simple_udaf.rs

@@ -596,7 +596,7 @@ impl TimeSum {
 }
 
 impl Accumulator for TimeSum {
-    fn state(&mut self) -> Result<Vec<ScalarValue>> {
+    fn state(&self) -> Result<Vec<ScalarValue>> {
         Ok(vec![self.evaluate()?])
     }
 
@@ -617,7 +617,7 @@ impl Accumulator for TimeSum {
         self.update_batch(states)
     }
 
-    fn evaluate(&mut self) -> Result<ScalarValue> {
+    fn evaluate(&self) -> Result<ScalarValue> {
         Ok(ScalarValue::TimestampNanosecond(Some(self.sum), None))
     }
 
@@ -740,12 +740,12 @@ impl FirstSelector {
 }
 
 impl Accumulator for FirstSelector {
-    fn state(&mut self) -> Result<Vec<ScalarValue>> {
+    fn state(&self) -> Result<Vec<ScalarValue>> {
         self.evaluate().map(|s| vec![s])
     }
 
     /// produce the output structure
-    fn evaluate(&mut self) -> Result<ScalarValue> {
+    fn evaluate(&self) -> Result<ScalarValue> {
         self.to_state()
     }
 
@@ -838,15 +838,15 @@ impl Accumulator for TestGroupsAccumulator {
         Ok(())
     }
 
-    fn evaluate(&mut self) -> Result<ScalarValue> {
+    fn evaluate(&self) -> Result<ScalarValue> {
         Ok(ScalarValue::from(self.result))
     }
 
     fn size(&self) -> usize {
         size_of::<u64>()
     }
 
-    fn state(&mut self) -> Result<Vec<ScalarValue>> {
+    fn state(&self) -> Result<Vec<ScalarValue>> {
         Ok(vec![ScalarValue::from(self.result)])
     }
 
@@ -977,7 +977,7 @@ impl Accumulator for MetadataBasedAccumulator {
         Ok(())
     }
 
-    fn evaluate(&mut self) -> Result<ScalarValue> {
+    fn evaluate(&self) -> Result<ScalarValue> {
         let v = match self.double_output {
             true => self.curr_sum * 2,
             false => self.curr_sum,
@@ -990,7 +990,7 @@ impl Accumulator for MetadataBasedAccumulator {
         9
     }
 
-    fn state(&mut self) -> Result<Vec<ScalarValue>> {
+    fn state(&self) -> Result<Vec<ScalarValue>> {
         Ok(vec![ScalarValue::from(self.curr_sum)])
     }
 

datafusion/core/tests/user_defined/user_defined_aggregates.rs

@@ -345,7 +345,7 @@ async fn udaf_as_window_func() -> Result<()> {
     struct MyAccumulator;
 
     impl Accumulator for MyAccumulator {
-        fn state(&mut self) -> Result<Vec<ScalarValue>> {
+        fn state(&self) -> Result<Vec<ScalarValue>> {
             unimplemented!()
         }
 
@@ -357,7 +357,7 @@ async fn udaf_as_window_func() -> Result<()> {
             unimplemented!()
         }
 
-        fn evaluate(&mut self) -> Result<ScalarValue> {
+        fn evaluate(&self) -> Result<ScalarValue> {
             unimplemented!()
         }
 

datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

@@ -52,7 +52,8 @@ impl<O: OffsetSizeTrait> Accumulator for BytesDistinctCountAccumulator<O> {
     }
 
     fn state(&self) -> datafusion_common::Result<Vec<ScalarValue>> {
-        exec_err!("immutable state not supported for BytesDistinctCount")
+        let arr = self.0.as_state();
+        Ok(vec![SingleRowListArrayBuilder::new(arr).build_list_scalar()])
     }
 
     fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion_common::Result<()> {

datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs

@@ -1076,7 +1076,7 @@ mod tests {
 
     fn merge(
         mut acc1: Box<dyn Accumulator>,
-        mut acc2: Box<dyn Accumulator>,
+        acc2: Box<dyn Accumulator>,
     ) -> Result<Box<dyn Accumulator>> {
         let intermediate_state = acc2.state().and_then(|e| {
             e.iter()

datafusion/functions-aggregate/src/array_agg.rs

@@ -499,7 +499,7 @@ mod tests {
 
     fn merge(
         mut acc1: Box<dyn Accumulator>,
-        mut acc2: Box<dyn Accumulator>,
+        acc2: Box<dyn Accumulator>,
     ) -> Result<Box<dyn Accumulator>> {
         let intermediate_state = acc2.state().and_then(|e| {
             e.iter()

datafusion/functions-aggregate/src/string_agg.rs

@@ -79,6 +79,10 @@ impl<O: OffsetSizeTrait> ArrowBytesSet<O> {
         self.0.into_state()
     }
 
+    pub fn as_state(&self) -> ArrayRef {
+        self.0.as_state()
+    }
+
     /// Returns the total number of distinct values (including nulls) seen so far
     pub fn len(&self) -> usize {
         self.0.len()
@@ -526,6 +530,59 @@ where
         }
     }
 
+    /// Converts this set into a `StringArray`, `LargeStringArray`,
+    /// `BinaryArray`, or `LargeBinaryArray` containing each distinct value
+    /// that was inserted. This is done without copying the values.
+    ///
+    /// The values are guaranteed to be returned in the same order in which
+    /// they were first seen.
+    pub fn as_state(&self) -> ArrayRef {
+        let Self {
+            output_type,
+            map: _,
+            map_size: _,
+            offsets,
+            buffer,
+            random_state: _,
+            hashes_buffer: _,
+            null,
+        } = self;
+
+        // Only make a `NullBuffer` if there was a null value
+        let nulls = null.map(|(_payload, null_index)| {
+            let num_values = offsets.len() - 1;
+            single_null_buffer(num_values, null_index)
+        });
+        // 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 values = buffer.as_slice().iter().map(|b| *b).collect();
+
+        match output_type {
+            OutputType::Binary => {
+                // SAFETY: the offsets were constructed correctly
+                Arc::new(unsafe {
+                    GenericBinaryArray::new_unchecked(offsets, values, nulls)
+                })
+            }
+            OutputType::Utf8 => {
+                // SAFETY:
+                // 1. the offsets were constructed safely
+                //
+                // 2. we asserted the input arrays were all the correct type and
+                // thus since all the values that went in were valid (e.g. utf8)
+                // so are all the values that come out
+                Arc::new(unsafe {
+                    GenericStringArray::new_unchecked(offsets, values, nulls)
+                })
+            }
+            _ => unreachable!("View types should use `ArrowBytesViewMap`"),
+        }
+    }
+
+
     /// Total number of entries (including null, if present)
     pub fn len(&self) -> usize {
         self.non_null_len() + self.null.map(|_| 1).unwrap_or(0)