Enhance Python Timestamp to be precision-variable up to nanos, and map it to Timestamp logical type

.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Direct.json

@@ -1,4 +1,4 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "modification": 17
+  "modification": 1
 }

.github/trigger_files/beam_PostCommit_Python_Xlang_IO_Direct.json

@@ -1,4 +1,4 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "modification": 7
+  "modification": 1
 }

CHANGES.md

@@ -73,6 +73,9 @@
 * (Java) Supported acknowledge mode for JmsIO ([#39253](https://github.com/apache/beam/issues/39253)).
 * (Python) Staged files directory is now automatically added to `sys.path` on the Python SDK worker at startup. This makes Python files provided via the '--files_to_stage' pipeline option importable in the pipeline code and makes it easier to initialize Python SDK harness at startup via the `--beam_plugins` pipeline option. For more information, see the [Staging Individual Files](https://beam.apache.org/documentation/sdks/python-pipeline-dependencies/#staging-files) section of the dependency management docs. This behavior can be disabled by passing the '--experiments=no_staged_dir_in_sys_path' pipeline option ([#39431](https://github.com/apache/beam/issues/39431)).
 * (Python) Added `equal_to_approx`, an `assert_that` matcher that compares numeric pipeline outputs with a configurable tolerance ([#18028](https://github.com/apache/beam/issues/18028)).
+* (Python) `Timestamp` now supports variable subsecond precision, up to nanoseconds. The portable
+  `beam:logical_type:timestamp:v1` logical type now maps to Python's `Timestamp` ([#39344](https://github.com/apache/beam/issues/39344)).
+* X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
 
 ## Breaking Changes
 

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

@@ -626,3 +626,34 @@ examples:
   "\u0001\u0003\u0061\u0062\u0063" : "abc"
   "\u0001\u000a\u006d\u006f\u0072\u0065\u0020\u0062\u0079\u0074\u0065\u0073" : "more bytes"
   "\u0000" : null
+
+---
+
+coder:
+  urn: "beam:coder:row:v1"
+  # f_timestamp: logical(timestamp:v1, precision=3, representation row<seconds:int64, subseconds:int16>)
+  payload: "\x0a\x88\x01\x0a\x0bf_timestamp\x1ay:w\x0a\x1ebeam:logical_type:timestamp:v1\x1aK2I\x0aG\x0a\x0d\x0a\x07seconds\x1a\x02\x10\x04\x0a\x10\x0a\x0asubseconds\x1a\x02\x10\x02\x12$5f2fbf6c-8b8e-4a30-a9be-e0f65f0f79c3\"\x02\x10\x03*\x04\x0a\x02\x18\x03\x12$9c7c1b6c-2a5e-4c53-9b32-91b71c1f8ac3"
+examples:
+  "\x01\x00\x02\x00\xb6\x95\xd5\xf9\x05\x03\xe7": {f_timestamp: {seconds: 1597328054, subseconds: 999}}
+  "\x01\x00\x02\x00\xfe\xff\xff\xff\xff\xff\xff\xff\xff\x01\x01\xf4": {f_timestamp: {seconds: -2, subseconds: 500}}
+  "\x01\x00\x02\x00\xb6\x95\xd5\xf9\x05\x00\x00": {f_timestamp: {seconds: 1597328054, subseconds: 0}}
+
+---
+
+coder:
+  urn: "beam:coder:row:v1"
+  # f_timestamp: logical(timestamp:v1, precision=6, representation row<seconds:int64, subseconds:int32>)
+  payload: "\x0a\x88\x01\x0a\x0bf_timestamp\x1ay:w\x0a\x1ebeam:logical_type:timestamp:v1\x1aK2I\x0aG\x0a\x0d\x0a\x07seconds\x1a\x02\x10\x04\x0a\x10\x0a\x0asubseconds\x1a\x02\x10\x03\x12$b8b3faf8-63f8-4b40-9f4d-3c1e83d4e07b\"\x02\x10\x03*\x04\x0a\x02\x18\x06\x12$e0a72e10-b214-4fd2-a4dc-fca0e7f7d924"
+examples:
+  "\x01\x00\x02\x00\xb6\x95\xd5\xf9\x05\xc0\xc4\x07": {f_timestamp: {seconds: 1597328054, subseconds: 123456}}
+  "\x01\x00\x02\x00\xfe\xff\xff\xff\xff\xff\xff\xff\xff\x01\xa0\xc2\x1e": {f_timestamp: {seconds: -2, subseconds: 500000}}
+
+---
+
+coder:
+  urn: "beam:coder:row:v1"
+  # f_timestamp: logical(timestamp:v1, precision=9, representation row<seconds:int64, subseconds:int32>)
+  payload: "\x0a\x88\x01\x0a\x0bf_timestamp\x1ay:w\x0a\x1ebeam:logical_type:timestamp:v1\x1aK2I\x0aG\x0a\x0d\x0a\x07seconds\x1a\x02\x10\x04\x0a\x10\x0a\x0asubseconds\x1a\x02\x10\x03\x12$3d6f7b81-2f13-4a41-8be3-4d8cd8f30f42\"\x02\x10\x03*\x04\x0a\x02\x18\x09\x12$cf03a3c1-4a3a-4f77-8e2b-6e5b19c19f27"
+examples:
+  "\x01\x00\x02\x00\xb6\x95\xd5\xf9\x05\x95\x9a\xef:": {f_timestamp: {seconds: 1597328054, subseconds: 123456789}}
+  "\x01\x00\x02\x00\xfe\xff\xff\xff\xff\xff\xff\xff\xff\x01\xff\x93\xeb\xdc\x03": {f_timestamp: {seconds: -2, subseconds: 999999999}}

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

@@ -234,6 +234,19 @@ message LogicalTypes {
     //     since the epoch.
     DATE = 8 [(org.apache.beam.model.pipeline.v1.beam_urn) =
       "beam:logical_type:date:v1"];
+
+    // A URN for Timestamp type
+    //   - Representation type: ROW<seconds: INT64, subseconds: INT16> for
+    //     precision < 5, otherwise ROW<seconds: INT64, subseconds: INT32>
+    //   - Argument type: INT32.
+    //     A timestamp without a timezone, represented as seconds since the
+    //     epoch plus a non-negative fraction of a second. The argument is
+    //     the precision: the number of decimal digits of the subseconds
+    //     field (e.g. 3 for milliseconds, 6 for micros, 9 for nanos).
+    //     subseconds stays non-negative even for timestamps before the epoch,
+    //     e.g. -1.5s at precision 6 is {seconds: -2, subseconds: 500000}.
+    TIMESTAMP = 9 [(org.apache.beam.model.pipeline.v1.beam_urn) =
+      "beam:logical_type:timestamp:v1"];
   }
 }
 

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

@@ -50,6 +50,7 @@
 import org.apache.beam.sdk.schemas.logicaltypes.MicrosInstant;
 import org.apache.beam.sdk.schemas.logicaltypes.PythonCallable;
 import org.apache.beam.sdk.schemas.logicaltypes.SchemaLogicalType;
+import org.apache.beam.sdk.schemas.logicaltypes.Timestamp;
 import org.apache.beam.sdk.schemas.logicaltypes.UnknownLogicalType;
 import org.apache.beam.sdk.schemas.logicaltypes.VariableBytes;
 import org.apache.beam.sdk.schemas.logicaltypes.VariableString;
@@ -115,6 +116,7 @@ private static String getLogicalTypeUrn(String identifier) {
           .put(FixedString.IDENTIFIER, FixedString.class)
           .put(VariableString.IDENTIFIER, VariableString.class)
           .put(Date.IDENTIFIER, Date.class)
+          .put(Timestamp.IDENTIFIER, Timestamp.class)
           .build();
 
   public static SchemaApi.Schema schemaToProto(Schema schema, boolean serializeLogicalType) {

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

@@ -22,6 +22,8 @@
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
 
 import java.time.Instant;
+import org.apache.beam.model.pipeline.v1.RunnerApi;
+import org.apache.beam.model.pipeline.v1.SchemaApi;
 import org.apache.beam.sdk.schemas.Schema;
 import org.apache.beam.sdk.values.Row;
 import org.checkerframework.checker.nullness.qual.NonNull;
@@ -56,7 +58,12 @@
  * }</pre>
  */
 public class Timestamp implements Schema.LogicalType<Instant, Row> {
-  public static final String IDENTIFIER = "beam:logical_type:timestamp:v1";
+  public static final String IDENTIFIER =
+      SchemaApi.LogicalTypes.Enum.TIMESTAMP
+          .getValueDescriptor()
+          .getOptions()
+          .getExtension(RunnerApi.beamUrn);
+
   static final int MIN_PRECISION = 0;
   static final int MAX_PRECISION = 9;
 

sdks/java/core/src/test/java/org/apache/beam/sdk/schemas/SchemaTranslationTest.java

@@ -51,6 +51,7 @@
 import org.apache.beam.sdk.schemas.logicaltypes.PythonCallable;
 import org.apache.beam.sdk.schemas.logicaltypes.SchemaLogicalType;
 import org.apache.beam.sdk.schemas.logicaltypes.SqlTypes;
+import org.apache.beam.sdk.schemas.logicaltypes.Timestamp;
 import org.apache.beam.sdk.schemas.logicaltypes.UnknownLogicalType;
 import org.apache.beam.sdk.schemas.logicaltypes.VariableBytes;
 import org.apache.beam.sdk.schemas.logicaltypes.VariableString;
@@ -460,6 +461,12 @@ public static Iterable<Schema.FieldType> data() {
           .add(FieldType.logicalType(FixedString.of(10)))
           .add(FieldType.logicalType(VariableString.of(10)))
           .add(FieldType.logicalType(FixedPrecisionNumeric.of(10)))
+          .add(FieldType.logicalType(Timestamp.MILLIS))
+          .add(FieldType.logicalType(Timestamp.MICROS))
+          .add(FieldType.logicalType(Timestamp.NANOS))
+          // INT16/INT32 subseconds representation boundary.
+          .add(FieldType.logicalType(Timestamp.of(4)))
+          .add(FieldType.logicalType(Timestamp.of(5)))
           .add(FieldType.logicalType(new PortableNullArgLogicalType()))
           .add(FieldType.logicalType(new NullArgumentLogicalType()))
           .build();

sdks/python/apache_beam/coders/row_coder_test.py

@@ -203,6 +203,60 @@ def test_create_row_coder_from_schema(self):
     for test_case in self.PEOPLE:
       self.assertEqual(test_case, coder.decode(coder.encode(test_case)))
 
+  def test_create_row_coder_from_schema_with_parameterized_timestamp(self):
+    # Simulates schemas produced by the Java SDK's parameterized Timestamp
+    # logical type. Subseconds field uses INT16 for
+    # precision < 5 and INT32 otherwise.
+    for precision, subseconds_type in [(3, schema_pb2.INT16),
+                                       (4, schema_pb2.INT16),
+                                       (5, schema_pb2.INT32),
+                                       (6, schema_pb2.INT32),
+                                       (9, schema_pb2.INT32)]:
+      schema = schema_pb2.Schema(
+          id="timestamp_precision_%d" % precision,
+          fields=[
+              schema_pb2.Field(
+                  name="event_time",
+                  type=schema_pb2.FieldType(
+                      logical_type=schema_pb2.LogicalType(
+                          urn="beam:logical_type:timestamp:v1",
+                          representation=schema_pb2.FieldType(
+                              row_type=schema_pb2.RowType(
+                                  schema=schema_pb2.Schema(
+                                      id="timestamp_repr_%d" % precision,
+                                      fields=[
+                                          schema_pb2.Field(
+                                              name="seconds",
+                                              type=schema_pb2.FieldType(
+                                                  atomic_type=schema_pb2.INT64)
+                                          ),
+                                          schema_pb2.Field(
+                                              name="subseconds",
+                                              type=schema_pb2.FieldType(
+                                                  atomic_type=subseconds_type)),
+                                      ]))),
+                          argument_type=schema_pb2.FieldType(
+                              atomic_type=schema_pb2.INT32),
+                          argument=schema_pb2.FieldValue(
+                              atomic_value=schema_pb2.AtomicTypeValue(
+                                  int32=precision))))),
+          ])
+      coder = RowCoder(schema)
+      row = named_tuple_from_schema(schema)
+
+      subseconds = 10**precision - 1
+      for value in [
+          Timestamp(seconds=1500000000,
+                    subseconds=subseconds,
+                    precision=precision),
+          Timestamp(seconds=-2, subseconds=subseconds, precision=precision),
+          Timestamp(seconds=1500000000, precision=precision),
+      ]:
+        test_case = row(event_time=value)
+        decoded = coder.decode(coder.encode(test_case))
+        self.assertEqual(test_case, decoded)
+        self.assertEqual(decoded.event_time.precision(), precision)
+
   def test_row_coder_negative_varint(self):
     schema = schema_pb2.Schema(
         id="negative",

sdks/python/apache_beam/portability/common_urns.py

@@ -93,3 +93,4 @@
 fixed_char = LogicalTypes.Enum.FIXED_CHAR
 var_char = LogicalTypes.Enum.VAR_CHAR
 date = LogicalTypes.Enum.DATE
+timestamp = LogicalTypes.Enum.TIMESTAMP

sdks/python/apache_beam/typehints/schemas.py

@@ -977,6 +977,92 @@ def _from_typing(cls, typ):
     'MicrosInstantRepresentation', [('seconds', np.int64),
                                     ('micros', np.int64)])
 
+ParameterizedTimestampRepresentation = NamedTuple(
+    'ParameterizedTimestampRepresentation', [('seconds', np.int64),
+                                             ('subseconds', np.int32)])
+
+# The subseconds field is INT16 for precision < 5
+_TIMESTAMP_SHORT_PRECISION_LIMIT = 5
+ParameterizedTimestampShortRepresentation = NamedTuple(
+    'ParameterizedTimestampShortRepresentation', [('seconds', np.int64),
+                                                  ('subseconds', np.int16)])
+
+
+@LogicalType._register_internal
+class ParameterizedTimestamp(LogicalType[Timestamp,
+                                         ParameterizedTimestampRepresentation,
+                                         np.int32]):
+  """Timestamp logical type parameterized by subsecond precision.
+
+  The argument is the precision: the number of decimal digits used to
+  represent the fraction of a second, e.g. 3 for milliseconds, 6 for
+  microseconds, 9 for nanoseconds.
+
+  Values are represented as a row of ``seconds`` (INT64, floored seconds
+  since the epoch) and ``subseconds`` (units of 10**-precision seconds,
+  always in ``[0, 10**precision)``.
+  ``subseconds`` is an INT16 field for precision < 5 and an INT32
+  field otherwise.
+
+  Note: Timestamp originating from Python to xlang still defaults to
+  MicrosInstant for backwards compatibility. To override the mapping of
+  Timestamp to this logical type, re-register using
+  :func:`~LogicalType.register_logical_type(ParameterizedTimestamp)`.
+  """
+  def __init__(self, precision: int = Timestamp.MICROS_PRECISION) -> None:
+    # The argument arrives as np.int32 when decoded from a schema proto.
+    precision = int(precision)
+    if not 0 <= precision <= Timestamp.NANOS_PRECISION:
+      raise ValueError(
+          'Timestamp precision must be between 0 and %d (inclusive), '
+          'but was %d.' % (Timestamp.NANOS_PRECISION, precision))
+    self._precision = precision
+
+  @classmethod
+  def urn(cls):
+    return common_urns.timestamp.urn
+
+  def representation_type(self) -> type:  # type: ignore[override]
+    # Unlike other logical types, the representation depends on the
+    # argument, so this is an instance method rather than a classmethod.
+    if self._precision < _TIMESTAMP_SHORT_PRECISION_LIMIT:
+      return ParameterizedTimestampShortRepresentation
+    return ParameterizedTimestampRepresentation
+
+  @classmethod
+  def language_type(cls):
+    return Timestamp
+
+  def to_representation_type(self, value: Timestamp):
+    # Verify that the value can be represented exactly at this type's precision
+    if value.precision() != self._precision:
+      value = value.to_precision(self._precision)
+    return self.representation_type()(value.seconds(), value.subseconds())
+
+  def to_language_type(self, value) -> Timestamp:
+    subseconds = int(value.subseconds)
+    # Match Java's toInputType: out-of-range subseconds indicate data
+    # corruption or a precision mismatch.
+    if not 0 <= subseconds < 10**self._precision:
+      raise ValueError(
+          'Invalid subseconds %d for Timestamp with precision %d.' %
+          (subseconds, self._precision))
+    return Timestamp(
+        seconds=int(value.seconds),
+        subseconds=subseconds,
+        precision=self._precision)
+
+  @classmethod
+  def argument_type(cls):
+    return np.int32
+
+  def argument(self):
+    return self._precision
+
+  @classmethod
+  def _from_typing(cls, typ):
+    return cls()
+
 
 @LogicalType._register_internal
 class MillisInstant(NoArgumentLogicalType[Timestamp, np.int64]):
@@ -1019,9 +1105,11 @@ def to_language_type(self, value):
     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.
+# Make sure MicrosInstant is registered after MillisInstant and
+# ParameterizedTimestamp so that it overwrites the mapping of Timestamp
+# language type representation choice: plain Timestamp typehints keep their
+# historical micros_instant encoding and do not lose microsecond precision
+# inside the python sdk.
 @LogicalType._register_internal
 class MicrosInstant(NoArgumentLogicalType[Timestamp,
                                           MicrosInstantRepresentation]):