diff --git a/sdks/java/core/jmh/src/main/java/org/apache/beam/sdk/jmh/util/RowJsonBenchmark.java b/sdks/java/core/jmh/src/main/java/org/apache/beam/sdk/jmh/util/RowJsonBenchmark.java new file mode 100644 index 000000000000..e1bf16d8f6e8 --- /dev/null +++ b/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 jsonBytesToRowFn; + SimpleFunction 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); + } +} diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/io/payloads/JsonPayloadSerializerProvider.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/io/payloads/JsonPayloadSerializerProvider.java index d3054e3689a5..dc9b6acdb6fe 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/io/payloads/JsonPayloadSerializerProvider.java +++ b/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); } } } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/utils/JsonUtils.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/utils/JsonUtils.java index 18f5813c6cc9..6896d9a6d310 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/utils/JsonUtils.java +++ b/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 getJsonBytesToRowFunction(Schema beamS return new JsonToRowFn(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 getRowToJsonBytesFunction(Schema beamS return new RowToJsonFn(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 extends SimpleFunction { - 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 extends SimpleFunction { - 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)); } } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/JsonToRow.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/JsonToRow.java index 9845af44a82e..9b7ee3f4d522 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/JsonToRow.java +++ b/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 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 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 failures = result.get(PARSE_ERROR).setRowSchema(errorSchema); return ParseResult.resultBuilder() .setCallingPipeline(jsonStrings.getPipeline()) @@ -292,18 +303,28 @@ protected abstract static class ParseWithError extends DoFn { 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; } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/RowJson.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/RowJson.java index ccfbd87d450d..7f5459c82af9 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/RowJson.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/RowJson.java @@ -41,28 +41,23 @@ import static org.apache.beam.sdk.util.RowJsonValueExtractors.shortValueExtractor; import static org.apache.beam.sdk.util.RowJsonValueExtractors.stringValueExtractor; import static org.apache.beam.sdk.util.RowJsonValueExtractors.timeValueExtractor; -import static org.apache.beam.sdk.values.Row.toRow; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList.toImmutableList; import com.fasterxml.jackson.core.JsonGenerator; import com.fasterxml.jackson.core.JsonParser; +import com.fasterxml.jackson.core.JsonToken; import com.fasterxml.jackson.databind.DeserializationContext; -import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.SerializerProvider; import com.fasterxml.jackson.databind.deser.std.StdDeserializer; -import com.fasterxml.jackson.databind.node.JsonNodeType; import com.fasterxml.jackson.databind.ser.std.StdSerializer; -import com.google.auto.value.AutoValue; import java.io.IOException; import java.math.BigDecimal; import java.time.LocalDate; import java.time.LocalDateTime; import java.time.LocalTime; +import java.util.Arrays; import java.util.HashMap; -import java.util.Iterator; import java.util.Map; -import java.util.stream.Stream; -import java.util.stream.StreamSupport; import org.apache.beam.sdk.schemas.Schema; import org.apache.beam.sdk.schemas.Schema.Field; import org.apache.beam.sdk.schemas.Schema.FieldType; @@ -73,6 +68,7 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.ReadableInstant; @@ -191,8 +187,6 @@ private static ImmutableList findUnsupportedFields( /** Jackson deserializer for parsing JSON into {@link Row Rows}. */ public static class RowJsonDeserializer extends StdDeserializer { - private static final boolean SEQUENTIAL = false; - /** * An enumeration type for specifying how {@link RowJsonDeserializer} should expect null values * to be represented. @@ -227,20 +221,26 @@ public enum NullBehavior { } private static final ImmutableMap> JSON_VALUE_GETTERS = - ImmutableMap.>builder() - .put(BYTE, byteValueExtractor()) - .put(INT16, shortValueExtractor()) - .put(INT32, intValueExtractor()) - .put(INT64, longValueExtractor()) - .put(FLOAT, floatValueExtractor()) - .put(DOUBLE, doubleValueExtractor()) - .put(BOOLEAN, booleanValueExtractor()) - .put(STRING, stringValueExtractor()) - .put(DECIMAL, decimalValueExtractor()) - .put(DATETIME, datetimeValueExtractor()) - .build(); - - private final Schema schema; + Maps.immutableEnumMap( + ImmutableMap.>builder() + .put(BYTE, byteValueExtractor()) + .put(INT16, shortValueExtractor()) + .put(INT32, intValueExtractor()) + .put(INT64, longValueExtractor()) + .put(FLOAT, floatValueExtractor()) + .put(DOUBLE, doubleValueExtractor()) + .put(BOOLEAN, booleanValueExtractor()) + .put(STRING, stringValueExtractor()) + .put(DECIMAL, decimalValueExtractor()) + .put(DATETIME, datetimeValueExtractor()) + .build()); + + private static final ValueExtractor DATE_EXTRACTOR = dateValueExtractor(); + private static final ValueExtractor TIME_EXTRACTOR = timeValueExtractor(); + private static final ValueExtractor LOCAL_DATETIME_EXTRACTOR = + localDatetimeValueExtractor(); + + private final FieldType rootRowType; private NullBehavior nullBehavior = NullBehavior.ACCEPT_MISSING_OR_NULL; /** Creates a deserializer for a {@link Row} {@link Schema}. */ @@ -251,7 +251,7 @@ public static RowJsonDeserializer forSchema(Schema schema) { private RowJsonDeserializer(Schema schema) { super(Row.class); - this.schema = schema; + this.rootRowType = FieldType.row(schema); } /** @@ -267,236 +267,216 @@ public RowJsonDeserializer withNullBehavior(NullBehavior behavior) { @Override public Row deserialize(JsonParser jsonParser, DeserializationContext deserializationContext) throws IOException { - + if (jsonParser.currentToken() == null) { + jsonParser.nextToken(); + } // Parse and convert the root object to Row as if it's a nested field with name 'root' - return (Row) - extractJsonNodeValue( - FieldValue.of("root", FieldType.row(schema), jsonParser.readValueAsTree())); + return (Row) extractJsonValue("root", rootRowType, jsonParser); } - private Object extractJsonNodeValue(FieldValue fieldValue) { - if (fieldValue.type().getNullable()) { - if (!fieldValue.isJsonValuePresent()) { + private @Nullable Object extractJsonValue( + String fieldName, FieldType fieldType, JsonParser jsonParser) throws IOException { + JsonToken token = jsonParser.currentToken(); + if (fieldType.getNullable()) { + if (token == null) { switch (this.nullBehavior) { case ACCEPT_MISSING_OR_NULL: case REQUIRE_MISSING: return null; case REQUIRE_NULL: throw new UnsupportedRowJsonException( - "Field '" + fieldValue.name() + "' is not present in the JSON object."); + "Field '" + fieldName + "' is not present in the JSON object."); } } - if (fieldValue.isJsonNull()) { + if (token == JsonToken.VALUE_NULL) { switch (this.nullBehavior) { case ACCEPT_MISSING_OR_NULL: case REQUIRE_NULL: return null; case REQUIRE_MISSING: throw new UnsupportedRowJsonException( - "Field '" + fieldValue.name() + "' has a null value in the JSON object."); + "Field '" + fieldName + "' has a null value in the JSON object."); } } } else { // field is not nullable - if (!fieldValue.isJsonValuePresent()) { + if (token == null) { throw new UnsupportedRowJsonException( - "Non-nullable field '" + fieldValue.name() + "' is not present in the JSON object."); - } else if (fieldValue.isJsonNull()) { + "Non-nullable field '" + fieldName + "' is not present in the JSON object."); + } else if (token == JsonToken.VALUE_NULL) { throw new UnsupportedRowJsonException( - "Non-nullable field '" + fieldValue.name() + "' has value null in the JSON object."); + "Non-nullable field '" + fieldName + "' has value null in the JSON object."); } } - if (fieldValue.isRowType()) { - return jsonObjectToRow(fieldValue); - } - - if (fieldValue.isArrayType()) { - return jsonArrayToList(fieldValue); + TypeName typeName = fieldType.getTypeName(); + switch (typeName) { + case ROW: + return jsonObjectToRow(fieldName, fieldType.getRowSchema(), jsonParser); + case ARRAY: + case ITERABLE: + return jsonArrayToList(fieldName, fieldType.getCollectionElementType(), jsonParser); + case MAP: + return jsonObjectToMap(fieldName, fieldType.getMapValueType(), jsonParser); + case LOGICAL_TYPE: + Schema.LogicalType logicalType = + (Schema.LogicalType) fieldType.getLogicalType(); + String identifier = logicalType.getIdentifier(); + if (SqlTypes.DATE.getIdentifier().equals(identifier)) { + return extractWithExtractor(fieldName, typeName, jsonParser, DATE_EXTRACTOR); + } else if (SqlTypes.TIME.getIdentifier().equals(identifier)) { + return extractWithExtractor(fieldName, typeName, jsonParser, TIME_EXTRACTOR); + } else if (SqlTypes.DATETIME.getIdentifier().equals(identifier)) { + return extractWithExtractor(fieldName, typeName, jsonParser, LOCAL_DATETIME_EXTRACTOR); + } else { + return logicalType.toInputType( + extractJsonValue(fieldName, logicalType.getBaseType(), jsonParser)); + } + default: + return extractJsonPrimitiveValue(fieldName, typeName, jsonParser); } + } - if (fieldValue.isMapType()) { - return jsonObjectToMap(fieldValue); + private Row jsonObjectToRow(String fieldName, Schema rowSchema, JsonParser jsonParser) + throws IOException { + if (jsonParser.currentToken() != JsonToken.START_OBJECT) { + throw new UnsupportedRowJsonException( + "Expected JSON object for field '" + + fieldName + + "'. Unable to convert '" + + jsonParser.getValueAsString("") + + "' to Beam Row, it is not a JSON object. Currently only JSON objects can be parsed to Beam Rows"); } - if (fieldValue.typeName().isLogicalType()) { - String identifier = fieldValue.type().getLogicalType().getIdentifier(); - if (SqlTypes.DATE.getIdentifier().equals(identifier)) { - return dateValueExtractor().extractValue(fieldValue.jsonValue()); - } else if (SqlTypes.TIME.getIdentifier().equals(identifier)) { - return timeValueExtractor().extractValue(fieldValue.jsonValue()); - } else if (SqlTypes.DATETIME.getIdentifier().equals(identifier)) { - return localDatetimeValueExtractor().extractValue(fieldValue.jsonValue()); + int fieldCount = rowSchema.getFieldCount(); + Object[] values = new Object[fieldCount]; + boolean[] fieldPresent = new boolean[fieldCount]; + + while (jsonParser.nextToken() != JsonToken.END_OBJECT) { + String propertyName = jsonParser.currentName(); + jsonParser.nextToken(); + if (rowSchema.hasField(propertyName)) { + int index = rowSchema.indexOf(propertyName); + Field schemaField = rowSchema.getField(index); + fieldPresent[index] = true; + values[index] = extractJsonValue(propertyName, schemaField.getType(), jsonParser); } else { - return extractJsonNodeValue( - FieldValue.of( - fieldValue.name(), - fieldValue.type().getLogicalType().getBaseType(), - fieldValue.jsonValue())); + jsonParser.skipChildren(); } } - return extractJsonPrimitiveValue(fieldValue); - } - - private Row jsonObjectToRow(FieldValue rowFieldValue) { - if (!rowFieldValue.isJsonObject()) { - throw new UnsupportedRowJsonException( - "Expected JSON object for field '" - + rowFieldValue.name() - + "'. Unable to convert '" - + rowFieldValue.jsonValue().asText() - + "' to Beam Row, it is not a JSON object. Currently only JSON objects can be parsed to Beam Rows"); + for (int i = 0; i < fieldCount; i++) { + if (!fieldPresent[i]) { + Field schemaField = rowSchema.getField(i); + if (!schemaField.getType().getNullable()) { + throw new UnsupportedRowJsonException( + "Non-nullable field '" + + schemaField.getName() + + "' is not present in the JSON object."); + } else if (this.nullBehavior == NullBehavior.REQUIRE_NULL) { + throw new UnsupportedRowJsonException( + "Field '" + schemaField.getName() + "' is not present in the JSON object."); + } + } } - return rowFieldValue.rowSchema().getFields().stream() - .map( - schemaField -> - extractJsonNodeValue( - FieldValue.of( - schemaField.getName(), - schemaField.getType(), - rowFieldValue.jsonFieldValue(schemaField.getName())))) - .collect(toRow(rowFieldValue.rowSchema())); + return Row.withSchema(rowSchema).addValues(Arrays.asList(values)).build(); } - private Object jsonArrayToList(FieldValue arrayFieldValue) { - if (!arrayFieldValue.isJsonArray()) { + private Object jsonArrayToList(String fieldName, FieldType elementType, JsonParser jsonParser) + throws IOException { + JsonToken token = jsonParser.currentToken(); + if (token != JsonToken.START_ARRAY) { throw new UnsupportedRowJsonException( "Expected JSON array for field '" - + arrayFieldValue.name() + + fieldName + "'. Instead got " - + arrayFieldValue.jsonNodeType().name()); + + tokenNodeTypeName(token)); } - return arrayFieldValue - .jsonArrayElements() - .map( - jsonArrayElement -> - extractJsonNodeValue( - FieldValue.of( - arrayFieldValue.name() + "[]", - arrayFieldValue.arrayElementType(), - jsonArrayElement))) - .collect(toImmutableList()); + if (jsonParser.nextToken() == JsonToken.END_ARRAY) { + return ImmutableList.of(); + } + String elementFieldName = fieldName + "[]"; + ImmutableList.Builder builder = ImmutableList.builder(); + do { + builder.add(extractJsonValue(elementFieldName, elementType, jsonParser)); + } while (jsonParser.nextToken() != JsonToken.END_ARRAY); + return builder.build(); } - private Map jsonObjectToMap(FieldValue mapFieldValue) { - if (!mapFieldValue.isJsonObject()) { + private Map jsonObjectToMap( + String fieldName, FieldType mapValueType, JsonParser jsonParser) throws IOException { + JsonToken token = jsonParser.currentToken(); + if (token != JsonToken.START_OBJECT) { throw new UnsupportedRowJsonException( "Expected JSON object for field '" - + mapFieldValue.name() + + fieldName + "'. Instead got " - + mapFieldValue.jsonNodeType().name()); + + tokenNodeTypeName(token)); } Map result = new HashMap<>(); - Iterator> fields = mapFieldValue.jsonValue().fields(); - while (fields.hasNext()) { - Map.Entry field = fields.next(); - String key = field.getKey(); - JsonNode value = field.getValue(); - + while (jsonParser.nextToken() != JsonToken.END_OBJECT) { + String key = jsonParser.currentName(); + jsonParser.nextToken(); Object extractedValue = - extractJsonNodeValue( - FieldValue.of( - mapFieldValue.name() + "['" + key + "']", mapFieldValue.mapValueType(), value)); - + extractJsonValue(fieldName + "['" + key + "']", mapValueType, jsonParser); result.put(key, extractedValue); } return result; } - private static Object extractJsonPrimitiveValue(FieldValue fieldValue) { + private static Object extractJsonPrimitiveValue( + String fieldName, TypeName typeName, JsonParser jsonParser) throws IOException { + return extractWithExtractor( + fieldName, typeName, jsonParser, JSON_VALUE_GETTERS.get(typeName)); + } + + private static Object extractWithExtractor( + String fieldName, TypeName typeName, JsonParser jsonParser, ValueExtractor extractor) + throws IOException { try { - return JSON_VALUE_GETTERS.get(fieldValue.typeName()).extractValue(fieldValue.jsonValue()); + return extractor.extractValue(jsonParser); } catch (RuntimeException e) { throw new UnsupportedRowJsonException( "Unable to get value from field '" - + fieldValue.name() + + fieldName + "'. Schema type '" - + fieldValue.typeName() + + typeName + "'. JSON node type " - + fieldValue.jsonNodeType().name(), + + tokenNodeTypeName(jsonParser.currentToken()), e); } } - /** - * Helper class to keep track of schema field type, name, and actual json value for the field. - */ - @AutoValue - abstract static class FieldValue { - abstract String name(); - - abstract FieldType type(); - - abstract @Nullable JsonNode jsonValue(); - - TypeName typeName() { - return type().getTypeName(); - } - - boolean isJsonValuePresent() { - return jsonValue() != null; - } - - boolean isJsonNull() { - return jsonValue().isNull(); - } - - JsonNodeType jsonNodeType() { - return jsonValue().getNodeType(); - } - - boolean isJsonArray() { - return jsonValue().isArray(); - } - - Stream jsonArrayElements() { - return StreamSupport.stream(jsonValue().spliterator(), SEQUENTIAL); - } - - boolean isArrayType() { - return TypeName.ARRAY.equals(type().getTypeName()) - || TypeName.ITERABLE.equals(type().getTypeName()); - } - - FieldType arrayElementType() { - return type().getCollectionElementType(); - } - - boolean isJsonObject() { - return jsonValue().isObject(); + private static String tokenNodeTypeName(@Nullable JsonToken token) { + if (token == null) { + return "MISSING"; } - - JsonNode jsonFieldValue(String fieldName) { - return jsonValue().get(fieldName); - } - - boolean isRowType() { - return TypeName.ROW.equals(type().getTypeName()); - } - - Schema rowSchema() { - return type().getRowSchema(); - } - - boolean isMapType() { - return TypeName.MAP.equals(type().getTypeName()); - } - - FieldType mapKeyType() { - return type().getMapKeyType(); - } - - FieldType mapValueType() { - return type().getMapValueType(); - } - - static FieldValue of(String name, FieldType type, JsonNode jsonValue) { - return new AutoValue_RowJson_RowJsonDeserializer_FieldValue(name, type, jsonValue); + switch (token) { + case START_OBJECT: + case END_OBJECT: + case FIELD_NAME: + return "OBJECT"; + case START_ARRAY: + case END_ARRAY: + return "ARRAY"; + case VALUE_STRING: + return "STRING"; + case VALUE_NUMBER_INT: + case VALUE_NUMBER_FLOAT: + return "NUMBER"; + case VALUE_TRUE: + case VALUE_FALSE: + return "BOOLEAN"; + case VALUE_NULL: + return "NULL"; + case VALUE_EMBEDDED_OBJECT: + return "POJO"; + default: + return token.name(); } } } @@ -505,7 +485,7 @@ static FieldValue of(String name, FieldType type, JsonNode jsonValue) { public static class RowJsonSerializer extends StdSerializer { private final Schema schema; - private Boolean dropNullsOnWrite = false; + private boolean dropNullsOnWrite = false; /** Creates a serializer for a {@link Row} {@link Schema}. */ public static RowJsonSerializer forSchema(Schema schema) { @@ -520,7 +500,7 @@ private RowJsonSerializer(Schema schema) { /** Serializer drops nulls on write if set to true instead of writing fieldName: null. */ public RowJsonSerializer withDropNullsOnWrite(Boolean dropNullsOnWrite) { - this.dropNullsOnWrite = dropNullsOnWrite; + this.dropNullsOnWrite = Boolean.TRUE.equals(dropNullsOnWrite); return this; } @@ -602,7 +582,9 @@ private void writeValue(JsonGenerator gen, FieldType type, Object value) throws gen.writeEndObject(); break; case LOGICAL_TYPE: - String identifier = type.getLogicalType().getIdentifier(); + Schema.LogicalType logicalType = + (Schema.LogicalType) type.getLogicalType(); + String identifier = logicalType.getIdentifier(); if (SqlTypes.DATE.getIdentifier().equals(identifier)) { gen.writeString(((LocalDate) value).toString()); // ISO 8601 format } else if (SqlTypes.TIME.getIdentifier().equals(identifier)) { @@ -610,7 +592,7 @@ private void writeValue(JsonGenerator gen, FieldType type, Object value) throws } else if (SqlTypes.DATETIME.getIdentifier().equals(identifier)) { gen.writeString(((LocalDateTime) value).toString()); // ISO 8601 format } else { - writeValue(gen, type.getLogicalType().getBaseType(), value); + writeValue(gen, logicalType.getBaseType(), logicalType.toBaseType(value)); } break; default: diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/RowJsonUtils.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/RowJsonUtils.java index ee41d0da28fe..fe710fe45bc8 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/RowJsonUtils.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/RowJsonUtils.java @@ -141,6 +141,19 @@ public static Row jsonToRow(ObjectMapper objectMapper, String jsonString) { } } + public static Row jsonToRow(ObjectMapper objectMapper, byte[] jsonBytes) { + try { + return objectMapper.readValue(jsonBytes, Row.class); + } catch (JsonParseException | JsonMappingException jsonException) { + throw new UnsupportedRowJsonException("Unable to parse Row", jsonException); + } catch (IOException e) { + throw new IllegalArgumentException( + "Unable to parse json object: " + + new String(jsonBytes, java.nio.charset.StandardCharsets.UTF_8), + e); + } + } + public static String rowToJson(ObjectMapper objectMapper, Row row) { try { return objectMapper.writeValueAsString(row); @@ -148,4 +161,12 @@ public static String rowToJson(ObjectMapper objectMapper, Row row) { throw new IllegalArgumentException("Unable to serialize row: " + row, e); } } + + public static byte[] rowToJsonAsBytes(ObjectMapper objectMapper, Row row) { + try { + return objectMapper.writeValueAsBytes(row); + } catch (JsonProcessingException e) { + throw new IllegalArgumentException("Unable to serialize row: " + row, e); + } + } } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/RowJsonValueExtractors.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/RowJsonValueExtractors.java index e769212fce76..a784234485cb 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/RowJsonValueExtractors.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/RowJsonValueExtractors.java @@ -17,8 +17,12 @@ */ package org.apache.beam.sdk.util; +import com.fasterxml.jackson.core.JsonParser; +import com.fasterxml.jackson.core.JsonParser.NumberType; +import com.fasterxml.jackson.core.JsonToken; import com.fasterxml.jackson.databind.JsonNode; import com.google.auto.value.AutoValue; +import java.io.IOException; import java.math.BigDecimal; import java.time.LocalDate; import java.time.LocalDateTime; @@ -29,7 +33,8 @@ import org.joda.time.DateTime; /** - * Contains utilities for extracting primitive values from JSON nodes. + * Contains utilities for extracting primitive values from JSON nodes and streaming {@link + * JsonParser} tokens. * *

Performs validation and rejects values which are out of bounds. */ @@ -37,6 +42,18 @@ class RowJsonValueExtractors { public interface ValueExtractor { V extractValue(JsonNode value); + + default V extractValue(JsonParser parser) throws IOException { + return extractValue((JsonNode) parser.readValueAsTree()); + } + } + + private static UnsupportedRowJsonException outOfRangeException(JsonParser parser) + throws IOException { + return new UnsupportedRowJsonException( + "Value \"" + + parser.getText() + + "\" is out of range for the type of the field defined in the row schema."); } /** @@ -53,6 +70,17 @@ static ValueExtractor byteValueExtractor() { && jsonNode.canConvertToInt() && jsonNode.intValue() >= Byte.MIN_VALUE && jsonNode.intValue() <= Byte.MAX_VALUE) + .setParserExtractor( + parser -> { + if (parser.currentToken() == JsonToken.VALUE_NUMBER_INT + && parser.getNumberType() == NumberType.INT) { + int value = parser.getIntValue(); + if (value >= Byte.MIN_VALUE && value <= Byte.MAX_VALUE) { + return (byte) value; + } + } + throw outOfRangeException(parser); + }) .build(); } @@ -70,6 +98,17 @@ static ValueExtractor shortValueExtractor() { && jsonNode.canConvertToInt() && jsonNode.intValue() >= Short.MIN_VALUE && jsonNode.intValue() <= Short.MAX_VALUE) + .setParserExtractor( + parser -> { + if (parser.currentToken() == JsonToken.VALUE_NUMBER_INT + && parser.getNumberType() == NumberType.INT) { + int value = parser.getIntValue(); + if (value >= Short.MIN_VALUE && value <= Short.MAX_VALUE) { + return (short) value; + } + } + throw outOfRangeException(parser); + }) .build(); } @@ -82,6 +121,14 @@ static ValueExtractor intValueExtractor() { return ValidatingValueExtractor.builder() .setExtractor(JsonNode::intValue) .setValidator(jsonNode -> jsonNode.isIntegralNumber() && jsonNode.canConvertToInt()) + .setParserExtractor( + parser -> { + if (parser.currentToken() == JsonToken.VALUE_NUMBER_INT + && parser.getNumberType() == NumberType.INT) { + return parser.getIntValue(); + } + throw outOfRangeException(parser); + }) .build(); } @@ -94,6 +141,16 @@ static ValueExtractor longValueExtractor() { return ValidatingValueExtractor.builder() .setExtractor(JsonNode::longValue) .setValidator(jsonNode -> jsonNode.isIntegralNumber() && jsonNode.canConvertToLong()) + .setParserExtractor( + parser -> { + if (parser.currentToken() == JsonToken.VALUE_NUMBER_INT) { + NumberType numberType = parser.getNumberType(); + if (numberType == NumberType.INT || numberType == NumberType.LONG) { + return parser.getLongValue(); + } + } + throw outOfRangeException(parser); + }) .build(); } @@ -120,10 +177,26 @@ static ValueExtractor floatValueExtractor() { // the float it went through is 2147483648. || (jsonNode.isIntegralNumber() && jsonNode.canConvertToInt() - && jsonNode - .decimalValue() - .compareTo(BigDecimal.valueOf(jsonNode.floatValue())) - == 0)) + && (double) (float) jsonNode.intValue() == (double) jsonNode.intValue())) + .setParserExtractor( + parser -> { + JsonToken token = parser.currentToken(); + if (token == JsonToken.VALUE_NUMBER_FLOAT) { + double d = parser.getDoubleValue(); + float f = (float) d; + if (d == (double) f) { + return f; + } + } else if (token == JsonToken.VALUE_NUMBER_INT + && parser.getNumberType() == NumberType.INT) { + int i = parser.getIntValue(); + float f = (float) i; + if ((double) f == (double) i) { + return f; + } + } + throw outOfRangeException(parser); + }) .build(); } @@ -146,10 +219,8 @@ static ValueExtractor doubleValueExtractor() { // would look like it survived the trip when it did not. || (jsonNode.isIntegralNumber() && jsonNode.canConvertToLong() - && jsonNode - .decimalValue() - .compareTo(BigDecimal.valueOf(jsonNode.doubleValue())) - == 0) + && jsonNode.longValue() != Long.MAX_VALUE + && (long) jsonNode.doubleValue() == jsonNode.longValue()) // Or a decimal number which allows lossless conversion to float || (jsonNode.isFloatingPointNumber() @@ -157,6 +228,30 @@ static ValueExtractor doubleValueExtractor() { .decimalValue() .compareTo(BigDecimal.valueOf(jsonNode.doubleValue())) == 0)) + .setParserExtractor( + parser -> { + JsonToken token = parser.currentToken(); + if (token == JsonToken.VALUE_NUMBER_FLOAT) { + double d = parser.getDoubleValue(); + if (!Double.isInfinite(d) + && !Double.isNaN(d) + && parser.getDecimalValue().compareTo(BigDecimal.valueOf(d)) == 0) { + return d; + } + } else if (token == JsonToken.VALUE_NUMBER_INT) { + NumberType numberType = parser.getNumberType(); + if (numberType == NumberType.INT) { + return (double) parser.getIntValue(); + } else if (numberType == NumberType.LONG) { + long l = parser.getLongValue(); + double d = (double) l; + if (l != Long.MAX_VALUE && (long) d == l) { + return d; + } + } + } + throw outOfRangeException(parser); + }) .build(); } @@ -169,6 +264,16 @@ static ValueExtractor booleanValueExtractor() { return ValidatingValueExtractor.builder() .setExtractor(JsonNode::booleanValue) .setValidator(JsonNode::isBoolean) + .setParserExtractor( + parser -> { + JsonToken token = parser.currentToken(); + if (token == JsonToken.VALUE_TRUE) { + return Boolean.TRUE; + } else if (token == JsonToken.VALUE_FALSE) { + return Boolean.FALSE; + } + throw outOfRangeException(parser); + }) .build(); } @@ -181,6 +286,13 @@ static ValueExtractor stringValueExtractor() { return ValidatingValueExtractor.builder() .setExtractor(JsonNode::textValue) .setValidator(JsonNode::isTextual) + .setParserExtractor( + parser -> { + if (parser.currentToken() == JsonToken.VALUE_STRING) { + return parser.getText(); + } + throw outOfRangeException(parser); + }) .build(); } @@ -193,6 +305,14 @@ static ValueExtractor decimalValueExtractor() { return ValidatingValueExtractor.builder() .setExtractor(JsonNode::decimalValue) .setValidator(jsonNode -> jsonNode.isNumber()) + .setParserExtractor( + parser -> { + JsonToken token = parser.currentToken(); + if (token == JsonToken.VALUE_NUMBER_INT || token == JsonToken.VALUE_NUMBER_FLOAT) { + return parser.getDecimalValue(); + } + throw outOfRangeException(parser); + }) .build(); } @@ -206,12 +326,23 @@ static ValueExtractor datetimeValueExtractor() { .setExtractor( jsonNode -> { String text = jsonNode.textValue(); - if (text.contains(" ")) { + if (text.indexOf(' ') >= 0) { text = text.replace(' ', 'T'); } return DateTime.parse(text); }) .setValidator(JsonNode::isTextual) + .setParserExtractor( + parser -> { + if (parser.currentToken() == JsonToken.VALUE_STRING) { + String text = parser.getText(); + if (text.indexOf(' ') >= 0) { + text = text.replace(' ', 'T'); + } + return DateTime.parse(text); + } + throw outOfRangeException(parser); + }) .build(); } @@ -224,6 +355,13 @@ static ValueExtractor dateValueExtractor() { return ValidatingValueExtractor.builder() .setExtractor(jsonNode -> LocalDate.parse(jsonNode.textValue())) .setValidator(JsonNode::isTextual) + .setParserExtractor( + parser -> { + if (parser.currentToken() == JsonToken.VALUE_STRING) { + return LocalDate.parse(parser.getText()); + } + throw outOfRangeException(parser); + }) .build(); } @@ -236,6 +374,13 @@ static ValueExtractor timeValueExtractor() { return ValidatingValueExtractor.builder() .setExtractor(jsonNode -> LocalTime.parse(jsonNode.textValue())) .setValidator(JsonNode::isTextual) + .setParserExtractor( + parser -> { + if (parser.currentToken() == JsonToken.VALUE_STRING) { + return LocalTime.parse(parser.getText()); + } + throw outOfRangeException(parser); + }) .build(); } @@ -248,9 +393,20 @@ static ValueExtractor localDatetimeValueExtractor() { return ValidatingValueExtractor.builder() .setExtractor(jsonNode -> LocalDateTime.parse(jsonNode.textValue())) .setValidator(JsonNode::isTextual) + .setParserExtractor( + parser -> { + if (parser.currentToken() == JsonToken.VALUE_STRING) { + return LocalDateTime.parse(parser.getText()); + } + throw outOfRangeException(parser); + }) .build(); } + interface ParserExtractor { + W extract(JsonParser parser) throws IOException; + } + @AutoValue public abstract static class ValidatingValueExtractor implements ValueExtractor { @@ -258,6 +414,8 @@ public abstract static class ValidatingValueExtractor implements ValueExtract abstract Function extractor(); + abstract ParserExtractor parserExtractor(); + static Builder builder() { return new AutoValue_RowJsonValueExtractors_ValidatingValueExtractor.Builder<>(); } @@ -275,12 +433,19 @@ public W extractValue(JsonNode value) { return extractor().apply(value); } + @Override + public W extractValue(JsonParser parser) throws IOException { + return parserExtractor().extract(parser); + } + @AutoValue.Builder abstract static class Builder { abstract Builder setValidator(Predicate validator); abstract Builder setExtractor(Function extractor); + abstract Builder setParserExtractor(ParserExtractor parserExtractor); + abstract ValidatingValueExtractor build(); } } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/JsonToRowTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/JsonToRowTest.java index 79918930e098..3c66937da0ea 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/JsonToRowTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/JsonToRowTest.java @@ -17,6 +17,9 @@ */ package org.apache.beam.sdk.transforms; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; + import java.io.Serializable; import java.util.Collection; import java.util.stream.Collectors; @@ -317,6 +320,8 @@ public void testParsesErrorWithErrorMsgRowsDeadLetterWithCustomFieldNames() thro PCollection personRows = results.getResults(); PCollection errorsWithMsg = results.getFailedToParseLines(); + assertThat(errorsWithMsg.getSchema(), equalTo(customSchema)); + PAssert.that(personRows).containsInAnyOrder(PERSON_ROWS); PAssert.that(errorsWithMsg) diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/RowJsonTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/RowJsonTest.java index fda8fea5480c..6476ff4d78ae 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/RowJsonTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/RowJsonTest.java @@ -34,6 +34,7 @@ import java.util.Map; import org.apache.beam.sdk.schemas.Schema; import org.apache.beam.sdk.schemas.Schema.FieldType; +import org.apache.beam.sdk.schemas.logicaltypes.EnumerationType; import org.apache.beam.sdk.schemas.logicaltypes.VariableString; import org.apache.beam.sdk.util.RowJson.RowJsonDeserializer; import org.apache.beam.sdk.util.RowJson.RowJsonDeserializer.NullBehavior; @@ -134,15 +135,18 @@ private static Object[] makeFlatRowTestCase() { } private static Object[] makeLogicalTypeTestCase() { + EnumerationType enumType = EnumerationType.create("ZERO", "ONE", "TWO"); Schema schema = Schema.builder() .addLogicalTypeField("f_passThroughString", VariableString.of(10)) + .addLogicalTypeField("f_enum", enumType) .build(); // fixed string will do padding - String rowString = "{\n" + "\"f_passThroughString\" : \"hello\"\n" + "}"; + String rowString = + "{\n" + "\"f_passThroughString\" : \"hello\",\n" + "\"f_enum\" : 1\n" + "}"; - Row expectedRow = Row.withSchema(schema).addValues("hello").build(); + Row expectedRow = Row.withSchema(schema).addValues("hello", enumType.valueOf("ONE")).build(); return new Object[] {"Logical Types", schema, rowString, expectedRow}; }