address comment

UBarney · UBarney · commit 71459283a93f · 2026-03-30T13:29:07.000Z
diff --git a/benchmarks/queries/clickbench/README.md b/benchmarks/queries/clickbench/README.md
@@ -228,6 +228,22 @@ Results look like
 Elapsed 30.195 seconds.
 ```
 
+
+### Q9-Q12: FIRST_VALUE Aggregation Performance
+
+These queries test the performance of the `FIRST_VALUE` aggregation function with different data types and grouping cardinalities.
+
+| Query | `FIRST_VALUE` Column | Column Type | Group By Column | Group By Type | Number of Groups |
+|-------|----------------------|-------------|-----------------|---------------|------------------|
+| Q9    | `URL`                | `Utf8`      | `UserID`        | `Int64`       | 17,630,976       |
+| Q10    | `URL`                | `Utf8`      | `OS`            | `Int16`       | 91               |
+| Q11   | `WatchID`            | `Int64`     | `UserID`        | `Int64`       | 17,630,976       |
+| Q12   | `WatchID`            | `Int64`     | `OS`            | `Int16`       | 91               |
+
+
+
+
+
 ## Data Notes
 
 Here are some interesting statistics about the data used in the queries
diff --git a/benchmarks/queries/clickbench/extended/q10.sql b/benchmarks/queries/clickbench/extended/q10.sql
@@ -0,0 +1,8 @@
+-- Must set for ClickBench hits_partitioned dataset. See https://github.com/apache/datafusion/issues/16591
+-- set datafusion.execution.parquet.binary_as_string = true
+
+SELECT MAX(len) FROM (
+    SELECT LENGTH(FIRST_VALUE("URL" ORDER BY "EventTime")) as len
+    FROM hits
+    GROUP BY "OS"
+);
diff --git a/benchmarks/queries/clickbench/extended/q11.sql b/benchmarks/queries/clickbench/extended/q11.sql
@@ -0,0 +1,8 @@
+-- Must set for ClickBench hits_partitioned dataset. See https://github.com/apache/datafusion/issues/16591
+-- set datafusion.execution.parquet.binary_as_string = true
+
+SELECT MAX(fv) FROM (
+    SELECT FIRST_VALUE("WatchID" ORDER BY "EventTime") as fv
+    FROM hits
+    GROUP BY "UserID"
+);
diff --git a/benchmarks/queries/clickbench/extended/q12.sql b/benchmarks/queries/clickbench/extended/q12.sql
@@ -0,0 +1,8 @@
+-- Must set for ClickBench hits_partitioned dataset. See https://github.com/apache/datafusion/issues/16591
+-- set datafusion.execution.parquet.binary_as_string = true
+
+SELECT MAX(fv) FROM (
+    SELECT FIRST_VALUE("WatchID" ORDER BY "EventTime") as fv
+    FROM hits
+    GROUP BY "OS"
+);
diff --git a/benchmarks/queries/clickbench/extended/q8.sql b/benchmarks/queries/clickbench/extended/q8.sql
@@ -0,0 +1,4 @@
+-- Must set for ClickBench hits_partitioned dataset. See https://github.com/apache/datafusion/issues/16591
+-- set datafusion.execution.parquet.binary_as_string = true
+
+SELECT "RegionID", "UserAgent", "OS", AVG(to_timestamp("ResponseEndTiming")-to_timestamp("ResponseStartTiming")) as avg_response_time, AVG(to_timestamp("ResponseEndTiming")-to_timestamp("ConnectTiming")) as avg_latency FROM hits GROUP BY "RegionID", "UserAgent", "OS" ORDER BY avg_latency DESC limit 10;
diff --git a/benchmarks/queries/clickbench/extended/q9.sql b/benchmarks/queries/clickbench/extended/q9.sql
@@ -0,0 +1,8 @@
+-- Must set for ClickBench hits_partitioned dataset. See https://github.com/apache/datafusion/issues/16591
+-- set datafusion.execution.parquet.binary_as_string = true
+
+SELECT MAX(len) FROM (
+    SELECT LENGTH(FIRST_VALUE("URL" ORDER BY "EventTime")) as len
+    FROM hits
+    GROUP BY "UserID"
+);
diff --git a/datafusion/functions-aggregate/src/first_last.rs b/datafusion/functions-aggregate/src/first_last.rs
@@ -23,9 +23,7 @@ use std::hash::Hash;
 use std::mem::size_of_val;
 use std::sync::Arc;
 
-use arrow::array::{
-    Array, ArrayRef, ArrowPrimitiveType, AsArray, BooleanArray, BooleanBufferBuilder,
-};
+use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, BooleanBufferBuilder};
 use arrow::buffer::BooleanBuffer;
 use arrow::compute::{self, LexicographicalComparator, SortColumn, SortOptions};
 use arrow::datatypes::{
@@ -79,31 +77,10 @@ pub fn last_value(expression: Expr, order_by: Vec<SortExpr>) -> Expr {
         .unwrap()
 }
 
-fn create_groups_primitive_accumulator<T: ArrowPrimitiveType + Send>(
-    args: &AccumulatorArgs,
-    is_first: bool,
-) -> Result<Box<dyn GroupsAccumulator>> {
-    let Some(ordering) = LexOrdering::new(args.order_bys.to_vec()) else {
-        return internal_err!("Groups accumulator must have an ordering.");
-    };
-
-    let ordering_dtypes = ordering
-        .iter()
-        .map(|e| e.expr.data_type(args.schema))
-        .collect::<Result<Vec<_>>>()?;
-
-    Ok(Box::new(FirstLastGroupsAccumulator::try_new(
-        PrimitiveValueState::<T>::new(args.return_field.data_type().clone()),
-        ordering,
-        args.ignore_nulls,
-        &ordering_dtypes,
-        is_first,
-    )?))
-}
-
-fn create_groups_bytes_accumulator(
+fn create_groups_accumulator_helper<S: ValueState + 'static>(
     args: &AccumulatorArgs,
     is_first: bool,
+    state: S,
 ) -> Result<Box<dyn GroupsAccumulator>> {
     let Some(ordering) = LexOrdering::new(args.order_bys.to_vec()) else {
         return internal_err!("Groups accumulator must have an ordering.");
@@ -115,7 +92,7 @@ fn create_groups_bytes_accumulator(
         .collect::<Result<Vec<_>>>()?;
 
     Ok(Box::new(FirstLastGroupsAccumulator::try_new(
-        BytesValueState::try_new(args.return_field.data_type().clone())?,
+        state,
         ordering,
         args.ignore_nulls,
         &ordering_dtypes,
@@ -128,99 +105,77 @@ fn create_groups_accumulator(
     is_first: bool,
     function_name: &str,
 ) -> Result<Box<dyn GroupsAccumulator>> {
-    match args.return_field.data_type() {
-        DataType::Int8 => create_groups_primitive_accumulator::<Int8Type>(args, is_first),
-        DataType::Int16 => {
-            create_groups_primitive_accumulator::<Int16Type>(args, is_first)
-        }
-        DataType::Int32 => {
-            create_groups_primitive_accumulator::<Int32Type>(args, is_first)
-        }
-        DataType::Int64 => {
-            create_groups_primitive_accumulator::<Int64Type>(args, is_first)
-        }
-        DataType::UInt8 => {
-            create_groups_primitive_accumulator::<UInt8Type>(args, is_first)
-        }
-        DataType::UInt16 => {
-            create_groups_primitive_accumulator::<UInt16Type>(args, is_first)
-        }
-        DataType::UInt32 => {
-            create_groups_primitive_accumulator::<UInt32Type>(args, is_first)
-        }
-        DataType::UInt64 => {
-            create_groups_primitive_accumulator::<UInt64Type>(args, is_first)
-        }
-        DataType::Float16 => {
-            create_groups_primitive_accumulator::<Float16Type>(args, is_first)
-        }
-        DataType::Float32 => {
-            create_groups_primitive_accumulator::<Float32Type>(args, is_first)
-        }
-        DataType::Float64 => {
-            create_groups_primitive_accumulator::<Float64Type>(args, is_first)
-        }
+    let data_type = args.return_field.data_type();
+
+    macro_rules! instantiate_primitive {
+        ($t:ty) => {
+            create_groups_accumulator_helper(
+                args,
+                is_first,
+                PrimitiveValueState::<$t>::new(data_type.clone()),
+            )
+        };
+    }
 
-        DataType::Decimal32(_, _) => {
-            create_groups_primitive_accumulator::<Decimal32Type>(args, is_first)
-        }
-        DataType::Decimal64(_, _) => {
-            create_groups_primitive_accumulator::<Decimal64Type>(args, is_first)
-        }
-        DataType::Decimal128(_, _) => {
-            create_groups_primitive_accumulator::<Decimal128Type>(args, is_first)
-        }
-        DataType::Decimal256(_, _) => {
-            create_groups_primitive_accumulator::<Decimal256Type>(args, is_first)
-        }
+    match data_type {
+        DataType::Int8 => instantiate_primitive!(Int8Type),
+        DataType::Int16 => instantiate_primitive!(Int16Type),
+        DataType::Int32 => instantiate_primitive!(Int32Type),
+        DataType::Int64 => instantiate_primitive!(Int64Type),
+        DataType::UInt8 => instantiate_primitive!(UInt8Type),
+        DataType::UInt16 => instantiate_primitive!(UInt16Type),
+        DataType::UInt32 => instantiate_primitive!(UInt32Type),
+        DataType::UInt64 => instantiate_primitive!(UInt64Type),
+        DataType::Float16 => instantiate_primitive!(Float16Type),
+        DataType::Float32 => instantiate_primitive!(Float32Type),
+        DataType::Float64 => instantiate_primitive!(Float64Type),
+
+        DataType::Decimal32(_, _) => instantiate_primitive!(Decimal32Type),
+        DataType::Decimal64(_, _) => instantiate_primitive!(Decimal64Type),
+        DataType::Decimal128(_, _) => instantiate_primitive!(Decimal128Type),
+        DataType::Decimal256(_, _) => instantiate_primitive!(Decimal256Type),
 
         DataType::Timestamp(TimeUnit::Second, _) => {
-            create_groups_primitive_accumulator::<TimestampSecondType>(args, is_first)
+            instantiate_primitive!(TimestampSecondType)
         }
         DataType::Timestamp(TimeUnit::Millisecond, _) => {
-            create_groups_primitive_accumulator::<TimestampMillisecondType>(
-                args, is_first,
-            )
+            instantiate_primitive!(TimestampMillisecondType)
         }
         DataType::Timestamp(TimeUnit::Microsecond, _) => {
-            create_groups_primitive_accumulator::<TimestampMicrosecondType>(
-                args, is_first,
-            )
+            instantiate_primitive!(TimestampMicrosecondType)
         }
         DataType::Timestamp(TimeUnit::Nanosecond, _) => {
-            create_groups_primitive_accumulator::<TimestampNanosecondType>(args, is_first)
+            instantiate_primitive!(TimestampNanosecondType)
         }
 
-        DataType::Date32 => {
-            create_groups_primitive_accumulator::<Date32Type>(args, is_first)
-        }
-        DataType::Date64 => {
-            create_groups_primitive_accumulator::<Date64Type>(args, is_first)
-        }
-        DataType::Time32(TimeUnit::Second) => {
-            create_groups_primitive_accumulator::<Time32SecondType>(args, is_first)
-        }
+        DataType::Date32 => instantiate_primitive!(Date32Type),
+        DataType::Date64 => instantiate_primitive!(Date64Type),
+        DataType::Time32(TimeUnit::Second) => instantiate_primitive!(Time32SecondType),
         DataType::Time32(TimeUnit::Millisecond) => {
-            create_groups_primitive_accumulator::<Time32MillisecondType>(args, is_first)
+            instantiate_primitive!(Time32MillisecondType)
         }
         DataType::Time64(TimeUnit::Microsecond) => {
-            create_groups_primitive_accumulator::<Time64MicrosecondType>(args, is_first)
+            instantiate_primitive!(Time64MicrosecondType)
         }
         DataType::Time64(TimeUnit::Nanosecond) => {
-            create_groups_primitive_accumulator::<Time64NanosecondType>(args, is_first)
+            instantiate_primitive!(Time64NanosecondType)
         }
 
         DataType::Utf8
         | DataType::LargeUtf8
         | DataType::Utf8View
         | DataType::Binary
         | DataType::LargeBinary
-        | DataType::BinaryView => create_groups_bytes_accumulator(args, is_first),
+        | DataType::BinaryView => create_groups_accumulator_helper(
+            args,
+            is_first,
+            BytesValueState::try_new(data_type.clone())?,
+        ),
 
         _ => internal_err!(
             "GroupsAccumulator not supported for {}({})",
             function_name,
-            args.return_field.data_type()
+            data_type
         ),
     }
 }
@@ -419,7 +374,7 @@ struct FirstLastGroupsAccumulator<S: ValueState> {
     // to avoid calling `ScalarValue::size_of_vec` by Self.size.
     size_of_orderings: usize,
 
-    // buffer for `get_filtered_min_of_each_group`
+    // buffer for `get_filtered_extreme_of_each_group`
     // filter_min_of_each_group_buf.0[group_idx] -> idx_in_val
     // only valid if filter_min_of_each_group_buf.1[group_idx] == true
     extreme_of_each_group_buf: (Vec<usize>, BooleanBufferBuilder),
diff --git a/datafusion/functions-aggregate/src/first_last/state.rs b/datafusion/functions-aggregate/src/first_last/state.rs
@@ -36,7 +36,6 @@ pub(crate) trait ValueState: Send + Sync {
     /// Note: While this is not a batch interface, it is not a performance bottleneck.
     /// In heavy aggregation benchmarks, the overhead of this method is typically less than 1%.
     ///
-    /// Benchmarked queries with < 1% `update` overhead:
     /// ```sql
     /// -- TPC-H SF10
     /// select l_shipmode, last_value(l_partkey order by l_orderkey, l_linenumber, l_comment, l_suppkey, l_tax)
