diff --git a/sdks/python/apache_beam/coders/coder_impl.pxd b/sdks/python/apache_beam/coders/coder_impl.pxd index e64177e6fd34..6330a9f7b4db 100644 --- a/sdks/python/apache_beam/coders/coder_impl.pxd +++ b/sdks/python/apache_beam/coders/coder_impl.pxd @@ -297,6 +297,7 @@ cdef class RowCoderImpl(StreamCoderImpl): cdef bint encoding_positions_are_trivial cdef list components cdef bint has_nullable_fields + cdef bint static_encoding @cython.locals(i=int, nvals=libc.stdint.int64_t, running=int, component_coder=CoderImpl, null_mask=bytes, null_mask_c=char_ptr) diff --git a/sdks/python/apache_beam/coders/coder_impl.py b/sdks/python/apache_beam/coders/coder_impl.py index 0bded25e05d2..1f5f54dace83 100644 --- a/sdks/python/apache_beam/coders/coder_impl.py +++ b/sdks/python/apache_beam/coders/coder_impl.py @@ -61,6 +61,7 @@ from apache_beam.coders.avro_record import AvroRecord from apache_beam.internal import cloudpickle_pickler from apache_beam.internal.cloudpickle import cloudpickle +from apache_beam.typehints.schemas import _SCHEMA_OPTION_STATIC_ENCODING from apache_beam.typehints.schemas import named_tuple_from_schema from apache_beam.utils import proto_utils from apache_beam.utils import windowed_value @@ -1955,31 +1956,35 @@ def __init__(self, schema, components): for i in self.encoding_positions) self.has_nullable_fields = any( field.type.nullable for field in self.schema.fields) + self.static_encoding = any( + opt.name == _SCHEMA_OPTION_STATIC_ENCODING and + opt.value.atomic_value.boolean for opt in self.schema.options) def encode_to_stream(self, value, out, nested): - out.write_var_int64(self.num_fields) attrs = [getattr(value, name) for name in self.field_names] - if self.has_nullable_fields: - any_nulls = False - for attr in attrs: - if attr is None: - any_nulls = True - break - if any_nulls: - out.write_var_int64((self.num_fields + 7) // 8) - # Pack the bits, little-endian, in consecutive bytes. - running = 0 - for i, attr in enumerate(attrs): - if i and i % 8 == 0: - out.write_byte(running) - running = 0 - running |= (attr is None) << (i % 8) - out.write_byte(running) + if not self.static_encoding: + out.write_var_int64(self.num_fields) + if self.has_nullable_fields: + any_nulls = False + for attr in attrs: + if attr is None: + any_nulls = True + break + if any_nulls: + out.write_var_int64((self.num_fields + 7) // 8) + # Pack the bits, little-endian, in consecutive bytes. + running = 0 + for i, attr in enumerate(attrs): + if i and i % 8 == 0: + out.write_byte(running) + running = 0 + running |= (attr is None) << (i % 8) + out.write_byte(running) + else: + out.write_byte(0) else: out.write_byte(0) - else: - out.write_byte(0) for i in range(self.num_fields): if not self.encoding_positions_are_trivial: @@ -2021,13 +2026,14 @@ def encode_batch_to_stream(self, columns: Dict[str, np.ndarray], out): has_null_bits = np.zeros((n, ), dtype=np.uint8) for k in range(n): - out.write_var_int64(self.num_fields) - if has_null_bits[k]: - out.write_byte(null_bits_len) - for i in range(null_bits_len): - out.write_byte(null_bits[k, i]) - else: - out.write_byte(0) + if not self.static_encoding: + out.write_var_int64(self.num_fields) + if has_null_bits[k]: + out.write_byte(null_bits_len) + for i in range(null_bits_len): + out.write_byte(null_bits[k, i]) + else: + out.write_byte(0) for i in range(self.num_fields): if not self.encoding_positions_are_trivial: i = self.encoding_positions_argsort[i] @@ -2040,11 +2046,15 @@ def encode_batch_to_stream(self, columns: Dict[str, np.ndarray], out): cython.cast(RowColumnEncoder, attrs[i]).encode_to_stream(k, out) def decode_from_stream(self, in_stream, nested): - nvals = in_stream.read_var_int64() - null_mask_len = in_stream.read_var_int64() - if null_mask_len: - # pylint: disable=unused-variable - null_mask_c = null_mask_py = in_stream.read(null_mask_len) + if self.static_encoding: + nvals = self.num_fields + null_mask_len = 0 + else: + nvals = in_stream.read_var_int64() + null_mask_len = in_stream.read_var_int64() + if null_mask_len: + # pylint: disable=unused-variable + null_mask_c = null_mask_py = in_stream.read(null_mask_len) # Note that if this coder's schema has *fewer* attributes than the encoded # value, we just need to ignore the additional values, which will occur @@ -2078,11 +2088,15 @@ def decode_batch_from_stream(self, dest: Dict[str, np.ndarray], in_stream): for k in range(n): if in_stream.size() == 0: break - nvals = in_stream.read_var_int64() - null_mask_len = in_stream.read_var_int64() - if null_mask_len: - # pylint: disable=unused-variable - null_mask_c = null_mask = in_stream.read(null_mask_len) + if self.static_encoding: + nvals = self.num_fields + null_mask_len = 0 + else: + nvals = in_stream.read_var_int64() + null_mask_len = in_stream.read_var_int64() + if null_mask_len: + # pylint: disable=unused-variable + null_mask_c = null_mask = in_stream.read(null_mask_len) for i in range(min(self.num_fields, nvals)): if not self.encoding_positions_are_trivial: diff --git a/sdks/python/apache_beam/coders/row_coder.py b/sdks/python/apache_beam/coders/row_coder.py index 29f85ba9cbf4..4c57de357b4c 100644 --- a/sdks/python/apache_beam/coders/row_coder.py +++ b/sdks/python/apache_beam/coders/row_coder.py @@ -208,5 +208,13 @@ def _create_impl(self): def is_deterministic(self): return self.representation_coder.is_deterministic() + def as_deterministic_coder(self, step_label, error_message=None): + if self.is_deterministic(): + return self + return LogicalTypeCoder( + self.logical_type, + self.representation_coder.as_deterministic_coder( + step_label, error_message)) + def to_type_hint(self): return self.logical_type.language_type() diff --git a/sdks/python/apache_beam/coders/row_coder_test.py b/sdks/python/apache_beam/coders/row_coder_test.py index 28170721c28b..3cc5dbf7f726 100644 --- a/sdks/python/apache_beam/coders/row_coder_test.py +++ b/sdks/python/apache_beam/coders/row_coder_test.py @@ -33,6 +33,8 @@ 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 _SCHEMA_OPTION_STATIC_ENCODING +from apache_beam.typehints.schemas import _static_encoding_option_pb2 from apache_beam.typehints.schemas import named_tuple_from_schema from apache_beam.typehints.schemas import typing_to_runner_api from apache_beam.utils.timestamp import Timestamp @@ -510,6 +512,47 @@ def test_batch_encode_decode(self): for field, a in columnar.items(): assert_array_equal(a[:n], dest[field][:n]) + def test_row_coder_with_tuples(self): + class TupleRecord(typing.NamedTuple): + key: str + fixed_tuple: typing.Tuple[str, int] + var_tuple: typing.Tuple[int, ...] + homo_tuple: typing.Tuple[str, str] + + coder = RowCoder(typing_to_runner_api(TupleRecord).row_type.schema) + record = TupleRecord("k1", ("hello", 42), (1, 2, 3), ("a", "b")) + encoded = coder.encode(record) + decoded = coder.decode(encoded) + + self.assertEqual(record, decoded) + self.assertIsInstance(decoded.fixed_tuple, tuple) + self.assertIsInstance(decoded.var_tuple, tuple) + self.assertIsInstance(decoded.homo_tuple, tuple) + # Verify hashability as dict keys + d = {decoded.homo_tuple: "val1", decoded.fixed_tuple: "val2"} + self.assertEqual(d[("a", "b")], "val1") + self.assertEqual(d[("hello", 42)], "val2") + + def test_static_encoding(self): + schema = schema_pb2.Schema( + fields=[ + schema_pb2.Field( + name="f_int32", + type=schema_pb2.FieldType(atomic_type=schema_pb2.INT32)), + schema_pb2.Field( + name="f_string", + type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING)), + ], + options=[_static_encoding_option_pb2()]) + RowType = named_tuple_from_schema(schema) + row = RowType(f_int32=42, f_string="hello world!") + coder = RowCoder(schema) + encoded = coder.encode(row) + # VarInt(42) = 1 byte, String("hello world!") = 1 byte len + 12 chars = 13 bytes. + # Total = 14 bytes (0 envelope overhead, matching TupleCoder) + self.assertEqual(14, len(encoded)) + self.assertEqual(row, coder.decode(encoded)) + if __name__ == "__main__": logging.getLogger().setLevel(logging.INFO) diff --git a/sdks/python/apache_beam/transforms/sql_test.py b/sdks/python/apache_beam/transforms/sql_test.py index 6649e210685a..0d3ac699a053 100644 --- a/sdks/python/apache_beam/transforms/sql_test.py +++ b/sdks/python/apache_beam/transforms/sql_test.py @@ -57,6 +57,10 @@ def __eq__(self, other): "UserTypeRow", [("id", int), ("arb", Aribitrary), ("complex", complex)]) coders.registry.register_coder(UserTypeRow, coders.RowCoder) +TupleRow = typing.NamedTuple( + "TupleRow", [("id", int), ("coords", typing.Tuple[str, int])]) +coders.registry.register_coder(TupleRow, coders.RowCoder) + @pytest.mark.xlang_sql_expansion_service @unittest.skipIf( @@ -66,19 +70,12 @@ def __eq__(self, other): class SqlTransformTest(unittest.TestCase): """Tests that exercise the cross-language SqlTransform (implemented in java). - Note this test must be executed with pipeline options that run jobs on a local - job server. The easiest way to accomplish this is to run the - `validatesCrossLanguageRunnerPythonUsingSql` gradle target for a particular - job server, which will start the runner and job server for you. For example, - `:runners:flink:1.13:job-server:validatesCrossLanguageRunnerPythonUsingSql` to - test on Flink 1.13. - - Alternatively, you may be able to iterate faster if you run the tests directly - using a runner like `FlinkRunner`, which can start a local Flink cluster and - job server for you: - $ pip install -e './sdks/python[gcp,test]' + To run these tests locally using PrismRunner, build the SQL expansion service + and prism binary first: + $ ./gradlew :sdks:java:extensions:sql:expansion-service:shadowJar + $ ./gradlew :runners:prism:build $ pytest apache_beam/transforms/sql_test.py \\ - --test-pipeline-options="--runner=FlinkRunner" + --test-pipeline-options="--runner=PrismRunner" """ _multiprocess_can_split_ = True @@ -229,6 +226,19 @@ def test_sql_ddl_set_option(self): # Verify the output matches the query (unaffected by the SET DDL) assert_that(out, equal_to([(3, 30)])) + def test_tuple_field(self): + with TestPipeline() as p: + out = ( + p + | beam.Create([ + TupleRow(1, ("foo", 100)), + TupleRow(2, ("bar", 200)), + ]) + | SqlTransform( + "SELECT t.id, t.coords.f0 AS `name`, t.coords.f1 AS `val`, t.coords " + "FROM PCOLLECTION t WHERE t.coords.f1 > 150")) + assert_that(out, equal_to([(2, "bar", 200, ("bar", 200))])) + if __name__ == "__main__": logging.getLogger().setLevel(logging.INFO) diff --git a/sdks/python/apache_beam/typehints/native_type_compatibility.py b/sdks/python/apache_beam/typehints/native_type_compatibility.py index 7f28d57b2f81..c63ba80f9dea 100644 --- a/sdks/python/apache_beam/typehints/native_type_compatibility.py +++ b/sdks/python/apache_beam/typehints/native_type_compatibility.py @@ -125,14 +125,16 @@ def _safe_issubclass(derived, parent): issubclass(derived, parent), or False if a TypeError was raised. """ try: - return issubclass(derived, parent) + if issubclass(derived, parent): + return True except (TypeError, AttributeError): - if hasattr(derived, '__origin__'): - try: - return issubclass(derived.__origin__, parent) - except TypeError: - pass - return False + pass + if hasattr(derived, '__origin__') and derived.__origin__ is not None: + try: + return issubclass(derived.__origin__, parent) + except (TypeError, AttributeError): + pass + return False def _match_issubclass(match_against): diff --git a/sdks/python/apache_beam/typehints/native_type_compatibility_test.py b/sdks/python/apache_beam/typehints/native_type_compatibility_test.py index 33d6051afc7a..8694cd08430b 100644 --- a/sdks/python/apache_beam/typehints/native_type_compatibility_test.py +++ b/sdks/python/apache_beam/typehints/native_type_compatibility_test.py @@ -41,6 +41,7 @@ from apache_beam.typehints.native_type_compatibility import convert_typing_to_builtin from apache_beam.typehints.native_type_compatibility import is_any from apache_beam.typehints.native_type_compatibility import match_dataclass_for_row +from apache_beam.typehints.native_type_compatibility import _safe_issubclass _TestNamedTuple = typing.NamedTuple( '_TestNamedTuple', [('age', int), ('name', bytes)]) @@ -573,6 +574,19 @@ class NonFrozenDC: self.assertEqual( compat_version == "2.73.0", match_dataclass_for_row(NonFrozenDC)) + def test_safe_issubclass(self): + # In Python <= 3.12, issubclass(types.GenericAlias, tuple) returns False + # directly without raising TypeError, whereas in Python >= 3.13 it raises + # TypeError. _safe_issubclass inspects __origin__ to ensure types.GenericAlias + # (e.g. tuple[...], list[...]) is recognized consistently across Python versions. + self.assertTrue(_safe_issubclass(tuple[int, str], tuple)) + self.assertTrue(_safe_issubclass(tuple[int, str], (str, tuple))) + self.assertTrue(_safe_issubclass(list[int], list)) + self.assertTrue(_safe_issubclass(typing.Tuple[int, str], tuple)) + self.assertTrue(_safe_issubclass(typing.List[int], list)) + self.assertFalse(_safe_issubclass(int, tuple)) + self.assertFalse(_safe_issubclass(typing.Union[int, str], tuple)) + if __name__ == '__main__': unittest.main() diff --git a/sdks/python/apache_beam/typehints/row_type_test.py b/sdks/python/apache_beam/typehints/row_type_test.py index 30bda0cd98ba..f8b1b01b4927 100644 --- a/sdks/python/apache_beam/typehints/row_type_test.py +++ b/sdks/python/apache_beam/typehints/row_type_test.py @@ -193,6 +193,18 @@ class MyDataClass: row_type.GeneratedClassRowTypeConstraint( fields=[('id', int), ('name', str), ('name_hash', int)])) + def test_tuple_preserved_in_row_after_serialization(self): + with TestPipeline() as p: + res = ( + p + | beam.Create([{"a": 1, "b": 2, "h": "h1"}]) + | beam.GroupBy( # group_by with custom field generates a Beam Row + row_field=lambda x: (x["a"], x["b"])) + | beam.MapTuple( + lambda k, vs: (type(k.row_field), k.row_field))) + + assert_that(res, equal_to([(tuple, (1, 2))])) + if __name__ == '__main__': unittest.main() diff --git a/sdks/python/apache_beam/typehints/schemas.py b/sdks/python/apache_beam/typehints/schemas.py index 80bee60ec952..6e80960ddf46 100644 --- a/sdks/python/apache_beam/typehints/schemas.py +++ b/sdks/python/apache_beam/typehints/schemas.py @@ -20,28 +20,31 @@ Imposes a mapping between common Python types and Beam portable schemas (https://s.apache.org/beam-schemas):: - Python Schema - np.int8 <-----> BYTE - np.int16 <-----> INT16 - np.int32 <-----> INT32 - np.int64 <-----> INT64 - int ------> INT64 - np.float32 <-----> FLOAT - np.float64 <-----> DOUBLE - float ------> DOUBLE - bool <-----> BOOLEAN - str <-----> STRING - bytes <-----> BYTES - ByteString ------> BYTES - Timestamp <-----> LogicalType(urn="beam:logical_type:micros_instant:v1") - datetime.date <---> LogicalType(urn="beam:logical_type:date:v1") - Decimal <-----> LogicalType(urn="beam:logical_type:fixed_decimal:v1") - Mapping <-----> MapType - Sequence <-----> ArrayType - NamedTuple <-----> RowType - beam.Row ------> RowType - -One direction mapping of Python types from Beam portable schemas: + Python Schema + np.int8 <-----> BYTE + np.int16 <-----> INT16 + np.int32 <-----> INT32 + np.int64 <-----> INT64 + int ------> INT64 + np.float32 <-----> FLOAT + np.float64 <-----> DOUBLE + float ------> DOUBLE + bool <-----> BOOLEAN + str <-----> STRING + bytes <-----> BYTES + ByteString ------> BYTES + Timestamp <-----> LogicalType(urn="beam:logical_type:micros_instant:v1") + datetime.date <-----> LogicalType(urn="beam:logical_type:date:v1") + datetime.time <-----> LogicalType(urn="beam:logical_type:time:v1") + Decimal <-----> LogicalType(urn="beam:logical_type:fixed_decimal:v1") + Tuple[T, ...] <-----> LogicalType(urn="beam:logical_type:var_tuple:v1") + Tuple[T1, T2] <-----> LogicalType(urn="beam:logical_type:fixed_tuple:v1") + Mapping <-----> MapType + Sequence <-----> ArrayType + NamedTuple <-----> RowType + beam.Row ------> RowType + +One direction mapping of Python types from Beam portable schemas:: bytes <------ LogicalType(urn="beam:logical_type:fixed_bytes:v1") @@ -110,6 +113,8 @@ _PYTHON_ANY_FIELD_TYPE_BYTE = "_pythonsdk_any_type_byte" _PYTHON_ANY_FIELD_PAYLOAD = "payload" _SCHEMA_OPTION_STATIC_ENCODING = "beam:option:row:static_encoding" +FIXED_TUPLE_URN = "beam:logical_type:fixed_tuple:v1" +VAR_TUPLE_URN = "beam:logical_type:var_tuple:v1" # Bi-directional mappings _PRIMITIVES = ( @@ -259,6 +264,14 @@ def schema_field( description=description) +def _static_encoding_option_pb2() -> schema_pb2.Option: + return schema_pb2.Option( + name=_SCHEMA_OPTION_STATIC_ENCODING, + type=schema_pb2.FieldType(atomic_type=schema_pb2.BOOLEAN), + value=schema_pb2.FieldValue( + atomic_value=schema_pb2.AtomicTypeValue(boolean=True))) + + def _python_any_schema_pb2(has_repr): # A portable schema matches FastPrimitivesCoder encoded values if has_repr: @@ -276,15 +289,7 @@ def _python_any_schema_pb2(has_repr): type=schema_pb2.FieldType( atomic_type=schema_pb2.BYTES, nullable=False)) ], - options=[ - schema_pb2.Option( - name=_SCHEMA_OPTION_STATIC_ENCODING, - type=schema_pb2.FieldType( - atomic_type=schema_pb2.BOOLEAN), - value=schema_pb2.FieldValue( - atomic_value=schema_pb2.AtomicTypeValue( - boolean=True))) - ]))) if has_repr else None + options=[_static_encoding_option_pb2()]))) if has_repr else None else: representation = None @@ -374,7 +379,17 @@ def typing_to_runner_api(self, type_: type) -> schema_pb2.FieldType: element_type=schema_pb2.FieldType( atomic_type=PRIMITIVE_TO_ATOMIC_TYPE[int]))) - elif _safe_issubclass(type_, Sequence) and not _safe_issubclass(type_, str): + elif _safe_issubclass(type_, tuple) and not match_is_named_tuple(type_): + arg_types = _get_args(type_) + if len(arg_types) == 2 and arg_types[1] is Ellipsis: # Tuple[typ, ...] + return self.typing_to_runner_api(VarTupleLogicalType(arg_types[0])) + elif len(arg_types) > 0: # Tuple[typ1, typ2, ...] + return self.typing_to_runner_api(FixedTupleLogicalType(arg_types)) + else: # tuple of unknown type, just fallback to Any + return _python_any_schema_pb2(has_repr=True) + + elif _safe_issubclass( + type_, Sequence) and not _safe_issubclass(type_, (str, tuple)): arg_types = _get_args(type_) if len(arg_types) > 0: element_type = self.typing_to_runner_api(arg_types[0]) @@ -386,7 +401,8 @@ def typing_to_runner_api(self, type_: type) -> schema_pb2.FieldType: return schema_pb2.FieldType( map_type=schema_pb2.MapType(key_type=key_type, value_type=value_type)) - elif _safe_issubclass(type_, Iterable) and not _safe_issubclass(type_, str): + elif _safe_issubclass( + type_, Iterable) and not _safe_issubclass(type_, (str, tuple)): arg_types = _get_args(type_) if len(arg_types) > 0: element_type = self.typing_to_runner_api(arg_types[0]) @@ -397,7 +413,9 @@ def typing_to_runner_api(self, type_: type) -> schema_pb2.FieldType: return _python_any_schema_pb2(has_repr=False) try: - if LogicalType.is_known_logical_type(type_): + if isinstance(type_, LogicalType): + logical_type = type_ + elif LogicalType.is_known_logical_type(type_): logical_type = type_ else: logical_type = LogicalType.from_typing(type_) @@ -910,23 +928,15 @@ def _from_typing(cls, typ): raise NotImplementedError() @classmethod - def from_runner_api(cls, logical_type_proto): + def _from_runner_api(cls, logical_type_proto): # type: (schema_pb2.LogicalType) -> LogicalType - """Construct an instance of a registered LogicalType implementation given a - proto LogicalType. - - Raises ValueError if no LogicalType registered for the given URN. + """Construct an instance of this LogicalType implementation given a proto. """ - logical_type = cls._known_logical_types.get_logical_type_by_urn( - logical_type_proto.urn) - if logical_type is None: - raise ValueError( - "No logical type registered for URN '%s'" % logical_type_proto.urn) if not logical_type_proto.HasField( "argument_type") or not logical_type_proto.HasField("argument"): # logical type_proto without argument - return logical_type() + return cls() else: try: argument = value_from_runner_api( @@ -939,8 +949,24 @@ def from_runner_api(cls, logical_type_proto): 'Logical type %s with argument is currently unsupported. ' 'Argument values are omitted', logical_type_proto.urn) - return logical_type() - return logical_type(argument) + return cls() + return cls(argument) + + @classmethod + def from_runner_api(cls, logical_type_proto): + # type: (schema_pb2.LogicalType) -> LogicalType + + """Construct an instance of a registered LogicalType implementation given a + proto LogicalType. + + Raises ValueError if no LogicalType registered for the given URN. + """ + logical_type = cls._known_logical_types.get_logical_type_by_urn( + logical_type_proto.urn) + if logical_type is None: + raise ValueError( + "No logical type registered for URN '%s'" % logical_type_proto.urn) + return logical_type._from_runner_api(logical_type_proto) @classmethod def is_known_logical_type(cls, logical_type): @@ -1511,3 +1537,88 @@ def argument(self): @classmethod def _from_typing(cls, typ): return cls() + + +_TUPLE_NAMEDTUPLE_CACHE: Dict[int, type] = {} + + +def _get_tuple_namedtuple(n: int) -> type: + cls = _TUPLE_NAMEDTUPLE_CACHE.get(n) + if cls is None: + cls = NamedTuple(f"_FixedTuple{n}", [(f"f{i}", object) for i in range(n)]) + _TUPLE_NAMEDTUPLE_CACHE[n] = cls + return cls + + +@LogicalType._register_internal +class FixedTupleLogicalType(NoArgumentLogicalType[tuple, Any]): + """Logical type representing fixed-length Python tuples backed by a Row.""" + def __init__(self, tuple_types: Sequence[type] = ()): + self._tuple_types = tuple(tuple_types) + + @classmethod + def urn(cls): + return FIXED_TUPLE_URN + + def language_type(self=None): + if self is None or not self._tuple_types: + return tuple + return Tuple[self._tuple_types] + + def representation_type(self): + if not self._tuple_types: + from apache_beam.pvalue import Row + return Row + fields = [(f"f{i}", t) for i, t in enumerate(self._tuple_types)] + options = [] + st = SchemaTranslation(schema_registry=SCHEMA_REGISTRY) + if not any(st.typing_to_runner_api(t).nullable for t in self._tuple_types): + options.append((_SCHEMA_OPTION_STATIC_ENCODING, True)) + return row_type.RowTypeConstraint.from_fields( + fields, schema_options=options) + + def to_representation_type(self, value): + cls = _get_tuple_namedtuple(len(value)) + return cls(*value) + + def to_language_type(self, value): + return tuple(value) + + @classmethod + def _from_runner_api(cls, logical_type_proto): + if logical_type_proto.HasField("representation"): + row_schema = logical_type_proto.representation.row_type.schema + return cls([typing_from_runner_api(f.type) for f in row_schema.fields]) + return cls() + + +@LogicalType._register_internal +class VarTupleLogicalType(NoArgumentLogicalType[tuple, Sequence]): + """Logical type representing variable-length Python tuples backed by an Array.""" + def __init__(self, elem_type: type = object): + self._elem_type = elem_type + + @classmethod + def urn(cls): + return VAR_TUPLE_URN + + def language_type(self=None): + if self is None or self._elem_type is object: + return tuple + return Tuple[self._elem_type, ...] + + def representation_type(self): + return Sequence[self._elem_type] + + def to_representation_type(self, value): + return value + + def to_language_type(self, value): + return tuple(value) + + @classmethod + def _from_runner_api(cls, logical_type_proto): + if logical_type_proto.HasField("representation"): + elem_type = logical_type_proto.representation.array_type.element_type + return cls(typing_from_runner_api(elem_type)) + return cls() diff --git a/sdks/python/apache_beam/typehints/schemas_test.py b/sdks/python/apache_beam/typehints/schemas_test.py index 5e66a491090d..5c202a3aa3df 100644 --- a/sdks/python/apache_beam/typehints/schemas_test.py +++ b/sdks/python/apache_beam/typehints/schemas_test.py @@ -32,6 +32,7 @@ from typing import NamedTuple from typing import Optional from typing import Sequence +from typing import Tuple import numpy as np import pytest @@ -79,6 +80,11 @@ all_primitives, all_primitives) ] +basic_tuple_types = [ + Tuple[str, np.int64], + Tuple[np.int64, ...], +] + class AllPrimitives(NamedTuple): field_int8: np.int8 @@ -109,6 +115,8 @@ class ComplexSchema(NamedTuple): array_optional: Sequence[Optional[bool]] timestamp: Timestamp date: datetime.date + fixed_tuple: Tuple[str, np.int64] + var_tuple: Tuple[np.int64, ...] def get_test_beam_fieldtype_protos(): @@ -390,7 +398,8 @@ class SchemaTest(unittest.TestCase): @parameterized.expand([(user_type,) for user_type in all_primitives + \ basic_array_types + \ - basic_map_types] + basic_map_types + \ + basic_tuple_types] ) def test_typing_survives_proto_roundtrip(self, user_type): self.assertEqual(