[Python] Complete Native read, write and update parity with Java (#10… · apache/paimon@247f7d8 · GitHub
Skip to content

Commit 247f7d8

Browse files
authored
[Python] Complete Native read, write and update parity with Java (#10674)
1 parent b7e9447 commit 247f7d8

49 files changed

Lines changed: 3284 additions & 125 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎paimon-python/pypaimon/common/predicate.py‎

Lines changed: 28 additions & 0 deletions

‎paimon-python/pypaimon/read/native_plan.py‎

Lines changed: 45 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -75,15 +75,55 @@ def _native_blob_view_supported(table, read_names) -> bool:
7575
return loader is None or _catalog_metastore(loader) == 'rest'
7676

7777

78-
def native_split_from_python(split):
79-
"""Convert a Python-planned DataSplit/IndexedSplit for the Rust reader.
78+
def _native_read_split(split, ranges_use_row_ids):
79+
"""Express Python's file slices with the existing indexed-split selector."""
80+
from pypaimon.globalindex.indexed_split import IndexedSplit
81+
from pypaimon.read.query_auth_split import QueryAuthSplit
82+
from pypaimon.read.sliced_split import SlicedSplit
83+
from pypaimon.utils.range import Range
84+
85+
def select_rows(base, ranges):
86+
if isinstance(base, QueryAuthSplit):
87+
return QueryAuthSplit(select_rows(base.split, ranges), base.auth_result)
88+
if isinstance(base, IndexedSplit):
89+
return select_rows(base.data_split(), Range.and_(base.row_ranges(), ranges))
90+
return IndexedSplit(base, ranges)
91+
92+
if isinstance(split, QueryAuthSplit):
93+
base = _native_read_split(split.split, ranges_use_row_ids)
94+
return split if base is split.split else QueryAuthSplit(base, split.auth_result)
95+
if isinstance(split, IndexedSplit):
96+
base = _native_read_split(split.data_split(), ranges_use_row_ids)
97+
return split if base is split.data_split() else select_rows(base, split.row_ranges())
98+
if not isinstance(split, SlicedSplit):
99+
return split
100+
if ranges_use_row_ids is None:
101+
raise ValueError('Native file-slice conversion requires the reader row-tracking mode')
102+
ranges = []
103+
offset = 0
104+
slices = split.shard_file_idx_map()
105+
for file in split.files:
106+
start, end = slices.get(file.file_name, (0, file.row_count))
107+
start, end = max(0, start), min(file.row_count, end)
108+
if start < end:
109+
base = file.first_row_id if ranges_use_row_ids else offset
110+
if base is None:
111+
raise ValueError("Row-tracked file '%s' is missing first_row_id" % file.file_name)
112+
ranges.append(Range(base + start, base + end - 1))
113+
offset += file.row_count
114+
return select_rows(_native_read_split(split.data_split(), ranges_use_row_ids), ranges)
115+
116+
117+
def native_split_from_python(split, *, ranges_use_row_ids=None):
118+
"""Convert a Python-planned split for the Rust reader.
80119
81120
Vector scores intentionally stay on the Python IndexedSplit. Rust needs
82-
only its row ranges to perform the physical read.
121+
only its row ranges to perform the physical read. File slices need the
122+
reader's tracking mode to map their local offsets to those ranges.
83123
"""
84124
from pypaimon_rust.datafusion import Split as NativeSplit
85-
return NativeSplit.deserialize(
86-
serialize_split_v1(split, include_scores=False))
125+
return NativeSplit.deserialize(serialize_split_v1(
126+
_native_read_split(split, ranges_use_row_ids), include_scores=False))
87127

88128

89129
def _partition_fields(table):

‎paimon-python/pypaimon/read/reader/data_file_batch_reader.py‎

Lines changed: 1 addition & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -454,18 +454,7 @@ def _assign_row_tracking(self, record_batch: RecordBatch) -> RecordBatch:
454454
# NULL entries from the manifest's max sequence number.
455455
arrays[idx] = pc.fill_null(arrays[idx], self.max_sequence_number)
456456

457-
names = record_batch.schema.names
458-
table = None
459-
for i, name in enumerate(names):
460-
field = pa.field(
461-
name, arrays[i].type,
462-
nullable=record_batch.schema.field(name).nullable
463-
)
464-
if table is None:
465-
table = pa.table({name: arrays[i]}, schema=pa.schema([field]))
466-
else:
467-
table = table.append_column(field, arrays[i])
468-
return table.to_batches()[0]
457+
return pa.RecordBatch.from_arrays(arrays, schema=record_batch.schema)
469458

470459
def close(self) -> None:
471460
self.format_reader.close()

‎paimon-python/pypaimon/read/scanner/file_scanner.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -833,8 +833,8 @@ def _init_bucket_selector(self):
833833
# Only the default hash function (Math.abs(hash % numBuckets)) is
834834
# supported for bucket pruning. Non-default functions (mod, hive)
835835
# use different algorithms and would produce wrong bucket sets.
836-
bucket_func = self.table.table_schema.options.get('bucket-function.type', 'default')
837-
if bucket_func.lower() != 'default':
836+
bucket_func = self.table.table_schema.options.get('bucket-function.type')
837+
if bucket_func is not None and bucket_func.lower() != 'default':
838838
return None
839839
try:
840840
bucket_key_fields = self.table.table_schema.logical_bucket_key_fields

‎paimon-python/pypaimon/read/split_read.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1670,6 +1670,7 @@ def _create_union_reader(self, need_merge_files: List[DataFileMeta], deletion_ve
16701670
row_offsets = [-1] * len(all_read_fields)
16711671
field_offsets = [-1] * len(all_read_fields)
16721672
schema_pos = {f.id: p for p, f in enumerate(self.table.fields)}
1673+
partition_ids = [field.id for field in self.table.fields if field.name in self.table.partition_keys]
16731674

16741675
for i, bunch in enumerate(fields_files):
16751676
first_file = bunch.files()[0]
@@ -1693,6 +1694,12 @@ def _create_union_reader(self, need_merge_files: List[DataFileMeta], deletion_ve
16931694
field_ids.append(SpecialFields.ROW_ID.id)
16941695
field_ids.append(SpecialFields.SEQUENCE_NUMBER.id)
16951696

1697+
if isinstance(bunch, DataBunch):
1698+
# Partition values belong to the split, even when write_cols
1699+
# omits them from every file. Let the normal reader's existing
1700+
# PartitionInfo provide them before the merged row is filtered.
1701+
field_ids.extend(field_id for field_id in partition_ids if field_id not in field_ids)
1702+
16961703
read_fields = []
16971704
for j, read_field_id in enumerate(read_field_index):
16981705
for field_id in field_ids:

‎paimon-python/pypaimon/read/streaming_table_scan.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,6 @@
3232
from pypaimon.common.options.core_options import ChangelogProducer
3333
from pypaimon.common.predicate import Predicate
3434
from pypaimon.consumer.consumer import Consumer
35-
from pypaimon.consumer.consumer_manager import ConsumerManager
3635
from pypaimon.manifest.manifest_file_manager import ManifestFileManager
3736
from pypaimon.manifest.manifest_list_manager import ManifestListManager
3837
from pypaimon.read.native_plan import _raise_if_native_fork_safety_error
@@ -115,7 +114,7 @@ def __init__(
115114
self._query_auth_fn = self.table.catalog_environment.table_query_auth(
116115
self.table.options, self.table.identifier)
117116
self._consumer_manager = (
118-
ConsumerManager(table.file_io, table.table_path)
117+
table.consumer_manager()
119118
if consumer_id else None
120119
)
121120

‎paimon-python/pypaimon/read/table_read.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -625,7 +625,8 @@ def _try_native_batches(
625625
rust_split = getattr(split, '_native_split', None)
626626
if rust_split is None:
627627
try:
628-
rust_split = native_split_from_python(split)
628+
rust_split = native_split_from_python(split, ranges_use_row_ids=(
629+
self.table.options.row_tracking_enabled() or self.table.options.data_evolution_enabled()))
629630
except Exception as e:
630631
return self._native_fallback(
631632
"native split conversion failed: %s" % e,

‎paimon-python/pypaimon/table/data_evolution_merge_into.py‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,10 @@ def merge_into(
109109
# unrestricted whole-table grant before any user callback or staging.
110110
from pypaimon.write.table_update import _authorize_update
111111
_authorize_update(target_table)
112+
if any(clause.condition is not None
113+
for clauses in (when_matched, when_not_matched) for clause in clauses):
114+
from pypaimon.ray.merge_condition import _load_datafusion
115+
_load_datafusion()
112116
target_on, source_on = _normalize_on(on)
113117
self_merge = _is_self_merge(target_table, source, target_on, source_on)
114118
from pypaimon.snapshot.time_travel_util import TimeTravelUtil
@@ -222,12 +226,7 @@ def _prepare(
222226
for c in list(when_matched) + list(when_not_matched)
223227
)
224228
if has_condition:
225-
from pypaimon.ray.merge_condition import (
226-
_load_datafusion,
227-
extract_target_columns,
228-
)
229-
230-
_load_datafusion()
229+
from pypaimon.ray.merge_condition import extract_target_columns
231230
for c in when_not_matched:
232231
if c.condition is not None:
233232
t_refs = extract_target_columns(c.condition)

‎paimon-python/pypaimon/table/row/generic_row.py‎

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -76,15 +76,13 @@ def _parse_type_precision_scale(data_type):
7676
def _datetime_to_millis_and_nanos(value: datetime):
7777
"""Convert datetime to (epoch_millis, nano_of_millisecond) without float arithmetic.
7878
79-
Python's datetime is microsecond-resolution, so nano_of_millisecond is
80-
always a multiple of 1000 and sub-microsecond precision is lost. For
81-
TIMESTAMP(7..9) columns the lower three nano digits will always serialise
82-
as zero — round-trip through pypaimon is microsecond-faithful, not
83-
nanosecond-faithful.
79+
Plain datetime values have microsecond resolution. Arrow returns pandas
80+
Timestamp values for nanosecond arrays, carrying the remaining three
81+
digits separately in nanosecond.
8482
"""
8583
epoch_seconds = calendar.timegm(value.timetuple())
8684
millis = epoch_seconds * 1000 + value.microsecond // 1000
87-
nano_of_millisecond = (value.microsecond % 1000) * 1000
85+
nano_of_millisecond = (value.microsecond % 1000) * 1000 + getattr(value, 'nanosecond', 0)
8886
return millis, nano_of_millisecond
8987

9088

@@ -100,14 +98,21 @@ def _is_ltz_type(type_name: str) -> bool:
10098

10199

102100
def _millis_nanos_to_datetime(millis: int, nano_of_millisecond: int = 0) -> datetime:
103-
"""Convert (epoch_millis, nano_of_millisecond) to datetime. Nanos truncated to micros."""
101+
"""Retain nanosecond tails while preserving ordinary datetime results."""
104102
total_micros = millis * 1000 + nano_of_millisecond // 1000
105103
seconds = total_micros // 1_000_000
106104
micros = total_micros % 1_000_000
107105
if micros < 0:
108106
seconds -= 1
109107
micros += 1_000_000
110-
return _EPOCH + timedelta(seconds=seconds, microseconds=micros)
108+
result = _EPOCH + timedelta(seconds=seconds, microseconds=micros)
109+
tail = nano_of_millisecond % 1000
110+
if tail:
111+
import pandas as pd
112+
# Construct the final instant directly: the floored microsecond before
113+
# Timestamp.min can be out of range even when these nanoseconds fit.
114+
return pd.Timestamp(millis * 1000000 + nano_of_millisecond, unit='ns')
115+
return result
111116

112117

113118
@dataclass

‎paimon-python/pypaimon/tests/commit_update_schema_identity_test.py‎

Lines changed: 48 additions & 0 deletions

0 commit comments

Comments
 (0)