From 814e171544c4f605369337c33e41f5df50b33d25 Mon Sep 17 00:00:00 2001
From: Rakesh Pai <41351936+developer-rpai@users.noreply.github.com>
Date: Fri, 25 Sep 2026 23:08:59 -0700
Subject: [PATCH 1/4] [Python] Fix ParquetIO silently corrupting data across
 windows

_RowDictionariesToArrowTable buffered rows from every window into a single
flat buffer and overwrote self._window on each process() call. At
finish_bundle() the whole mixed-window table was emitted in one
WindowedValue carrying only the last window seen. In unbounded pipelines a
bundle routinely spans windows (the convert_fn ParDo runs before the
window-grouping GroupByKey in WriteImpl._expand_unbounded), so rows from
earlier windows were silently attributed to the latest window: data loss
for the earlier windows and corrupted downstream aggregations.

Buffer record batches per window and emit one WindowedValue per window at
finish_bundle() (and on mid-bundle row-group flushes), mirroring how
_WriteWindowedBundleDoFn keys its writers by window. Pane info is
intentionally not propagated: rows from many panes may land in one table,
and the downstream file sink keys writers by window only.

Adds regression tests driving the DoFn directly with interleaved windows.

Fixes #40284
---
 sdks/python/apache_beam/io/parquetio.py      | 89 ++++++++++++--------
 sdks/python/apache_beam/io/parquetio_test.py | 68 +++++++++++++++
 2 files changed, 121 insertions(+), 36 deletions(-)

diff --git a/sdks/python/apache_beam/io/parquetio.py b/sdks/python/apache_beam/io/parquetio.py
index e5dce15e71ef..af0f053c63b6 100644
--- a/sdks/python/apache_beam/io/parquetio.py
+++ b/sdks/python/apache_beam/io/parquetio.py
@@ -105,46 +105,62 @@ def __init__(
       record_batch_size=1000):
     self._schema = schema
     self._row_group_buffer_size = row_group_buffer_size
-    self._buffer = [[] for _ in range(len(schema.names))]
     self._buffer_size = record_batch_size
-    self._record_batches = []
-    self._record_batches_byte_size = 0
-    self._window = None
+    self._reset_window_state()
+
+  def _reset_window_state(self):
+    # Buffers are keyed by window: a single bundle may contain elements from
+    # multiple windows (the convert_fn ParDo runs before the window-grouping
+    # GroupByKey in WriteImpl._expand_unbounded), and rows must not leak
+    # across windows.
+    self._buffers = {}
+    self._record_batches = {}
+    self._record_batches_byte_size = {}
+
+  def _buffer_for_window(self, w):
+    if w not in self._buffers:
+      self._buffers[w] = [[] for _ in range(len(self._schema.names))]
+      self._record_batches[w] = []
+      self._record_batches_byte_size[w] = 0
+    return self._buffers[w]
+
+  def start_bundle(self):
+    self._reset_window_state()
 
   def process(self, row, w=DoFn.WindowParam, pane=DoFn.PaneInfoParam):
-    self._window = w
-    if len(self._buffer[0]) >= self._buffer_size:
-      self._flush_buffer()
+    buffer = self._buffer_for_window(w)
+    if len(buffer[0]) >= self._buffer_size:
+      self._flush_buffer(w)
 
-    if self._record_batches_byte_size >= self._row_group_buffer_size:
-      table = self._create_table()
-      yield table
+    if self._record_batches_byte_size[w] >= self._row_group_buffer_size:
+      yield self._windowed_table(self._create_table(w), w)
 
     # reorder the data in columnar format.
     for i, n in enumerate(self._schema.names):
       # Handle missing nullable fields by using None as default value
       field = self._schema.field(i)
       if field.nullable and n not in row:
-        self._buffer[i].append(None)
+        buffer[i].append(None)
       else:
-        self._buffer[i].append(row[n])
+        buffer[i].append(row[n])
 
   def finish_bundle(self):
