From a041c280c172de0cdfff73803f112be5b35316aa Mon Sep 17 00:00:00 2001 From: Andrew Pilloud Date: Wed, 5 Feb 2020 17:20:05 -0800 Subject: [PATCH] [BEAM-8630] Add logical types, make public --- .../sdk/extensions/sql/zetasql/ZetaSqlUtils.java | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlUtils.java b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlUtils.java index efbbad1993a8..a4641f024e80 100644 --- a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlUtils.java +++ b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlUtils.java @@ -29,6 +29,8 @@ import java.util.List; import java.util.stream.Collectors; import org.apache.beam.sdk.annotations.Internal; +import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils.DateType; +import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils.TimeType; import org.apache.beam.sdk.schemas.Schema; import org.apache.beam.sdk.schemas.Schema.Field; import org.apache.beam.sdk.schemas.Schema.FieldType; @@ -68,6 +70,16 @@ public static Type beamFieldTypeToZetaSqlType(FieldType fieldType) { return createZetaSqlArrayTypeFromBeamElementFieldType(fieldType.getCollectionElementType()); case ROW: return createZetaSqlStructTypeFromBeamSchema(fieldType.getRowSchema()); + case LOGICAL_TYPE: + switch (fieldType.getLogicalType().getIdentifier()) { + case DateType.IDENTIFIER: + return TypeFactory.createSimpleType(TypeKind.TYPE_DATE); + case TimeType.IDENTIFIER: + return TypeFactory.createSimpleType(TypeKind.TYPE_TIME); + default: + throw new IllegalArgumentException( + "Unsupported Beam logical type: " + fieldType.getLogicalType().getIdentifier()); + } default: throw new UnsupportedOperationException( "Unsupported Beam fieldType: " + fieldType.getTypeName()); @@ -79,7 +91,7 @@ private static ArrayType createZetaSqlArrayTypeFromBeamElementFieldType( return TypeFactory.createArrayType(beamFieldTypeToZetaSqlType(elementFieldType)); } - private static StructType createZetaSqlStructTypeFromBeamSchema(Schema schema) { + public static StructType createZetaSqlStructTypeFromBeamSchema(Schema schema) { return TypeFactory.createStructType( schema.getFields().stream() .map(ZetaSqlUtils::beamFieldToZetaSqlStructField) @@ -139,7 +151,7 @@ private static Value javaListToZetaSqlArrayValue(List elements, FieldTyp createZetaSqlArrayTypeFromBeamElementFieldType(elementType), values); } - private static Value beamRowToZetaSqlStructValue(Row row, Schema schema) { + public static Value beamRowToZetaSqlStructValue(Row row, Schema schema) { List values = new ArrayList<>(row.getFieldCount()); for (int i = 0; i < row.getFieldCount(); i++) {