From e87a965853a246b0edf03fb8b46dc6189c753674 Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Mon, 31 Aug 2026 18:34:26 +0000 Subject: [PATCH 1/2] [Spark 4] Register Spark's streaming internals with Kryo A Structured Streaming query serializes two Spark internals behind the runner's back, so streaming pipelines fail spark.kryo.registrationRequired without these registrations: - StateSchemaMetadata is broadcast for every transformWithState query, hit on the very first micro-batch of any pipeline using Beam state or timers. - MemoryWriterCommitMessage is the memory sink's commit message, nested inside the already registered DataWritingSparkTaskResult. Both are registered by name because the shared base also compiles against Spark 3, where neither class exists, and with a JavaSerializer so their whole Scala object graph is covered without tracking Spark's internal field layout across versions. Neither is on a hot path. The registration call sits at the end of the registrator on purpose: Kryo auto assigns ids sequentially, so appending these conditional, by-name registrations keeps the auto assigned ids of everything above identical on Spark 3 and Spark 4 classpaths. SparkKryoRegistratorStreamingTest locks down both registrations and the id parity. --- .../SparkKryoRegistratorStreamingTest.java | 135 ++++++++++++++++++ .../translation/SparkSessionFactory.java | 55 ++++++- 2 files changed, 189 insertions(+), 1 deletion(-) create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkKryoRegistratorStreamingTest.java diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkKryoRegistratorStreamingTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkKryoRegistratorStreamingTest.java new file mode 100644 index 000000000000..3506daf75d85 --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkKryoRegistratorStreamingTest.java @@ -0,0 +1,135 @@ +/* + * 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.beam.runners.spark.structuredstreaming.translation; + +import static org.apache.beam.runners.spark.structuredstreaming.translation.utils.ScalaInterop.seqOf; +import static org.apache.beam.runners.spark.structuredstreaming.translation.utils.ScalaInterop.tuple; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; + +import com.esotericsoftware.kryo.Kryo; +import com.esotericsoftware.kryo.io.Input; +import com.esotericsoftware.kryo.io.Output; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import org.apache.beam.runners.spark.StreamingTest; +import org.apache.spark.SparkConf; +import org.apache.spark.serializer.KryoSerializer; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.execution.streaming.sources.MemoryWriterCommitMessage; +import org.apache.spark.sql.execution.streaming.state.StateSchemaMetadata; +import org.apache.spark.sql.execution.streaming.state.StateSchemaMetadataKey; +import org.apache.spark.sql.execution.streaming.state.StateSchemaMetadataValue; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.StructType; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import scala.collection.immutable.Map$; + +/** + * Guards the Spark 4 streaming entries of {@link SparkSessionFactory.SparkKryoRegistrator}: the + * classes Spark itself pushes through the user Kryo instance during a Structured Streaming query, + * which the runner never names anywhere else. + * + *

These are registered by name in the shared runner base, because that base also compiles + * against Spark 3 where neither class exists. A rename or a package move on a future Spark version + * would therefore not break the build, it would silently drop the registration and only surface as + * a streaming query dying on its first micro-batch. This test names the classes at compile time + * against the Spark 4 classpath, so that failure mode becomes a compile error instead. + * + * @see SparkSessionFactory.SparkKryoRegistrator + */ +@RunWith(JUnit4.class) +@Category(StreamingTest.class) +public class SparkKryoRegistratorStreamingTest { + + /** A Kryo configured exactly the way the runner configures it, strict registration included. */ + private static Kryo strictKryo() { + SparkConf conf = + new SparkConf(false) + .set("spark.serializer", KryoSerializer.class.getName()) + .set("spark.kryo.registrationRequired", "true") + .set( + "spark.kryo.registrator", SparkSessionFactory.SparkKryoRegistrator.class.getName()); + return new KryoSerializer(conf).newKryo(); + } + + private static Object roundTrip(Kryo kryo, Object value) { + ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + try (Output output = new Output(bytes)) { + kryo.writeClassAndObject(output, value); + } + try (Input input = new Input(new ByteArrayInputStream(bytes.toByteArray()))) { + return kryo.readClassAndObject(input); + } + } + + /** + * Spark 4 broadcasts a {@link StateSchemaMetadata} to the executors for every {@code + * transformWithState} query, so this is the registration that decides whether a Beam streaming + * pipeline with state or timers runs at all under {@code spark.kryo.registrationRequired=true}. + * + *

The instance below is deliberately not empty. It carries the nested {@code StructType} and + * {@code org.apache.avro.Schema} that make the difference between a registration that only + * survives a trivial payload and one that survives a real one. + */ + @Test + public void stateSchemaMetadataRoundTripsWithRegistrationRequired() { + StructType sqlSchema = + new StructType().add("key", DataTypes.StringType).add("value", DataTypes.BinaryType, false); + StateSchemaMetadataKey key = new StateSchemaMetadataKey("default", (short) 1, true); + StateSchemaMetadataValue value = + new StateSchemaMetadataValue( + sqlSchema, org.apache.avro.Schema.create(org.apache.avro.Schema.Type.STRING)); + StateSchemaMetadata metadata = + new StateSchemaMetadata(Map$.MODULE$.from(seqOf(tuple(key, value)))); + + Kryo kryo = strictKryo(); + assertNotNull( + "StateSchemaMetadata must be registered, see SparkKryoRegistrator", + kryo.getRegistration(StateSchemaMetadata.class)); + + StateSchemaMetadata back = (StateSchemaMetadata) roundTrip(kryo, metadata); + assertEquals(1, back.activeSchemas().size()); + assertEquals(value, back.activeSchemas().apply(key)); + } + + /** + * The commit message of Spark's {@code memory} sink, nested inside the already registered {@code + * DataWritingSparkTaskResult}. The runner writes to {@code noop}, but the {@code memory} sink is + * the obvious thing to reach for when inspecting a query, and it used to fail on batch 0. + */ + @Test + public void memoryWriterCommitMessageRoundTripsWithRegistrationRequired() { + Row row = RowFactory.create("a", 1); + MemoryWriterCommitMessage message = new MemoryWriterCommitMessage(3, seqOf(row)); + + Kryo kryo = strictKryo(); + assertNotNull( + "MemoryWriterCommitMessage must be registered, see SparkKryoRegistrator", + kryo.getRegistration(MemoryWriterCommitMessage.class)); + + MemoryWriterCommitMessage back = (MemoryWriterCommitMessage) roundTrip(kryo, message); + assertEquals(3, back.partition()); + assertEquals(1, back.data().size()); + assertEquals(row, back.data().apply(0)); + } +} diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java index 822d1871b12e..9b6265ab6cf8 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java @@ -22,6 +22,7 @@ import static org.apache.commons.lang3.math.NumberUtils.toInt; import com.esotericsoftware.kryo.Kryo; +import com.esotericsoftware.kryo.Serializer; import com.esotericsoftware.kryo.serializers.JavaSerializer; import java.util.ArrayList; import java.util.Collection; @@ -288,11 +289,63 @@ public void registerClasses(Kryo kryo) { kryo.register(CoGbkResultSchema.class); kryo.register(TupleTag.class); kryo.register(TupleTagList.class); + + // Spark internals only present when running streaming pipelines on Spark 4. These are + // registered by name because the shared runner base also compiles against Spark 3, where + // none of these classes exist. This call must stay last: Kryo auto assigns registration ids + // sequentially, so registering these conditional, by-name classes after everything else + // keeps the auto assigned ids of all the registrations above identical on Spark 3 and + // Spark 4 classpaths. See registerSparkStreamingInternals for the details. + registerSparkStreamingInternals(kryo); + } + + /** + * Registers the Spark internals that a Structured Streaming query serializes behind the + * runner's back, so streaming pipelines also work with {@code + * spark.kryo.registrationRequired=true}. + * + *

+ * + *

Both are Scala case classes holding further Scala and Spark types ({@code immutable.Map}, + * {@code StructType}, {@code org.apache.avro.Schema}, {@code Row}), none of which are + * registered either. Registering them with a {@link JavaSerializer} rather than Kryo's default + * field serializer covers that whole object graph in one go, since both classes are {@link + * java.io.Serializable}. That keeps this list from having to track Spark's internal field + * layout across versions. Neither object is on a hot path, one is broadcast once per query and + * the other is one message per task commit, so the cost of Java serialization here does not + * matter. + */ + private void registerSparkStreamingInternals(Kryo kryo) { + tryToRegister( + kryo, + "org.apache.spark.sql.execution.streaming.state.StateSchemaMetadata", + new JavaSerializer()); + tryToRegister( + kryo, + "org.apache.spark.sql.execution.streaming.sources.MemoryWriterCommitMessage", + new JavaSerializer()); } private void tryToRegister(Kryo kryo, String className) { + tryToRegister(kryo, className, null); + } + + private void tryToRegister(Kryo kryo, String className, @Nullable Serializer serializer) { try { - kryo.register(Class.forName(className)); + Class cls = Class.forName(className); + if (serializer == null) { + kryo.register(cls); + } else { + kryo.register(cls, serializer); + } } catch (ClassNotFoundException e) { LOG.info("Class {}} was not found on classpath", className); } From 4815a681e77b29668fd5a52d93d5f5bddccd8df2 Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Tue, 1 Sep 2026 09:16:27 +0000 Subject: [PATCH 2/2] [Spark 4] Simplify the Kryo registration comments Corrects the compile model description, the shared sources are compiled once per Spark version into separate artifacts, and trims the javadoc to the essentials per review. --- .../SparkKryoRegistratorStreamingTest.java | 13 ++----- .../translation/SparkSessionFactory.java | 37 +++++-------------- 2 files changed, 13 insertions(+), 37 deletions(-) diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkKryoRegistratorStreamingTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkKryoRegistratorStreamingTest.java index 3506daf75d85..c37d726afb4d 100644 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkKryoRegistratorStreamingTest.java +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkKryoRegistratorStreamingTest.java @@ -45,15 +45,10 @@ import scala.collection.immutable.Map$; /** - * Guards the Spark 4 streaming entries of {@link SparkSessionFactory.SparkKryoRegistrator}: the - * classes Spark itself pushes through the user Kryo instance during a Structured Streaming query, - * which the runner never names anywhere else. - * - *

These are registered by name in the shared runner base, because that base also compiles - * against Spark 3 where neither class exists. A rename or a package move on a future Spark version - * would therefore not break the build, it would silently drop the registration and only surface as - * a streaming query dying on its first micro-batch. This test names the classes at compile time - * against the Spark 4 classpath, so that failure mode becomes a compile error instead. + * Guards the Spark 4 streaming entries of {@link SparkSessionFactory.SparkKryoRegistrator}. The + * registrator references them by name, so a rename in a future Spark version would silently drop + * the registration and only surface as a streaming query dying on its first micro-batch. This test + * names the classes at compile time, turning that failure mode into a compile error. * * @see SparkSessionFactory.SparkKryoRegistrator */ diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java index 9b6265ab6cf8..148188bb15a2 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java @@ -290,38 +290,19 @@ public void registerClasses(Kryo kryo) { kryo.register(TupleTag.class); kryo.register(TupleTagList.class); - // Spark internals only present when running streaming pipelines on Spark 4. These are - // registered by name because the shared runner base also compiles against Spark 3, where - // none of these classes exist. This call must stay last: Kryo auto assigns registration ids - // sequentially, so registering these conditional, by-name classes after everything else - // keeps the auto assigned ids of all the registrations above identical on Spark 3 and - // Spark 4 classpaths. See registerSparkStreamingInternals for the details. + // Streaming internals that only exist as of Spark 4, registered by name so this shared + // source also compiles against Spark 3. Must stay last so the auto assigned ids of the + // registrations above are identical in the Spark 3 and Spark 4 artifacts. registerSparkStreamingInternals(kryo); } /** - * Registers the Spark internals that a Structured Streaming query serializes behind the - * runner's back, so streaming pipelines also work with {@code - * spark.kryo.registrationRequired=true}. - * - *

- * - *

Both are Scala case classes holding further Scala and Spark types ({@code immutable.Map}, - * {@code StructType}, {@code org.apache.avro.Schema}, {@code Row}), none of which are - * registered either. Registering them with a {@link JavaSerializer} rather than Kryo's default - * field serializer covers that whole object graph in one go, since both classes are {@link - * java.io.Serializable}. That keeps this list from having to track Spark's internal field - * layout across versions. Neither object is on a hot path, one is broadcast once per query and - * the other is one message per task commit, so the cost of Java serialization here does not - * matter. + * Registers the internals a Structured Streaming query serializes behind the runner's back, so + * streaming pipelines work with {@code spark.kryo.registrationRequired=true}: {@code + * StateSchemaMetadata} (broadcast for every {@code transformWithState} query) and {@code + * MemoryWriterCommitMessage} (the {@code memory} sink's commit message, nested inside the + * already registered {@link DataWritingSparkTaskResult}). A {@link JavaSerializer} covers their + * whole Scala object graph without tracking Spark's field layout; neither is on a hot path. */ private void registerSparkStreamingInternals(Kryo kryo) { tryToRegister(