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 3b889d239f [python] Reduce memory usage for overlapping temporal
windows (#10000)
3b889d239f is described below
commit 3b889d239fe54a8673d87936c2f3fe00c11474f9
Author: QuakeWang <[email protected]>
AuthorDate: Sun Sep 20 15:58:59 2026 +0800
[python] Reduce memory usage for overlapping temporal windows (#10000)
---
paimon-python/pypaimon/multimodal/temporal.py | 45 ++++++------
.../pypaimon/tests/multimodal_temporal_test.py | 79 +++++++++++++++++++---
2 files changed, 95 insertions(+), 29 deletions(-)
diff --git a/paimon-python/pypaimon/multimodal/temporal.py
b/paimon-python/pypaimon/multimodal/temporal.py
index 93f118d94b..50af264bfb 100644
--- a/paimon-python/pypaimon/multimodal/temporal.py
+++ b/paimon-python/pypaimon/multimodal/temporal.py
@@ -591,11 +591,11 @@ class _WindowJoinRight(_AsOfJoinRight):
metadata=source.metadata))
return fields
- def match(self, anchor_row):
+ def _match_range(self, anchor_row):
key = tuple(anchor_row[name] for name in self.by)
bounds = self._index.get(key)
if bounds is None:
- return []
+ return 0, 0
target = anchor_row[_TIME_KEY]
start, end = bounds
left = target - self._preceding_key
@@ -616,7 +616,7 @@ class _WindowJoinRight(_AsOfJoinRight):
first_key = max(first_key, _python_scalar(self._time_keys[start]))
last_key = min(last_key, _python_scalar(self._time_keys[end - 1]))
if first_key > last_key:
- return []
+ return 0, 0
first = bisect_left(
self._time_keys, first_key, start, end)
last = bisect_right(
@@ -632,21 +632,29 @@ class _WindowJoinRight(_AsOfJoinRight):
if self.closed in ("both", "right")
else bisect_left(self._time_keys, right, first, end)
)
- return [self._row_ids[index].as_py()
- for index in range(first, last)]
+ return first, last
def build_arrays(self, anchor_rows, fetcher):
- matches = [self.match(row) for row in anchor_rows]
- unique_ids = list(dict.fromkeys(
- row_id for match in matches for row_id in match))
+ ranges = [self._match_range(row) for row in anchor_rows]
+ slices = [(0, 0)] * len(ranges)
+ unique_ids = []
+ covered_end = 0
+ offset = 0
+ # Ranges refer to sorted metadata positions, not contiguous row IDs.
+ # Fetch their union in that order so each window becomes an Arrow
slice.
+ for index, (start, end) in sorted(
+ enumerate(ranges), key=lambda item: item[1]):
+ if start == end:
+ continue
+ if start >= covered_end:
+ offset = len(unique_ids) - start
+ unread_start = max(start, covered_end)
+ if end > unread_start:
+ unique_ids.extend(
+ self._row_ids.slice(unread_start, end -
unread_start).to_pylist())
+ covered_end = end
+ slices[index] = (start + offset, end - start)
values = fetcher.fetch(unique_ids)
- positions = {
- row_id: index for index, row_id in enumerate(unique_ids)
- }
- indices = [
- [positions[row_id] for row_id in match]
- for match in matches
- ]
aggregations_by_source = {}
for index, (_, source_name, aggregation) in
enumerate(self.aggregations):
aggregations_by_source.setdefault(source_name, []).append(
@@ -661,11 +669,8 @@ class _WindowJoinRight(_AsOfJoinRight):
]
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 start, length in slices:
+ selected = source_values.slice(start, length)
for (_, aggregation), result in zip(aggregations, results):
result.append(_aggregate_values(selected, aggregation))
del selected
diff --git a/paimon-python/pypaimon/tests/multimodal_temporal_test.py
b/paimon-python/pypaimon/tests/multimodal_temporal_test.py
index 50d8e92bcb..6ffbc80368 100644
--- a/paimon-python/pypaimon/tests/multimodal_temporal_test.py
+++ b/paimon-python/pypaimon/tests/multimodal_temporal_test.py
@@ -19,6 +19,7 @@ import os
import shutil
import sys
import tempfile
+import tracemalloc
import unittest
from datetime import datetime, timedelta
from unittest import mock
@@ -376,7 +377,7 @@ class MultimodalTemporalTest(unittest.TestCase):
self.assertIsNone(rows[2]["average"])
self.assertEqual(0, rows[2]["valid_count"])
- def test_window_join_reuses_gathered_values(self):
+ def test_window_join_reuses_arrow_slices(self):
source = temporal._WindowJoinRight.__new__(temporal._WindowJoinRight)
source.by = ("group",)
source._index = {(1,): (0, 3), (2,): (3, 4)}
@@ -425,14 +426,74 @@ class MultimodalTemporalTest(unittest.TestCase):
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)
+ take.assert_not_called()
+
+ def test_window_join_preserves_order_across_disjoint_ranges(self):
+ anchors = self._table("range_anchors", {
+ "episode_id": pa.int32(), "event_time": pa.int64(),
+ })
+ samples = self._table("range_samples", {
+ "episode_id": pa.int32(), "event_time": pa.int64(),
+ "label": pa.string(),
+ })
+ # Physical row IDs are neither contiguous nor ordered within a window.
+ samples.add([
+ {"episode_id": 1, "event_time": 12, "label": "late"},
+ {"episode_id": 2, "event_time": 2, "label": "other"},
+ {"episode_id": 1, "event_time": 2, "label": "first"},
+ {"episode_id": 1, "event_time": 6, "label": "excluded"},
+ {"episode_id": 1, "event_time": 2, "label": None},
+ {"episode_id": 1, "event_time": 2, "label": "last"},
+ {"episode_id": 1, "event_time": 1, "label": "early"},
+ ])
+ anchors.add([
+ {"episode_id": group, "event_time": time}
+ for group, time in [(1, 12), (1, 3), (2, 2), (1, 2),
+ (1, 2), (3, 2), (1, 20)]
+ ])
+ query = pmm.join_window(
+ anchors.scan(), samples.scan().select("label"),
+ on="event_time", by="episode_id", preceding=1,
+ aggregations={
+ "first": ("label", "first"),
+ "last": ("label", "last"),
+ "count": ("label", "count"),
+ })
+ rows = query.to_list()
+ self.assertEqual([
+ ("late", "late", 1), ("first", "last", 2),
+ ("other", "other", 1), ("early", "last", 3),
+ ("early", "last", 3), (None, None, 0), (None, None, 0),
+ ], [(row["first"], row["last"], row["count"]) for row in rows])
+
+ def test_window_join_bounds_overlapping_window_memory(self):
+ size = 1024
+ source = temporal._WindowJoinRight.__new__(temporal._WindowJoinRight)
+ source.by = ("group",)
+ source._index = {(1,): (0, size)}
+ source._time_keys = np.arange(size, dtype=np.int64)
+ source._row_ids = pa.array(range(10000, 10000 + size), type=pa.int64())
+ source.time_type = pa.int64()
+ source._preceding_key = size - 1
+ source._following_key = 0
+ source.closed = "both"
+ source.aggregations = (("count", "value", "count"),)
+ values = pa.table({"value": pa.array(range(size), type=pa.int64())})
+ fetcher = mock.Mock(schema=values.schema)
+ fetcher.fetch.return_value = values
+ anchors = [{"group": 1, temporal._TIME_KEY: size - 1}] * 256
+
+ tracemalloc.start()
+ try:
+ arrays = source.build_arrays(anchors, fetcher)
+ _, peak = tracemalloc.get_traced_memory()
+ finally:
+ tracemalloc.stop()
+
+ self.assertEqual([size] * len(anchors), arrays[0].to_pylist())
+ fetcher.fetch.assert_called_once_with(list(range(10000, 10000 + size)))
+ # Per-window row-ID and position lists alone exceed this budget.
+ self.assertLess(peak, 4 * 1024 * 1024)
def test_window_join_supports_asymmetric_timestamp_bounds(self):
anchors = self._table("window_timestamp_anchors", {
