This is an automated email from the ASF dual-hosted git repository.
loneylee pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/incubator-gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 07b8679a5c [GLUTEN-8748][CH] Support function
monotonically_increasing_id (#8771)
07b8679a5c is described below
commit 07b8679a5c9240441978b8b18c9da311e6eeba4d
Author: Shuai li <[email protected]>
AuthorDate: Wed Feb 19 15:35:38 2025 +0800
[GLUTEN-8748][CH] Support function monotonically_increasing_id (#8771)
* [GLUTEN-8748][CH] Support function monotonically_increasing_id
---
.../vectorized/CHNativeExpressionEvaluator.java | 14 ++-
.../vectorized/ExpressionEvaluatorJniWrapper.java | 3 +-
.../backendsapi/clickhouse/CHIteratorApi.scala | 20 +++-
.../execution/NativeFileScanColumnarRDD.scala | 3 +-
.../org/apache/gluten/utils/CHExpressionUtil.scala | 1 -
.../metrics/GlutenClickHouseMetricsUTUtils.scala | 3 +-
.../SparkFunctionMonotonicallyIncreasingID.cpp | 106 +++++++++++++++++++++
cpp-ch/local-engine/Parser/ParserContext.cpp | 8 +-
cpp-ch/local-engine/Parser/ParserContext.h | 7 +-
.../monotonicallyIncreasingId.cpp | 49 ++++++++++
cpp-ch/local-engine/local_engine_jni.cpp | 5 +-
11 files changed, 201 insertions(+), 18 deletions(-)
diff --git
a/backends-clickhouse/src/main/java/org/apache/gluten/vectorized/CHNativeExpressionEvaluator.java
b/backends-clickhouse/src/main/java/org/apache/gluten/vectorized/CHNativeExpressionEvaluator.java
index 12b79a54ce..31bb2e16fb 100644
---
a/backends-clickhouse/src/main/java/org/apache/gluten/vectorized/CHNativeExpressionEvaluator.java
+++
b/backends-clickhouse/src/main/java/org/apache/gluten/vectorized/CHNativeExpressionEvaluator.java
@@ -64,7 +64,8 @@ public class CHNativeExpressionEvaluator extends
ExpressionEvaluatorJniWrapper {
byte[] wsPlan,
byte[][] splitInfo,
List<ColumnarNativeIterator> iterList,
- boolean materializeInput) {
+ boolean materializeInput,
+ int partitionIndex) {
CHThreadGroup.registerNewThreadGroup();
long handle =
nativeCreateKernelWithIterator(
@@ -72,13 +73,17 @@ public class CHNativeExpressionEvaluator extends
ExpressionEvaluatorJniWrapper {
splitInfo,
iterList.toArray(new ColumnarNativeIterator[0]),
ConfigUtil.serialize(getNativeBackendConf()),
- materializeInput);
+ materializeInput,
+ partitionIndex);
return createBatchIterator(handle);
}
// Only for UT.
public static BatchIterator createKernelWithBatchIterator(
- byte[] wsPlan, byte[][] splitInfo, List<ColumnarNativeIterator>
iterList) {
+ byte[] wsPlan,
+ byte[][] splitInfo,
+ List<ColumnarNativeIterator> iterList,
+ int partitionIndex) {
CHThreadGroup.registerNewThreadGroup();
long handle =
nativeCreateKernelWithIterator(
@@ -86,7 +91,8 @@ public class CHNativeExpressionEvaluator extends
ExpressionEvaluatorJniWrapper {
splitInfo,
iterList.toArray(new ColumnarNativeIterator[0]),
ConfigUtil.serialize(getNativeBackendConf()),
- false);
+ false,
+ partitionIndex);
return createBatchIterator(handle);
}
diff --git
a/backends-clickhouse/src/main/java/org/apache/gluten/vectorized/ExpressionEvaluatorJniWrapper.java
b/backends-clickhouse/src/main/java/org/apache/gluten/vectorized/ExpressionEvaluatorJniWrapper.java
index 76fddcafdd..20a4241798 100644
---
a/backends-clickhouse/src/main/java/org/apache/gluten/vectorized/ExpressionEvaluatorJniWrapper.java
+++
b/backends-clickhouse/src/main/java/org/apache/gluten/vectorized/ExpressionEvaluatorJniWrapper.java
@@ -41,7 +41,8 @@ public class ExpressionEvaluatorJniWrapper {
byte[][] splitInfo,
ColumnarNativeIterator[] batchItr,
byte[] confArray,
- boolean materializeInput);
+ boolean materializeInput,
+ int partitionIndex);
public static native void updateQueryRuntimeSettings(byte[] settings);
}
diff --git
a/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHIteratorApi.scala
b/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHIteratorApi.scala
index a4e641a98c..308b3fa4f5 100644
---
a/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHIteratorApi.scala
+++
b/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHIteratorApi.scala
@@ -73,6 +73,7 @@ class CHIteratorApi extends IteratorApi with Logging with
LogLevelUtil {
splitInfoByteArray: Array[Array[Byte]],
wsPlan: Array[Byte],
materializeInput: Boolean,
+ partitionIndex: Int,
inputIterators: Seq[Iterator[ColumnarBatch]]): BatchIterator = {
/** Generate closeable ColumnBatch iterator. */
@@ -88,7 +89,8 @@ class CHIteratorApi extends IteratorApi with Logging with
LogLevelUtil {
wsPlan,
splitInfoByteArray,
listIterator,
- materializeInput
+ materializeInput,
+ partitionIndex
)
}
@@ -290,7 +292,13 @@ class CHIteratorApi extends IteratorApi with Logging with
LogLevelUtil {
pipelineTime,
updateNativeMetrics,
Some(updateInputMetrics),
- createNativeIterator(splitInfoByteArray, wsPlan, materializeInput,
inputIterators))
+ createNativeIterator(
+ splitInfoByteArray,
+ wsPlan,
+ materializeInput,
+ partitionIndex,
+ inputIterators)
+ )
)
}
@@ -318,7 +326,13 @@ class CHIteratorApi extends IteratorApi with Logging with
LogLevelUtil {
pipelineTime,
updateNativeMetrics,
None,
- createNativeIterator(splitInfoByteArray, wsPlan, materializeInput,
inputIterators))
+ createNativeIterator(
+ splitInfoByteArray,
+ wsPlan,
+ materializeInput,
+ partitionIndex,
+ inputIterators)
+ )
}
}
diff --git
a/backends-clickhouse/src/main/scala/org/apache/gluten/execution/NativeFileScanColumnarRDD.scala
b/backends-clickhouse/src/main/scala/org/apache/gluten/execution/NativeFileScanColumnarRDD.scala
index fe29b6d2f2..abe8ea0b94 100644
---
a/backends-clickhouse/src/main/scala/org/apache/gluten/execution/NativeFileScanColumnarRDD.scala
+++
b/backends-clickhouse/src/main/scala/org/apache/gluten/execution/NativeFileScanColumnarRDD.scala
@@ -54,7 +54,8 @@ class NativeFileScanColumnarRDD(
inputPartition.plan,
splitInfoByteArray,
inBatchIters,
- false
+ false,
+ split.index
)
}
TaskContext
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 dc2d8716b2..4a9ce71c79 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
@@ -186,7 +186,6 @@ object CHExpressionUtil {
DECODE -> EncodeDecodeValidator(),
ENCODE -> EncodeDecodeValidator(),
DATE_FROM_UNIX_DATE -> DefaultValidator(),
- MONOTONICALLY_INCREASING_ID -> DefaultValidator(),
SPARK_PARTITION_ID -> DefaultValidator(),
AT_LEAST_N_NON_NULLS -> DefaultValidator(),
URL_DECODE -> DefaultValidator(),
diff --git
a/backends-clickhouse/src/test/scala/org/apache/gluten/execution/metrics/GlutenClickHouseMetricsUTUtils.scala
b/backends-clickhouse/src/test/scala/org/apache/gluten/execution/metrics/GlutenClickHouseMetricsUTUtils.scala
index 52a68e2d62..3a2be45ff8 100644
---
a/backends-clickhouse/src/test/scala/org/apache/gluten/execution/metrics/GlutenClickHouseMetricsUTUtils.scala
+++
b/backends-clickhouse/src/test/scala/org/apache/gluten/execution/metrics/GlutenClickHouseMetricsUTUtils.scala
@@ -47,7 +47,8 @@ object GlutenClickHouseMetricsUTUtils {
val resIter = CHNativeExpressionEvaluator.createKernelWithBatchIterator(
substraitPlan.toByteArray,
new Array[Array[Byte]](0),
- inBatchIters)
+ inBatchIters,
+ 0)
val iter = new Iterator[Any] {
private var outputRowCount = 0L
private var outputVectorCount = 0L
diff --git
a/cpp-ch/local-engine/Functions/SparkFunctionMonotonicallyIncreasingID.cpp
b/cpp-ch/local-engine/Functions/SparkFunctionMonotonicallyIncreasingID.cpp
new file mode 100644
index 0000000000..6c3e654bf8
--- /dev/null
+++ b/cpp-ch/local-engine/Functions/SparkFunctionMonotonicallyIncreasingID.cpp
@@ -0,0 +1,106 @@
+/*
+ * 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 <atomic>
+#include <Columns/ColumnsNumber.h>
+#include <DataTypes/DataTypesNumber.h>
+#include <Functions/FunctionFactory.h>
+#include <Functions/IFunction.h>
+
+namespace DB
+{
+namespace ErrorCodes
+{
+extern const int ILLEGAL_TYPE_OF_ARGUMENT;
+extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
+}
+}
+
+namespace local_engine
+{
+/**
+ * Returns monotonically increasing 64-bit integers. The generated ID is
guaranteed
+ * to be monotonically increasing and unique, but not consecutive. The
current implementation
+ * puts the partition ID in the upper 31 bits, and the lower 33 bits
represent the record number
+ * within each partition. The assumption is that the data frame has less
than 1 billion
+ * partitions, and each partition has less than 8 billion records.
+ * The function is non-deterministic because its result depends on
partition IDs.
+ */
+class SparkFunctionMonotonicallyIncreasingID : public DB::IFunction
+{
+public:
+ static constexpr auto name = "sparkMonotonicallyIncreasingId";
+ static DB::FunctionPtr create(DB::ContextPtr) { return
std::make_shared<SparkFunctionMonotonicallyIncreasingID>(); }
+
+ String getName() const override { return name; }
+
+ bool isStateful() const override { return true; }
+
+ bool isSuitableForShortCircuitArgumentsExecution(const
DB::DataTypesWithConstInfo & /*arguments*/) const override { return false; }
+
+ size_t getNumberOfArguments() const override { return 1; }
+ DB::DataTypePtr getReturnTypeImpl(const DB::DataTypes & arguments) const
override
+ {
+ if (arguments.size() != 1)
+ throw DB::Exception(
+ DB::ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH,
+ "Number of arguments for function {} doesn't match: passed {},
should be 1.",
+ getName(),
+ arguments.size());
+
+ if (!isInteger(arguments[0]))
+ throw DB::Exception(DB::ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
"Argument for function {} must be integer", getName());
+
+ return std::make_shared<DB::DataTypeInt64>();
+ }
+
+ DB::ColumnPtr
+ executeImplDryRun(const DB::ColumnsWithTypeAndName &, const
DB::DataTypePtr & result_type, size_t input_rows_count) const override
+ {
+ return DB::ColumnInt64::create(input_rows_count);
+ }
+
+ DB::ColumnPtr
+ executeImpl(const DB::ColumnsWithTypeAndName & arguments, const
DB::DataTypePtr & result_type, size_t input_rows_count) const override
+ {
+ const auto partitionIndex = arguments[0].column->getInt(0);
+ const auto partitionMask = partitionIndex << 33;
+
+ auto result = result_type->createColumn();
+ result->reserve(input_rows_count);
+
+ for (size_t i = 0; i < input_rows_count; ++i)
+ {
+ auto a=partitionMask + count.load(std::memory_order_relaxed);
+ result->insert(a);
+ count.fetch_add(1, std::memory_order_relaxed);
+ }
+
+ return result;
+ }
+
+private:
+ mutable std::atomic<size_t> count{0};
+};
+
+
+REGISTER_FUNCTION(SparkFunctionMonotonicallyIncreasingID)
+{
+ factory.registerFunction<SparkFunctionMonotonicallyIncreasingID>();
+}
+
+}
\ No newline at end of file
diff --git a/cpp-ch/local-engine/Parser/ParserContext.cpp
b/cpp-ch/local-engine/Parser/ParserContext.cpp
index 3e5d6fffe7..884d8152fc 100644
--- a/cpp-ch/local-engine/Parser/ParserContext.cpp
+++ b/cpp-ch/local-engine/Parser/ParserContext.cpp
@@ -39,11 +39,13 @@ void ParserContext::addSelfDefinedFunctionMapping()
legacy_function_mapping[std::to_string(SelfDefinedFunctionReference::GET_JSON_OBJECT)]
= sig.signature;
}
-ParserContextPtr ParserContext::build(DB::ContextPtr context_, const
std::unordered_map<String, String> & function_mapping_)
+ParserContextPtr
+ParserContext::build(DB::ContextPtr context_, const std::unordered_map<String,
String> & function_mapping_, const size_t & partition_index_)
{
auto res = std::make_shared<ParserContext>();
res->query_context = context_;
res->legacy_function_mapping = function_mapping_;
+ res->partition_index = partition_index_;
auto & function_mapping = res->function_mapping;
for (const auto & fm : function_mapping_)
{
@@ -66,7 +68,7 @@ ParserContextPtr ParserContext::build(DB::ContextPtr context_)
return res;
}
-ParserContextPtr ParserContext::build(DB::ContextPtr context_, const
substrait::Plan & plan_)
+ParserContextPtr ParserContext::build(DB::ContextPtr context_, const
substrait::Plan & plan_, const size_t & partition_index_)
{
std::unordered_map<String, String> function_mapping;
for (const auto & extension : plan_.extensions())
@@ -77,7 +79,7 @@ ParserContextPtr ParserContext::build(DB::ContextPtr
context_, const substrait::
std::to_string(extension.extension_function().function_anchor()),
extension.extension_function().name());
}
}
- return build(context_, function_mapping);
+ return build(context_, function_mapping, partition_index_);
}
DB::ContextPtr ParserContext::queryContext() const
diff --git a/cpp-ch/local-engine/Parser/ParserContext.h
b/cpp-ch/local-engine/Parser/ParserContext.h
index c5c40199a2..2b591146cc 100644
--- a/cpp-ch/local-engine/Parser/ParserContext.h
+++ b/cpp-ch/local-engine/Parser/ParserContext.h
@@ -38,9 +38,10 @@ public:
ParserContext() = default;
using Ptr = std::shared_ptr<const ParserContext>;
- static Ptr build(DB::ContextPtr context_, const std::unordered_map<String,
String> & function_mapping_);
+ static Ptr
+ build(DB::ContextPtr context_, const std::unordered_map<String, String> &
function_mapping_, const size_t & partition_index_ = 0);
static Ptr build(DB::ContextPtr context_);
- static Ptr build(DB::ContextPtr context_, const substrait::Plan & plan_);
+ static Ptr build(DB::ContextPtr context_, const substrait::Plan & plan_,
const size_t & partition_index_ = 0);
DB::ContextPtr queryContext() const;
@@ -49,11 +50,13 @@ public:
std::optional<String> getFunctionNameInSignature(Int32 ref) const;
std::optional<String> getFunctionNameInSignature(const
substrait::Expression_ScalarFunction & func) const;
const std::unordered_map<String, String> & getLegacyFunctionMapping()
const;
+ size_t getPartitionIndex() const { return partition_index; }
private:
DB::ContextPtr query_context;
std::unordered_map<String, String> legacy_function_mapping;
std::unordered_map<Int32, FunctionSignature> function_mapping;
+ size_t partition_index = 0;
void addSelfDefinedFunctionMapping();
};
diff --git
a/cpp-ch/local-engine/Parser/scalar_function_parser/monotonicallyIncreasingId.cpp
b/cpp-ch/local-engine/Parser/scalar_function_parser/monotonicallyIncreasingId.cpp
new file mode 100644
index 0000000000..e9bdf60e4d
--- /dev/null
+++
b/cpp-ch/local-engine/Parser/scalar_function_parser/monotonicallyIncreasingId.cpp
@@ -0,0 +1,49 @@
+/*
+ * 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 <DataTypes/DataTypesNumber.h>
+#include <Parser/FunctionParser.h>
+#include <Common/Exception.h>
+
+namespace local_engine
+{
+class SparkFunctionMonotonicallyIncreasingIDParser : public FunctionParser
+{
+public:
+ explicit SparkFunctionMonotonicallyIncreasingIDParser(ParserContextPtr
parser_context_) : FunctionParser(parser_context_) { }
+ ~SparkFunctionMonotonicallyIncreasingIDParser() override = default;
+ static constexpr auto name = "monotonically_increasing_id";
+ String getName() const { return name; }
+ String getCHFunctionName(const substrait::Expression_ScalarFunction &)
const override { return "sparkMonotonicallyIncreasingId"; }
+
+ const DB::ActionsDAG::Node *
+ parse(const substrait::Expression_ScalarFunction & substrait_func,
DB::ActionsDAG & actions_dag) const override
+ {
+ const DB::DataTypePtr type = std::make_shared<DB::DataTypeInt64>();
+ auto partition_index = parser_context->getPartitionIndex();
+
+ DB::ActionsDAG::NodeRawConstPtrs args;
+ args.reserve(1);
+ args.emplace_back(&actions_dag.addColumn(
+ DB::ColumnWithTypeAndName(type->createColumnConst(1,
partition_index), type, getUniqueName("partition_index"))));
+
+ auto ch_function_name = getCHFunctionName(substrait_func);
+ const auto * func_node = toFunctionNode(actions_dag, ch_function_name,
args);
+ return convertNodeTypeIfNeeded(substrait_func, func_node, actions_dag);
+ }
+};
+static FunctionParserRegister<SparkFunctionMonotonicallyIncreasingIDParser>
register_monotonically_increasing_id;
+}
diff --git a/cpp-ch/local-engine/local_engine_jni.cpp
b/cpp-ch/local-engine/local_engine_jni.cpp
index 6f3df78fb6..3d742e274d 100644
--- a/cpp-ch/local-engine/local_engine_jni.cpp
+++ b/cpp-ch/local-engine/local_engine_jni.cpp
@@ -229,7 +229,8 @@ JNIEXPORT jlong
Java_org_apache_gluten_vectorized_ExpressionEvaluatorJniWrapper_
jobjectArray split_infos,
jobjectArray iter_arr,
jbyteArray conf_plan,
- jboolean materialize_input)
+ jboolean materialize_input,
+ jint partition_index)
{
LOCAL_ENGINE_JNI_METHOD_START
auto query_context =
local_engine::QueryContext::instance().currentQueryContext();
@@ -243,7 +244,7 @@ JNIEXPORT jlong
Java_org_apache_gluten_vectorized_ExpressionEvaluatorJniWrapper_
const std::string::size_type plan_size = plan_a.length();
auto plan_pb =
local_engine::BinaryToMessage<substrait::Plan>({reinterpret_cast<const char
*>(plan_a.elems()), plan_size});
- auto parser_context = local_engine::ParserContext::build(query_context,
plan_pb);
+ auto parser_context = local_engine::ParserContext::build(query_context,
plan_pb, partition_index);
local_engine::SerializedPlanParser parser(parser_context);
jsize iter_num = env->GetArrayLength(iter_arr);
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
