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 d9599117a7 [python] Add Mosaic row-group stats skipping (#8419)
d9599117a7 is described below
commit d9599117a7f329c4f765ea7c1976e4e321cc3ed0
Author: QuakeWang <[email protected]>
AuthorDate: Fri Jul 3 13:12:43 2026 +0800
[python] Add Mosaic row-group stats skipping (#8419)
Python Mosaic reads currently apply pushed predicates only after
`read_row_group`, by filtering the in-memory Arrow batch. This keeps
results correct but misses Mosaic row-group statistics pruning.
This PR passes the structured Paimon predicate to the Mosaic format
reader and evaluates it against Mosaic row-group statistics before
reading a row group. The final Arrow predicate filter is still applied
after reading, so stats pruning remains an optimization. If stats are
missing, conversion fails, or the file schema is older than the table
schema, the reader fails open and keeps the previous behavior.
---
.../pypaimon/read/reader/format_mosaic_reader.py | 143 ++++++++++++++++++++-
paimon-python/pypaimon/read/split_read.py | 19 ++-
.../pypaimon/tests/format_mosaic_reader_test.py | 55 ++++++++
.../tests/test_format_mosaic_reader_writer.py | 140 +++++++++++++++++++-
4 files changed, 345 insertions(+), 12 deletions(-)
diff --git a/paimon-python/pypaimon/read/reader/format_mosaic_reader.py
b/paimon-python/pypaimon/read/reader/format_mosaic_reader.py
index 4ab7d3b698..bfe592f456 100644
--- a/paimon-python/pypaimon/read/reader/format_mosaic_reader.py
+++ b/paimon-python/pypaimon/read/reader/format_mosaic_reader.py
@@ -15,6 +15,10 @@
# specific language governing permissions and limitations
# under the License.
+import re
+import struct
+from datetime import datetime, timedelta
+from decimal import Decimal
from typing import Any, List, Optional
import pyarrow as pa
@@ -22,17 +26,23 @@ import pyarrow.dataset as ds
from pyarrow import RecordBatch
from pypaimon.common.file_io import FileIO, supports_pread, pread
+from pypaimon.common.predicate import Predicate
+from pypaimon.manifest.schema.simple_stats import SimpleStats
+from pypaimon.read.push_down_utils import rewrite_predicate_indices
from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader
from pypaimon.schema.data_types import DataField, PyarrowFieldParser
+from pypaimon.table.row.generic_row import GenericRow
from pypaimon.table.special_fields import SpecialFields
class FormatMosaicReader(RecordBatchReader):
def __init__(self, file_io: FileIO, file_path: str, read_fields:
List[DataField],
- push_down_predicate: Any, batch_size: int = 1024):
+ push_down_predicate: Any, batch_size: int = 1024,
+ row_group_predicate: Optional[Predicate] = None):
from mosaic import MosaicReader
+ self._read_fields = read_fields
self._read_field_names = [f.name for f in read_fields]
self._batch_size = batch_size
@@ -75,13 +85,21 @@ class FormatMosaicReader(RecordBatchReader):
self._num_row_groups = self._reader.num_row_groups
self._current_batches = None
- if push_down_predicate is not None:
- self._predicate = push_down_predicate
- else:
- self._predicate = None
+ self._predicate = push_down_predicate
+ self._row_group_predicate = None
+ if row_group_predicate is not None:
+ try:
+ self._row_group_predicate = rewrite_predicate_indices(
+ row_group_predicate, read_fields)
+ except ValueError:
+ pass
def _next_row_group_batches(self):
while self._current_rg < self._num_row_groups:
+ if not self._matches_row_group(self._current_rg):
+ self._current_rg += 1
+ continue
+
batch = self._reader.read_row_group(self._current_rg)
self._current_rg += 1
@@ -99,6 +117,46 @@ class FormatMosaicReader(RecordBatchReader):
max_chunksize=self._batch_size))
return None
+ def _matches_row_group(self, row_group_index: int) -> bool:
+ if self._row_group_predicate is None:
+ return True
+
+ try:
+ row_count = self._reader.row_group_num_rows(row_group_index)
+ stats_map = self._reader.get_row_group_statistics(row_group_index)
+ except Exception:
+ return True
+
+ if not stats_map:
+ return True
+
+ try:
+ stats = self._to_simple_stats(stats_map)
+ return self._row_group_predicate.test_by_simple_stats(stats,
row_count)
+ except Exception:
+ return True
+
+ def _to_simple_stats(self, stats_map) -> SimpleStats:
+ min_values = [None] * len(self._read_fields)
+ max_values = [None] * len(self._read_fields)
+ null_counts = [None] * len(self._read_fields)
+
+ for i, field in enumerate(self._read_fields):
+ stats = stats_map.get(field.name)
+ if stats is None:
+ continue
+
+ null_counts[i] = stats.null_count
+ if stats.has_min_max:
+ min_values[i] = _convert_stats_value(stats.min, field.type)
+ max_values[i] = _convert_stats_value(
+ stats.max, field.type, round_up_timestamp=True)
+
+ return SimpleStats(
+ GenericRow(min_values, self._read_fields),
+ GenericRow(max_values, self._read_fields),
+ null_counts)
+
def _fill_missing_fields(self, batch: RecordBatch) -> RecordBatch:
if not self.missing_fields:
return batch
@@ -138,3 +196,78 @@ class FormatMosaicReader(RecordBatchReader):
self._stream = None
self._reader = None
self._current_batches = None
+
+
+_EPOCH = datetime(1970, 1, 1)
+
+
+def _convert_stats_value(
+ value: Optional[bytes], data_type, round_up_timestamp: bool = False)
-> Any:
+ if value is None:
+ return None
+
+ type_name = str(data_type).upper().strip()
+ type_name = re.sub(r"\s+NOT\s+NULL$", "", type_name)
+ type_name = re.sub(r"\s+NULL$", "", type_name)
+ base_type = type_name.split("(", 1)[0].strip()
+
+ if base_type in ("CHAR", "VARCHAR", "STRING"):
+ return value.decode("utf-8")
+ if base_type in ("BINARY", "VARBINARY", "BYTES", "BLOB"):
+ return value
+ if len(value) == 0:
+ return None
+
+ if base_type == "BOOLEAN":
+ return value[0] != 0
+ if base_type == "TINYINT":
+ return struct.unpack(">b", value[:1])[0]
+ if base_type == "SMALLINT":
+ return struct.unpack(">h", value)[0]
+ if base_type in ("INT", "INTEGER", "DATE", "TIME"):
+ return struct.unpack(">i", value)[0]
+ if base_type == "BIGINT":
+ return struct.unpack(">q", value)[0]
+ if base_type == "FLOAT":
+ return struct.unpack(">f", value)[0]
+ if base_type == "DOUBLE":
+ return struct.unpack(">d", value)[0]
+ if base_type in ("DECIMAL", "NUMERIC"):
+ scale = _type_scale(type_name)
+ return Decimal(int.from_bytes(value, byteorder="big",
signed=True)).scaleb(-scale)
+ if base_type in ("TIMESTAMP", "TIMESTAMP_LTZ"):
+ precision = _type_precision(type_name, default=6)
+ if precision <= 3:
+ return _EPOCH + timedelta(milliseconds=struct.unpack(">q",
value)[0])
+ if precision <= 6:
+ return _EPOCH + timedelta(microseconds=struct.unpack(">q",
value)[0])
+ millis, nanos_of_milli = struct.unpack(">qi", value)
+ microseconds = nanos_of_milli // 1000
+ if round_up_timestamp and nanos_of_milli % 1000 != 0:
+ microseconds += 1
+ return _EPOCH + timedelta(
+ milliseconds=millis,
+ microseconds=microseconds)
+
+ return None
+
+
+def _type_precision(type_name: str, default: int = 0) -> int:
+ params = _type_params(type_name)
+ if not params:
+ return default
+ return int(params[0])
+
+
+def _type_scale(type_name: str) -> int:
+ params = _type_params(type_name)
+ if len(params) < 2:
+ return 0
+ return int(params[1])
+
+
+def _type_params(type_name: str) -> List[str]:
+ match = re.search(r"\(([^)]*)\)", type_name)
+ if match is None:
+ return []
+ return [p.strip() for p in match.group(1).split(",") if p.strip()]
diff --git a/paimon-python/pypaimon/read/split_read.py
b/paimon-python/pypaimon/read/split_read.py
index dd9058c55f..9167b7824a 100644
--- a/paimon-python/pypaimon/read/split_read.py
+++ b/paimon-python/pypaimon/read/split_read.py
@@ -202,7 +202,11 @@ class SplitRead(ABC):
read_fields: List[str], row_tracking_enabled:
bool,
row_ranges: Optional[List[Range]] = None,
shard_range: Optional[Tuple[int, int]] = None) ->
RecordBatchReader:
- (read_file_fields, read_arrow_predicate) =
self._get_fields_and_predicate(file.schema_id, read_fields)
+ (
+ read_file_fields,
+ read_arrow_predicate,
+ read_paimon_predicate,
+ ) = self._get_fields_and_predicate(file.schema_id, read_fields)
# Use external_path if available, otherwise use file_path
file_path = file.external_path if file.external_path else
file.file_path
@@ -316,8 +320,13 @@ class SplitRead(ABC):
raise NotImplementedError(
"Nested-field projection is not supported on Mosaic files")
ordered_read_fields = [name_to_field[n] for n in read_file_fields
if n in name_to_field]
+ row_group_predicate = (
+ read_paimon_predicate
+ if file.schema_id == self.table.table_schema.id else None
+ )
format_reader = FormatMosaicReader(self.table.file_io, file_path,
ordered_read_fields,
- read_arrow_predicate,
batch_size=batch_size)
+ read_arrow_predicate,
batch_size=batch_size,
+
row_group_predicate=row_group_predicate)
elif file_format == CoreOptions.FILE_FORMAT_PARQUET or file_format ==
CoreOptions.FILE_FORMAT_ORC:
ordered_read_fields = [name_to_field[n] for n in read_file_fields
if n in name_to_field]
ordered_nested_paths = (
@@ -490,7 +499,11 @@ class SplitRead(ABC):
]
read_predicate =
trim_predicate_by_fields(self.push_down_predicate, read_file_fields)
read_arrow_predicate = read_predicate.to_arrow() if read_predicate
else None
- self.schema_id_2_fields[key] = (read_file_fields,
read_arrow_predicate)
+ self.schema_id_2_fields[key] = (
+ read_file_fields,
+ read_arrow_predicate,
+ read_predicate,
+ )
return self.schema_id_2_fields[key]
@abstractmethod
diff --git a/paimon-python/pypaimon/tests/format_mosaic_reader_test.py
b/paimon-python/pypaimon/tests/format_mosaic_reader_test.py
new file mode 100644
index 0000000000..83493eb5af
--- /dev/null
+++ b/paimon-python/pypaimon/tests/format_mosaic_reader_test.py
@@ -0,0 +1,55 @@
+# 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.
+
+import struct
+from datetime import datetime
+
+from pypaimon.common.predicate_builder import PredicateBuilder
+from pypaimon.manifest.schema.simple_stats import SimpleStats
+from pypaimon.read.reader.format_mosaic_reader import _convert_stats_value
+from pypaimon.schema.data_types import AtomicType, DataField
+from pypaimon.table.row.generic_row import GenericRow
+
+
+def test_timestamp9_stats_rounds_max_up_for_row_group_skipping():
+ fields = [
+ DataField(0, "ts", AtomicType("TIMESTAMP(9)")),
+ ]
+ stats_bytes = struct.pack(">qi", 123, 456789)
+ min_value = _convert_stats_value(stats_bytes, fields[0].type)
+ max_value = _convert_stats_value(
+ stats_bytes, fields[0].type, round_up_timestamp=True)
+ stats = SimpleStats(
+ GenericRow([min_value], fields),
+ GenericRow([max_value], fields),
+ [0])
+
+ predicate = PredicateBuilder(fields).greater_than(
+ "ts", datetime(1970, 1, 1, 0, 0, 0, 123456))
+
+ assert min_value == datetime(1970, 1, 1, 0, 0, 0, 123456)
+ assert max_value == datetime(1970, 1, 1, 0, 0, 0, 123457)
+ assert predicate.test_by_simple_stats(stats, 1)
+
+
+def test_timestamp9_stats_keeps_exact_microsecond_max():
+ value = _convert_stats_value(
+ struct.pack(">qi", 123, 456000),
+ AtomicType("TIMESTAMP(9)"),
+ round_up_timestamp=True)
+
+ assert value == datetime(1970, 1, 1, 0, 0, 0, 123456)
diff --git a/paimon-python/pypaimon/tests/test_format_mosaic_reader_writer.py
b/paimon-python/pypaimon/tests/test_format_mosaic_reader_writer.py
index 154d56cd3c..67c4d63999 100644
--- a/paimon-python/pypaimon/tests/test_format_mosaic_reader_writer.py
+++ b/paimon-python/pypaimon/tests/test_format_mosaic_reader_writer.py
@@ -22,6 +22,7 @@ import pyarrow as pa
import pytest
import mosaic
+from pypaimon.common.predicate_builder import PredicateBuilder
from pypaimon.read.reader.format_mosaic_reader import FormatMosaicReader
from pypaimon.schema.data_types import AtomicType, DataField
@@ -36,15 +37,17 @@ class SimpleFileIO:
return open(path, 'rb')
-def _write_mosaic_file(path, data: pa.Table):
+def _write_mosaic_file(path, data: pa.Table, options=None):
with open(path, 'wb') as f:
- mosaic.write_table(data, f)
+ mosaic.write_table(data, f, options=options)
-def _read_mosaic_file(path, read_fields, push_down_predicate=None):
+def _read_mosaic_file(path, read_fields, push_down_predicate=None,
+ row_group_predicate=None):
file_io = SimpleFileIO()
reader = FormatMosaicReader(file_io, path, read_fields,
- push_down_predicate, batch_size=1024)
+ push_down_predicate, batch_size=1024,
+ row_group_predicate=row_group_predicate)
batches = []
while True:
batch = reader.read_arrow_batch()
@@ -253,6 +256,135 @@ class TestFormatMosaicReaderWriter:
finally:
os.unlink(path)
+ def test_predicate_skips_row_groups_by_stats(self, monkeypatch):
+ import pyarrow.compute as pc
+
+ fields = [
+ DataField(0, "id", AtomicType("INT")),
+ DataField(1, "name", AtomicType("STRING")),
+ ]
+ num_rows = 5000
+ data = pa.table({
+ "id": pa.array(list(range(num_rows)), type=pa.int32()),
+ "name": pa.array([f"user_{i}" for i in range(num_rows)],
type=pa.string()),
+ })
+
+ with tempfile.NamedTemporaryFile(suffix=".mosaic", delete=False) as
tmp:
+ path = tmp.name
+
+ original_from_input_file = mosaic.MosaicReader.from_input_file
+ readers = []
+
+ class CountingReader:
+ def __init__(self, reader):
+ self.reader = reader
+ self.read_row_groups = []
+
+ def __getattr__(self, name):
+ return getattr(self.reader, name)
+
+ def read_row_group(self, row_group_index):
+ self.read_row_groups.append(row_group_index)
+ return self.reader.read_row_group(row_group_index)
+
+ def from_input_file(read_at, file_length):
+ reader = CountingReader(original_from_input_file(read_at,
file_length))
+ readers.append(reader)
+ return reader
+
+ monkeypatch.setattr(
+ mosaic.MosaicReader, "from_input_file",
staticmethod(from_input_file))
+
+ try:
+ options = mosaic.WriterOptions(
+ compression=mosaic.WriterOptions.COMPRESSION_NONE,
+ num_buckets=1,
+ row_group_max_size=1024,
+ stats_columns=["id"])
+ with open(path, 'wb') as f:
+ with mosaic.MosaicWriter(f, data.schema, options) as writer:
+ for batch in data.to_batches(max_chunksize=250):
+ writer.write(batch)
+
+ predicate = PredicateBuilder(fields).greater_than("id", num_rows +
1)
+ result = _read_mosaic_file(
+ path,
+ fields,
+ push_down_predicate=pc.field("id") > num_rows + 1,
+ row_group_predicate=predicate)
+
+ assert result.num_rows == 0
+ assert len(readers) == 1
+ assert readers[0].num_row_groups > 1
+ assert readers[0].read_row_groups == []
+ finally:
+ os.unlink(path)
+
+ def test_predicate_skips_non_matching_row_groups_by_stats(self,
monkeypatch):
+ import pyarrow.compute as pc
+
+ fields = [
+ DataField(0, "id", AtomicType("INT")),
+ DataField(1, "name", AtomicType("STRING")),
+ ]
+ num_rows = 5000
+ data = pa.table({
+ "id": pa.array(list(range(num_rows)), type=pa.int32()),
+ "name": pa.array([f"user_{i}" for i in range(num_rows)],
type=pa.string()),
+ })
+
+ with tempfile.NamedTemporaryFile(suffix=".mosaic", delete=False) as
tmp:
+ path = tmp.name
+
+ original_from_input_file = mosaic.MosaicReader.from_input_file
+ readers = []
+
+ class CountingReader:
+ def __init__(self, reader):
+ self.reader = reader
+ self.read_row_groups = []
+
+ def __getattr__(self, name):
+ return getattr(self.reader, name)
+
+ def read_row_group(self, row_group_index):
+ self.read_row_groups.append(row_group_index)
+ return self.reader.read_row_group(row_group_index)
+
+ def from_input_file(read_at, file_length):
+ reader = CountingReader(original_from_input_file(read_at,
file_length))
+ readers.append(reader)
+ return reader
+
+ monkeypatch.setattr(
+ mosaic.MosaicReader, "from_input_file",
staticmethod(from_input_file))
+
+ try:
+ options = mosaic.WriterOptions(
+ compression=mosaic.WriterOptions.COMPRESSION_NONE,
+ num_buckets=1,
+ row_group_max_size=1024,
+ stats_columns=["id"])
+ with open(path, 'wb') as f:
+ with mosaic.MosaicWriter(f, data.schema, options) as writer:
+ for batch in data.to_batches(max_chunksize=250):
+ writer.write(batch)
+
+ predicate = PredicateBuilder(fields).greater_than("id", 4900)
+ result = _read_mosaic_file(
+ path,
+ fields,
+ push_down_predicate=pc.field("id") > 4900,
+ row_group_predicate=predicate)
+
+ assert result.column("id").to_pylist() == list(range(4901,
num_rows))
+ assert len(readers) == 1
+ assert readers[0].num_row_groups > 1
+ assert 0 not in readers[0].read_row_groups
+ assert len(readers[0].read_row_groups) < readers[0].num_row_groups
+ finally:
+ os.unlink(path)
+
def test_large_dataset(self):
fields = [
DataField(0, "id", AtomicType("INT")),
