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 4b275a4ef1 [python] Expand native Rust scan planning (#9000)
4b275a4ef1 is described below
commit 4b275a4ef1a6aae03d40f9638e293325f657181b
Author: XiaoHongbo <[email protected]>
AuthorDate: Tue Aug 4 11:38:29 2026 +0800
[python] Expand native Rust scan planning (#9000)
---
paimon-python/pypaimon/read/native_plan.py | 93 ++++++++++--
paimon-python/pypaimon/read/table_scan.py | 79 ++++++----
.../pypaimon/tests/native_plan_integration_test.py | 143 +++++++++++++++++-
paimon-python/pypaimon/tests/native_plan_test.py | 161 +++++++++++++++++++--
4 files changed, 423 insertions(+), 53 deletions(-)
diff --git a/paimon-python/pypaimon/read/native_plan.py
b/paimon-python/pypaimon/read/native_plan.py
index 91f9827940..54c87e3594 100644
--- a/paimon-python/pypaimon/read/native_plan.py
+++ b/paimon-python/pypaimon/read/native_plan.py
@@ -18,8 +18,8 @@
"""Plan splits with pypaimon_rust, decoded for the normal pypaimon reader.
Optional, lazily-imported dependency; enabled by ``scan.native-plan.enabled``.
-The predicate is applied pypaimon-side (partition pruning + row/limit filter),
-so results match the normal path.
+Predicates and limits are pushed into Rust planning. The normal pypaimon reader
+still applies them while reading, so pushdown remains an optimization.
"""
from typing import List, Optional
@@ -27,6 +27,7 @@ from typing import List, Optional
from pypaimon.common.options.config import CatalogOptions
from pypaimon.common.options.core_options import CoreOptions
from pypaimon.common.options.options_utils import OptionsUtils
+from pypaimon.common.predicate import Predicate
from pypaimon.read.split import Split
from pypaimon.read.split_serializer import deserialize_split_v1
@@ -96,20 +97,85 @@ def _catalog_options(table) -> dict:
def _read_options(table) -> dict:
- """Effective split-shaping options, including FileStoreTable.copy
overrides."""
- return {
+ """Effective Rust read options, including FileStoreTable.copy overrides."""
+ options = {
CoreOptions.SOURCE_SPLIT_TARGET_SIZE.key(): str(
table.options.source_split_target_size()),
CoreOptions.SOURCE_SPLIT_OPEN_FILE_COST.key(): str(
table.options.source_split_open_file_cost()),
}
+ table_options = table.options.options
+ for option in (
+ CoreOptions.SCAN_SNAPSHOT_ID,
+ CoreOptions.SCAN_TAG_NAME,
+ CoreOptions.SCAN_TIMESTAMP_MILLIS):
+ if table_options.contains_key(option.key()):
+ options[option.key()] = _option_value_to_string(
+ table_options.get(option))
+
+ # Rust takes epoch millis but PyPaimon also accepts a timestamp string.
+ if table_options.contains_key(CoreOptions.SCAN_TIMESTAMP.key()):
+ from pypaimon.snapshot.time_travel_util import
_parse_timestamp_to_millis
+ options[CoreOptions.SCAN_TIMESTAMP_MILLIS.key()] = str(
+ _parse_timestamp_to_millis(
+ table_options.get(CoreOptions.SCAN_TIMESTAMP)))
+ return options
+
+
+def _predicate_to_native(predicate: Predicate) -> dict:
+ """Convert PyPaimon's predicate tree to pypaimon-rust's dict API."""
+ if predicate.method in ('and', 'or'):
+ children = predicate.literals or []
+ if not children:
+ raise ValueError("Native compound predicate requires children")
+ return {
+ 'method': predicate.method,
+ 'children': [_predicate_to_native(child) for child in children],
+ }
+ return {
+ 'method': predicate.method,
+ 'field': predicate.field,
+ 'literals': list(predicate.literals or []),
+ }
-def native_plan(table) -> List[Split]:
+def _restore_python_partition_paths(table, splits: List[Split]) -> None:
+ """Restore legacy PyPaimon paths with one listing per bucket."""
+ if not table.partition_keys:
+ return
+ path_factory = table.path_factory()
+ bucket_files = {}
+ for split in splits:
+ bucket_path = path_factory.bucket_path(
+ tuple(split.partition.values), split.bucket)
+ candidates = []
+ for data_file in split.files:
+ python_path = "%s/%s" % (
+ bucket_path.rstrip('/'), data_file.file_name)
+ if (not data_file.external_path
+ and python_path != data_file.file_path):
+ candidates.append((data_file, python_path))
+ if not candidates:
+ continue
+ if bucket_path not in bucket_files:
+ bucket_files[bucket_path] = {
+ status.base_name
+ for status in table.file_io.list_status(bucket_path)
+ }
+ for data_file, python_path in candidates:
+ if data_file.file_name in bucket_files[bucket_path]:
+ data_file.file_path = python_path
+
+
+def native_plan(
+ table,
+ predicate: Optional[Predicate] = None,
+ limit: Optional[int] = None,
+ projection: Optional[List[str]] = None) -> List[Split]:
"""Plan with pypaimon_rust and return the decoded pypaimon splits.
- Predicate/limit are not pushed to the native planner (pushdown is a
- follow-up); pypaimon applies them at read time.
+ Native conversion or planning failures are handled by TableScan, which
+ falls back to the Python planner.
"""
if not native_runtime_available():
raise RuntimeError(
@@ -117,8 +183,17 @@ def native_plan(table) -> List[Split]:
from pypaimon_rust.datafusion import PaimonCatalog
rt =
PaimonCatalog(_catalog_options(table)).get_table(table.identifier.get_full_name())
- rust_splits =
rt.new_read_builder(_read_options(table)).new_scan().plan().splits()
+ builder = rt.new_read_builder(_read_options(table))
+ if projection is not None:
+ builder = builder.with_projection(projection)
+ if predicate is not None:
+ builder = builder.with_filter(_predicate_to_native(predicate))
+ if limit is not None:
+ builder = builder.with_limit(limit)
+ rust_splits = builder.new_scan().plan().splits()
pfields = _partition_fields(table)
# Trimmed primary keys decode per-file min/max keys (PK merge-on-read).
kfields = table.trimmed_primary_keys_fields
- return [deserialize_split_v1(s.serialize(), pfields, kfields) for s in
rust_splits]
+ splits = [deserialize_split_v1(s.serialize(), pfields, kfields) for s in
rust_splits]
+ _restore_python_partition_paths(table, splits)
+ return splits
diff --git a/paimon-python/pypaimon/read/table_scan.py
b/paimon-python/pypaimon/read/table_scan.py
index 313046aecb..410cae9a4e 100755
--- a/paimon-python/pypaimon/read/table_scan.py
+++ b/paimon-python/pypaimon/read/table_scan.py
@@ -37,6 +37,16 @@ _NATIVE_FORWARDED_OPTIONS = frozenset({
CoreOptions.SCAN_NATIVE_PLAN_ENABLED.key(),
CoreOptions.SOURCE_SPLIT_TARGET_SIZE.key(),
CoreOptions.SOURCE_SPLIT_OPEN_FILE_COST.key(),
+ CoreOptions.SCAN_SNAPSHOT_ID.key(),
+ CoreOptions.SCAN_TAG_NAME.key(),
+ CoreOptions.SCAN_TIMESTAMP.key(),
+ CoreOptions.SCAN_TIMESTAMP_MILLIS.key(),
+})
+_NATIVE_TIME_TRAVEL_OPTIONS = frozenset({
+ CoreOptions.SCAN_SNAPSHOT_ID.key(),
+ CoreOptions.SCAN_TAG_NAME.key(),
+ CoreOptions.SCAN_TIMESTAMP.key(),
+ CoreOptions.SCAN_TIMESTAMP_MILLIS.key(),
})
@@ -86,31 +96,24 @@ class TableScan:
def _native_plan_supported_impl(self) -> bool:
"""Fall back to the Python scanner for scans native can't carry:
shard/slice, chunk-shuffle, global-index, first-row merge-engine (Rust
- drops L0), deletion vectors (Python drops L0), data evolution
- (dedicated split generator), postpone bucket (drops synthetic buckets),
+ drops L0), deletion vectors, postpone bucket (drops synthetic buckets),
a primary-key table whose trimmed PK is empty (PK equals the partition
key; native may mark splits raw-convertible and skip merge), dynamic
- bucket / cross-partition PK tables (unconfirmed Rust parity), any
- partitioned table (Rust bucket_path vs the writer's str(value) can
- diverge), a stale schema (Rust reloads the latest), copy() overrides
- Rust does not see (e.g. a removed scan.snapshot-id), a row limit
- (native has no plan-time limit pushdown), query auth, non-main branch,
- time-travel, scan.version, incremental, a missing/old pypaimon-rust, or
- a catalog / identifier Rust cannot reconstruct. Keep this capability
- gate in sync when adding scan features."""
+ bucket / cross-partition PK tables (unconfirmed Rust parity), a stale
+ schema without time travel, copy() overrides Rust does not see (notably
+ removing a persisted scan option), unsupported time travel selectors,
+ query auth, non-main branch, incremental scans, a missing/old
+ pypaimon-rust, or a catalog / identifier Rust cannot reconstruct. Keep
+ this capability gate in sync when adding scan features."""
from pypaimon.read.native_plan import native_runtime_available
if not native_runtime_available():
return False
- # Native has no plan-time limit pushdown; Python trims splits before
read.
- if self.limit is not None:
- return False
fs = self.file_scanner
if (getattr(fs, 'idx_of_this_subtask', None) is not None
or getattr(fs, 'start_pos_of_this_subtask', None) is not None
or getattr(fs, 'chunk_shuffle', None) is not None
or getattr(fs, '_global_index_result', None) is not None
or getattr(fs, 'deletion_vectors_enabled', False)
- or getattr(fs, 'data_evolution', False)
or getattr(fs, 'only_read_real_buckets', False)):
return False
loader = getattr(
@@ -146,20 +149,24 @@ class TableScan:
from pypaimon.table.bucket_mode import BucketMode
if self.table.bucket_mode() in (BucketMode.HASH_DYNAMIC,
BucketMode.CROSS_PARTITION):
return False
- # Rust bucket_path vs the writer's unescaped str(value) can diverge ->
fall back.
- if self.table.partition_keys:
- return False
- # Rust reloads the latest schema; fall back if this table's schema is
stale.
+ options = self.table.options.options
+ supported_time_travel = any(
+ options.contains_key(key) for key in _NATIVE_TIME_TRAVEL_OPTIONS)
+ # Time travel intentionally carries a historical schema; other stale
+ # table objects must still fall back because Rust reloads the latest.
latest_schema = self.table.schema_manager.latest()
- if latest_schema is not None and latest_schema.id !=
self.table.table_schema.id:
+ if (not supported_time_travel and latest_schema is not None
+ and latest_schema.id != self.table.table_schema.id):
return False
- # copy() overrides Rust can't see (e.g. removed scan.snapshot-id) ->
fall back.
- overrides = set(getattr(self.table, '_applied_dynamic_options', {}) or
{})
- if overrides - _NATIVE_FORWARDED_OPTIONS:
+ # Rust cannot remove an option persisted in the catalog-loaded schema.
+ applied_options = getattr(self.table, '_applied_dynamic_options', {})
or {}
+ if (set(applied_options) - _NATIVE_FORWARDED_OPTIONS
+ or any(key in _NATIVE_TIME_TRAVEL_OPTIONS and value is None
+ for key, value in applied_options.items())):
return False
from pypaimon.snapshot.time_travel_util import SCAN_KEYS
- options = self.table.options.options
- if any(options.contains_key(k) for k in SCAN_KEYS) \
+ unsupported_scan_keys = set(SCAN_KEYS) - _NATIVE_TIME_TRAVEL_OPTIONS
+ if any(options.contains_key(k) for k in unsupported_scan_keys) \
or options.contains_key('scan.version'):
return False
return not options.contains(CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP)
@@ -167,15 +174,29 @@ class TableScan:
def _try_native_plan(self) -> Optional[Plan]:
"""Plan via pypaimon_rust, then drop partitions the predicate rejects.
- The predicate is not pushed to the native planner, so this may read
more
- files; the reader's row filter/limit still apply, so results match.
Return
- None when Rust finds no splits so the caller can use the matching
Python
- fallback (with scan stats when requested).
+ Predicate and limit are pushed into Rust planning and are still
enforced
+ by the reader. Return None when Rust finds no splits so the caller can
use
+ the matching Python fallback (with scan stats when requested).
"""
from pypaimon.read.native_plan import native_plan
try:
- splits = native_plan(self.table)
+ native_predicate = self.predicate
+ if self.partition_predicate is not None:
+ native_predicate = PredicateBuilder.and_predicates([
+ predicate for predicate in (
+ native_predicate,
+ self.file_scanner.partition_key_predicate,
+ ) if predicate is not None
+ ])
+ splits = native_plan(
+ self.table,
+ predicate=native_predicate,
+ limit=self.limit,
+ projection=(
+ [field.name for field in self._read_type]
+ if self._read_type is not None else None),
+ )
if not splits:
return None
snapshot_id = splits[0].snapshot_id
diff --git a/paimon-python/pypaimon/tests/native_plan_integration_test.py
b/paimon-python/pypaimon/tests/native_plan_integration_test.py
index a78e1b55a0..e308f3a611 100644
--- a/paimon-python/pypaimon/tests/native_plan_integration_test.py
+++ b/paimon-python/pypaimon/tests/native_plan_integration_test.py
@@ -120,6 +120,138 @@ class NativePlanIntegrationTest(unittest.TestCase):
self._write('ap_t', [{'k': 3, 'v': 'c'}])
self._assert_matches('ap_t')
+ def test_data_evolution_blob_projection_filter_limit(self):
+ schema = pa.schema([
+ ('k', pa.int64()),
+ ('v', pa.string()),
+ ('media.camera', pa.large_binary()),
+ ])
+ self.cat.create_table('default.de_t', Schema.from_pyarrow_schema(
+ schema, options={
+ 'row-tracking.enabled': 'true',
+ 'data-evolution.enabled': 'true',
+ }), False)
+ table = self.cat.get_table('default.de_t')
+ write_builder = table.new_batch_write_builder()
+ write = write_builder.new_write()
+ write.write_arrow(pa.Table.from_pylist([
+ {'k': 1, 'v': 'a', 'media.camera': b'a'},
+ {'k': 2, 'v': 'b', 'media.camera': b'b'},
+ {'k': 3, 'v': 'c', 'media.camera': b'c'},
+ ], schema=schema))
+ write_builder.new_commit().commit(write.prepare_commit())
+ write.close()
+
+ update_builder = table.new_batch_write_builder()
+ update = update_builder.new_update().with_update_type(['v'])
+ messages = update.update_by_arrow_with_row_id(pa.Table.from_pydict({
+ '_ROW_ID': pa.array([1], type=pa.int64()),
+ 'v': pa.array(['b2'], type=pa.string()),
+ }))
+ update_builder.new_commit().commit(messages)
+
+ self._assert_matches('de_t')
+
+ native_table = self.cat.get_table('default.de_t').copy(
+ {'scan.native-plan.enabled': 'true'})
+ predicate =
native_table.new_read_builder().new_predicate_builder().equal(
+ 'v', 'b2')
+ builder = (native_table.new_read_builder()
+ .with_projection(['k'])
+ .with_filter(predicate)
+ .with_limit(1))
+ plan = builder.new_scan().plan()
+ rows = builder.new_read().to_arrow(plan.splits()).to_pylist()
+
+ self.assertEqual(rows, [{'k': 2}])
+ self.assertTrue(builder.explain().native_planned)
+
+ blob_builder = (native_table.new_read_builder()
+ .with_projection(['media.camera'])
+ .with_limit(1))
+ blob_plan = blob_builder.new_scan().plan()
+ blob_rows = blob_builder.new_read().to_arrow(
+ blob_plan.splits()).to_pylist()
+ self.assertEqual(blob_rows, [{'media.camera': b'a'}])
+ self.assertTrue(any(
+ data_file.file_name.endswith('.blob')
+ for split in blob_plan.splits()
+ for data_file in split.files
+ ))
+
+ def test_filter_is_pushed_to_native_plan(self):
+ options = {
+ 'source.split.target-size': '1b',
+ 'source.split.open-file-cost': '1b',
+ }
+ self.cat.create_table('default.filter_t', Schema.from_pyarrow_schema(
+ self.schema, options=options), False)
+ for k in range(1, 4):
+ self._write('filter_t', [{'k': k, 'v': 'v%d' % k}])
+
+ table = self.cat.get_table('default.filter_t')
+ normal_builder = table.new_read_builder()
+ predicate = normal_builder.new_predicate_builder().equal('k', 2)
+ normal_builder.with_filter(predicate)
+ normal_plan = normal_builder.new_scan().plan()
+
+ native_builder = table.copy(
+ {'scan.native-plan.enabled': 'true'}).new_read_builder()
+ predicate = native_builder.new_predicate_builder().equal('k', 2)
+ native_builder.with_filter(predicate)
+ native_plan = native_builder.new_scan().plan()
+ rows =
native_builder.new_read().to_arrow(native_plan.splits()).to_pylist()
+
+ self.assertEqual(rows, [{'k': 2, 'v': 'v2'}])
+ self.assertEqual(len(native_plan.splits()), len(normal_plan.splits()))
+ self.assertTrue(native_builder.explain().native_planned)
+
+ def test_limit_is_pushed_to_native_plan(self):
+ options = {
+ 'source.split.target-size': '1b',
+ 'source.split.open-file-cost': '1b',
+ }
+ self.cat.create_table('default.limit_t', Schema.from_pyarrow_schema(
+ self.schema, options=options), False)
+ for k in range(1, 4):
+ self._write('limit_t', [{'k': k, 'v': 'v%d' % k}])
+
+ table = self.cat.get_table('default.limit_t')
+ normal = table.new_read_builder().with_limit(1).new_scan().plan()
+ native_builder = table.copy(
+ {'scan.native-plan.enabled':
'true'}).new_read_builder().with_limit(1)
+ native = native_builder.new_scan().plan()
+ rows = native_builder.new_read().to_arrow(native.splits()).to_pylist()
+
+ self.assertEqual(len(rows), 1)
+ self.assertEqual(len(native.splits()), len(normal.splits()))
+ self.assertEqual(len(native.splits()), 1)
+ self.assertTrue(native_builder.explain().native_planned)
+
+ def test_snapshot_time_travel_matches_normal_plan(self):
+ self.cat.create_table(
+ 'default.travel_t', Schema.from_pyarrow_schema(self.schema), False)
+ self._write('travel_t', [{'k': 1, 'v': 'a'}])
+ self._write('travel_t', [{'k': 2, 'v': 'b'}])
+ options = {'scan.snapshot-id': '1'}
+
+ normal_table = self.cat.get_table('default.travel_t').copy(options)
+ normal_builder = normal_table.new_read_builder()
+ normal_plan = normal_builder.new_scan().plan()
+ normal_rows = normal_builder.new_read().to_arrow(
+ normal_plan.splits()).to_pylist()
+
+ native_table = normal_table.copy({'scan.native-plan.enabled': 'true'})
+ native_builder = native_table.new_read_builder()
+ native_plan = native_builder.new_scan().plan()
+ native_rows = native_builder.new_read().to_arrow(
+ native_plan.splits()).to_pylist()
+
+ self.assertEqual(native_plan.snapshot_id, 1)
+ self.assertEqual(native_rows, normal_rows)
+ self.assertEqual(native_rows, [{'k': 1, 'v': 'a'}])
+ self.assertTrue(native_builder.explain().native_planned)
+
def test_dynamic_split_target_size_matches_normal_plan(self):
self.cat.create_table(
'default.split_t', Schema.from_pyarrow_schema(self.schema), False)
@@ -182,8 +314,8 @@ class NativePlanIntegrationTest(unittest.TestCase):
self.assertEqual(native.split_count, len(normal.splits()))
self.assertEqual(native.split_count, 1)
- def test_partitioned_table_falls_back(self):
- # Rust bucket_path can diverge from the writer's str(value) partition
dir -> fall back.
+ def test_partitioned_table_matches_normal_plan(self):
+ # Native decoding restores PyPaimon's legacy unescaped partition path.
schema = pa.schema([('k', pa.int64()), ('p', pa.string())])
self.cat.create_table('default.pt_t', Schema.from_pyarrow_schema(
schema, partition_keys=['p']), False)
@@ -191,14 +323,13 @@ class NativePlanIntegrationTest(unittest.TestCase):
wb = t.new_batch_write_builder()
w, c = wb.new_write(), wb.new_commit()
w.write_arrow(pa.Table.from_pylist(
- [{'k': 1, 'p': 'a'}, {'k': 2, 'p': 'a'}, {'k': 3, 'p': 'b'}],
schema=schema))
+ [{'k': 1, 'p': 'a/b'}, {'k': 2, 'p': 'a/b'}, {'k': 3, 'p': 'c'}],
+ schema=schema))
c.commit(w.prepare_commit())
w.close()
c.close()
- native_table = self.cat.get_table('default.pt_t').copy(
- {'scan.native-plan.enabled': 'true'})
-
self.assertFalse(native_table.new_read_builder().explain().native_planned)
+ self._assert_matches('pt_t')
def test_explain_reflects_native_plan(self):
self.cat.create_table(
diff --git a/paimon-python/pypaimon/tests/native_plan_test.py
b/paimon-python/pypaimon/tests/native_plan_test.py
index b21ac808a6..b590bb4f90 100644
--- a/paimon-python/pypaimon/tests/native_plan_test.py
+++ b/paimon-python/pypaimon/tests/native_plan_test.py
@@ -26,7 +26,15 @@ from pypaimon.catalog.jdbc_catalog_loader import
JdbcCatalogLoader
from pypaimon.catalog.rest.rest_catalog_loader import RESTCatalogLoader
from pypaimon.common.options.core_options import CoreOptions
from pypaimon.common.options.options import Options
-from pypaimon.read.native_plan import _catalog_options, native_plan
+from pypaimon.common.predicate import Predicate
+from pypaimon.common.predicate_builder import PredicateBuilder
+from pypaimon.read.native_plan import (
+ _catalog_options,
+ _predicate_to_native,
+ _read_options,
+ _restore_python_partition_paths,
+ native_plan,
+)
from pypaimon.read.scan_stats import ScanStats
from pypaimon.read.table_scan import TableScan
from pypaimon.table.bucket_mode import BucketMode
@@ -60,6 +68,8 @@ def _scan(native_enabled, file_scanner):
file_scanner.data_evolution = False # no data evolution
file_scanner.only_read_real_buckets = False # not postpone bucket
scan.file_scanner = file_scanner
+ scan.predicate = None
+ scan.partition_predicate = None
scan._query_auth_fn = None # no query-auth restrictions
scan._read_type = None
scan.limit = None # no row limit
@@ -109,7 +119,8 @@ class NativePlanTest(unittest.TestCase):
with patch('pypaimon.read.native_plan.native_plan',
return_value=[keep, drop]) as np:
plan = scan.plan()
- np.assert_called_once_with(scan.table)
+ np.assert_called_once_with(
+ scan.table, predicate=None, limit=None, projection=None)
fs.scan.assert_not_called()
self.assertEqual(plan.splits(), [keep])
@@ -132,9 +143,39 @@ class NativePlanTest(unittest.TestCase):
with patch('pypaimon.read.native_plan.native_plan',
return_value=splits):
self.assertEqual(scan.plan().splits(), splits)
+ def test_plan_forwards_filter_limit_partition_and_time_travel(self):
+ fs = Mock(partition_key_predicate=None)
+ scan = _scan(native_enabled=True, file_scanner=fs)
+ scan.table.partition_keys = ['dt']
+ scan.limit = 5
+ scan._read_type = [Mock(name='k'), Mock(name='dt')]
+ scan._read_type[0].name = 'k'
+ scan._read_type[1].name = 'dt'
+ scan.predicate = PredicateBuilder.and_predicates([
+ Predicate('equal', 0, 'k', [7]),
+ Predicate('equal', 1, 'dt', ['2026-08-02']),
+ ])
+ scan.table._applied_dynamic_options = {'scan.snapshot-id': '3'}
+ scan.table.options.options.contains_key.side_effect = (
+ lambda key: key == 'scan.snapshot-id')
+ scan.table.table_schema.id = 2
+ scan.table.schema_manager.latest.return_value.id = 3
+ split = Mock(partition=Mock(values=['2026-08-02']), snapshot_id=3)
+
+ with patch('pypaimon.read.native_plan.native_plan',
return_value=[split]) as np:
+ plan = scan.plan()
+
+ self.assertEqual(plan.snapshot_id, 3)
+ np.assert_called_once_with(
+ scan.table,
+ predicate=scan.predicate,
+ limit=5,
+ projection=['k', 'dt'],
+ )
+
def test_plan_falls_back_when_scan_is_not_plain(self):
# Native planning does not carry shard/slice, global-index, or
- # time-travel/incremental scans -> must fall back to the file scanner.
+ # incremental scans -> must fall back to the file scanner.
def check(setup):
fs = Mock(partition_key_predicate=None)
sentinel = object()
@@ -151,7 +192,6 @@ class NativePlanTest(unittest.TestCase):
check(lambda s, fs: setattr(fs, 'chunk_shuffle', (1, 100)))
check(lambda s, fs: setattr(fs, '_global_index_result', object()))
check(lambda s, fs: setattr(fs, 'deletion_vectors_enabled', True))
- check(lambda s, fs: setattr(fs, 'data_evolution', True))
check(lambda s, fs: setattr(fs, 'only_read_real_buckets', True))
check(lambda s, fs: (setattr(s.table, 'is_primary_key_table', True),
setattr(s.table, 'trimmed_primary_keys', [])))
@@ -161,11 +201,9 @@ class NativePlanTest(unittest.TestCase):
'return_value', BucketMode.CROSS_PARTITION))
check(lambda s, fs: setattr(
s.table, '_applied_dynamic_options', {'scan.snapshot-id': None}))
- check(lambda s, fs: setattr(s.table, 'partition_keys', ['dt']))
check(lambda s, fs:
setattr(s.table.schema_manager.latest.return_value, 'id', 2))
check(lambda s, fs: s.table.schema_manager.latest.__setattr__(
'side_effect', RuntimeError('metadata read failed')))
- check(lambda s, fs: setattr(s, 'limit', 5))
check(lambda s, fs: s.table.options.options.contains_key.__setattr__(
'side_effect', lambda k: k == 'scan.version'))
check(lambda s, fs: s.table.options.merge_engine.__setattr__(
@@ -181,8 +219,6 @@ class NativePlanTest(unittest.TestCase):
for attr in ('hadoop_conf', 'prefer_io_loader', 'fallback_io_loader'):
check(lambda s, fs, attr=attr: setattr(
s.table.catalog_environment.catalog_loader.context(), attr,
object()))
- check(lambda s, fs: s.table.options.options.contains_key.__setattr__(
- 'return_value', True)) # time-travel
check(lambda s, fs: s.table.options.options.contains.__setattr__(
'return_value', True)) # incremental
@@ -271,7 +307,8 @@ class NativePlanTest(unittest.TestCase):
self.assertIs(plan, fallback_plan)
self.assertIs(stats, fallback_stats)
- np.assert_called_once_with(scan.table)
+ np.assert_called_once_with(
+ scan.table, predicate=None, limit=None, projection=None)
fs.scan_with_stats.assert_called_once_with()
fs.scan.assert_not_called()
@@ -316,14 +353,120 @@ class NativePlanTest(unittest.TestCase):
with self.assertRaisesRegex(ValueError, 'exact built-in catalog
loader'):
_catalog_options(table)
+ def test_predicate_and_time_travel_are_converted_for_rust(self):
+ predicate = PredicateBuilder.and_predicates([
+ Predicate('greaterOrEqual', 0, 'k', [10]),
+ Predicate('in', 1, 'v', ['a', 'b']),
+ ])
+ self.assertEqual(_predicate_to_native(predicate), {
+ 'method': 'and',
+ 'children': [
+ {'method': 'greaterOrEqual', 'field': 'k', 'literals': [10]},
+ {'method': 'in', 'field': 'v', 'literals': ['a', 'b']},
+ ],
+ })
+
+ table = Mock()
+ table.options.source_split_target_size.return_value = 1024
+ table.options.source_split_open_file_cost.return_value = 128
+ table.options.options = Options({
+ 'scan.snapshot-id': '9',
+ })
+ self.assertEqual(_read_options(table), {
+ 'source.split.target-size': '1024',
+ 'source.split.open-file-cost': '128',
+ 'scan.snapshot-id': '9',
+ })
+
+ def test_partition_path_prefers_existing_python_legacy_path(self):
+ table = Mock(partition_keys=['p'])
+ table.path_factory.return_value.bucket_path.return_value = (
+ '/warehouse/t/p=a/b/bucket-0')
+ table.file_io.list_status.return_value =
[Mock(base_name='data.parquet')]
+ data_file = Mock(
+ external_path=None,
+ file_name='data.parquet',
+ file_path='/warehouse/t/p=a%2Fb/bucket-0/data.parquet',
+ )
+ split = Mock(
+ partition=Mock(values=['a/b']), bucket=0, files=[data_file])
+
+ _restore_python_partition_paths(table, [split])
+
+ self.assertEqual(
+ data_file.file_path,
+ '/warehouse/t/p=a/b/bucket-0/data.parquet',
+ )
+
+ def test_partition_path_keeps_existing_rust_path(self):
+ table = Mock(partition_keys=['p'])
+ table.path_factory.return_value.bucket_path.return_value = (
+ '/warehouse/t/p=a/b/bucket-0')
+ table.file_io.list_status.return_value = []
+ rust_path = '/warehouse/t/p=a%2Fb/bucket-0/data.parquet'
+ data_file = Mock(
+ external_path=None,
+ file_name='data.parquet',
+ file_path=rust_path,
+ )
+ split = Mock(
+ partition=Mock(values=['a/b']), bucket=0, files=[data_file])
+
+ _restore_python_partition_paths(table, [split])
+
+ self.assertEqual(data_file.file_path, rust_path)
+
+ def test_partition_path_lists_each_bucket_once(self):
+ table = Mock(partition_keys=['p'])
+ table.path_factory.return_value.bucket_path.return_value = (
+ '/warehouse/t/p=a/b/bucket-0')
+ table.file_io.list_status.return_value = [
+ Mock(base_name='a.parquet'), Mock(base_name='b.parquet')]
+ splits = [
+ Mock(partition=Mock(values=['a/b']), bucket=0, files=[Mock(
+ external_path=None,
+ file_name=name,
+ file_path='/warehouse/t/p=a%%2Fb/bucket-0/%s' % name,
+ )])
+ for name in ('a.parquet', 'b.parquet')
+ ]
+
+ _restore_python_partition_paths(table, splits)
+
+ table.file_io.list_status.assert_called_once_with(
+ '/warehouse/t/p=a/b/bucket-0')
+ self.assertEqual(
+ [split.files[0].file_path for split in splits],
+ [
+ '/warehouse/t/p=a/b/bucket-0/a.parquet',
+ '/warehouse/t/p=a/b/bucket-0/b.parquet',
+ ],
+ )
+
+ def test_partition_path_listing_failure_is_not_hidden(self):
+ table = Mock(partition_keys=['p'])
+ table.path_factory.return_value.bucket_path.return_value = (
+ '/warehouse/t/p=a/b/bucket-0')
+ table.file_io.list_status.side_effect = PermissionError('denied')
+ split = Mock(partition=Mock(values=['a/b']), bucket=0, files=[Mock(
+ external_path=None,
+ file_name='data.parquet',
+ file_path='/warehouse/t/p=a%2Fb/bucket-0/data.parquet',
+ )])
+
+ with self.assertRaises(PermissionError):
+ _restore_python_partition_paths(table, [split])
+
def test_native_plan_threads_trimmed_keys_to_deserializer(self):
# PK tables route through: the trimmed primary keys must reach the
# deserializer so per-file min/max keys are decoded for merge-on-read.
kfields = [object()]
table = Mock(trimmed_primary_keys_fields=kfields)
table.table_schema = Mock(fields=[], partition_keys=[])
+ table.partition_keys = []
table.options.source_split_target_size.return_value = 1024
table.options.source_split_open_file_cost.return_value = 128
+ table.options.options.contains_key.return_value = False
split = Mock()
split.serialize.return_value = b'bytes'
rt = Mock()
