This is an automated email from the ASF dual-hosted git repository. XiaoHongbo-Hope pushed a commit to branch release-2.1 in repository https://gitbox.apache.org/repos/asf/paimon.git
commit 237cde76ad0d808e9ae7a781a945a15401ba1599 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(),
