Preserve original logical type in Beam->Calcite->Beam trip

Author: Abacn

sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/utils/CalciteUtils.java

@@ -22,6 +22,7 @@
 import java.lang.reflect.Type;
 import java.util.Date;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.stream.IntStream;
 import org.apache.beam.sdk.schemas.Schema;
 import org.apache.beam.sdk.schemas.Schema.FieldType;
@@ -169,6 +170,12 @@ public static boolean isStringType(FieldType fieldType) {
           FieldType.DATETIME, SqlTypeName.TIMESTAMP,
           FieldType.STRING, SqlTypeName.VARCHAR);
 
+  // Associating FieldType to generated RelDataType objects for Beam logical types. Used for
+  // recovering the original type in output schema after full Beam FieldType->Calcite Type->Beam
+  // FieldType trip
+  private static final Map<RelDataType, FieldType> LOGICAL_TYPE_REL_DATA_MAPPING =
+      new ConcurrentHashMap<>();
+
   /** Generate {@link Schema} from {@code RelDataType} which is used to create table. */
   public static Schema toSchema(RelDataType tableInfo) {
     return tableInfo.getFieldList().stream().map(CalciteUtils::toField).collect(Schema.toSchema());
@@ -254,6 +261,9 @@ public static Schema.Field toField(String name, RelDataType calciteType) {
   }
 
   public static FieldType toFieldType(RelDataType calciteType) {
+    if (LOGICAL_TYPE_REL_DATA_MAPPING.containsKey(calciteType)) {
+      return LOGICAL_TYPE_REL_DATA_MAPPING.get(calciteType);
+    }
     switch (calciteType.getSqlTypeName()) {
       case ARRAY:
       case MULTISET:
@@ -317,10 +327,27 @@ public static RelDataType toRelDataType(RelDataTypeFactory dataTypeFactory, Fiel
         return toCalciteRowType(schema, dataTypeFactory);
       case LOGICAL_TYPE:
         Schema.LogicalType<?, ?> logicalType = fieldType.getLogicalType();
+        RelDataType relDataType;
         if (logicalType instanceof PassThroughLogicalType) {
-          return toRelDataType(dataTypeFactory, logicalType.getBaseType());
+          relDataType =
+              toRelDataType(
+                  dataTypeFactory, logicalType.getBaseType().withNullable(fieldType.getNullable()));
+        } else {
+          relDataType = dataTypeFactory.createSqlType(toSqlTypeName(fieldType));
         }
-        return dataTypeFactory.createSqlType(toSqlTypeName(fieldType));
+        // For backward-compatibility, exclude logical types registered in
+        // CALCITE_TO_BEAM_TYPE_MAPPING,
+        // e.g., primitive types, date time types, etc.
+        SqlTypeName typeName = relDataType.getSqlTypeName();
+        if (typeName != null && !CALCITE_TO_BEAM_TYPE_MAPPING.containsKey(typeName)) {
+          // register both nullable and non-nullable variants.
+          boolean flipNullable = !relDataType.isNullable();
+          LOGICAL_TYPE_REL_DATA_MAPPING.put(relDataType, fieldType);
+          LOGICAL_TYPE_REL_DATA_MAPPING.put(
+              dataTypeFactory.createTypeWithNullability(relDataType, flipNullable),
+              fieldType.withNullable(flipNullable));
+        }
+        return relDataType;
       default:
         return dataTypeFactory.createSqlType(toSqlTypeName(fieldType));
     }

sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamComplexTypeTest.java

@@ -17,6 +17,8 @@
  */
 package org.apache.beam.sdk.extensions.sql;
 
+import static org.junit.Assert.assertEquals;
+
 import java.nio.charset.StandardCharsets;
 import java.time.LocalDate;
 import java.time.LocalDateTime;
@@ -792,6 +794,7 @@ public void testUnknownLogicalType() {
             .apply(SqlTransform.query("select * from PCOLLECTION"));
 
     PAssert.that(outputRow).containsInAnyOrder(inputRow);
+    assertEquals(inputRow.getSchema(), outputRow.getSchema());
     pipeline.run().waitUntilFinish(Duration.standardMinutes(1));
   }
 }

sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/utils/CalciteUtilsTest.java

@@ -24,6 +24,8 @@
 import java.util.Map;
 import java.util.stream.Collectors;
 import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.logicaltypes.PassThroughLogicalType;
+import org.apache.beam.sdk.values.Row;
 import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.type.RelDataType;
 import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.type.RelDataTypeFactory;
 import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.type.RelDataTypeSystem;
@@ -184,9 +186,8 @@ public void testToRelDataTypeWithRowBackedLogicalType() {
     Schema nestedSchema = Schema.builder().addField("nested_f1", Schema.FieldType.INT32).build();
     Schema.FieldType rowType = Schema.FieldType.row(nestedSchema);
 
-    Schema.LogicalType<org.apache.beam.sdk.values.Row, org.apache.beam.sdk.values.Row> logicalType =
-        new org.apache.beam.sdk.schemas.logicaltypes.PassThroughLogicalType<
-            org.apache.beam.sdk.values.Row>(
+    Schema.LogicalType<Row, Row> logicalType =
+        new PassThroughLogicalType<Row>(
             "RowBackedLogicalType", Schema.FieldType.STRING, "", rowType) {};
 
     Schema.FieldType logicalFieldType = Schema.FieldType.logicalType(logicalType);

sdks/python/apache_beam/coders/coder_impl.py

@@ -841,7 +841,7 @@ def encode_to_stream(self, value, out, nested):
     out.write_bigendian_int16(value)
 
   def decode_from_stream(self, in_stream, nested):
-    # type: (create_InputStream, bool) -> float
+    # type: (create_InputStream, bool) -> int
     return in_stream.read_bigendian_int16()
 
   def estimate_size(self, unused_value, nested=False):
@@ -857,12 +857,12 @@ def encode_to_stream(self, value, out, nested):
     out.write_byte(value)
 
   def decode_from_stream(self, in_stream, nested):
-    # type: (create_InputStream, bool) -> float
+    # type: (create_InputStream, bool) -> int
     return in_stream.read_byte()
 
   def estimate_size(self, unused_value, nested=False):
     # type: (Any, bool) -> int
-    # A short is encoded as 2 bytes, regardless of nesting.
+    # A byte is encoded as 1 byte, regardless of nesting.
     return 1
 
 

sdks/python/apache_beam/transforms/sql_test.py

@@ -162,19 +162,6 @@ def test_row(self):
           | SqlTransform("SELECT a*a as s, LENGTH(b) AS c FROM PCOLLECTION"))
       assert_that(out, equal_to([(1, 1), (4, 1), (100, 2)]))
 
-  @staticmethod
-  def recover_to_python_type(input):
-    fields = []
-    for field in input:
-      print(field)
-      if hasattr(field, 'type_byte') and hasattr(field, 'payload'):
-        obj = coders.FastPrimitivesCoder().decode(
-            field.type_byte.to_bytes() + field.payload)
-        fields.append(obj)
-      else:
-        fields.append(field)
-    return tuple(fields)
-
   def test_row_user_type(self):
     with TestPipeline() as p:
       out = (
@@ -183,9 +170,7 @@ def test_row_user_type(self):
               UserTypeRow(1, Aribitrary("abc"), -1j),
           ])
           | SqlTransform("SELECT arb, complex FROM PCOLLECTION")
-          # TODO: recover to user type. Currently pipeline can run,
-          # but elements returned back to Python are generated rows
-          | beam.Map(self.recover_to_python_type))
+          | beam.Map(tuple))
       assert_that(
           out,
           equal_to([(Aribitrary(1.0), 1 + 2.5j), (Aribitrary("abc"), -1j)]))

