This is an automated email from the ASF dual-hosted git repository.
zhouyuan pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 5ada13649a [VL] Translate Kafka Substrait reads (#13236)
5ada13649a is described below
commit 5ada13649a5248d76a0c1e397b991f3adc7ee568
Author: Reema <[email protected]>
AuthorDate: Fri Oct 23 11:05:19 2026 +0300
[VL] Translate Kafka Substrait reads (#13236)
---
.github/workflows/velox_kafka.yml | 9 +-
cpp/velox/compute/VeloxPlanConverter.cc | 53 ++++++--
cpp/velox/compute/VeloxPlanConverter.h | 11 ++
cpp/velox/compute/VeloxRuntime.cc | 25 ++--
cpp/velox/compute/VeloxRuntime.h | 2 +
cpp/velox/compute/WholeStageResultIterator.cc | 14 ++
cpp/velox/operators/reader/KafkaSplit.h | 4 +-
cpp/velox/operators/reader/KafkaSplitInfo.h | 42 ++++++
cpp/velox/substrait/SubstraitToVeloxPlan.cc | 81 +++++++++++-
cpp/velox/substrait/SubstraitToVeloxPlan.h | 7 +-
.../substrait/SubstraitToVeloxPlanValidator.h | 1 +
.../tests/Substrait2VeloxPlanConversionTest.cc | 145 +++++++++++++++++++++
.../tests/Substrait2VeloxPlanValidatorTest.cc | 14 ++
13 files changed, 373 insertions(+), 35 deletions(-)
diff --git a/.github/workflows/velox_kafka.yml
b/.github/workflows/velox_kafka.yml
index 3f2a367e13..f4e050866d 100644
--- a/.github/workflows/velox_kafka.yml
+++ b/.github/workflows/velox_kafka.yml
@@ -82,11 +82,10 @@ jobs:
--enable_s3=OFF --enable_gcs=OFF --enable_hdfs=OFF
--enable_abfs=OFF
- name: Run native Kafka connector tests
run: ./cpp/build/velox/tests/velox_kafka_connector_test
- # TODO: Re-enable when the Kafka plan-conversion tests are added in
#13236.
- # - name: Run Kafka and file plan conversion tests
- # run: |
- # ./cpp/build/velox/tests/velox_plan_conversion_test \
- #
--gtest_filter='Substrait2VeloxPlanConversionTest.serializedFileSplit:Substrait2VeloxPlanConversionTest.kafkaReadValidation:Substrait2VeloxPlanConversionTest.mixedKafkaAndFileSplits:Substrait2VeloxPlanValidatorTest.kafkaReadRequiresNativeSupport'
+ - name: Run Kafka and file plan conversion tests
+ run: |
+ ./cpp/build/velox/tests/velox_plan_conversion_test \
+
--gtest_filter='Substrait2VeloxPlanConversionTest.serializedFileSplit:Substrait2VeloxPlanConversionTest.kafka*:Substrait2VeloxPlanConversionTest.mixedKafkaAndFileSplits:Substrait2VeloxPlanValidatorTest.kafkaReadRequiresNativeSupport'
- uses: actions/upload-artifact@v7
with:
name: velox-kafka-native-${{ github.sha }}
diff --git a/cpp/velox/compute/VeloxPlanConverter.cc
b/cpp/velox/compute/VeloxPlanConverter.cc
index 2b5fb8bd9a..3cd8ff43ca 100644
--- a/cpp/velox/compute/VeloxPlanConverter.cc
+++ b/cpp/velox/compute/VeloxPlanConverter.cc
@@ -22,10 +22,14 @@
#include <google/protobuf/any.pb.h>
#include <google/protobuf/wrappers.pb.h>
+#include "compute/ProtobufUtils.h"
#include "config/GlutenConfig.h"
#include "delta/DeltaSplitInfo.h"
#include "iceberg/IcebergPlanConverter.h"
#include "operators/plannodes/IteratorSplit.h"
+#ifdef ENABLE_KAFKA
+#include "operators/reader/KafkaSplitInfo.h"
+#endif
namespace gluten {
@@ -148,7 +152,9 @@ std::shared_ptr<DeltaSplitInfo> parseDeltaSplitInfo(
return deltaSplitInfo;
}
-std::shared_ptr<SplitInfo> parseScanSplitInfo(
+} // namespace
+
+std::shared_ptr<SplitInfo> VeloxPlanConverter::parseScanSplitInfo(
const facebook::velox::config::ConfigBase* veloxCfg,
const substrait::ReadRel_LocalFiles& localFiles) {
using SubstraitFileFormatCase =
::substrait::ReadRel_LocalFiles_FileOrFiles::FileFormatCase;
@@ -254,25 +260,48 @@ std::shared_ptr<SplitInfo> parseScanSplitInfo(
return splitInfo;
}
-void parseLocalFileNodes(
- SubstraitToVeloxPlanConverter* planConverter,
+std::shared_ptr<SplitInfo> VeloxPlanConverter::parseSplitInfo(
const facebook::velox::config::ConfigBase* veloxCfg,
- std::vector<::substrait::ReadRel_LocalFiles>& localFiles) {
- std::vector<std::shared_ptr<SplitInfo>> splitInfos;
- splitInfos.reserve(localFiles.size());
- for (const auto& localFile : localFiles) {
- splitInfos.push_back(parseScanSplitInfo(veloxCfg, localFile));
+ const uint8_t* data,
+ int32_t size) {
+#ifdef ENABLE_KAFKA
+ // Kafka splits carry an explicit Any type URL. Parsing success alone cannot
distinguish an
+ // ExtensionTable from LocalFiles because both messages use field 1 for an
embedded message.
+ ::substrait::ReadRel_ExtensionTable extensionTable;
+ if (parseProtobuf(data, size, &extensionTable) &&
extensionTable.detail().Is<StreamKafka>()) {
+ auto splitInfo = std::make_shared<KafkaSplitInfo>();
+ splitInfo->leafType = SplitInfo::LeafType::TABLE_SCAN;
+ VELOX_CHECK(
+ extensionTable.detail().UnpackTo(&splitInfo->streamKafka),
+ "Kafka split info must contain a gluten.StreamKafka payload");
+ return splitInfo;
}
-
- planConverter->setSplitInfos(std::move(splitInfos));
+#endif
+ ::substrait::ReadRel_LocalFiles localFiles;
+ VELOX_CHECK(parseProtobuf(data, size, &localFiles), "Parse split info
failed");
+ return parseScanSplitInfo(veloxCfg, localFiles);
}
-} // namespace
std::shared_ptr<const facebook::velox::core::PlanNode>
VeloxPlanConverter::toVeloxPlan(
const ::substrait::Plan& substraitPlan,
std::vector<::substrait::ReadRel_LocalFiles> localFiles) {
+ std::vector<std::shared_ptr<SplitInfo>> splitInfos;
+ if (!validationMode_) {
+ splitInfos.reserve(localFiles.size());
+ for (const auto& localFile : localFiles) {
+ splitInfos.push_back(parseScanSplitInfo(veloxCfg_, localFile));
+ }
+ }
+ return toVeloxPlanWithSplitInfos(substraitPlan, std::move(splitInfos));
+}
+
+std::shared_ptr<const facebook::velox::core::PlanNode>
VeloxPlanConverter::toVeloxPlanWithSplitInfos(
+ const ::substrait::Plan& substraitPlan,
+ std::vector<std::shared_ptr<SplitInfo>> splitInfos) {
if (!validationMode_) {
- parseLocalFileNodes(&substraitVeloxPlanConverter_, veloxCfg_, localFiles);
+ // Runtime splits are already decoded as file or Kafka SplitInfo objects
in scan traversal order.
+ // Transfer the shared pointers without copying their payloads; each
physical scan consumes one.
+ substraitVeloxPlanConverter_.setSplitInfos(std::move(splitInfos));
}
return substraitVeloxPlanConverter_.toVeloxPlan(substraitPlan);
diff --git a/cpp/velox/compute/VeloxPlanConverter.h
b/cpp/velox/compute/VeloxPlanConverter.h
index 1aee2c36bd..bebc6f04a4 100644
--- a/cpp/velox/compute/VeloxPlanConverter.h
+++ b/cpp/velox/compute/VeloxPlanConverter.h
@@ -43,6 +43,17 @@ class VeloxPlanConverter {
const ::substrait::Plan& substraitPlan,
std::vector<::substrait::ReadRel_LocalFiles> localFiles);
+ std::shared_ptr<const facebook::velox::core::PlanNode>
toVeloxPlanWithSplitInfos(
+ const ::substrait::Plan& substraitPlan,
+ std::vector<std::shared_ptr<SplitInfo>> splitInfos);
+
+ static std::shared_ptr<SplitInfo>
+ parseSplitInfo(const facebook::velox::config::ConfigBase* veloxCfg, const
uint8_t* data, int32_t size);
+
+ static std::shared_ptr<SplitInfo> parseScanSplitInfo(
+ const facebook::velox::config::ConfigBase* veloxCfg,
+ const ::substrait::ReadRel_LocalFiles& localFiles);
+
const std::unordered_map<facebook::velox::core::PlanNodeId,
std::shared_ptr<SplitInfo>>& splitInfos() {
return substraitVeloxPlanConverter_.splitInfos();
}
diff --git a/cpp/velox/compute/VeloxRuntime.cc
b/cpp/velox/compute/VeloxRuntime.cc
index 4420976572..2dbb53bcb2 100644
--- a/cpp/velox/compute/VeloxRuntime.cc
+++ b/cpp/velox/compute/VeloxRuntime.cc
@@ -66,6 +66,10 @@
DECLARE_bool(velox_memory_pool_capacity_transfer_across_tasks);
#include "operators/serializer/VeloxGpuColumnarBatchSerializer.h"
#endif
+#ifdef ENABLE_KAFKA
+#include "operators/reader/KafkaSplitInfo.h"
+#endif
+
using namespace facebook;
namespace gluten {
@@ -391,22 +395,27 @@ void VeloxRuntime::parsePlan(const uint8_t* data, int32_t
size) {
}
void VeloxRuntime::parseSplitInfo(const uint8_t* data, int32_t size, int32_t
splitIndex) {
+ auto splitInfo = VeloxPlanConverter::parseSplitInfo(veloxCfg_.get(), data,
size);
if (debugModeEnabled_ || dumper_ != nullptr) {
+ std::string_view splitType = "ReadRel.LocalFiles";
+#ifdef ENABLE_KAFKA
+ if (std::dynamic_pointer_cast<KafkaSplitInfo>(splitInfo)) {
+ splitType = "ReadRel.ExtensionTable";
+ }
+#endif
try {
- auto splitJson = substraitFromPbToJson("ReadRel.LocalFiles", data, size);
+ auto splitJson = substraitFromPbToJson(splitType, data, size);
if (dumper_ != nullptr) {
dumper_->dumpInputSplit(splitIndex, splitJson);
}
LOG_IF(INFO, debugModeEnabled_ && taskInfo_.has_value())
- << std::string(50, '#') << " received substrait::ReadRel.LocalFiles:
" << taskInfo_.value() << std::endl
+ << std::string(50, '#') << " received substrait::" << splitType <<
": " << taskInfo_.value() << std::endl
<< splitJson;
} catch (const std::exception& e) {
- LOG(WARNING) << "Error converting substrait::ReadRel.LocalFiles to JSON:
" << e.what();
+ LOG(WARNING) << "Error converting substrait::" << splitType << " to
JSON: " << e.what();
}
}
- ::substrait::ReadRel_LocalFiles localFile;
- GLUTEN_CHECK(parseProtobuf(data, size, &localFile) == true, "Parse substrait
plan failed");
- localFiles_.push_back(localFile);
+ scanSplitInfos_.emplace_back(std::move(splitInfo));
}
void VeloxRuntime::getInfoAndIds(
@@ -444,7 +453,7 @@ std::string VeloxRuntime::planString(bool details, const
std::unordered_map<std:
auto veloxMemoryPool = gluten::defaultLeafVeloxMemoryPool();
VeloxPlanConverter veloxPlanConverter(
veloxMemoryPool.get(), veloxCfg_.get(), {}, connectorIds_, std::nullopt,
std::nullopt, true);
- auto veloxPlan = veloxPlanConverter.toVeloxPlan(substraitPlan_, localFiles_);
+ auto veloxPlan = veloxPlanConverter.toVeloxPlan(substraitPlan_, {});
return veloxPlan->toString(details, true);
}
@@ -466,7 +475,7 @@ std::shared_ptr<ResultIterator>
VeloxRuntime::createResultIterator(
connectorIds_,
*localWriteFilesTempPath(),
*localWriteFileName());
- veloxPlan_ = veloxPlanConverter.toVeloxPlan(substraitPlan_,
std::move(localFiles_));
+ veloxPlan_ = veloxPlanConverter.toVeloxPlanWithSplitInfos(substraitPlan_,
std::move(scanSplitInfos_));
LOG_IF(INFO, debugModeEnabled_ && taskInfo_.has_value())
<< "############### Velox plan for task " << taskInfo_.value() << "
###############" << std::endl
<< veloxPlan_->toString(true, true);
diff --git a/cpp/velox/compute/VeloxRuntime.h b/cpp/velox/compute/VeloxRuntime.h
index c854fb2461..bebf3c7e69 100644
--- a/cpp/velox/compute/VeloxRuntime.h
+++ b/cpp/velox/compute/VeloxRuntime.h
@@ -153,6 +153,8 @@ class VeloxRuntime final : public Runtime {
std::unique_ptr<folly::Executor> spillExecutor_;
std::unique_ptr<folly::Executor> ioExecutor_;
VeloxConnectorIds connectorIds_;
+ // File SplitInfo and KafkaSplitInfo objects in JVM scan traversal order.
+ std::vector<std::shared_ptr<SplitInfo>> scanSplitInfos_;
std::unordered_map<int32_t, std::shared_ptr<VeloxColumnarBatch>>
emptySchemaBatchLoopUp_;
};
diff --git a/cpp/velox/compute/WholeStageResultIterator.cc
b/cpp/velox/compute/WholeStageResultIterator.cc
index 5fd4f8b878..71e0d20a94 100644
--- a/cpp/velox/compute/WholeStageResultIterator.cc
+++ b/cpp/velox/compute/WholeStageResultIterator.cc
@@ -30,6 +30,9 @@
#include "utils/ConfigExtractor.h"
#include "velox/connectors/hive/HiveConfig.h"
#include "velox/connectors/hive/HiveConnectorSplit.h"
+#ifdef ENABLE_KAFKA
+#include "operators/reader/KafkaSplitInfo.h"
+#endif
#include "velox/exec/PlanNodeStats.h"
#include "velox/functions/sparksql/SparkQueryConfig.h"
#ifdef GLUTEN_ENABLE_GPU
@@ -144,6 +147,14 @@ WholeStageResultIterator::WholeStageResultIterator(
for (size_t scanInfoIdx = 0; scanInfoIdx < scanInfos.size(); ++scanInfoIdx) {
const auto& scanInfo = scanInfos[scanInfoIdx];
+#ifdef ENABLE_KAFKA
+ if (auto kafkaSplitInfo =
std::dynamic_pointer_cast<KafkaSplitInfo>(scanInfo)) {
+ std::vector<velox::exec::Split> kafkaSplits;
+
kafkaSplits.emplace_back(kafkaSplitInfo->toConnectorSplit(connectorIds_.kafka));
+ splits_.emplace_back(std::move(kafkaSplits));
+ continue;
+ }
+#endif
// Get the information for TableScan.
// Partition index in scan info is not used.
const auto& paths = scanInfo->paths;
@@ -268,6 +279,9 @@ std::shared_ptr<velox::core::QueryCtx>
WholeStageResultIterator::createNewVeloxQ
connectorConfigs[connectorIds_.iceberg] = hiveSessionConfig;
connectorConfigs[connectorIds_.delta] = hiveSessionConfig;
connectorConfigs[connectorIds_.iterator] = hiveSessionConfig;
+#ifdef ENABLE_KAFKA
+ connectorConfigs[connectorIds_.kafka] = hiveSessionConfig;
+#endif
#ifdef GLUTEN_ENABLE_GPU
if (!connectorIds_.cudfHive.empty()) {
connectorConfigs[connectorIds_.cudfHive] = hiveSessionConfig;
diff --git a/cpp/velox/operators/reader/KafkaSplit.h
b/cpp/velox/operators/reader/KafkaSplit.h
index df1e48a3c7..231e6c3889 100644
--- a/cpp/velox/operators/reader/KafkaSplit.h
+++ b/cpp/velox/operators/reader/KafkaSplit.h
@@ -27,8 +27,8 @@
namespace gluten {
/// Split for one Kafka topic partition, covering offsets [startOffset,
endOffset).
-/// Built from the substrait ReadRel.StreamKafka split info sent by the JVM,
which
-/// mirrors Spark's KafkaBatchInputPartition.
+/// Built from the gluten.StreamKafka payload in the ReadRel.ExtensionTable
split
+/// info sent by the JVM, which mirrors Spark's KafkaBatchInputPartition.
class KafkaConnectorSplit : public facebook::velox::connector::ConnectorSplit {
public:
KafkaConnectorSplit(
diff --git a/cpp/velox/operators/reader/KafkaSplitInfo.h
b/cpp/velox/operators/reader/KafkaSplitInfo.h
new file mode 100644
index 0000000000..036ce2b586
--- /dev/null
+++ b/cpp/velox/operators/reader/KafkaSplitInfo.h
@@ -0,0 +1,42 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include "kafka.pb.h"
+#include "operators/reader/KafkaSplit.h"
+#include "substrait/SubstraitToVeloxPlan.h"
+
+namespace gluten {
+
+struct KafkaSplitInfo : SplitInfo {
+ StreamKafka streamKafka;
+
+ std::shared_ptr<KafkaConnectorSplit> toConnectorSplit(const std::string&
connectorId) const {
+ return std::make_shared<KafkaConnectorSplit>(
+ connectorId,
+ streamKafka.topic_partition().topic(),
+ streamKafka.topic_partition().partition(),
+ streamKafka.start_offset(),
+ streamKafka.end_offset(),
+ streamKafka.poll_timeout_ms(),
+ streamKafka.fail_on_data_loss(),
+ std::unordered_map<std::string,
std::string>(streamKafka.params().begin(), streamKafka.params().end()));
+ }
+};
+
+} // namespace gluten
diff --git a/cpp/velox/substrait/SubstraitToVeloxPlan.cc
b/cpp/velox/substrait/SubstraitToVeloxPlan.cc
index dc9833663a..52e62bfadd 100644
--- a/cpp/velox/substrait/SubstraitToVeloxPlan.cc
+++ b/cpp/velox/substrait/SubstraitToVeloxPlan.cc
@@ -25,6 +25,11 @@
#include "jni/JniHashTable.h"
#include "operators/hashjoin/HashTableBuilder.h"
#include "operators/plannodes/RowVectorStream.h"
+#ifdef ENABLE_KAFKA
+#include "kafka.pb.h"
+#include "operators/reader/KafkaConnector.h"
+#include "operators/reader/KafkaSplitInfo.h"
+#endif
#include "velox/connectors/hive/HiveDataSink.h"
#include "velox/connectors/hive/iceberg/IcebergColumnHandle.h"
#include "velox/exec/TableWriter.h"
@@ -52,12 +57,15 @@ using namespace cudf_velox::connector::hive;
namespace gluten {
namespace {
-bool useCudfTableHandle(const std::vector<std::shared_ptr<SplitInfo>>&
splitInfos) {
+#ifdef ENABLE_KAFKA
+bool isKafkaRead(const ::substrait::ReadRel& readRel) {
+ return readRel.has_extension_table() &&
readRel.extension_table().detail().Is<StreamKafka>();
+}
+#endif
+
+bool useCudfTableHandle(const std::shared_ptr<SplitInfo>& splitInfo) {
#ifdef GLUTEN_ENABLE_GPU
- if (splitInfos.empty()) {
- return false;
- }
- return splitInfos[0]->canUseCudfConnector();
+ return splitInfo->canUseCudfConnector();
#else
return false;
#endif
@@ -1511,6 +1519,52 @@ core::PlanNodePtr
SubstraitToVeloxPlanConverter::constructCudfValueStreamNode(
}
#endif
+#ifdef ENABLE_KAFKA
+core::PlanNodePtr
SubstraitToVeloxPlanConverter::constructKafkaStreamNode(const
::substrait::ReadRel& readRel) {
+ VELOX_USER_CHECK(!readRel.has_filter(), "Filter pushdown is not supported
for Kafka reads.");
+ // File and Kafka scans consume the same typed split sequence in scan
traversal order.
+ auto splitInfo = std::make_shared<KafkaSplitInfo>();
+ splitInfo->leafType = SplitInfo::LeafType::TABLE_SCAN;
+ if (!validationMode_) {
+ VELOX_CHECK_LT(splitInfoIdx_, splitInfos_.size(), "Plan must have readRel
and related split info.");
+ splitInfo =
std::dynamic_pointer_cast<KafkaSplitInfo>(splitInfos_[splitInfoIdx_++]);
+ VELOX_CHECK_NOT_NULL(splitInfo, "Kafka read requires Kafka split info");
+ }
+
+ std::vector<std::string> colNameList;
+ std::vector<TypePtr> veloxTypeList;
+ bool asLowerCase = !veloxCfg_->get<bool>(kCaseSensitive, false);
+ if (readRel.has_base_schema()) {
+ const auto& baseSchema = readRel.base_schema();
+ colNameList.reserve(baseSchema.names().size());
+ for (const auto& name : baseSchema.names()) {
+ std::string fieldName = name;
+ if (asLowerCase) {
+ folly::toLowerAscii(fieldName);
+ }
+ colNameList.emplace_back(fieldName);
+ }
+ veloxTypeList = SubstraitParser::parseNamedStruct(baseSchema, asLowerCase);
+ }
+
+ std::vector<std::string> outNames;
+ outNames.reserve(colNameList.size());
+ connector::ColumnHandleMap assignments;
+ for (int idx = 0; idx < colNameList.size(); idx++) {
+ auto outName = SubstraitParser::makeNodeName(planNodeId_, idx);
+ assignments[outName] =
std::make_shared<KafkaColumnHandle>(colNameList[idx], veloxTypeList[idx]);
+ outNames.emplace_back(outName);
+ }
+ auto outputType = ROW(std::move(outNames), std::move(veloxTypeList));
+
+ auto tableHandle = std::make_shared<KafkaTableHandle>(connectorIds_.kafka);
+ auto tableScanNode = std::make_shared<core::TableScanNode>(
+ nextPlanNodeId(), std::move(outputType), std::move(tableHandle),
assignments);
+ splitInfoMap_[tableScanNode->id()] = splitInfo;
+ return tableScanNode;
+}
+#endif
+
core::PlanNodePtr SubstraitToVeloxPlanConverter::constructValuesNode(
const ::substrait::ReadRel& readRel,
int32_t streamIdx) {
@@ -1539,6 +1593,16 @@ core::PlanNodePtr
SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait::
!readRel.common().has_emit(), "Emit not supported for ValuesNode and
TableScanNode related Substrait plans.");
}
+#ifdef ENABLE_KAFKA
+ if (isKafkaRead(readRel)) {
+ return constructKafkaStreamNode(readRel);
+ }
+#else
+ if (readRel.has_extension_table()) {
+ VELOX_USER_FAIL("The plan contains an extension table read, but Gluten was
built without ENABLE_KAFKA.");
+ }
+#endif
+
auto streamIdx = getStreamIndex(readRel);
if (streamIdx >= 0) {
// Check if the ReadRel specifies an input of stream. If yes, build
TableScanNode with iterator connector.
@@ -1561,6 +1625,9 @@ core::PlanNodePtr
SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait::
if (!validationMode_) {
VELOX_CHECK_LT(splitInfoIdx_, splitInfos_.size(), "Plan must have readRel
and related split info.");
splitInfo = splitInfos_[splitInfoIdx_++];
+#ifdef ENABLE_KAFKA
+ VELOX_CHECK(!std::dynamic_pointer_cast<KafkaSplitInfo>(splitInfo), "File
read requires local file split info");
+#endif
}
// Get output names and types.
@@ -1625,8 +1692,8 @@ core::PlanNodePtr
SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait::
if (std::dynamic_pointer_cast<IcebergSplitInfo>(splitInfo)) {
connectorId = connectorIds_.iceberg;
}
- if ((connectorId == connectorIds_.hive || connectorId ==
connectorIds_.iceberg) && useCudfTableHandle(splitInfos_) &&
- veloxCfg_->get<bool>(kCudfEnableTableScan, kCudfEnableTableScanDefault)
&&
+ if (!validationMode_ && (connectorId == connectorIds_.hive || connectorId ==
connectorIds_.iceberg) &&
+ useCudfTableHandle(splitInfo) &&
veloxCfg_->get<bool>(kCudfEnableTableScan, kCudfEnableTableScanDefault) &&
veloxCfg_->get<bool>(kCudfEnabled, kCudfEnabledDefault)) {
#ifdef GLUTEN_ENABLE_GPU
connectorId = connectorIds_.cudfHive;
diff --git a/cpp/velox/substrait/SubstraitToVeloxPlan.h
b/cpp/velox/substrait/SubstraitToVeloxPlan.h
index 2464881052..281cdd22bd 100644
--- a/cpp/velox/substrait/SubstraitToVeloxPlan.h
+++ b/cpp/velox/substrait/SubstraitToVeloxPlan.h
@@ -159,6 +159,11 @@ class SubstraitToVeloxPlanConverter {
// Construct a cuDF value stream node.
core::PlanNodePtr constructCudfValueStreamNode(const ::substrait::ReadRel&
sRead, int32_t streamIdx);
+#ifdef ENABLE_KAFKA
+ // Construct a Kafka stream node for reading from Kafka topics.
+ core::PlanNodePtr constructKafkaStreamNode(const ::substrait::ReadRel&
sRead);
+#endif
+
// This is only used in benchmark and enable query trace, which will load
all the data to ValuesNode.
core::PlanNodePtr constructValuesNode(const ::substrait::ReadRel& sRead,
int32_t streamIdx);
@@ -201,7 +206,7 @@ class SubstraitToVeloxPlanConverter {
}
void setSplitInfos(std::vector<std::shared_ptr<SplitInfo>> splitInfos) {
- splitInfos_ = splitInfos;
+ splitInfos_ = std::move(splitInfos);
}
/// The input iterators not inlined to VeloxPlan. They should be then
manually added to the Velox task
diff --git a/cpp/velox/substrait/SubstraitToVeloxPlanValidator.h
b/cpp/velox/substrait/SubstraitToVeloxPlanValidator.h
index 3f83c4c32c..cb8b322301 100644
--- a/cpp/velox/substrait/SubstraitToVeloxPlanValidator.h
+++ b/cpp/velox/substrait/SubstraitToVeloxPlanValidator.h
@@ -45,6 +45,7 @@ class SubstraitToVeloxPlanValidator {
.hive = kHiveConnectorId,
.iceberg = kIcebergConnectorId,
.iterator = kIteratorConnectorId,
+ .kafka = kKafkaConnectorId,
.cudfHive = kCudfHiveConnectorId},
std::nullopt,
std::nullopt,
diff --git a/cpp/velox/tests/Substrait2VeloxPlanConversionTest.cc
b/cpp/velox/tests/Substrait2VeloxPlanConversionTest.cc
index aa02a3f4e8..323b6873fe 100644
--- a/cpp/velox/tests/Substrait2VeloxPlanConversionTest.cc
+++ b/cpp/velox/tests/Substrait2VeloxPlanConversionTest.cc
@@ -29,6 +29,10 @@
#include "FilePathGenerator.h"
#include "compute/VeloxBackend.h"
+#ifdef ENABLE_KAFKA
+#include "kafka.pb.h"
+#include "operators/reader/KafkaSplitInfo.h"
+#endif
using namespace facebook::velox;
using namespace facebook::velox::test;
@@ -39,6 +43,13 @@ namespace gluten {
class Substrait2VeloxPlanConversionTest : public
exec::test::HiveConnectorTestBase {
protected:
+#ifdef ENABLE_KAFKA
+ std::shared_ptr<SplitInfo> parseSerializedSplitInfo(const std::string&
serialized) {
+ return VeloxPlanConverter::parseSplitInfo(
+ veloxCfg_.get(), reinterpret_cast<const uint8_t*>(serialized.data()),
static_cast<int32_t>(serialized.size()));
+ }
+#endif
+
std::vector<std::shared_ptr<facebook::velox::connector::ConnectorSplit>>
makeSplits(
std::shared_ptr<const core::PlanNode> planNode) {
const auto& splitInfos = planConverter_->splitInfos();
@@ -427,4 +438,138 @@ TEST_F(Substrait2VeloxPlanConversionTest,
aggregateMaskMustBeTopLevelField) {
"Aggregation Operator only supports a top-level field mask.");
}
+#ifdef ENABLE_KAFKA
+TEST_F(Substrait2VeloxPlanConversionTest, serializedFileSplit) {
+ ::substrait::Plan plan;
+ auto* read = plan.add_relations()->mutable_rel()->mutable_read();
+ read->mutable_base_schema()->add_names("offset");
+ read->mutable_base_schema()->mutable_struct_()->add_types()->mutable_i64();
+
+ ::substrait::ReadRel_LocalFiles files;
+ auto* file = files.add_items();
+ file->set_uri_file("file:///tmp/kafka-join.parquet");
+ file->set_start(10);
+ file->set_length(20);
+ file->mutable_parquet();
+
+ auto splitInfo = parseSerializedSplitInfo(files.SerializeAsString());
+ auto node = planConverter_->toVeloxPlanWithSplitInfos(plan, {splitInfo});
+ const auto& split = planConverter_->splitInfos().at(node->id());
+ EXPECT_EQ(split, splitInfo);
+ EXPECT_EQ(split->paths,
std::vector<std::string>{"file:///tmp/kafka-join.parquet"});
+ EXPECT_EQ(split->starts, std::vector<uint64_t>{10});
+ EXPECT_EQ(split->lengths, std::vector<uint64_t>{20});
+}
+
+TEST_F(Substrait2VeloxPlanConversionTest, kafkaReadValidation) {
+ ::substrait::ReadRel read;
+ read.mutable_extension_table()->mutable_detail()->PackFrom(StreamKafka{},
"");
+ read.mutable_base_schema()->add_names("offset");
+ read.mutable_base_schema()->mutable_struct_()->add_types()->mutable_i64();
+ SubstraitToVeloxPlanConverter converter(
+ pool(), veloxCfg_.get(), {}, VeloxConnectorIds{.kafka = "kafka-test"},
std::nullopt, std::nullopt, true);
+ auto node = converter.toVeloxPlan(read);
+ auto scan = std::dynamic_pointer_cast<const core::TableScanNode>(node);
+ ASSERT_NE(scan, nullptr);
+ EXPECT_EQ(scan->tableHandle()->connectorId(), "kafka-test");
+ EXPECT_EQ(converter.splitInfos().at(node->id())->leafType,
SplitInfo::LeafType::TABLE_SCAN);
+}
+
+TEST_F(Substrait2VeloxPlanConversionTest, mixedKafkaAndFileSplits) {
+ ::substrait::ReadRel kafkaRead;
+
kafkaRead.mutable_extension_table()->mutable_detail()->PackFrom(StreamKafka{},
"");
+ kafkaRead.mutable_base_schema()->add_names("offset");
+
kafkaRead.mutable_base_schema()->mutable_struct_()->add_types()->mutable_i64();
+ auto fileRead = kafkaRead;
+ fileRead.clear_extension_table();
+
+ StreamKafka kafka;
+ kafka.mutable_topic_partition()->set_topic("test-topic");
+ kafka.mutable_topic_partition()->set_partition(2);
+ kafka.set_start_offset(5);
+ kafka.set_end_offset(9);
+ kafka.set_poll_timeout_ms(1000);
+ kafka.set_fail_on_data_loss(true);
+ (*kafka.mutable_params())["bootstrap.servers"] = "localhost:9092";
+ ::substrait::ReadRel_ExtensionTable kafkaSplit;
+ kafkaSplit.mutable_detail()->PackFrom(kafka, "");
+
+ ::substrait::ReadRel_LocalFiles files;
+ files.add_items()->set_uri_file("file:///tmp/kafka-join.parquet");
+ files.mutable_items(0)->mutable_parquet();
+
+ for (const bool kafkaFirst : {true, false}) {
+ SubstraitToVeloxPlanConverter converter(
+ pool(),
+ veloxCfg_.get(),
+ {},
+ VeloxConnectorIds{.hive =
facebook::velox::exec::test::kHiveConnectorId, .kafka = "kafka-test"});
+ auto parsedKafkaSplit =
parseSerializedSplitInfo(kafkaSplit.SerializeAsString());
+ auto parsedFileSplit = parseSerializedSplitInfo(files.SerializeAsString());
+ converter.setSplitInfos(
+ kafkaFirst ? std::vector<std::shared_ptr<SplitInfo>>{parsedKafkaSplit,
parsedFileSplit}
+ : std::vector<std::shared_ptr<SplitInfo>>{parsedFileSplit,
parsedKafkaSplit});
+ auto first = converter.toVeloxPlan(kafkaFirst ? kafkaRead : fileRead);
+ auto second = converter.toVeloxPlan(kafkaFirst ? fileRead : kafkaRead);
+ const auto& kafkaNode = kafkaFirst ? first : second;
+ const auto& fileNode = kafkaFirst ? second : first;
+ EXPECT_EQ(converter.splitInfos().at(kafkaNode->id()), parsedKafkaSplit);
+ EXPECT_EQ(converter.splitInfos().at(fileNode->id()), parsedFileSplit);
+ auto splitInfo =
std::dynamic_pointer_cast<KafkaSplitInfo>(converter.splitInfos().at(kafkaNode->id()));
+ ASSERT_NE(splitInfo, nullptr);
+ auto split = splitInfo->toConnectorSplit("kafka-test");
+ EXPECT_EQ(split->topic, "test-topic");
+ EXPECT_EQ(split->partition, 2);
+ EXPECT_EQ(split->startOffset, 5);
+ EXPECT_EQ(split->endOffset, 9);
+ EXPECT_EQ(split->pollTimeoutMs, 1000);
+ EXPECT_TRUE(split->failOnDataLoss);
+ EXPECT_EQ(split->kafkaParams.at("bootstrap.servers"), "localhost:9092");
+ EXPECT_EQ(
+ converter.splitInfos().at(fileNode->id())->paths,
std::vector<std::string>{"file:///tmp/kafka-join.parquet"});
+ }
+}
+
+TEST_F(Substrait2VeloxPlanConversionTest, kafkaSplitValidation) {
+ ::substrait::ReadRel read;
+ read.mutable_extension_table()->mutable_detail()->PackFrom(StreamKafka{});
+ read.mutable_base_schema()->add_names("offset");
+ read.mutable_base_schema()->mutable_struct_()->add_types()->mutable_i64();
+
+ auto fileRead = read;
+ fileRead.clear_extension_table();
+
+ const auto makeConverter = [&](std::vector<std::shared_ptr<SplitInfo>>
splits) {
+ auto converter = std::make_unique<SubstraitToVeloxPlanConverter>(
+ pool(),
+ veloxCfg_.get(),
+ std::vector<std::shared_ptr<ResultIterator>>{},
+ VeloxConnectorIds{.hive =
facebook::velox::exec::test::kHiveConnectorId, .kafka = "kafka-test"});
+ converter->setSplitInfos(std::move(splits));
+ return converter;
+ };
+
+ VELOX_ASSERT_THROW(makeConverter({})->toVeloxPlan(read), "Plan must have
readRel and related split info.");
+ VELOX_ASSERT_THROW(parseSerializedSplitInfo(std::string(1, '\xff')), "Parse
split info failed");
+
+ ::substrait::ReadRel_ExtensionTable split;
+ split.mutable_detail()->PackFrom(StreamKafka{});
+ split.mutable_detail()->set_value(std::string(1, '\xff'));
+ VELOX_ASSERT_THROW(
+ parseSerializedSplitInfo(split.SerializeAsString()),
+ "Kafka split info must contain a gluten.StreamKafka payload");
+
+ ::substrait::ReadRel_LocalFiles files;
+ files.add_items()->set_uri_file("file:///tmp/kafka-join.parquet");
+ files.mutable_items(0)->mutable_parquet();
+ VELOX_ASSERT_THROW(
+
makeConverter({parseSerializedSplitInfo(files.SerializeAsString())})->toVeloxPlan(read),
+ "Kafka read requires Kafka split info");
+ split.mutable_detail()->PackFrom(StreamKafka{});
+ VELOX_ASSERT_THROW(
+
makeConverter({parseSerializedSplitInfo(split.SerializeAsString())})->toVeloxPlan(fileRead),
+ "File read requires local file split info");
+}
+#endif
+
} // namespace gluten
diff --git a/cpp/velox/tests/Substrait2VeloxPlanValidatorTest.cc
b/cpp/velox/tests/Substrait2VeloxPlanValidatorTest.cc
index c8c41e4e16..c48bcb1a33 100644
--- a/cpp/velox/tests/Substrait2VeloxPlanValidatorTest.cc
+++ b/cpp/velox/tests/Substrait2VeloxPlanValidatorTest.cc
@@ -17,6 +17,9 @@
#include "FilePathGenerator.h"
#include "JsonToProtoConverter.h"
+#ifdef ENABLE_KAFKA
+#include "kafka.pb.h"
+#endif
#include "memory/VeloxMemoryManager.h"
#include "substrait/SubstraitToVeloxPlan.h"
@@ -83,6 +86,17 @@ TEST_F(Substrait2VeloxPlanValidatorTest, group) {
ASSERT_FALSE(validatePlan(substraitPlan));
}
+#ifdef ENABLE_KAFKA
+TEST_F(Substrait2VeloxPlanValidatorTest, kafkaReadRequiresNativeSupport) {
+ ::substrait::Plan plan;
+ auto* read = plan.add_relations()->mutable_rel()->mutable_read();
+ read->mutable_extension_table()->mutable_detail()->PackFrom(StreamKafka{},
"");
+ read->mutable_base_schema()->add_names("offset");
+ read->mutable_base_schema()->mutable_struct_()->add_types()->mutable_i64();
+ EXPECT_TRUE(validatePlan(plan));
+}
+#endif
+
TEST_F(Substrait2VeloxPlanValidatorTest, expandSelectionMustBeTopLevelField) {
const auto makePlan = [](bool nestedSelection) {
::substrait::Plan plan;
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
