[BEAM-9424] Allow grouping by LogicalType

apache · Mar 3, 2020 · 73576ac · 73576ac
1 parent d924f72
commit 73576ac
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 0 deletions.
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/SchemaCoder.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/SchemaCoder.java
@@ -133,6 +133,8 @@ public static <T> Coder<T> coderForFieldType(FieldType fieldType) {
             MapCoder.of(
                 coderForFieldType(fieldType.getMapKeyType()),
                 coderForFieldType(fieldType.getMapValueType()));
+      case LOGICAL_TYPE:
+        return coderForFieldType(fieldType.getLogicalType().getBaseType());
       default:
         return (Coder<T>) CODER_MAP.get(fieldType.getTypeName());
     }

diff --git a/.../extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamComplexTypeTest.java b/.../extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamComplexTypeTest.java
@@ -373,6 +373,41 @@ public void testNullInnerRow() {
     pipeline.run().waitUntilFinish(Duration.standardMinutes(2));
   }
 
+  @Test
+  public void testLogicalTypes() {
+    DateTime dateTime = DateTime.parse("2020-02-02T00:00:00");
+
+    Schema inputRowSchema =
+        Schema.builder()
+            .addField("timeTypeField", FieldType.logicalType(new DummySqlTimeType()))
+            .addField("dateTypeField", FieldType.logicalType(new DummySqlDateType()))
+            .build();
+
+    Row row =
+        Row.withSchema(inputRowSchema)
+            .addValues(dateTime.getMillis(), dateTime.getMillis())
+            .build();
+
+    Schema outputRowSchema =
+        Schema.builder()
+            .addField("timeTypeField", FieldType.DATETIME)
+            .addNullableField("dateTypeField", FieldType.DATETIME)
+            .build();
+
+    PCollection<Row> outputRow =
+        pipeline
+            .apply(Create.of(row))
+            .setRowSchema(outputRowSchema)
+            .apply(
+                SqlTransform.query(
+                    "SELECT timeTypeField, dateTypeField FROM PCOLLECTION GROUP BY timeTypeField, dateTypeField"));
+
+    PAssert.that(outputRow)
+        .containsInAnyOrder(Row.withSchema(outputRowSchema).addValues(dateTime, dateTime).build());
+
+    pipeline.run().waitUntilFinish(Duration.standardMinutes(2));
+  }
+
   private static class DummySqlTimeType implements Schema.LogicalType<Long, Instant> {
     @Override
     public String getIdentifier() {