@@ -151,6 +150,10 @@ impl ValueState for BytesValueState {
     }
 
     fn update(&mut self, group_idx: usize, array: &ArrayRef, idx: usize) -> Result<()> {
+        if let Some(v) = &self.vals[group_idx] {
+            self.total_capacity -= v.capacity();
+        }
+
         if array.is_null(idx) {
             self.vals[group_idx] = None;
         } else {
@@ -170,16 +173,16 @@ impl ValueState for BytesValueState {
             };
 
             if let Some(v) = &mut self.vals[group_idx] {
-                self.total_capacity -= v.capacity();
                 v.clear();
                 v.extend_from_slice(val);
-
-                self.total_capacity += v.capacity();
             } else {
                 let v = val.to_vec();
-                self.total_capacity += v.capacity();
                 self.vals[group_idx] = Some(v);
             }
+
+            self.vals[group_idx]
+                .as_ref()
+                .inspect(|x| self.total_capacity += x.capacity());
         }
         Ok(())
     }
@@ -436,4 +439,28 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_bytes_value_state_update_null() -> Result<()> {
+        let mut state = BytesValueState::try_new(DataType::Utf8)?;
+        state.resize(1);
+
+        let array: ArrayRef = Arc::new(StringArray::from(vec![Some("hello"), None]));
+
+        // group 0 = "hello"
+        state.update(0, &array, 0)?;
+        assert_eq!(state.total_capacity, state.total_capacity_calculated());
+        assert!(state.total_capacity > 0);
+
+        // group 0 = NULL
+        state.update(0, &array, 1)?;
+        assert_eq!(
+            state.total_capacity,
+            state.total_capacity_calculated(),
+            "total_capacity should match calculated capacity after update(NULL)"
+        );
+        assert_eq!(state.total_capacity, 0);
+
+        Ok(())
+    }
 }
diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt
