Skip to content

Exposed source window to AssignContext in window_mapping_fn for side inputs #36722

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

Merged
damccorm merged 14 commits into from
Mar 26, 2026
Merged

Exposed source window to AssignContext in window_mapping_fn for side inputs #36722

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
1dc8447
added source window to assign context window mapping fn
alexmalao1 Nov 4, 2025
a30b97f
yapf
alexmalao1 Nov 4, 2025
d6d570d
added default window mapping test
alexmalao1 Nov 18, 2025
35ef010
Merge branch 'master' into issue-36719
alexmalao1 Nov 18, 2025
fedded5
side input test
alexmalao1 Nov 19, 2025
41661a8
Merge branch 'master' into issue-36719
alexmalao1 Feb 11, 2026
8dabb36
added window and windowing fn required methods
alexmalao1 Feb 11, 2026
450d8d3
added window and windowing fn required methods
alexmalao1 Feb 11, 2026
e41566c
updated typehint for assign in sideinputs_test.py
alexmalao1 Feb 24, 2026
27a2857
added none assert
alexmalao1 Mar 16, 2026
bbd1d95
removed unused mock import in sideinput_test.py
alexmalao1 Mar 17, 2026
9398bdf
removed extraneous pdb file
alexmalao1 Mar 24, 2026
46b0190
removed 1igm.pdb
alexmalao1 Mar 24, 2026
13aed8b
fixed import order
alexmalao1 Mar 24, 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
3 changes: 2 additions & 1 deletion sdks/python/apache_beam/transforms/sideinputs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,8 @@ def default_window_mapping_fn(
def map_via_end(source_window: window.BoundedWindow) -> window.BoundedWindow:
return list(
target_window_fn.assign(
window.WindowFn.AssignContext(source_window.max_timestamp())))[-1]
window.WindowFn.AssignContext(
source_window.max_timestamp(), window=source_window)))[-1]

return map_via_end

Expand Down
35 changes: 35 additions & 0 deletions sdks/python/apache_beam/transforms/sideinputs_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@
from apache_beam.testing.util import equal_to
from apache_beam.testing.util import equal_to_per_window
from apache_beam.transforms import Map
from apache_beam.transforms import sideinputs
from apache_beam.transforms import trigger
from apache_beam.transforms import window
from apache_beam.utils.timestamp import Timestamp
Expand Down Expand Up @@ -489,6 +490,40 @@ def process(
assert_that(results, equal_to([(num_records, expected_fingerprint)]))
pipeline.run()

def test_default_window_mapping_fn_source_window(self):
Copy link
Copy Markdown
Contributor

@damccorm damccorm Jan 29, 2026

Choose a reason for hiding this comment

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

Looks like this is failing in CI with:

<testcase classname="apache_beam.transforms.sideinputs_test.SideInputsTest" name="test_default_window_mapping_fn_source_window" time="0.001">
<failure message="TypeError: Can't instantiate abstract class StringIDWindows with abstract method get_window_coder">self = <apache_beam.transforms.sideinputs_test.SideInputsTest testMethod=test_default_window_mapping_fn_source_window> def test_default_window_mapping_fn_source_window(self): """Test that the default window mapping function will propagate the source window when attempting to assign context. """ class StringIDWindow(window.BoundedWindow): """A window defined by an arbitrary string ID.""" def __init__(self, window_id: str): super().__init__(self._getTimestampFromProto()) self.id = window_id class StringIDWindows(window.NonMergingWindowFn): """ A windowing function that assigns each element a window with ID.""" def assign( self, assign_context: window.WindowFn.AssignContext ) -> Iterable[StringIDWindow]: if assign_context.element is None: return [assign_context.window] return [StringIDWindow(str(assign_context.element))] > mapping_fn = sideinputs.default_window_mapping_fn(StringIDWindows()) E TypeError: Can't instantiate abstract class StringIDWindows with abstract method get_window_coder apache_beam/transforms/sideinputs_test.py:513: TypeError</failure>
</testcase>

Could you please take a look?

Copy link
Copy Markdown
Contributor Author

@alexmalao1 alexmalao1 Feb 11, 2026

Choose a reason for hiding this comment

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

Added the required dummy _getTimestampFromProto and get_window_coder methods for the string id window examples.

Test should pass now.

"""Test that the default window mapping function will propagate the
source window when attempting to assign context.
"""
class StringIDWindow(window.BoundedWindow):
"""A window defined by an arbitrary string ID."""
def __init__(self, window_id: str):
super().__init__(self._getTimestampFromProto())
self.id = window_id

@staticmethod
def _getTimestampFromProto() -> Timestamp:
return Timestamp(micros=0)

class StringIDWindows(window.NonMergingWindowFn):
""" A windowing function that assigns each element a window with ID."""
def assign(
self, assign_context: window.WindowFn.AssignContext
) -> Iterable[window.BoundedWindow]:
if assign_context.element is None:
assert assign_context.window is not None
return [assign_context.window]
return [StringIDWindow(str(assign_context.element))]

def get_window_coder(self):
return None

mapping_fn = sideinputs.default_window_mapping_fn(StringIDWindows())
source_window = StringIDWindows().assign(
window.WindowFn.AssignContext(Timestamp(10), element='element'))[0]
bounded_window = mapping_fn(source_window)
assert bounded_window is not None
assert bounded_window.id == 'element'


if __name__ == '__main__':
logging.getLogger().setLevel(logging.DEBUG)
Expand Down
Loading