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

Match Spark's NaN handling in collect_set #8909

Merged
merged 9 commits into from
Aug 9, 2023
Merged

Match Spark's NaN handling in collect_set #8909

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
0fbc47f
Match Spark's NaN handling in collect_set
thirtiseven Aug 2, 2023
e9e482b
Revert "Fix collect_set_on_nested_type tests failed (#8783)"
thirtiseven Aug 2, 2023
04e62e3
clean up and add comments
thirtiseven Aug 2, 2023
250a89a
remove xfail
thirtiseven Aug 2, 2023
f500a8a
update cudfmergesets too
thirtiseven Aug 3, 2023
db500a5
edit comment
thirtiseven Aug 3, 2023
cd587cf
remove related nan datagens
thirtiseven Aug 9, 2023
3b4abd7
clean up
thirtiseven Aug 9, 2023
7b8af24
clean up
thirtiseven Aug 9, 2023
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
12 changes: 6 additions & 6 deletions docs/supported_ops.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16595,9 +16595,9 @@ are limited.
<td>S</td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><em>PS<br/>Arrays containing floats/doubles will not be supported.;<br/>UTC is only supported TZ for child TIMESTAMP;<br/>unsupported child types BINARY, CALENDAR, MAP, UDT</em></td>
<td><em>PS<br/>UTC is only supported TZ for child TIMESTAMP;<br/>unsupported child types BINARY, CALENDAR, MAP, UDT</em></td>
<td><b>NS</b></td>
<td><em>PS<br/>Structs containing float/double will not be supported.;<br/>UTC is only supported TZ for child TIMESTAMP;<br/>unsupported child types BINARY, CALENDAR, MAP, UDT</em></td>
<td><em>PS<br/>UTC is only supported TZ for child TIMESTAMP;<br/>unsupported child types BINARY, CALENDAR, MAP, UDT</em></td>
<td><b>NS</b></td>
</tr>
<tr>
Expand Down Expand Up @@ -16638,9 +16638,9 @@ are limited.
<td>S</td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><em>PS<br/>Arrays containing floats/doubles will not be supported.;<br/>UTC is only supported TZ for child TIMESTAMP;<br/>unsupported child types BINARY, CALENDAR, MAP, UDT</em></td>
<td><em>PS<br/>UTC is only supported TZ for child TIMESTAMP;<br/>unsupported child types BINARY, CALENDAR, MAP, UDT</em></td>
<td><b>NS</b></td>
<td><em>PS<br/>Structs containing float/double will not be supported.;<br/>UTC is only supported TZ for child TIMESTAMP;<br/>unsupported child types BINARY, CALENDAR, MAP, UDT</em></td>
<td><em>PS<br/>UTC is only supported TZ for child TIMESTAMP;<br/>unsupported child types BINARY, CALENDAR, MAP, UDT</em></td>
<td><b>NS</b></td>
</tr>
<tr>
Expand Down Expand Up @@ -16681,9 +16681,9 @@ are limited.
<td>S</td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><em>PS<br/>Arrays containing floats/doubles will not be supported.;<br/>UTC is only supported TZ for child TIMESTAMP;<br/>unsupported child types BINARY, CALENDAR, MAP, UDT</em></td>
<td><em>PS<br/>UTC is only supported TZ for child TIMESTAMP;<br/>unsupported child types BINARY, CALENDAR, MAP, UDT</em></td>
<td><b>NS</b></td>
<td><em>PS<br/>Structs containing float/double will not be supported.;<br/>UTC is only supported TZ for child TIMESTAMP;<br/>unsupported child types BINARY, CALENDAR, MAP, UDT</em></td>
<td><em>PS<br/>UTC is only supported TZ for child TIMESTAMP;<br/>unsupported child types BINARY, CALENDAR, MAP, UDT</em></td>
<td><b>NS</b></td>
</tr>
<tr>
Expand Down
13 changes: 1 addition & 12 deletions integration_tests/src/main/python/data_gen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1013,9 +1013,6 @@ def gen_scalars_for_sql(data_gen, count, seed=0, force_no_nulls=False):
all_basic_gens_no_nan = [byte_gen, short_gen, int_gen, long_gen, FloatGen(no_nans=True), DoubleGen(no_nans=True),
string_gen, boolean_gen, date_gen, timestamp_gen, null_gen]

all_basic_gens_no_floats = [byte_gen, short_gen, int_gen, long_gen,
string_gen, boolean_gen, date_gen, timestamp_gen, null_gen]

# Many Spark versions have issues sorting large decimals,
# see https://issues.apache.org/jira/browse/SPARK-40089.
orderable_decimal_gen_128bit = decimal_gen_128bit
Expand Down Expand Up @@ -1048,8 +1045,6 @@ def gen_scalars_for_sql(data_gen, count, seed=0, force_no_nulls=False):

single_level_array_gens_no_nan = [ArrayGen(sub_gen) for sub_gen in all_basic_gens_no_nan + decimal_gens]

single_level_array_gens_no_floats = [ArrayGen(sub_gen) for sub_gen in all_basic_gens_no_floats + decimal_gens]

single_level_array_gens_no_decimal = [ArrayGen(sub_gen) for sub_gen in all_basic_gens]

map_string_string_gen = [MapGen(StringGen(pattern='key_[0-9]', nullable=False), StringGen())]
Expand All @@ -1068,13 +1063,7 @@ def gen_scalars_for_sql(data_gen, count, seed=0, force_no_nulls=False):

all_basic_struct_gen_no_nan = StructGen([['child'+str(ind), sub_gen] for ind, sub_gen in enumerate(all_basic_gens_no_nan)])

all_basic_struct_gen_no_floats = StructGen([['child'+str(ind), sub_gen] for ind, sub_gen in enumerate(all_basic_gens_no_floats)])

struct_array_gen = StructGen([['child'+str(ind), sub_gen] for ind, sub_gen in enumerate(single_level_array_gens)])

struct_array_gen_no_nans = StructGen([['child'+str(ind), sub_gen] for ind, sub_gen in enumerate(single_level_array_gens_no_nan)])

struct_array_gen_no_floats = StructGen([['child'+str(ind), sub_gen] for ind, sub_gen in enumerate(single_level_array_gens_no_floats)])
struct_array_gen_no_nans = StructGen([['child'+str(ind), sub_gen] for ind, sub_gen in enumerate(single_level_array_gens_no_nan)])

# Some struct gens, but not all because of nesting
nonempty_struct_gens_sample = [all_basic_struct_gen,
Expand Down
66 changes: 8 additions & 58 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 @@ -603,52 +603,27 @@ def test_hash_pivot_groupby_duplicates_fallback(data_gen):
('b', value_gen)] for value_gen in _repeat_agg_column_for_collect_list_op]

_repeat_agg_column_for_collect_set_op = [
RepeatSeqGen(all_basic_struct_gen_no_floats, length=15),
RepeatSeqGen(all_basic_struct_gen, length=15),
RepeatSeqGen(StructGen([
['c0', all_basic_struct_gen_no_floats], ['c1', int_gen]]), length=15)]

struct_array_gen_no_nans = StructGen([['child'+str(ind), sub_gen] for ind, sub_gen in enumerate(single_level_array_gens_no_nan)])

struct_array_gen_no_floats = StructGen([['child'+str(ind), sub_gen] for ind, sub_gen in enumerate(single_level_array_gens_no_floats)])
['c0', all_basic_struct_gen], ['c1', int_gen]]), length=15)]

# data generating for collect_set based-nested Struct[Array] types
_repeat_agg_column_for_collect_set_op_nested = [
RepeatSeqGen(struct_array_gen_no_floats, length=15),
RepeatSeqGen(struct_array_gen_no_nans, length=15),
RepeatSeqGen(StructGen([
['c0', struct_array_gen_no_floats], ['c1', int_gen]]), length=15),
RepeatSeqGen(ArrayGen(all_basic_struct_gen_no_floats), length=15)]

struct_array_gen_floats = StructGen([['child'+str(ind), sub_gen] for ind, sub_gen
in enumerate([ArrayGen(double_gen), ArrayGen(float_gen)])])
['c0', struct_array_gen_no_nans], ['c1', int_gen]]), length=15),
RepeatSeqGen(ArrayGen(all_basic_struct_gen_no_nan), length=15)]

struct_gen_floats = StructGen([['child'+str(ind), sub_gen] for ind, sub_gen in enumerate([float_gen, double_gen])])

