Reduce string conversions and allocations in RowJsonDeserializer and RowJsonSerializer

sdks/java/core/jmh/src/main/java/org/apache/beam/sdk/jmh/util/RowJsonBenchmark.java

@@ -0,0 +1,187 @@
+/*
+ * 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.util;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import java.math.BigDecimal;
+import java.nio.charset.StandardCharsets;
+import java.time.LocalDate;
+import java.time.LocalDateTime;
+import java.time.LocalTime;
+import java.util.Arrays;
+import java.util.concurrent.TimeUnit;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.schemas.logicaltypes.SqlTypes;
+import org.apache.beam.sdk.schemas.utils.JsonUtils;
+import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.apache.beam.sdk.util.RowJson.RowJsonDeserializer;
+import org.apache.beam.sdk.util.RowJson.RowJsonSerializer;
+import org.apache.beam.sdk.util.RowJsonUtils;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+import org.joda.time.DateTime;
+import org.joda.time.DateTimeZone;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+
+/** Benchmarks for {@link org.apache.beam.sdk.util.RowJson} and {@link JsonUtils}. */
+@BenchmarkMode(Mode.Throughput)
+@OutputTimeUnit(TimeUnit.MICROSECONDS)
+@Warmup(iterations = 3, time = 1, timeUnit = TimeUnit.SECONDS)
+@Measurement(iterations = 5, time = 1, timeUnit = TimeUnit.SECONDS)
+@Fork(1)
+public class RowJsonBenchmark {
+
+  @State(Scope.Benchmark)
+  public static class BenchmarkState {
+    Schema flatSchema;
+    String flatJsonString;
+    Row flatRow;
+    ObjectMapper flatDeserializerMapper;
+    ObjectMapper flatSerializerMapper;
+
+    Schema complexSchema;
+    String complexJsonString;
+    byte[] complexJsonBytes;
+    Row complexRow;
+    ObjectMapper complexDeserializerMapper;
+    ObjectMapper complexSerializerMapper;
+    SimpleFunction<byte[], Row> jsonBytesToRowFn;
+    SimpleFunction<Row, byte[]> rowToJsonBytesFn;
+
+    @Setup
+    public void setup() {
+      flatSchema =
+          Schema.builder()
+              .addByteField("f_byte")
+              .addInt16Field("f_int16")
+              .addInt32Field("f_int32")
+              .addInt64Field("f_int64")
+              .addFloatField("f_float")
+              .addDoubleField("f_double")
+              .addBooleanField("f_boolean")
+              .addStringField("f_string")
+              .addDecimalField("f_decimal")
+              .addDateTimeField("f_timestamp")
+              .build();
+
+      flatRow =
+          Row.withSchema(flatSchema)
+              .addValues(
+                  (byte) 12,
+                  (short) 22,
+                  32,
+                  42L,
+                  1.02E5f,
+                  62.2d,
+                  true,
+                  "hello",
+                  new BigDecimal("123.12"),
+                  new DateTime(2021, 11, 18, 8, 45, 38, DateTimeZone.UTC).toInstant())
+              .build();
+
+      flatDeserializerMapper =
+          RowJsonUtils.newObjectMapperWith(RowJsonDeserializer.forSchema(flatSchema));
+      flatSerializerMapper =
+          RowJsonUtils.newObjectMapperWith(RowJsonSerializer.forSchema(flatSchema));
+      flatJsonString = RowJsonUtils.rowToJson(flatSerializerMapper, flatRow);
+
+      Schema nestedSchema =
+          Schema.builder()
+              .addStringField("nested_str")
+              .addInt64Field("nested_long")
+              .addNullableField("nested_nullable", FieldType.STRING)
+              .build();
+
+      complexSchema =
+          Schema.builder()
+              .addStringField("id")
+              .addInt32Field("count")
+              .addDoubleField("score")
+              .addBooleanField("active")
+              .addLogicalTypeField("date", SqlTypes.DATE)
+              .addLogicalTypeField("time", SqlTypes.TIME)
+              .addLogicalTypeField("datetime", SqlTypes.DATETIME)
+              .addArrayField("tags", FieldType.STRING)
+              .addMapField("metrics", FieldType.STRING, FieldType.INT64)
+              .addRowField("nested", nestedSchema)
+              .build();
+
+      Row nestedRow = Row.withSchema(nestedSchema).addValues("inner", 999L, null).build();
+
+      complexRow =
+          Row.withSchema(complexSchema)
+              .addValues(
+                  "user-12345",
+                  42,
+                  98.6d,
+                  true,
+                  LocalDate.of(2025, 10, 24),
+                  LocalTime.of(11, 30, 0),
+                  LocalDateTime.of(2025, 10, 24, 11, 30, 0),
+                  Arrays.asList("alpha", "beta", "gamma", "delta"),
+                  ImmutableMap.of("k1", 10L, "k2", 20L, "k3", 30L),
+                  nestedRow)
+              .build();
+
+      complexDeserializerMapper =
+          RowJsonUtils.newObjectMapperWith(RowJsonDeserializer.forSchema(complexSchema));
+      complexSerializerMapper =
+          RowJsonUtils.newObjectMapperWith(RowJsonSerializer.forSchema(complexSchema));
+      complexJsonString = RowJsonUtils.rowToJson(complexSerializerMapper, complexRow);
+      complexJsonBytes = complexJsonString.getBytes(StandardCharsets.UTF_8);
+
+      jsonBytesToRowFn = JsonUtils.getJsonBytesToRowFunction(complexSchema);
+      rowToJsonBytesFn = JsonUtils.getRowToJsonBytesFunction(complexSchema);
+    }
+  }
+
+  @Benchmark
+  public Row deserializeFlatRowString(BenchmarkState state) {
+    return RowJsonUtils.jsonToRow(state.flatDeserializerMapper, state.flatJsonString);
+  }
+
+  @Benchmark
+  public Row deserializeComplexRowString(BenchmarkState state) {
+    return RowJsonUtils.jsonToRow(state.complexDeserializerMapper, state.complexJsonString);
+  }
+
+  @Benchmark
+  public Row deserializeComplexRowBytes(BenchmarkState state) {
+    return state.jsonBytesToRowFn.apply(state.complexJsonBytes);
+  }
+
+  @Benchmark
+  public String serializeComplexRowString(BenchmarkState state) {
+    return RowJsonUtils.rowToJson(state.complexSerializerMapper, state.complexRow);
+  }
+
+  @Benchmark
+  public byte[] serializeComplexRowBytes(BenchmarkState state) {
+    return state.rowToJsonBytesFn.apply(state.complexRow);
+  }
+}

sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/io/payloads/JsonPayloadSerializerProvider.java

@@ -17,8 +17,6 @@
  */
 package org.apache.beam.sdk.schemas.io.payloads;
 
-import static java.nio.charset.StandardCharsets.UTF_8;
-
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.auto.service.AutoService;
 import java.util.Map;
@@ -77,12 +75,12 @@ private ObjectMapper getSerializeMapper() {
 
     @Override
     public byte[] serialize(Row row) {
-      return RowJsonUtils.rowToJson(getSerializeMapper(), row).getBytes(UTF_8);
+      return RowJsonUtils.rowToJsonAsBytes(getSerializeMapper(), row);
     }
 
     @Override
     public Row deserialize(byte[] bytes) {
-      return RowJsonUtils.jsonToRow(getDeserializeMapper(), new String(bytes, UTF_8));
+      return RowJsonUtils.jsonToRow(getDeserializeMapper(), bytes);
     }
   }
 }

sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/utils/JsonUtils.java

@@ -18,6 +18,8 @@
 package org.apache.beam.sdk.schemas.utils;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+import java.io.IOException;
+import java.io.ObjectInputStream;
 import java.nio.charset.StandardCharsets;
 import java.util.HashMap;
 import java.util.Map;
@@ -94,8 +96,7 @@ public static SimpleFunction<byte[], Row> getJsonBytesToRowFunction(Schema beamS
     return new JsonToRowFn<byte[]>(beamSchema) {
       @Override
       public Row apply(byte[] input) {
-        String jsonString = byteArrayToJsonString(input);
-        return RowJsonUtils.jsonToRow(objectMapper, jsonString);
+        return RowJsonUtils.jsonToRow(objectMapper, input);
       }
     };
   }
@@ -115,8 +116,7 @@ public static SimpleFunction<Row, byte[]> getRowToJsonBytesFunction(Schema beamS
     return new RowToJsonFn<byte[]>(beamSchema) {
       @Override
       public byte[] apply(Row input) {
-        String jsonString = RowJsonUtils.rowToJson(objectMapper, input);
-        return jsonStringToByteArray(jsonString);
+        return RowJsonUtils.rowToJsonAsBytes(objectMapper, input);
       }
     };
   }
