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 08417cfb18 [VL] Add native Kafka connector (#13183)
08417cfb18 is described below
commit 08417cfb18459a6f3f40e1427745a2a78f5d6b06
Author: Reema <[email protected]>
AuthorDate: Tue Oct 20 11:28:21 2026 +0300
[VL] Add native Kafka connector (#13183)
---
.github/workflows/velox_kafka.yml | 9 +-
cpp/velox/CMakeLists.txt | 18 +-
cpp/velox/compute/VeloxBackend.cc | 11 +
cpp/velox/compute/VeloxBackend.h | 4 +
cpp/velox/compute/VeloxConnectorIds.h | 2 +
cpp/velox/compute/VeloxRuntime.cc | 16 ++
cpp/velox/config/VeloxConfig.h | 1 +
cpp/velox/operators/reader/KafkaConnector.cc | 58 ++++
cpp/velox/operators/reader/KafkaConnector.h | 87 ++++++
cpp/velox/operators/reader/KafkaDataSource.cc | 377 ++++++++++++++++++++++++++
cpp/velox/operators/reader/KafkaDataSource.h | 96 +++++++
cpp/velox/operators/reader/KafkaSplit.h | 72 +++++
cpp/velox/tests/CMakeLists.txt | 4 +
cpp/velox/tests/KafkaConnectorTest.cc | 367 +++++++++++++++++++++++++
14 files changed, 1113 insertions(+), 9 deletions(-)
diff --git a/.github/workflows/velox_kafka.yml
b/.github/workflows/velox_kafka.yml
index 33a5303dbe..3f2a367e13 100644
--- a/.github/workflows/velox_kafka.yml
+++ b/.github/workflows/velox_kafka.yml
@@ -80,10 +80,9 @@ jobs:
./dev/builddeps-veloxbe.sh --enable_vcpkg=ON --enable_kafka=ON \
--build_arrow=OFF --build_tests=ON --build_benchmarks=OFF
--build_examples=OFF \
--enable_s3=OFF --enable_gcs=OFF --enable_hdfs=OFF
--enable_abfs=OFF
- # TODO: Re-enable after https://github.com/apache/gluten/pull/13183 is
merged.
- # - 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.
+ - 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 \
@@ -147,7 +146,7 @@ jobs:
# ./build/mvn -B -ntp test -pl gluten-ut/spark35 -Pbackends-velox
-Pspark-3.5 \
# -Pscala-2.12 -Pjava-17 -Pspark-ut -Pfast-build \
#
-Dsuites=org.apache.gluten.execution.VeloxKafkaScanValidationSuite
- # TODO: Re-enable when VeloxGlutenKafkaScanSuite is added.
+ # TODO: Re-enable when VeloxGlutenKafkaScanSuite is added in #13237.
# - name: Check native micro-batch reads and headers fallback
# run: |
# ./build/mvn -B -ntp test -pl backends-velox -Pbackends-velox
-Pspark-3.5 \
diff --git a/cpp/velox/CMakeLists.txt b/cpp/velox/CMakeLists.txt
index 86f23171f5..fd5ff0a9d4 100644
--- a/cpp/velox/CMakeLists.txt
+++ b/cpp/velox/CMakeLists.txt
@@ -219,6 +219,11 @@ if(ENABLE_S3)
find_package(ZLIB)
endif()
+if(ENABLE_KAFKA)
+ list(APPEND VELOX_SRCS operators/reader/KafkaDataSource.cc
+ operators/reader/KafkaConnector.cc)
+endif()
+
if(ENABLE_GPU)
list(
APPEND
@@ -508,11 +513,16 @@ set(CMAKE_FIND_LIBRARY_SUFFIXES
${CMAKE_FIND_LIBRARY_SUFFIXES_BCK})
# Find and link librdkafka for Kafka streaming support
if(ENABLE_KAFKA)
- target_compile_definitions(velox PUBLIC ENABLE_KAFKA)
+ target_compile_definitions(gluten_velox_impl PUBLIC ENABLE_KAFKA)
+ # Native tests link the implementation target directly and use both Kafka
+ # APIs.
find_package(RdKafka CONFIG QUIET)
- if(RdKafka_FOUND AND TARGET RdKafka::rdkafka++)
+ if(RdKafka_FOUND
+ AND TARGET RdKafka::rdkafka++
+ AND TARGET RdKafka::rdkafka)
message(STATUS "Found RdKafka: ${RdKafka_DIR}")
- target_link_libraries(velox PUBLIC RdKafka::rdkafka++)
+ target_link_libraries(gluten_velox_impl PUBLIC RdKafka::rdkafka++
+ RdKafka::rdkafka)
else()
# Link the C++ client and C mock API.
find_package(PkgConfig QUIET)
@@ -520,7 +530,7 @@ if(ENABLE_KAFKA)
pkg_check_modules(RDKAFKA QUIET IMPORTED_TARGET rdkafka++ rdkafka)
if(RDKAFKA_FOUND)
message(STATUS "Found librdkafka via pkg-config: ${RDKAFKA_LIBRARIES}")
- target_link_libraries(velox PUBLIC PkgConfig::RDKAFKA)
+ target_link_libraries(gluten_velox_impl PUBLIC PkgConfig::RDKAFKA)
else()
message(
FATAL_ERROR
diff --git a/cpp/velox/compute/VeloxBackend.cc
b/cpp/velox/compute/VeloxBackend.cc
index c131c008f4..7633216249 100644
--- a/cpp/velox/compute/VeloxBackend.cc
+++ b/cpp/velox/compute/VeloxBackend.cc
@@ -27,6 +27,10 @@
#include "operators/plannodes/RowVectorStream.h"
#include "utils/ConfigExtractor.h"
+#ifdef ENABLE_KAFKA
+#include "operators/reader/KafkaConnector.h"
+#endif
+
#ifdef GLUTEN_ENABLE_QAT
#include "utils/qat/QatCodec.h"
#endif
@@ -426,6 +430,13 @@ std::shared_ptr<facebook::velox::connector::Connector>
VeloxBackend::createValue
return std::make_shared<ValueStreamConnector>(connectorId,
hiveConnectorConfig_, dynamicFilterEnabled);
}
+#ifdef ENABLE_KAFKA
+std::shared_ptr<facebook::velox::connector::Connector>
VeloxBackend::createKafkaConnector(
+ const std::string& connectorId) const {
+ return std::make_shared<KafkaConnector>(connectorId, hiveConnectorConfig_);
+}
+#endif
+
#ifdef GLUTEN_ENABLE_GPU
std::shared_ptr<facebook::velox::connector::Connector>
VeloxBackend::createCudfHiveConnector(
const std::string& connectorId,
diff --git a/cpp/velox/compute/VeloxBackend.h b/cpp/velox/compute/VeloxBackend.h
index 5e129b9059..a7592d718b 100644
--- a/cpp/velox/compute/VeloxBackend.h
+++ b/cpp/velox/compute/VeloxBackend.h
@@ -96,6 +96,10 @@ class VeloxBackend {
const std::string& connectorId,
bool dynamicFilterEnabled) const;
+#ifdef ENABLE_KAFKA
+ std::shared_ptr<facebook::velox::connector::Connector>
createKafkaConnector(const std::string& connectorId) const;
+#endif
+
#ifdef GLUTEN_ENABLE_GPU
std::shared_ptr<facebook::velox::connector::Connector>
createCudfHiveConnector(
const std::string& connectorId,
diff --git a/cpp/velox/compute/VeloxConnectorIds.h
b/cpp/velox/compute/VeloxConnectorIds.h
index 73b526229f..751d890fec 100644
--- a/cpp/velox/compute/VeloxConnectorIds.h
+++ b/cpp/velox/compute/VeloxConnectorIds.h
@@ -26,11 +26,13 @@ struct VeloxConnectorIds {
std::string iceberg;
std::string delta;
std::string iterator;
+ std::string kafka;
std::string cudfHive;
bool hiveRegistered{false};
bool icebergRegistered{false};
bool deltaRegistered{false};
bool iteratorRegistered{false};
+ bool kafkaRegistered{false};
bool cudfHiveRegistered{false};
};
diff --git a/cpp/velox/compute/VeloxRuntime.cc
b/cpp/velox/compute/VeloxRuntime.cc
index 5c4621a209..4420976572 100644
--- a/cpp/velox/compute/VeloxRuntime.cc
+++ b/cpp/velox/compute/VeloxRuntime.cc
@@ -230,6 +230,7 @@ VeloxConnectorIds makeScopedConnectorIds(uint64_t
runtimeId) {
.iceberg = makeScopedConnectorId(kIcebergConnectorId, runtimeId),
.delta =
makeScopedConnectorId(delta::DeltaConnectorFactory::kDeltaConnectorName,
runtimeId),
.iterator = makeScopedConnectorId(kIteratorConnectorId, runtimeId),
+ .kafka = makeScopedConnectorId(kKafkaConnectorId, runtimeId),
.cudfHive = makeScopedConnectorId(kCudfHiveConnectorId, runtimeId)};
}
@@ -316,6 +317,15 @@ void VeloxRuntime::registerConnectors() {
velox::connector::hasConnector(connectorIds_.iterator),
"Scoped iterator connector not found after registration: " +
connectorIds_.iterator);
+#ifdef ENABLE_KAFKA
+ connectorIds_.kafkaRegistered =
+
velox::connector::registerConnector(backend->createKafkaConnector(connectorIds_.kafka));
+ GLUTEN_CHECK(connectorIds_.kafkaRegistered, "Failed to register scoped Kafka
connector: " + connectorIds_.kafka);
+ GLUTEN_CHECK(
+ velox::connector::hasConnector(connectorIds_.kafka),
+ "Scoped Kafka connector not found after registration: " +
connectorIds_.kafka);
+#endif
+
#ifdef GLUTEN_ENABLE_GPU
if (veloxCfg_->get<bool>(kCudfEnableTableScan, kCudfEnableTableScanDefault)
&&
veloxCfg_->get<bool>(kCudfEnabled, kCudfEnabledDefault)) {
@@ -336,6 +346,12 @@ void VeloxRuntime::unregisterConnectors() {
velox::connector::unregisterConnector(connectorIds_.cudfHive);
connectorIds_.cudfHiveRegistered = false;
}
+#endif
+#ifdef ENABLE_KAFKA
+ if (connectorIds_.kafkaRegistered) {
+ velox::connector::unregisterConnector(connectorIds_.kafka);
+ connectorIds_.kafkaRegistered = false;
+ }
#endif
if (connectorIds_.iteratorRegistered) {
velox::connector::unregisterConnector(connectorIds_.iterator);
diff --git a/cpp/velox/config/VeloxConfig.h b/cpp/velox/config/VeloxConfig.h
index 91069e86a2..8d9a1ed2e0 100644
--- a/cpp/velox/config/VeloxConfig.h
+++ b/cpp/velox/config/VeloxConfig.h
@@ -134,6 +134,7 @@ const uint64_t kVeloxMemReclaimMaxWaitMsDefault = 3600000;
// 60min
const std::string kHiveConnectorId = "test-hive";
const std::string kIcebergConnectorId = "test-iceberg";
+const std::string kKafkaConnectorId = "kafka-stream";
const std::string kVeloxCacheEnabled =
"spark.gluten.sql.columnar.backend.velox.cacheEnabled";
diff --git a/cpp/velox/operators/reader/KafkaConnector.cc
b/cpp/velox/operators/reader/KafkaConnector.cc
new file mode 100644
index 0000000000..03469b1018
--- /dev/null
+++ b/cpp/velox/operators/reader/KafkaConnector.cc
@@ -0,0 +1,58 @@
+/*
+ * 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 "KafkaConnector.h"
+#include "KafkaDataSource.h"
+
+#include <folly/String.h>
+
+namespace gluten {
+
+KafkaColumnHandle::KafkaColumnHandle(std::string name,
facebook::velox::TypePtr type)
+ : name_(std::move(name)), type_(std::move(type)) {
+ using namespace facebook::velox;
+ static const std::unordered_map<std::string, std::pair<KafkaField, TypePtr>>
kFields = {
+ {"key", {KafkaField::kKey, VARBINARY()}},
+ {"value", {KafkaField::kValue, VARBINARY()}},
+ {"topic", {KafkaField::kTopic, VARCHAR()}},
+ {"partition", {KafkaField::kPartition, INTEGER()}},
+ {"offset", {KafkaField::kOffset, BIGINT()}},
+ {"timestamp", {KafkaField::kTimestamp, TIMESTAMP()}},
+ {"timestamptype", {KafkaField::kTimestampType, INTEGER()}},
+ };
+ auto lowerName = name_;
+ folly::toLowerAscii(lowerName);
+ auto it = kFields.find(lowerName);
+ VELOX_USER_CHECK(it != kFields.end(), "Unsupported Kafka column: {}", name_);
+ VELOX_USER_CHECK(
+ type_->equivalent(*it->second.second),
+ "Kafka column {} must be of type {}, got {}",
+ name_,
+ it->second.second->toString(),
+ type_->toString());
+ field_ = it->second.first;
+}
+
+std::unique_ptr<facebook::velox::connector::DataSource>
KafkaConnector::createDataSource(
+ const facebook::velox::RowTypePtr& outputType,
+ const facebook::velox::connector::ConnectorTableHandlePtr& /*tableHandle*/,
+ const facebook::velox::connector::ColumnHandleMap& columnHandles,
+ facebook::velox::connector::ConnectorQueryCtx* connectorQueryCtx) {
+ return std::make_unique<KafkaDataSource>(outputType, columnHandles,
connectorQueryCtx->memoryPool());
+}
+
+} // namespace gluten
diff --git a/cpp/velox/operators/reader/KafkaConnector.h
b/cpp/velox/operators/reader/KafkaConnector.h
new file mode 100644
index 0000000000..34e75e415c
--- /dev/null
+++ b/cpp/velox/operators/reader/KafkaConnector.h
@@ -0,0 +1,87 @@
+/*
+ * 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 "velox/connectors/Connector.h"
+
+namespace gluten {
+
+class KafkaTableHandle : public
facebook::velox::connector::ConnectorTableHandle {
+ public:
+ explicit KafkaTableHandle(std::string connectorId) :
ConnectorTableHandle(std::move(connectorId)) {}
+
+ const std::string& name() const override {
+ static const std::string kName = "KafkaTableHandle";
+ return kName;
+ }
+
+ folly::dynamic serialize() const override {
+ VELOX_NYI();
+ }
+};
+
+/// A column of Spark's Kafka source schema.
+enum class KafkaField { kKey, kValue, kTopic, kPartition, kOffset, kTimestamp,
kTimestampType };
+
+/// Column handle naming one of the Kafka source columns: key, value, topic,
partition, offset,
+/// timestamp or timestampType. Throws a user error for other columns, e.g.
headers.
+class KafkaColumnHandle : public facebook::velox::connector::ColumnHandle {
+ public:
+ KafkaColumnHandle(std::string name, facebook::velox::TypePtr type);
+
+ const std::string& name() const override {
+ return name_;
+ }
+
+ const facebook::velox::TypePtr& type() const {
+ return type_;
+ }
+
+ KafkaField field() const {
+ return field_;
+ }
+
+ private:
+ const std::string name_;
+ const facebook::velox::TypePtr type_;
+ KafkaField field_;
+};
+
+/// Connector reading Kafka topic partitions through librdkafka. Each
KafkaConnectorSplit is read
+/// by a KafkaDataSource.
+class KafkaConnector : public facebook::velox::connector::Connector {
+ public:
+ KafkaConnector(const std::string& id, std::shared_ptr<const
facebook::velox::config::ConfigBase> config)
+ : Connector(id, std::move(config)) {}
+
+ std::unique_ptr<facebook::velox::connector::DataSource> createDataSource(
+ const facebook::velox::RowTypePtr& outputType,
+ const facebook::velox::connector::ConnectorTableHandlePtr& tableHandle,
+ const facebook::velox::connector::ColumnHandleMap& columnHandles,
+ facebook::velox::connector::ConnectorQueryCtx* connectorQueryCtx)
override;
+
+ std::unique_ptr<facebook::velox::connector::DataSink> createDataSink(
+ facebook::velox::RowTypePtr inputType,
+ facebook::velox::connector::ConnectorInsertTableHandlePtr
connectorInsertTableHandle,
+ facebook::velox::connector::ConnectorQueryCtx* connectorQueryCtx,
+ facebook::velox::connector::CommitStrategy commitStrategy) override {
+ VELOX_UNSUPPORTED("KafkaConnector does not support data sinks");
+ }
+};
+
+} // namespace gluten
diff --git a/cpp/velox/operators/reader/KafkaDataSource.cc
b/cpp/velox/operators/reader/KafkaDataSource.cc
new file mode 100644
index 0000000000..20c1ae575d
--- /dev/null
+++ b/cpp/velox/operators/reader/KafkaDataSource.cc
@@ -0,0 +1,377 @@
+/*
+ * 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 "KafkaDataSource.h"
+
+#include <algorithm>
+#include <chrono>
+#include <limits>
+
+#include <glog/logging.h>
+#include <librdkafka/rdkafka.h>
+#include <librdkafka/rdkafkacpp.h>
+
+#include "velox/vector/FlatVector.h"
+
+using namespace facebook::velox;
+
+namespace gluten {
+
+namespace {
+
+// Spark's default for kafkaConsumer.pollTimeoutMs.
+constexpr int64_t kDefaultPollTimeoutMs = 120'000;
+constexpr int kConsumeTimeoutMs = 100;
+// Spark requires a group id on executors; librdkafka's KafkaConsumer does too.
+const std::string kDefaultGroupId = "gluten-kafka-executor";
+
+// Java client properties that have no librdkafka counterpart, or that the
reader must control.
+bool isIgnoredProperty(const std::string& key) {
+ return key == "key.deserializer" || key == "value.deserializer" || key ==
"auto.offset.reset" ||
+ key == "enable.auto.commit";
+}
+
+void setProperty(RdKafka::Conf& conf, const std::string& key, const
std::string& value) {
+ std::string errstr;
+ VELOX_CHECK(
+ conf.set(key, value, errstr) == RdKafka::Conf::CONF_OK, "Failed to set
Kafka property {}: {}", key, errstr);
+}
+
+int64_t getLogEndOffset(RdKafka::KafkaConsumer& consumer, const std::string&
topic, int32_t partition, int timeoutMs) {
+ auto* client = consumer.c_ptr();
+ std::unique_ptr<rd_kafka_AdminOptions_t,
decltype(&rd_kafka_AdminOptions_destroy)> options(
+ rd_kafka_AdminOptions_new(client, RD_KAFKA_ADMIN_OP_LISTOFFSETS),
rd_kafka_AdminOptions_destroy);
+ char errstr[512];
+ const auto err = rd_kafka_AdminOptions_set_request_timeout(options.get(),
timeoutMs, errstr, sizeof(errstr));
+ VELOX_CHECK(err == RD_KAFKA_RESP_ERR_NO_ERROR, "Failed to set Kafka offset
request timeout: {}", errstr);
+ std::unique_ptr<rd_kafka_error_t, decltype(&rd_kafka_error_destroy)>
isolationError(
+ rd_kafka_AdminOptions_set_isolation_level(options.get(),
RD_KAFKA_ISOLATION_LEVEL_READ_UNCOMMITTED),
+ rd_kafka_error_destroy);
+ if (isolationError) {
+ VELOX_FAIL("Failed to set Kafka offset request isolation: {}",
rd_kafka_error_string(isolationError.get()));
+ }
+ std::unique_ptr<rd_kafka_topic_partition_list_t,
decltype(&rd_kafka_topic_partition_list_destroy)> partitions(
+ rd_kafka_topic_partition_list_new(1),
rd_kafka_topic_partition_list_destroy);
+ rd_kafka_topic_partition_list_add(partitions.get(), topic.c_str(),
partition)->offset = RD_KAFKA_OFFSET_SPEC_LATEST;
+ std::unique_ptr<rd_kafka_queue_t, decltype(&rd_kafka_queue_destroy)> queue(
+ rd_kafka_queue_new(client), rd_kafka_queue_destroy);
+ rd_kafka_ListOffsets(client, partitions.get(), options.get(), queue.get());
+ std::unique_ptr<rd_kafka_event_t, decltype(&rd_kafka_event_destroy)> event(
+ rd_kafka_queue_poll(queue.get(), timeoutMs), rd_kafka_event_destroy);
+ VELOX_CHECK_NOT_NULL(event, "Timed out getting Kafka log end offset for
{}-{}", topic, partition);
+ VELOX_CHECK(
+ rd_kafka_event_error(event.get()) == RD_KAFKA_RESP_ERR_NO_ERROR,
+ "Failed to get Kafka log end offset for {}-{}: {}",
+ topic,
+ partition,
+ rd_kafka_event_error_string(event.get()));
+ const auto* result = rd_kafka_event_ListOffsets_result(event.get());
+ VELOX_CHECK_NOT_NULL(result, "Missing Kafka log end offset result for
{}-{}", topic, partition);
+ size_t count = 0;
+ const auto** infos = rd_kafka_ListOffsets_result_infos(result, &count);
+ VELOX_CHECK_EQ(count, 1, "Expected one Kafka log end offset result");
+ const auto* info = rd_kafka_ListOffsetsResultInfo_topic_partition(infos[0]);
+ VELOX_CHECK(
+ info->err == RD_KAFKA_RESP_ERR_NO_ERROR,
+ "Failed to get Kafka log end offset for {}-{}: {}",
+ topic,
+ partition,
+ rd_kafka_err2str(info->err));
+ VELOX_CHECK_GE(info->offset, 0, "Invalid Kafka log end offset for {}-{}",
topic, partition);
+ return info->offset;
+}
+
+int32_t toSparkTimestampType(RdKafka::MessageTimestamp::MessageTimestampType
type) {
+ // Matches org.apache.kafka.common.record.TimestampType ids.
+ switch (type) {
+ case RdKafka::MessageTimestamp::MSG_TIMESTAMP_CREATE_TIME:
+ return 0;
+ case RdKafka::MessageTimestamp::MSG_TIMESTAMP_LOG_APPEND_TIME:
+ return 1;
+ default:
+ return -1;
+ }
+}
+
+} // namespace
+
+KafkaDataSource::KafkaDataSource(
+ const RowTypePtr& outputType,
+ const connector::ColumnHandleMap& columnHandles,
+ memory::MemoryPool* pool)
+ : outputType_(outputType), pool_(pool) {
+ fields_.reserve(outputType_->size());
+ for (const auto& outputName : outputType_->names()) {
+ auto it = columnHandles.find(outputName);
+ VELOX_CHECK(it != columnHandles.end(), "Missing column handle for Kafka
output column {}", outputName);
+ auto handle = std::dynamic_pointer_cast<const
KafkaColumnHandle>(it->second);
+ VELOX_CHECK_NOT_NULL(handle, "Wrong type of column handle for Kafka output
column {}", outputName);
+ fields_.push_back(handle->field());
+ }
+}
+
+KafkaDataSource::~KafkaDataSource() {
+ closeConsumer();
+}
+
+void KafkaDataSource::addSplit(std::shared_ptr<connector::ConnectorSplit>
split) {
+ VELOX_CHECK_NULL(split_, "Previous Kafka split has not been fully
processed");
+ VELOX_CHECK_NOT_NULL(split, "Kafka split must not be null");
+ split_ = std::dynamic_pointer_cast<const KafkaConnectorSplit>(split);
+ VELOX_CHECK_NOT_NULL(split_, "Wrong type of split for KafkaDataSource: {}",
split->toString());
+ nextOffset_ = split_->startOffset;
+ splitDone_ = nextOffset_ >= split_->endOffset;
+ if (splitDone_) {
+ return;
+ }
+
+ std::unique_ptr<RdKafka::Conf>
conf(RdKafka::Conf::create(RdKafka::Conf::CONF_GLOBAL));
+ std::string errstr;
+ for (const auto& [key, value] : split_->kafkaParams) {
+ if (isIgnoredProperty(key)) {
+ continue;
+ }
+ // The params are Java client properties; skip the ones librdkafka does
not know.
+ const auto result = conf->set(key, value, errstr);
+ if (result == RdKafka::Conf::CONF_UNKNOWN) {
+ VLOG(1) << "Ignoring unsupported Java Kafka property " << key;
+ } else {
+ VELOX_USER_CHECK(result == RdKafka::Conf::CONF_OK, "Invalid Kafka
property {}: {}", key, errstr);
+ }
+ }
+ if (split_->kafkaParams.find("group.id") == split_->kafkaParams.end()) {
+ setProperty(*conf, "group.id", kDefaultGroupId);
+ }
+ // Match the Java client default; Spark computes end offsets from the high
watermark.
+ if (split_->kafkaParams.find("isolation.level") ==
split_->kafkaParams.end()) {
+ setProperty(*conf, "isolation.level", "read_uncommitted");
+ }
+ // Offsets are tracked by Spark, never commit them to the consumer group.
+ setProperty(*conf, "enable.auto.commit", "false");
+ setProperty(*conf, "enable.auto.offset.store", "false");
+ // Report aged-out offsets as errors instead of silently resetting.
+ setProperty(*conf, "auto.offset.reset", "error");
+ // Detects reaching the log end when trailing offsets (e.g. transaction
markers) carry no record.
+ setProperty(*conf, "enable.partition.eof", "true");
+
+ consumer_.reset(RdKafka::KafkaConsumer::create(conf.get(), errstr));
+ VELOX_CHECK_NOT_NULL(consumer_, "Failed to create Kafka consumer: {}",
errstr);
+
+ std::unique_ptr<RdKafka::TopicPartition> topicPartition(
+ RdKafka::TopicPartition::create(split_->topic, split_->partition,
split_->startOffset));
+ std::vector<RdKafka::TopicPartition*> partitions{topicPartition.get()};
+ const auto err = consumer_->assign(partitions);
+ VELOX_CHECK(err == RdKafka::ERR_NO_ERROR, "Failed to assign {}: {}",
split_->toString(), RdKafka::err2str(err));
+}
+
+std::optional<RowVectorPtr> KafkaDataSource::next(uint64_t size,
ContinueFuture& /*future*/) {
+ if (split_ == nullptr) {
+ return nullptr;
+ }
+ if (!splitDone_) {
+ auto messages = poll(std::max<uint64_t>(size, 1));
+ if (!messages.empty()) {
+ return toRowVector(messages);
+ }
+ }
+ VELOX_CHECK(splitDone_);
+ closeConsumer();
+ split_.reset();
+ return nullptr;
+}
+
+std::vector<std::unique_ptr<RdKafka::Message>> KafkaDataSource::poll(uint64_t
maxRows) {
+ std::vector<std::unique_ptr<RdKafka::Message>> messages;
+ const auto pollTimeoutMs = split_->pollTimeoutMs > 0 ? split_->pollTimeoutMs
: kDefaultPollTimeoutMs;
+ const auto deadline = std::chrono::steady_clock::now() +
std::chrono::milliseconds(pollTimeoutMs);
+
+ while (!splitDone_ && messages.size() < maxRows) {
+ // Only wait for more records while the batch is empty.
+ std::unique_ptr<RdKafka::Message> msg(consumer_->consume(messages.empty()
? kConsumeTimeoutMs : 0));
+ switch (msg->err()) {
+ case RdKafka::ERR_NO_ERROR: {
+ const auto offset = msg->offset();
+ if (offset < nextOffset_) {
+ break;
+ }
+ if (offset >= split_->endOffset) {
+ splitDone_ = true;
+ break;
+ }
+ completedBytes_ += msg->len() + msg->key_len();
+ messages.push_back(std::move(msg));
+ nextOffset_ = offset + 1;
+ splitDone_ = nextOffset_ >= split_->endOffset;
+ break;
+ }
+ case RdKafka::ERR__PARTITION_EOF: {
+ // Offsets below EOF without a record are compacted records or
transaction markers.
+ splitDone_ = msg->offset() >= split_->endOffset;
+ if (!splitDone_) {
+ // A read_committed EOF can be the last stable offset, or a follower
can lag the leader.
+ // Query the leader's physical log end before treating the missing
tail as data loss.
+ const auto timeoutMs =
static_cast<int>(std::min<int64_t>(pollTimeoutMs,
std::numeric_limits<int>::max()));
+ const auto logEndOffset = getLogEndOffset(*consumer_, split_->topic,
split_->partition, timeoutMs);
+ if (logEndOffset < split_->endOffset) {
+ VELOX_USER_CHECK(
+ !split_->failOnDataLoss,
+ "Kafka offsets [{}, {}) of {}-{} are no longer available. Some
data may have been lost. "
+ "Set the source option failOnDataLoss to false to ignore
this.",
+ logEndOffset,
+ split_->endOffset,
+ split_->topic,
+ split_->partition);
+ LOG(WARNING) << "Kafka offsets [" << logEndOffset << ", " <<
split_->endOffset << ") of " << split_->topic
+ << "-" << split_->partition << " are no longer
available";
+ splitDone_ = true;
+ }
+ }
+ break;
+ }
+ case RdKafka::ERR__TIMED_OUT:
+ if (!messages.empty()) {
+ return messages;
+ }
+ break;
+ case RdKafka::ERR__AUTO_OFFSET_RESET:
+ case RdKafka::ERR_OFFSET_OUT_OF_RANGE: {
+ VELOX_USER_CHECK(
+ !split_->failOnDataLoss,
+ "Kafka offset {} of {}-{} is no longer available: {}. Some data
may have been lost. "
+ "Set the source option failOnDataLoss to false to ignore this.",
+ nextOffset_,
+ split_->topic,
+ split_->partition,
+ msg->errstr());
+ LOG(WARNING) << "Kafka offset " << nextOffset_ << " of " <<
split_->topic << "-" << split_->partition
+ << " is no longer available: " << msg->errstr();
+ // Retention may have removed only the beginning of the split.
Preserve the remaining
+ // available records when Spark permits data loss instead of dropping
the entire split.
+ int64_t lowOffset;
+ int64_t highOffset;
+ const auto timeoutMs =
static_cast<int>(std::min<int64_t>(pollTimeoutMs,
std::numeric_limits<int>::max()));
+ auto err =
+ consumer_->query_watermark_offsets(split_->topic,
split_->partition, &lowOffset, &highOffset, timeoutMs);
+ VELOX_CHECK(err == RdKafka::ERR_NO_ERROR, "Failed to get Kafka offset
range: {}", RdKafka::err2str(err));
+ nextOffset_ = std::max(nextOffset_, lowOffset);
+ splitDone_ = nextOffset_ >= std::min(highOffset, split_->endOffset);
+ if (!splitDone_) {
+ std::unique_ptr<RdKafka::TopicPartition> partition(
+ RdKafka::TopicPartition::create(split_->topic,
split_->partition, nextOffset_));
+ err = consumer_->seek(*partition, timeoutMs);
+ VELOX_CHECK(err == RdKafka::ERR_NO_ERROR, "Failed to seek Kafka
offset: {}", RdKafka::err2str(err));
+ }
+ break;
+ }
+ default:
+ // librdkafka retries transient errors itself; the poll timeout bounds
how long we wait.
+ LOG(WARNING) << "Kafka consume error on " << split_->toString() << ":
" << msg->errstr();
+ break;
+ }
+
+ if (!splitDone_ && std::chrono::steady_clock::now() >= deadline) {
+ if (!messages.empty()) {
+ return messages;
+ }
+ VELOX_FAIL(
+ "Timed out after {} ms fetching offset {} of {}-{}",
+ pollTimeoutMs,
+ nextOffset_,
+ split_->topic,
+ split_->partition);
+ }
+ }
+ return messages;
+}
+
+RowVectorPtr KafkaDataSource::toRowVector(const
std::vector<std::unique_ptr<RdKafka::Message>>& messages) {
+ std::vector<VectorPtr> children;
+ children.reserve(fields_.size());
+ for (size_t i = 0; i < fields_.size(); ++i) {
+ children.push_back(makeColumn(fields_[i], outputType_->childAt(i),
messages));
+ }
+ completedRows_ += messages.size();
+ return std::make_shared<RowVector>(pool_, outputType_, BufferPtr(nullptr),
messages.size(), std::move(children));
+}
+
+VectorPtr KafkaDataSource::makeColumn(
+ KafkaField field,
+ const TypePtr& type,
+ const std::vector<std::unique_ptr<RdKafka::Message>>& messages) {
+ const auto numRows = messages.size();
+ switch (field) {
+ case KafkaField::kKey:
+ case KafkaField::kValue:
+ case KafkaField::kTopic: {
+ auto vector = BaseVector::create<FlatVector<StringView>>(type, numRows,
pool_);
+ for (size_t i = 0; i < numRows; ++i) {
+ const auto& msg = *messages[i];
+ // FlatVector<StringView>::set copies non-inlined strings into the
vector's buffers.
+ if (field == KafkaField::kTopic) {
+ const auto topic = msg.topic_name();
+ vector->set(i, StringView(topic));
+ } else if (field == KafkaField::kKey) {
+ if (msg.key_pointer() == nullptr) {
+ vector->setNull(i, true);
+ } else {
+ vector->set(i, StringView(static_cast<const
char*>(msg.key_pointer()), msg.key_len()));
+ }
+ } else if (msg.payload() == nullptr) {
+ vector->setNull(i, true);
+ } else {
+ vector->set(i, StringView(static_cast<const char*>(msg.payload()),
msg.len()));
+ }
+ }
+ return vector;
+ }
+ case KafkaField::kPartition:
+ case KafkaField::kTimestampType: {
+ auto vector = BaseVector::create<FlatVector<int32_t>>(type, numRows,
pool_);
+ for (size_t i = 0; i < numRows; ++i) {
+ vector->set(
+ i,
+ field == KafkaField::kPartition ? messages[i]->partition()
+ :
toSparkTimestampType(messages[i]->timestamp().type));
+ }
+ return vector;
+ }
+ case KafkaField::kOffset: {
+ auto vector = BaseVector::create<FlatVector<int64_t>>(type, numRows,
pool_);
+ for (size_t i = 0; i < numRows; ++i) {
+ vector->set(i, messages[i]->offset());
+ }
+ return vector;
+ }
+ case KafkaField::kTimestamp: {
+ auto vector = BaseVector::create<FlatVector<Timestamp>>(type, numRows,
pool_);
+ for (size_t i = 0; i < numRows; ++i) {
+ vector->set(i,
Timestamp::fromMillis(messages[i]->timestamp().timestamp));
+ }
+ return vector;
+ }
+ }
+ VELOX_UNREACHABLE();
+}
+
+void KafkaDataSource::closeConsumer() {
+ if (consumer_ != nullptr) {
+ consumer_->close();
+ consumer_.reset();
+ }
+}
+
+} // namespace gluten
diff --git a/cpp/velox/operators/reader/KafkaDataSource.h
b/cpp/velox/operators/reader/KafkaDataSource.h
new file mode 100644
index 0000000000..59484b5e2a
--- /dev/null
+++ b/cpp/velox/operators/reader/KafkaDataSource.h
@@ -0,0 +1,96 @@
+/*
+ * 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 <memory>
+#include <string>
+#include <vector>
+
+#include "operators/reader/KafkaConnector.h"
+#include "operators/reader/KafkaSplit.h"
+#include "velox/connectors/Connector.h"
+
+namespace RdKafka {
+class KafkaConsumer;
+class Message;
+} // namespace RdKafka
+
+namespace gluten {
+
+/// Reads the offset range [startOffset, endOffset) of one Kafka topic
partition per split and
+/// produces rows in Spark's Kafka source schema. Offsets are never committed
back to Kafka; Spark
+/// tracks progress in its own checkpoint.
+class KafkaDataSource : public facebook::velox::connector::DataSource {
+ public:
+ KafkaDataSource(
+ const facebook::velox::RowTypePtr& outputType,
+ const facebook::velox::connector::ColumnHandleMap& columnHandles,
+ facebook::velox::memory::MemoryPool* pool);
+
+ ~KafkaDataSource() override;
+
+ void addSplit(std::shared_ptr<facebook::velox::connector::ConnectorSplit>
split) override;
+
+ std::optional<facebook::velox::RowVectorPtr> next(uint64_t size,
facebook::velox::ContinueFuture& future) override;
+
+ void addDynamicFilter(
+ facebook::velox::column_index_t /*outputChannel*/,
+ const std::shared_ptr<facebook::velox::common::Filter>& /*filter*/)
override {
+ VELOX_UNSUPPORTED("KafkaDataSource does not support dynamic filters");
+ }
+
+ uint64_t getCompletedBytes() override {
+ return completedBytes_;
+ }
+
+ uint64_t getCompletedRows() override {
+ return completedRows_;
+ }
+
+ private:
+ friend class KafkaConnectorTest;
+
+ // Polls up to 'maxRows' messages of the current split. Sets 'splitDone_'
once the end offset is
+ // reached.
+ std::vector<std::unique_ptr<RdKafka::Message>> poll(uint64_t maxRows);
+
+ facebook::velox::RowVectorPtr toRowVector(const
std::vector<std::unique_ptr<RdKafka::Message>>& messages);
+
+ facebook::velox::VectorPtr makeColumn(
+ KafkaField field,
+ const facebook::velox::TypePtr& type,
+ const std::vector<std::unique_ptr<RdKafka::Message>>& messages);
+
+ void closeConsumer();
+
+ const facebook::velox::RowTypePtr outputType_;
+ facebook::velox::memory::MemoryPool* const pool_;
+ // The Kafka field for each output column.
+ std::vector<KafkaField> fields_;
+
+ std::shared_ptr<const KafkaConnectorSplit> split_;
+ std::unique_ptr<RdKafka::KafkaConsumer> consumer_;
+ // The next offset expected from the current split.
+ int64_t nextOffset_{0};
+ bool splitDone_{false};
+
+ uint64_t completedBytes_{0};
+ uint64_t completedRows_{0};
+};
+
+} // namespace gluten
diff --git a/cpp/velox/operators/reader/KafkaSplit.h
b/cpp/velox/operators/reader/KafkaSplit.h
new file mode 100644
index 0000000000..df1e48a3c7
--- /dev/null
+++ b/cpp/velox/operators/reader/KafkaSplit.h
@@ -0,0 +1,72 @@
+/*
+ * 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 <string>
+#include <unordered_map>
+
+#include <fmt/format.h>
+
+#include "velox/connectors/Connector.h"
+
+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.
+class KafkaConnectorSplit : public facebook::velox::connector::ConnectorSplit {
+ public:
+ KafkaConnectorSplit(
+ const std::string& connectorId,
+ std::string topic,
+ int32_t partition,
+ int64_t startOffset,
+ int64_t endOffset,
+ int64_t pollTimeoutMs,
+ bool failOnDataLoss,
+ std::unordered_map<std::string, std::string> kafkaParams)
+ : ConnectorSplit(connectorId),
+ topic(std::move(topic)),
+ partition(partition),
+ startOffset(startOffset),
+ endOffset(endOffset),
+ pollTimeoutMs(pollTimeoutMs),
+ failOnDataLoss(failOnDataLoss),
+ kafkaParams(std::move(kafkaParams)) {}
+
+ std::string toString() const override {
+ return fmt::format(
+ "KafkaConnectorSplit[topic={}, partition={}, startOffset={},
endOffset={}]",
+ topic,
+ partition,
+ startOffset,
+ endOffset);
+ }
+
+ const std::string topic;
+ const int32_t partition;
+ const int64_t startOffset;
+ // Exclusive.
+ const int64_t endOffset;
+ const int64_t pollTimeoutMs;
+ const bool failOnDataLoss;
+ // Kafka consumer properties with the "kafka." prefix already stripped, e.g.
bootstrap.servers.
+ const std::unordered_map<std::string, std::string> kafkaParams;
+};
+
+} // namespace gluten
diff --git a/cpp/velox/tests/CMakeLists.txt b/cpp/velox/tests/CMakeLists.txt
index f8afb308cd..9abc4b6aa5 100644
--- a/cpp/velox/tests/CMakeLists.txt
+++ b/cpp/velox/tests/CMakeLists.txt
@@ -176,3 +176,7 @@ endif()
if(ENABLE_ENHANCED_FEATURES)
add_velox_test(velox_iceberg_test SOURCES iceberg/IcebergWriteTest.cc)
endif()
+
+if(ENABLE_KAFKA)
+ add_velox_test(velox_kafka_connector_test SOURCES KafkaConnectorTest.cc)
+endif()
diff --git a/cpp/velox/tests/KafkaConnectorTest.cc
b/cpp/velox/tests/KafkaConnectorTest.cc
new file mode 100644
index 0000000000..026da1b421
--- /dev/null
+++ b/cpp/velox/tests/KafkaConnectorTest.cc
@@ -0,0 +1,367 @@
+/*
+ * 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 <gtest/gtest.h>
+#include <librdkafka/rdkafka.h>
+#include <librdkafka/rdkafka_mock.h>
+#include <librdkafka/rdkafkacpp.h>
+
+#include "operators/reader/KafkaConnector.h"
+#include "operators/reader/KafkaDataSource.h"
+#include "velox/common/base/tests/GTestUtils.h"
+#include "velox/vector/FlatVector.h"
+#include "velox/vector/tests/utils/VectorTestBase.h"
+
+using namespace facebook::velox;
+
+namespace gluten {
+namespace {
+
+constexpr char kConnectorId[] = "kafka-test";
+constexpr char kTopic[] = "native-kafka-test";
+
+// Supply metadata independently of the mock broker's timestamp support.
+class KafkaMetadataMessage final : public RdKafka::Message {
+ public:
+ KafkaMetadataMessage(
+ int32_t partition,
+ RdKafka::MessageTimestamp::MessageTimestampType timestampType,
+ int64_t timestamp)
+ : partition_(partition), timestamp_{timestampType, timestamp} {}
+
+ int32_t partition() const override {
+ return partition_;
+ }
+
+ RdKafka::MessageTimestamp timestamp() const override {
+ return timestamp_;
+ }
+
+ std::string errstr() const override {
+ VELOX_UNREACHABLE();
+ }
+
+ RdKafka::ErrorCode err() const override {
+ VELOX_UNREACHABLE();
+ }
+
+ RdKafka::Topic* topic() const override {
+ VELOX_UNREACHABLE();
+ }
+
+ std::string topic_name() const override {
+ VELOX_UNREACHABLE();
+ }
+
+ void* payload() const override {
+ VELOX_UNREACHABLE();
+ }
+
+ size_t len() const override {
+ VELOX_UNREACHABLE();
+ }
+
+ const std::string* key() const override {
+ VELOX_UNREACHABLE();
+ }
+
+ const void* key_pointer() const override {
+ VELOX_UNREACHABLE();
+ }
+
+ size_t key_len() const override {
+ VELOX_UNREACHABLE();
+ }
+
+ int64_t offset() const override {
+ VELOX_UNREACHABLE();
+ }
+
+ void* msg_opaque() const override {
+ VELOX_UNREACHABLE();
+ }
+
+ int64_t latency() const override {
+ VELOX_UNREACHABLE();
+ }
+
+ rd_kafka_message_t* c_ptr() override {
+ VELOX_UNREACHABLE();
+ }
+
+ Status status() const override {
+ VELOX_UNREACHABLE();
+ }
+
+ RdKafka::Headers* headers() override {
+ VELOX_UNREACHABLE();
+ }
+
+ RdKafka::Headers* headers(RdKafka::ErrorCode* /*error*/) override {
+ VELOX_UNREACHABLE();
+ }
+
+ int32_t broker_id() const override {
+ VELOX_UNREACHABLE();
+ }
+
+ int32_t leader_epoch() const override {
+ VELOX_UNREACHABLE();
+ }
+
+ RdKafka::Error* offset_store() override {
+ VELOX_UNREACHABLE();
+ }
+
+ private:
+ const int32_t partition_;
+ const RdKafka::MessageTimestamp timestamp_;
+};
+
+} // namespace
+
+class KafkaConnectorTest : public ::testing::Test, public test::VectorTestBase
{
+ protected:
+ static void SetUpTestCase() {
+
memory::MemoryManager::testingSetInstance(memory::MemoryManager::Options{});
+ }
+
+ void SetUp() override {
+ auto* conf = rd_kafka_conf_new();
+ char error[512];
+ ASSERT_EQ(rd_kafka_conf_set(conf, "test.mock.num.brokers", "1", error,
sizeof(error)), RD_KAFKA_CONF_OK);
+ producer_ = rd_kafka_new(RD_KAFKA_PRODUCER, conf, error, sizeof(error));
+ ASSERT_NE(producer_, nullptr) << error;
+ cluster_ = rd_kafka_handle_mock_cluster(producer_);
+ ASSERT_NE(cluster_, nullptr);
+ ASSERT_EQ(rd_kafka_mock_topic_create(cluster_, kTopic, 1, 1),
RD_KAFKA_RESP_ERR_NO_ERROR);
+ }
+
+ void TearDown() override {
+ if (producer_ != nullptr) {
+ rd_kafka_destroy(producer_);
+ }
+ }
+
+ void produce(const std::optional<std::string>& key, const
std::optional<std::string>& value) {
+ // librdkafka's macros cast arguments without grouping them, so
parenthesize each conditional.
+ ASSERT_EQ(
+ rd_kafka_producev(
+ producer_,
+ RD_KAFKA_V_TOPIC(kTopic),
+ RD_KAFKA_V_PARTITION(0),
+ RD_KAFKA_V_MSGFLAGS(RD_KAFKA_MSG_F_COPY),
+ RD_KAFKA_V_KEY((key ? key->data() : nullptr), (key ? key->size() :
0)),
+ RD_KAFKA_V_VALUE((value ? const_cast<char*>(value->data()) :
nullptr), (value ? value->size() : 0)),
+ RD_KAFKA_V_END),
+ RD_KAFKA_RESP_ERR_NO_ERROR);
+ ASSERT_EQ(rd_kafka_flush(producer_, 10'000), RD_KAFKA_RESP_ERR_NO_ERROR);
+ }
+
+ std::shared_ptr<KafkaConnectorSplit> split(
+ int64_t start,
+ int64_t end,
+ bool failOnDataLoss = true,
+ std::unordered_map<std::string, std::string> kafkaParams = {}) {
+ kafkaParams.emplace("bootstrap.servers",
rd_kafka_mock_cluster_bootstraps(cluster_));
+ return std::make_shared<KafkaConnectorSplit>(
+ kConnectorId, kTopic, 0, start, end, 10'000, failOnDataLoss,
std::move(kafkaParams));
+ }
+
+ std::unique_ptr<KafkaDataSource> source(const RowTypePtr& type) {
+ connector::ColumnHandleMap columns;
+ for (size_t i = 0; i < type->size(); ++i) {
+ columns.emplace(type->nameOf(i),
std::make_shared<KafkaColumnHandle>(type->nameOf(i), type->childAt(i)));
+ }
+ return std::make_unique<KafkaDataSource>(type, columns, pool());
+ }
+
+ RowVectorPtr toRowVector(KafkaDataSource& reader, const
std::vector<std::unique_ptr<RdKafka::Message>>& messages) {
+ return reader.toRowVector(messages);
+ }
+
+ rd_kafka_t* producer_{nullptr};
+ rd_kafka_mock_cluster_t* cluster_{nullptr};
+};
+
+TEST_F(KafkaConnectorTest, columnValidation) {
+ EXPECT_EQ(KafkaColumnHandle("timestampType", INTEGER()).field(),
KafkaField::kTimestampType);
+ VELOX_ASSERT_THROW(KafkaColumnHandle("offset", VARCHAR()), "Kafka column
offset must be of type BIGINT");
+ VELOX_ASSERT_THROW(KafkaColumnHandle("headers", VARBINARY()), "Unsupported
Kafka column: headers");
+}
+
+TEST_F(KafkaConnectorTest, projectedPartitionAndTimestampMetadata) {
+ using TimestampType = RdKafka::MessageTimestamp;
+ std::vector<std::unique_ptr<RdKafka::Message>> messages;
+ messages.push_back(
+ std::make_unique<KafkaMetadataMessage>(3,
TimestampType::MSG_TIMESTAMP_CREATE_TIME, 1'700'000'000'123));
+ messages.push_back(
+ std::make_unique<KafkaMetadataMessage>(7,
TimestampType::MSG_TIMESTAMP_LOG_APPEND_TIME, 1'700'000'456'789));
+ messages.push_back(std::make_unique<KafkaMetadataMessage>(11,
TimestampType::MSG_TIMESTAMP_NOT_AVAILABLE, -1));
+
+ // Prune and reorder the metadata columns from Spark's Kafka schema.
+ auto reader = source(ROW({"timestampType", "partition", "timestamp"},
{INTEGER(), INTEGER(), TIMESTAMP()}));
+ auto batch = toRowVector(*reader, messages);
+ ASSERT_EQ(batch->size(), 3);
+
+ const std::vector<int32_t> expectedTypes{0, 1, -1};
+ const std::vector<int32_t> expectedPartitions{3, 7, 11};
+ // Spark converts the unavailable -1 ms sentinel to a non-null timestamp.
+ const std::vector<Timestamp> expectedTimestamps{
+ Timestamp(1'700'000'000, 123'000'000), Timestamp(1'700'000'456,
789'000'000), Timestamp(-1, 999'000'000)};
+ auto* timestampTypes = batch->childAt(0)->as<FlatVector<int32_t>>();
+ auto* partitions = batch->childAt(1)->as<FlatVector<int32_t>>();
+ auto* timestamps = batch->childAt(2)->as<FlatVector<Timestamp>>();
+ for (vector_size_t row = 0; row < batch->size(); ++row) {
+ SCOPED_TRACE(row);
+ ASSERT_FALSE(timestampTypes->isNullAt(row));
+ ASSERT_FALSE(partitions->isNullAt(row));
+ ASSERT_FALSE(timestamps->isNullAt(row));
+ EXPECT_EQ(timestampTypes->valueAt(row), expectedTypes[row]);
+ EXPECT_EQ(partitions->valueAt(row), expectedPartitions[row]);
+ EXPECT_EQ(timestamps->valueAt(row), expectedTimestamps[row]);
+ }
+}
+
+TEST_F(KafkaConnectorTest, emptyAndConsecutiveSplits) {
+ auto reader = source(ROW({"offset"}, {BIGINT()}));
+ ContinueFuture future = ContinueFuture::makeEmpty();
+ // Empty ranges do not need consumer properties or a Kafka connection.
+ auto empty = std::make_shared<KafkaConnectorSplit>(
+ kConnectorId, kTopic, 0, 4, 4, 10'000, true,
std::unordered_map<std::string, std::string>{});
+ reader->addSplit(empty);
+ VELOX_ASSERT_THROW(reader->addSplit(empty), "Previous Kafka split has not
been fully processed");
+ EXPECT_EQ(reader->next(10, future).value(), nullptr);
+ reader->addSplit(empty);
+ EXPECT_EQ(reader->next(10, future).value(), nullptr);
+ EXPECT_EQ(reader->getCompletedRows(), 0);
+ EXPECT_EQ(reader->getCompletedBytes(), 0);
+ VELOX_ASSERT_THROW(reader->addSplit(nullptr), "Kafka split must not be
null");
+}
+
+TEST_F(KafkaConnectorTest, exclusiveOffsetsAndProjectedBinaryRows) {
+ const std::string key("binary\0key", 10);
+ const std::string value("a long binary\0value preserved after polling", 43);
+ produce("excluded", "before range");
+ produce(key, value);
+ produce(std::nullopt, std::nullopt);
+ produce("", "");
+ produce("excluded", "after range");
+
+ // Prune and reorder the Spark Kafka schema. Reading in single-row batches
also checks that
+ // strings remain owned by the output vector after librdkafka messages have
been destroyed.
+ auto reader = source(ROW({"value", "offset", "key", "topic"}, {VARBINARY(),
BIGINT(), VARBINARY(), VARCHAR()}));
+ reader->addSplit(split(1, 3));
+ ContinueFuture future = ContinueFuture::makeEmpty();
+ auto first = reader->next(1, future).value();
+ ASSERT_NE(first, nullptr);
+ ASSERT_EQ(first->size(), 1);
+ EXPECT_EQ(first->childAt(0)->as<FlatVector<StringView>>()->valueAt(0).str(),
value);
+ EXPECT_EQ(first->childAt(1)->as<FlatVector<int64_t>>()->valueAt(0), 1);
+ EXPECT_EQ(first->childAt(2)->as<FlatVector<StringView>>()->valueAt(0).str(),
key);
+ EXPECT_EQ(first->childAt(3)->as<FlatVector<StringView>>()->valueAt(0).str(),
kTopic);
+
+ auto second = reader->next(1, future).value();
+ ASSERT_NE(second, nullptr);
+ ASSERT_EQ(second->size(), 1);
+ EXPECT_TRUE(second->childAt(0)->isNullAt(0));
+ EXPECT_TRUE(second->childAt(2)->isNullAt(0));
+ EXPECT_EQ(second->childAt(1)->as<FlatVector<int64_t>>()->valueAt(0), 2);
+ EXPECT_EQ(reader->next(1, future).value(), nullptr);
+
+ reader->addSplit(split(3, 4));
+ auto third = reader->next(10, future).value();
+ ASSERT_NE(third, nullptr);
+ ASSERT_EQ(third->size(), 1);
+ EXPECT_FALSE(third->childAt(0)->isNullAt(0));
+ EXPECT_FALSE(third->childAt(2)->isNullAt(0));
+ EXPECT_EQ(third->childAt(0)->as<FlatVector<StringView>>()->valueAt(0).str(),
"");
+ EXPECT_EQ(third->childAt(2)->as<FlatVector<StringView>>()->valueAt(0).str(),
"");
+ EXPECT_EQ(third->childAt(1)->as<FlatVector<int64_t>>()->valueAt(0), 3);
+ EXPECT_EQ(reader->next(10, future).value(), nullptr);
+ EXPECT_EQ(reader->getCompletedRows(), 3);
+ EXPECT_EQ(reader->getCompletedBytes(), key.size() + value.size());
+}
+
+TEST_F(KafkaConnectorTest, offsetOutOfRange) {
+ produce("key", "value");
+ ContinueFuture future = ContinueFuture::makeEmpty();
+ auto reader = source(ROW({"offset"}, {BIGINT()}));
+ reader->addSplit(split(100, 101));
+ VELOX_ASSERT_THROW(reader->next(1, future), "is no longer available");
+
+ auto ignoreDataLoss = source(ROW({"offset"}, {BIGINT()}));
+ ignoreDataLoss->addSplit(split(100, 101, false));
+ EXPECT_EQ(ignoreDataLoss->next(1, future).value(), nullptr);
+ EXPECT_EQ(ignoreDataLoss->getCompletedRows(), 0);
+}
+
+TEST_F(KafkaConnectorTest, missingLogTail) {
+ produce("key", "value");
+ ContinueFuture future = ContinueFuture::makeEmpty();
+ for (const auto* isolationLevel : {"read_uncommitted", "read_committed"}) {
+ SCOPED_TRACE(isolationLevel);
+ const std::unordered_map<std::string, std::string>
params{{"isolation.level", isolationLevel}};
+ auto reader = source(ROW({"offset"}, {BIGINT()}));
+ // Emulate a split planned before the log was truncated: offset 0 still
exists, but 1 does not.
+ reader->addSplit(split(0, 2, false, params));
+ auto batch = reader->next(10, future).value();
+ ASSERT_NE(batch, nullptr);
+ ASSERT_EQ(batch->size(), 1);
+ EXPECT_EQ(batch->childAt(0)->as<FlatVector<int64_t>>()->valueAt(0), 0);
+ EXPECT_EQ(reader->next(1, future).value(), nullptr);
+ EXPECT_EQ(reader->getCompletedRows(), 1);
+
+ // Completing the truncated range must release the split and allow another
one.
+ reader->addSplit(split(0, 1, true, params));
+ ASSERT_NE(reader->next(1, future).value(), nullptr);
+ EXPECT_EQ(reader->next(1, future).value(), nullptr);
+
+ auto failOnDataLoss = source(ROW({"offset"}, {BIGINT()}));
+ failOnDataLoss->addSplit(split(0, 2, true, params));
+ ASSERT_NE(failOnDataLoss->next(1, future).value(), nullptr);
+ VELOX_ASSERT_THROW(failOnDataLoss->next(1, future), "Some data may have
been lost");
+ }
+}
+
+TEST_F(KafkaConnectorTest, missingLogTailWithoutRecords) {
+ ContinueFuture future = ContinueFuture::makeEmpty();
+ auto reader = source(ROW({"offset"}, {BIGINT()}));
+ reader->addSplit(split(0, 1, false));
+ EXPECT_EQ(reader->next(1, future).value(), nullptr);
+ EXPECT_EQ(reader->getCompletedRows(), 0);
+
+ auto failOnDataLoss = source(ROW({"offset"}, {BIGINT()}));
+ failOnDataLoss->addSplit(split(0, 1));
+ VELOX_ASSERT_THROW(failOnDataLoss->next(1, future), "Some data may have been
lost");
+}
+
+TEST_F(KafkaConnectorTest, invalidKafkaProperty) {
+ auto reader = source(ROW({"offset"}, {BIGINT()}));
+ auto invalid = std::make_shared<KafkaConnectorSplit>(
+ kConnectorId,
+ kTopic,
+ 0,
+ 0,
+ 1,
+ 10'000,
+ true,
+ std::unordered_map<std::string, std::string>{{"isolation.level",
"invalid"}});
+ VELOX_ASSERT_THROW(reader->addSplit(invalid), "Invalid Kafka property
isolation.level");
+}
+
+} // namespace gluten
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
