From be57db84f54ddf4f609733eab9bafeebd3ed8785 Mon Sep 17 00:00:00 2001 From: kevinyhzou <37431499+KevinyhZou@users.noreply.github.com> Date: Tue, 9 Jul 2024 11:44:20 +0800 Subject: [PATCH] [GLUTEN-6304][CH]Support array_join (#6305) What changes were proposed in this pull request? (Please fill in changes proposed in this fix) (Fixes: #6304) How was this patch tested? test by spark ut --- .../gluten/utils/CHExpressionUtil.scala | 9 +- .../Functions/SparkFunctionArrayJoin.cpp | 204 ++++++++++++++++++ .../Parser/SerializedPlanParser.h | 3 +- .../clickhouse/ClickHouseTestSettings.scala | 1 - .../clickhouse/ClickHouseTestSettings.scala | 1 - .../clickhouse/ClickHouseTestSettings.scala | 1 - .../clickhouse/ClickHouseTestSettings.scala | 1 - 7 files changed, 214 insertions(+), 6 deletions(-) create mode 100644 cpp-ch/local-engine/Functions/SparkFunctionArrayJoin.cpp diff --git a/backends-clickhouse/src/main/scala/org/apache/gluten/utils/CHExpressionUtil.scala b/backends-clickhouse/src/main/scala/org/apache/gluten/utils/CHExpressionUtil.scala index 14f0ff489188..ac03a7a5b0b9 100644 --- a/backends-clickhouse/src/main/scala/org/apache/gluten/utils/CHExpressionUtil.scala +++ b/backends-clickhouse/src/main/scala/org/apache/gluten/utils/CHExpressionUtil.scala @@ -159,6 +159,13 @@ case class EncodeDecodeValidator() extends FunctionValidator { } } +case class ArrayJoinValidator() extends FunctionValidator { + override def doValidate(expr: Expression): Boolean = expr match { + case t: ArrayJoin => !t.children.head.isInstanceOf[Literal] + case _ => true + } +} + object CHExpressionUtil { final val CH_AGGREGATE_FUNC_BLACKLIST: Map[String, FunctionValidator] = Map( @@ -167,7 +174,7 @@ object CHExpressionUtil { ) final val CH_BLACKLIST_SCALAR_FUNCTION: Map[String, FunctionValidator] = Map( - ARRAY_JOIN -> DefaultValidator(), + ARRAY_JOIN -> ArrayJoinValidator(), SPLIT_PART -> DefaultValidator(), TO_UNIX_TIMESTAMP -> UnixTimeStampValidator(), UNIX_TIMESTAMP -> UnixTimeStampValidator(), diff --git a/cpp-ch/local-engine/Functions/SparkFunctionArrayJoin.cpp b/cpp-ch/local-engine/Functions/SparkFunctionArrayJoin.cpp new file mode 100644 index 000000000000..ed99c0904272 --- /dev/null +++ b/cpp-ch/local-engine/Functions/SparkFunctionArrayJoin.cpp @@ -0,0 +1,204 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include +#include +#include +#include + +using namespace DB; + +namespace DB +{ +namespace ErrorCodes +{ + extern const int ILLEGAL_TYPE_OF_ARGUMENT; + extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH; +} +} + +namespace local_engine +{ +class SparkFunctionArrayJoin : public IFunction +{ +public: + static constexpr auto name = "sparkArrayJoin"; + static FunctionPtr create(ContextPtr) { return std::make_shared(); } + SparkFunctionArrayJoin() = default; + ~SparkFunctionArrayJoin() override = default; + bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo &) const override { return true; } + size_t getNumberOfArguments() const override { return 0; } + String getName() const override { return name; } + bool isVariadic() const override { return true; } + bool useDefaultImplementationForNulls() const override { return false; } + + DB::DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName &) const override + { + auto data_type = std::make_shared(); + return makeNullable(data_type); + } + + ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t) const override + { + if (arguments.size() != 2 && arguments.size() != 3) + throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH, "Function {} must have 2 or 3 arguments", getName()); + + const auto * arg_null_col = checkAndGetColumn(arguments[0].column.get()); + const ColumnArray * array_col; + if (!arg_null_col) + array_col = checkAndGetColumn(arguments[0].column.get()); + else + array_col = checkAndGetColumn(arg_null_col->getNestedColumnPtr().get()); + if (!array_col) + throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Function {} 1st argument must be array type", getName()); + + auto res_col = ColumnString::create(); + auto null_col = ColumnUInt8::create(array_col->size(), 0); + PaddedPODArray & null_result = null_col->getData(); + std::pair delim_p, null_replacement_p; + bool return_result = false; + auto checkAndGetConstString = [&](const ColumnPtr & col) -> std::pair + { + StringRef res; + const auto * str_null_col = checkAndGetColumnConstData(col.get()); + if (str_null_col) + { + if (str_null_col->isNullAt(0)) + { + for (size_t i = 0; i < array_col->size(); ++i) + { + res_col->insertDefault(); + null_result[i] = 1; + } + return_result = true; + return std::pair(false, res); + } + } + else + { + const auto * string_col = checkAndGetColumnConstData(col.get()); + if (!string_col) + return std::pair(false, res); + else + return std::pair(true, string_col->getDataAt(0)); + } + }; + delim_p = checkAndGetConstString(arguments[1].column); + if (return_result) + return ColumnNullable::create(std::move(res_col), std::move(null_col)); + + if (arguments.size() == 3) + { + null_replacement_p = checkAndGetConstString(arguments[2].column); + if (return_result) + return ColumnNullable::create(std::move(res_col), std::move(null_col)); + } + const ColumnNullable * array_nested_col = checkAndGetColumn(&array_col->getData()); + const ColumnString * string_col; + if (array_nested_col) + string_col = checkAndGetColumn(array_nested_col->getNestedColumnPtr().get()); + else + string_col = checkAndGetColumn(&array_col->getData()); + const ColumnArray::Offsets & array_offsets = array_col->getOffsets(); + const ColumnString::Offsets & string_offsets = string_col->getOffsets(); + const ColumnString::Chars & string_data = string_col->getChars(); + const ColumnNullable * delim_col = checkAndGetColumn(arguments[1].column.get()); + const ColumnNullable * null_replacement_col = arguments.size() == 3 ? checkAndGetColumn(arguments[2].column.get()) : nullptr; + size_t current_offset = 0, array_pos = 0; + for (size_t i = 0; i < array_col->size(); ++i) + { + String res; + auto setResultNull = [&]() -> void + { + res_col->insertDefault(); + null_result[i] = 1; + current_offset = array_offsets[i]; + }; + auto getDelimiterOrNullReplacement = [&](const std::pair & s, const ColumnNullable * col) -> StringRef + { + if (s.first) + return s.second; + else + { + if (col->isNullAt(i)) + return StringRef(nullptr, 0); + else + { + const ColumnString * col_string = checkAndGetColumn(col->getNestedColumnPtr().get()); + return col_string->getDataAt(i); + } + } + }; + if (arg_null_col->isNullAt(i)) + { + setResultNull(); + continue; + } + const StringRef delim = getDelimiterOrNullReplacement(delim_p, delim_col); + if (!delim.data) + { + setResultNull(); + continue; + } + StringRef null_replacement; + if (arguments.size() == 3) + { + null_replacement = getDelimiterOrNullReplacement(null_replacement_p, null_replacement_col); + if (!null_replacement.data) + { + setResultNull(); + continue; + } + } + + size_t array_size = array_offsets[i] - current_offset; + size_t data_pos = array_pos == 0 ? 0 : string_offsets[array_pos - 1]; + for (size_t j = 0; j < array_size; ++j) + { + if (array_nested_col && array_nested_col->isNullAt(j + array_pos)) + { + if (null_replacement.data) + { + res += null_replacement.toString(); + if (j != array_size - 1) + res += delim.toString(); + } + } + else + { + const StringRef s(&string_data[data_pos], string_offsets[j + array_pos] - data_pos - 1); + res += s.toString(); + if (j != array_size - 1) + res += delim.toString(); + } + data_pos = string_offsets[j + array_pos]; + } + array_pos += array_size; + res_col->insertData(res.data(), res.length()); + current_offset = array_offsets[i]; + } + return ColumnNullable::create(std::move(res_col), std::move(null_col)); + } +}; + +REGISTER_FUNCTION(SparkArrayJoin) +{ + factory.registerFunction(); +} +} \ No newline at end of file diff --git a/cpp-ch/local-engine/Parser/SerializedPlanParser.h b/cpp-ch/local-engine/Parser/SerializedPlanParser.h index 90086ea28649..477fdb1f6d44 100644 --- a/cpp-ch/local-engine/Parser/SerializedPlanParser.h +++ b/cpp-ch/local-engine/Parser/SerializedPlanParser.h @@ -177,7 +177,8 @@ static const std::map SCALAR_FUNCTIONS {"array", "array"}, {"shuffle", "arrayShuffle"}, {"range", "range"}, /// dummy mapping - {"flatten", "sparkArrayFlatten"}, + {"flatten", "sparkArrayFlatten"}, + {"array_join", "sparkArrayJoin"}, // map functions {"map", "map"}, diff --git a/gluten-ut/spark32/src/test/scala/org/apache/gluten/utils/clickhouse/ClickHouseTestSettings.scala b/gluten-ut/spark32/src/test/scala/org/apache/gluten/utils/clickhouse/ClickHouseTestSettings.scala index a7ffbc9fa7f6..8fd68d5170b9 100644 --- a/gluten-ut/spark32/src/test/scala/org/apache/gluten/utils/clickhouse/ClickHouseTestSettings.scala +++ b/gluten-ut/spark32/src/test/scala/org/apache/gluten/utils/clickhouse/ClickHouseTestSettings.scala @@ -665,7 +665,6 @@ class ClickHouseTestSettings extends BackendTestSettings { .exclude("Map Concat") .exclude("MapFromEntries") .exclude("ArraysOverlap") - .exclude("ArrayJoin") .exclude("ArraysZip") .exclude("Sequence of numbers") .exclude("Sequence of timestamps") diff --git a/gluten-ut/spark33/src/test/scala/org/apache/gluten/utils/clickhouse/ClickHouseTestSettings.scala b/gluten-ut/spark33/src/test/scala/org/apache/gluten/utils/clickhouse/ClickHouseTestSettings.scala index ceb0d8a87e21..f69598adf555 100644 --- a/gluten-ut/spark33/src/test/scala/org/apache/gluten/utils/clickhouse/ClickHouseTestSettings.scala +++ b/gluten-ut/spark33/src/test/scala/org/apache/gluten/utils/clickhouse/ClickHouseTestSettings.scala @@ -656,7 +656,6 @@ class ClickHouseTestSettings extends BackendTestSettings { .exclude("Map Concat") .exclude("MapFromEntries") .exclude("ArraysOverlap") - .exclude("ArrayJoin") .exclude("ArraysZip") .exclude("Sequence of numbers") .exclude("Sequence of timestamps") diff --git a/gluten-ut/spark34/src/test/scala/org/apache/gluten/utils/clickhouse/ClickHouseTestSettings.scala b/gluten-ut/spark34/src/test/scala/org/apache/gluten/utils/clickhouse/ClickHouseTestSettings.scala index 66007a36770f..ab288e835b12 100644 --- a/gluten-ut/spark34/src/test/scala/org/apache/gluten/utils/clickhouse/ClickHouseTestSettings.scala +++ b/gluten-ut/spark34/src/test/scala/org/apache/gluten/utils/clickhouse/ClickHouseTestSettings.scala @@ -544,7 +544,6 @@ class ClickHouseTestSettings extends BackendTestSettings { .exclude("Map Concat") .exclude("MapFromEntries") .exclude("ArraysOverlap") - .exclude("ArrayJoin") .exclude("ArraysZip") .exclude("Sequence of numbers") .exclude("Sequence of timestamps") diff --git a/gluten-ut/spark35/src/test/scala/org/apache/gluten/utils/clickhouse/ClickHouseTestSettings.scala b/gluten-ut/spark35/src/test/scala/org/apache/gluten/utils/clickhouse/ClickHouseTestSettings.scala index 66007a36770f..ab288e835b12 100644 --- a/gluten-ut/spark35/src/test/scala/org/apache/gluten/utils/clickhouse/ClickHouseTestSettings.scala +++ b/gluten-ut/spark35/src/test/scala/org/apache/gluten/utils/clickhouse/ClickHouseTestSettings.scala @@ -544,7 +544,6 @@ class ClickHouseTestSettings extends BackendTestSettings { .exclude("Map Concat") .exclude("MapFromEntries") .exclude("ArraysOverlap") - .exclude("ArrayJoin") .exclude("ArraysZip") .exclude("Sequence of numbers") .exclude("Sequence of timestamps")