diff --git a/sdks/python/apache_beam/io/parquetio.py b/sdks/python/apache_beam/io/parquetio.py index e5dce15e71ef..691b68d789b7 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): @@ -831,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 c7bf687f12cd..3b2eec559f3a 100644 --- a/sdks/python/apache_beam/io/parquetio_test.py +++ b/sdks/python/apache_beam/io/parquetio_test.py @@ -1013,6 +1013,114 @@ 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 + # 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[\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): + """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)