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

Fix collect_set_on_nested_type tests failed #8783

Merged
merged 6 commits into from
Jul 28, 2023
Merged

Fix collect_set_on_nested_type tests failed #8783

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
fb1c522
Fix collect_set_on_nested_type test failed
thirtiseven Jul 24, 2023
1d6c327
Fallback to cpu if double/float in nested type
thirtiseven Jul 26, 2023
7170209
Add fallback tests
thirtiseven Jul 26, 2023
bd9791c
Update fallback tests
thirtiseven Jul 27, 2023
33b01ea
add psnote and tests
thirtiseven Jul 28, 2023
6b91e5e
verify
thirtiseven Jul 28, 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
13 changes: 12 additions & 1 deletion 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,6 +1013,9 @@ 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 @@ -1045,6 +1048,8 @@ 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 @@ -1063,7 +1068,13 @@ 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)])

struct_array_gen_no_nans = StructGen([['child'+str(ind), sub_gen] for ind, sub_gen in enumerate(single_level_array_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)])

# Some struct gens, but not all because of nesting
nonempty_struct_gens_sample = [all_basic_struct_gen,
Expand Down
50 changes: 40 additions & 10 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,23 +603,31 @@ 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, length=15),
RepeatSeqGen(all_basic_struct_gen_no_floats, length=15),
RepeatSeqGen(StructGen([
['c0', all_basic_struct_gen], ['c1', int_gen]]), length=15)]
['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)])

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

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

_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]
Expand Down Expand Up @@ -688,13 +696,25 @@ def test_hash_groupby_collect_set(data_gen):

@ignore_order(local=True)
@pytest.mark.parametrize('data_gen', _gen_data_for_collect_set_op, ids=idfn)
@pytest.mark.xfail(reason='https://github.com/NVIDIA/spark-rapids/issues/8716')
def test_hash_groupby_collect_set_on_nested_type(data_gen):
assert_gpu_and_cpu_are_equal_collect(
lambda spark: gen_df(spark, data_gen, length=100)
.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, 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 @@ -731,12 +751,22 @@ def test_hash_reduction_collect_set(data_gen):

@ignore_order(local=True)
@pytest.mark.parametrize('data_gen', _gen_data_for_collect_set_op, ids=idfn)
@pytest.mark.xfail(reason='https://github.com/NVIDIA/spark-rapids/issues/8716')
def test_hash_reduction_collect_set_on_nested_type(data_gen):
assert_gpu_and_cpu_are_equal_collect(
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, 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 @@ -1177,8 +1207,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,
StructGen([['child0', byte_gen], ['child1', all_basic_struct_gen]]),
_struct_only_nested_gens = [all_basic_struct_gen_no_floats,
StructGen([['child0', byte_gen], ['child1', all_basic_struct_gen_no_floats]]),
StructGen([])]
@pytest.mark.parametrize('data_gen',
_no_neg_zero_all_basic_gens + decimal_gens + _struct_only_nested_gens,
Expand Down
46 changes: 35 additions & 11 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 @@ -1120,10 +1120,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_nans, length=15)),
('c_struct_array_1', RepeatSeqGen(struct_array_gen_no_floats, length=15)),
('c_struct_array_2', RepeatSeqGen(StructGen([
['c0', struct_array_gen_no_nans], ['c1', int_gen]]), length=14)),
('c_array_struct', RepeatSeqGen(ArrayGen(all_basic_struct_gen_no_nan), length=15)),
['c0', struct_array_gen_no_floats], ['c1', int_gen]]), length=14)),
Copy link
Collaborator

Choose a reason for hiding this comment

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

nit: Can we verify that we fallback for window operations too?

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Added.

('c_array_struct', RepeatSeqGen(ArrayGen(all_basic_struct_gen_no_floats), 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 @@ -1132,13 +1132,19 @@ 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 @@ -1235,10 +1241,6 @@ 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 @@ -1264,15 +1266,37 @@ 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
19 changes: 7 additions & 12 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 @@ -3264,18 +3264,13 @@ object GpuOverrides extends Logging {
TypeSig.all))),
(c, conf, p, r) => new TypedImperativeAggExprMeta[CollectSet](c, conf, p, r) {

private def isNestedArrayType(dt: DataType): Boolean = {
dt match {
case StructType(fields) =>
fields.exists { field =>
field.dataType match {
case sdt: StructType => isNestedArrayType(sdt)
case _: ArrayType => true
case _ => false
}
}
case ArrayType(et, _) => et.isInstanceOf[ArrayType] || et.isInstanceOf[StructType]
case _ => false
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)) {
willNotWorkOnGpu("Float/Double in nested type is not supported")
}
}

Expand Down