Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Support Collect-like Reduction Aggregations #4992

Merged
merged 6 commits into from
Apr 11, 2022
Merged

Support Collect-like Reduction Aggregations #4992

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
91350f5
init
sperlingxx Mar 21, 2022
af2033a
update
sperlingxx Mar 21, 2022
8bf99f1
update doc
sperlingxx Mar 21, 2022
065afe8
fix
sperlingxx Mar 21, 2022
c016739
Merge remote-tracking branch 'origin/branch-22.06' into reduce_collects
sperlingxx Apr 2, 2022
7c3f66d
Merge remote-tracking branch 'origin/branch-22.06' into reduce_collects
sperlingxx Apr 11, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 34 additions & 34 deletions docs/supported_ops.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14824,25 +14824,25 @@ are limited.
<td rowSpan="6">`collect_list`</td>
<td rowSpan="6">Collect a list of non-unique elements, not supported in reduction</td>
<td rowSpan="6">None</td>
<td rowSpan="2">reduction</td>
<td rowSpan="2">aggregation</td>
<td>input</td>
<td>S</td>
<td>S</td>
<td>S</td>
<td>S</td>
<td>S</td>
<td>S</td>
<td>S</td>
<td>S</td>
<td><em>PS<br/>UTC is only supported TZ for TIMESTAMP</em></td>
<td>S</td>
<td>S</td>
<td>S</td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><em>PS<br/>UTC is only supported TZ for child TIMESTAMP;<br/>unsupported child types BINARY, CALENDAR, UDT</em></td>
<td><em>PS<br/>UTC is only supported TZ for child TIMESTAMP;<br/>unsupported child types BINARY, CALENDAR, UDT</em></td>
<td><em>PS<br/>UTC is only supported TZ for child TIMESTAMP;<br/>unsupported child types BINARY, CALENDAR, UDT</em></td>
<td><b>NS</b></td>
</tr>
<tr>
Expand All @@ -14861,13 +14861,13 @@ are limited.
<td> </td>
<td> </td>
<td> </td>
<td><b>NS</b></td>
<td><em>PS<br/>UTC is only supported TZ for child TIMESTAMP;<br/>unsupported child types BINARY, CALENDAR, UDT</em></td>
<td> </td>
<td> </td>
<td> </td>
</tr>
<tr>
<td rowSpan="2">aggregation</td>
<td rowSpan="2">reduction</td>
<td>input</td>
<td>S</td>
<td>S</td>
Expand Down Expand Up @@ -14983,25 +14983,25 @@ are limited.
<td rowSpan="6">`collect_set`</td>
<td rowSpan="6">Collect a set of unique elements, not supported in reduction</td>
<td rowSpan="6">None</td>
<td rowSpan="2">reduction</td>
<td rowSpan="2">aggregation</td>
<td>input</td>
<td>S</td>
<td>S</td>
<td>S</td>
<td>S</td>
<td>S</td>
<td>S</td>
<td>S</td>
<td>S</td>
<td><em>PS<br/>UTC is only supported TZ for TIMESTAMP</em></td>
<td>S</td>
<td>S</td>
<td>S</td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><em>PS<br/>UTC is only supported TZ for child TIMESTAMP;<br/>unsupported child types BINARY, CALENDAR, ARRAY, MAP, UDT</em></td>
<td><b>NS</b></td>
</tr>
<tr>
Expand All @@ -15020,13 +15020,13 @@ are limited.
<td> </td>
<td> </td>
<td> </td>
<td><b>NS</b></td>
<td><em>PS<br/>UTC is only supported TZ for child TIMESTAMP;<br/>unsupported child types BINARY, CALENDAR, ARRAY, MAP, UDT</em></td>
<td> </td>
<td> </td>
<td> </td>
</tr>
<tr>
<td rowSpan="2">aggregation</td>
<td rowSpan="2">reduction</td>
<td>input</td>
<td>S</td>
<td>S</td>
Expand Down
49 changes: 37 additions & 12 deletions integration_tests/src/main/python/hash_aggregate_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1038,18 +1038,18 @@ def test_first_last_reductions_nested_types(data_gen):
def test_generic_reductions(data_gen):
local_conf = copy_and_update(_no_nans_float_conf, {'spark.sql.legacy.allowParameterlessCount': 'true'})
assert_gpu_and_cpu_are_equal_collect(
# Coalesce and sort are to make sure that first and last, which are non-deterministic
# become deterministic
lambda spark : unary_op_df(spark, data_gen)\
.coalesce(1).selectExpr(
'min(a)',
'max(a)',
'first(a)',
'last(a)',
'count(a)',
'count()',
'count(1)'),
conf = local_conf)
# Coalesce and sort are to make sure that first and last, which are non-deterministic
# become deterministic
lambda spark : unary_op_df(spark, data_gen) \
.coalesce(1).selectExpr(
'min(a)',
'max(a)',
'first(a)',
'last(a)',
'count(a)',
'count()',
'count(1)'),
conf=local_conf)