sdks/python/apache_beam/typehints/schemas.py

@@ -106,6 +106,8 @@
 from apache_beam.utils.timestamp import Timestamp
 
 PYTHON_ANY_URN = "beam:logical:pythonsdk_any:v1"
+_PYTHON_ANY_FIELD_TYPE_BYTE = "_pythonsdk_any_type_byte"
+_PYTHON_ANY_FIELD_PAYLOAD = "payload"
 _SCHEMA_OPTION_STATIC_ENCODING = "beam:option:row:static_encoding"
 
 # Bi-directional mappings
@@ -257,6 +259,7 @@ def schema_field(
 
 
 def _python_any_schema_pb2():
+  # A portable schema matches FastPrimitivesCoder encoded values
   return schema_pb2.FieldType(
       logical_type=schema_pb2.LogicalType(
           urn=PYTHON_ANY_URN,
@@ -266,11 +269,11 @@ def _python_any_schema_pb2():
                   schema=schema_pb2.Schema(
                       fields=[
                           schema_pb2.Field(
-                              name="type_byte",
+                              name=_PYTHON_ANY_FIELD_TYPE_BYTE,
                               type=schema_pb2.FieldType(
                                   atomic_type=schema_pb2.BYTE, nullable=False)),
                           schema_pb2.Field(
-                              name="payload",
+                              name=_PYTHON_ANY_FIELD_PAYLOAD,
                               type=schema_pb2.FieldType(
                                   atomic_type=schema_pb2.BYTES, nullable=False))
                       ],