You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "gabotechs (via GitHub)" <gi...@apache.org> on 2025/02/02 17:27:40 UTC
[PR] String agg missing functionality [datafusion]
gabotechs opened a new pull request, #14412:
URL: https://github.com/apache/datafusion/pull/14412
## Which issue does this PR close?
Closes #8260.
## Rationale for this change
Complete the missing functionality of the STRING_AGG function.
## What changes are included in this PR?
Adds support for DISTINCT and ORDER_BY clauses by reusing the existing ARRAY_AGG functionality and building the whole STRING_AGG aggregation function on top of it. This way, the full STRING_AGG functionality is automatically implemented [almost] for free.
The rationale for reusing the ARRAY_AGG functionality is because both functions are very similar, with just two minor diferences:
- STRING_AGG works only with strings, while ARRAY_AGG works with any type.
- The return of STRING_AGG is the same as ARRAY_AGG, but with the resulting array of strings joined by a delimiter.
In order to have the full STRING_AGG functionality, some small addition is also needed for the ARRAY_AGG function, as the current implementation is missing support for DISTINCT + ORDER BY. See <insert PR>.
## Are these changes tested?
Yes, both in unit tests and sqllogictests.
## Are there any user-facing changes?
Users will be able to issue STRING_AGG calls with DISTINCT and ORDER BY clauses.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] STRING_AGG missing functionality [datafusion]
Posted by "geoffreyclaude (via GitHub)" <gi...@apache.org>.
geoffreyclaude commented on code in PR #14412:
URL: https://github.com/apache/datafusion/pull/14412#discussion_r1962046382
##########
datafusion/functions-aggregate/src/string_agg.rs:
##########
@@ -106,20 +113,40 @@ impl AggregateUDFImpl for StringAgg {
Ok(DataType::LargeUtf8)
}
+ fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<Field>> {
+ self.array_agg.state_fields(args)
+ }
+
fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
- if let Some(lit) = acc_args.exprs[1].as_any().downcast_ref::<Literal>() {
- return match lit.value().try_as_str() {
- Some(Some(delimiter)) => {
- Ok(Box::new(StringAggAccumulator::new(delimiter)))
- }
- Some(None) => Ok(Box::new(StringAggAccumulator::new(""))),
- None => {
- not_impl_err!("StringAgg not supported for delimiter {}", lit.value())
- }
- };
- }
+ let Some(lit) = acc_args.exprs[1].as_any().downcast_ref::<Literal>() else {
+ return not_impl_err!(
+ "The second argument of the string_agg function must be a string literal"
+ );
+ };
+
+ let delimiter = if lit.value().is_null() {
+ // If the second argument (the delimiter that joins strings) is NULL, join
+ // on an empty string. (e.g. [a, b, c] => "abc").
+ ""
+ } else if let Some(lit_string) = lit.value().try_as_str() {
+ lit_string.unwrap_or("")
+ } else {
+ return not_impl_err!(
+ "StringAgg not supported for delimiter {}",
Review Comment:
Could you just add brackets or quotes around the `{}`? So unprintable or other weirdly formatted delimiters don't produce the error `"StringAgg not supported for delimiter "`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] STRING_AGG missing functionality [datafusion]
Posted by "geoffreyclaude (via GitHub)" <gi...@apache.org>.
geoffreyclaude commented on PR #14412:
URL: https://github.com/apache/datafusion/pull/14412#issuecomment-2783476941
@alamb all good for me! You can give the final review.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] STRING_AGG missing functionality [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #14412:
URL: https://github.com/apache/datafusion/pull/14412#issuecomment-2778282626
@geoffreyclaude when you think this PR is ready, can you please ping me and I'll give it a review?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] STRING_AGG missing functionality [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb merged PR #14412:
URL: https://github.com/apache/datafusion/pull/14412
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] STRING_AGG missing functionality [datafusion]
Posted by "gabotechs (via GitHub)" <gi...@apache.org>.
gabotechs commented on code in PR #14412:
URL: https://github.com/apache/datafusion/pull/14412#discussion_r1963954221
##########
datafusion/sqllogictest/test_files/aggregate.slt:
##########
@@ -5568,6 +5573,16 @@ SELECT STRING_AGG(x,',') FROM strings WHERE g > 100
----
NULL
+query T
+SELECT STRING_AGG(DISTINCT x,',') FROM strings WHERE g > 100
Review Comment:
🤷♂️ It's the format that all the previous tests were following, I kept it for consistency
##########
datafusion/sqllogictest/test_files/aggregate.slt:
##########
@@ -5568,6 +5573,16 @@ SELECT STRING_AGG(x,',') FROM strings WHERE g > 100
----
NULL
+query T
+SELECT STRING_AGG(DISTINCT x,',') FROM strings WHERE g > 100
Review Comment:
🤷♂️ It's the format that all the previous STRING_AGG tests were following, I kept it for consistency
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] STRING_AGG missing functionality [datafusion]
Posted by "geoffreyclaude (via GitHub)" <gi...@apache.org>.
geoffreyclaude commented on code in PR #14412:
URL: https://github.com/apache/datafusion/pull/14412#discussion_r1962048387
##########
datafusion/sqllogictest/test_files/aggregate.slt:
##########
@@ -5568,6 +5573,16 @@ SELECT STRING_AGG(x,',') FROM strings WHERE g > 100
----
NULL
+query T
+SELECT STRING_AGG(DISTINCT x,',') FROM strings WHERE g > 100
Review Comment:
nit: missing space between the `x` and `','`.
##########
datafusion/sqllogictest/test_files/aggregate.slt:
##########
@@ -5547,6 +5547,11 @@ SELECT STRING_AGG(column1, '|') FROM (values (''), (null), (''));
----
|
+query T
+SELECT STRING_AGG(DISTINCT column1, '|') FROM (values (''), (null), (''));
Review Comment:
Same comment as in the [`array_agg` PR](https://github.com/apache/datafusion/pull/14413), can you add one or two tests triggering errors? Just to validate that invalid delimiters for instance are correctly detected.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] STRING_AGG missing functionality [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #14412:
URL: https://github.com/apache/datafusion/pull/14412#issuecomment-2634783601
Close/reopen to rerun CI checks
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] STRING_AGG missing functionality [datafusion]
Posted by "geoffreyclaude (via GitHub)" <gi...@apache.org>.
geoffreyclaude commented on code in PR #14412:
URL: https://github.com/apache/datafusion/pull/14412#discussion_r1981567196
##########
datafusion/sqllogictest/test_files/aggregate.slt:
##########
@@ -5568,6 +5573,16 @@ SELECT STRING_AGG(x,',') FROM strings WHERE g > 100
----
NULL
+query T
+SELECT STRING_AGG(DISTINCT x,',') FROM strings WHERE g > 100
Review Comment:
consistency always wins over correctness!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] STRING_AGG missing functionality [datafusion]
Posted by "gabotechs (via GitHub)" <gi...@apache.org>.
gabotechs commented on code in PR #14412:
URL: https://github.com/apache/datafusion/pull/14412#discussion_r2026461993
##########
datafusion/functions-aggregate/src/string_agg.rs:
##########
@@ -129,52 +172,326 @@ impl AggregateUDFImpl for StringAgg {
#[derive(Debug)]
pub(crate) struct StringAggAccumulator {
- values: Option<String>,
+ array_agg_acc: Box<dyn Accumulator>,
delimiter: String,
}
impl StringAggAccumulator {
- pub fn new(delimiter: &str) -> Self {
+ pub fn new(array_agg_acc: Box<dyn Accumulator>, delimiter: &str) -> Self {
Self {
- values: None,
+ array_agg_acc,
delimiter: delimiter.to_string(),
}
}
}
impl Accumulator for StringAggAccumulator {
fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
- let string_array: Vec<_> = as_generic_string_array::<i64>(&values[0])?
- .iter()
- .filter_map(|v| v.as_ref().map(ToString::to_string))
- .collect();
- if !string_array.is_empty() {
- let s = string_array.join(self.delimiter.as_str());
- let v = self.values.get_or_insert("".to_string());
- if !v.is_empty() {
- v.push_str(self.delimiter.as_str());
+ self.array_agg_acc.update_batch(&filter_index(values, 1))
+ }
+
+ fn evaluate(&mut self) -> Result<ScalarValue> {
+ let scalar = self.array_agg_acc.evaluate()?;
+
+ let ScalarValue::List(list) = scalar else {
+ return internal_err!("Expected a DataType::List while evaluating underlying ArrayAggAccumulator, but got {}", scalar.data_type());
+ };
+
+ let string_arr: Vec<_> = match list.value_type() {
+ DataType::LargeUtf8 => as_generic_string_array::<i64>(list.values())?
+ .iter()
+ .flatten()
+ .collect(),
+ DataType::Utf8 => as_generic_string_array::<i32>(list.values())?
+ .iter()
+ .flatten()
+ .collect(),
+ _ => {
+ return internal_err!(
+ "Expected elements to of type Utf8 or LargeUtf8, but got {}",
+ list.value_type()
+ )
}
- v.push_str(s.as_str());
+ };
+
+ if string_arr.is_empty() {
+ return Ok(ScalarValue::LargeUtf8(None));
}
- Ok(())
+
+ Ok(ScalarValue::LargeUtf8(Some(
+ string_arr.join(&self.delimiter),
+ )))
+ }
+
+ fn size(&self) -> usize {
+ size_of_val(self) - size_of_val(&self.array_agg_acc)
+ + self.array_agg_acc.size()
+ + self.delimiter.capacity()
+ }
+
+ fn state(&mut self) -> Result<Vec<ScalarValue>> {
+ self.array_agg_acc.state()
}
fn merge_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
- self.update_batch(values)?;
+ self.array_agg_acc.merge_batch(values)
+ }
+}
+
+fn filter_index<T: Clone>(values: &[T], index: usize) -> Vec<T> {
+ values
+ .iter()
+ .enumerate()
+ .filter(|(i, _)| *i != index)
+ .map(|(_, v)| v)
+ .cloned()
+ .collect::<Vec<_>>()
+}
+
+#[cfg(test)]
+mod tests {
Review Comment:
I'm finding that having unit tests near the function definitions itself is significantly more ergonomic than writing sql logic tests for a couple of reasons:
- We can tests the accumulators in isolation, allowing for finer grained control about batch updating, state generation, merging different states from different accumulators, etc...
- The time it takes since a developer makes a code change, until the appropriate test is run is reduced significantly:
```shell
cargo test --lib string_agg::tests --manifest-path 1.40s user 1.53s system 132% cpu 2.204 total
```
VS
```shell
cargo test --test sqllogictests 33.61s user 7.17s system 365% cpu 11.164 total
```
Measured on a Mac M3
I see that this is not an stablished pattern though, and I'm wondering what are people's take on this
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] STRING_AGG missing functionality [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #14412:
URL: https://github.com/apache/datafusion/pull/14412#discussion_r2028567883
##########
datafusion/functions-aggregate/src/string_agg.rs:
##########
@@ -129,52 +172,326 @@ impl AggregateUDFImpl for StringAgg {
#[derive(Debug)]
pub(crate) struct StringAggAccumulator {
- values: Option<String>,
+ array_agg_acc: Box<dyn Accumulator>,
delimiter: String,
}
impl StringAggAccumulator {
- pub fn new(delimiter: &str) -> Self {
+ pub fn new(array_agg_acc: Box<dyn Accumulator>, delimiter: &str) -> Self {
Self {
- values: None,
+ array_agg_acc,
delimiter: delimiter.to_string(),
}
}
}
impl Accumulator for StringAggAccumulator {
fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
- let string_array: Vec<_> = as_generic_string_array::<i64>(&values[0])?
- .iter()
- .filter_map(|v| v.as_ref().map(ToString::to_string))
- .collect();
- if !string_array.is_empty() {
- let s = string_array.join(self.delimiter.as_str());
- let v = self.values.get_or_insert("".to_string());
- if !v.is_empty() {
- v.push_str(self.delimiter.as_str());
+ self.array_agg_acc.update_batch(&filter_index(values, 1))
+ }
+
+ fn evaluate(&mut self) -> Result<ScalarValue> {
+ let scalar = self.array_agg_acc.evaluate()?;
+
+ let ScalarValue::List(list) = scalar else {
+ return internal_err!("Expected a DataType::List while evaluating underlying ArrayAggAccumulator, but got {}", scalar.data_type());
+ };
+
+ let string_arr: Vec<_> = match list.value_type() {
+ DataType::LargeUtf8 => as_generic_string_array::<i64>(list.values())?
+ .iter()
+ .flatten()
+ .collect(),
+ DataType::Utf8 => as_generic_string_array::<i32>(list.values())?
+ .iter()
+ .flatten()
+ .collect(),
+ _ => {
+ return internal_err!(
+ "Expected elements to of type Utf8 or LargeUtf8, but got {}",
+ list.value_type()
+ )
}
- v.push_str(s.as_str());
+ };
+
+ if string_arr.is_empty() {
+ return Ok(ScalarValue::LargeUtf8(None));
}
- Ok(())
+
+ Ok(ScalarValue::LargeUtf8(Some(
+ string_arr.join(&self.delimiter),
+ )))
+ }
+
+ fn size(&self) -> usize {
+ size_of_val(self) - size_of_val(&self.array_agg_acc)
+ + self.array_agg_acc.size()
+ + self.delimiter.capacity()
+ }
+
+ fn state(&mut self) -> Result<Vec<ScalarValue>> {
+ self.array_agg_acc.state()
}
fn merge_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
- self.update_batch(values)?;
+ self.array_agg_acc.merge_batch(values)
+ }
+}
+
+fn filter_index<T: Clone>(values: &[T], index: usize) -> Vec<T> {
+ values
+ .iter()
+ .enumerate()
+ .filter(|(i, _)| *i != index)
+ .map(|(_, v)| v)
+ .cloned()
+ .collect::<Vec<_>>()
+}
+
+#[cfg(test)]
+mod tests {
Review Comment:
In general I think that we should strive that sqllogictests cover all the "end user" visible behavior -- since they function as integration style test and make sure the functionality is all hooked up together correctly (not just working in isolation)
Unit tests / in module tests are good to cover cases that are hard to cover in sqllogictests
I think there is some more back story here: https://datafusion.apache.org/contributor-guide/testing.html#sqllogictests-tests
In terms of cycle times, sqllogictests do have the benefit you can update them without any code changes (so writing / updating them is sometimes faster than code), though you are right that testing code changes requires a recompilation
@logan-keede has been working on improving the build performance recently. Hopefully this will get better
- https://github.com/apache/datafusion/issues/13814
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] STRING_AGG missing functionality [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb closed pull request #14412: STRING_AGG missing functionality
URL: https://github.com/apache/datafusion/pull/14412
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] STRING_AGG missing functionality [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb closed pull request #14412: STRING_AGG missing functionality
URL: https://github.com/apache/datafusion/pull/14412
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] STRING_AGG missing functionality [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #14412:
URL: https://github.com/apache/datafusion/pull/14412#discussion_r2038100274
##########
datafusion/functions-aggregate/src/string_agg.rs:
##########
@@ -65,6 +83,7 @@ make_udaf_expr_and_func!(
#[derive(Debug)]
pub struct StringAgg {
signature: Signature,
+ array_agg: ArrayAgg,
}
Review Comment:
seems good to me 👍
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] STRING_AGG missing functionality [datafusion]
Posted by "gabotechs (via GitHub)" <gi...@apache.org>.
gabotechs commented on code in PR #14412:
URL: https://github.com/apache/datafusion/pull/14412#discussion_r2026415116
##########
datafusion/functions-aggregate/src/string_agg.rs:
##########
@@ -65,6 +83,7 @@ make_udaf_expr_and_func!(
#[derive(Debug)]
pub struct StringAgg {
signature: Signature,
+ array_agg: ArrayAgg,
}
Review Comment:
I don't think I've seen this pattern much in DataFusion codebase about reusing another function as a building block. Any opinions about this approach?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] STRING_AGG missing functionality [datafusion]
Posted by "gabotechs (via GitHub)" <gi...@apache.org>.
gabotechs commented on code in PR #14412:
URL: https://github.com/apache/datafusion/pull/14412#discussion_r2026461993
##########
datafusion/functions-aggregate/src/string_agg.rs:
##########
@@ -129,52 +172,326 @@ impl AggregateUDFImpl for StringAgg {
#[derive(Debug)]
pub(crate) struct StringAggAccumulator {
- values: Option<String>,
+ array_agg_acc: Box<dyn Accumulator>,
delimiter: String,
}
impl StringAggAccumulator {
- pub fn new(delimiter: &str) -> Self {
+ pub fn new(array_agg_acc: Box<dyn Accumulator>, delimiter: &str) -> Self {
Self {
- values: None,
+ array_agg_acc,
delimiter: delimiter.to_string(),
}
}
}
impl Accumulator for StringAggAccumulator {
fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
- let string_array: Vec<_> = as_generic_string_array::<i64>(&values[0])?
- .iter()
- .filter_map(|v| v.as_ref().map(ToString::to_string))
- .collect();
- if !string_array.is_empty() {
- let s = string_array.join(self.delimiter.as_str());
- let v = self.values.get_or_insert("".to_string());
- if !v.is_empty() {
- v.push_str(self.delimiter.as_str());
+ self.array_agg_acc.update_batch(&filter_index(values, 1))
+ }
+
+ fn evaluate(&mut self) -> Result<ScalarValue> {
+ let scalar = self.array_agg_acc.evaluate()?;
+
+ let ScalarValue::List(list) = scalar else {
+ return internal_err!("Expected a DataType::List while evaluating underlying ArrayAggAccumulator, but got {}", scalar.data_type());
+ };
+
+ let string_arr: Vec<_> = match list.value_type() {
+ DataType::LargeUtf8 => as_generic_string_array::<i64>(list.values())?
+ .iter()
+ .flatten()
+ .collect(),
+ DataType::Utf8 => as_generic_string_array::<i32>(list.values())?
+ .iter()
+ .flatten()
+ .collect(),
+ _ => {
+ return internal_err!(
+ "Expected elements to of type Utf8 or LargeUtf8, but got {}",
+ list.value_type()
+ )
}
- v.push_str(s.as_str());
+ };
+
+ if string_arr.is_empty() {
+ return Ok(ScalarValue::LargeUtf8(None));
}
- Ok(())
+
+ Ok(ScalarValue::LargeUtf8(Some(
+ string_arr.join(&self.delimiter),
+ )))
+ }
+
+ fn size(&self) -> usize {
+ size_of_val(self) - size_of_val(&self.array_agg_acc)
+ + self.array_agg_acc.size()
+ + self.delimiter.capacity()
+ }
+
+ fn state(&mut self) -> Result<Vec<ScalarValue>> {
+ self.array_agg_acc.state()
}
fn merge_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
- self.update_batch(values)?;
+ self.array_agg_acc.merge_batch(values)
+ }
+}
+
+fn filter_index<T: Clone>(values: &[T], index: usize) -> Vec<T> {
+ values
+ .iter()
+ .enumerate()
+ .filter(|(i, _)| *i != index)
+ .map(|(_, v)| v)
+ .cloned()
+ .collect::<Vec<_>>()
+}
+
+#[cfg(test)]
+mod tests {
Review Comment:
I'm finding that having unit tests near the function definitions itself is significantly more ergonomic than writing sql logic tests for a couple of reasons:
- We can tests the accumulators in isolation, allowing for finer grained control about batch updating, state generation, merging different states from different accumulators, etc...
- The time it takes since a developer makes a code change, until the appropriate test is run is reduced significantly:
```shell
cargo test --lib string_agg::tests --manifest-path 1.40s user 1.53s system 132% cpu 2.204 total
```
VS
```shell
cargo test --test sqllogictests 33.61s user 7.17s system 365% cpu 11.164 total
```
Measured on a Mac M3
Not saying that we should not have sql logic tests, I think those are a must, but maybe having some testing tooling for folks to be able to contribute unit tests also here could improve DX.
I see that this is not an stablished pattern though, and I'm wondering what are people's take on this
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] STRING_AGG missing functionality [datafusion]
Posted by "gabotechs (via GitHub)" <gi...@apache.org>.
gabotechs commented on code in PR #14412:
URL: https://github.com/apache/datafusion/pull/14412#discussion_r2028725710
##########
datafusion/functions-aggregate/src/string_agg.rs:
##########
@@ -129,52 +172,326 @@ impl AggregateUDFImpl for StringAgg {
#[derive(Debug)]
pub(crate) struct StringAggAccumulator {
- values: Option<String>,
+ array_agg_acc: Box<dyn Accumulator>,
delimiter: String,
}
impl StringAggAccumulator {
- pub fn new(delimiter: &str) -> Self {
+ pub fn new(array_agg_acc: Box<dyn Accumulator>, delimiter: &str) -> Self {
Self {
- values: None,
+ array_agg_acc,
delimiter: delimiter.to_string(),
}
}
}
impl Accumulator for StringAggAccumulator {
fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
- let string_array: Vec<_> = as_generic_string_array::<i64>(&values[0])?
- .iter()
- .filter_map(|v| v.as_ref().map(ToString::to_string))
- .collect();
- if !string_array.is_empty() {
- let s = string_array.join(self.delimiter.as_str());
- let v = self.values.get_or_insert("".to_string());
- if !v.is_empty() {
- v.push_str(self.delimiter.as_str());
+ self.array_agg_acc.update_batch(&filter_index(values, 1))
+ }
+
+ fn evaluate(&mut self) -> Result<ScalarValue> {
+ let scalar = self.array_agg_acc.evaluate()?;
+
+ let ScalarValue::List(list) = scalar else {
+ return internal_err!("Expected a DataType::List while evaluating underlying ArrayAggAccumulator, but got {}", scalar.data_type());
+ };
+
+ let string_arr: Vec<_> = match list.value_type() {
+ DataType::LargeUtf8 => as_generic_string_array::<i64>(list.values())?
+ .iter()
+ .flatten()
+ .collect(),
+ DataType::Utf8 => as_generic_string_array::<i32>(list.values())?
+ .iter()
+ .flatten()
+ .collect(),
+ _ => {
+ return internal_err!(
+ "Expected elements to of type Utf8 or LargeUtf8, but got {}",
+ list.value_type()
+ )
}
- v.push_str(s.as_str());
+ };
+
+ if string_arr.is_empty() {
+ return Ok(ScalarValue::LargeUtf8(None));
}
- Ok(())
+
+ Ok(ScalarValue::LargeUtf8(Some(
+ string_arr.join(&self.delimiter),
+ )))
+ }
+
+ fn size(&self) -> usize {
+ size_of_val(self) - size_of_val(&self.array_agg_acc)
+ + self.array_agg_acc.size()
+ + self.delimiter.capacity()
+ }
+
+ fn state(&mut self) -> Result<Vec<ScalarValue>> {
+ self.array_agg_acc.state()
}
fn merge_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
- self.update_batch(values)?;
+ self.array_agg_acc.merge_batch(values)
+ }
+}
+
+fn filter_index<T: Clone>(values: &[T], index: usize) -> Vec<T> {
+ values
+ .iter()
+ .enumerate()
+ .filter(|(i, _)| *i != index)
+ .map(|(_, v)| v)
+ .cloned()
+ .collect::<Vec<_>>()
+}
+
+#[cfg(test)]
+mod tests {
Review Comment:
🤔 it looks there's still some cards that can be played for making compilation times faster then. 👍 thanks for all that info!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org