You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/04/04 10:00:55 UTC
(datafusion) branch gh-readonly-queue/main/pr-21090-3d2e6b2fd096a17f74224fcccccb4cedd1e38a82 created (now 587f4c0c74)
This is an automated email from the ASF dual-hosted git repository.
github-bot pushed a change to branch gh-readonly-queue/main/pr-21090-3d2e6b2fd096a17f74224fcccccb4cedd1e38a82
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 587f4c0c74 feat: support GroupsAccumulator for first_value and last_value with string/binary types (#21090)
This branch includes the following new commits:
new 587f4c0c74 feat: support GroupsAccumulator for first_value and last_value with string/binary types (#21090)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: feat: support GroupsAccumulator for first_value and last_value with string/binary types (#21090)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-bot pushed a commit to branch gh-readonly-queue/main/pr-21090-3d2e6b2fd096a17f74224fcccccb4cedd1e38a82
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 587f4c0c74350769c71a67bf3bb80d0917ec24e4
Author: UBarney <lv...@gmail.com>
AuthorDate: Sat Apr 4 18:00:34 2026 +0800
feat: support GroupsAccumulator for first_value and last_value with string/binary types (#21090)
## Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax. For example
`Closes #123` indicates that this PR will close issue #123.
-->
- Closes https://github.com/apache/datafusion/issues/17899.
## Rationale for this change
```
┏━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Query ┃ main ┃ first_val_group_acc ┃ Change ┃
┡━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ QQuery 0 │ 872.01 ms │ 904.35 ms │ no change │
│ QQuery 1 │ 156.37 ms │ 164.83 ms │ 1.05x slower │
│ QQuery 2 │ 448.58 ms │ 497.05 ms │ 1.11x slower │
│ QQuery 3 │ 233.99 ms │ 274.10 ms │ 1.17x slower │
│ QQuery 4 │ 1448.99 ms │ 1556.95 ms │ 1.07x slower │
│ QQuery 5 │ 10816.83 ms │ 11315.69 ms │ no change │
│ QQuery 6 │ 2053.16 ms │ 2030.02 ms │ no change │
│ QQuery 7 │ 2154.74 ms │ 2274.63 ms │ 1.06x slower │
│ QQuery 8 │ 405.62 ms │ 405.72 ms │ no change │
│ QQuery 9 │ 17160.65 ms │ 4167.34 ms │ +4.12x faster │
│ QQuery 10 │ 1206.03 ms │ 1090.69 ms │ +1.11x faster │
│ QQuery 11 │ 2437.12 ms │ 2446.51 ms │ no change │
│ QQuery 12 │ 331.27 ms │ 317.73 ms │ no change │
└───────────┴─────────────┴─────────────────────┴───────────────┘
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┓
┃ Benchmark Summary ┃ ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━┩
│ Total Time (main) │ 39725.35ms │
│ Total Time (first_val_group_acc) │ 27445.62ms │
│ Average Time (main) │ 3055.80ms │
│ Average Time (first_val_group_acc) │ 2111.20ms │
│ Queries Faster │ 2 │
│ Queries Slower │ 5 │
│ Queries with No Change │ 6 │
│ Queries with Failure │ 0 │
└────────────────────────────────────┴────────────┘
```
Previously, the `first_value` and `last_value` aggregate functions only
supported GroupsAccumulator for primitive types. For string or binary
types (Utf8, LargeUtf8, Binary, etc.), they fell back to the slower
row-based Accumulator path.
This change implements a specialized state management for byte-based
types, enabling high-performance grouped aggregation for strings and
binary data, especially when used with `ORDER BY`.
<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->
## What changes are included in this PR?
- New `ValueState` Trait: Abstracted the state management for
`first_value` and `last_value` to support different storage backends.
- `PrimitiveValueState` : Re-implemented the existing primitive handling
using the new trait.
- `BytesValueState`: Added a new state implementation for Utf8,
LargeUtf8, Utf8View, Binary, LargeBinary, and BinaryView. It
optimizes memory by reusing `Vec<u8>` buffers for group updates.
- Refactored `FirstLastGroupsAccumulato`r: Migrated the accumulator to
use the generic ValueState trait, allowing it to handle both primitive
and byte types uniformly.
<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->
## Are these changes tested?
YES
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
-->
## Are there any user-facing changes?
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
-->
<!--
If there are any breaking changes to public APIs, please add the `api
change` label.
-->
---------
Co-authored-by: Andrew Lamb <an...@nerdnetworks.org>
---
benchmarks/queries/clickbench/README.md | 16 +
benchmarks/queries/clickbench/extended/q10.sql | 8 +
benchmarks/queries/clickbench/extended/q11.sql | 8 +
benchmarks/queries/clickbench/extended/q12.sql | 8 +
benchmarks/queries/clickbench/extended/q8.sql | 4 +
benchmarks/queries/clickbench/extended/q9.sql | 8 +
datafusion/functions-aggregate/src/first_last.rs | 527 ++++++++-------------
.../functions-aggregate/src/first_last/state.rs | 466 ++++++++++++++++++
datafusion/sqllogictest/test_files/aggregate.slt | 118 +++++
9 files changed, 838 insertions(+), 325 deletions(-)
diff --git a/benchmarks/queries/clickbench/README.md b/benchmarks/queries/clickbench/README.md
index 877ea0e0c3..cead49586b 100644
--- 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
new file mode 100644
index 0000000000..d6019de178
--- /dev/null
+++ 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
new file mode 100644
index 0000000000..bca38f836b
--- /dev/null
+++ 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
new file mode 100644
index 0000000000..fa062ac1f5
--- /dev/null
+++ 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
new file mode 100644
index 0000000000..e580807841
--- /dev/null
+++ 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;
\ No newline at end of file
diff --git a/benchmarks/queries/clickbench/extended/q9.sql b/benchmarks/queries/clickbench/extended/q9.sql
new file mode 100644
index 0000000000..53952ebec2
--- /dev/null
+++ 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
index 91964eba1a..8e6bf9205b 100644
--- a/datafusion/functions-aggregate/src/first_last.rs
+++ b/datafusion/functions-aggregate/src/first_last.rs
@@ -22,11 +22,8 @@ use std::hash::Hash;
use std::mem::size_of_val;
use std::sync::Arc;
-use arrow::array::{
- Array, ArrayRef, ArrowPrimitiveType, AsArray, BooleanArray, BooleanBufferBuilder,
- PrimitiveArray,
-};
-use arrow::buffer::{BooleanBuffer, NullBuffer};
+use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, BooleanBufferBuilder};
+use arrow::buffer::BooleanBuffer;
use arrow::compute::{self, LexicographicalComparator, SortColumn, SortOptions};
use arrow::datatypes::{
DataType, Date32Type, Date64Type, Decimal32Type, Decimal64Type, Decimal128Type,
@@ -52,6 +49,10 @@ use datafusion_functions_aggregate_common::utils::get_sort_options;
use datafusion_macros::user_doc;
use datafusion_physical_expr_common::sort_expr::LexOrdering;
+mod state;
+
+use state::{BytesValueState, PrimitiveValueState, ValueState};
+
create_func!(FirstValue, first_value_udaf);
create_func!(LastValue, last_value_udaf);
@@ -75,6 +76,142 @@ pub fn last_value(expression: Expr, order_by: Vec<SortExpr>) -> Expr {
.unwrap()
}
+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.");
+ };
+
+ let ordering_dtypes = ordering
+ .iter()
+ .map(|e| e.expr.data_type(args.schema))
+ .collect::<Result<Vec<_>>>()?;
+
+ Ok(Box::new(FirstLastGroupsAccumulator::try_new(
+ state,
+ ordering,
+ args.ignore_nulls,
+ &ordering_dtypes,
+ is_first,
+ )?))
+}
+
+fn create_groups_accumulator(
+ args: &AccumulatorArgs,
+ is_first: bool,
+ function_name: &str,
+) -> Result<Box<dyn GroupsAccumulator>> {
+ 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()),
+ )
+ };
+ }
+
+ 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, _) => {
+ instantiate_primitive!(TimestampSecondType)
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ instantiate_primitive!(TimestampMillisecondType)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ instantiate_primitive!(TimestampMicrosecondType)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+ instantiate_primitive!(TimestampNanosecondType)
+ }
+
+ DataType::Date32 => instantiate_primitive!(Date32Type),
+ DataType::Date64 => instantiate_primitive!(Date64Type),
+ DataType::Time32(TimeUnit::Second) => instantiate_primitive!(Time32SecondType),
+ DataType::Time32(TimeUnit::Millisecond) => {
+ instantiate_primitive!(Time32MillisecondType)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ instantiate_primitive!(Time64MicrosecondType)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ instantiate_primitive!(Time64NanosecondType)
+ }
+
+ DataType::Utf8
+ | DataType::LargeUtf8
+ | DataType::Utf8View
+ | DataType::Binary
+ | DataType::LargeBinary
+ | DataType::BinaryView => create_groups_accumulator_helper(
+ args,
+ is_first,
+ BytesValueState::try_new(data_type.clone())?,
+ ),
+
+ _ => internal_err!(
+ "GroupsAccumulator not supported for {}({})",
+ function_name,
+ data_type
+ ),
+ }
+}
+
+fn groups_accumulator_supported(args: &AccumulatorArgs) -> bool {
+ use DataType::*;
+ !args.order_bys.is_empty()
+ && matches!(
+ args.return_field.data_type(),
+ Int8 | Int16
+ | Int32
+ | Int64
+ | UInt8
+ | UInt16
+ | UInt32
+ | UInt64
+ | Float16
+ | Float32
+ | Float64
+ | Decimal32(_, _)
+ | Decimal64(_, _)
+ | Decimal128(_, _)
+ | Decimal256(_, _)
+ | Date32
+ | Date64
+ | Time32(_)
+ | Time64(_)
+ | Timestamp(_, _)
+ | Utf8
+ | LargeUtf8
+ | Utf8View
+ | Binary
+ | LargeBinary
+ | BinaryView
+ )
+}
+
#[user_doc(
doc_section(label = "General Functions"),
description = "Returns the first element in an aggregation group according to the requested ordering. If no ordering is given, returns an arbitrary element from the group.",
@@ -178,110 +315,14 @@ impl AggregateUDFImpl for FirstValue {
}
fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool {
- use DataType::*;
- !args.order_bys.is_empty()
- && matches!(
- args.return_field.data_type(),
- Int8 | Int16
- | Int32
- | Int64
- | UInt8
- | UInt16
- | UInt32
- | UInt64
- | Float16
- | Float32
- | Float64
- | Decimal32(_, _)
- | Decimal64(_, _)
- | Decimal128(_, _)
- | Decimal256(_, _)
- | Date32
- | Date64
- | Time32(_)
- | Time64(_)
- | Timestamp(_, _)
- )
+ groups_accumulator_supported(&args)
}
fn create_groups_accumulator(
&self,
args: AccumulatorArgs,
) -> Result<Box<dyn GroupsAccumulator>> {
- fn create_accumulator<T: ArrowPrimitiveType + Send>(
- args: &AccumulatorArgs,
- ) -> 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<_>>>()?;
-
- FirstPrimitiveGroupsAccumulator::<T>::try_new(
- ordering,
- args.ignore_nulls,
- args.return_field.data_type(),
- &ordering_dtypes,
- true,
- )
- .map(|acc| Box::new(acc) as _)
- }
-
- match args.return_field.data_type() {
- DataType::Int8 => create_accumulator::<Int8Type>(&args),
- DataType::Int16 => create_accumulator::<Int16Type>(&args),
- DataType::Int32 => create_accumulator::<Int32Type>(&args),
- DataType::Int64 => create_accumulator::<Int64Type>(&args),
- DataType::UInt8 => create_accumulator::<UInt8Type>(&args),
- DataType::UInt16 => create_accumulator::<UInt16Type>(&args),
- DataType::UInt32 => create_accumulator::<UInt32Type>(&args),
- DataType::UInt64 => create_accumulator::<UInt64Type>(&args),
- DataType::Float16 => create_accumulator::<Float16Type>(&args),
- DataType::Float32 => create_accumulator::<Float32Type>(&args),
- DataType::Float64 => create_accumulator::<Float64Type>(&args),
-
- DataType::Decimal32(_, _) => create_accumulator::<Decimal32Type>(&args),
- DataType::Decimal64(_, _) => create_accumulator::<Decimal64Type>(&args),
- DataType::Decimal128(_, _) => create_accumulator::<Decimal128Type>(&args),
- DataType::Decimal256(_, _) => create_accumulator::<Decimal256Type>(&args),
-
- DataType::Timestamp(TimeUnit::Second, _) => {
- create_accumulator::<TimestampSecondType>(&args)
- }
- DataType::Timestamp(TimeUnit::Millisecond, _) => {
- create_accumulator::<TimestampMillisecondType>(&args)
- }
- DataType::Timestamp(TimeUnit::Microsecond, _) => {
- create_accumulator::<TimestampMicrosecondType>(&args)
- }
- DataType::Timestamp(TimeUnit::Nanosecond, _) => {
- create_accumulator::<TimestampNanosecondType>(&args)
- }
-
- DataType::Date32 => create_accumulator::<Date32Type>(&args),
- DataType::Date64 => create_accumulator::<Date64Type>(&args),
- DataType::Time32(TimeUnit::Second) => {
- create_accumulator::<Time32SecondType>(&args)
- }
- DataType::Time32(TimeUnit::Millisecond) => {
- create_accumulator::<Time32MillisecondType>(&args)
- }
-
- DataType::Time64(TimeUnit::Microsecond) => {
- create_accumulator::<Time64MicrosecondType>(&args)
- }
- DataType::Time64(TimeUnit::Nanosecond) => {
- create_accumulator::<Time64NanosecondType>(&args)
- }
-
- _ => internal_err!(
- "GroupsAccumulator not supported for first_value({})",
- args.return_field.data_type()
- ),
- }
+ create_groups_accumulator(&args, true, self.name())
}
fn with_beneficial_ordering(
@@ -311,13 +352,9 @@ impl AggregateUDFImpl for FirstValue {
}
}
-// TODO: rename to PrimitiveGroupsAccumulator
-struct FirstPrimitiveGroupsAccumulator<T>
-where
- T: ArrowPrimitiveType + Send,
-{
+struct FirstLastGroupsAccumulator<S: ValueState> {
// ================ state ===========
- vals: Vec<T::Native>,
+ state: S,
// Stores ordering values, of the aggregator requirement corresponding to first value
// of the aggregator.
// The `orderings` are stored row-wise, meaning that `orderings[group_idx]`
@@ -326,19 +363,16 @@ where
// At the beginning, `is_sets[group_idx]` is false, which means `first` is not seen yet.
// Once we see the first value, we set the `is_sets[group_idx]` flag
is_sets: BooleanBufferBuilder,
- // null_builder[group_idx] == false => vals[group_idx] is null
- null_builder: BooleanBufferBuilder,
// size of `self.orderings`
// Calculating the memory usage of `self.orderings` using `ScalarValue::size_of_vec` is quite costly.
// Therefore, we cache it and compute `size_of` only after each update
// 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
- // TODO: rename to extreme_of_each_group_buf
- min_of_each_group_buf: (Vec<usize>, BooleanBufferBuilder),
+ extreme_of_each_group_buf: (Vec<usize>, BooleanBufferBuilder),
// =========== option ============
@@ -351,19 +385,14 @@ where
sort_options: Vec<SortOptions>,
// Ignore null values.
ignore_nulls: bool,
- /// The output type
- data_type: DataType,
default_orderings: Vec<ScalarValue>,
}
-impl<T> FirstPrimitiveGroupsAccumulator<T>
-where
- T: ArrowPrimitiveType + Send,
-{
+impl<S: ValueState> FirstLastGroupsAccumulator<S> {
fn try_new(
+ state: S,
ordering_req: LexOrdering,
ignore_nulls: bool,
- data_type: &DataType,
ordering_dtypes: &[DataType],
pick_first_in_group: bool,
) -> Result<Self> {
@@ -375,17 +404,15 @@ where
let sort_options = get_sort_options(&ordering_req);
Ok(Self {
- null_builder: BooleanBufferBuilder::new(0),
ordering_req,
sort_options,
ignore_nulls,
default_orderings,
- data_type: data_type.clone(),
- vals: Vec::new(),
+ state,
orderings: Vec::new(),
is_sets: BooleanBufferBuilder::new(0),
size_of_orderings: 0,
- min_of_each_group_buf: (Vec::new(), BooleanBufferBuilder::new(0)),
+ extreme_of_each_group_buf: (Vec::new(), BooleanBufferBuilder::new(0)),
pick_first_in_group,
})
}
@@ -424,32 +451,8 @@ where
result
}
- fn take_need(
- bool_buf_builder: &mut BooleanBufferBuilder,
- emit_to: EmitTo,
- ) -> BooleanBuffer {
- let bool_buf = bool_buf_builder.finish();
- match emit_to {
- EmitTo::All => bool_buf,
- EmitTo::First(n) => {
- // split off the first N values in seen_values
- //
- // TODO make this more efficient rather than two
- // copies and bitwise manipulation
- let first_n: BooleanBuffer = bool_buf.iter().take(n).collect();
- // reset the existing buffer
- for b in bool_buf.iter().skip(n) {
- bool_buf_builder.append(b);
- }
- first_n
- }
- }
- }
-
fn resize_states(&mut self, new_size: usize) {
- self.vals.resize(new_size, T::default_value());
-
- self.null_builder.resize(new_size);
+ self.state.resize(new_size);
if self.orderings.len() < new_size {
let current_len = self.orderings.len();
@@ -468,44 +471,43 @@ where
self.is_sets.resize(new_size);
- self.min_of_each_group_buf.0.resize(new_size, 0);
- self.min_of_each_group_buf.1.resize(new_size);
+ self.extreme_of_each_group_buf.0.resize(new_size, 0);
+ self.extreme_of_each_group_buf.1.resize(new_size);
}
fn update_state(
&mut self,
group_idx: usize,
orderings: &[ScalarValue],
- new_val: T::Native,
- is_null: bool,
- ) {
- self.vals[group_idx] = new_val;
+ array: &ArrayRef,
+ idx: usize,
+ ) -> Result<()> {
+ self.state.update(group_idx, array, idx)?;
self.is_sets.set_bit(group_idx, true);
- self.null_builder.set_bit(group_idx, !is_null);
-
assert!(orderings.len() == self.ordering_req.len());
let old_size = ScalarValue::size_of_vec(&self.orderings[group_idx]);
self.orderings[group_idx].clear();
self.orderings[group_idx].extend_from_slice(orderings);
let new_size = ScalarValue::size_of_vec(&self.orderings[group_idx]);
self.size_of_orderings = self.size_of_orderings - old_size + new_size;
+ Ok(())
}
fn take_state(
&mut self,
emit_to: EmitTo,
- ) -> (ArrayRef, Vec<Vec<ScalarValue>>, BooleanBuffer) {
- emit_to.take_needed(&mut self.min_of_each_group_buf.0);
- self.min_of_each_group_buf
+ ) -> Result<(ArrayRef, Vec<Vec<ScalarValue>>, BooleanBuffer)> {
+ emit_to.take_needed(&mut self.extreme_of_each_group_buf.0);
+ self.extreme_of_each_group_buf
.1
- .truncate(self.min_of_each_group_buf.0.len());
+ .truncate(self.extreme_of_each_group_buf.0.len());
- (
- self.take_vals_and_null_buf(emit_to),
+ Ok((
+ self.state.take(emit_to)?,
self.take_orderings(emit_to),
- Self::take_need(&mut self.is_sets, emit_to),
- )
+ state::take_need(&mut self.is_sets, emit_to),
+ ))
}
// should be used in test only
@@ -519,20 +521,19 @@ where
/// Returns a vector of tuples `(group_idx, idx_in_val)` representing the index of the
/// minimum value in `orderings` for each group, using lexicographical comparison.
/// Values are filtered using `opt_filter` and `is_set_arr` if provided.
- /// TODO: rename to get_filtered_extreme_of_each_group
- fn get_filtered_min_of_each_group(
+ fn get_filtered_extreme_of_each_group(
&mut self,
orderings: &[ArrayRef],
group_indices: &[usize],
opt_filter: Option<&BooleanArray>,
- vals: &PrimitiveArray<T>,
+ vals: &ArrayRef,
is_set_arr: Option<&BooleanArray>,
) -> Result<Vec<(usize, usize)>> {
// Set all values in min_of_each_group_buf.1 to false.
- self.min_of_each_group_buf.1.truncate(0);
- self.min_of_each_group_buf
+ self.extreme_of_each_group_buf.1.truncate(0);
+ self.extreme_of_each_group_buf
.1
- .append_n(self.vals.len(), false);
+ .append_n(self.is_sets.len(), false);
// No need to call `clear` since `self.min_of_each_group_buf.0[group_idx]`
// is only valid when `self.min_of_each_group_buf.1[group_idx] == true`.
@@ -565,48 +566,35 @@ where
continue;
}
- let is_valid = self.min_of_each_group_buf.1.get_bit(group_idx);
+ let is_valid = self.extreme_of_each_group_buf.1.get_bit(group_idx);
if !is_valid {
- self.min_of_each_group_buf.1.set_bit(group_idx, true);
- self.min_of_each_group_buf.0[group_idx] = idx_in_val;
+ self.extreme_of_each_group_buf.1.set_bit(group_idx, true);
+ self.extreme_of_each_group_buf.0[group_idx] = idx_in_val;
} else {
let ordering = comparator
- .compare(self.min_of_each_group_buf.0[group_idx], idx_in_val);
+ .compare(self.extreme_of_each_group_buf.0[group_idx], idx_in_val);
if (ordering.is_gt() && self.pick_first_in_group)
|| (ordering.is_lt() && !self.pick_first_in_group)
{
- self.min_of_each_group_buf.0[group_idx] = idx_in_val;
+ self.extreme_of_each_group_buf.0[group_idx] = idx_in_val;
}
}
}
Ok(self
- .min_of_each_group_buf
+ .extreme_of_each_group_buf
.0
.iter()
.enumerate()
- .filter(|(group_idx, _)| self.min_of_each_group_buf.1.get_bit(*group_idx))
+ .filter(|(group_idx, _)| self.extreme_of_each_group_buf.1.get_bit(*group_idx))
.map(|(group_idx, idx_in_val)| (group_idx, *idx_in_val))
.collect::<Vec<_>>())
}
-
- fn take_vals_and_null_buf(&mut self, emit_to: EmitTo) -> ArrayRef {
- let r = emit_to.take_needed(&mut self.vals);
-
- let null_buf = NullBuffer::new(Self::take_need(&mut self.null_builder, emit_to));
-
- let values = PrimitiveArray::<T>::new(r.into(), Some(null_buf)) // no copy
- .with_data_type(self.data_type.clone());
- Arc::new(values)
- }
}
-impl<T> GroupsAccumulator for FirstPrimitiveGroupsAccumulator<T>
-where
- T: ArrowPrimitiveType + Send,
-{
+impl<S: ValueState + 'static> GroupsAccumulator for FirstLastGroupsAccumulator<S> {
fn update_batch(
&mut self,
// e.g. first_value(a order by b): values_and_order_cols will be [a, b]
@@ -617,13 +605,13 @@ where
) -> Result<()> {
self.resize_states(total_num_groups);
- let vals = values_and_order_cols[0].as_primitive::<T>();
+ let vals = &values_and_order_cols[0];
let mut ordering_buf = Vec::with_capacity(self.ordering_req.len());
// The overhead of calling `extract_row_at_idx_to_buf` is somewhat high, so we need to minimize its calls as much as possible.
for (group_idx, idx) in self
- .get_filtered_min_of_each_group(
+ .get_filtered_extreme_of_each_group(
&values_and_order_cols[1..],
group_indices,
opt_filter,
@@ -639,12 +627,7 @@ where
)?;
if self.should_update_state(group_idx, &ordering_buf)? {
- self.update_state(
- group_idx,
- &ordering_buf,
- vals.value(idx),
- vals.is_null(idx),
- );
+ self.update_state(group_idx, &ordering_buf, vals, idx)?;
}
}
@@ -652,11 +635,11 @@ where
}
fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
- Ok(self.take_state(emit_to).0)
+ Ok(self.take_state(emit_to)?.0)
}
fn state(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
- let (val_arr, orderings, is_sets) = self.take_state(emit_to);
+ let (val_arr, orderings, is_sets) = self.take_state(emit_to)?;
let mut result = Vec::with_capacity(self.orderings.len() + 2);
result.push(val_arr);
@@ -702,9 +685,9 @@ where
let is_set_arr = as_boolean_array(is_set_arr)?;
- let vals = values[0].as_primitive::<T>();
+ let vals = &values[0];
// The overhead of calling `extract_row_at_idx_to_buf` is somewhat high, so we need to minimize its calls as much as possible.
- let groups = self.get_filtered_min_of_each_group(
+ let groups = self.get_filtered_extreme_of_each_group(
&val_and_order_cols[1..],
group_indices,
opt_filter,
@@ -716,12 +699,7 @@ where
extract_row_at_idx_to_buf(&val_and_order_cols[1..], idx, &mut ordering_buf)?;
if self.should_update_state(group_idx, &ordering_buf)? {
- self.update_state(
- group_idx,
- &ordering_buf,
- vals.value(idx),
- vals.is_null(idx),
- );
+ self.update_state(group_idx, &ordering_buf, vals, idx)?;
}
}
@@ -729,12 +707,11 @@ where
}
fn size(&self) -> usize {
- self.vals.capacity() * size_of::<T::Native>()
- + self.null_builder.capacity() / 8 // capacity is in bits, so convert to bytes
- + self.is_sets.capacity() / 8
+ self.state.size()
+ + self.is_sets.capacity() / 8 // capacity is in bits, so convert to bytes
+ self.size_of_orderings
- + self.min_of_each_group_buf.0.capacity() * size_of::<usize>()
- + self.min_of_each_group_buf.1.capacity() / 8
+ + self.extreme_of_each_group_buf.0.capacity() * size_of::<usize>()
+ + self.extreme_of_each_group_buf.1.capacity() / 8
}
fn supports_convert_to_state(&self) -> bool {
@@ -1140,114 +1117,14 @@ impl AggregateUDFImpl for LastValue {
}
fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool {
- use DataType::*;
- !args.order_bys.is_empty()
- && matches!(
- args.return_field.data_type(),
- Int8 | Int16
- | Int32
- | Int64
- | UInt8
- | UInt16
- | UInt32
- | UInt64
- | Float16
- | Float32
- | Float64
- | Decimal32(_, _)
- | Decimal64(_, _)
- | Decimal128(_, _)
- | Decimal256(_, _)
- | Date32
- | Date64
- | Time32(_)
- | Time64(_)
- | Timestamp(_, _)
- )
+ groups_accumulator_supported(&args)
}
fn create_groups_accumulator(
&self,
args: AccumulatorArgs,
) -> Result<Box<dyn GroupsAccumulator>> {
- fn create_accumulator<T>(
- args: &AccumulatorArgs,
- ) -> Result<Box<dyn GroupsAccumulator>>
- where
- T: ArrowPrimitiveType + Send,
- {
- 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(FirstPrimitiveGroupsAccumulator::<T>::try_new(
- ordering,
- args.ignore_nulls,
- args.return_field.data_type(),
- &ordering_dtypes,
- false,
- )?))
- }
-
- match args.return_field.data_type() {
- DataType::Int8 => create_accumulator::<Int8Type>(&args),
- DataType::Int16 => create_accumulator::<Int16Type>(&args),
- DataType::Int32 => create_accumulator::<Int32Type>(&args),
- DataType::Int64 => create_accumulator::<Int64Type>(&args),
- DataType::UInt8 => create_accumulator::<UInt8Type>(&args),
- DataType::UInt16 => create_accumulator::<UInt16Type>(&args),
- DataType::UInt32 => create_accumulator::<UInt32Type>(&args),
- DataType::UInt64 => create_accumulator::<UInt64Type>(&args),
- DataType::Float16 => create_accumulator::<Float16Type>(&args),
- DataType::Float32 => create_accumulator::<Float32Type>(&args),
- DataType::Float64 => create_accumulator::<Float64Type>(&args),
-
- DataType::Decimal32(_, _) => create_accumulator::<Decimal32Type>(&args),
- DataType::Decimal64(_, _) => create_accumulator::<Decimal64Type>(&args),
- DataType::Decimal128(_, _) => create_accumulator::<Decimal128Type>(&args),
- DataType::Decimal256(_, _) => create_accumulator::<Decimal256Type>(&args),
-
- DataType::Timestamp(TimeUnit::Second, _) => {
- create_accumulator::<TimestampSecondType>(&args)
- }
- DataType::Timestamp(TimeUnit::Millisecond, _) => {
- create_accumulator::<TimestampMillisecondType>(&args)
- }
- DataType::Timestamp(TimeUnit::Microsecond, _) => {
- create_accumulator::<TimestampMicrosecondType>(&args)
- }
- DataType::Timestamp(TimeUnit::Nanosecond, _) => {
- create_accumulator::<TimestampNanosecondType>(&args)
- }
-
- DataType::Date32 => create_accumulator::<Date32Type>(&args),
- DataType::Date64 => create_accumulator::<Date64Type>(&args),
- DataType::Time32(TimeUnit::Second) => {
- create_accumulator::<Time32SecondType>(&args)
- }
- DataType::Time32(TimeUnit::Millisecond) => {
- create_accumulator::<Time32MillisecondType>(&args)
- }
-
- DataType::Time64(TimeUnit::Microsecond) => {
- create_accumulator::<Time64MicrosecondType>(&args)
- }
- DataType::Time64(TimeUnit::Nanosecond) => {
- create_accumulator::<Time64NanosecondType>(&args)
- }
-
- _ => {
- internal_err!(
- "GroupsAccumulator not supported for last_value({})",
- args.return_field.data_type()
- )
- }
- }
+ create_groups_accumulator(&args, false, self.name())
}
}
@@ -1541,7 +1418,7 @@ mod tests {
use std::iter::repeat_with;
use arrow::{
- array::{BooleanArray, Int64Array, ListArray, StringArray},
+ array::{BooleanArray, Int64Array, ListArray, PrimitiveArray, StringArray},
compute::SortOptions,
datatypes::Schema,
};
@@ -1668,10 +1545,10 @@ mod tests {
options: SortOptions::default(),
}];
- let mut group_acc = FirstPrimitiveGroupsAccumulator::<Int64Type>::try_new(
+ let mut group_acc = FirstLastGroupsAccumulator::try_new(
+ PrimitiveValueState::<Int64Type>::new(DataType::Int64),
sort_keys.into(),
true,
- &DataType::Int64,
&[DataType::Int64],
true,
)?;
@@ -1762,10 +1639,10 @@ mod tests {
options: SortOptions::default(),
}];
- let mut group_acc = FirstPrimitiveGroupsAccumulator::<Int64Type>::try_new(
+ let mut group_acc = FirstLastGroupsAccumulator::try_new(
+ PrimitiveValueState::<Int64Type>::new(DataType::Int64),
sort_keys.into(),
true,
- &DataType::Int64,
&[DataType::Int64],
true,
)?;
@@ -1843,10 +1720,10 @@ mod tests {
options: SortOptions::default(),
}];
- let mut group_acc = FirstPrimitiveGroupsAccumulator::<Int64Type>::try_new(
+ let mut group_acc = FirstLastGroupsAccumulator::try_new(
+ PrimitiveValueState::<Int64Type>::new(DataType::Int64),
sort_keys.into(),
true,
- &DataType::Int64,
&[DataType::Int64],
false,
)?;
diff --git a/datafusion/functions-aggregate/src/first_last/state.rs b/datafusion/functions-aggregate/src/first_last/state.rs
new file mode 100644
index 0000000000..3aa578c74c
--- /dev/null
+++ b/datafusion/functions-aggregate/src/first_last/state.rs
@@ -0,0 +1,466 @@
+// 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.
+
+use std::mem::size_of;
+use std::sync::Arc;
+
+use arrow::array::{
+ Array, ArrayRef, ArrowPrimitiveType, AsArray, BinaryBuilder, BinaryViewBuilder,
+ BooleanBufferBuilder, LargeBinaryBuilder, LargeStringBuilder, PrimitiveArray,
+ StringBuilder, StringViewBuilder,
+};
+use arrow::buffer::{BooleanBuffer, NullBuffer};
+use arrow::datatypes::DataType;
+use datafusion_common::{Result, internal_err};
+use datafusion_expr::EmitTo;
+
+pub(crate) trait ValueState: Send + Sync {
+ /// Resizes the state to accommodate `new_size` groups.
+ fn resize(&mut self, new_size: usize);
+ /// Updates the state for the specified `group_idx` using the value at `idx` from the provided `array`.
+ ///
+ /// 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%.
+ ///
+ /// ```sql
+ /// -- TPC-H SF10
+ /// select l_shipmode, last_value(l_partkey order by l_orderkey, l_linenumber, l_comment, l_suppkey, l_tax)
+ /// from 'benchmarks/data/tpch_sf10/lineitem'
+ /// group by l_shipmode;
+ ///
+ /// -- H2O G1_1e8
+ /// select t.id1, first_value(t.id3 order by t.id2, t.id4) as r2
+ /// from 'benchmarks/data/h2o/G1_1e8_1e8_100_0.parquet' as t
+ /// group by t.id1, t.v1;
+ /// ```
+ fn update(&mut self, group_idx: usize, array: &ArrayRef, idx: usize) -> Result<()>;
+ /// Takes the accumulated state and returns it as an [`ArrayRef`], respecting the `emit_to` strategy.
+ fn take(&mut self, emit_to: EmitTo) -> Result<ArrayRef>;
+ /// Returns the estimated memory size of the state in bytes.
+ fn size(&self) -> usize;
+}
+
+pub(crate) struct PrimitiveValueState<T: ArrowPrimitiveType + Send> {
+ /// Values data
+ vals: Vec<T::Native>,
+ nulls: BooleanBufferBuilder,
+ data_type: DataType,
+}
+
+impl<T: ArrowPrimitiveType + Send> PrimitiveValueState<T> {
+ pub(crate) fn new(data_type: DataType) -> Self {
+ Self {
+ vals: vec![],
+ nulls: BooleanBufferBuilder::new(0),
+ data_type,
+ }
+ }
+}
+
+impl<T: ArrowPrimitiveType + Send> ValueState for PrimitiveValueState<T> {
+ fn resize(&mut self, new_size: usize) {
+ self.vals.resize(new_size, T::default_value());
+ self.nulls.resize(new_size);
+ }
+
+ fn update(&mut self, group_idx: usize, array: &ArrayRef, idx: usize) -> Result<()> {
+ let array = array.as_primitive::<T>();
+ self.vals[group_idx] = array.value(idx);
+ self.nulls.set_bit(group_idx, !array.is_null(idx));
+ Ok(())
+ }
+
+ fn take(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
+ let values = emit_to.take_needed(&mut self.vals);
+ let null_buf = NullBuffer::new(take_need(&mut self.nulls, emit_to));
+ let array: PrimitiveArray<T> =
+ PrimitiveArray::<T>::new(values.into(), Some(null_buf))
+ .with_data_type(self.data_type.clone());
+ Ok(Arc::new(array))
+ }
+
+ fn size(&self) -> usize {
+ self.vals.capacity() * size_of::<T::Native>() + self.nulls.capacity() / 8
+ }
+}
+
+/// Stores internal state for "bytes" types (Utf8, Binary, etc.).
+///
+/// This implementation is similar to `MinMaxBytesState` in `min_max_bytes.rs`, but
+/// it does not reuse it for two main reasons:
+///
+/// 1. **Direct Overwrite**: `MinMaxBytesState::update_batch` is tightly coupled
+/// with min/max comparison logic, whereas `FirstLast` performs its own comparisons
+/// externally (using ordering columns) and only needs a simple interface to
+/// unconditionally set/overwrite values for specific groups.
+/// 2. **Different NULL Handling**: `MinMaxBytesState` always ignores `NULL` values
+/// in the input, while `BytesValueState` needs to support setting `NULL` values
+/// to correctly implement `RESPECT NULLS` behavior.
+///
+pub(crate) struct BytesValueState {
+ vals: Vec<Option<Vec<u8>>>,
+ data_type: DataType,
+ /// The sum of the capacities of all vectors in `vals`.
+ total_capacity: usize,
+}
+
+impl BytesValueState {
+ pub(crate) fn try_new(data_type: DataType) -> Result<Self> {
+ if !matches!(
+ data_type,
+ DataType::Utf8
+ | DataType::LargeUtf8
+ | DataType::Utf8View
+ | DataType::Binary
+ | DataType::LargeBinary
+ | DataType::BinaryView
+ ) {
+ return internal_err!("BytesValueState does not support {}", data_type);
+ }
+ Ok(Self {
+ vals: vec![],
+ data_type,
+ total_capacity: 0,
+ })
+ }
+}
+
+impl ValueState for BytesValueState {
+ fn resize(&mut self, new_size: usize) {
+ if new_size < self.vals.len() {
+ for v in self.vals[new_size..].iter().flatten() {
+ self.total_capacity -= v.capacity();
+ }
+ }
+ self.vals.resize(new_size, None);
+ }
+
+ 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 {
+ let val = match self.data_type {
+ DataType::Utf8 => array.as_string::<i32>().value(idx).as_bytes(),
+ DataType::LargeUtf8 => array.as_string::<i64>().value(idx).as_bytes(),
+ DataType::Utf8View => array.as_string_view().value(idx).as_bytes(),
+ DataType::Binary => array.as_binary::<i32>().value(idx),
+ DataType::LargeBinary => array.as_binary::<i64>().value(idx),
+ DataType::BinaryView => array.as_binary_view().value(idx),
+ _ => {
+ return internal_err!(
+ "Unsupported data type for BytesValueState: {}",
+ self.data_type
+ );
+ }
+ };
+
+ if let Some(v) = &mut self.vals[group_idx] {
+ v.clear();
+ v.extend_from_slice(val);
+ } else {
+ let v = val.to_vec();
+ self.vals[group_idx] = Some(v);
+ }
+
+ self.vals[group_idx]
+ .as_ref()
+ .inspect(|x| self.total_capacity += x.capacity());
+ }
+ Ok(())
+ }
+
+ fn take(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
+ let values = emit_to.take_needed(&mut self.vals);
+
+ let (total_len, taken_capacity) = values
+ .iter()
+ .flatten()
+ .fold((0, 0), |(len_acc, cap_acc), v| {
+ (len_acc + v.len(), cap_acc + v.capacity())
+ });
+ self.total_capacity -= taken_capacity;
+
+ match self.data_type {
+ DataType::Utf8 => {
+ let mut builder = StringBuilder::with_capacity(values.len(), total_len);
+ for val in values {
+ match val {
+ Some(v) => builder.append_value(
+ // SAFETY: The bytes were originally from a valid UTF-8 array in `update`
+ unsafe { std::str::from_utf8_unchecked(&v) },
+ ),
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()))
+ }
+ DataType::LargeUtf8 => {
+ let mut builder =
+ LargeStringBuilder::with_capacity(values.len(), total_len);
+ for val in values {
+ match val {
+ Some(v) => builder.append_value(
+ // SAFETY: The bytes were originally from a valid UTF-8 array in `update`
+ unsafe { std::str::from_utf8_unchecked(&v) },
+ ),
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()))
+ }
+ DataType::Utf8View => {
+ let mut builder = StringViewBuilder::with_capacity(values.len());
+ for val in values {
+ match val {
+ Some(v) => builder.append_value(
+ // SAFETY: The bytes were originally from a valid UTF-8 array in `update`
+ unsafe { std::str::from_utf8_unchecked(&v) },
+ ),
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()))
+ }
+ DataType::Binary => {
+ let mut builder = BinaryBuilder::with_capacity(values.len(), total_len);
+ for val in values {
+ match val {
+ Some(v) => builder.append_value(&v),
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()))
+ }
+ DataType::LargeBinary => {
+ let mut builder =
+ LargeBinaryBuilder::with_capacity(values.len(), total_len);
+ for val in values {
+ match val {
+ Some(v) => builder.append_value(&v),
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()))
+ }
+ DataType::BinaryView => {
+ let mut builder = BinaryViewBuilder::with_capacity(values.len());
+ for val in values {
+ match val {
+ Some(v) => builder.append_value(&v),
+ None => builder.append_null(),
+ }
+ }
+ Ok(Arc::new(builder.finish()))
+ }
+ _ => internal_err!(
+ "Unsupported data type for BytesValueState: {}",
+ self.data_type
+ ),
+ }
+ }
+
+ fn size(&self) -> usize {
+ self.vals.capacity() * size_of::<Option<Vec<u8>>>() + self.total_capacity
+ }
+}
+
+impl BytesValueState {
+ #[cfg(test)]
+ /// For testing only: strictly calculate the sum of capacities of all vectors in `vals`.
+ fn total_capacity_calculated(&self) -> usize {
+ self.vals.iter().flatten().map(|v| v.capacity()).sum()
+ }
+}
+
+pub(crate) fn take_need(
+ bool_buf_builder: &mut BooleanBufferBuilder,
+ emit_to: EmitTo,
+) -> BooleanBuffer {
+ let bool_buf = bool_buf_builder.finish();
+ match emit_to {
+ EmitTo::All => bool_buf,
+ EmitTo::First(n) => {
+ // split off the first N values in seen_values
+ //
+ // TODO make this more efficient rather than two
+ // copies and bitwise manipulation
+ let first_n: BooleanBuffer = bool_buf.iter().take(n).collect();
+ // reset the existing buffer
+ for b in bool_buf.iter().skip(n) {
+ bool_buf_builder.append(b);
+ }
+ first_n
+ }
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow::array::{
+ BinaryArray, BinaryViewArray, LargeBinaryArray, LargeStringArray, StringArray,
+ StringViewArray,
+ };
+
+ #[test]
+ fn test_bytes_value_state_utf8() -> Result<()> {
+ let mut state = BytesValueState::try_new(DataType::Utf8)?;
+ state.resize(2);
+
+ let array: ArrayRef = Arc::new(StringArray::from(vec![
+ Some("hello"),
+ Some("world"),
+ Some("longer_string_than_hello"),
+ ]));
+
+ state.update(0, &array, 0)?; // group 0 = "hello"
+ state.update(1, &array, 1)?; // group 1 = "world"
+
+ assert_eq!(state.total_capacity, state.total_capacity_calculated());
+
+ // Overwrite group 0 with a longer string (checks capacity update logic)
+ state.update(0, &array, 2)?;
+ assert_eq!(state.total_capacity, state.total_capacity_calculated());
+
+ let result = state.take(EmitTo::All)?;
+ let result = result.as_string::<i32>();
+ assert_eq!(result.len(), 2);
+ assert_eq!(result.value(0), "longer_string_than_hello");
+ assert_eq!(result.value(1), "world");
+
+ // After take all, size should be 0 (excluding vals vector capacity)
+ assert_eq!(state.total_capacity, 0);
+ assert_eq!(state.total_capacity, state.total_capacity_calculated());
+
+ Ok(())
+ }
+
+ #[test]
+ fn test_bytes_value_state_large_utf8() -> Result<()> {
+ let mut state = BytesValueState::try_new(DataType::LargeUtf8)?;
+ state.resize(1);
+ let array: ArrayRef = Arc::new(LargeStringArray::from(vec!["large_utf8"]));
+ state.update(0, &array, 0)?;
+ let result = state.take(EmitTo::All)?;
+ assert_eq!(result.as_string::<i64>().value(0), "large_utf8");
+ Ok(())
+ }
+
+ #[test]
+ fn test_bytes_value_state_utf8_view() -> Result<()> {
+ let mut state = BytesValueState::try_new(DataType::Utf8View)?;
+ state.resize(1);
+ let array: ArrayRef = Arc::new(StringViewArray::from(vec!["Utf8View"]));
+ state.update(0, &array, 0)?;
+ let result = state.take(EmitTo::All)?;
+ assert_eq!(result.as_string_view().value(0), "Utf8View");
+ Ok(())
+ }
+
+ #[test]
+ fn test_bytes_value_state_binary() -> Result<()> {
+ let mut state = BytesValueState::try_new(DataType::Binary)?;
+ state.resize(1);
+ let array: ArrayRef = Arc::new(BinaryArray::from(vec![b"binary" as &[u8]]));
+ state.update(0, &array, 0)?;
+ let result = state.take(EmitTo::All)?;
+ assert_eq!(result.as_binary::<i32>().value(0), b"binary");
+ Ok(())
+ }
+
+ #[test]
+ fn test_bytes_value_state_large_binary() -> Result<()> {
+ let mut state = BytesValueState::try_new(DataType::LargeBinary)?;
+ state.resize(1);
+ let array: ArrayRef =
+ Arc::new(LargeBinaryArray::from(vec![b"large_binary" as &[u8]]));
+ state.update(0, &array, 0)?;
+ let result = state.take(EmitTo::All)?;
+ assert_eq!(result.as_binary::<i64>().value(0), b"large_binary");
+ Ok(())
+ }
+
+ #[test]
+ fn test_bytes_value_state_binary_view() -> Result<()> {
+ let mut state = BytesValueState::try_new(DataType::BinaryView)?;
+ state.resize(1);
+
+ let data: Vec<Option<&[u8]>> = vec![Some(b"long_binary_value_to_test_view")];
+ let array: ArrayRef = Arc::new(BinaryViewArray::from(data));
+
+ state.update(0, &array, 0)?;
+
+ let result = state.take(EmitTo::All)?;
+ let result = result.as_binary_view();
+ assert_eq!(result.value(0), b"long_binary_value_to_test_view");
+
+ Ok(())
+ }
+
+ #[test]
+ fn test_bytes_value_state_emit_first() -> Result<()> {
+ let mut state = BytesValueState::try_new(DataType::Utf8)?;
+ state.resize(3);
+
+ let array: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c"]));
+ state.update(0, &array, 0)?;
+ state.update(1, &array, 1)?;
+ state.update(2, &array, 2)?;
+
+ let result = state.take(EmitTo::First(2))?;
+ let result = result.as_string::<i32>();
+ assert_eq!(result.len(), 2);
+ assert_eq!(result.value(0), "a");
+ assert_eq!(result.value(1), "b");
+
+ // Remaining should be "c"
+ let result = state.take(EmitTo::All)?;
+ let result = result.as_string::<i32>();
+ assert_eq!(result.len(), 1);
+ assert_eq!(result.value(0), "c");
+
+ 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
index 4eb33f0cb5..0232367163 100644
--- a/datafusion/sqllogictest/test_files/aggregate.slt
+++ b/datafusion/sqllogictest/test_files/aggregate.slt
@@ -8817,3 +8817,121 @@ set datafusion.execution.target_partitions = 4;
statement ok
DROP TABLE stream_test;
+
+
+#################
+# first_value on strings/binary with groups and ordering
+#################
+
+statement ok
+CREATE TABLE first_last_value_str_tests(id INT, sort_key INT, val TEXT) AS VALUES
+(1, 10, 'apple'),
+(1, 2, 'banana'),
+(1, 5, 'cherry'),
+(2, 100, 'dog'),
+(2, 200, 'elephant');
+
+# Utf8
+query IT
+SELECT id, first_value(val ORDER BY sort_key)
+FROM (SELECT id, sort_key, arrow_cast(val, 'Utf8') as val FROM first_last_value_str_tests)
+GROUP BY id ORDER BY id;
+----
+1 banana
+2 dog
+
+query IT
+SELECT id, last_value(val ORDER BY sort_key)
+FROM (SELECT id, sort_key, arrow_cast(val, 'Utf8') as val FROM first_last_value_str_tests)
+GROUP BY id ORDER BY id;
+----
+1 apple
+2 elephant
+
+# LargeUtf8
+query IT
+SELECT id, first_value(val ORDER BY sort_key)
+FROM (SELECT id, sort_key, arrow_cast(val, 'LargeUtf8') as val FROM first_last_value_str_tests)
+GROUP BY id ORDER BY id;
+----
+1 banana
+2 dog
+
+query IT
+SELECT id, last_value(val ORDER BY sort_key)
+FROM (SELECT id, sort_key, arrow_cast(val, 'LargeUtf8') as val FROM first_last_value_str_tests)
+GROUP BY id ORDER BY id;
+----
+1 apple
+2 elephant
+
+# Utf8View
+query IT
+SELECT id, first_value(val ORDER BY sort_key)
+FROM (SELECT id, sort_key, arrow_cast(val, 'Utf8View') as val FROM first_last_value_str_tests)
+GROUP BY id ORDER BY id;
+----
+1 banana
+2 dog
+
+query IT
+SELECT id, last_value(val ORDER BY sort_key)
+FROM (SELECT id, sort_key, arrow_cast(val, 'Utf8View') as val FROM first_last_value_str_tests)
+GROUP BY id ORDER BY id;
+----
+1 apple
+2 elephant
+
+# Binary
+query I?
+SELECT id, first_value(val ORDER BY sort_key)
+FROM (SELECT id, sort_key, arrow_cast(val, 'Binary') as val FROM first_last_value_str_tests)
+GROUP BY id ORDER BY id;
+----
+1 62616e616e61
+2 646f67
+
+query I?
+SELECT id, last_value(val ORDER BY sort_key)
+FROM (SELECT id, sort_key, arrow_cast(val, 'Binary') as val FROM first_last_value_str_tests)
+GROUP BY id ORDER BY id;
+----
+1 6170706c65
+2 656c657068616e74
+
+# LargeBinary
+query I?
+SELECT id, first_value(val ORDER BY sort_key)
+FROM (SELECT id, sort_key, arrow_cast(val, 'LargeBinary') as val FROM first_last_value_str_tests)
+GROUP BY id ORDER BY id;
+----
+1 62616e616e61
+2 646f67
+
+query I?
+SELECT id, last_value(val ORDER BY sort_key)
+FROM (SELECT id, sort_key, arrow_cast(val, 'LargeBinary') as val FROM first_last_value_str_tests)
+GROUP BY id ORDER BY id;
+----
+1 6170706c65
+2 656c657068616e74
+
+# BinaryView
+query I?
+SELECT id, first_value(val ORDER BY sort_key)
+FROM (SELECT id, sort_key, arrow_cast(arrow_cast(val, 'Binary'), 'BinaryView') as val FROM first_last_value_str_tests)
+GROUP BY id ORDER BY id;
+----
+1 62616e616e61
+2 646f67
+
+query I?
+SELECT id, last_value(val ORDER BY sort_key)
+FROM (SELECT id, sort_key, arrow_cast(arrow_cast(val, 'Binary'), 'BinaryView') as val FROM first_last_value_str_tests)
+GROUP BY id ORDER BY id;
+----
+1 6170706c65
+2 656c657068616e74
+
+statement ok
+DROP TABLE first_last_value_str_tests;
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org