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

Support Union on Struct of Map #3061

Merged
merged 1 commit into from
Jul 29, 2021
Merged

Support Union on Struct of Map #3061

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
2 changes: 1 addition & 1 deletion docs/supported_ops.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -422,7 +422,7 @@ Accelerator supports are described below.
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><b>NS</b></td>
<td><em>PS* (unionByName will not optionally impute nulls for missing struct fields when the column is a struct and there are non-overlapping fields; missing nested BINARY, CALENDAR, ARRAY, MAP, UDT)</em></td>
<td><em>PS* (unionByName will not optionally impute nulls for missing struct fields when the column is a struct and there are non-overlapping fields; missing nested BINARY, CALENDAR, ARRAY, UDT)</em></td>
<td><b>NS</b></td>
</tr>
<tr>
Expand Down
27 changes: 22 additions & 5 deletions integration_tests/src/main/python/repart_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,8 @@
from marks import ignore_order, allow_non_gpu
import pyspark.sql.functions as f

nested_scalar_mark=pytest.mark.xfail(reason="https://github.com/NVIDIA/spark-rapids/issues/1459")
nested_scalar_mark = pytest.mark.xfail(reason="https://github.com/NVIDIA/spark-rapids/issues/1459")

# 4 level nested struct
# each level has a different number of children to avoid a bug in spark < 3.1
nested_struct = StructGen([
Expand All @@ -37,6 +38,15 @@
])],
['child1', BooleanGen()]
])]])

# map generators without ArrayType value, since Union on ArrayType is not supported
map_gens = [simple_string_to_string_map_gen,
MapGen(RepeatSeqGen(IntegerGen(nullable=False), 10), long_gen, max_length=10),
MapGen(BooleanGen(nullable=False), boolean_gen, max_length=2),
MapGen(StringGen(pattern='key_[0-9]', nullable=False), simple_string_to_string_map_gen)]
struct_of_maps = StructGen([['child0', BooleanGen()]] + [
['child%d' % (i + 1), gen] for i, gen in enumerate(map_gens)])

@pytest.mark.parametrize('data_gen', [pytest.param((StructGen([['child0', DecimalGen(7, 2)]]),
StructGen([['child1', IntegerGen()]])), marks=nested_scalar_mark),
# left_struct(child0 = 4 level nested struct, child1 = Int)
Expand All @@ -51,7 +61,11 @@
StructGen([['child0', DecimalGen(7, 2)]])]])]])]])]], nullable=False),
StructGen([['child1', IntegerGen()]], nullable=False)),
(StructGen([['child0', DecimalGen(7, 2)]], nullable=False),
StructGen([['child1', IntegerGen()]], nullable=False))], ids=idfn)
StructGen([['child1', IntegerGen()]], nullable=False)),
# left_struct(child0 = Map[String, String], child1 = missing map)
# right_struct(child0 = missing map, child1 = Map[Boolean, Boolean])
(StructGen([['child0', simple_string_to_string_map_gen]], nullable=False),
StructGen([['child1', MapGen(BooleanGen(nullable=False), boolean_gen)]], nullable=False))], ids=idfn)
@pytest.mark.skipif(is_before_spark_311(), reason="This is supported only in Spark 3.1.1+")
# This tests the union of DF of structs with different types of cols as long as the struct itself
# isn't null. This is a limitation in cudf because we don't support nested types as literals
Expand All @@ -62,7 +76,8 @@ def test_union_struct_missing_children(data_gen):
spark, right_gen), True))

@pytest.mark.parametrize('data_gen', all_gen + [all_basic_struct_gen, StructGen([['child0', DecimalGen(7, 2)]]),
nested_struct], ids=idfn)
nested_struct,
struct_of_maps], ids=idfn)
# This tests union of two DFs of two cols each. The types of the left col and right col is the same
def test_union(data_gen):
assert_gpu_and_cpu_are_equal_collect(
Expand All @@ -72,7 +87,8 @@ def test_union(data_gen):
pytest.param(StructGen([[ 'child0', DecimalGen(7, 2)]]), marks=nested_scalar_mark),
nested_struct,
StructGen([['child0', StructGen([['child0', StructGen([['child0', StructGen([['child0',
StructGen([['child0', DecimalGen(7, 2)]])]])]])]])], ['child1', IntegerGen()]])], ids=idfn)
StructGen([['child0', DecimalGen(7, 2)]])]])]])]])], ['child1', IntegerGen()]]),
struct_of_maps], ids=idfn)
@pytest.mark.skipif(is_before_spark_311(), reason="This is supported only in Spark 3.1.1+")
# This tests the union of two DFs of structs with missing child column names. The missing child
# column will be replaced by nulls in the output DF. This is a feature added in 3.1+
Expand All @@ -82,7 +98,8 @@ def test_union_by_missing_col_name(data_gen):
.unionByName(binary_op_df(spark, data_gen).withColumnRenamed("a", "y"), True))

@pytest.mark.parametrize('data_gen', all_gen + [all_basic_struct_gen, StructGen([['child0', DecimalGen(7, 2)]]),
nested_struct], ids=idfn)
nested_struct,
struct_of_maps], ids=idfn)
def test_union_by_name(data_gen):
assert_gpu_and_cpu_are_equal_collect(
lambda spark : binary_op_df(spark, data_gen).unionByName(binary_op_df(spark, data_gen)))
Expand Down
2 changes: 1 addition & 1 deletion 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 @@ -2980,7 +2980,7 @@ object GpuOverrides {
"The backend for the union operator",
ExecChecks(TypeSig.commonCudfTypes + TypeSig.NULL + TypeSig.DECIMAL +
TypeSig.STRUCT.nested(TypeSig.commonCudfTypes + TypeSig.NULL + TypeSig.DECIMAL
+ TypeSig.STRUCT)
+ TypeSig.STRUCT + TypeSig.MAP)
.withPsNote(TypeEnum.STRUCT,
"unionByName will not optionally impute nulls for missing struct fields " +
"when the column is a struct and there are non-overlapping fields"), TypeSig.all),
Expand Down