_repeat_agg_column_for_collect_set_op_nested_floats = [
RepeatSeqGen(struct_array_gen_floats, length=15),
RepeatSeqGen(StructGen([
['c0', struct_array_gen_floats], ['c1', int_gen]]), length=15),
RepeatSeqGen(ArrayGen(struct_gen_floats), length=15),
RepeatSeqGen(ArrayGen(ArrayGen(float_gen)), length=15),
RepeatSeqGen(ArrayGen(ArrayGen(double_gen)), length=15)]

_array_of_array_gen = [RepeatSeqGen(ArrayGen(sub_gen), length=15) for sub_gen in single_level_array_gens_no_floats]
_array_of_array_gen = [RepeatSeqGen(ArrayGen(sub_gen), length=15) for sub_gen in single_level_array_gens_no_nan]

_gen_data_for_collect_set_op = [[
('a', RepeatSeqGen(LongGen(), length=20)),
('b', value_gen)] for value_gen in _repeat_agg_column_for_collect_set_op]

_gen_data_for_collect_set_op_floats = [[
('a', RepeatSeqGen(LongGen(), length=20)),
('b', RepeatSeqGen(struct_array_gen, length=15))]]

_gen_data_for_collect_set_op_nested = [[
('a', RepeatSeqGen(LongGen(), length=20)),
('b', value_gen)] for value_gen in _repeat_agg_column_for_collect_set_op_nested + _array_of_array_gen]

_gen_data_for_collect_set_op_nested_floats = [[
('a', RepeatSeqGen(LongGen(), length=20)),
('b', value_gen)] for value_gen in _repeat_agg_column_for_collect_set_op_nested_floats]

_all_basic_gens_with_all_nans_cases = all_basic_gens + [SetValuesGen(t, [math.nan, None]) for t in [FloatType(), DoubleType()]]

# very simple test for just a count on decimals 128 values until we can support more with them
Expand Down Expand Up @@ -719,20 +694,6 @@ def test_hash_groupby_collect_set_on_nested_type(data_gen):
.groupby('a')
.agg(f.sort_array(f.collect_set('b'))))

# Fall back to CPU if type is nested and contains floats or doubles in the type tree.
# Because NaNs in nested types are not supported yet.
# See https://github.com/NVIDIA/spark-rapids/issues/8808
@ignore_order(local=True)
@allow_non_gpu('ObjectHashAggregateExec', 'ShuffleExchangeExec', 'CollectSet')
@pytest.mark.parametrize('data_gen', _gen_data_for_collect_set_op_floats +
_gen_data_for_collect_set_op_nested_floats, ids=idfn)
def test_hash_groupby_collect_set_fallback_on_nested_floats(data_gen):
assert_gpu_fallback_collect(
lambda spark: gen_df(spark, data_gen, length=100)
.groupby('a')
.agg(f.sort_array(f.collect_set('b'))),
'CollectSet')


# Note, using sort_array() on the CPU, because sort_array() does not yet
# support sorting certain nested/arbitrary types on the GPU
Expand Down Expand Up @@ -774,17 +735,6 @@ def test_hash_reduction_collect_set_on_nested_type(data_gen):
lambda spark: gen_df(spark, data_gen, length=100)
.agg(f.sort_array(f.collect_set('b'))))

# Fall back to CPU if type is nested and contains floats or doubles in the type tree.
# Because NaNs in nested types are not supported yet.
# See https://github.com/NVIDIA/spark-rapids/issues/8808
@ignore_order(local=True)
@allow_non_gpu('ObjectHashAggregateExec', 'ShuffleExchangeExec', 'CollectSet')
@pytest.mark.parametrize('data_gen', _gen_data_for_collect_set_op_floats + _gen_data_for_collect_set_op_nested_floats, ids=idfn)
def test_hash_reduction_collect_set_fallback_on_nested_floats(data_gen):
assert_gpu_fallback_collect(
lambda spark: gen_df(spark, data_gen, length=100)
.agg(f.sort_array(f.collect_set('b'))),
'CollectSet')

# Note, using sort_array() on the CPU, because sort_array() does not yet
# support sorting certain nested/arbitrary types on the GPU
Expand Down Expand Up @@ -1225,8 +1175,8 @@ def test_collect_list_reductions(data_gen):
FloatGen(special_cases=[FLOAT_MIN, FLOAT_MAX, 0.0, 1.0, -1.0]), DoubleGen(special_cases=[]),
string_gen, boolean_gen, date_gen, timestamp_gen]

