This is an automated email from the ASF dual-hosted git repository.
claudevdm pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new f7b52028e13 [Python] Fix ParquetIO silently corrupting data across
windows in unbounded pipelines (#40293)
f7b52028e13 is described below
commit f7b52028e13d351370203477725ee6595ae7a1dc
Author: Rakesh Pai <[email protected]>
AuthorDate: Sun Oct 11 19:20:22 2026 -0700
[Python] Fix ParquetIO silently corrupting data across windows in unbounded
pipelines (#40293)
* [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
* style: apply Python formatter
* 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)
* Fix formatting: remove extra blank line before new streaming test
---
sdks/python/apache_beam/io/parquetio.py | 125 ++++++++++++++++-----------
sdks/python/apache_beam/io/parquetio_test.py | 108 +++++++++++++++++++++++
2 files changed, 181 insertions(+), 52 deletions(-)
diff --git a/sdks/python/apache_beam/io/parquetio.py
b/sdks/python/apache_beam/io/parquetio.py
index e5dce15e71e..691b68d789b 100644
--- a/sdks/python/apache_beam/io/parquetio.py
+++ b/sdks/python/apache_beam/io/parquetio.py
@@ -105,46 +105,62 @@ class _RowDictionariesToArrowTable(DoFn):
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 @@ class _RowDictionariesToArrowTable(DoFn):
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):
@@ -831,31 +848,35 @@ class _ParquetSink(filebasedsink.FileBasedSink):
"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 c7bf687f12c..3b2eec559f3 100644
--- a/sdks/python/apache_beam/io/parquetio_test.py
+++ b/sdks/python/apache_beam/io/parquetio_test.py
@@ -1013,6 +1013,114 @@ class WriteStreamingTest(unittest.TestCase):
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
+ # 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)
+
+
[email protected](pa is None, "PyArrow is not installed.")
[email protected]_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)
