[GLUTEN-12597][CORE] Remodel Kafka read onto ReadRel.ExtensionTable (Substrait 0.98)

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());
   }
 }

cpp-ch/local-engine/Parser/RelParsers/ReadRelParser.cpp

@@ -35,6 +35,7 @@
 #include <Storages/SubstraitSource/SubstraitFileSourceStep.h>
 #include <google/protobuf/wrappers.pb.h>
 #include <rapidjson/document.h>
+#include <kafka.pb.h>
 #include <Common/BlockTypeUtils.h>
 #include <Common/DebugUtils.h>
 
@@ -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<gluten::StreamKafka>();
 }
 
 DB::QueryPlanStepPtr ReadRelParser::parseReadRelWithJavaIter(const substrait::ReadRel & rel)

cpp-ch/local-engine/Parser/RelParsers/StreamKafkaRelParser.cpp

@@ -20,7 +20,9 @@
 #include <Parser/SubstraitParserUtils.h>
 #include <Parser/TypeParser.h>
 #include <Storages/Kafka/ReadFromGlutenStorageKafka.h>
+#include <kafka.pb.h>
 #include <Common/BlockTypeUtils.h>
+#include <Common/DebugUtils.h>
 #include <Common/logger_useful.h>
 
 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<substrait::ReadRel::ExtensionTable>(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<substrait::ReadRel::StreamKafka>(split_info);
     auto topic = kafka_task.topic_partition().topic();
     auto partition = kafka_task.topic_partition().partition();
     auto start_offset = kafka_task.start_offset();

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());
         }
     }

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})

cpp-ch/local-engine/proto/kafka.proto

@@ -0,0 +1 @@
+../../../gluten-core/src/main/resources/org/apache/gluten/proto/kafka.proto

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

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<string, string> params = 4;
+  int64 poll_timeout_ms = 5;
+  bool fail_on_data_loss = 6;
+  bool include_headers = 7;
+}

gluten-kafka/pom.xml

@@ -52,6 +52,12 @@
       <version>${spark.version}</version>
       <scope>provided</scope>
     </dependency>
+    <dependency>
+      <groupId>com.google.protobuf</groupId>
+      <artifactId>protobuf-java</artifactId>
+      <version>${protobuf.version}</version>
+      <scope>provided</scope>
+    </dependency>
 
     <!-- For test -->
     <dependency>
@@ -88,12 +94,6 @@
       <version>${hadoop.version}</version>
       <scope>test</scope>
     </dependency>
-    <dependency>
-      <groupId>com.google.protobuf</groupId>
-      <artifactId>protobuf-java</artifactId>
-      <version>${protobuf.version}</version>
-      <scope>test</scope>
-    </dependency>
     <dependency>
       <groupId>org.scalatest</groupId>
       <artifactId>scalatest_${scala.binary.version}</artifactId>

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
   }
 }

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<ColumnTypeNode> 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());

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<String> 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());
   }
 }

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/<ver>-<uuid>.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<string, string> 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;

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 = {

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],