@pytest.mark.parametrize('data_gen', all_gen + _nested_gens, ids=idfn)
def test_count(data_gen):
Expand Down Expand Up @@ -1083,6 +1083,31 @@ def test_arithmetic_reductions(data_gen):
'avg(a)'),
conf = _no_nans_float_conf)

@pytest.mark.parametrize('data_gen',
non_nan_all_basic_gens + decimal_gens + _nested_gens,
ids=idfn)
def test_collect_list_reductions(data_gen):
assert_gpu_and_cpu_are_equal_collect(
lambda spark: unary_op_df(spark, data_gen).selectExpr('collect_list(a)'),
conf=_no_nans_float_conf)

_struct_only_nested_gens = [all_basic_struct_gen,
StructGen([['child0', byte_gen], ['child1', all_basic_struct_gen]]),
StructGen([])]
@pytest.mark.parametrize('data_gen',
non_nan_all_basic_gens + decimal_gens + _struct_only_nested_gens,
ids=idfn)
def test_collect_set_reductions(data_gen):
assert_gpu_and_cpu_are_equal_collect(
lambda spark: unary_op_df(spark, data_gen).selectExpr('sort_array(collect_set(a))'),
conf=_no_nans_float_conf)

def test_collect_empty():
assert_gpu_and_cpu_are_equal_collect(
lambda spark: spark.sql("select collect_list(null)"))
assert_gpu_and_cpu_are_equal_collect(
lambda spark: spark.sql("select collect_set(null)"))

