Skip to content

Commit 5307c71

Browse files
damccormdamccorm
authored
Support MAP type in RowJson and fix datetime parsing with spaces (#38865)
1 parent 03f4f79 commit 5307c71

3 files changed

Lines changed: 97 additions & 3 deletions

File tree

‎sdks/java/core/src/main/java/org/apache/beam/sdk/util/RowJson.java‎

Lines changed: 64 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import static org.apache.beam.sdk.schemas.Schema.TypeName.INT16;
2727
import static org.apache.beam.sdk.schemas.Schema.TypeName.INT32;
2828
import static org.apache.beam.sdk.schemas.Schema.TypeName.INT64;
29+
import static org.apache.beam.sdk.schemas.Schema.TypeName.MAP;
2930
import static org.apache.beam.sdk.schemas.Schema.TypeName.STRING;
3031
import static org.apache.beam.sdk.util.RowJsonValueExtractors.booleanValueExtractor;
3132
import static org.apache.beam.sdk.util.RowJsonValueExtractors.byteValueExtractor;
@@ -57,6 +58,9 @@
5758
import java.time.LocalDate;
5859
import java.time.LocalDateTime;
5960
import java.time.LocalTime;
61+
import java.util.HashMap;
62+
import java.util.Iterator;
63+
import java.util.Map;
6064
import java.util.stream.Stream;
6165
import java.util.stream.StreamSupport;
6266
import org.apache.beam.sdk.schemas.Schema;
@@ -95,7 +99,8 @@
9599
})
96100
public class RowJson {
97101
private static final ImmutableSet<TypeName> SUPPORTED_TYPES =
98-
ImmutableSet.of(BYTE, INT16, INT32, INT64, FLOAT, DOUBLE, BOOLEAN, STRING, DECIMAL, DATETIME);
102+
ImmutableSet.of(
103+
BYTE, INT16, INT32, INT64, FLOAT, DOUBLE, BOOLEAN, STRING, DECIMAL, DATETIME, MAP);
99104
private static final ImmutableSet<String> KNOWN_LOGICAL_TYPE_IDENTIFIERS =
100105
ImmutableSet.of(
101106
SqlTypes.DATE.getIdentifier(),
@@ -160,6 +165,14 @@ private static ImmutableList<UnsupportedField> findUnsupportedFields(
160165
return findUnsupportedFields(fieldType.getCollectionElementType(), fieldName + "[]");
161166
}
162167

168+
if (fieldTypeName.isMapType()) {
169+
if (!STRING.equals(fieldType.getMapKeyType().getTypeName())) {
170+
return ImmutableList.of(
171+
new UnsupportedField(fieldName + ".key", fieldType.getMapKeyType().getTypeName()));
172+
}
173+
return findUnsupportedFields(fieldType.getMapValueType(), fieldName + "{}");
174+
}
175+
163176
if (fieldTypeName.isLogicalType()) {
164177
if (KNOWN_LOGICAL_TYPE_IDENTIFIERS.contains(fieldType.getLogicalType().getIdentifier())) {
165178
return ImmutableList.of();
@@ -303,6 +316,10 @@ private Object extractJsonNodeValue(FieldValue fieldValue) {
303316
return jsonArrayToList(fieldValue);
304317
}
305318

319+
if (fieldValue.isMapType()) {
320+
return jsonObjectToMap(fieldValue);
321+
}
322+
306323
if (fieldValue.typeName().isLogicalType()) {
307324
String identifier = fieldValue.type().getLogicalType().getIdentifier();
308325
if (SqlTypes.DATE.getIdentifier().equals(identifier)) {
@@ -365,6 +382,32 @@ private Object jsonArrayToList(FieldValue arrayFieldValue) {
365382
.collect(toImmutableList());
366383
}
367384

385+
private Map<String, Object> jsonObjectToMap(FieldValue mapFieldValue) {
386+
if (!mapFieldValue.isJsonObject()) {
387+
throw new UnsupportedRowJsonException(
388+
"Expected JSON object for field '"
389+
+ mapFieldValue.name()
390+
+ "'. Instead got "
391+
+ mapFieldValue.jsonNodeType().name());
392+
}
393+
394+
Map<String, Object> result = new HashMap<>();
395+
Iterator<Map.Entry<String, JsonNode>> fields = mapFieldValue.jsonValue().fields();
396+
while (fields.hasNext()) {
397+
Map.Entry<String, JsonNode> field = fields.next();
398+
String key = field.getKey();
399+
JsonNode value = field.getValue();
400+
401+
Object extractedValue =
402+
extractJsonNodeValue(
403+
FieldValue.of(
404+
mapFieldValue.name() + "['" + key + "']", mapFieldValue.mapValueType(), value));
405+
406+
result.put(key, extractedValue);
407+
}
408+
return result;
409+
}
410+
368411
private static Object extractJsonPrimitiveValue(FieldValue fieldValue) {
369412
try {
370413
return JSON_VALUE_GETTERS.get(fieldValue.typeName()).extractValue(fieldValue.jsonValue());
@@ -440,6 +483,18 @@ Schema rowSchema() {
440483
return type().getRowSchema();
441484
}
442485

486+
boolean isMapType() {
487+
return TypeName.MAP.equals(type().getTypeName());
488+
}
489+
490+
FieldType mapKeyType() {
491+
return type().getMapKeyType();
492+
}
493+
494+
FieldType mapValueType() {
495+
return type().getMapValueType();
496+
}
497+
443498
static FieldValue of(String name, FieldType type, JsonNode jsonValue) {
444499
return new AutoValue_RowJson_RowJsonDeserializer_FieldValue(name, type, jsonValue);
445500
}
@@ -538,6 +593,14 @@ private void writeValue(JsonGenerator gen, FieldType type, Object value) throws
538593
case ROW:
539594
writeRow((Row) value, type.getRowSchema(), gen);
540595
break;
596+
case MAP:
597+
gen.writeStartObject();
598+
for (Map.Entry<Object, Object> entry : ((Map<Object, Object>) value).entrySet()) {
599+
gen.writeFieldName(entry.getKey().toString());
600+
writeValue(gen, type.getMapValueType(), entry.getValue());
601+
}
602+
gen.writeEndObject();
603+
break;
541604
case LOGICAL_TYPE:
542605
String identifier = type.getLogicalType().getIdentifier();
543606
if (SqlTypes.DATE.getIdentifier().equals(identifier)) {

‎sdks/java/core/src/main/java/org/apache/beam/sdk/util/RowJsonValueExtractors.java‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -189,7 +189,14 @@ static ValueExtractor<BigDecimal> decimalValueExtractor() {
189189
*/
190190
static ValueExtractor<DateTime> datetimeValueExtractor() {
191191
return ValidatingValueExtractor.<DateTime>builder()
192-
.setExtractor(jsonNode -> DateTime.parse(jsonNode.textValue()))
192+
.setExtractor(
193+
jsonNode -> {
194+
String text = jsonNode.textValue();
195+
if (text.contains(" ")) {
196+
text = text.replace(' ', 'T');
197+
}
198+
return DateTime.parse(text);
199+
})
193200
.setValidator(JsonNode::isTextual)
194201
.build();
195202
}

‎sdks/java/core/src/test/java/org/apache/beam/sdk/util/RowJsonTest.java‎

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@
3030
import java.math.BigDecimal;
3131
import java.util.Arrays;
3232
import java.util.Collection;
33+
import java.util.HashMap;
34+
import java.util.Map;
3335
import org.apache.beam.sdk.schemas.Schema;
3436
import org.apache.beam.sdk.schemas.Schema.FieldType;
3537
import org.apache.beam.sdk.schemas.logicaltypes.VariableString;
@@ -80,7 +82,8 @@ public static Collection<Object[]> data() {
8082
makeArrayOfArraysTestCase(),
8183
makeNestedRowTestCase(),
8284
makeDoublyNestedRowTestCase(),
83-
makeNullsTestCase());
85+
makeNullsTestCase(),
86+
makeMapFieldTestCase());
8487
}
8588

8689
private static Object[] makeFlatRowTestCase() {
@@ -244,6 +247,21 @@ private static Object[] makeNullsTestCase() {
244247
return new Object[] {"Nulls", schema, rowString, expectedRow};
245248
}
246249

250+
private static Object[] makeMapFieldTestCase() {
251+
Schema schema =
252+
Schema.builder().addMapField("f_map", FieldType.STRING, FieldType.INT32).build();
253+
254+
String rowString = "{\n" + "\"f_map\" : {\"key1\": 1, \"key2\": 2}\n" + "}";
255+
256+
Map<String, Integer> expectedMap = new HashMap<>();
257+
expectedMap.put("key1", 1);
258+
expectedMap.put("key2", 2);
259+
260+
Row expectedRow = Row.withSchema(schema).addValues(expectedMap).build();
261+
262+
return new Object[] {"Map field", schema, rowString, expectedRow};
263+
}
264+
247265
@Test
248266
public void testDeserialize() throws IOException {
249267
Row parsedRow =
@@ -564,6 +582,12 @@ public void testSupportedDatetimeConversions() throws Exception {
564582
testSupportedConversion(FieldType.DATETIME, quoted(DATETIME_STRING), DATETIME_VALUE);
565583
}
566584

585+
@Test
586+
public void testSupportedDatetimeWithSpaceConversions() throws Exception {
587+
String datetimeWithSpace = DATETIME_STRING.replace('T', ' ');
588+
testSupportedConversion(FieldType.DATETIME, quoted(datetimeWithSpace), DATETIME_VALUE);
589+
}
590+
567591
private void testSupportedConversion(
568592
FieldType fieldType, String jsonFieldValue, Object expectedRowFieldValue) throws Exception {
569593

0 commit comments

Comments
 (0)