From 17936ca3a3f6a50994adbff448f958efd76d907d Mon Sep 17 00:00:00 2001 From: Niels Pardon Date: Thu, 20 Aug 2026 07:48:53 +0200 Subject: [PATCH 1/2] [GLUTEN-12597][CORE] Migrate ReadRel read_type to Substrait 0.98 (add iceberg_table, relocate stream_kafka) Substrait 0.98 added `IcebergTable iceberg_table = 9` to the `ReadRel.read_type` oneof, exactly where Gluten's local `stream_kafka` graft sat. This vendors the 0.98 `iceberg_table` field and `IcebergTable` message verbatim and relocates the `stream_kafka` graft off the collision, as one step of the Substrait v0.23.0 -> 0.98.0 proto rebase (#12597). The graft moves to field 1000, following the "Gluten-local fields start at 1000" convention established for WriteRel's bucket_spec (#12746); it stays inside the read_type oneof. The enclosing Rel.read oneof tag is unchanged. All accessors are name-based (setStreamKafka/hasStreamKafka), so the field renumber needs no source change, and the new iceberg_table field is unreferenced by any producer or consumer, so no source is touched. A descriptor-level `ReadRelProtoSuite` pins the `read_type` field numbers (and the vendored `IcebergTable` layout), since a renumber round-trips cleanly through the shared schema and would otherwise be invisible to tests. This is the first of three ReadRel slices (read_type / text options / VirtualTable); the other two are separate follow-ups. Part of #12597 Generated-by: Claude Code (Claude Opus 4.8) --- .../substrait/proto/substrait/algebra.proto | 29 ++++++++- .../substrait/rel/ReadRelProtoSuite.scala | 64 +++++++++++++++++++ 2 files changed, 92 insertions(+), 1 deletion(-) create mode 100644 gluten-substrait/src/test/scala/org/apache/gluten/substrait/rel/ReadRelProtoSuite.scala 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 e239691208..d590ba2dc0 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,11 @@ message ReadRel { LocalFiles local_files = 6; NamedTable named_table = 7; ExtensionTable extension_table = 8; - bool stream_kafka = 9; + IcebergTable iceberg_table = 9; + // Gluten addition: streaming Kafka source. Relocated from field 9 to + // field 1000 so the official Substrait iceberg_table can occupy its + // 0.98 slot. + bool stream_kafka = 1000; } // A base table. The list of string is used to represent namespacing (e.g., mydb.mytable). @@ -79,6 +83,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; 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 0000000000..0f458692ee --- /dev/null +++ b/gluten-substrait/src/test/scala/org/apache/gluten/substrait/rel/ReadRelProtoSuite.scala @@ -0,0 +1,64 @@ +/* + * 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 after rebasing it onto upstream + * Substrait v0.98.0: adding the official `iceberg_table = 9` and relocating Gluten's `stream_kafka` + * graft off the field-9 collision into the 1000+ range. 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. See docs/developers/SubstraitModifications.md for + * the numbering convention. + */ +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 plus the relocated graft") { + assertFieldNumbers( + ReadRel.getDescriptor, + "virtual_table" -> 5, + "local_files" -> 6, + "named_table" -> 7, + "extension_table" -> 8, + // Official Substrait 0.98 addition; must own field 9. + "iceberg_table" -> 9, + // Gluten-local graft, relocated off upstream's field 9 to the 1000+ range so that + // iceberg_table can take its 0.98 slot. + "stream_kafka" -> 1000 + ) + } + + 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) + } +} From 3264a7ff72d8c12570f07446f26c6718dda89c6b Mon Sep 17 00:00:00 2001 From: Niels Pardon Date: Thu, 20 Aug 2026 17:54:25 +0200 Subject: [PATCH 2/2] [GLUTEN-12597][CORE] Remodel Kafka read onto ReadRel.ExtensionTable (Substrait 0.98) Part of #12597. Stacked on #12832 (the iceberg_table half of the ReadRel read_type migration). Makes ReadRel.read_type verbatim upstream Substrait 0.98 by removing Gluten's last graft on it -- the `bool stream_kafka = 1000` discriminator and the nested `ReadRel.StreamKafka` message. read_type now matches upstream exactly, which never allocated field 1000, so the retired tag is not reserved. Gluten's Kafka streaming read is remodeled onto the official `extension_table = 8` path, the mechanism MergeTree and Range already use. The StreamKafka payload moves to a new Gluten-owned `kafka.proto` (`package gluten`, `org.apache.gluten.proto`), packed into a `google.protobuf.Any` and carried in `ReadRel.ExtensionTable.detail`. The native consumer discriminates a Kafka read by the detail's type_url (`detail().Is()`) -- mirroring the already-merged Velox Iceberg idiom (`enhancement().Is()`). This Any-in-an-official-extension-field pattern is preferred over grafting new fields; `WriteRel.bucket_spec` (field 1000) remains the one legacy graft. Gluten plans are transient and the JAR + native library are generated from one proto source and ship together, so there is no wire-compatibility constraint. Two side effects: the Kafka split-info payload changes from a bare StreamKafka to a ReadRel.ExtensionTable wrapping it, and read_type -- previously set on every ReadRel by the unconditional stream_kafka flag -- is now left unset for non-Kafka scans (no native code reads read_type_case). A JAR and native library must therefore be rebuilt together. ClickHouse-only; Kafka has no Velox path. Generated-by: Claude Code (Claude Opus 4.8) --- .../clickhouse/ExtensionTableNode.java | 8 ++--- .../Parser/RelParsers/ReadRelParser.cpp | 6 ++-- .../RelParsers/StreamKafkaRelParser.cpp | 18 ++++++++++-- .../Parser/SerializedPlanParser.cpp | 6 ++++ cpp-ch/local-engine/proto/CMakeLists.txt | 2 +- cpp-ch/local-engine/proto/kafka.proto | 1 + docs/developers/SubstraitModifications.md | 1 - .../org/apache/gluten/proto/kafka.proto | 29 +++++++++++++++++++ gluten-kafka/pom.xml | 12 ++++---- .../MicroBatchScanExecTransformer.scala | 10 ++++++- .../gluten/substrait/rel/ReadRelNode.java | 13 ++++++--- .../substrait/rel/StreamKafkaSourceNode.java | 12 ++++---- .../substrait/proto/substrait/algebra.proto | 20 ------------- .../utils/SubstraitPlanPrinterUtil.scala | 5 ++++ .../apache/gluten/utils/SubstraitUtil.scala | 10 ++++++- .../substrait/rel/ReadRelProtoSuite.scala | 20 +++++-------- 16 files changed, 111 insertions(+), 62 deletions(-) create mode 120000 cpp-ch/local-engine/proto/kafka.proto create mode 100644 gluten-core/src/main/resources/org/apache/gluten/proto/kafka.proto 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 7e5de6eb46..dc88abe672 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 17bdbc5cfd..91e5c79d7c 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 cb45e4f3eb..2f2dc379c3 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 5ca985575d..594484adae 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 0a12af465d..f06dc0a13e 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 0000000000..dbd46d6d5d --- /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 11aa46eb1c..3849df7c04 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 0000000000..91ef725383 --- /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 eabfc1119a..20e10e4285 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 d2777f04ae..11071a0f50 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 33332f1867..5a5c8f4139 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 5974e3bb0f..3ae4844459 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 d590ba2dc0..f19015d419 100644 --- a/gluten-substrait/src/main/resources/substrait/proto/substrait/algebra.proto +++ b/gluten-substrait/src/main/resources/substrait/proto/substrait/algebra.proto @@ -70,10 +70,6 @@ message ReadRel { NamedTable named_table = 7; ExtensionTable extension_table = 8; IcebergTable iceberg_table = 9; - // Gluten addition: streaming Kafka source. Relocated from field 9 to - // field 1000 so the official Substrait iceberg_table can occupy its - // 0.98 slot. - bool stream_kafka = 1000; } // A base table. The list of string is used to represent namespacing (e.g., mydb.mytable). @@ -117,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 a6ec7cb21f..72be380691 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 c2032aa592..c804a9ef89 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 index 0f458692ee..18602b92cf 100644 --- 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 @@ -21,12 +21,11 @@ import io.substrait.proto.ReadRel import org.scalatest.funsuite.AnyFunSuite /** - * Pins the wire tags of the vendored `ReadRel.read_type` oneof after rebasing it onto upstream - * Substrait v0.98.0: adding the official `iceberg_table = 9` and relocating Gluten's `stream_kafka` - * graft off the field-9 collision into the 1000+ range. 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. See docs/developers/SubstraitModifications.md for - * the numbering convention. + * 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 { @@ -38,18 +37,15 @@ class ReadRelProtoSuite extends AnyFunSuite { assert(field.getNumber === number, s"${descriptor.getName} field $name changed its number") } - test("ReadRel.read_type field numbers match upstream v0.98.0 plus the relocated graft") { + 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; must own field 9. - "iceberg_table" -> 9, - // Gluten-local graft, relocated off upstream's field 9 to the 1000+ range so that - // iceberg_table can take its 0.98 slot. - "stream_kafka" -> 1000 + // Official Substrait 0.98 addition. + "iceberg_table" -> 9 ) }