@ignore_order(local=True)
@pytest.mark.parametrize('data_gen', all_gen + _nested_gens, ids=idfn)
def test_groupby_first_last(data_gen):
Expand Down
8 changes: 2 additions & 6 deletions sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3218,8 +3218,7 @@ object GpuOverrides extends Logging {
}),
expr[CollectList](
"Collect a list of non-unique elements, not supported in reduction",
// GpuCollectList is not yet supported in Reduction context.
ExprChecks.aggNotReduction(
ExprChecks.fullAgg(
TypeSig.ARRAY.nested(TypeSig.commonCudfTypes + TypeSig.DECIMAL_128 +
TypeSig.NULL + TypeSig.STRUCT + TypeSig.ARRAY + TypeSig.MAP),
TypeSig.ARRAY.nested(TypeSig.all),
Expand Down Expand Up @@ -3249,10 +3248,7 @@ object GpuOverrides extends Logging {
}),
expr[CollectSet](
"Collect a set of unique elements, not supported in reduction",
// GpuCollectSet is not yet supported in Reduction context.
// Compared to CollectList, ArrayType and MapType are NOT supported in GpuCollectSet
// because underlying cuDF operator drop_list_duplicates doesn't support LIST type.
ExprChecks.aggNotReduction(
ExprChecks.fullAgg(
TypeSig.ARRAY.nested(TypeSig.commonCudfTypes + TypeSig.DECIMAL_128 +
TypeSig.NULL + TypeSig.STRUCT),
TypeSig.ARRAY.nested(TypeSig.all),
Expand Down
30 changes: 19 additions & 11 deletions sql-plugin/src/main/scala/org/apache/spark/sql/rapids/AggregateFunctions.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.TypeCheckSuccess
import org.apache.spark.sql.catalyst.expressions.{AttributeReference, AttributeSet, Expression, ExprId, ImplicitCastInputTypes, UnsafeProjection, UnsafeRow}
import org.apache.spark.sql.catalyst.expressions.aggregate._
import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
import org.apache.spark.sql.catalyst.util.{ArrayData, TypeUtils}
import org.apache.spark.sql.catalyst.util.{ArrayData, GenericArrayData, TypeUtils}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types._
import org.apache.spark.sql.vectorized.ColumnarBatch
Expand Down Expand Up @@ -352,32 +352,39 @@ class CudfMin(override val dataType: DataType) extends CudfAggregate {
}

class CudfCollectList(override val dataType: DataType) extends CudfAggregate {
override lazy val reductionAggregate: cudf.ColumnVector => cudf.Scalar = _ =>
throw new UnsupportedOperationException("CollectList is not yet supported in reduction")
override lazy val reductionAggregate: cudf.ColumnVector => cudf.Scalar =
(col: cudf.ColumnVector) => col.reduce(ReductionAggregation.collectList(), DType.LIST)
override lazy val groupByAggregate: GroupByAggregation =
GroupByAggregation.collectList()
override val name: String = "CudfCollectList"
}

class CudfMergeLists(override val dataType: DataType) extends CudfAggregate {
override lazy val reductionAggregate: cudf.ColumnVector => cudf.Scalar = _ =>
throw new UnsupportedOperationException("MergeLists is not yet supported in reduction")
override lazy val reductionAggregate: cudf.ColumnVector => cudf.Scalar =
(col: cudf.ColumnVector) => col.reduce(ReductionAggregation.mergeLists(), DType.LIST)
override lazy val groupByAggregate: GroupByAggregation =
GroupByAggregation.mergeLists()
override val name: String = "CudfMergeLists"
}

class CudfCollectSet(override val dataType: DataType) extends CudfAggregate {
override lazy val reductionAggregate: cudf.ColumnVector => cudf.Scalar = _ =>
throw new UnsupportedOperationException("CollectSet is not yet supported in reduction")
override lazy val reductionAggregate: cudf.ColumnVector => cudf.Scalar =
(col: cudf.ColumnVector) => {
val collectSet = ReductionAggregation.collectSet(
NullPolicy.EXCLUDE, NullEquality.EQUAL, NaNEquality.UNEQUAL)
col.reduce(collectSet, DType.LIST)
}
override lazy val groupByAggregate: GroupByAggregation =
GroupByAggregation.collectSet(NullPolicy.EXCLUDE, NullEquality.EQUAL, NaNEquality.UNEQUAL)
override val name: String = "CudfCollectSet"
}

class CudfMergeSets(override val dataType: DataType) extends CudfAggregate {
override lazy val reductionAggregate: cudf.ColumnVector => cudf.Scalar = _ =>
throw new UnsupportedOperationException("CudfMergeSets is not yet supported in reduction")
override lazy val reductionAggregate: cudf.ColumnVector => cudf.Scalar =
(col: cudf.ColumnVector) => {
val mergeSets = ReductionAggregation.mergeSets(NullEquality.EQUAL, NaNEquality.UNEQUAL)
col.reduce(mergeSets, DType.LIST)
}
override lazy val groupByAggregate: GroupByAggregation =
GroupByAggregation.mergeSets(NullEquality.EQUAL, NaNEquality.UNEQUAL)
override val name: String = "CudfMergeSets"
Expand Down Expand Up @@ -1578,8 +1585,9 @@ trait GpuCollectBase
// WINDOW FUNCTION
override val windowInputProjection: Seq[Expression] = Seq(child)

// Make them lazy to avoid being initialized when creating a GpuCollectOp.
override lazy val initialValues: Seq[Expression] = throw new UnsupportedOperationException
override val initialValues: Seq[Expression] = {
Seq(GpuLiteral.create(new GenericArrayData(Array.empty[Any]), dataType))
}

override val inputProjection: Seq[Expression] = Seq(child)

Expand Down