[python] Reuse gathered window values across aggregations

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,18 +710,17 @@ 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":
         items = [item for item in selected.to_pylist()
                  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()

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