[Python] Fix ParquetIO silently corrupting data across windows in unbounded pipelines

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()

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<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):
+  """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)