`string_agg` aggregate function is 1000x slower than duckdb (SQLStorm)

[2010YOUY01]

Is your feature request related to a problem or challenge?

Reproducer

1. Generate dataset with https://github.com/clflushopt/tpchgen-rs/tree/main/tpchgen-cli
2. In datafusion-cli

DataFusion CLI v50.0.0
> CREATE EXTERNAL TABLE partsupp
STORED AS PARQUET
LOCATION '/Users/yongting/Code/datafusion-sqlstorm/data/partsupp.parquet';

> select ps_partkey, string_agg(ps_comment, ';')
from partsupp
group by ps_partkey;

...
20000 row(s) fetched. (First 40 displayed. Use --maxrows to adjust)
Elapsed 53.737 seconds.

Also the memory usage is quite high.
DuckDB finishes in 0.03s

Describe the solution you'd like

No response

Describe alternatives you've considered

No response

Additional context

Found by SQLStorm #17698

[vegarsti]

This will likely be improved by a GroupsAccumulator for StringAgg. Since there isn't a specialized version, the query in the issue goes through the GroupsAccumulatorAdapter:

datafusion/datafusion/physical-plan/src/aggregates/row_hash.rs

Lines 622 to 640 in 5cc0be5

	/// Create an accumulator for `agg_expr` -- a [`GroupsAccumulator`] if
	/// that is supported by the aggregate, or a
	/// [`GroupsAccumulatorAdapter`] if not.
	pub(crate) fn create_group_accumulator(
	agg_expr: &Arc<AggregateFunctionExpr>,
	) -> Result<Box<dyn GroupsAccumulator>> {
	if agg_expr.groups_accumulator_supported() {
	agg_expr.create_groups_accumulator()
	} else {
	// Note in the log when the slow path is used
	debug!(
	"Creating GroupsAccumulatorAdapter for {}: {agg_expr:?}",
	agg_expr.name()
	);
	let agg_expr_captured = Arc::clone(agg_expr);
	let factory = move || agg_expr_captured.create_accumulator();
	Ok(Box::new(GroupsAccumulatorAdapter::new(factory)))
	}
	}

[vegarsti]

Would love to try implementing GroupsAccumulator for StringAgg.

[vegarsti]

take

[vegarsti]

StringAgg is implemented using ArrayAgg:

datafusion/datafusion/functions-aggregate/src/string_agg.rs

Lines 144 to 178 in 5cc0be5

	fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
	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 \"{}\"",
	lit.value()
	);
	};
	let array_agg_acc = self.array_agg.accumulator(AccumulatorArgs {
	return_field: Field::new(
	"f",
	DataType::new_list(acc_args.return_field.data_type().clone(), true),
	true,
	)
	.into(),
	exprs: &filter_index(acc_args.exprs, 1),
	..acc_args
	})?;
	Ok(Box::new(StringAggAccumulator::new(
	array_agg_acc,
	delimiter,
	)))

So the work to implement GroupsAccumulator for that should help here. #10145 (also related: #17446)

[vegarsti]

Would be good to first profile with samply like here #17446 (comment)

[vegarsti]

@2010YOUY01 What specific tpchgen-cli command did you run to generate the data here? (scale factor)

[2010YOUY01]

@2010YOUY01 What specific tpchgen-cli command did you run to generate the data here? (scale factor)

I'm using sf 0.1

[2010YOUY01]

I did a PR #17837 to add a specialized accumulator for no DISTINCT/ORDER case.

I think it's also necessary in addition to the GroupsAccumulator-based solution because

• Single group (no group-by) case like select string_agg(str, ',') from t1 is still using the Accumulator interface instead of GroupsAccumulator
• It's very simple to implement, and can be used to compare the performance with the future GroupsAccumulator solution

The remaining work to close this issue is

• Figure out the inefficiency inside array_agg, and make the general case (with DISTINCT/ORDER) faster
• Further speed up with GroupsAccumulator

[vegarsti]

I did a PR #17837 to add a specialized accumulator for no DISTINCT/ORDER case.

I think it's also necessary in addition to the GroupsAccumulator-based solution because

• Single group (no group-by) case like select string_agg(str, ',') from t1 is still using the Accumulator interface instead of GroupsAccumulator
• It's very simple to implement, and can be used to compare the performance with the future GroupsAccumulator solution

The remaining work to close this issue is

• Figure out the inefficiency inside array_agg, and make the general case (with DISTINCT/ORDER) faster
• Further speed up with GroupsAccumulator

Awesome!

[vegarsti]

Here's a Samply profile as well (from main):

From the root of the datafusion repo

• Installed samply cargo install --locked samply
• Installed tpchgen-cli RUSTFLAGS='-C target-cpu=native' cargo install tpchgen-cli
• Ran tpchgen-cli -s 0.1 --format=parquet --output-dir tpch-data
• Created string-agg.sql, see below
• Compiled the CLI with the profiling profile cargo build --profile profiling --bin datafusion-cli
• Ran samply record -- ./target/profiling/datafusion-cli -f string-agg.sql

string-agg.sql:

CREATE EXTERNAL TABLE partsupp
STORED AS PARQUET
LOCATION 'tpch-data/partsupp.parquet';

select ps_partkey, string_agg(ps_comment, ';')
from partsupp
group by ps_partkey;

[vegarsti]

Well done with that PR!! At this point I'm not sure the GroupsAccumulator is the right thing to do, what do you think? Once your PR is merged, getting rid of the usage of ArrayAggAccumulator in the general case for StringAgg might be the best thing to do?