From d70691cf31f710e5cf042457e1fe789520fe9521 Mon Sep 17 00:00:00 2001 From: Bruno Volpato Date: Tue, 4 Aug 2026 17:35:08 -0400 Subject: [PATCH] [Java] Avoid RowCoder bytecode generation --- .../sdk/jmh/coders/RowCoderBenchmark.java | 85 ++++++ .../coders/RowCoderGenerationBenchmark.java | 53 ++++ .../beam/sdk/jmh/coders/package-info.java | 20 ++ .../beam/sdk/coders/RowCoderGenerator.java | 289 ++---------------- .../apache/beam/sdk/schemas/SchemaCoder.java | 2 +- 5 files changed, 185 insertions(+), 264 deletions(-) create mode 100644 sdks/java/core/jmh/src/main/java/org/apache/beam/sdk/jmh/coders/RowCoderBenchmark.java create mode 100644 sdks/java/core/jmh/src/main/java/org/apache/beam/sdk/jmh/coders/RowCoderGenerationBenchmark.java create mode 100644 sdks/java/core/jmh/src/main/java/org/apache/beam/sdk/jmh/coders/package-info.java diff --git a/sdks/java/core/jmh/src/main/java/org/apache/beam/sdk/jmh/coders/RowCoderBenchmark.java b/sdks/java/core/jmh/src/main/java/org/apache/beam/sdk/jmh/coders/RowCoderBenchmark.java new file mode 100644 index 000000000000..44e7153df58b --- /dev/null +++ b/sdks/java/core/jmh/src/main/java/org/apache/beam/sdk/jmh/coders/RowCoderBenchmark.java @@ -0,0 +1,85 @@ +/* + * 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.sdk.jmh.coders; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.util.concurrent.TimeUnit; +import org.apache.beam.sdk.coders.RowCoder; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.schemas.Schema.FieldType; +import org.apache.beam.sdk.values.Row; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.Warmup; + +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(TimeUnit.NANOSECONDS) +@Warmup(iterations = 5, time = 500, timeUnit = TimeUnit.MILLISECONDS) +@Measurement(iterations = 10, time = 500, timeUnit = TimeUnit.MILLISECONDS) +public class RowCoderBenchmark { + @State(Scope.Thread) + public static class CoderState { + @Param({"false", "true"}) + boolean staticEncoding; + + RowCoder coder; + Row row; + ByteArrayOutputStream output; + ByteArrayInputStream input; + + @Setup(Level.Trial) + public void setup() throws Exception { + Schema.Builder builder = + Schema.builder().addByteField("_pythonsdk_any_type_byte").addByteArrayField("payload"); + if (staticEncoding) { + builder.setOptions( + Schema.Options.builder() + .setOption("beam:option:row:static_encoding", FieldType.BOOLEAN, true) + .build()); + } + Schema schema = builder.build(); + coder = RowCoder.of(schema); + row = Row.withSchema(schema).addValues((byte) 5, new byte[] {1, 2, 3, 4}).build(); + output = new ByteArrayOutputStream(64); + coder.encode(row, output); + input = new ByteArrayInputStream(output.toByteArray()); + } + } + + @Benchmark + public int encode(CoderState state) throws Exception { + state.output.reset(); + state.coder.encode(state.row, state.output); + return state.output.size(); + } + + @Benchmark + public Row decode(CoderState state) throws Exception { + state.input.reset(); + return state.coder.decode(state.input); + } +} diff --git a/sdks/java/core/jmh/src/main/java/org/apache/beam/sdk/jmh/coders/RowCoderGenerationBenchmark.java b/sdks/java/core/jmh/src/main/java/org/apache/beam/sdk/jmh/coders/RowCoderGenerationBenchmark.java new file mode 100644 index 000000000000..7de1cc7053e1 --- /dev/null +++ b/sdks/java/core/jmh/src/main/java/org/apache/beam/sdk/jmh/coders/RowCoderGenerationBenchmark.java @@ -0,0 +1,53 @@ +/* + * 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.sdk.jmh.coders; + +import java.util.concurrent.TimeUnit; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.RowCoderGenerator; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.values.Row; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Warmup; + +@BenchmarkMode(Mode.SingleShotTime) +@OutputTimeUnit(TimeUnit.MICROSECONDS) +@Warmup(iterations = 5) +@Measurement(iterations = 10) +public class RowCoderGenerationBenchmark { + @Benchmark + public Schema buildSchema() { + return newSchema(); + } + + @Benchmark + public Coder generate() { + return RowCoderGenerator.generate(newSchema()); + } + + private static Schema newSchema() { + return Schema.builder() + .addByteField("_pythonsdk_any_type_byte") + .addByteArrayField("payload") + .build(); + } +} diff --git a/sdks/java/core/jmh/src/main/java/org/apache/beam/sdk/jmh/coders/package-info.java b/sdks/java/core/jmh/src/main/java/org/apache/beam/sdk/jmh/coders/package-info.java new file mode 100644 index 000000000000..dc62de3f53ea --- /dev/null +++ b/sdks/java/core/jmh/src/main/java/org/apache/beam/sdk/jmh/coders/package-info.java @@ -0,0 +1,20 @@ +/* + * 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. + */ + +/** Benchmarks for coders. */ +package org.apache.beam.sdk.jmh.coders; diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/RowCoderGenerator.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/RowCoderGenerator.java index dc6a28fdf6b8..d73956a5e478 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/RowCoderGenerator.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/RowCoderGenerator.java @@ -17,45 +17,22 @@ */ package org.apache.beam.sdk.coders; -import static org.apache.beam.sdk.util.ByteBuddyUtils.getClassLoadingStrategy; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; -import java.lang.reflect.InvocationTargetException; -import java.lang.reflect.Modifier; import java.util.Arrays; import java.util.BitSet; import java.util.Map; import java.util.UUID; import javax.annotation.Nullable; import javax.annotation.concurrent.GuardedBy; -import net.bytebuddy.ByteBuddy; -import net.bytebuddy.description.modifier.FieldManifestation; -import net.bytebuddy.description.modifier.Ownership; -import net.bytebuddy.description.modifier.Visibility; -import net.bytebuddy.description.type.TypeDescription; -import net.bytebuddy.description.type.TypeDescription.ForLoadedType; -import net.bytebuddy.dynamic.DynamicType; -import net.bytebuddy.dynamic.scaffold.InstrumentedType; -import net.bytebuddy.implementation.FixedValue; -import net.bytebuddy.implementation.Implementation; -import net.bytebuddy.implementation.bytecode.ByteCodeAppender; -import net.bytebuddy.implementation.bytecode.ByteCodeAppender.Size; -import net.bytebuddy.implementation.bytecode.Duplication; -import net.bytebuddy.implementation.bytecode.StackManipulation; -import net.bytebuddy.implementation.bytecode.member.FieldAccess; -import net.bytebuddy.implementation.bytecode.member.MethodInvocation; -import net.bytebuddy.implementation.bytecode.member.MethodReturn; -import net.bytebuddy.implementation.bytecode.member.MethodVariableAccess; -import net.bytebuddy.matcher.ElementMatchers; import org.apache.beam.sdk.schemas.Schema; import org.apache.beam.sdk.schemas.Schema.Field; import org.apache.beam.sdk.schemas.Schema.FieldType; import org.apache.beam.sdk.schemas.SchemaCoder; import org.apache.beam.sdk.util.StringUtils; -import org.apache.beam.sdk.util.common.ReflectHelpers; import org.apache.beam.sdk.values.Row; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; @@ -63,56 +40,18 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -/** - * A utility for automatically generating a {@link Coder} for {@link Row} objects corresponding to a - * specific schema. The resulting coder is loaded into the default ClassLoader and returned. - * - *

When {@link RowCoderGenerator#generate(Schema)} is called, a new subclass of {@literal - * Coder} is generated for the specified schema. This class is generated using low-level - * bytecode generation, and hardcodes encodings for all fields of the Schema. Empirically, this is - * 30-40% faster than a coder that introspects the schema. - * - *

The generated class corresponds to the following Java class: - * - * 


- * class SchemaRowCoder extends{@literal Coder} {
- *   // Generated array containing a coder for each field in the Schema.
- *   private static final Coder[] FIELD_CODERS;
- *
- *   // Generated method to return the schema this class corresponds to. Used during code
- *   // generation.
- *   private static getSchema() {
- *     return schema;
- *   }
- *
- *  {@literal @}Override
- *   public void encode(T value, OutputStream outStream) {
- *     // Delegate to a method that evaluates each coder in the static array.
- *     encodeDelegate(FIELD_CODERS, value, outStream);
- *   }
- *
- *  {@literal @}Override
- *   public abstract T decode(InputStream inStream) {
- *     // Delegate to a method that evaluates each coder in the static array.
- *     return decodeDelegate(FIELD_CODERS, inStream);
- *   }
- * }
- *
- */ +/** Creates and caches a {@link Coder} for {@link Row} objects corresponding to a schema. */ @SuppressWarnings({ "nullness", // TODO(https://github.com/apache/beam/issues/20497) "rawtypes" }) public abstract class RowCoderGenerator { - private static final ByteBuddy BYTE_BUDDY = new ByteBuddy(); private static final BitSetCoder NULL_LIST_CODER = BitSetCoder.of(); private static final VarIntCoder VAR_INT_CODER = VarIntCoder.of(); // BitSet.get(n) will return false for any n >= nbits, so a BitSet with 0 bits will return false // for all calls to get. private static final BitSet EMPTY_BIT_SET = new BitSet(0); - private static final String CODERS_FIELD_NAME = "FIELD_CODERS"; - private static final String POSITIONS_FIELD_NAME = "FIELD_ENCODING_POSITIONS"; private static final String SCHEMA_OPTION_STATIC_ENCODING = "beam:option:row:static_encoding"; static class WithStackTrace { @@ -133,7 +72,7 @@ public String getStackTrace() { } } - // Cache for Coder class that are already generated. + // Cache for coders that are already created. @GuardedBy("cacheLock") private static final Map>> GENERATED_CODERS = Maps.newHashMap(); @@ -198,7 +137,6 @@ static void clearRowCoderCache() { } } - @SuppressWarnings("unchecked") public static Coder generate(Schema schema) { UUID uuid = Preconditions.checkNotNull(schema.getUUID()); // Avoid using computeIfAbsent which may cause issues with nested schemas. @@ -207,12 +145,6 @@ public static Coder generate(Schema schema) { if (existingRowCoder != null) { return existingRowCoder.getValue(); } - TypeDescription.Generic coderType = - TypeDescription.Generic.Builder.parameterizedType(Coder.class, Row.class).build(); - DynamicType.Builder builder = - (DynamicType.Builder) BYTE_BUDDY.subclass(coderType); - builder = implementMethods(schema, builder); - int[] encodingPosToRowIndex = new int[schema.getFieldCount()]; @Nullable WithStackTrace> existingEncodingPositions = @@ -241,33 +173,12 @@ public static Coder generate(Schema schema) { SchemaCoder.coderForFieldType(schema.getField(rowIndex).getType().withNullable(false)); } - builder = - builder - .defineField( - CODERS_FIELD_NAME, Coder[].class, Visibility.PRIVATE, FieldManifestation.FINAL) - .defineField( - POSITIONS_FIELD_NAME, int[].class, Visibility.PRIVATE, FieldManifestation.FINAL) - .defineConstructor(Modifier.PUBLIC) - .withParameters(Coder[].class, int[].class) - .intercept(new GeneratedCoderConstructor()); - - Coder rowCoder; - try { - rowCoder = - builder - .make() - .load( - ReflectHelpers.findClassLoader(Coder.class.getClassLoader()), - getClassLoadingStrategy(Coder.class)) - .getLoaded() - .getDeclaredConstructor(Coder[].class, int[].class) - .newInstance((Object) componentCoders, (Object) encodingPosToRowIndex); - } catch (InstantiationException - | IllegalAccessException - | NoSuchMethodException - | InvocationTargetException e) { - throw new RuntimeException("Unable to generate coder for schema " + schema, e); - } + Coder rowCoder = + new RowCoderImpl( + schema, + componentCoders, + encodingPosToRowIndex, + schema.getFields().stream().map(Field::getType).anyMatch(FieldType::getNullable)); String stackTrace = getStackTrace(); GENERATED_CODERS.put(uuid, new WithStackTrace<>(rowCoder, stackTrace)); LOG.debug( @@ -279,126 +190,32 @@ public static Coder generate(Schema schema) { } } - private static class GeneratedCoderConstructor implements Implementation { - @Override - public InstrumentedType prepare(InstrumentedType instrumentedType) { - return instrumentedType; - } - - @Override - public ByteCodeAppender appender(final Target implementationTarget) { - return (methodVisitor, implementationContext, instrumentedMethod) -> { - int numLocals = 1 + instrumentedMethod.getParameters().size(); - StackManipulation stackManipulation = - new StackManipulation.Compound( - // Call the base constructor. - MethodVariableAccess.loadThis(), - Duplication.SINGLE, - MethodInvocation.invoke( - new ForLoadedType(Coder.class) - .getDeclaredMethods() - .filter( - ElementMatchers.isConstructor().and(ElementMatchers.takesArguments(0))) - .getOnly()), - Duplication.SINGLE, - // Store the list of Coders as a member variable. - MethodVariableAccess.REFERENCE.loadFrom(1), - FieldAccess.forField( - implementationTarget - .getInstrumentedType() - .getDeclaredFields() - .filter(ElementMatchers.named(CODERS_FIELD_NAME)) - .getOnly()) - .write(), - // Store the list of encoding offsets as a member variable. - MethodVariableAccess.REFERENCE.loadFrom(2), - FieldAccess.forField( - implementationTarget - .getInstrumentedType() - .getDeclaredFields() - .filter(ElementMatchers.named(POSITIONS_FIELD_NAME)) - .getOnly()) - .write(), - MethodReturn.VOID); - StackManipulation.Size size = stackManipulation.apply(methodVisitor, implementationContext); - return new Size(size.getMaximalSize(), numLocals); - }; + private static final class RowCoderImpl extends CustomCoder { + private final Schema schema; + private final Coder[] coders; + private final int[] encodingPosToIndex; + private final boolean hasNullableFields; + + private RowCoderImpl( + Schema schema, Coder[] coders, int[] encodingPosToIndex, boolean hasNullableFields) { + this.schema = schema; + this.coders = coders; + this.encodingPosToIndex = encodingPosToIndex; + this.hasNullableFields = hasNullableFields; } - } - - private static DynamicType.Builder implementMethods( - Schema schema, DynamicType.Builder builder) { - boolean hasNullableFields = - schema.getFields().stream().map(Field::getType).anyMatch(FieldType::getNullable); - return builder - .defineMethod("getSchema", Schema.class, Visibility.PRIVATE, Ownership.STATIC) - .intercept(FixedValue.reference(schema)) - .defineMethod("hasNullableFields", boolean.class, Visibility.PRIVATE, Ownership.STATIC) - .intercept(FixedValue.reference(hasNullableFields)) - .method(ElementMatchers.named("encode")) - .intercept(new EncodeInstruction()) - .method(ElementMatchers.named("decode")) - .intercept(new DecodeInstruction()); - } - - private static class EncodeInstruction implements Implementation { - static final ForLoadedType LOADED_TYPE = new ForLoadedType(EncodeInstruction.class); @Override - public ByteCodeAppender appender(Target implementationTarget) { - return (methodVisitor, implementationContext, instrumentedMethod) -> { - StackManipulation manipulation = - new StackManipulation.Compound( - // Array of coders. - MethodVariableAccess.loadThis(), - FieldAccess.forField( - implementationContext - .getInstrumentedType() - .getDeclaredFields() - .filter(ElementMatchers.named(CODERS_FIELD_NAME)) - .getOnly()) - .read(), - MethodVariableAccess.loadThis(), - FieldAccess.forField( - implementationContext - .getInstrumentedType() - .getDeclaredFields() - .filter(ElementMatchers.named(POSITIONS_FIELD_NAME)) - .getOnly()) - .read(), - // Element to encode. (offset 1, as offset 0 is always "this"). - MethodVariableAccess.REFERENCE.loadFrom(1), - // OutputStream. - MethodVariableAccess.REFERENCE.loadFrom(2), - // hasNullableFields - MethodInvocation.invoke( - implementationContext - .getInstrumentedType() - .getDeclaredMethods() - .filter(ElementMatchers.named("hasNullableFields")) - .getOnly()), - // Call EncodeInstruction.encodeDelegate - MethodInvocation.invoke( - LOADED_TYPE - .getDeclaredMethods() - .filter( - ElementMatchers.isStatic().and(ElementMatchers.named("encodeDelegate"))) - .getOnly()), - MethodReturn.VOID); - StackManipulation.Size size = manipulation.apply(methodVisitor, implementationContext); - return new ByteCodeAppender.Size(size.getMaximalSize(), instrumentedMethod.getStackSize()); - }; + public void encode(Row value, OutputStream outputStream) throws IOException { + encodeDelegate(coders, encodingPosToIndex, value, outputStream, hasNullableFields); } @Override - public InstrumentedType prepare(InstrumentedType instrumentedType) { - return instrumentedType; + public Row decode(InputStream inputStream) throws IOException { + return decodeDelegate(schema, coders, encodingPosToIndex, inputStream); } - // The encode method of the generated Coder delegates to this method to evaluate all of the - // per-field Coders. @SuppressWarnings("unchecked") - static void encodeDelegate( + private static void encodeDelegate( Coder[] coders, int[] encodingPosToIndex, Row value, @@ -463,62 +280,8 @@ private static BitSet scanNullFields(Object[] fieldValues, int[] encodingPosToIn } return nullFields; } - } - - private static class DecodeInstruction implements Implementation { - static final ForLoadedType LOADED_TYPE = new ForLoadedType(DecodeInstruction.class); - - @Override - public ByteCodeAppender appender(Target implementationTarget) { - return (methodVisitor, implementationContext, instrumentedMethod) -> { - StackManipulation manipulation = - new StackManipulation.Compound( - // Schema. Used in generation of DecodeInstruction. - MethodInvocation.invoke( - implementationContext - .getInstrumentedType() - .getDeclaredMethods() - .filter(ElementMatchers.named("getSchema")) - .getOnly()), - // Array of coders. - MethodVariableAccess.loadThis(), - FieldAccess.forField( - implementationContext - .getInstrumentedType() - .getDeclaredFields() - .filter(ElementMatchers.named(CODERS_FIELD_NAME)) - .getOnly()) - .read(), - MethodVariableAccess.loadThis(), - FieldAccess.forField( - implementationContext - .getInstrumentedType() - .getDeclaredFields() - .filter(ElementMatchers.named(POSITIONS_FIELD_NAME)) - .getOnly()) - .read(), - // read the InputStream. (offset 1, as offset 0 is always "this"). - MethodVariableAccess.REFERENCE.loadFrom(1), - MethodInvocation.invoke( - LOADED_TYPE - .getDeclaredMethods() - .filter( - ElementMatchers.isStatic().and(ElementMatchers.named("decodeDelegate"))) - .getOnly()), - MethodReturn.REFERENCE); - StackManipulation.Size size = manipulation.apply(methodVisitor, implementationContext); - return new ByteCodeAppender.Size(size.getMaximalSize(), instrumentedMethod.getStackSize()); - }; - } - - @Override - public InstrumentedType prepare(InstrumentedType instrumentedType) { - return instrumentedType; - } - // The decode method of the generated Coder delegates to this method to evaluate all of the - // per-field Coders. - static Row decodeDelegate( + private static Row decodeDelegate( Schema schema, Coder[] coders, int[] encodingPosToIndex, InputStream inputStream) throws IOException { int fieldCount; diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/SchemaCoder.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/SchemaCoder.java index 5223cab8f7ca..089bcd2919a1 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/SchemaCoder.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/SchemaCoder.java @@ -110,7 +110,7 @@ public SerializableFunction getToRowFunction() { private Coder getDelegateCoder() { if (delegateCoder == null) { // RowCoderGenerator caches based on id, so if a new instance of this RowCoder is - // deserialized, we don't need to run ByteBuddy again to construct the class. + // deserialized, we don't need to construct the delegate again. delegateCoder = RowCoderGenerator.generate(schema); } return delegateCoder;