Support Timestamp type in xlang JDBC Read and Write (#22561)

Support Timestamp type in xlang JDBC Read and Write

CHANGES.md

@@ -77,6 +77,8 @@
 ## Bugfixes
 
 * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
+* Fixed Python cross-language JDBC IO Connector cannot read or write rows containing Timestamp type values [19817](https://github.com/apache/beam/issues/19817).
+
 ## Known Issues
 
 * ([#X](https://github.com/apache/beam/issues/X)).

model/fn-execution/src/main/resources/org/apache/beam/model/fnexecution/v1/standard_coders.yaml

@@ -391,7 +391,7 @@ examples:
 # import typing
 # import apache_beam as beam
 # class Test(typing.NamedTuple):
-# f_map: typing.Mapping[str,int]
+# f_map: typing.Mapping[str,typing.Optional[int]]
 # schema = beam.typehints.schemas.named_tuple_to_schema(Test)
 # coder = beam.coders.row_coder.RowCoder(schema)
 # print("payload = %s" % schema.SerializeToString())
@@ -453,6 +453,13 @@ coder:
 examples:
  "\x03\x00\x02\x00\xb6\x95\xd5\xf9\x05\xc0\xc4\x07\x1b2020-08-13T14:14:14.123456Z\xc0\xf7\x85\xda\xae\x98\xeb\x02": {f_timestamp: {seconds: 1597328054, micros: 123456}, f_string: "2020-08-13T14:14:14.123456Z", f_int: 1597328054123456}
 
+coder:
+ urn: "beam:coder:row:v1"
+ # f_timestamp: logical(millis_instant), f_string: string, f_int: int64
+ payload: "\n:\n\x0bf_timestamp\x1a+:)\n#beam:logical_type:millis_instant:v1\x1a\x02\x10\x04\n\x0e\n\x08f_string\x1a\x02\x10\x07\n\x0b\n\x05f_int\x1a\x02\x10\x04\x12$80be749a-5700-4ede-89d8-dd9a4433a3f8"
+examples:
+ "\x03\x00\x80\x00\x01s\xe8+\xd7k\x182020-08-13T14:14:14.123Z\xeb\xae\xaf\xc1\xbe.": {f_timestamp: -9223370439526721685, f_string: "2020-08-13T14:14:14.123Z", f_int: 1597328054123}
+
 ---
 
 coder:

model/pipeline/src/main/proto/org/apache/beam/model/pipeline/v1/schema.proto

@@ -130,6 +130,15 @@ message LogicalTypes {
  // amount of time since the epoch.
  MICROS_INSTANT = 1 [(org.apache.beam.model.pipeline.v1.beam_urn) =
  "beam:logical_type:micros_instant:v1"];
+
+ // A URN for MillisInstant type
+ // - Representation type: INT64
+ // - A timestamp without a timezone represented by the number of
+ // milliseconds since the epoch. The INT64 value is encoded with
+ // big-endian shifted such that lexicographic ordering of the bytes
+ // corresponds to chronological order.
+ MILLIS_INSTANT = 2 [(org.apache.beam.model.pipeline.v1.beam_urn) =
+ "beam:logical_type:millis_instant:v1"];
  }
 }
 

runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/wire/CommonCoderTest.java

@@ -422,6 +422,9 @@ private static Object parseField(Object value, Schema.FieldType fieldType) {
  return (String) value;
  case BOOLEAN:
  return (Boolean) value;
+ case DATETIME:
+ // convert shifted millis to epoch millis as in InstantCoder
+ return new Instant((Long) value + -9223372036854775808L);
  case BYTES:
  // extract String as byte[]
  return ((String) value).getBytes(StandardCharsets.ISO_8859_1);
@@ -465,7 +468,7 @@ private static Object parseField(Object value, Schema.FieldType fieldType) {
  return fieldType
  .getLogicalType()
  .toInputType(parseField(value, fieldType.getLogicalType().getBaseType()));
- default: // DECIMAL, DATETIME
+ default: // DECIMAL
  throw new IllegalArgumentException("Unsupported type name: " + fieldType.getTypeName());
  }
  }

sdks/go/test/regression/coders/fromyaml/fromyaml.go

@@ -51,6 +51,7 @@ var unimplementedCoders = map[string]bool{
 var filteredCases = []struct{ filter, reason string }{
  {"logical", "BEAM-9615: Support logical types"},
  {"30ea5a25-dcd8-4cdb-abeb-5332d15ab4b9", "https://github.com/apache/beam/issues/21206: Support encoding position."},
+ {"80be749a-5700-4ede-89d8-dd9a4433a3f8", "https://github.com/apache/beam/issues/19817: Support millis_instant."},
 }
 
 // Coder is a representation a serialized beam coder.

sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/SchemaTranslation.java

@@ -28,6 +28,7 @@
 import java.util.Objects;
 import java.util.UUID;
 import java.util.stream.Collectors;
+import org.apache.beam.model.pipeline.v1.RunnerApi;
 import org.apache.beam.model.pipeline.v1.SchemaApi;
 import org.apache.beam.model.pipeline.v1.SchemaApi.ArrayTypeValue;
 import org.apache.beam.model.pipeline.v1.SchemaApi.AtomicTypeValue;
@@ -64,9 +65,13 @@
 })
 public class SchemaTranslation {
 
- private static final String URN_BEAM_LOGICAL_DATETIME = "beam:logical_type:datetime:v1";
  private static final String URN_BEAM_LOGICAL_DECIMAL = "beam:logical_type:decimal:v1";
  private static final String URN_BEAM_LOGICAL_JAVASDK = "beam:logical_type:javasdk:v1";
+ private static final String URN_BEAM_LOGICAL_MILLIS_INSTANT =
+ SchemaApi.LogicalTypes.Enum.MILLIS_INSTANT
+ .getValueDescriptor()
+ .getOptions()
+ .getExtension(RunnerApi.beamUrn);
 
  // TODO(https://github.com/apache/beam/issues/19715): Populate this with a LogicalTypeRegistrar,
  // which includes a way to construct
@@ -198,7 +203,7 @@ private static SchemaApi.FieldType fieldTypeToProto(
  case DATETIME:
  builder.setLogicalType(
  SchemaApi.LogicalType.newBuilder()
- .setUrn(URN_BEAM_LOGICAL_DATETIME)
+ .setUrn(URN_BEAM_LOGICAL_MILLIS_INSTANT)
  .setRepresentation(fieldTypeToProto(FieldType.INT64, serializeLogicalType))
  .build());
  break;
@@ -358,7 +363,7 @@ private static FieldType fieldTypeFromProtoWithoutNullable(SchemaApi.FieldType p
  }
  // Special-case for DATETIME and DECIMAL which are logical types in portable representation,
  // but not yet in Java. (https://github.com/apache/beam/issues/19817)
- if (urn.equals(URN_BEAM_LOGICAL_DATETIME)) {
+ if (urn.equals(URN_BEAM_LOGICAL_MILLIS_INSTANT)) {
  return FieldType.DATETIME;
  } else if (urn.equals(URN_BEAM_LOGICAL_DECIMAL)) {
  return FieldType.DECIMAL;

sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcUtil.java

@@ -39,6 +39,7 @@
 import org.apache.beam.sdk.io.jdbc.JdbcIO.ReadWithPartitions;
 import org.apache.beam.sdk.io.jdbc.JdbcIO.RowMapper;
 import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.logicaltypes.MicrosInstant;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.Row;
@@ -189,6 +190,18 @@ static JdbcIO.PreparedStatementSetCaller getPreparedStatementSetCaller(
  }
 
  String logicalTypeName = fieldType.getLogicalType().getIdentifier();
+
+ if (logicalTypeName.equals(MicrosInstant.IDENTIFIER)) {
+ // Process timestamp of MicrosInstant kind, which should only be passed from other type
+ // systems such as SQL and other Beam SDKs.
+ return (element, ps, i, fieldWithIndex) -> {
+ // MicrosInstant uses native java.time.Instant instead of joda.Instant.
+ java.time.Instant value =
+ element.getLogicalTypeValue(fieldWithIndex.getIndex(), java.time.Instant.class);
+ ps.setTimestamp(i + 1, value == null ? null : new Timestamp(value.toEpochMilli()));
+ };
+ }
+
  JDBCType jdbcType = JDBCType.valueOf(logicalTypeName);
  switch (jdbcType) {
  case DATE:

sdks/python/apache_beam/coders/row_coder.py

@@ -32,6 +32,7 @@
 from apache_beam.coders.coders import NullableCoder
 from apache_beam.coders.coders import SinglePrecisionFloatCoder
 from apache_beam.coders.coders import StrUtf8Coder
+from apache_beam.coders.coders import TimestampCoder
 from apache_beam.coders.coders import VarIntCoder
 from apache_beam.portability import common_urns
 from apache_beam.portability.api import schema_pb2
@@ -168,10 +169,15 @@ def _nonnull_coder_from_type(field_type):
  _coder_from_type(field_type.map_type.key_type),
  _coder_from_type(field_type.map_type.value_type))
  elif type_info == "logical_type":
- # Special case for the Any logical type. Just use the default coder for an
- # unknown Python object.
  if field_type.logical_type.urn == PYTHON_ANY_URN:
+ # Special case for the Any logical type. Just use the default coder for an
+ # unknown Python object.
  return typecoders.registry.get_coder(object)
+ elif field_type.logical_type.urn == common_urns.millis_instant.urn:
+ # Special case for millis instant logical type used to handle Java sdk's
+ # millis Instant. It explicitly uses TimestampCoder which deals with fix
+ # length 8-bytes big-endian-long instead of VarInt coder.
+ return TimestampCoder()
 
  logical_type = LogicalType.from_runner_api(field_type.logical_type)
  return LogicalTypeCoder(

sdks/python/apache_beam/io/external/xlang_jdbcio_it_test.py

@@ -17,12 +17,15 @@
 
 # pytype: skip-file
 
+import datetime
 import logging
+import time
 import typing
 import unittest
 from typing import Callable
 from typing import Union
 
+import pytz
 from parameterized import parameterized
 
 import apache_beam as beam
@@ -33,6 +36,9 @@
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.util import assert_that
 from apache_beam.testing.util import equal_to
+from apache_beam.typehints.schemas import LogicalType
+from apache_beam.typehints.schemas import MillisInstant
+from apache_beam.utils.timestamp import Timestamp
 
 # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports
 try:
@@ -53,19 +59,14 @@
 
 JdbcReadTestRow = typing.NamedTuple(
  "JdbcReadTestRow",
- [
- ("f_int", int),
- ],
+ [("f_int", int), ("f_timestamp", Timestamp)],
 )
 coders.registry.register_coder(JdbcReadTestRow, coders.RowCoder)
 
 JdbcWriteTestRow = typing.NamedTuple(
  "JdbcWriteTestRow",
- [
- ("f_id", int),
- ("f_real", float),
- ("f_string", str),
- ],
+ [("f_id", int), ("f_real", float), ("f_string", str),
+ ("f_timestamp", Timestamp)],
 )
 coders.registry.register_coder(JdbcWriteTestRow, coders.RowCoder)
 
@@ -127,10 +128,15 @@ def test_xlang_jdbc_write(self, database):
  self._setUpTestCase(container_init, db_string, driver)
  table_name = 'jdbc_external_test_write'
  self.engine.execute(
- "CREATE TABLE {}(f_id INTEGER, f_real FLOAT, f_string VARCHAR(100))".
- format(table_name))
+ "CREATE TABLE {}(f_id INTEGER, f_real FLOAT, f_string VARCHAR(100), f_timestamp TIMESTAMP(3))" # pylint: disable=line-too-long
+ .format(table_name))
  inserted_rows = [
- JdbcWriteTestRow(i, i + 0.1, 'Test{}'.format(i))
+ JdbcWriteTestRow(
+ i,
+ i + 0.1,
+ 'Test{}'.format(i),
+ # In alignment with Java Instant which supports milli precision.
+ Timestamp.of(seconds=round(time.time(), 3)))
  for i in range(ROW_COUNT)
  ]
 
@@ -152,7 +158,11 @@ def test_xlang_jdbc_write(self, database):
 
  fetched_data = self.engine.execute("SELECT * FROM {}".format(table_name))
  fetched_rows = [
- JdbcWriteTestRow(int(row[0]), float(row[1]), str(row[2]))
+ JdbcWriteTestRow(
+ int(row[0]),
+ float(row[1]),
+ str(row[2]),
+ Timestamp.from_utc_datetime(row[3].replace(tzinfo=pytz.UTC)))
  for row in fetched_data
  ]
 
@@ -168,10 +178,25 @@ def test_xlang_jdbc_read(self, database):
  CrossLanguageJdbcIOTest.DB_CONTAINER_CLASSPATH_STRING[database])
  self._setUpTestCase(container_init, db_string, driver)
  table_name = 'jdbc_external_test_read'
- self.engine.execute("CREATE TABLE {}(f_int INTEGER)".format(table_name))
+ self.engine.execute(
+ "CREATE TABLE {}(f_int INTEGER, f_timestamp TIMESTAMP)".format(
+ table_name))
 
+ all_timestamps = []
  for i in range(ROW_COUNT):
- self.engine.execute("INSERT INTO {} VALUES({})".format(table_name, i))
+ # prepare timestamp
+ strtime = Timestamp.now().to_utc_datetime().strftime('%Y-%m-%dT%H:%M:%S')
+ dttime = datetime.datetime.strptime(
+ strtime, '%Y-%m-%dT%H:%M:%S').replace(tzinfo=pytz.UTC)
+ all_timestamps.append(Timestamp.from_utc_datetime(dttime))
+
+ # write records using sqlalchemy engine
+ self.engine.execute(
+ "INSERT INTO {} VALUES({},'{}')".format(table_name, i, strtime))
+
+ # Register MillisInstant logical type to override the mapping from Timestamp
+ # originally handled by MicrosInstant.
+ LogicalType.register_logical_type(MillisInstant)
 
  with TestPipeline() as p:
  p.not_use_test_runner_api = True
@@ -188,7 +213,10 @@ def test_xlang_jdbc_read(self, database):
  classpath=classpath))
 
  assert_that(
- result, equal_to([JdbcReadTestRow(i) for i in range(ROW_COUNT)]))
+ result,
+ equal_to([
+ JdbcReadTestRow(i, all_timestamps[i]) for i in range(ROW_COUNT)
+ ]))
 
  # Creating a container with testcontainers sometimes raises ReadTimeout
  # error. In java there are 2 retries set by default.

sdks/python/apache_beam/portability/common_urns.py

@@ -80,4 +80,5 @@
 java_class_lookup = ExpansionMethods.Enum.JAVA_CLASS_LOOKUP
 
 micros_instant = LogicalTypes.Enum.MICROS_INSTANT
+millis_instant = LogicalTypes.Enum.MILLIS_INSTANT
 python_callable = LogicalTypes.Enum.PYTHON_CALLABLE

sdks/python/apache_beam/typehints/schemas.py

@@ -34,6 +34,7 @@
  bytes <-----> BYTES
  ByteString ------> BYTES
  Timestamp <-----> LogicalType(urn="beam:logical_type:micros_instant:v1")
+ Timestamp <------ LogicalType(urn="beam:logical_type:millis_instant:v1")
  Mapping <-----> MapType
  Sequence <-----> ArrayType
  NamedTuple <-----> RowType
@@ -571,13 +572,13 @@ def argument(self):
  """Return the argument for this instance of the LogicalType."""
  raise NotImplementedError()
 
- def to_representation_type(value):
+ def to_representation_type(self, value):
  # type: (LanguageT) -> RepresentationT
 
  """Convert an instance of LanguageT to RepresentationT."""
  raise NotImplementedError()
 
- def to_language_type(value):
+ def to_language_type(self, value):
  # type: (RepresentationT) -> LanguageT
 
  """Convert an instance of RepresentationT to LanguageT."""
@@ -587,6 +588,7 @@ def to_language_type(value):
  def register_logical_type(cls, logical_type_cls):
  """Register an implementation of LogicalType."""
  cls._known_logical_types.add(logical_type_cls.urn(), logical_type_cls)
+ return logical_type_cls
 
  @classmethod
  def from_typing(cls, typ):
@@ -655,9 +657,54 @@ def _from_typing(cls, typ):
  ('micros', np.int64)])
 
 
+@LogicalType.register_logical_type
+class MillisInstant(NoArgumentLogicalType[Timestamp, np.int64]):
+ """Millisecond-precision instant logical type handles values consistent with
+ that encoded by ``InstantCoder`` in the Java SDK.
+
+ This class handles :class:`apache_beam.utils.timestamp.Timestamp` language
+ type as :class:`MicrosInstant`, but it only provides millisecond precision,
+ because it is aimed to handle data encoded by Java sdk's InstantCoder which
+ has same precision level.
+
+ Timestamp is handled by `MicrosInstant` by default. In some scenario, such as
+ read from cross-language transform with rows containing InstantCoder encoded
+ timestamps, one may need to override the mapping of Timetamp to MillisInstant.
+ To do this, re-register this class with
+ :func:`~LogicalType.register_logical_type`.
+ """
+ @classmethod
+ def representation_type(cls):
+ # type: () -> type
+ return np.int64
+
+ @classmethod
+ def urn(cls):
+ return common_urns.millis_instant.urn
+
+ @classmethod
+ def language_type(cls):
+ return Timestamp
+
+ def to_language_type(self, value):
+ # type: (np.int64) -> Timestamp
+
+ # value shifted as in apache_beams.coders.coder_impl.TimestampCoderImpl
+ if value < 0:
+ millis = int(value) + (1 << 63)
+ else:
+ millis = int(value) - (1 << 63)
+
+ return Timestamp(micros=millis * 1000)
+
+
+# Make sure MicrosInstant is registered after MillisInstant so that it
+# overwrites the mapping of Timestamp language type representation choice and
+# thus does not lose microsecond precision inside python sdk.
 @LogicalType.register_logical_type
 class MicrosInstant(NoArgumentLogicalType[Timestamp,
  MicrosInstantRepresentation]):
+ """Microsecond-precision instant logical type that handles ``Timestamp``."""
  @classmethod
  def urn(cls):
  return common_urns.micros_instant.urn
@@ -683,6 +730,7 @@ def to_language_type(self, value):
 
 @LogicalType.register_logical_type
 class PythonCallable(NoArgumentLogicalType[PythonCallableWithSource, str]):
+ """A logical type for PythonCallableSource objects."""
  @classmethod
  def urn(cls):
  return common_urns.python_callable.urn