NVIDIA / NVIDIA/cudf

[FEA] aggregation each list in a column to a single value using a user supplied function

Open
#8,020 5 comments 0 reactions 0 assignees View on GitHub
0 - Backlog feature request libcudf Spark
Dominant language
C++
Stars
9.8k
Forks
1.1k
Avg merge
3d 6m
Merged PRs (30d)
278

Description

**Is your feature request related to a problem? Please describe.**
Spark supports the aggregate function in SQL (Not really standard but we have customers who use it)

https://spark.apache.org/docs/latest/api/sql/index.html#aggregate

It takes 4 arguments.
* **argument** an array/list column to do the aggregation on
* **initial** an initial value for accumulation
* **merge** a higher order function that takes two arguments an accumulation value and the current value in the list
* **finish** an optional higher order function that takes the output of merge and transforms it into a final value

A higher order function is a function that is written in SQL like `(a, b) -> a + b` to add two things together.

Even though **finish** is a higher order function we don't have to treat it that way so we can ignore it for now.

**merge** however is something new that CUDF has not really supported before. It allows the user to specify how they want an aggregation to happen instead of having a declarative aggregation like most SQL does. So for example if I wanted to do the equivalent of SUM it would look something like

```
SELECT aggregate(list_of_int_column, 0, (acc, x) -> acc + x) as sum_of_ints_in_list
```

For each list in the column it would do essentially the equivalent of

```
MERGE_OUTPUT_TYPE acc = initial_value;
for (ELEMENT_TYPE & elem : list_data) {
acc = merge(acc, elem);
}
```

The problem we are running into is that our customers have rather complicated operations, where the higher order function can reach out to other columns in the same row.

```
(acc, x) -> (
CASE WHEN other_column - x.struct_sub_column >= acc
AND other_column - x.struct_sub_column < 100
AND x.struct_string_column = 'FOO'
AND yet_another_column <> x.third_struct_sub_column
THEN other_column - x.struct_sub_column
ELSE acc
END))
```

**Describe the solution you'd like**

I would love something where we could build up an AST tree that represents the higher order function and have cudf provide a list_aggregation function that would do what we need/want. But we know that there are potentially issues with the AST in terms or performance when there are too many operators so this is all open to discussion.

**Describe alternatives you've considered**

We have thought about trying to do pattern matching to decompose the higher order function into something more manageable for CUDF to support.

i.e.

```
(acc, x) -> acc + x
```
could be translated into an SUM aggregation across the values in the list, or

```
(acc, x) -> CASE WHEN x > acc THEN x ELSE acc END
```

could be translated into a MAX aggregation across the values in the list.

We could even use pattern matching for things like

```
(acc, x) -> acc + x.first - x.second
```

To translate it into first doing a `x.first - x.second` for all of the struct values within the list, and then doing a SUM aggregation on that resulting list. But things get much more difficult when we try to support pulling in other columns, and struct columns, etc.

```
(acc, x) -> acc + x.first + foo
```

In this case we would have to do essentially an `explode` on `foo` and the list column so we could execute `x.first + foo` and then finally do the SUM aggregation. This is a bit problematic because of potential memory issues that explode can cause.

So if the AST is not a workable solution we would like to request a generic list aggregation operation instead.

```
cudf::column list_aggregate(cudf::lists_column_view list, std::unique_ptr & aggregation);
```

With at a minimum supporting MAX, SUM, and MIN aggregations initially.

It would probably be ideal to expand it out to multiple aggregations at once like with `groupby`, but it is not a requirement.

I also need to add that null handling would have to be a bit different than other aggregations. If the list itself is a null, then the output should be a null, but if a value in the list is a null, then the output should also be a null.

```
scala> spark.sql("SELECT aggregate(array(1, 2, 3), 0, (acc, x) -> acc + x) as A").show
+---+
| A|
+---+
| 6|
+---+

scala> spark.sql("SELECT aggregate(array(1, 2, 3, null), 0, (acc, x) -> acc + x) as A").show
+----+
| A|
+----+
|null|
+----+
```

Ideally we would also love to have some kind of explode that would not make a copy of the array we are exploding on, but instead just do the explode on the columns that need it.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.