Skip to content

Use registered type for Row #38108

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Abacn wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Use registered type for Row #38108

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d6172eb
Use registered type for Row
Abacn Apr 8, 2026
1b75112
Allow decorator usage; document register_row is preferred
Abacn Apr 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Dataflow.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to run",
"modification": 1
"modification": 16
}
2 changes: 1 addition & 1 deletion .github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Direct.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to run",
"modification": 1
"modification": 17
}
2 changes: 1 addition & 1 deletion .github/trigger_files/beam_PostCommit_Python_Xlang_IO_Direct.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to run",
"modification": 1
"modification": 3
}
4 changes: 2 additions & 2 deletions sdks/python/apache_beam/coders/row_coder_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,13 +60,13 @@
("one_more_field", typing.Optional[str])])


@coders_registry.register_row
class People(typing.NamedTuple):
primary: Person
partner: typing.Optional[Person]


coders_registry.register_coder(Person, RowCoder)
coders_registry.register_coder(People, RowCoder)
coders_registry.register_row(Person)


class RowCoderTest(unittest.TestCase):
Expand Down
41 changes: 41 additions & 0 deletions sdks/python/apache_beam/coders/typecoders.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -124,6 +124,19 @@ def _normalize_typehint_type(typehint_type):
def register_coder(
self, typehint_type: Any,
typehint_coder_class: Type[coders.Coder]) -> None:
"""
Register a user type with a coder.

Typical usage::

class MyCustomType:
pass

coders.registry.register_coder(MyCustomType, MyCustomCoder)

To register a supported user type (data class or named tuple) with Beam Row,
use :meth:`register_row` instead, as it registers both coder and schema.
"""
if not isinstance(typehint_coder_class, type):
raise TypeError(
'Coder registration requires a coder class object. '
Expand All @@ -133,6 +146,34 @@ def register_coder(
self._register_coder_internal(
self._normalize_typehint_type(typehint_type), typehint_coder_class)

def register_row(self, typehint_type: type[Any]) -> type[Any]:
"""
Register a user type with a Beam Row.

This registers the type with a RowCoder and register its schema.

Register a dataclass::

@coders.registry.register_row
@dataclass
class MyDataClass:
id: int
name: str

Register a named tuple::

coders.registry.register_row(MyNamedTuple)
"""
from apache_beam.coders import RowCoder
from apache_beam.typehints.schemas import typing_to_runner_api

# Register with row coder
self.register_coder(typehint_type, RowCoder)
# This call generated a schema id for the type and register it with
# schema registry
typing_to_runner_api(typehint_type)
return typehint_type

def get_coder(self, typehint: Any) -> coders.Coder:
if typehint and typehint.__module__ == '__main__':
# See https://github.com/apache/beam/issues/21541
Expand Down
13 changes: 12 additions & 1 deletion sdks/python/apache_beam/internal/cloudpickle_pickler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -256,20 +256,27 @@ def dump_session(file_path):
# dump supported Beam Registries (currently only logical type registry)
from apache_beam.coders import typecoders
from apache_beam.typehints import schemas
from apache_beam.typehints.schema_registry import SCHEMA_REGISTRY

with _pickle_lock, open(file_path, 'wb') as file:
coder_reg = typecoders.registry.get_custom_type_coder_tuples()
logical_type_reg = schemas.LogicalType._known_logical_types.copy_custom()
schema_reg = SCHEMA_REGISTRY.get_registered_typings()

pickler = cloudpickle.CloudPickler(file)
# TODO(https://github.com/apache/beam/issues/18500) add file system registry
# once implemented
pickler.dump({"coder": coder_reg, "logical_type": logical_type_reg})
pickler.dump({
"coder": coder_reg,
"logical_type": logical_type_reg,
"schema": schema_reg
})


def load_session(file_path):
from apache_beam.coders import typecoders
from apache_beam.typehints import schemas
from apache_beam.typehints.schema_registry import SCHEMA_REGISTRY

with _pickle_lock, open(file_path, 'rb') as file:
registries = cloudpickle.load(file)
Expand All @@ -284,3 +291,7 @@ def load_session(file_path):
schemas.LogicalType._known_logical_types.load(registries["logical_type"])
else:
_LOGGER.warning('No logical type registry found in saved session')
if "schema" in registries:
SCHEMA_REGISTRY.load_registered_typings(registries["schema"])
else:
_LOGGER.warning('No schema registry found in saved session')
13 changes: 10 additions & 3 deletions sdks/python/apache_beam/io/external/xlang_jdbcio_it_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@

import apache_beam as beam
from apache_beam import coders
from apache_beam.io import jdbc
from apache_beam.io.jdbc import ReadFromJdbc
from apache_beam.io.jdbc import WriteToJdbc
from apache_beam.options.pipeline_options import StandardOptions
Expand Down Expand Up @@ -64,7 +65,7 @@
("f_timestamp", Timestamp), ("f_decimal", Decimal),
("f_date", datetime.date), ("f_time", datetime.time)],
)
coders.registry.register_coder(JdbcTestRow, coders.RowCoder)
coders.registry.register_row(JdbcTestRow)