_struct_only_nested_gens = [all_basic_struct_gen_no_floats,
StructGen([['child0', byte_gen], ['child1', all_basic_struct_gen_no_floats]]),
_struct_only_nested_gens = [all_basic_struct_gen,
StructGen([['child0', byte_gen], ['child1', all_basic_struct_gen]]),
StructGen([])]
@pytest.mark.parametrize('data_gen',
_no_neg_zero_all_basic_gens + decimal_gens + _struct_only_nested_gens,
Expand Down
46 changes: 11 additions & 35 deletions integration_tests/src/main/python/window_function_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1124,10 +1124,10 @@ def do_it(spark):
('a', RepeatSeqGen(LongGen(), length=20)),
('b', UniqueLongGen()),
('c_int', RepeatSeqGen(IntegerGen(), length=15)),
('c_struct_array_1', RepeatSeqGen(struct_array_gen_no_floats, length=15)),
('c_struct_array_1', RepeatSeqGen(struct_array_gen_no_nans, length=15)),
('c_struct_array_2', RepeatSeqGen(StructGen([
['c0', struct_array_gen_no_floats], ['c1', int_gen]]), length=14)),
('c_array_struct', RepeatSeqGen(ArrayGen(all_basic_struct_gen_no_floats), length=15)),
['c0', struct_array_gen_no_nans], ['c1', int_gen]]), length=14)),
('c_array_struct', RepeatSeqGen(ArrayGen(all_basic_struct_gen_no_nan), length=15)),
('c_array_array_bool', RepeatSeqGen(ArrayGen(ArrayGen(BooleanGen())), length=15)),
('c_array_array_int', RepeatSeqGen(ArrayGen(ArrayGen(IntegerGen())), length=15)),
('c_array_array_long', RepeatSeqGen(ArrayGen(ArrayGen(LongGen())), length=15)),
Expand All @@ -1136,19 +1136,13 @@ def do_it(spark):
('c_array_array_timestamp', RepeatSeqGen(ArrayGen(ArrayGen(TimestampGen())), length=15)),
('c_array_array_byte', RepeatSeqGen(ArrayGen(ArrayGen(ByteGen())), length=15)),
('c_array_array_string', RepeatSeqGen(ArrayGen(ArrayGen(StringGen())), length=15)),
('c_array_array_float', RepeatSeqGen(ArrayGen(ArrayGen(FloatGen(no_nans=True))), length=15)),
('c_array_array_double', RepeatSeqGen(ArrayGen(ArrayGen(DoubleGen(no_nans=True))), length=15)),
('c_array_array_decimal_32', RepeatSeqGen(ArrayGen(ArrayGen(DecimalGen(precision=8, scale=3))), length=15)),
('c_array_array_decimal_64', RepeatSeqGen(ArrayGen(ArrayGen(decimal_gen_64bit)), length=15)),
('c_array_array_decimal_128', RepeatSeqGen(ArrayGen(ArrayGen(decimal_gen_128bit)), length=15)),
]

_gen_data_for_collect_set_nest_floats_fallback = [
('a', RepeatSeqGen(LongGen(), length=20)),
('b', UniqueLongGen()),
('c_int', RepeatSeqGen(IntegerGen(), length=15)),
('c_array_array_float', RepeatSeqGen(ArrayGen(ArrayGen(FloatGen(no_nans=True))), length=15)),
('c_array_array_double', RepeatSeqGen(ArrayGen(ArrayGen(DoubleGen(no_nans=True))), length=15)),
]

