You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@paimon.apache.org by "XiaoHongbo-Hope (via GitHub)" <gi...@apache.org> on 2026/09/01 19:21:01 UTC
[PR] [python] Add temporal alignment for multimodal scans [paimon]
XiaoHongbo-Hope opened a new pull request, #9536:
URL: https://github.com/apache/paimon/pull/9536
### Purpose
Multimodal training data is often stored as independently sampled streams. A
caller currently has to collect those streams and implement timestamp matching,
episode isolation, snapshot pinning, and payload lookup itself.
This draft proposes a typed, bounded temporal-alignment API for
`MultimodalTable.scan()`:
```python
steps = align(
actions.scan().select(["episode_id", "event_time", "action"]),
on="event_time",
by="episode_id",
camera=nearest(images.scan().select("image"), tolerance=20),
state=backward(states.scan().select("state"), tolerance=50),
)
for batch in steps.to_arrow_batch_reader(batch_size=128):
train(batch)
```
The implementation:
- requires an explicit group boundary so matches never cross episodes/clips;
- supports exact, backward, forward, and nearest matching with inclusive
tolerance and deterministic earlier-frame tie breaking;
- pins every input scan to its current snapshot;
- plans only group keys, timestamps, and `_ROW_ID`, then fetches selected
payload rows in batches;
- keeps BLOB values as descriptors so existing image/video readers can fetch or
decode only selected samples;
- emits validity, matched timestamp, and signed delta columns for auditing;
- exposes `ScanQuery.to_arrow_batch_reader()` as the high-level streaming entry
point.
This is deliberately a bounded MVP. It does not add interpolation, training
windows, watermark/streaming semantics, or a customer-specific JSON spec. The
metadata index is currently coordinator-local. Since this is a draft, feedback
on the API shape and whether the alignment executor belongs in PyPaimon is
especially welcome.
### Tests
```text
python -m pytest \
pypaimon/tests/multimodal_temporal_test.py \
pypaimon/tests/multimodal_table_test.py -q
81 passed
```
```text
python -m flake8 \
pypaimon/multimodal/temporal.py \
pypaimon/tests/multimodal_temporal_test.py
```
A broader local run reached 4306 passed / 102 skipped. Remaining failures were
in suites requiring unavailable Vortex, Mosaic, DuckDB, Java golden data, or
existing Ray/Pandas dtype expectations; no temporal or multimodal test failed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [python] Add as-of joins for multimodal scans [paimon]
Posted by "JingsongLi (via GitHub)" <gi...@apache.org>.
JingsongLi merged PR #9536:
URL: https://github.com/apache/paimon/pull/9536
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [python] Add as-of joins for multimodal scans [paimon]
Posted by "JingsongLi (via GitHub)" <gi...@apache.org>.
JingsongLi commented on code in PR #9536:
URL: https://github.com/apache/paimon/pull/9536#discussion_r3986846365
##########
paimon-python/pypaimon/multimodal/temporal.py:
##########
@@ -0,0 +1,993 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Temporal alignment for multimodal table scans."""
+
+from bisect import bisect_left, bisect_right
+from datetime import timedelta
+import json
+import math
+from numbers import Integral, Real
+
+import pyarrow as pa
+import pyarrow.compute as pc
+
+from pypaimon.catalog.table_query_auth import TableQueryAuthResult
+from pypaimon.common.options.core_options import CoreOptions, StartupMode
+from pypaimon.common.predicate_json_parser import (
+ _apply_predicate_transform,
+ _collect_all_field_refs_from_transform,
+)
+from pypaimon.globalindex.indexed_split import IndexedSplit
+from pypaimon.multimodal.query import ScanQuery
+from pypaimon.read.query_auth_split import QueryAuthSplit
+from pypaimon.read.reader.format_pyarrow_reader import _DecodedRowGroupCache
+from pypaimon.read.table_read import _ClosableArrowBatchReader
+from pypaimon.schema.data_types import PyarrowFieldParser
+from pypaimon.snapshot.time_travel_util import TimeTravelUtil
+from pypaimon.table.special_fields import SpecialFields
+from pypaimon.table.source.global_index_live_row_filter import (
+ table_at_snapshot,
+)
+from pypaimon.utils.range import Range
+
+
+_ROW_ID = SpecialFields.ROW_ID.name
+_MAX_INT64 = (1 << 63) - 1
+_TIME_KEY = object()
+_TEMPORAL_ROW_GROUP_CACHE_MAX_SIZE = 64 * 1024 * 1024
+
+
+def join_asof(left, right, *, on, by, direction="backward", tolerance=None,
+ right_on=None, suffix="_right") -> "AsOfJoin":
+ """Join each left row with at most one time-aligned right row."""
+ if not isinstance(on, str) or not on:
+ raise ValueError("on must be a non-empty column name.")
+ if isinstance(by, str):
+ by = (by,)
+ else:
+ try:
+ by = tuple(by)
+ except TypeError as error:
+ raise ValueError(
+ "by must be a column name or sequence.") from error
+ if not by:
+ raise ValueError(
+ "join_asof requires at least one grouping column in by.")
+ if (any(not isinstance(name, str) or not name for name in by)
+ or len(set(by)) != len(by)):
+ raise ValueError("by must contain unique, non-empty column names.")
+ return AsOfJoin(left, on, by).join_asof(
+ right,
+ direction=direction,
+ tolerance=tolerance,
+ right_on=right_on,
+ suffix=suffix,
+ )
+
+
+class AsOfJoin:
+ """Lazy, chainable result of :func:`join_asof`."""
+
+ def __init__(self, left, on, by):
+ self._anchor = _pin_scan_to_snapshot(_require_scan(left, "left"))
+ self._on = on
+ self._by = by
+ self._sources = ()
+ self._anchor_schema = _query_schema(self._anchor)
+ self._anchor_table_schema = _table_schema(self._anchor)
+ self._validate_anchor()
+ self.schema = self._output_schema()
+
+ def join_asof(self, right, *, direction="backward", tolerance=None,
+ right_on=None, suffix="_right") -> "AsOfJoin":
+ """Append a right-side as-of join without materializing this scan."""
+ position = len(self._sources) + 1
+ label = "right source %d" % position
+ source = _AsOfJoinRight(
+ label,
+ right,
+ self._on,
+ self._by,
+ direction,
+ tolerance,
+ right_on,
+ suffix,
+ )
+
+ result = object.__new__(AsOfJoin)
+ result._anchor = self._anchor
+ result._on = self._on
+ result._by = self._by
+ result._sources = self._sources + (source,)
+ result._anchor_schema = self._anchor_schema
+ result._anchor_table_schema = self._anchor_table_schema
+ result._validate_anchor()
+ result.schema = result._output_schema()
+ return result
+
+ def to_arrow_batch_reader(self, *, batch_size=1024):
+ """Index right-side timestamps, then stream aligned rows in batches."""
+ if (isinstance(batch_size, bool)
+ or not isinstance(batch_size, int)
+ or batch_size <= 0):
+ raise ValueError("batch_size must be a positive integer.")
+
+ anchor_metadata = _metadata_batches(
+ self._anchor, self._on, self._by, batch_size)
+ row_group_cache = _DecodedRowGroupCache(
+ _TEMPORAL_ROW_GROUP_CACHE_MAX_SIZE)
+ anchor_fetcher = _RowIdFetcher(self._anchor, row_group_cache)
+ source_fetchers = []
+ for source in self._sources:
+ source.plan()
+ source_fetchers.append(
+ _RowIdFetcher(source.query, row_group_cache))
+ schema = self._output_schema(anchor_fetcher.schema, source_fetchers)
+ self.schema = schema
+
+ def batches():
+ try:
+ for metadata in anchor_metadata:
+ rows = _metadata_rows(
+ metadata, self._on,
+ self._anchor_table_schema.field(self._on).type)
+ yield from self._build_batches(
+ rows, anchor_fetcher, source_fetchers, schema)
+ finally:
+ anchor_metadata.close()
+
+ batch_iterator = batches()
+ reader = pa.ipc.RecordBatchReader.from_batches(schema, batch_iterator)
+ return _ClosableArrowBatchReader(reader, batch_iterator)
+
+ def to_arrow(self):
+ reader = self.to_arrow_batch_reader()
+ try:
+ return reader.read_all()
+ finally:
+ close = getattr(reader, "close", None)
+ if close is not None:
+ close()
+
+ def to_pandas(self):
+ return self.to_arrow().to_pandas()
+
+ def to_list(self):
+ return _arrow_rows(self.to_arrow())
+
+ @property
+ def resolved_snapshots(self):
+ """Return the table snapshots pinned by this alignment."""
+ snapshots = {
+ "left": _resolved_snapshot(self._anchor),
+ }
+ snapshots.update({
+ "right_%d" % position: _resolved_snapshot(source.query)
+ for position, source in enumerate(self._sources, 1)
+ })
+ return snapshots
+
+ def _validate_anchor(self):
+ _require_columns(
+ self._anchor_table_schema, self._by + (self._on,), "anchor")
+ anchor_type = self._anchor_table_schema.field(self._on).type
+ _delta_type(anchor_type)
+ for name in self._by:
+ _validate_group_type(
+ name, self._anchor_table_schema.field(name).type)
+ for source in self._sources:
+ if source.time_type != anchor_type:
+ raise TypeError(
+ "Left and %s temporal columns must have the same "
+ "type; got %s and %s."
+ % (source.label, anchor_type, source.time_type)
+ )
+ for name in self._by:
+ anchor_group_type = self._anchor_table_schema.field(name).type
+ source_group_type = source.table_schema.field(name).type
+ if source_group_type != anchor_group_type:
+ raise TypeError(
+ "Left and %s grouping column %r must have "
+ "the same type; got %s and %s."
+ % (source.label, name, anchor_group_type,
+ source_group_type)
+ )
+
+ def _output_schema(self, anchor_schema=None, source_fetchers=None):
+ if anchor_schema is None:
+ anchor_schema = self._anchor_schema
+ fields = list(anchor_schema)
+ names = set(anchor_schema.names)
+ for position, source in enumerate(self._sources):
+ payload_schema = (
+ source.payload_schema if source_fetchers is None
+ else source_fetchers[position].schema
+ )
+ for name in source.payload_schema.names:
+ field = payload_schema.field(name)
+ output_name = field.name
+ if output_name in names:
+ output_name += source.suffix
+ if output_name in names:
+ raise ValueError(
+ "%s column %r conflicts after applying suffix %r."
+ % (source.label, field.name, source.suffix)
+ )
+ output = pa.field(
+ output_name, field.type, nullable=True,
+ metadata=field.metadata)
+ fields.append(output)
+ names.add(output.name)
+ return pa.schema(fields, metadata=anchor_schema.metadata)
+
+ def _build_batch(
+ self, anchor_rows, anchor_fetcher, source_fetchers, schema):
+ anchor_ids = [row[_ROW_ID] for row in anchor_rows]
+ anchor = anchor_fetcher.fetch(anchor_ids)
+ anchor.validate()
+ arrays = [anchor[name] for name in self._anchor_schema.names]
+
+ for source, fetcher in zip(self._sources, source_fetchers):
+ matches = [source.match(row) for row in anchor_rows]
+ matched_ids = [match for match in matches if match is not None]
+ unique_ids = list(dict.fromkeys(matched_ids))
+ values = fetcher.fetch(unique_ids)
+ positions = {
+ row_id: index for index, row_id in enumerate(unique_ids)
+ }
+ take = pa.array([
+ None if match is None else positions[match]
+ for match in matches
+ ], type=pa.int64())
+ for field in source.payload_schema:
+ array = pc.take(values[field.name], take)
+ array.validate()
+ arrays.append(array)
+
+ if arrays:
+ table = pa.Table.from_arrays(
+ arrays, schema=schema).combine_chunks()
+ return table.to_batches(max_chunksize=len(anchor_rows))[0]
+ batch = pa.RecordBatch.from_struct_array(pa.array(
+ [{}] * len(anchor_rows), type=pa.struct([])))
+ return batch.replace_schema_metadata(schema.metadata)
+
+ def _build_batches(
+ self, anchor_rows, anchor_fetcher, source_fetchers, schema):
+ try:
+ yield self._build_batch(
+ anchor_rows, anchor_fetcher, source_fetchers, schema)
+ except pa.ArrowInvalid as error:
+ if len(anchor_rows) < 2 or "offset" not in str(error).lower():
+ raise
+ middle = len(anchor_rows) // 2
+ yield from self._build_batches(
+ anchor_rows[:middle], anchor_fetcher,
+ source_fetchers, schema)
+ yield from self._build_batches(
+ anchor_rows[middle:], anchor_fetcher,
+ source_fetchers, schema)
+
+
+class _AsOfJoinRight:
+
+ def __init__(self, label, query, anchor_on, by, direction, tolerance,
+ right_on, suffix):
+ _validate_join_options(direction, tolerance, right_on, suffix)
+ self.label = label
+ self.query = _pin_scan_to_snapshot(_require_scan(query, label))
+ self.direction = direction
+ self.suffix = suffix
+ self.anchor_on = anchor_on
+ self.on = anchor_on if right_on is None else right_on
+ self.by = by
+ self.table_schema = _table_schema(self.query)
+ _require_columns(
+ self.table_schema, by + (self.on,), label)
+ self.time_type = self.table_schema.field(self.on).type
+ _delta_type(self.time_type)
+ _validate_tolerance(tolerance, self.time_type)
+ self._tolerance_key = _time_tolerance_key(tolerance, self.time_type)
+ schema = _query_schema(self.query)
+ projection = self.query._effective_projection()
+ excluded = set(by)
+ if projection is None or self.on not in projection:
+ excluded.add(self.on)
+ self.payload_schema = pa.schema([
+ field for field in schema
+ if field.name not in excluded
+ ])
+ self._index = None
+
+ def plan(self):
+ metadata = _metadata_table(self.query, self.on, self.by)
+ self._times = metadata[self.on].combine_chunks()
+ self._time_keys = _time_search_keys(self._times, self.time_type)
+ self._row_ids = metadata[_ROW_ID].combine_chunks()
+ self._index = {}
+ group_columns = [metadata[name].combine_chunks() for name in self.by]
+ previous = None
+ start = 0
+ for position in range(len(metadata)):
+ key = tuple(column[position].as_py() for column in group_columns)
+ if position and key != previous:
+ self._index[previous] = (start, position)
+ start = position
+ previous = key
+ if len(metadata):
+ self._index[previous] = (start, len(metadata))
+
+ def match(self, anchor_row):
+ key = tuple(anchor_row[name] for name in self.by)
+ bounds = self._index.get(key)
+ if bounds is None:
+ return None
+ target_key = anchor_row[_TIME_KEY]
+ index = _match_index(
+ self._time_keys, target_key, self.direction, *bounds)
+ if index is None:
+ return None
+ matched_key = _python_scalar(self._time_keys[index])
+ if (self._tolerance_key is not None
+ and abs(matched_key - target_key) > self._tolerance_key):
+ return None
+ return self._row_ids[index].as_py()
+
+
+def _validate_join_options(direction, tolerance, right_on, suffix):
+ if direction not in ("backward", "forward", "nearest"):
+ raise ValueError(
+ "direction must be 'backward', 'forward', or 'nearest'.")
+ if right_on is not None and (
+ not isinstance(right_on, str) or not right_on):
+ raise ValueError("right_on must be a non-empty column name.")
+ if not isinstance(suffix, str):
+ raise TypeError("suffix must be a string.")
+ if tolerance is not None:
+ if isinstance(tolerance, bool) or not isinstance(
+ tolerance, (Real, timedelta)):
+ raise TypeError("tolerance must be numeric or datetime.timedelta.")
+ if (isinstance(tolerance, Real)
+ and not isinstance(tolerance, Integral)
+ and not math.isfinite(tolerance)):
+ raise ValueError("tolerance must be finite.")
+ zero = timedelta(0) if isinstance(tolerance, timedelta) else 0
+ if tolerance < zero:
+ raise ValueError("tolerance must be non-negative.")
+
+
+def _require_scan(query, label):
+ if (type(query) is not ScanQuery
+ or getattr(query, "_result_factory", None) is not None):
+ raise TypeError("%s must be a MultimodalTable.scan() query." % label)
+ return query
+
+
+def _pin_scan_to_snapshot(query):
+ table = query._table
+ options = table.options
+ if not options.row_tracking_enabled(False):
+ raise ValueError(
+ "join_asof requires 'row-tracking.enabled' = 'true'.")
+ if (options.scan_mode() == StartupMode.INCREMENTAL
+ or options.options.contains(
+ CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP)):
+ raise ValueError(
+ "join_asof does not support incremental scans; inputs must "
+ "represent a complete point-in-time snapshot.")
+ # Validate the original scan configuration before replacing it with a
+ # pinned snapshot. Otherwise an invalid or unsupported scan mode can be
+ # silently converted into a latest-full scan.
+ table.new_read_builder().new_scan()
+ snapshot = TimeTravelUtil.try_travel_to_snapshot(
+ options.options, table.tag_manager(), table.snapshot_manager())
+ if snapshot is None:
+ snapshot = table.snapshot_manager().get_latest_snapshot()
+ empty = snapshot is None
+ tag_name = (
+ options.scan_tag_name()
+ if options.options.contains_key(CoreOptions.SCAN_TAG_NAME.key())
+ else None
+ )
+ if snapshot is not None and tag_name is None:
+ table = table_at_snapshot(table, snapshot)
+ pinned = ScanQuery(table)
+ pinned._predicate = query._predicate
+ pinned._projection = query._projection
+ pinned._limit = query._limit
+ pinned._include_row_id = query._include_row_id
+ pinned._temporal_empty = empty
+ pinned._temporal_snapshot_id = (
+ None if snapshot is None else snapshot.id)
+ pinned._temporal_tag_name = tag_name
+ return pinned
+
+
+def _resolved_snapshot(query):
+ resolved = {
+ "table": query._table.identifier.get_full_name(),
+ "snapshot_id": query._temporal_snapshot_id,
+ }
+ tag_name = getattr(query, "_temporal_tag_name", None)
+ if tag_name is not None:
+ resolved["tag_name"] = tag_name
+ return resolved
+
+
+def _query_schema(query):
+ table = query._table.copy_without_time_travel({
+ CoreOptions.BLOB_AS_DESCRIPTOR.key(): "true",
+ })
+ builder = query._configured_read_builder(table)
+ return PyarrowFieldParser.from_paimon_schema(builder.read_type())
+
+
+def _table_schema(query):
+ return PyarrowFieldParser.from_paimon_schema(query._table.fields)
+
+
+def _metadata_table(query, on, by):
+ read_builder, splits, key_columns, output_columns = (
+ _metadata_builders(query, on, by))
+ if getattr(query, "_temporal_empty", False):
+ return _empty_metadata(query, key_columns, output_columns)
+ arrow = read_builder.new_read().to_arrow(splits)
+ metadata = arrow.select(output_columns).combine_chunks()
+ _validate_metadata(query, metadata, key_columns)
+ sort_keys = [(name, "ascending") for name in output_columns]
+ return metadata.take(pc.sort_indices(metadata, sort_keys=sort_keys))
+
+
+def _metadata_batches(query, on, by, batch_size):
+ read_builder, splits, key_columns, output_columns = (
+ _metadata_builders(query, on, by))
+ if getattr(query, "_temporal_empty", False):
+ return
+ reader = read_builder.new_read()._to_managed_arrow_batch_reader(splits)
+ try:
+ for batch in reader:
+ metadata = pa.Table.from_batches([batch]).select(output_columns)
+ _validate_metadata(query, metadata, key_columns)
+ for start in range(0, len(metadata), batch_size):
+ yield metadata.slice(start, batch_size)
+ finally:
+ reader.close()
+
+
+def _metadata_builders(query, on, by):
+ _validate_pinned_tag(query)
+ key_columns = list(dict.fromkeys(by + (on,)))
+ output_columns = key_columns + [_ROW_ID]
+ table = query._table.copy_without_time_travel({
+ CoreOptions.BLOB_AS_DESCRIPTOR.key(): "true",
+ })
+ plan_builder = table.new_read_builder()
+ if query._predicate is not None:
+ plan_builder = plan_builder.with_filter(query._predicate)
+ plan_builder = plan_builder.with_projection(key_columns)
+ if query._limit is not None:
+ plan_builder = plan_builder.with_limit(query._limit)
+ splits, masking = _plan_with_internal_row_id(
+ plan_builder, query._temporal_snapshot_id)
+ key_masking = {
+ name: masking[name] for name in key_columns if name in masking
+ }
+ dependencies = _mask_dependencies(
+ key_masking, key_columns, _table_schema(query))
+ splits = _with_active_masking(splits, key_columns)
+ read_projection = list(dict.fromkeys(
+ output_columns + dependencies))
+ read_builder = table.new_read_builder().with_projection(read_projection)
+ if query._predicate is not None:
+ read_builder = read_builder.with_filter(query._predicate)
+ if query._limit is not None:
+ read_builder = read_builder.with_limit(query._limit)
+
+ physical_schema = PyarrowFieldParser.from_paimon_schema(
+ read_builder.read_type())
+ effective_schema = _effective_masked_schema(
+ physical_schema, key_masking)
+ for name in key_columns:
+ physical_type = physical_schema.field(name).type
+ effective_type = effective_schema.field(name).type
+ if effective_type != physical_type:
+ raise TypeError(
+ "Temporal key %r must preserve its type after column "
+ "masking; got %s instead of %s."
+ % (name, effective_type, physical_type)
+ )
+ return read_builder, splits, key_columns, output_columns
+
+
+def _empty_metadata(query, key_columns, output_columns):
+ return pa.Table.from_arrays([
+ pa.array([], type=_table_schema(query).field(name).type)
+ for name in key_columns
+ ] + [pa.array([], type=pa.int64())], names=output_columns)
+
+
+def _validate_metadata(query, metadata, key_columns):
+ for name in key_columns:
+ column = metadata[name]
+ if column.null_count:
+ raise ValueError("Temporal key %r cannot be null." % name)
+ if pa.types.is_floating(column.type):
+ for scalar in column:
+ if not math.isfinite(scalar.as_py()):
+ raise ValueError(
+ "Temporal key %r must be finite." % name)
+
+
+class _RowIdFetcher:
+
+ def __init__(self, query, row_group_cache):
+ _validate_pinned_tag(query)
+ self._schema = _query_schema(query)
+ table = query._table.copy_without_time_travel({
+ CoreOptions.BLOB_AS_DESCRIPTOR.key(): "true",
+ })
+ visible_projection = query._effective_projection()
+ plan_builder = table.new_read_builder()
+ if visible_projection is not None:
+ plan_projection = visible_projection
+ projected = table.new_read_builder().with_projection(
+ visible_projection)
+ projected_paths = projected._nested_name_paths()
+ if projected_paths is not None:
+ plan_projection = list(dict.fromkeys(
+ path[0] for path in projected_paths))
+ plan_builder = plan_builder.with_projection(plan_projection)
+ if query._predicate is not None:
+ plan_builder = plan_builder.with_filter(query._predicate)
+ self._splits, masking = _plan_with_internal_row_id(
+ plan_builder, query._temporal_snapshot_id)
+
+ read_projection = (
+ [field.name for field in table.fields]
+ if not visible_projection else list(visible_projection)
+ )
+ projected_builder = table.new_read_builder().with_projection(
+ list(dict.fromkeys(read_projection + [_ROW_ID])))
+ projected_schema = PyarrowFieldParser.from_paimon_schema(
+ projected_builder.read_type())
+ projected_paths = projected_builder._nested_name_paths()
+ if projected_paths is not None:
+ for field, path in zip(projected_schema, projected_paths):
+ if field.name in masking and field.name != path[0]:
+ raise ValueError(
+ "Temporal alignment cannot safely apply column "
+ "masking to nested projection %r."
+ % ".".join(path)
+ )
+ if projected_paths is None:
+ active_targets = projected_schema.names
+ else:
+ active_targets = list(dict.fromkeys(
+ path[0] for path in projected_paths))
+ dependencies = _mask_dependencies(
+ masking, active_targets, _table_schema(query))
+ read_projection = list(dict.fromkeys(
+ read_projection + dependencies + [_ROW_ID]))
+ visible_builder = table.new_read_builder().with_projection(
+ read_projection)
+ self._fetch_schema = PyarrowFieldParser.from_paimon_schema(
+ visible_builder.read_type())
+ self._name_paths = visible_builder._nested_name_paths()
+ if self._name_paths is None:
+ builder = visible_builder
+ else:
+ top_level = list(dict.fromkeys(
+ path[0] for path in self._name_paths))
+ builder = table.new_read_builder().with_projection(top_level)
+ if query._predicate is not None:
+ builder = builder.with_filter(query._predicate)
+ self._read = builder.new_read()
+ self._read._parquet_row_group_cache = row_group_cache
+ physical_schema = PyarrowFieldParser.from_paimon_schema(
+ builder.read_type())
+ effective_schema = _effective_masked_schema(
+ physical_schema, masking)
+ visible_paths = (
+ None if self._name_paths is None
+ else self._name_paths[:len(self._schema)]
+ )
+ self._schema = _project_effective_schema(
+ self._schema, visible_paths, effective_schema, masking)
+ self._fetch_schema = _project_effective_schema(
+ self._fetch_schema, self._name_paths,
+ effective_schema, masking)
+ self._split_ranges = [
+ self._row_ranges(split) for split in self._splits]
+ self._range_intervals = sorted(
+ (row_range.from_, row_range.to, split_index)
+ for split_index, ranges in enumerate(self._split_ranges)
+ for row_range in ranges
+ )
+ self._range_starts = [
+ interval[0] for interval in self._range_intervals]
+ self._range_max_ends = []
+ max_end = -1
+ for _, end, _ in self._range_intervals:
+ max_end = max(max_end, end)
+ self._range_max_ends.append(max_end)
+
+ @property
+ def schema(self):
+ return self._schema
+
+ @staticmethod
+ def _row_ranges(split):
+ if isinstance(split, QueryAuthSplit):
+ split = split.split
+ if isinstance(split, IndexedSplit):
+ ranges = split.row_ranges()
+ else:
+ ranges = [
+ data_file.row_id_range()
+ for data_file in split.files
+ if data_file.row_id_range() is not None
+ ]
+ return Range.sort_and_merge_overlap(ranges, True)
+
+ def fetch(self, row_ids):
+ if not row_ids:
+ return pa.Table.from_arrays(
+ [pa.array([], type=field.type) for field in self._schema],
+ schema=self._schema,
+ )
+
+ wanted = Range.sort_and_merge_overlap(
+ [Range(row_id, row_id) for row_id in set(row_ids)], True)
+ selected_splits = []
+ for split_index in self._find_splits(wanted):
+ original = self._splits[split_index]
+ auth_result = None
+ split = original
+ if isinstance(split, QueryAuthSplit):
+ auth_result = split.auth_result
+ split = split.split
+ if isinstance(split, IndexedSplit):
+ split = split.data_split()
+ allowed = Range.and_(wanted, self._split_ranges[split_index])
+ if not allowed:
+ continue
+ indexed = IndexedSplit(
+ split,
+ allowed,
+ exact_merged_row_count=sum(r.count() for r in allowed),
+ )
+ if auth_result is not None:
+ indexed = QueryAuthSplit(indexed, auth_result)
+ selected_splits.append(indexed)
+
+ arrow = self._project_fetch(self._read.to_arrow(selected_splits))
+ found = arrow[_ROW_ID].to_pylist()
Review Comment:
**[P2] Preserve the physical row ID before flattening projections**
A valid nested projection such as `left.scan().select(['_ROW.ID', 't'])` flattens the payload field to `_ROW_ID`, so the actual internal row ID is renamed to `_ROW_ID__0`. After `_project_fetch`, this lookup therefore builds the position map from payload values. With left rows `(t=10, _ROW.ID=1)` and `(t=20, _ROW.ID=0)` in one group, and right values `10` and `20` at the corresponding timestamps, an exact join returns `(t=20, value=10)` and `(t=10, value=20)`: the left rows are silently paired with the wrong right payloads. Please extract the physical row IDs before flattening, or track the internal field by its original identity rather than its projected name.
##########
paimon-python/pypaimon/multimodal/temporal.py:
##########
@@ -0,0 +1,993 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Temporal alignment for multimodal table scans."""
+
+from bisect import bisect_left, bisect_right
+from datetime import timedelta
+import json
+import math
+from numbers import Integral, Real
+
+import pyarrow as pa
+import pyarrow.compute as pc
+
+from pypaimon.catalog.table_query_auth import TableQueryAuthResult
+from pypaimon.common.options.core_options import CoreOptions, StartupMode
+from pypaimon.common.predicate_json_parser import (
+ _apply_predicate_transform,
+ _collect_all_field_refs_from_transform,
+)
+from pypaimon.globalindex.indexed_split import IndexedSplit
+from pypaimon.multimodal.query import ScanQuery
+from pypaimon.read.query_auth_split import QueryAuthSplit
+from pypaimon.read.reader.format_pyarrow_reader import _DecodedRowGroupCache
+from pypaimon.read.table_read import _ClosableArrowBatchReader
+from pypaimon.schema.data_types import PyarrowFieldParser
+from pypaimon.snapshot.time_travel_util import TimeTravelUtil
+from pypaimon.table.special_fields import SpecialFields
+from pypaimon.table.source.global_index_live_row_filter import (
+ table_at_snapshot,
+)
+from pypaimon.utils.range import Range
+
+
+_ROW_ID = SpecialFields.ROW_ID.name
+_MAX_INT64 = (1 << 63) - 1
+_TIME_KEY = object()
+_TEMPORAL_ROW_GROUP_CACHE_MAX_SIZE = 64 * 1024 * 1024
+
+
+def join_asof(left, right, *, on, by, direction="backward", tolerance=None,
+ right_on=None, suffix="_right") -> "AsOfJoin":
+ """Join each left row with at most one time-aligned right row."""
+ if not isinstance(on, str) or not on:
+ raise ValueError("on must be a non-empty column name.")
+ if isinstance(by, str):
+ by = (by,)
+ else:
+ try:
+ by = tuple(by)
+ except TypeError as error:
+ raise ValueError(
+ "by must be a column name or sequence.") from error
+ if not by:
+ raise ValueError(
+ "join_asof requires at least one grouping column in by.")
+ if (any(not isinstance(name, str) or not name for name in by)
+ or len(set(by)) != len(by)):
+ raise ValueError("by must contain unique, non-empty column names.")
+ return AsOfJoin(left, on, by).join_asof(
+ right,
+ direction=direction,
+ tolerance=tolerance,
+ right_on=right_on,
+ suffix=suffix,
+ )
+
+
+class AsOfJoin:
+ """Lazy, chainable result of :func:`join_asof`."""
+
+ def __init__(self, left, on, by):
+ self._anchor = _pin_scan_to_snapshot(_require_scan(left, "left"))
+ self._on = on
+ self._by = by
+ self._sources = ()
+ self._anchor_schema = _query_schema(self._anchor)
+ self._anchor_table_schema = _table_schema(self._anchor)
+ self._validate_anchor()
+ self.schema = self._output_schema()
+
+ def join_asof(self, right, *, direction="backward", tolerance=None,
+ right_on=None, suffix="_right") -> "AsOfJoin":
+ """Append a right-side as-of join without materializing this scan."""
+ position = len(self._sources) + 1
+ label = "right source %d" % position
+ source = _AsOfJoinRight(
+ label,
+ right,
+ self._on,
+ self._by,
+ direction,
+ tolerance,
+ right_on,
+ suffix,
+ )
+
+ result = object.__new__(AsOfJoin)
+ result._anchor = self._anchor
+ result._on = self._on
+ result._by = self._by
+ result._sources = self._sources + (source,)
+ result._anchor_schema = self._anchor_schema
+ result._anchor_table_schema = self._anchor_table_schema
+ result._validate_anchor()
+ result.schema = result._output_schema()
+ return result
+
+ def to_arrow_batch_reader(self, *, batch_size=1024):
+ """Index right-side timestamps, then stream aligned rows in batches."""
+ if (isinstance(batch_size, bool)
+ or not isinstance(batch_size, int)
+ or batch_size <= 0):
+ raise ValueError("batch_size must be a positive integer.")
+
+ anchor_metadata = _metadata_batches(
+ self._anchor, self._on, self._by, batch_size)
+ row_group_cache = _DecodedRowGroupCache(
+ _TEMPORAL_ROW_GROUP_CACHE_MAX_SIZE)
+ anchor_fetcher = _RowIdFetcher(self._anchor, row_group_cache)
+ source_fetchers = []
+ for source in self._sources:
+ source.plan()
+ source_fetchers.append(
+ _RowIdFetcher(source.query, row_group_cache))
+ schema = self._output_schema(anchor_fetcher.schema, source_fetchers)
+ self.schema = schema
+
+ def batches():
+ try:
+ for metadata in anchor_metadata:
+ rows = _metadata_rows(
+ metadata, self._on,
+ self._anchor_table_schema.field(self._on).type)
+ yield from self._build_batches(
+ rows, anchor_fetcher, source_fetchers, schema)
+ finally:
+ anchor_metadata.close()
+
+ batch_iterator = batches()
+ reader = pa.ipc.RecordBatchReader.from_batches(schema, batch_iterator)
+ return _ClosableArrowBatchReader(reader, batch_iterator)
+
+ def to_arrow(self):
+ reader = self.to_arrow_batch_reader()
+ try:
+ return reader.read_all()
+ finally:
+ close = getattr(reader, "close", None)
+ if close is not None:
+ close()
+
+ def to_pandas(self):
+ return self.to_arrow().to_pandas()
+
+ def to_list(self):
+ return _arrow_rows(self.to_arrow())
+
+ @property
+ def resolved_snapshots(self):
+ """Return the table snapshots pinned by this alignment."""
+ snapshots = {
+ "left": _resolved_snapshot(self._anchor),
+ }
+ snapshots.update({
+ "right_%d" % position: _resolved_snapshot(source.query)
+ for position, source in enumerate(self._sources, 1)
+ })
+ return snapshots
+
+ def _validate_anchor(self):
+ _require_columns(
+ self._anchor_table_schema, self._by + (self._on,), "anchor")
+ anchor_type = self._anchor_table_schema.field(self._on).type
+ _delta_type(anchor_type)
+ for name in self._by:
+ _validate_group_type(
+ name, self._anchor_table_schema.field(name).type)
+ for source in self._sources:
+ if source.time_type != anchor_type:
+ raise TypeError(
+ "Left and %s temporal columns must have the same "
+ "type; got %s and %s."
+ % (source.label, anchor_type, source.time_type)
+ )
+ for name in self._by:
+ anchor_group_type = self._anchor_table_schema.field(name).type
+ source_group_type = source.table_schema.field(name).type
+ if source_group_type != anchor_group_type:
+ raise TypeError(
+ "Left and %s grouping column %r must have "
+ "the same type; got %s and %s."
+ % (source.label, name, anchor_group_type,
+ source_group_type)
+ )
+
+ def _output_schema(self, anchor_schema=None, source_fetchers=None):
+ if anchor_schema is None:
+ anchor_schema = self._anchor_schema
+ fields = list(anchor_schema)
+ names = set(anchor_schema.names)
+ for position, source in enumerate(self._sources):
+ payload_schema = (
+ source.payload_schema if source_fetchers is None
+ else source_fetchers[position].schema
+ )
+ for name in source.payload_schema.names:
+ field = payload_schema.field(name)
+ output_name = field.name
+ if output_name in names:
+ output_name += source.suffix
+ if output_name in names:
+ raise ValueError(
+ "%s column %r conflicts after applying suffix %r."
+ % (source.label, field.name, source.suffix)
+ )
+ output = pa.field(
+ output_name, field.type, nullable=True,
+ metadata=field.metadata)
+ fields.append(output)
+ names.add(output.name)
+ return pa.schema(fields, metadata=anchor_schema.metadata)
+
+ def _build_batch(
+ self, anchor_rows, anchor_fetcher, source_fetchers, schema):
+ anchor_ids = [row[_ROW_ID] for row in anchor_rows]
+ anchor = anchor_fetcher.fetch(anchor_ids)
+ anchor.validate()
+ arrays = [anchor[name] for name in self._anchor_schema.names]
+
+ for source, fetcher in zip(self._sources, source_fetchers):
+ matches = [source.match(row) for row in anchor_rows]
+ matched_ids = [match for match in matches if match is not None]
+ unique_ids = list(dict.fromkeys(matched_ids))
+ values = fetcher.fetch(unique_ids)
+ positions = {
+ row_id: index for index, row_id in enumerate(unique_ids)
+ }
+ take = pa.array([
+ None if match is None else positions[match]
+ for match in matches
+ ], type=pa.int64())
+ for field in source.payload_schema:
+ array = pc.take(values[field.name], take)
+ array.validate()
+ arrays.append(array)
+
+ if arrays:
+ table = pa.Table.from_arrays(
+ arrays, schema=schema).combine_chunks()
+ return table.to_batches(max_chunksize=len(anchor_rows))[0]
+ batch = pa.RecordBatch.from_struct_array(pa.array(
+ [{}] * len(anchor_rows), type=pa.struct([])))
+ return batch.replace_schema_metadata(schema.metadata)
+
+ def _build_batches(
+ self, anchor_rows, anchor_fetcher, source_fetchers, schema):
+ try:
+ yield self._build_batch(
+ anchor_rows, anchor_fetcher, source_fetchers, schema)
+ except pa.ArrowInvalid as error:
+ if len(anchor_rows) < 2 or "offset" not in str(error).lower():
+ raise
+ middle = len(anchor_rows) // 2
+ yield from self._build_batches(
+ anchor_rows[:middle], anchor_fetcher,
+ source_fetchers, schema)
+ yield from self._build_batches(
+ anchor_rows[middle:], anchor_fetcher,
+ source_fetchers, schema)
+
+
+class _AsOfJoinRight:
+
+ def __init__(self, label, query, anchor_on, by, direction, tolerance,
+ right_on, suffix):
+ _validate_join_options(direction, tolerance, right_on, suffix)
+ self.label = label
+ self.query = _pin_scan_to_snapshot(_require_scan(query, label))
+ self.direction = direction
+ self.suffix = suffix
+ self.anchor_on = anchor_on
+ self.on = anchor_on if right_on is None else right_on
+ self.by = by
+ self.table_schema = _table_schema(self.query)
+ _require_columns(
+ self.table_schema, by + (self.on,), label)
+ self.time_type = self.table_schema.field(self.on).type
+ _delta_type(self.time_type)
+ _validate_tolerance(tolerance, self.time_type)
+ self._tolerance_key = _time_tolerance_key(tolerance, self.time_type)
+ schema = _query_schema(self.query)
+ projection = self.query._effective_projection()
+ excluded = set(by)
+ if projection is None or self.on not in projection:
+ excluded.add(self.on)
+ self.payload_schema = pa.schema([
+ field for field in schema
+ if field.name not in excluded
+ ])
+ self._index = None
+
+ def plan(self):
+ metadata = _metadata_table(self.query, self.on, self.by)
+ self._times = metadata[self.on].combine_chunks()
+ self._time_keys = _time_search_keys(self._times, self.time_type)
+ self._row_ids = metadata[_ROW_ID].combine_chunks()
+ self._index = {}
+ group_columns = [metadata[name].combine_chunks() for name in self.by]
+ previous = None
+ start = 0
+ for position in range(len(metadata)):
+ key = tuple(column[position].as_py() for column in group_columns)
+ if position and key != previous:
+ self._index[previous] = (start, position)
+ start = position
+ previous = key
+ if len(metadata):
+ self._index[previous] = (start, len(metadata))
+
+ def match(self, anchor_row):
+ key = tuple(anchor_row[name] for name in self.by)
+ bounds = self._index.get(key)
+ if bounds is None:
+ return None
+ target_key = anchor_row[_TIME_KEY]
+ index = _match_index(
+ self._time_keys, target_key, self.direction, *bounds)
+ if index is None:
+ return None
+ matched_key = _python_scalar(self._time_keys[index])
+ if (self._tolerance_key is not None
+ and abs(matched_key - target_key) > self._tolerance_key):
+ return None
+ return self._row_ids[index].as_py()
+
+
+def _validate_join_options(direction, tolerance, right_on, suffix):
+ if direction not in ("backward", "forward", "nearest"):
+ raise ValueError(
+ "direction must be 'backward', 'forward', or 'nearest'.")
+ if right_on is not None and (
+ not isinstance(right_on, str) or not right_on):
+ raise ValueError("right_on must be a non-empty column name.")
+ if not isinstance(suffix, str):
+ raise TypeError("suffix must be a string.")
+ if tolerance is not None:
+ if isinstance(tolerance, bool) or not isinstance(
+ tolerance, (Real, timedelta)):
+ raise TypeError("tolerance must be numeric or datetime.timedelta.")
+ if (isinstance(tolerance, Real)
+ and not isinstance(tolerance, Integral)
+ and not math.isfinite(tolerance)):
+ raise ValueError("tolerance must be finite.")
+ zero = timedelta(0) if isinstance(tolerance, timedelta) else 0
+ if tolerance < zero:
+ raise ValueError("tolerance must be non-negative.")
+
+
+def _require_scan(query, label):
+ if (type(query) is not ScanQuery
+ or getattr(query, "_result_factory", None) is not None):
+ raise TypeError("%s must be a MultimodalTable.scan() query." % label)
+ return query
+
+
+def _pin_scan_to_snapshot(query):
+ table = query._table
+ options = table.options
+ if not options.row_tracking_enabled(False):
+ raise ValueError(
+ "join_asof requires 'row-tracking.enabled' = 'true'.")
+ if (options.scan_mode() == StartupMode.INCREMENTAL
+ or options.options.contains(
+ CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP)):
+ raise ValueError(
+ "join_asof does not support incremental scans; inputs must "
+ "represent a complete point-in-time snapshot.")
+ # Validate the original scan configuration before replacing it with a
+ # pinned snapshot. Otherwise an invalid or unsupported scan mode can be
+ # silently converted into a latest-full scan.
+ table.new_read_builder().new_scan()
+ snapshot = TimeTravelUtil.try_travel_to_snapshot(
+ options.options, table.tag_manager(), table.snapshot_manager())
+ if snapshot is None:
+ snapshot = table.snapshot_manager().get_latest_snapshot()
+ empty = snapshot is None
+ tag_name = (
+ options.scan_tag_name()
+ if options.options.contains_key(CoreOptions.SCAN_TAG_NAME.key())
+ else None
+ )
+ if snapshot is not None and tag_name is None:
+ table = table_at_snapshot(table, snapshot)
+ pinned = ScanQuery(table)
+ pinned._predicate = query._predicate
+ pinned._projection = query._projection
+ pinned._limit = query._limit
+ pinned._include_row_id = query._include_row_id
+ pinned._temporal_empty = empty
+ pinned._temporal_snapshot_id = (
+ None if snapshot is None else snapshot.id)
+ pinned._temporal_tag_name = tag_name
+ return pinned
+
+
+def _resolved_snapshot(query):
+ resolved = {
+ "table": query._table.identifier.get_full_name(),
+ "snapshot_id": query._temporal_snapshot_id,
+ }
+ tag_name = getattr(query, "_temporal_tag_name", None)
+ if tag_name is not None:
+ resolved["tag_name"] = tag_name
+ return resolved
+
+
+def _query_schema(query):
+ table = query._table.copy_without_time_travel({
+ CoreOptions.BLOB_AS_DESCRIPTOR.key(): "true",
+ })
+ builder = query._configured_read_builder(table)
+ return PyarrowFieldParser.from_paimon_schema(builder.read_type())
+
+
+def _table_schema(query):
+ return PyarrowFieldParser.from_paimon_schema(query._table.fields)
+
+
+def _metadata_table(query, on, by):
+ read_builder, splits, key_columns, output_columns = (
+ _metadata_builders(query, on, by))
+ if getattr(query, "_temporal_empty", False):
+ return _empty_metadata(query, key_columns, output_columns)
Review Comment:
**[P2] Preserve an input pinned before its first snapshot**
The empty-input check happens after `_metadata_builders()` has already planned a scan. If an alignment is constructed with a nonempty left table and a newly created empty right table, then the right table receives its first append before the alignment is consumed, `to_list()` raises `Temporal input changed from snapshot None to 1 during planning`. The same alignment returns null matches before that append. This breaks the documented snapshot-pinned behavior for an initially empty optional stream. Please bypass physical snapshot planning for `_temporal_empty` while preserving authorization and the effective schema; `_RowIdFetcher` also plans unconditionally and needs the same treatment.
##########
paimon-python/pypaimon/multimodal/temporal.py:
##########
@@ -0,0 +1,993 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Temporal alignment for multimodal table scans."""
+
+from bisect import bisect_left, bisect_right
+from datetime import timedelta
+import json
+import math
+from numbers import Integral, Real
+
+import pyarrow as pa
+import pyarrow.compute as pc
+
+from pypaimon.catalog.table_query_auth import TableQueryAuthResult
+from pypaimon.common.options.core_options import CoreOptions, StartupMode
+from pypaimon.common.predicate_json_parser import (
+ _apply_predicate_transform,
+ _collect_all_field_refs_from_transform,
+)
+from pypaimon.globalindex.indexed_split import IndexedSplit
+from pypaimon.multimodal.query import ScanQuery
+from pypaimon.read.query_auth_split import QueryAuthSplit
+from pypaimon.read.reader.format_pyarrow_reader import _DecodedRowGroupCache
+from pypaimon.read.table_read import _ClosableArrowBatchReader
+from pypaimon.schema.data_types import PyarrowFieldParser
+from pypaimon.snapshot.time_travel_util import TimeTravelUtil
+from pypaimon.table.special_fields import SpecialFields
+from pypaimon.table.source.global_index_live_row_filter import (
+ table_at_snapshot,
+)
+from pypaimon.utils.range import Range
+
+
+_ROW_ID = SpecialFields.ROW_ID.name
+_MAX_INT64 = (1 << 63) - 1
+_TIME_KEY = object()
+_TEMPORAL_ROW_GROUP_CACHE_MAX_SIZE = 64 * 1024 * 1024
+
+
+def join_asof(left, right, *, on, by, direction="backward", tolerance=None,
+ right_on=None, suffix="_right") -> "AsOfJoin":
+ """Join each left row with at most one time-aligned right row."""
+ if not isinstance(on, str) or not on:
+ raise ValueError("on must be a non-empty column name.")
+ if isinstance(by, str):
+ by = (by,)
+ else:
+ try:
+ by = tuple(by)
+ except TypeError as error:
+ raise ValueError(
+ "by must be a column name or sequence.") from error
+ if not by:
+ raise ValueError(
+ "join_asof requires at least one grouping column in by.")
+ if (any(not isinstance(name, str) or not name for name in by)
+ or len(set(by)) != len(by)):
+ raise ValueError("by must contain unique, non-empty column names.")
+ return AsOfJoin(left, on, by).join_asof(
+ right,
+ direction=direction,
+ tolerance=tolerance,
+ right_on=right_on,
+ suffix=suffix,
+ )
+
+
+class AsOfJoin:
+ """Lazy, chainable result of :func:`join_asof`."""
+
+ def __init__(self, left, on, by):
+ self._anchor = _pin_scan_to_snapshot(_require_scan(left, "left"))
+ self._on = on
+ self._by = by
+ self._sources = ()
+ self._anchor_schema = _query_schema(self._anchor)
+ self._anchor_table_schema = _table_schema(self._anchor)
+ self._validate_anchor()
+ self.schema = self._output_schema()
+
+ def join_asof(self, right, *, direction="backward", tolerance=None,
+ right_on=None, suffix="_right") -> "AsOfJoin":
+ """Append a right-side as-of join without materializing this scan."""
+ position = len(self._sources) + 1
+ label = "right source %d" % position
+ source = _AsOfJoinRight(
+ label,
+ right,
+ self._on,
+ self._by,
+ direction,
+ tolerance,
+ right_on,
+ suffix,
+ )
+
+ result = object.__new__(AsOfJoin)
+ result._anchor = self._anchor
+ result._on = self._on
+ result._by = self._by
+ result._sources = self._sources + (source,)
+ result._anchor_schema = self._anchor_schema
+ result._anchor_table_schema = self._anchor_table_schema
+ result._validate_anchor()
+ result.schema = result._output_schema()
+ return result
+
+ def to_arrow_batch_reader(self, *, batch_size=1024):
+ """Index right-side timestamps, then stream aligned rows in batches."""
+ if (isinstance(batch_size, bool)
+ or not isinstance(batch_size, int)
+ or batch_size <= 0):
+ raise ValueError("batch_size must be a positive integer.")
+
+ anchor_metadata = _metadata_batches(
+ self._anchor, self._on, self._by, batch_size)
+ row_group_cache = _DecodedRowGroupCache(
+ _TEMPORAL_ROW_GROUP_CACHE_MAX_SIZE)
+ anchor_fetcher = _RowIdFetcher(self._anchor, row_group_cache)
+ source_fetchers = []
+ for source in self._sources:
+ source.plan()
+ source_fetchers.append(
+ _RowIdFetcher(source.query, row_group_cache))
+ schema = self._output_schema(anchor_fetcher.schema, source_fetchers)
+ self.schema = schema
+
+ def batches():
+ try:
+ for metadata in anchor_metadata:
+ rows = _metadata_rows(
+ metadata, self._on,
+ self._anchor_table_schema.field(self._on).type)
+ yield from self._build_batches(
+ rows, anchor_fetcher, source_fetchers, schema)
+ finally:
+ anchor_metadata.close()
+
+ batch_iterator = batches()
+ reader = pa.ipc.RecordBatchReader.from_batches(schema, batch_iterator)
+ return _ClosableArrowBatchReader(reader, batch_iterator)
+
+ def to_arrow(self):
+ reader = self.to_arrow_batch_reader()
+ try:
+ return reader.read_all()
+ finally:
+ close = getattr(reader, "close", None)
+ if close is not None:
+ close()
+
+ def to_pandas(self):
+ return self.to_arrow().to_pandas()
+
+ def to_list(self):
+ return _arrow_rows(self.to_arrow())
+
+ @property
+ def resolved_snapshots(self):
+ """Return the table snapshots pinned by this alignment."""
+ snapshots = {
+ "left": _resolved_snapshot(self._anchor),
+ }
+ snapshots.update({
+ "right_%d" % position: _resolved_snapshot(source.query)
+ for position, source in enumerate(self._sources, 1)
+ })
+ return snapshots
+
+ def _validate_anchor(self):
+ _require_columns(
+ self._anchor_table_schema, self._by + (self._on,), "anchor")
+ anchor_type = self._anchor_table_schema.field(self._on).type
+ _delta_type(anchor_type)
+ for name in self._by:
+ _validate_group_type(
+ name, self._anchor_table_schema.field(name).type)
+ for source in self._sources:
+ if source.time_type != anchor_type:
+ raise TypeError(
+ "Left and %s temporal columns must have the same "
+ "type; got %s and %s."
+ % (source.label, anchor_type, source.time_type)
+ )
+ for name in self._by:
+ anchor_group_type = self._anchor_table_schema.field(name).type
+ source_group_type = source.table_schema.field(name).type
+ if source_group_type != anchor_group_type:
+ raise TypeError(
+ "Left and %s grouping column %r must have "
+ "the same type; got %s and %s."
+ % (source.label, name, anchor_group_type,
+ source_group_type)
+ )
+
+ def _output_schema(self, anchor_schema=None, source_fetchers=None):
+ if anchor_schema is None:
+ anchor_schema = self._anchor_schema
+ fields = list(anchor_schema)
+ names = set(anchor_schema.names)
+ for position, source in enumerate(self._sources):
+ payload_schema = (
+ source.payload_schema if source_fetchers is None
+ else source_fetchers[position].schema
+ )
+ for name in source.payload_schema.names:
+ field = payload_schema.field(name)
+ output_name = field.name
+ if output_name in names:
+ output_name += source.suffix
+ if output_name in names:
+ raise ValueError(
+ "%s column %r conflicts after applying suffix %r."
+ % (source.label, field.name, source.suffix)
+ )
+ output = pa.field(
+ output_name, field.type, nullable=True,
+ metadata=field.metadata)
+ fields.append(output)
+ names.add(output.name)
+ return pa.schema(fields, metadata=anchor_schema.metadata)
+
+ def _build_batch(
+ self, anchor_rows, anchor_fetcher, source_fetchers, schema):
+ anchor_ids = [row[_ROW_ID] for row in anchor_rows]
+ anchor = anchor_fetcher.fetch(anchor_ids)
+ anchor.validate()
+ arrays = [anchor[name] for name in self._anchor_schema.names]
+
+ for source, fetcher in zip(self._sources, source_fetchers):
+ matches = [source.match(row) for row in anchor_rows]
+ matched_ids = [match for match in matches if match is not None]
+ unique_ids = list(dict.fromkeys(matched_ids))
+ values = fetcher.fetch(unique_ids)
+ positions = {
+ row_id: index for index, row_id in enumerate(unique_ids)
+ }
+ take = pa.array([
+ None if match is None else positions[match]
+ for match in matches
+ ], type=pa.int64())
+ for field in source.payload_schema:
+ array = pc.take(values[field.name], take)
+ array.validate()
+ arrays.append(array)
+
+ if arrays:
+ table = pa.Table.from_arrays(
+ arrays, schema=schema).combine_chunks()
+ return table.to_batches(max_chunksize=len(anchor_rows))[0]
+ batch = pa.RecordBatch.from_struct_array(pa.array(
+ [{}] * len(anchor_rows), type=pa.struct([])))
+ return batch.replace_schema_metadata(schema.metadata)
+
+ def _build_batches(
+ self, anchor_rows, anchor_fetcher, source_fetchers, schema):
+ try:
+ yield self._build_batch(
+ anchor_rows, anchor_fetcher, source_fetchers, schema)
+ except pa.ArrowInvalid as error:
+ if len(anchor_rows) < 2 or "offset" not in str(error).lower():
+ raise
+ middle = len(anchor_rows) // 2
+ yield from self._build_batches(
+ anchor_rows[:middle], anchor_fetcher,
+ source_fetchers, schema)
+ yield from self._build_batches(
+ anchor_rows[middle:], anchor_fetcher,
+ source_fetchers, schema)
+
+
+class _AsOfJoinRight:
+
+ def __init__(self, label, query, anchor_on, by, direction, tolerance,
+ right_on, suffix):
+ _validate_join_options(direction, tolerance, right_on, suffix)
+ self.label = label
+ self.query = _pin_scan_to_snapshot(_require_scan(query, label))
+ self.direction = direction
+ self.suffix = suffix
+ self.anchor_on = anchor_on
+ self.on = anchor_on if right_on is None else right_on
+ self.by = by
+ self.table_schema = _table_schema(self.query)
+ _require_columns(
+ self.table_schema, by + (self.on,), label)
+ self.time_type = self.table_schema.field(self.on).type
+ _delta_type(self.time_type)
+ _validate_tolerance(tolerance, self.time_type)
+ self._tolerance_key = _time_tolerance_key(tolerance, self.time_type)
+ schema = _query_schema(self.query)
+ projection = self.query._effective_projection()
+ excluded = set(by)
+ if projection is None or self.on not in projection:
+ excluded.add(self.on)
+ self.payload_schema = pa.schema([
+ field for field in schema
+ if field.name not in excluded
Review Comment:
**[P2] Exclude join keys by their original projection paths**
These names are already flattened, while `by` and `on` refer to original table columns. For a right table containing the grouping column `payload_value` and a separate `payload: struct<value: int>`, using `by='payload_value'` with `right.scan().select('payload.value')` drops the selected payload because it also flattens to `payload_value`. I reproduced the payload disappearing entirely; selecting `['payload.value', 'payload_value']` instead retains the renamed grouping value `1` rather than the requested payload value `7`. Please decide which fields are join keys from their original paths or field identities, then retain the corresponding projected output fields.
##########
paimon-python/pypaimon/read/reader/format_pyarrow_reader.py:
##########
@@ -397,18 +414,32 @@ def _iter_reader_batches(reader):
def _iter_row_group_batches(self):
columns = self._row_group_read_columns()
for row_group in self._surviving_row_group_ids():
- for batch in self._parquet_file.iter_batches(
- row_groups=[row_group],
- columns=columns,
- batch_size=self._scan_batch_size):
+ if (self._row_group_cache is not None
+ and self._selected_parquet_row_groups is not None):
+ key = (
+ self._row_group_cache_filesystem,
+ self._row_group_cache_path,
+ row_group,
+ tuple(columns),
+ self._scan_batch_size,
+ )
+ batches = self._row_group_cache.get_or_load(
+ key,
+ lambda: list(self._read_parquet_row_group_batches(
+ row_group, columns)),
Review Comment:
**[P2] Bound decoding memory before materializing a cached row group**
`list(...)` retains the entire decoded row group before the cache checks its 64 MiB limit, so a small output batch no longer bounds payload memory. In an end-to-end join using a normally written right table with 4,096 BYTES values of 32 KiB each, `read.batch-size=32` and output `batch_size=1`, peak Arrow allocation was about 128 MiB with this cache versus 3 MiB with the cache disabled. The oversized entry is immediately evicted, and each output batch repeats the full materialization. Larger payload groups can therefore exhaust memory despite the cache limit. Please use bounded incremental caching, or keep groups that cannot fit on a streaming read path.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org