From be4b21a62bf6463554ee2f8cb88ad7c3fe640c16 Mon Sep 17 00:00:00 2001 From: Yann Date: Sat, 17 Oct 2026 14:53:00 +0800 Subject: [PATCH] perf(python): avoid recursive copies of flat CWD vectors Copy flat lists of builtin immutable values without recursively copying each scalar. Keep deepcopy for nested and custom values so overlapping windows and padding remain isolated before transforms. Co-Authored-By: Codex AI-Model: gpt-6 Co-Authored-By: Codex AI-Contributed/Feature: 15/15 AI-Contributed/UT: 35/35 --- .../pypaimon/multimodal/window_dataset.py | 15 +++++++- .../tests/contiguous_window_dataset_test.py | 35 ++++++++++++++----- 2 files changed, 40 insertions(+), 10 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/window_dataset.py b/paimon-python/pypaimon/multimodal/window_dataset.py index 4a8209c955ee..51ffe7ecf715 100644 --- a/paimon-python/pypaimon/multimodal/window_dataset.py +++ b/paimon-python/pypaimon/multimodal/window_dataset.py @@ -46,6 +46,19 @@ from pypaimon.utils.range import Range +_IMMUTABLE_CELL_TYPES = frozenset((bool, int, float, str, bytes, type(None))) + + +def _copy_window_value(value): + # Arrow numeric vectors are flat lists. Copy the container without walking + # deepcopy's memo machinery for every scalar in overlapping windows. + if type(value) in _IMMUTABLE_CELL_TYPES: + return value + if type(value) is list and all(type(item) in _IMMUTABLE_CELL_TYPES for item in value): + return value.copy() + return copy.deepcopy(value) + + class ContiguousWindowDataset(Dataset): """Map-style training samples from a single table of contiguous frames. @@ -310,7 +323,7 @@ def _sample(self, anchor, position, reads): for name in self.columns: plans, rows = reads[name] row_ids, padding = plans[position] - values = [copy.deepcopy( + values = [_copy_window_value( self.pad_values[name] if row_id is None else rows[row_id][name]) for row_id in row_ids] if self.window_size is None: diff --git a/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py index ac3296295d83..5aad9ab7fd1c 100644 --- a/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py +++ b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py @@ -614,11 +614,13 @@ def test_plural_access_isolates_mutable_cells_between_samples(self): pa.field("episode", pa.string(), nullable=False), pa.field("step", pa.int32(), nullable=False), pa.field("values", pa.list_(pa.int32()), nullable=False), + pa.field("nested", pa.list_(pa.list_(pa.int32())), nullable=False), ]), options=_TABLE_OPTIONS, ) table.add([ - {"episode": "episode-a", "step": step, "values": [step]} + {"episode": "episode-a", "step": step, + "values": [step], "nested": [[step]]} for step in range(3) ]) @@ -627,21 +629,36 @@ def mutate(values): value.append(99) return values + def mutate_nested(values): + for value in values: + value[0].append(99) + return values + + padding = {"values": [-1], "nested": [[-1]]} dataset = table.scan().to_contiguous_window_dataset( - window_size=2, - columns=["values"], + columns=["values", "nested"], group_key="episode", order_key="step", - column_transforms={"values": mutate}, + frame_offsets={name: [0, 1, 1, 3, 3] for name in padding}, + boundary="pad", + pad_values=padding, + column_transforms={"values": mutate, "nested": mutate_nested}, ) batched = dataset.__getitems__([0, 1, 0]) singles = [dataset[index] for index in (0, 1, 0)] - - self.assertEqual( - [sample["values"] for sample in singles], - [sample["values"] for sample in batched], - ) + expected = [ + [[0, 99], [1, 99], [1, 99], [-1, 99], [-1, 99]], + [[1, 99], [2, 99], [2, 99], [-1, 99], [-1, 99]], + [[0, 99], [1, 99], [1, 99], [-1, 99], [-1, 99]], + ] + for samples in (batched, singles): + self.assertEqual(expected, [sample["values"] for sample in samples]) + self.assertEqual( + [[[value] for value in window] for window in expected], + [sample["nested"] for sample in samples], + ) + self.assertEqual({"values": [-1], "nested": [[-1]]}, padding) def test_pad_tail_repeats_last_row_and_marks_real_padding(self): dataset = self._dataset(