Copilot commented on code in PR #13183: URL: https://github.com/apache/gluten/pull/13183#discussion_r4274325758
########## 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) { Review Comment: `tableHandle` is ignored entirely. This makes it impossible to detect mismatched/invalid table handles early (e.g., a non-Kafka table handle routed to this connector), and can lead to harder-to-debug failures later. Consider validating the handle type (e.g., dynamic cast to `KafkaTableHandle`) and failing with a user-facing error if it’s not the expected type. ########## cpp/velox/operators/reader/KafkaDataSource.cc: ########## @@ -0,0 +1,313 @@ +/* + * 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/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); +} + +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: + // The offset of an EOF event is the log end offset. Offsets below it without a record are + // compacted records or transaction markers. + splitDone_ = msg->offset() >= split_->endOffset; + 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)); Review Comment: `RdKafka::Message::timestamp()` can return `MSG_TIMESTAMP_NOT_AVAILABLE` with a sentinel timestamp value (commonly `-1`). Unconditionally converting that value via `Timestamp::fromMillis(...)` can yield invalid/incorrect output rather than Spark-compatible semantics (often null/absent). Handle the “not available” timestamp type explicitly (e.g., set nulls for `timestamp` and/or align with Spark’s expected behavior for missing timestamps). ########## 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(); Review Comment: `KafkaTableHandle::serialize()` is currently `VELOX_NYI()`, which will crash if Velox/framework code attempts to serialize table handles (common for plan caching, exchange, or debugging). Implement a minimal serialization payload (and corresponding deserialization if required by the surrounding framework), or return a clear `VELOX_UNSUPPORTED` message until the connector is fully wired for serialized plans. ########## cpp/velox/operators/reader/KafkaSplit.h: ########## @@ -0,0 +1,70 @@ +/* + * 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 "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); + } Review Comment: `KafkaSplit.h` uses `fmt::format` but does not include a fmt header. This can break compilation depending on transitive includes. Add an explicit include (e.g. `<fmt/format.h>` or the project’s preferred fmt header) in this file. ########## cpp/velox/operators/reader/KafkaDataSource.cc: ########## @@ -0,0 +1,313 @@ +/* + * 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/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); +} + +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: + // The offset of an EOF event is the log end offset. Offsets below it without a record are + // compacted records or transaction markers. + splitDone_ = msg->offset() >= split_->endOffset; + 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) { Review Comment: `makeColumn()` introduces non-trivial behavior for `partition`, `timestamp`, and `timestampType` (including mapping to Spark’s ids and timestamp conversion), but the added unit tests primarily exercise key/value/topic/offset and offset-range handling. Add test coverage for `timestamp` + `timestampType` (including the “timestamp not available” case) to prevent regressions across librdkafka/platform versions. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
