diff --git a/backends-clickhouse/src/main/java/org/apache/spark/sql/execution/datasources/clickhouse/ExtensionTableNode.java b/backends-clickhouse/src/main/java/org/apache/spark/sql/execution/datasources/clickhouse/ExtensionTableNode.java index 7e5de6eb463..dc88abe6722 100644 --- a/backends-clickhouse/src/main/java/org/apache/spark/sql/execution/datasources/clickhouse/ExtensionTableNode.java +++ b/backends-clickhouse/src/main/java/org/apache/spark/sql/execution/datasources/clickhouse/ExtensionTableNode.java @@ -16,8 +16,8 @@ */ package org.apache.spark.sql.execution.datasources.clickhouse; -import org.apache.gluten.backendsapi.BackendsApiManager; import org.apache.gluten.substrait.rel.SplitInfo; +import org.apache.gluten.utils.SubstraitUtil; import com.google.protobuf.StringValue; import io.substrait.proto.ReadRel; @@ -59,10 +59,6 @@ public String getExtensionTableStr() { } public static ReadRel.ExtensionTable toProtobuf(String result) { - ReadRel.ExtensionTable.Builder extensionTableBuilder = ReadRel.ExtensionTable.newBuilder(); - StringValue extensionTable = StringValue.newBuilder().setValue(result).build(); - extensionTableBuilder.setDetail( - BackendsApiManager.getTransformerApiInstance().packPBMessage(extensionTable)); - return extensionTableBuilder.build(); + return SubstraitUtil.packExtensionTable(StringValue.newBuilder().setValue(result).build()); } } diff --git a/cpp-ch/local-engine/Parser/RelParsers/ReadRelParser.cpp b/cpp-ch/local-engine/Parser/RelParsers/ReadRelParser.cpp index 17bdbc5cfd2..91e5c79d7c6 100644 --- a/cpp-ch/local-engine/Parser/RelParsers/ReadRelParser.cpp +++ b/cpp-ch/local-engine/Parser/RelParsers/ReadRelParser.cpp @@ -35,6 +35,7 @@ #include #include #include +#include #include #include @@ -122,7 +123,8 @@ bool ReadRelParser::isReadRelFromLocalFile(const substrait::ReadRel & rel) if (rel.has_local_files()) return !isReadRelFromJavaIter(rel); else - return !rel.has_extension_table() && !isReadRelFromMergeTree(rel) && !isReadRelFromRange(rel) && !isReadFromStreamKafka(rel); + /// Kafka reads are identified by an extension table, so !has_extension_table() already excludes them. + return !rel.has_extension_table() && !isReadRelFromMergeTree(rel) && !isReadRelFromRange(rel); } bool ReadRelParser::isReadRelFromMergeTree(const substrait::ReadRel & rel) @@ -161,7 +163,7 @@ bool ReadRelParser::isReadRelFromRange(const substrait::ReadRel & rel) bool ReadRelParser::isReadFromStreamKafka(const substrait::ReadRel & rel) { - return rel.has_stream_kafka() && rel.stream_kafka(); + return rel.has_extension_table() && rel.extension_table().detail().Is(); } DB::QueryPlanStepPtr ReadRelParser::parseReadRelWithJavaIter(const substrait::ReadRel & rel) diff --git a/cpp-ch/local-engine/Parser/RelParsers/StreamKafkaRelParser.cpp b/cpp-ch/local-engine/Parser/RelParsers/StreamKafkaRelParser.cpp index cb45e4f3ebb..2f2dc379c3c 100644 --- a/cpp-ch/local-engine/Parser/RelParsers/StreamKafkaRelParser.cpp +++ b/cpp-ch/local-engine/Parser/RelParsers/StreamKafkaRelParser.cpp @@ -20,7 +20,9 @@ #include #include #include +#include #include +#include #include namespace DB @@ -30,6 +32,7 @@ namespace ErrorCodes { extern const int NO_SUCH_DATA_PART; extern const int LOGICAL_ERROR; +extern const int CANNOT_PARSE_PROTOBUF_SCHEMA; extern const int UNKNOWN_FUNCTION; extern const int UNKNOWN_TYPE; } @@ -50,10 +53,19 @@ StreamKafkaRelParser::parse(DB::QueryPlanPtr query_plan, const substrait::Rel & DB::QueryPlanPtr StreamKafkaRelParser::parseRelImpl(DB::QueryPlanPtr query_plan, const substrait::ReadRel & read_rel) { - if (!read_rel.has_stream_kafka()) - throw DB::Exception(DB::ErrorCodes::LOGICAL_ERROR, "Can't not parse kafka rel, because of read rel don't contained stream kafka"); + if (split_info.empty()) + throw DB::Exception(DB::ErrorCodes::LOGICAL_ERROR, "Can't parse kafka rel, because no split info was supplied for it"); + + auto extension_table = BinaryToMessage(split_info); + debug::dumpMessage(extension_table, "extension_table"); + + gluten::StreamKafka kafka_task; + if (!extension_table.detail().UnpackTo(&kafka_task)) + throw DB::Exception( + DB::ErrorCodes::CANNOT_PARSE_PROTOBUF_SCHEMA, + "Can't parse kafka rel, expected an extension table detail of type gluten.StreamKafka but got '{}'", + extension_table.detail().type_url()); - auto kafka_task = BinaryToMessage(split_info); auto topic = kafka_task.topic_partition().topic(); auto partition = kafka_task.topic_partition().partition(); auto start_offset = kafka_task.start_offset(); diff --git a/cpp-ch/local-engine/Parser/SerializedPlanParser.cpp b/cpp-ch/local-engine/Parser/SerializedPlanParser.cpp index 5ca985575dd..594484adaed 100644 --- a/cpp-ch/local-engine/Parser/SerializedPlanParser.cpp +++ b/cpp-ch/local-engine/Parser/SerializedPlanParser.cpp @@ -270,6 +270,12 @@ QueryPlanPtr SerializedPlanParser::parseOp(const substrait::Rel & rel, std::list } else if (read_rel_parser->isReadFromStreamKafka(read)) { + /// Unlike MergeTree/Range above, a Kafka read is *identified* by an in-plan extension_table + /// (its detail's type_url is gluten.StreamKafka), so has_extension_table() is always true here + /// yet the real per-partition payload still rides split-info -- we must consume a split + /// unconditionally. Do NOT add a `!read.has_extension_table()` guard like the siblings have: + /// that would stop split_info_index from advancing and hand later leaves the wrong split. + chassert(read.has_extension_table()); read_rel_parser->setSplitInfo(nextSplitInfo()); } } diff --git a/cpp-ch/local-engine/proto/CMakeLists.txt b/cpp-ch/local-engine/proto/CMakeLists.txt index 0a12af465df..f06dc0a13e5 100644 --- a/cpp-ch/local-engine/proto/CMakeLists.txt +++ b/cpp-ch/local-engine/proto/CMakeLists.txt @@ -12,7 +12,7 @@ # 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. -file(GLOB protobuf_files ./*.proto substrait/*.proto +file(GLOB protobuf_files CONFIGURE_DEPENDS ./*.proto substrait/*.proto substrait/extensions/*.proto) foreach(FIL ${protobuf_files}) diff --git a/cpp-ch/local-engine/proto/kafka.proto b/cpp-ch/local-engine/proto/kafka.proto new file mode 120000 index 00000000000..dbd46d6d5d3 --- /dev/null +++ b/cpp-ch/local-engine/proto/kafka.proto @@ -0,0 +1 @@ +../../../gluten-core/src/main/resources/org/apache/gluten/proto/kafka.proto \ No newline at end of file diff --git a/docs/developers/SubstraitModifications.md b/docs/developers/SubstraitModifications.md index 11aa46eb1cf..3849df7c04d 100644 --- a/docs/developers/SubstraitModifications.md +++ b/docs/developers/SubstraitModifications.md @@ -35,7 +35,6 @@ changed `Unbounded` in `WindowFunction` into `Unbounded_Preceding` and `Unbounde * Added `TopNRel` ([#5409](https://github.com/apache/gluten/pull/5409)). * Added `ref` field in window bound `Preceding` and `Following` ([#5626](https://github.com/apache/gluten/pull/5626)). * Added `BucketSpec` field in `WriteRel`([#8386](https://github.com/apache/gluten/pull/8386)) -* Added `StreamKafka` in `ReadRel`([#8321](https://github.com/apache/gluten/pull/8321)) * Rebased the `WriteRel` body onto upstream `v0.98.0`: field 7 is now `common`, with `create_mode` (8) and `advanced_extension` (9) added, and the `OutputMode` value `OUTPUT_MODE_MODIFIED_TUPLES` renamed to `OUTPUT_MODE_MODIFIED_RECORDS`. Gluten's `BucketSpec` field moved off field 7 to 1000. The enclosing diff --git a/gluten-core/src/main/resources/org/apache/gluten/proto/kafka.proto b/gluten-core/src/main/resources/org/apache/gluten/proto/kafka.proto new file mode 100644 index 00000000000..91ef7253832 --- /dev/null +++ b/gluten-core/src/main/resources/org/apache/gluten/proto/kafka.proto @@ -0,0 +1,29 @@ +// SPDX-License-Identifier: Apache-2.0 +syntax = "proto3"; + +package gluten; + +option java_package = "org.apache.gluten.proto"; +option java_multiple_files = true; + +// Kafka streaming-source payloads for Gluten's ClickHouse backend. A payload is +// packed into a `google.protobuf.Any` and carried in the official +// `substrait.ReadRel.ExtensionTable.detail` field; the native consumer +// discriminates a Kafka read by the `Any`'s type_url (`gluten.StreamKafka`). + +// Streaming Kafka source, used for KafkaBatch / KafkaContinuous reads. +// Produced per-partition as the detail of a ReadRel.ExtensionTable. +message StreamKafka { + message TopicPartition { + string topic = 1; + int32 partition = 2; + } + + TopicPartition topic_partition = 1; + int64 start_offset = 2; + int64 end_offset = 3; + map params = 4; + int64 poll_timeout_ms = 5; + bool fail_on_data_loss = 6; + bool include_headers = 7; +} diff --git a/gluten-kafka/pom.xml b/gluten-kafka/pom.xml index eabfc1119a7..20e10e4285c 100644 --- a/gluten-kafka/pom.xml +++ b/gluten-kafka/pom.xml @@ -52,6 +52,12 @@ ${spark.version} provided + + com.google.protobuf + protobuf-java + ${protobuf.version} + provided + @@ -88,12 +94,6 @@ ${hadoop.version} test - - com.google.protobuf - protobuf-java - ${protobuf.version} - test - org.scalatest scalatest_${scala.binary.version} diff --git a/gluten-kafka/src/main/scala/org/apache/gluten/execution/MicroBatchScanExecTransformer.scala b/gluten-kafka/src/main/scala/org/apache/gluten/execution/MicroBatchScanExecTransformer.scala index d2777f04ae3..11071a0f50e 100644 --- a/gluten-kafka/src/main/scala/org/apache/gluten/execution/MicroBatchScanExecTransformer.scala +++ b/gluten-kafka/src/main/scala/org/apache/gluten/execution/MicroBatchScanExecTransformer.scala @@ -16,6 +16,8 @@ */ package org.apache.gluten.execution +import org.apache.gluten.backendsapi.BackendsApiManager +import org.apache.gluten.proto.StreamKafka import org.apache.gluten.substrait.SubstraitContext import org.apache.gluten.substrait.rel.{ReadRelNode, SplitInfo} import org.apache.gluten.substrait.rel.LocalFilesNode.ReadFileFormat @@ -105,7 +107,13 @@ case class MicroBatchScanExecTransformer( override protected def doTransform(context: SubstraitContext): TransformContext = { val ctx = super.doTransform(context) - ctx.root.asInstanceOf[ReadRelNode].setStreamKafka(true) + // Mark this read as a Kafka stream by stamping an empty gluten.StreamKafka into the in-plan + // ReadRel.ExtensionTable. The native consumer discriminates on the detail's type_url; the real + // per-partition offsets/params still ride the split-info payload (see StreamKafkaSourceNode). + ctx.root + .asInstanceOf[ReadRelNode] + .setExtensionTableDetail( + BackendsApiManager.getTransformerApiInstance.packPBMessage(StreamKafka.getDefaultInstance)) ctx } } diff --git a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/ReadRelNode.java b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/ReadRelNode.java index 33332f18672..5a5c8f41394 100644 --- a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/ReadRelNode.java +++ b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/ReadRelNode.java @@ -44,7 +44,7 @@ public class ReadRelNode implements RelNode, Serializable { private final List columnTypeNodes = new ArrayList<>(); private final ExpressionNode filterNode; private final AdvancedExtensionNode extensionNode; - private boolean streamKafka = false; + private Any extensionTableDetail; private BigInt rowCount; private InputStats inputStats; @@ -78,8 +78,9 @@ public void setInputStats(InputStats inputStats) { this.inputStats = inputStats; } - public void setStreamKafka(boolean streamKafka) { - this.streamKafka = streamKafka; + /** Sets the detail of the {@code read_type} oneof's {@code extension_table} member. */ + public void setExtensionTableDetail(Any extensionTableDetail) { + this.extensionTableDetail = extensionTableDetail; } @Override @@ -93,7 +94,11 @@ public Rel toProtobuf() { ReadRel.Builder readBuilder = ReadRel.newBuilder(); readBuilder.setCommon(relCommonBuilder.build()); readBuilder.setBaseSchema(nStructBuilder.build()); - readBuilder.setStreamKafka(streamKafka); + + if (extensionTableDetail != null) { + readBuilder.setExtensionTable( + ReadRel.ExtensionTable.newBuilder().setDetail(extensionTableDetail).build()); + } if (filterNode != null) { readBuilder.setFilter(filterNode.toProtobuf()); diff --git a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/StreamKafkaSourceNode.java b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/StreamKafkaSourceNode.java index 5974e3bb0f4..3ae4844459c 100644 --- a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/StreamKafkaSourceNode.java +++ b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/StreamKafkaSourceNode.java @@ -16,6 +16,9 @@ */ package org.apache.gluten.substrait.rel; +import org.apache.gluten.proto.StreamKafka; +import org.apache.gluten.utils.SubstraitUtil; + import io.substrait.proto.ReadRel; import java.util.Collections; @@ -59,11 +62,10 @@ public List preferredLocations() { } @Override - public ReadRel.StreamKafka toProtobuf() { - ReadRel.StreamKafka.Builder builder = ReadRel.StreamKafka.newBuilder(); + public ReadRel.ExtensionTable toProtobuf() { + StreamKafka.Builder builder = StreamKafka.newBuilder(); - ReadRel.StreamKafka.TopicPartition.Builder topicPartition = - ReadRel.StreamKafka.TopicPartition.newBuilder(); + StreamKafka.TopicPartition.Builder topicPartition = StreamKafka.TopicPartition.newBuilder(); topicPartition.setTopic(topic); topicPartition.setPartition(partition); @@ -75,6 +77,6 @@ public ReadRel.StreamKafka toProtobuf() { builder.setIncludeHeaders(includeHeaders); kafkaParams.forEach((k, v) -> builder.putParams(k, v.toString())); - return builder.build(); + return SubstraitUtil.packExtensionTable(builder.build()); } } diff --git a/gluten-substrait/src/main/resources/substrait/proto/substrait/algebra.proto b/gluten-substrait/src/main/resources/substrait/proto/substrait/algebra.proto index e2396912085..f19015d419a 100644 --- a/gluten-substrait/src/main/resources/substrait/proto/substrait/algebra.proto +++ b/gluten-substrait/src/main/resources/substrait/proto/substrait/algebra.proto @@ -69,7 +69,7 @@ message ReadRel { LocalFiles local_files = 6; NamedTable named_table = 7; ExtensionTable extension_table = 8; - bool stream_kafka = 9; + IcebergTable iceberg_table = 9; } // A base table. The list of string is used to represent namespacing (e.g., mydb.mytable). @@ -79,6 +79,29 @@ message ReadRel { substrait.extensions.AdvancedExtension advanced_extension = 10; } + // Read an Iceberg Table + message IcebergTable { + oneof table_type { + MetadataFileRead direct = 1; + // future: add catalog table types (e.g. rest api, latest metadata in path, etc) + } + + // Read an Iceberg table using a metadata file. Implicit assumption: required credentials are already known by plan consumer. + message MetadataFileRead { + // the specific uri of a metadata file (e.g. s3://mybucket/mytable/-.metadata.json) + string metadata_uri = 1; + + // snapshot options. if none set, uses the current snapshot listed in the metadata file + oneof snapshot { + // the snapshot id to read. + string snapshot_id = 2; + + // the timestamp that should be used to select the snapshot (Time passed in microseconds since 1970-01-01 00:00:00.000000 in UTC) + int64 snapshot_timestamp = 3; + } + } + } + // A table composed of literals. message VirtualTable { repeated Expression.Literal.Struct values = 1; @@ -90,22 +113,6 @@ message ReadRel { google.protobuf.Any detail = 1; } - // Used to KafkaBatch or KafkaContinuous source - message StreamKafka { - message TopicPartition { - string topic = 1; - int32 partition = 2; - } - - TopicPartition topic_partition = 1; - int64 start_offset = 2; - int64 end_offset = 3; - map params = 4; - int64 poll_timeout_ms = 5; - bool fail_on_data_loss = 6; - bool include_headers = 7; - } - // Represents a list of files in input of a scan operation message LocalFiles { repeated FileOrFiles items = 1; diff --git a/gluten-substrait/src/main/scala/org/apache/gluten/utils/SubstraitPlanPrinterUtil.scala b/gluten-substrait/src/main/scala/org/apache/gluten/utils/SubstraitPlanPrinterUtil.scala index a6ec7cb21fb..72be380691e 100644 --- a/gluten-substrait/src/main/scala/org/apache/gluten/utils/SubstraitPlanPrinterUtil.scala +++ b/gluten-substrait/src/main/scala/org/apache/gluten/utils/SubstraitPlanPrinterUtil.scala @@ -16,6 +16,8 @@ */ package org.apache.gluten.utils +import org.apache.gluten.proto.Kafka + import org.apache.spark.internal.Logging import com.google.protobuf.WrappersProto @@ -31,6 +33,9 @@ object SubstraitPlanPrinterUtil extends Logging { .newBuilder() .add(d) .add(defaultRegistry) + // Gluten's own payloads (e.g. StreamKafka) ride in Any fields of the Substrait plan, and + // nothing imports kafka.proto, so its messages are not reachable from the plan descriptor. + .add(Kafka.getDescriptor.getMessageTypes) .build() } private def MessageToJson(message: com.google.protobuf.Message): String = { diff --git a/gluten-substrait/src/main/scala/org/apache/gluten/utils/SubstraitUtil.scala b/gluten-substrait/src/main/scala/org/apache/gluten/utils/SubstraitUtil.scala index c2032aa592c..c804a9ef899 100644 --- a/gluten-substrait/src/main/scala/org/apache/gluten/utils/SubstraitUtil.scala +++ b/gluten-substrait/src/main/scala/org/apache/gluten/utils/SubstraitUtil.scala @@ -26,7 +26,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression} import org.apache.spark.sql.catalyst.plans.{ExistenceJoin, FullOuter, InnerLike, JoinType, LeftAnti, LeftOuter, LeftSemi, RightOuter} import com.google.protobuf.{Any, DoubleValue, Int32Value, Int64Value, Message, StringValue} -import io.substrait.proto.{JoinRel, NamedStruct, NestedLoopJoinRel, Type} +import io.substrait.proto.{JoinRel, NamedStruct, NestedLoopJoinRel, ReadRel, Type} import java.lang.{Double => JDouble, Long => JLong} import java.util.{Collections, List => JList} @@ -80,6 +80,14 @@ object SubstraitUtil { TypeBuilder.makeStruct(false, inputTypeNodes.asJava).toProtobuf) } + /** Wrap a Gluten payload as the detail of a Substrait `ReadRel.ExtensionTable`. */ + def packExtensionTable(detail: Message): ReadRel.ExtensionTable = { + ReadRel.ExtensionTable + .newBuilder() + .setDetail(BackendsApiManager.getTransformerApiInstance.packPBMessage(detail)) + .build() + } + def toSubstraitExpression( expr: Expression, attributeSeq: Seq[Attribute], diff --git a/gluten-substrait/src/test/scala/org/apache/gluten/substrait/rel/ReadRelProtoSuite.scala b/gluten-substrait/src/test/scala/org/apache/gluten/substrait/rel/ReadRelProtoSuite.scala new file mode 100644 index 00000000000..18602b92cfd --- /dev/null +++ b/gluten-substrait/src/test/scala/org/apache/gluten/substrait/rel/ReadRelProtoSuite.scala @@ -0,0 +1,60 @@ +/* + * 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. + */ +package org.apache.gluten.substrait.rel + +import com.google.protobuf.Descriptors.Descriptor +import io.substrait.proto.ReadRel +import org.scalatest.funsuite.AnyFunSuite + +/** + * Pins the wire tags of the vendored `ReadRel.read_type` oneof to upstream Substrait v0.98.0, + * including the official `iceberg_table = 9`. + * + * Producer and consumer share one schema, so a renumber round-trips cleanly through the generated + * classes and cannot be caught by exercising them; these assert on the descriptors instead. + */ +class ReadRelProtoSuite extends AnyFunSuite { + + private def assertFieldNumbers(descriptor: Descriptor, expected: (String, Int)*): Unit = + expected.foreach { + case (name, number) => + val field = descriptor.findFieldByName(name) + assert(field != null, s"${descriptor.getName} has no field named $name") + assert(field.getNumber === number, s"${descriptor.getName} field $name changed its number") + } + + test("ReadRel.read_type field numbers match upstream v0.98.0 verbatim") { + assertFieldNumbers( + ReadRel.getDescriptor, + "virtual_table" -> 5, + "local_files" -> 6, + "named_table" -> 7, + "extension_table" -> 8, + // Official Substrait 0.98 addition. + "iceberg_table" -> 9 + ) + } + + test("ReadRel.IcebergTable structure matches the vendored upstream layout") { + assertFieldNumbers(ReadRel.IcebergTable.getDescriptor, "direct" -> 1) + assertFieldNumbers( + ReadRel.IcebergTable.MetadataFileRead.getDescriptor, + "metadata_uri" -> 1, + "snapshot_id" -> 2, + "snapshot_timestamp" -> 3) + } +}