-    if len(self._buffer[0]) > 0:
-      self._flush_buffer()
-    if self._record_batches_byte_size > 0:
-      table = self._create_table()
-      if self._window is None or isinstance(self._window, window.GlobalWindow):
-        # bounded input
-        yield window.GlobalWindows.windowed_value_at_end_of_window(table)
-      else:
-        # unbounded input
-        yield WindowedValue(
-            table,
-            timestamp=self._window.
-            end,  #or it could be max of timestamp of the rows processed
-            windows=[self._window]  # TODO(pabloem) HOW DO WE GET THE PANE
-        )
+    for w in list(self._buffers):
+      if len(self._buffers[w][0]) > 0:
+        self._flush_buffer(w)
+      if self._record_batches_byte_size[w] > 0:
+        yield self._windowed_table(self._create_table(w), w)
+    self._reset_window_state()
+
+  @staticmethod
+  def _windowed_table(table, w):
+    if w is None or isinstance(w, window.GlobalWindow):
+      # bounded input
+      return window.GlobalWindows.windowed_value_at_end_of_window(table)
+    # unbounded input. Rows from many panes may be batched into one table, so
+    # pane info is intentionally not propagated; the downstream file sink keys
+    # its writers by window only.
+    return WindowedValue(table, timestamp=w.end, windows=[w])
 
   def display_data(self):
     res = super().display_data()
@@ -153,25 +169,26 @@ def display_data(self):
 
     return res
 
-  def _create_table(self):
-    table = pa.Table.from_batches(self._record_batches, schema=self._schema)
-    self._record_batches = []
-    self._record_batches_byte_size = 0
+  def _create_table(self, w):
+    table = pa.Table.from_batches(self._record_batches[w], schema=self._schema)
+    self._record_batches[w] = []
+    self._record_batches_byte_size[w] = 0
     return table
 
-  def _flush_buffer(self):
+  def _flush_buffer(self, w):
     arrays = [[] for _ in range(len(self._schema.names))]
-    for x, y in enumerate(self._buffer):
+    buffer = self._buffers[w]
+    for x, y in enumerate(buffer):
       arrays[x] = pa.array(y, type=self._schema.types[x])
-      self._buffer[x] = []
+      buffer[x] = []
     rb = pa.RecordBatch.from_arrays(arrays, schema=self._schema)
-    self._record_batches.append(rb)
+    self._record_batches[w].append(rb)
     size = 0
     for x in arrays:
       for b in x.buffers():
         if b is not None:
           size = size + b.size
-    self._record_batches_byte_size = self._record_batches_byte_size + size
+    self._record_batches_byte_size[w] = self._record_batches_byte_size[w] + size
 
 
 class _ArrowTableToBeamRows(DoFn):
diff --git a/sdks/python/apache_beam/io/parquetio_test.py b/sdks/python/apache_beam/io/parquetio_test.py
index c7bf687f12cd..fc361ea599a8 100644
--- a/sdks/python/apache_beam/io/parquetio_test.py
+++ b/sdks/python/apache_beam/io/parquetio_test.py
@@ -1014,6 +1014,74 @@ def test_write_streaming_undef_shards_default_shard_name_template_windowed_pcoll
         "expected %d files, but got: %d" % (1 * 3, len(file_names)))
 
 
