This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 763dc4a5ef [python] Reuse gathered window values across aggregations
(#9804)
763dc4a5ef is described below
commit 763dc4a5ef57e78adb5fa5da113732dd43ca7e83
Author: QuakeWang <[email protected]>
AuthorDate: Tue Sep 15 11:15:48 2026 +0800
[python] Reuse gathered window values across aggregations (#9804)
---
paimon-python/pypaimon/multimodal/temporal.py | 40 ++++++++++-----
.../pypaimon/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 bb42053ef3..7ccf119ca9 100644
--- a/paimon-python/pypaimon/multimodal/temporal.py
+++ b/paimon-python/pypaimon/multimodal/temporal.py
@@ -633,16 +633,31 @@ class _WindowJoinRight(_AsOfJoinRight):
[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 5ee34556d4..7cba7dfaa6 100644
--- a/paimon-python/pypaimon/tests/multimodal_temporal_test.py
+++ b/paimon-python/pypaimon/tests/multimodal_temporal_test.py
@@ -248,6 +248,64 @@ class MultimodalTemporalTest(unittest.TestCase):
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(),