CustomSchemaRow = typing.NamedTuple(
"CustomSchemaRow",
Expand All @@ -82,11 +83,17 @@
("renamed_time", datetime.time),
],
)
coders.registry.register_coder(CustomSchemaRow, coders.RowCoder)

# Need to put inside enforce_millis_instant_for_timestamp context to align
# with the same setup in ReadFromJdbc.__init__. Remove once Beam moved to
# micros instant for timestamp
# Alternatively, use coders.registry.register_coder(CustomSchemaRow, RowCoder)
with jdbc.enforce_millis_instant_for_timestamp():
coders.registry.register_row(CustomSchemaRow)

SimpleRow = typing.NamedTuple(
"SimpleRow", [("id", int), ("name", str), ("value", float)])
coders.registry.register_coder(SimpleRow, coders.RowCoder)
coders.registry.register_row(SimpleRow)


@pytest.mark.uses_gcp_java_expansion_service
Expand Down
1 change: 1 addition & 0 deletions sdks/python/apache_beam/io/jdbc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -264,6 +264,7 @@ def enforce_millis_instant_for_timestamp():
LogicalType._known_logical_types = old_registry.copy()
try:
LogicalType.register_logical_type(MillisInstant)
LogicalType.register_logical_type(JdbcDateType)
Copy link
Copy Markdown
Contributor Author

@Abacn Abacn Apr 20, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the test was passing before, but this line is needed after #38078

yield
finally:
LogicalType._known_logical_types = old_registry
Expand Down
18 changes: 8 additions & 10 deletions sdks/python/apache_beam/typehints/row_type_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,17 +33,18 @@ class RowTypeTest(unittest.TestCase):
@staticmethod
def _check_key_type_and_count(x) -> int:
key_type = type(x[0])
if not row_type._user_type_is_generated(key_type):
raise RuntimeError("Expect type after GBK to be generated user type")
if row_type._user_type_is_generated(key_type):
raise RuntimeError("Type after GBK not preserved, get generated type")
if not hasattr(key_type, row_type._BEAM_SCHEMA_ID):
raise RuntimeError("Type after GBK missing Beam schema ID")

return len(x[1])

def test_group_by_key_namedtuple(self):
MyNamedTuple = typing.NamedTuple(
"MyNamedTuple", [("id", int), ("name", str)])

beam.coders.typecoders.registry.register_coder(
MyNamedTuple, beam.coders.RowCoder)
beam.coders.typecoders.registry.register_row(MyNamedTuple)
Copy link
Copy Markdown
Collaborator

@claudevdm claudevdm Apr 19, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am not sure the divegence between beam.coders.typecoders.registry.register_coder and beam.coders.typecoders.registry.register_row is a good thing. Most users probably want the benefit of this PR but rely on register_coder.

Can we either

  1. make register_coder act the same as register_row OR
  2. Deprecate register_coder(T, RowCoder) with a warning. DeprecationWarning: use register_row(T) instead when the coder class is RowCoder

Copy link
Copy Markdown
Contributor Author

@Abacn Abacn Apr 20, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

register_row calls register_coder so it doesn't diverge that much. My thought introducing "register_row" was on