+@unittest.skipIf(pa is None, "PyArrow is not installed.")
+@pytest.mark.uses_pyarrow
+class RowDictionariesToArrowTableTest(unittest.TestCase):
+  """Unit tests for _RowDictionariesToArrowTable windowing behavior.
+
+  Regression tests for https://github.com/apache/beam/issues/40284: a single
+  bundle may contain elements from multiple windows (the convert_fn ParDo runs
+  before the window-grouping GroupByKey in WriteImpl._expand_unbounded), so
+  the DoFn must buffer and emit tables per window instead of attributing
+  every row to the last window seen.
+  """
+
+  def _make_dofn(self, **kwargs):
+    from apache_beam.io.parquetio import _RowDictionariesToArrowTable
+    schema = pa.schema([('id', pa.int64()), ('name', pa.string())])
+    return _RowDictionariesToArrowTable(schema, **kwargs)
+
+  def _drain(self, dofn, elements):
+    outputs = []
+    dofn.start_bundle()
+    for row, w in elements:
+      outputs.extend(dofn.process(row, w=w, pane=None))
+    outputs.extend(dofn.finish_bundle())
+    return outputs
+
+  def test_rows_are_buffered_per_window(self):
+    window = beam.transforms.window
+    w_a = window.IntervalWindow(0, 10)
+    w_b = window.IntervalWindow(10, 20)
+    outputs = self._drain(
+        self._make_dofn(),
+        [({'id': 1, 'name': 'a1'}, w_a),
+         ({'id': 2, 'name': 'b1'}, w_b),
+         ({'id': 3, 'name': 'a2'}, w_a)])
+    self.assertEqual(2, len(outputs))
+    by_window = {o.windows[0]: o.value.to_pydict() for o in outputs}
+    self.assertEqual({'id': [1, 3], 'name': ['a1', 'a2']}, by_window[w_a])
+    self.assertEqual({'id': [2], 'name': ['b1']}, by_window[w_b])
+    for o in outputs:
+      self.assertEqual(1, len(o.windows))
+      self.assertEqual(o.windows[0].end, o.timestamp)
+
+  def test_mid_bundle_flush_stays_within_window(self):
+    window = beam.transforms.window
+    w_a = window.IntervalWindow(0, 10)
+    w_b = window.IntervalWindow(10, 20)
+    outputs = self._drain(
+        self._make_dofn(record_batch_size=2),
+        [({'id': 1, 'name': 'a1'}, w_a),
+         ({'id': 2, 'name': 'a2'}, w_a),
+         ({'id': 3, 'name': 'a3'}, w_a),
+         ({'id': 4, 'name': 'b1'}, w_b)])
+    ids_by_window = {}
+    for o in outputs:
+      ids_by_window.setdefault(o.windows[0], []).extend(o.value.to_pydict()['id'])
+    self.assertEqual([1, 2, 3], ids_by_window[w_a])
+    self.assertEqual([4], ids_by_window[w_b])
+
+  def test_global_window_emits_single_windowed_value(self):
+    window = beam.transforms.window
+    outputs = self._drain(
+        self._make_dofn(),
+        [({'id': i, 'name': 'n%d' % i}, window.GlobalWindow()) for i in range(3)])
+    self.assertEqual(1, len(outputs))
+    self.assertIsInstance(outputs[0].windows[0], window.GlobalWindow)
+    self.assertEqual([0, 1, 2], outputs[0].value.to_pydict()['id'])
+
+
 if __name__ == '__main__':
   logging.getLogger().setLevel(logging.INFO)
   unittest.main()

From cd77dc99baccf466cf486038cc63c04754d1fe86 Mon Sep 17 00:00:00 2001
From: Rakesh Pai <41351936+developer-rpai@users.noreply.github.com>
Date: Wed, 30 Sep 2026 15:34:13 -0700
Subject: [PATCH 2/4] style: apply Python formatter

---
 sdks/python/apache_beam/io/parquetio_test.py | 31 +++++++++++++-------
 1 file changed, 21 insertions(+), 10 deletions(-)

