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(