# SortExec does not support array type, so sort the result locally.
@ignore_order(local=True)
def test_window_aggs_for_rows_collect_set():
Expand Down Expand Up @@ -1245,6 +1239,10 @@ def do_it(spark):
(partition by a order by b,c_int rows between CURRENT ROW and UNBOUNDED FOLLOWING) as cc_array_array_byte,
collect_set(c_array_array_string) over
(partition by a order by b,c_int rows between CURRENT ROW and UNBOUNDED FOLLOWING) as cc_array_array_str,
collect_set(c_array_array_float) over
(partition by a order by b,c_int rows between CURRENT ROW and UNBOUNDED FOLLOWING) as cc_array_array_float,
collect_set(c_array_array_double) over
(partition by a order by b,c_int rows between CURRENT ROW and UNBOUNDED FOLLOWING) as cc_array_array_double,
collect_set(c_array_array_decimal_32) over
(partition by a order by b,c_int rows between CURRENT ROW and UNBOUNDED FOLLOWING) as cc_array_array_decimal_32,
collect_set(c_array_array_decimal_64) over
Expand All @@ -1270,37 +1268,15 @@ def do_it(spark):
sort_array(cc_array_array_ts),
sort_array(cc_array_array_byte),
sort_array(cc_array_array_str),
sort_array(cc_array_array_float),
sort_array(cc_array_array_double),
sort_array(cc_array_array_decimal_32),
sort_array(cc_array_array_decimal_64),
sort_array(cc_array_array_decimal_128)
from window_collect_table_2
""")
assert_gpu_and_cpu_are_equal_collect(do_it, conf=conf)

# Fall back to CPU if type is nested and contains floats or doubles in the type tree.
# Because NaNs in nested types are not supported yet.
# See https://github.com/NVIDIA/spark-rapids/issues/8808
@ignore_order(local=True)
@allow_non_gpu("ProjectExec", "CollectSet", "WindowExec")
def test_window_aggs_for_rows_collect_set_nested_array():
conf = copy_and_update(_float_conf, {
"spark.rapids.sql.castFloatToString.enabled": "true",
"spark.rapids.sql.expression.SortArray": "false"
})

def do_it(spark):
df = gen_df(spark, _gen_data_for_collect_set_nest_floats_fallback, length=512)
df.createOrReplaceTempView("window_collect_table")
return spark.sql(
"""select a, b,
collect_set(c_array_array_float) over
(partition by a order by b,c_int rows between CURRENT ROW and UNBOUNDED FOLLOWING) as cc_array_array_float,
collect_set(c_array_array_double) over
(partition by a order by b,c_int rows between CURRENT ROW and UNBOUNDED FOLLOWING) as cc_array_array_double
from window_collect_table
""")

assert_gpu_fallback_collect(do_it, 'CollectSet', conf=conf)

# In a distributed setup the order of the partitions returned might be different, so we must ignore the order
# but small batch sizes can make sort very slow, so do the final order by locally
Expand Down
16 changes: 2 additions & 14 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 @@ -3259,23 +3259,11 @@ object GpuOverrides extends Logging {
Seq(ParamCheck("input",
(TypeSig.commonCudfTypes + TypeSig.DECIMAL_128 +
TypeSig.NULL +
TypeSig.STRUCT.withPsNote(TypeEnum.STRUCT, "Structs containing " +
s"float/double will not be supported.") +
TypeSig.ARRAY.withPsNote(TypeEnum.ARRAY, "Arrays containing " +
s"floats/doubles will not be supported.")).nested(),
TypeSig.STRUCT +
TypeSig.ARRAY).nested(),
TypeSig.all))),
(c, conf, p, r) => new TypedImperativeAggExprMeta[CollectSet](c, conf, p, r) {

override def tagAggForGpu(): Unit = {
// Fall back to CPU if type is nested and contains floats or doubles in the type tree.
// Because NaNs in nested types are not supported yet.
// See https://github.com/NVIDIA/spark-rapids/issues/8808
if (c.child.dataType != FloatType && c.child.dataType != DoubleType &&
isOrContainsFloatingPoint(c.child.dataType)) {
Copy link
Collaborator

@abellina abellina Aug 3, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this check performed a recursive check for floats/doubles at all nested levels, but your changes seem to say that Array[Double] and Array[Float] are UNEQUAL. Should your code handle Array[Array[Double]] (and any level of nesting, including things like Struct[Array[Double]] for example)? Are those cases UNEQUAL or are they somehow compared differently. If so we should add comments clarifying why an array/struct a 2nd+ level is ok.

Copy link
Collaborator Author

@thirtiseven thirtiseven Aug 3, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh yes it's a bit confusing. I added a comment.

In GpuCollectBase we have:

def child: Expression
override def dataType: DataType = ArrayType(child.dataType, containsNull = false)

My change wants to make NaNs UNEQUAL only if child.dataType is double or float, otherwise ALL_EQUAL.

But when we call CudfCollectSet or CudfMergeSets we need to pass dataType to it. So when I check if datatype is Array(Double) or Array(Float), what I really want to check is if child.dataType is float or double.

(I hope I made it clear...)

willNotWorkOnGpu("Float/Double in nested type is not supported")
}
}

override def convertToGpu(childExprs: Seq[Expression]): GpuExpression =
GpuCollectSet(childExprs.head, c.mutableAggBufferOffset, c.inputAggBufferOffset)

Expand Down
Loading