From cbdbc3dcd9306cae414a491eae5ac479b456ac0c Mon Sep 17 00:00:00 2001 From: Bobby Wang Date: Wed, 19 May 2021 15:51:23 +0800 Subject: [PATCH] rebase --- integration_tests/src/main/python/array_test.py | 13 ++----------- .../com/nvidia/spark/rapids/GpuOverrides.scala | 12 +++++++++++- .../spark/sql/rapids/complexTypeCreator.scala | 4 ++-- 3 files changed, 15 insertions(+), 14 deletions(-) diff --git a/integration_tests/src/main/python/array_test.py b/integration_tests/src/main/python/array_test.py index f46f16840ea8..b07777553139 100644 --- a/integration_tests/src/main/python/array_test.py +++ b/integration_tests/src/main/python/array_test.py @@ -55,17 +55,8 @@ def test_make_array(data_gen): assert_gpu_and_cpu_are_equal_collect( lambda spark : binary_op_df(spark, data_gen).selectExpr( 'array(a, b)', - 'array(b, a, null, {}, {})'.format(s1, s2))) - - -@pytest.mark.parametrize('data_gen', all_basic_gens + [decimal_gen_default, decimal_gen_scale_precision], ids=idfn) -def test_make_array_of_array(data_gen): - (s1, s2) = gen_scalars_for_sql(data_gen, 2, force_no_nulls=not isinstance(data_gen, NullGen)) - assert_gpu_and_cpu_are_equal_collect( - lambda spark : binary_op_df(spark, data_gen).selectExpr( - 'array(array(1, 2, 3), array(2), array(null), array())', - 'array(array(), array(null), array(a, b))', - 'array(array(b, a, null, {}, {}), array(a, b), array(), array(null))'.format(s1, s2))) + 'array(b, a, null, {}, {})'.format(s1, s2), + 'array(array(b, a, null, {}, {}), array(a), array(null))'.format(s1, s2))) @pytest.mark.parametrize('data_gen', single_level_array_gens, ids=idfn) diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala index 2db7288a68e0..0c386dbe17b6 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala @@ -2366,9 +2366,19 @@ object GpuOverrides { TypeSig.numeric + TypeSig.NULL + TypeSig.STRING + TypeSig.BOOLEAN + TypeSig.DATE + TypeSig.TIMESTAMP + TypeSig.ARRAY.nested(TypeSig.numeric + TypeSig.NULL + TypeSig.STRING + - TypeSig.BOOLEAN + TypeSig.DATE + TypeSig.TIMESTAMP), + TypeSig.BOOLEAN + TypeSig.DATE + TypeSig.TIMESTAMP + TypeSig.ARRAY), TypeSig.all))), (in, conf, p, r) => new ExprMeta[CreateArray](in, conf, p, r) { + + override def tagExprForGpu(): Unit = { + wrapped.dataType match { + case ArrayType(ArrayType(ArrayType(_, _), _), _) => + willNotWorkOnGpu("Only support to create array or array of array, Found: " + + s"${wrapped.dataType}") + case _ => + } + } + override def convertToGpu(): GpuExpression = GpuCreateArray(childExprs.map(_.convertToGpu()), wrapped.useStringTypeWhenEmpty) }), diff --git a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/complexTypeCreator.scala b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/complexTypeCreator.scala index eb7d33e3f916..5e8318fa03ba 100644 --- a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/complexTypeCreator.scala +++ b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/complexTypeCreator.scala @@ -17,9 +17,9 @@ package org.apache.spark.sql.rapids import ai.rapids.cudf.{ColumnVector, DType} - -import com.nvidia.spark.rapids.{GpuColumnVector, GpuExpression, GpuExpressionsUtils, GpuScalar} +import com.nvidia.spark.rapids.{GpuColumnVector, GpuExpression, GpuExpressionsUtils} import com.nvidia.spark.rapids.RapidsPluginImplicits.ReallyAGpuExpression + import org.apache.spark.sql.catalyst.analysis.{TypeCheckResult, TypeCoercion} import org.apache.spark.sql.catalyst.analysis.FunctionRegistry.FUNC_ALIAS import org.apache.spark.sql.catalyst.expressions.{EmptyRow, Expression, NamedExpression}