@@ -319,22 +319,40 @@ private static org.everit.json.schema.ObjectSchema jsonSchemaFromString(String j
   }
 
   private abstract static class JsonToRowFn<T> extends SimpleFunction<T, Row> {
-    final RowJson.RowJsonDeserializer deserializer;
-    final ObjectMapper objectMapper;
+    final Schema beamSchema;
+    transient ObjectMapper objectMapper;
 
     private JsonToRowFn(Schema beamSchema) {
-      deserializer = RowJson.RowJsonDeserializer.forSchema(beamSchema);
-      objectMapper = RowJsonUtils.newObjectMapperWith(deserializer);
+      this.beamSchema = beamSchema;
+      objectMapper = newObjectMapper(beamSchema);
+    }
+
+    private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException {
+      in.defaultReadObject();
+      objectMapper = newObjectMapper(beamSchema);
+    }
+
+    private static ObjectMapper newObjectMapper(Schema schema) {
+      return RowJsonUtils.newObjectMapperWith(RowJson.RowJsonDeserializer.forSchema(schema));
     }
   }
 
   private abstract static class RowToJsonFn<T> extends SimpleFunction<Row, T> {
-    final RowJson.RowJsonSerializer serializer;
-    final ObjectMapper objectMapper;
+    final Schema beamSchema;
+    transient ObjectMapper objectMapper;
 
     private RowToJsonFn(Schema beamSchema) {
-      serializer = RowJson.RowJsonSerializer.forSchema(beamSchema);
-      objectMapper = RowJsonUtils.newObjectMapperWith(serializer);
+      this.beamSchema = beamSchema;
+      objectMapper = newObjectMapper(beamSchema);
+    }
+
+    private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException {
+      in.defaultReadObject();
+      objectMapper = newObjectMapper(beamSchema);
+    }
+
+    private static ObjectMapper newObjectMapper(Schema schema) {
+      return RowJsonUtils.newObjectMapperWith(RowJson.RowJsonSerializer.forSchema(schema));
     }
   }
 

sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/JsonToRow.java

@@ -261,22 +261,33 @@ public JsonToRowWithErrFn withNullBehavior(NullBehavior nullBehavior) {
       return this.toBuilder().setNullBehavior(nullBehavior).build();
     }
 
+    Schema getErrorSchema() {
+      if (getExtendedErrorInfo()) {
+        if (LINE_FIELD_NAME.equals(getLineFieldName())
+            && ERROR_FIELD_NAME.equals(getErrorFieldName())) {
+          return ERROR_ROW_WITH_ERR_MSG_SCHEMA;
+        }
+        return Schema.of(
+            Field.of(getLineFieldName(), FieldType.STRING),
+            Field.of(getErrorFieldName(), FieldType.STRING));
+      } else {
+        if (LINE_FIELD_NAME.equals(getLineFieldName())) {
+          return ERROR_ROW_SCHEMA;
+        }
+        return Schema.of(Field.of(getLineFieldName(), FieldType.STRING));
+      }
+    }
+
     @Override
     public ParseResult expand(PCollection<String> jsonStrings) {
+      Schema errorSchema = getErrorSchema();
 
       PCollectionTuple result =
           jsonStrings.apply(
-              ParDo.of(ParseWithError.create(this))
+              ParDo.of(ParseWithError.create(this, errorSchema))
                   .withOutputTags(PARSED_LINE, TupleTagList.of(PARSE_ERROR)));
 
-      PCollection<Row> failures;
-
-      if (getExtendedErrorInfo()) {
-        failures =
-            result.get(PARSE_ERROR).setRowSchema(JsonToRowWithErrFn.ERROR_ROW_WITH_ERR_MSG_SCHEMA);
-      } else {
-        failures = result.get(PARSE_ERROR).setRowSchema(JsonToRowWithErrFn.ERROR_ROW_SCHEMA);
-      }
+      PCollection<Row> failures = result.get(PARSE_ERROR).setRowSchema(errorSchema);
 
       return ParseResult.resultBuilder()
           .setCallingPipeline(jsonStrings.getPipeline())
@@ -292,18 +303,28 @@ protected abstract static class ParseWithError extends DoFn<String, Row> {
 
       public abstract JsonToRowWithErrFn getJsonToRowWithErrFn();
 
+      public abstract Schema getErrorSchema();
+
       public abstract Builder toBuilder();
 
       @AutoValue.Builder
       public abstract static class Builder {
         public abstract Builder setJsonToRowWithErrFn(JsonToRowWithErrFn value);
 
+        public abstract Builder setErrorSchema(Schema value);
+
         public abstract ParseWithError build();
       }
 
       public static ParseWithError create(JsonToRowWithErrFn jsonToRowWithErrFn) {
+        return create(jsonToRowWithErrFn, jsonToRowWithErrFn.getErrorSchema());
+      }
+
+      public static ParseWithError create(
+          JsonToRowWithErrFn jsonToRowWithErrFn, Schema errorSchema) {
         return new AutoValue_JsonToRow_JsonToRowWithErrFn_ParseWithError.Builder()
             .setJsonToRowWithErrFn(jsonToRowWithErrFn)
+            .setErrorSchema(errorSchema)
             .build();
       }
 
@@ -317,14 +338,14 @@ public void processElement(@Element String element, MultiOutputReceiver output)
             output
                 .get(PARSE_ERROR)
                 .output(
-                    Row.withSchema(ERROR_ROW_WITH_ERR_MSG_SCHEMA)
+                    Row.withSchema(getErrorSchema())
                         .addValue(element)
                         .addValue(ex.getMessage())
                         .build());
           } else {
             output
                 .get(PARSE_ERROR)
-                .output(Row.withSchema(ERROR_ROW_SCHEMA).addValue(element).build());
+                .output(Row.withSchema(getErrorSchema()).addValue(element).build());
           }
           return;
         }