diff --git a/sdks/python/apache_beam/io/parquetio_test.py b/sdks/python/apache_beam/io/parquetio_test.py
index fc361ea599a8..a2389804c4c8 100644
--- a/sdks/python/apache_beam/io/parquetio_test.py
+++ b/sdks/python/apache_beam/io/parquetio_test.py
@@ -1025,7 +1025,6 @@ class RowDictionariesToArrowTableTest(unittest.TestCase):
   the DoFn must buffer and emit tables per window instead of attributing
   every row to the last window seen.
   """
-
   def _make_dofn(self, **kwargs):
     from apache_beam.io.parquetio import _RowDictionariesToArrowTable
     schema = pa.schema([('id', pa.int64()), ('name', pa.string())])
@@ -1045,9 +1044,13 @@ def test_rows_are_buffered_per_window(self):
     w_b = window.IntervalWindow(10, 20)
     outputs = self._drain(
         self._make_dofn(),
-        [({'id': 1, 'name': 'a1'}, w_a),
-         ({'id': 2, 'name': 'b1'}, w_b),
-         ({'id': 3, 'name': 'a2'}, w_a)])
+        [({
+            'id': 1, 'name': 'a1'
+        }, w_a), ({
+            'id': 2, 'name': 'b1'
+        }, w_b), ({
+            'id': 3, 'name': 'a2'
+        }, w_a)])
     self.assertEqual(2, len(outputs))
     by_window = {o.windows[0]: o.value.to_pydict() for o in outputs}
     self.assertEqual({'id': [1, 3], 'name': ['a1', 'a2']}, by_window[w_a])
@@ -1062,13 +1065,19 @@ def test_mid_bundle_flush_stays_within_window(self):
     w_b = window.IntervalWindow(10, 20)
     outputs = self._drain(
         self._make_dofn(record_batch_size=2),
-        [({'id': 1, 'name': 'a1'}, w_a),
-         ({'id': 2, 'name': 'a2'}, w_a),
-         ({'id': 3, 'name': 'a3'}, w_a),
-         ({'id': 4, 'name': 'b1'}, w_b)])
+        [({
+            'id': 1, 'name': 'a1'
+        }, w_a), ({
+            'id': 2, 'name': 'a2'
+        }, w_a), ({
+            'id': 3, 'name': 'a3'
+        }, w_a), ({
+            'id': 4, 'name': 'b1'
+        }, w_b)])
     ids_by_window = {}
     for o in outputs:
-      ids_by_window.setdefault(o.windows[0], []).extend(o.value.to_pydict()['id'])
+      ids_by_window.setdefault(o.windows[0],
+                               []).extend(o.value.to_pydict()['id'])
     self.assertEqual([1, 2, 3], ids_by_window[w_a])
     self.assertEqual([4], ids_by_window[w_b])
 
@@ -1076,7 +1085,9 @@ def test_global_window_emits_single_windowed_value(self):
     window = beam.transforms.window
     outputs = self._drain(
         self._make_dofn(),
-        [({'id': i, 'name': 'n%d' % i}, window.GlobalWindow()) for i in range(3)])
+        [({
+            'id': i, 'name': 'n%d' % i
+        }, window.GlobalWindow()) for i in range(3)])
     self.assertEqual(1, len(outputs))
     self.assertIsInstance(outputs[0].windows[0], window.GlobalWindow)
     self.assertEqual([0, 1, 2], outputs[0].value.to_pydict()['id'])

From c5e6fb7849f0c1f885a8dd24b61b137720cdf3cf Mon Sep 17 00:00:00 2001
From: Rakesh Pai <41351936+developer-rpai@users.noreply.github.com>
Date: Wed, 30 Sep 2026 21:41:09 -0700
Subject: [PATCH 3/4] Fix _ParquetSink file-handle leak across windows + add
 streaming e2e test

The _RowDictionariesToArrowTable fix exposed a second bug further down the
write path: _ParquetSink stored the open file handle on the sink
(self._file_handle), but several writers may be open on one sink at once
(one per window in streaming writes). The wrong handle got closed, producing
Parquet files with missing footers ("Parquet magic bytes not found").

Make the file handle travel with its writer: open() now returns a
(writer, file_handle) tuple, and write_record()/close() unpack it.

Also adds the end-to-end regression test from review: one bundle spanning
3 windows must land each row in its own window's file.

Co-authored-by: claudevdm (review feedback)
---
 sdks/python/apache_beam/io/parquetio.py      | 36 +++++++++++---------
 sdks/python/apache_beam/io/parquetio_test.py | 30 ++++++++++++++++
 2 files changed, 50 insertions(+), 16 deletions(-)

diff --git a/sdks/python/apache_beam/io/parquetio.py b/sdks/python/apache_beam/io/parquetio.py
index af0f053c63b6..691b68d789b7 100644
--- a/sdks/python/apache_beam/io/parquetio.py
+++ b/sdks/python/apache_beam/io/parquetio.py
@@ -848,31 +848,35 @@ def __init__(
           "pyarrow version >= 4.x, please use a different pyarrow version. "
           f"Your pyarrow version: {pa.__version__}")
     self._use_compliant_nested_type = use_compliant_nested_type
-    self._file_handle = None
 
   def open(self, temp_path):
-    self._file_handle = super().open(temp_path)
+    # Several writers may be open on this sink at once (one per window in
+    # streaming writes), so the file handle travels with its writer rather
+    # than being stored on the sink.
+    file_handle = super().open(temp_path)
     if ARROW_MAJOR_VERSION < 4:
-      return pq.ParquetWriter(
-          self._file_handle,
+      writer = pq.ParquetWriter(
+          file_handle,
           self._schema,
           compression=self._codec,
           use_deprecated_int96_timestamps=self._use_deprecated_int96_timestamps)
-    return pq.ParquetWriter(
-        self._file_handle,
-        self._schema,
-        compression=self._codec,
-        use_deprecated_int96_timestamps=self._use_deprecated_int96_timestamps,
-        use_compliant_nested_type=self._use_compliant_nested_type)
-
-  def write_record(self, writer, table: paTable):
+    else:
+      writer = pq.ParquetWriter(
+          file_handle,
+          self._schema,
+          compression=self._codec,
+          use_deprecated_int96_timestamps=self._use_deprecated_int96_timestamps,
+          use_compliant_nested_type=self._use_compliant_nested_type)
+    return writer, file_handle
+
+  def write_record(self, writer_and_handle, table: paTable):
+    writer, _ = writer_and_handle
     writer.write_table(table)
 
-  def close(self, writer):
+  def close(self, writer_and_handle):
+    writer, file_handle = writer_and_handle
     writer.close()
-    if self._file_handle:
-      self._file_handle.close()
-      self._file_handle = None
+    file_handle.close()
 
   def display_data(self):
     res = super().display_data()
diff --git a/sdks/python/apache_beam/io/parquetio_test.py b/sdks/python/apache_beam/io/parquetio_test.py
index a2389804c4c8..e76b154b9306 100644
--- a/sdks/python/apache_beam/io/parquetio_test.py
+++ b/sdks/python/apache_beam/io/parquetio_test.py
@@ -1014,6 +1014,36 @@ def test_write_streaming_undef_shards_default_shard_name_template_windowed_pcoll
         "expected %d files, but got: %d" % (1 * 3, len(file_names)))
 
 
+  def test_write_streaming_rows_land_in_their_own_window(self):
+    # One add_elements call is delivered as one bundle spanning 3 windows.
+    # Regression test for the _ParquetSink file-handle bug: several writers
+    # may be open on the sink at once (one per window), so each writer must
+    # close its own file handle or the windowed files come out corrupted.
+    base = datetime(2021, 3, 1, tzinfo=pytz.UTC).timestamp()
+    offsets = [1, 11, 21, 3, 13, 23]
+    stream = TestStream().add_elements([
+        beam.window.TimestampedValue({'offset': o}, base + o) for o in offsets
+    ]).advance_watermark_to_infinity()
+    with TestPipeline() as p:
+      _ = (
+          p
+          | stream
+          | beam.WindowInto(beam.window.FixedWindows(10))
+          | beam.io.WriteToParquet(
+              file_path_prefix=self.tempdir + '/out',
+              file_name_suffix='.parquet',
+              num_shards=1,
+              schema=pa.schema([('offset', pa.int64())])))
+
+    pattern = re.compile(r'.*-\[(?P<start>[\d\.]+), [\d\.]+\)-.*\.parquet$')
+    rows_by_window = {}
+    for file_name in glob.glob(self.tempdir + '/out*'):
+      start = float(pattern.match(file_name).group('start')) - base
+      rows = pq.read_table(file_name).column('offset').to_pylist()
+      rows_by_window[start] = sorted(rows)
+    self.assertEqual({0: [1, 3], 10: [11, 13], 20: [21, 23]}, rows_by_window)
+
+
 @unittest.skipIf(pa is None, "PyArrow is not installed.")
 @pytest.mark.uses_pyarrow
 class RowDictionariesToArrowTableTest(unittest.TestCase):

From 6b1f0f6f968aa95198e472fba5eaeadefb75eb02 Mon Sep 17 00:00:00 2001
From: Rakesh Pai <41351936+developer-rpai@users.noreply.github.com>
Date: Thu, 1 Oct 2026 08:28:20 -0700
Subject: [PATCH 4/4] Fix formatting: remove extra blank line before new
 streaming test

---
 sdks/python/apache_beam/io/parquetio_test.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/sdks/python/apache_beam/io/parquetio_test.py b/sdks/python/apache_beam/io/parquetio_test.py
index e76b154b9306..3b2eec559f3a 100644
--- a/sdks/python/apache_beam/io/parquetio_test.py
+++ b/sdks/python/apache_beam/io/parquetio_test.py
@@ -1013,7 +1013,6 @@ def test_write_streaming_undef_shards_default_shard_name_template_windowed_pcoll
         1 * 3,  #25s of data covered by 3 10s windows
         "expected %d files, but got: %d" % (1 * 3, len(file_names)))
 
-
   def test_write_streaming_rows_land_in_their_own_window(self):
     # One add_elements call is delivered as one bundle spanning 3 windows.
     # Regression test for the _ParquetSink file-handle bug: several writers