From 3de2fe3742d44d361ff3852e367468bec5cf8157 Mon Sep 17 00:00:00 2001 From: QuakeWang Date: Mon, 14 Sep 2026 11:30:34 +0800 Subject: [PATCH] [python] Reuse gathered window values across aggregations Gather each source column once per window and reuse it across aggregations, avoiding repeated Arrow takes and index allocations. Preserve aggregation semantics and output order without caching gathered windows. Signed-off-by: QuakeWang --- paimon-python/pypaimon/multimodal/temporal.py | 40 ++++++++----- .../tests/multimodal_temporal_test.py | 58 +++++++++++++++++++ 2 files changed, 85 insertions(+), 13 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/temporal.py b/paimon-python/pypaimon/multimodal/temporal.py index bb42053ef399..7ccf119ca99b 100644 --- a/paimon-python/pypaimon/multimodal/temporal.py +++ b/paimon-python/pypaimon/multimodal/temporal.py @@ -633,16 +633,31 @@ def build_arrays(self, anchor_rows, fetcher): [positions[row_id] for row_id in match] for match in matches ] - arrays = [] - for _, source_name, aggregation in self.aggregations: + aggregations_by_source = {} + for index, (_, source_name, aggregation) in enumerate(self.aggregations): + aggregations_by_source.setdefault(source_name, []).append( + (index, aggregation)) + + arrays = [None] * len(self.aggregations) + for source_name, aggregations in aggregations_by_source.items(): effective = fetcher.schema.field(source_name) - output_type = _aggregate_output_type( - effective.type, aggregation) - arrays.append(pa.array([ - _aggregate_values( - values[source_name], row_indices, aggregation) - for row_indices in indices - ], type=output_type)) + output_types = [ + _aggregate_output_type(effective.type, aggregation) + for _, aggregation in aggregations + ] + source_values = values[source_name] + results = [[] for _ in aggregations] + for row_indices in indices: + selected = ( + pc.take(source_values, pa.array(row_indices, type=pa.int64())) + if row_indices else source_values.slice(0, 0) + ) + for (_, aggregation), result in zip(aggregations, results): + result.append(_aggregate_values(selected, aggregation)) + del selected + for (index, _), result, output_type in zip( + aggregations, results, output_types): + arrays[index] = pa.array(result, type=output_type) return arrays @@ -695,10 +710,9 @@ def _aggregate_output_type(data_type, aggregation): return data_type -def _aggregate_values(values, indices, aggregation): - if not indices: +def _aggregate_values(selected, aggregation): + if len(selected) == 0: return 0 if aggregation == "count" else None - selected = pc.take(values, pa.array(indices, type=pa.int64())) if aggregation == "count": return pc.count(selected).as_py() if aggregation == "mean": @@ -706,7 +720,7 @@ def _aggregate_values(values, indices, aggregation): if item is not None] if not items: return None - if pa.types.is_integer(values.type): + if pa.types.is_integer(selected.type): return sum(items) / len(items) if not all(math.isfinite(item) for item in items): return pc.mean(selected).as_py() diff --git a/paimon-python/pypaimon/tests/multimodal_temporal_test.py b/paimon-python/pypaimon/tests/multimodal_temporal_test.py index 5ee34556d4d1..7cba7dfaa689 100644 --- a/paimon-python/pypaimon/tests/multimodal_temporal_test.py +++ b/paimon-python/pypaimon/tests/multimodal_temporal_test.py @@ -248,6 +248,64 @@ def test_window_join_stays_in_group_and_skips_nulls(self): self.assertIsNone(rows[2]["average"]) self.assertEqual(0, rows[2]["valid_count"]) + def test_window_join_reuses_gathered_values(self): + source = temporal._WindowJoinRight.__new__(temporal._WindowJoinRight) + source.by = ("group",) + source._index = {(1,): (0, 3), (2,): (3, 4)} + source._time_keys = np.array([5, 10, 15, 10], dtype=np.int64) + source._row_ids = pa.array([11, 12, 13, 14], type=pa.int64()) + source.time_type = pa.int64() + source._preceding_key = 5 + source._following_key = 0 + source.closed = "both" + specifications = ( + ("average", "value", "mean"), + ("first_label", "label", "first"), + ("minimum", "value", "min"), + ("last_label", "label", "last"), + ("maximum", "value", "max"), + ("valid_count", "label", "count"), + ) + values = pa.table({ + "value": pa.array([1, None, 5, None], type=pa.int32()), + "label": pa.array(["a", None, "c", None], type=pa.string()), + }) + anchors = [ + {"group": group, temporal._TIME_KEY: time} + for group, time in [(1, 10), (1, 15), (2, 10), (3, 10)] + ] + expected = { + "mean": pa.array([1.0, 5.0, None, None], type=pa.float64()), + "min": pa.array([1, 5, None, None], type=pa.int32()), + "max": pa.array([1, 5, None, None], type=pa.int32()), + "first": pa.array(["a", "c", None, None], type=pa.string()), + "last": pa.array(["a", "c", None, None], type=pa.string()), + "count": pa.array([1, 1, 0, 0], type=pa.int64()), + } + for aggregations in (specifications[:1], specifications[::2], + specifications): + with self.subTest(aggregations=aggregations): + source.aggregations = aggregations + fetcher = mock.Mock(schema=values.schema) + fetcher.fetch.return_value = values + with mock.patch.object( + temporal.pc, "take", wraps=temporal.pc.take) as take: + arrays = source.build_arrays(anchors, fetcher) + + fetcher.fetch.assert_called_once_with([11, 12, 13, 14]) + self.assertEqual(len(aggregations), len(arrays)) + for (_, _, operation), array in zip(aggregations, arrays): + self.assertEqual(expected[operation].type, array.type) + self.assertEqual(expected[operation], array) + source_names = {name for _, name, _ in aggregations} + self.assertEqual(3 * len(source_names), take.call_count) + for name in source_names: + gathered_indices = [ + call[0][1].to_pylist() for call in take.call_args_list + if call[0][0].equals(values[name]) + ] + self.assertEqual([[0, 1], [1, 2], [3]], gathered_indices) + def test_window_join_supports_asymmetric_timestamp_bounds(self): anchors = self._table("window_timestamp_anchors", { "episode_id": pa.int32(),