DeprecationWarning can be added after we migrate all register_coder(..., RowCoder) usage over Beam tests. For now, just added a documentation recommending register_row for RowCoders


def generate(num: int):
for i in range(100):
Expand All @@ -67,8 +68,7 @@ class MyDataClass:
id: int
name: str

beam.coders.typecoders.registry.register_coder(
MyDataClass, beam.coders.RowCoder)
beam.coders.typecoders.registry.register_row(MyDataClass)

def generate(num: int):
for i in range(100):
Expand Down Expand Up @@ -120,10 +120,8 @@ class DataClassInt:
class DataClassStr(DataClassInt):
name: str

beam.coders.typecoders.registry.register_coder(
DataClassInt, beam.coders.RowCoder)
beam.coders.typecoders.registry.register_coder(
DataClassStr, beam.coders.RowCoder)
beam.coders.typecoders.registry.register_row(DataClassInt)
beam.coders.typecoders.registry.register_row(DataClassStr)

def generate(num: int):
for i in range(10):
Expand Down
11 changes: 10 additions & 1 deletion sdks/python/apache_beam/typehints/schema_registry.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@
class SchemaTypeRegistry(object):
def __init__(self):
self.by_id = {}
self.by_typing = {}
self.by_typing = {} # currently not used

def generate_new_id(self):
for _ in range(100):
Expand All @@ -43,6 +43,15 @@ def add(self, typing, schema):
if schema.id:
self.by_id[schema.id] = (typing, schema)

def load_registered_typings(self, by_id):
for id, typing in by_id.items():
if id not in self.by_id:
self.by_id[id] = (typing, None)
Comment on lines +46 to +49
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot Apr 19, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The variable name id shadows the built-in Python function id(). Additionally, the argument name by_id is slightly confusing because it contains a mapping of IDs to typings, whereas the class attribute self.by_id stores tuples of (typing, schema). Renaming these to id_to_typing and schema_id improves clarity and avoids potential conflicts with built-in names.

Suggested change
def load_registered_typings(self, by_id):
for id, typing in by_id.items():
if id not in self.by_id:
self.by_id[id] = (typing, None)
def load_registered_typings(self, id_to_typing):
for schema_id, typing in id_to_typing.items():
if schema_id not in self.by_id:
self.by_id[schema_id] = (typing, None)

Copy link
Copy Markdown
Contributor Author

@Abacn Abacn Apr 20, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No actual conflict happening. schema object has an "id" field (schema.id).


def get_registered_typings(self):
# Used by save_main_session, as pb2.schema isn't picklable
return {k: v[0] for k, v in self.by_id.items()}

def get_typing_by_id(self, unique_id):
if not unique_id:
return None
Expand Down
29 changes: 16 additions & 13 deletions sdks/python/apache_beam/typehints/schemas.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -588,19 +588,22 @@ def named_tuple_from_schema(self, schema: schema_pb2.Schema) -> type:
descriptions[field.name] = field.description
subfields.append((field.name, field_py_type))

user_type = NamedTuple(type_name, subfields)

# Define a reduce function, otherwise these types can't be pickled
# (See BEAM-9574)
setattr(
user_type,
'__reduce__',
_named_tuple_reduce_method(schema.SerializeToString()))
setattr(user_type, "_field_descriptions", descriptions)
setattr(user_type, row_type._BEAM_SCHEMA_ID, schema.id)

self.schema_registry.add(user_type, schema)
coders.registry.register_coder(user_type, coders.RowCoder)
if schema.id in self.schema_registry.by_id:
user_type = self.schema_registry.by_id[schema.id][0]
else:
user_type = NamedTuple(type_name, subfields)

# Define a reduce function, otherwise these types can't be pickled
# (See BEAM-9574)
setattr(
user_type,
'__reduce__',
_named_tuple_reduce_method(schema.SerializeToString()))
setattr(user_type, "_field_descriptions", descriptions)
setattr(user_type, row_type._BEAM_SCHEMA_ID, schema.id)

self.schema_registry.add(user_type, schema)
coders.registry.register_coder(user_type, coders.RowCoder)

return user_type

Expand Down
Loading