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..82def4b105ff 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 @@ -321,7 +321,9 @@ private Object extractJsonNodeValue(FieldValue fieldValue) { } if (fieldValue.typeName().isLogicalType()) { - String identifier = fieldValue.type().getLogicalType().getIdentifier(); + Schema.LogicalType logicalType = + (Schema.LogicalType) fieldValue.type().getLogicalType(); + String identifier = logicalType.getIdentifier(); if (SqlTypes.DATE.getIdentifier().equals(identifier)) { return dateValueExtractor().extractValue(fieldValue.jsonValue()); } else if (SqlTypes.TIME.getIdentifier().equals(identifier)) { @@ -329,11 +331,10 @@ private Object extractJsonNodeValue(FieldValue fieldValue) { } else if (SqlTypes.DATETIME.getIdentifier().equals(identifier)) { return localDatetimeValueExtractor().extractValue(fieldValue.jsonValue()); } else { - return extractJsonNodeValue( - FieldValue.of( - fieldValue.name(), - fieldValue.type().getLogicalType().getBaseType(), - fieldValue.jsonValue())); + return logicalType.toInputType( + extractJsonNodeValue( + FieldValue.of( + fieldValue.name(), logicalType.getBaseType(), fieldValue.jsonValue()))); } } @@ -602,7 +603,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 +613,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/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}; }