This is an automated email from the ASF dual-hosted git repository.
zhouyuan pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 2496e4c083 [VL] Offload Kafka micro-batch scans (#13237)
2496e4c083 is described below
commit 2496e4c083eebdd86a5a9ecd175511b19463c87d
Author: Reema <[email protected]>
AuthorDate: Thu Nov 5 15:19:47 2026 +0300
[VL] Offload Kafka micro-batch scans (#13237)
---
.github/workflows/velox_kafka.yml | 24 ++--
backends-velox/pom.xml | 19 +++
...org.apache.gluten.component.VeloxKafkaComponent | 0
.../gluten/component/VeloxKafkaComponent.scala | 37 ++++++
.../kafka/VeloxGlutenKafkaScanSuite.scala | 129 +++++++++++++++++++++
.../gluten/backendsapi/velox/VeloxBackend.scala | 19 ++-
.../execution/VeloxKafkaScanValidationSuite.scala | 66 +++++++++++
7 files changed, 281 insertions(+), 13 deletions(-)
diff --git a/.github/workflows/velox_kafka.yml
b/.github/workflows/velox_kafka.yml
index b900636493..c4821a34a6 100644
--- a/.github/workflows/velox_kafka.yml
+++ b/.github/workflows/velox_kafka.yml
@@ -141,18 +141,18 @@ jobs:
run: |
./build/mvn -B -ntp install -Pbackends-velox -Pspark-3.5
-Pscala-2.12 \
-Pjava-17 -Pkafka -Pspark-ut -Pfast-build -DskipTests
- # TODO: Re-enable when VeloxKafkaScanValidationSuite is added.
- # - name: Check Kafka schema validation without a broker
- # run: |
- # ./build/mvn -B -ntp test -pl gluten-ut/spark35 -Pbackends-velox
-Pspark-3.5 \
- # -Pscala-2.12 -Pjava-17 -Pspark-ut -Pfast-build \
- #
-Dsuites=org.apache.gluten.execution.VeloxKafkaScanValidationSuite
- # TODO: Re-enable when VeloxGlutenKafkaScanSuite is added in #13237.
- # - name: Check native micro-batch reads and headers fallback
- # run: |
- # ./build/mvn -B -ntp test -pl backends-velox -Pbackends-velox
-Pspark-3.5 \
- # -Pscala-2.12 -Pjava-17 -Pkafka -Pfast-build \
- #
-Dsuites=org.apache.gluten.execution.kafka.VeloxGlutenKafkaScanSuite
+
+ - name: Check Kafka schema validation without a broker
+ run: |
+ ./build/mvn -B -ntp test -pl gluten-ut/spark35 -Pbackends-velox
-Pspark-3.5 \
+ -Pscala-2.12 -Pjava-17 -Pspark-ut -Pfast-build \
+ -Dsuites=org.apache.gluten.execution.VeloxKafkaScanValidationSuite
+
+ - name: Check native micro-batch reads and headers fallback
+ run: |
+ ./build/mvn -B -ntp test -pl backends-velox -Pbackends-velox
-Pspark-3.5 \
+ -Pscala-2.12 -Pjava-17 -Pkafka -Pfast-build \
+
-Dsuites=org.apache.gluten.execution.kafka.VeloxGlutenKafkaScanSuite
- name: Upload test results
if: always()
uses: actions/upload-artifact@v7
diff --git a/backends-velox/pom.xml b/backends-velox/pom.xml
index 028dc79a44..fff0f1fdfe 100644
--- a/backends-velox/pom.xml
+++ b/backends-velox/pom.xml
@@ -878,5 +878,24 @@
</dependency>
</dependencies>
</profile>
+ <profile>
+ <id>kafka</id>
+ <activation>
+ <activeByDefault>false</activeByDefault>
+ </activation>
+ <dependencies>
+ <dependency>
+ <groupId>org.apache.gluten</groupId>
+ <artifactId>gluten-kafka</artifactId>
+ <version>${project.version}</version>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-sql-kafka-0-10_${scala.binary.version}</artifactId>
+ <version>${spark.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ </dependencies>
+ </profile>
</profiles>
</project>
diff --git
a/backends-velox/src-kafka/main/resources/META-INF/gluten-components/org.apache.gluten.component.VeloxKafkaComponent
b/backends-velox/src-kafka/main/resources/META-INF/gluten-components/org.apache.gluten.component.VeloxKafkaComponent
new file mode 100644
index 0000000000..e69de29bb2
diff --git
a/backends-velox/src-kafka/main/scala/org/apache/gluten/component/VeloxKafkaComponent.scala
b/backends-velox/src-kafka/main/scala/org/apache/gluten/component/VeloxKafkaComponent.scala
new file mode 100644
index 0000000000..e2f46946d8
--- /dev/null
+++
b/backends-velox/src-kafka/main/scala/org/apache/gluten/component/VeloxKafkaComponent.scala
@@ -0,0 +1,37 @@
+/*
+ * 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.component
+
+import org.apache.gluten.backendsapi.velox.VeloxBackend
+import org.apache.gluten.execution.OffloadKafkaScan
+import org.apache.gluten.extension.injector.Injector
+
+import org.apache.spark.util.SparkReflectionUtil
+
+class VeloxKafkaComponent extends Component {
+ override def name(): String = "velox-kafka"
+
+ override def dependencies(): Seq[Class[_ <: Component]] =
classOf[VeloxBackend] :: Nil
+
+ override def isRuntimeCompatible: Boolean = {
+
SparkReflectionUtil.isClassPresent("org.apache.spark.sql.kafka010.KafkaSourceProvider")
+ }
+
+ override def injectRules(injector: Injector): Unit = {
+ OffloadKafkaScan.inject(injector)
+ }
+}
diff --git
a/backends-velox/src-kafka/test/scala/org/apache/gluten/execution/kafka/VeloxGlutenKafkaScanSuite.scala
b/backends-velox/src-kafka/test/scala/org/apache/gluten/execution/kafka/VeloxGlutenKafkaScanSuite.scala
new file mode 100644
index 0000000000..f9f20ceaea
--- /dev/null
+++
b/backends-velox/src-kafka/test/scala/org/apache/gluten/execution/kafka/VeloxGlutenKafkaScanSuite.scala
@@ -0,0 +1,129 @@
+/*
+ * 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.execution.kafka
+
+import org.apache.gluten.execution.{MicroBatchScanExecTransformer,
VeloxWholeStageTransformerSuite}
+
+import org.apache.spark.SparkConf
+import org.apache.spark.sql.Row
+import org.apache.spark.sql.execution.streaming.StreamingQueryWrapper
+import org.apache.spark.sql.streaming.Trigger
+
+import org.apache.kafka.clients.admin.{AdminClient, AdminClientConfig,
NewTopic}
+
+import java.time.Duration
+import java.util.{Collections, Properties, UUID}
+import java.util.concurrent.TimeUnit
+
+/**
+ * Requires ENABLE_KAFKA in the native build and a broker at
KAFKA_BOOTSTRAP_SERVERS (default:
+ * localhost:9092). Run with the kafka Maven profile.
+ */
+class VeloxGlutenKafkaScanSuite extends VeloxWholeStageTransformerSuite {
+ override protected val resourcePath: String = "/tpch-data-parquet"
+ override protected val fileFormat: String = "parquet"
+
+ private val kafkaBootstrapServers =
+ sys.env.getOrElse("KAFKA_BOOTSTRAP_SERVERS", "localhost:9092")
+
+ override protected def sparkConf: SparkConf = {
+ super.sparkConf
+ .set("spark.shuffle.manager",
"org.apache.spark.shuffle.sort.ColumnarShuffleManager")
+ .set("spark.sql.shuffle.partitions", "2")
+ }
+
+ private def withTopic(func: String => Unit): Unit = {
+ val topic = "velox_kafka_" + UUID.randomUUID().toString.replace("-", "")
+ val props = new Properties()
+ props.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG,
kafkaBootstrapServers)
+ props.put(AdminClientConfig.DEFAULT_API_TIMEOUT_MS_CONFIG, "30000")
+ val adminClient = AdminClient.create(props)
+ try {
+ adminClient
+ .createTopics(Collections.singletonList(new NewTopic(topic, 2,
1.toShort)))
+ .all()
+ .get(30, TimeUnit.SECONDS)
+ try {
+ func(topic)
+ } finally {
+
adminClient.deleteTopics(Collections.singletonList(topic)).all().get(30,
TimeUnit.SECONDS)
+ }
+ } finally {
+ adminClient.close(Duration.ofSeconds(5))
+ }
+ }
+
+ private def checkStreamingRead(includeHeaders: Boolean): Unit = {
+ withTopic {
+ topic =>
+ withTempDir {
+ dir =>
+ spark
+ .range(1000)
+ .selectExpr("cast(id as string) as value")
+ .write
+ .format("kafka")
+ .option("kafka.bootstrap.servers", kafkaBootstrapServers)
+ .option("topic", topic)
+ .save()
+
+ val stream = spark.readStream
+ .format("kafka")
+ .option("kafka.bootstrap.servers", kafkaBootstrapServers)
+ .option("subscribe", topic)
+ .option("startingOffsets", "earliest")
+ .option("includeHeaders", includeHeaders.toString)
+ .load()
+ // Keep headers in the fallback query so column pruning cannot
hide the unsupported type.
+ val result = if (includeHeaders) {
+ stream.selectExpr("cast(cast(value as string) as int) as id",
"headers")
+ } else {
+ stream.selectExpr("cast(cast(value as string) as int) as id")
+ }
+ val streamQuery = result.writeStream
+ .format("parquet")
+ .option("checkpointLocation", dir.getCanonicalPath +
"/checkpoint")
+ .trigger(Trigger.AvailableNow())
+ .start(dir.getCanonicalPath + "/data")
+
+ try {
+ assert(streamQuery.awaitTermination(60000), "Kafka micro-batch
did not finish")
+ val scans = streamQuery
+ .asInstanceOf[StreamingQueryWrapper]
+ .streamingQuery
+ .lastExecution
+ .executedPlan
+ .collect { case p: MicroBatchScanExecTransformer => p }
+ assert(scans.size == (if (includeHeaders) 0 else 1))
+ checkAnswer(
+ spark.read.parquet(dir.getCanonicalPath +
"/data").select("id"),
+ (0 until 1000).map(Row(_)))
+ } finally {
+ streamQuery.stop()
+ }
+ }
+ }
+ }
+
+ test("kafka micro-batch scan offloads and returns every record") {
+ checkStreamingRead(includeHeaders = false)
+ }
+
+ test("kafka micro-batch scan with headers falls back and returns every
record") {
+ checkStreamingRead(includeHeaders = true)
+ }
+}
diff --git
a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxBackend.scala
b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxBackend.scala
index 5426471f37..2e75321156 100644
---
a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxBackend.scala
+++
b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxBackend.scala
@@ -27,7 +27,7 @@ import
org.apache.gluten.extension.columnar.cost.{LegacyCoster, LongCoster, Roug
import org.apache.gluten.extension.columnar.transition.{Convention,
ConventionFunc}
import org.apache.gluten.substrait.rel.LocalFilesNode
import org.apache.gluten.substrait.rel.LocalFilesNode.ReadFileFormat
-import
org.apache.gluten.substrait.rel.LocalFilesNode.ReadFileFormat.{DwrfReadFormat,
OrcReadFormat, ParquetReadFormat}
+import
org.apache.gluten.substrait.rel.LocalFilesNode.ReadFileFormat.{DwrfReadFormat,
KafkaReadFormat, OrcReadFormat, ParquetReadFormat}
import org.apache.gluten.utils._
import org.apache.spark.sql.catalyst.catalog.BucketSpec
@@ -48,6 +48,8 @@ import org.apache.spark.sql.types._
import org.apache.hadoop.conf.Configuration
import org.apache.hadoop.fs.Path
+import java.util.Locale
+
import scala.util.control.Breaks.breakable
class VeloxBackend extends SubstraitBackend {
@@ -100,6 +102,17 @@ object VeloxBackendSettings extends BackendSettingsApi {
// unit: SECONDS, default 1 day
val GLUTEN_VELOX_BROADCAST_CACHE_EXPIRED_TIME_DEFAULT: Int = 86400
+ // Columns produced by the native Kafka reader. Kafka headers are not
supported.
+ private val kafkaColumns = Map(
+ "key" -> BinaryType,
+ "value" -> BinaryType,
+ "topic" -> StringType,
+ "partition" -> IntegerType,
+ "offset" -> LongType,
+ "timestamp" -> TimestampType,
+ "timestamptype" -> IntegerType
+ )
+
override def primaryBatchType: Convention.BatchType = VeloxBatchType
override def supportTimestampNtz: Boolean = true
@@ -188,6 +201,10 @@ object VeloxBackendSettings extends BackendSettingsApi {
validateTypes(fieldTypeValidator, fields)
.orElse(validateTypes(schemaTypeValidator, dataSchema.fields))
}
+ case KafkaReadFormat =>
+ fields
+ .find(f =>
!kafkaColumns.get(f.name.toLowerCase(Locale.ROOT)).contains(f.dataType))
+ .map(f => s"Unsupported Kafka column ${f.name}:
${f.dataType.simpleString}.")
case _ => Some(s"Unsupported file format $format.")
}
}
diff --git
a/gluten-ut/spark35/src/test/backends-velox/org/apache/gluten/execution/VeloxKafkaScanValidationSuite.scala
b/gluten-ut/spark35/src/test/backends-velox/org/apache/gluten/execution/VeloxKafkaScanValidationSuite.scala
new file mode 100644
index 0000000000..6bd56f608f
--- /dev/null
+++
b/gluten-ut/spark35/src/test/backends-velox/org/apache/gluten/execution/VeloxKafkaScanValidationSuite.scala
@@ -0,0 +1,66 @@
+/*
+ * 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.execution
+
+import org.apache.gluten.backendsapi.velox.VeloxBackendSettings
+import
org.apache.gluten.substrait.rel.LocalFilesNode.ReadFileFormat.KafkaReadFormat
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.types._
+
+import org.apache.hadoop.conf.Configuration
+
+/** Schema validation must also work without Kafka jars, a broker, or a
Kafka-enabled native build. */
+class VeloxKafkaScanValidationSuite extends SparkFunSuite {
+ private val kafkaFields = Array(
+ StructField("key", BinaryType),
+ StructField("value", BinaryType),
+ StructField("topic", StringType),
+ StructField("partition", IntegerType),
+ StructField("offset", LongType),
+ StructField("timestamp", TimestampType),
+ StructField("timestampType", IntegerType))
+
+ private def validate(fields: Array[StructField]): ValidationResult = {
+ VeloxBackendSettings.validateScanExec(
+ KafkaReadFormat,
+ fields,
+ StructType(fields),
+ Seq.empty,
+ Map.empty,
+ new Configuration(false),
+ Set(KafkaReadFormat))
+ }
+
+ test("native Kafka scan accepts the Spark Kafka schema and pruned columns") {
+ assert(validate(kafkaFields).ok())
+ assert(validate(Array(kafkaFields(4), kafkaFields(1))).ok())
+ assert(validate(Array.empty[StructField]).ok())
+ }
+
+ test("native Kafka scan rejects headers, unknown columns, and incompatible
types") {
+ val headers = StructField(
+ "headers",
+ ArrayType(StructType(Seq(StructField("key", StringType),
StructField("value", BinaryType)))))
+ Seq(headers, StructField("unknown", StringType), StructField("value",
StringType)).foreach {
+ field =>
+ val result = validate(kafkaFields :+ field)
+ assert(!result.ok())
+ assert(result.reason().contains(s"Unsupported Kafka column
${field.name}"))
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
