[python] Filter manifest entry by advance to reduce memory (#6428)

Author: JingsongLi

paimon-python/pypaimon/common/predicate_builder.py

@@ -101,17 +101,27 @@ def between(self, field: str, included_lower_bound: Any, included_upper_bound: A
         """Create a between predicate."""
         return self._build_predicate('between', field, [included_lower_bound, included_upper_bound])
 
-    def and_predicates(self, predicates: List[Predicate]) -> Predicate:
+    @staticmethod
+    def and_predicates(predicates: List[Predicate]) -> Optional[Predicate]:
         """Create an AND predicate from multiple predicates."""
+        if len(predicates) == 0:
+            return None
+        if len(predicates) == 1:
+            return predicates[0]
         return Predicate(
             method='and',
             index=None,
             field=None,
             literals=predicates
         )
 
-    def or_predicates(self, predicates: List[Predicate]) -> Predicate:
+    @staticmethod
+    def or_predicates(predicates: List[Predicate]) -> Optional[Predicate]:
         """Create an OR predicate from multiple predicates."""
+        if len(predicates) == 0:
+            return None
+        if len(predicates) == 1:
+            return predicates[0]
         return Predicate(
             method='or',
             index=None,

paimon-python/pypaimon/manifest/manifest_file_manager.py

@@ -41,7 +41,7 @@ def __init__(self, table):
         self.primary_key_fields = self.table.table_schema.get_primary_key_fields()
         self.trimmed_primary_key_fields = self.table.table_schema.get_trimmed_primary_key_fields()
 
-    def read(self, manifest_file_name: str, bucket_filter=None) -> List[ManifestEntry]:
+    def read(self, manifest_file_name: str, manifest_entry_filter=None) -> List[ManifestEntry]:
         manifest_file_path = self.manifest_path / manifest_file_name
 
         entries = []
@@ -105,7 +105,7 @@ def read(self, manifest_file_name: str, bucket_filter=None) -> List[ManifestEntr
                 total_buckets=record['_TOTAL_BUCKETS'],
                 file=file_meta
             )
-            if bucket_filter is not None and not bucket_filter(entry):
+            if manifest_entry_filter is not None and not manifest_entry_filter(entry):
                 continue
             entries.append(entry)
         return entries

paimon-python/pypaimon/read/push_down_utils.py

@@ -19,31 +19,29 @@
 from typing import Dict, List, Set
 
 from pypaimon.common.predicate import Predicate
+from pypaimon.common.predicate_builder import PredicateBuilder
 
 
-def to_partition_predicate(input_predicate: 'Predicate', all_fields: List[str], partition_keys: List[str]):
-    if not input_predicate or not partition_keys:
-        return None
-
-    predicates: list['Predicate'] = _split_and(input_predicate)
-    predicates = [element for element in predicates if _get_all_fields(element).issubset(partition_keys)]
-    new_predicate = Predicate(
-        method='and',
-        index=None,
-        field=None,
-        literals=predicates
-    )
-
-    part_to_index = {element: idx for idx, element in enumerate(partition_keys)}
+def filter_and_transform_predicate(input_predicate: Predicate, all_fields: List[str], fields: List[str]):
+    new_predicate = filter_predicate_by_fields(input_predicate, fields)
+    part_to_index = {element: idx for idx, element in enumerate(fields)}
     mapping: Dict[int, int] = {
         i: part_to_index.get(all_fields[i], -1)
         for i in range(len(all_fields))
     }
-
     return _change_index(new_predicate, mapping)
 
 
-def _split_and(input_predicate: 'Predicate'):
+def filter_predicate_by_fields(input_predicate: Predicate, fields: List[str]):
+    if not input_predicate or not fields:
+        return None
+
+    predicates: list[Predicate] = _split_and(input_predicate)
+    predicates = [element for element in predicates if _get_all_fields(element).issubset(fields)]
+    return PredicateBuilder.and_predicates(predicates)
+
+
+def _split_and(input_predicate: Predicate):
     if not input_predicate:
         return list()
 
@@ -53,65 +51,23 @@ def _split_and(input_predicate: 'Predicate'):
     return [input_predicate]
 
 
-def _change_index(input_predicate: 'Predicate', mapping: Dict[int, int]):
+def _change_index(input_predicate: Predicate, mapping: Dict[int, int]):
     if not input_predicate:
         return None
 
     if input_predicate.method == 'and' or input_predicate.method == 'or':
-        predicates: list['Predicate'] = input_predicate.literals
+        predicates: list[Predicate] = input_predicate.literals
         new_predicates = [_change_index(element, mapping) for element in predicates]
         return input_predicate.new_literals(new_predicates)
 
     return input_predicate.new_index(mapping[input_predicate.index])
 
 
-def extract_predicate_to_list(result: list, input_predicate: 'Predicate', keys: List[str]):
-    if not input_predicate or not keys:
-        return
-
-    if input_predicate.method == 'and':
-        for sub_predicate in input_predicate.literals:
-            extract_predicate_to_list(result, sub_predicate, keys)
-        return
-    elif input_predicate.method == 'or':
-        # condition: involved keys all belong to primary keys
-        involved_fields = _get_all_fields(input_predicate)
-        if involved_fields and involved_fields.issubset(keys):
-            result.append(input_predicate)
-        return
-
-    if input_predicate.field in keys:
-        result.append(input_predicate)
-
-
-def _get_all_fields(predicate: 'Predicate') -> Set[str]:
+def _get_all_fields(predicate: Predicate) -> Set[str]:
     if predicate.field is not None:
         return {predicate.field}
     involved_fields = set()
     if predicate.literals:
         for sub_predicate in predicate.literals:
             involved_fields.update(_get_all_fields(sub_predicate))
     return involved_fields
-
-
-def extract_predicate_to_dict(result: Dict, input_predicate: 'Predicate', keys: List[str]):
-    if not input_predicate or not keys:
-        return
-
-    if input_predicate.method == 'and':
-        for sub_predicate in input_predicate.literals:
-            extract_predicate_to_dict(result, sub_predicate, keys)
-        return
-    elif input_predicate.method == 'or':
-        # ensure no recursive and/or
-        if not input_predicate.literals or any(p.field is None for p in input_predicate.literals):
-            return
-        # condition: only one key for 'or', and the key belongs to keys
-        involved_fields = {p.field for p in input_predicate.literals}
-        field = involved_fields.pop() if len(involved_fields) == 1 else None
-        if field is not None and field in keys:
-            result[field].append(input_predicate)
-        return
-
-    if input_predicate.field in keys:
-        result[input_predicate.field].append(input_predicate)

paimon-python/pypaimon/read/scanner/full_starting_scanner.py

@@ -20,17 +20,14 @@
 
 from pypaimon.common.core_options import CoreOptions
 from pypaimon.common.predicate import Predicate
-from pypaimon.common.predicate_builder import PredicateBuilder
 from pypaimon.manifest.manifest_file_manager import ManifestFileManager
 from pypaimon.manifest.manifest_list_manager import ManifestListManager
 from pypaimon.manifest.schema.data_file_meta import DataFileMeta
 from pypaimon.manifest.schema.manifest_entry import ManifestEntry
 from pypaimon.manifest.schema.manifest_file_meta import ManifestFileMeta
 from pypaimon.read.interval_partition import IntervalPartition, SortedRun
 from pypaimon.read.plan import Plan
-from pypaimon.read.push_down_utils import (extract_predicate_to_dict,
-                                           extract_predicate_to_list,
-                                           to_partition_predicate)
+from pypaimon.read.push_down_utils import (filter_and_transform_predicate)
 from pypaimon.read.scanner.starting_scanner import StartingScanner
 from pypaimon.read.split import Split
 from pypaimon.snapshot.snapshot_manager import SnapshotManager
@@ -49,14 +46,11 @@ def __init__(self, table, predicate: Optional[Predicate], limit: Optional[int]):
         self.manifest_list_manager = ManifestListManager(table)
         self.manifest_file_manager = ManifestFileManager(table)
 
-        pk_conditions = []
-        trimmed_pk = [field.name for field in self.table.table_schema.get_trimmed_primary_key_fields()]
-        extract_predicate_to_list(pk_conditions, self.predicate, trimmed_pk)
-        self.primary_key_predicate = PredicateBuilder(self.table.fields).and_predicates(pk_conditions)
+        self.primary_key_predicate = filter_and_transform_predicate(
+            self.predicate, self.table.field_names, self.table.table_schema.get_trimmed_primary_keys())
 
-        partition_conditions = defaultdict(list)
-        extract_predicate_to_dict(partition_conditions, self.predicate, self.table.partition_keys)
-        self.partition_key_predicate = partition_conditions
+        self.partition_key_predicate = filter_and_transform_predicate(
+            self.predicate, self.table.field_names, self.table.partition_keys)
 
         self.target_split_size = 128 * 1024 * 1024
         self.open_file_cost = 4 * 1024 * 1024
@@ -82,29 +76,29 @@ def scan(self) -> Plan:
         splits = self._apply_push_down_limit(splits)
         return Plan(splits)
 
-    def _read_manifest_files(self) -> List[ManifestFileMeta]:
+    def plan_files(self) -> List[ManifestEntry]:
         latest_snapshot = self.snapshot_manager.get_latest_snapshot()
         if not latest_snapshot:
             return []
         manifest_files = self.manifest_list_manager.read_all(latest_snapshot)
-        partition_predicate = to_partition_predicate(self.predicate, self.table.field_names, self.table.partition_keys)
+        return self.read_manifest_entries(manifest_files)
 
-        def test_predicate(file: ManifestFileMeta) -> bool:
-            if not partition_predicate:
+    def read_manifest_entries(self, manifest_files: List[ManifestFileMeta]) -> List[ManifestEntry]:
+        def filter_manifest_file(file: ManifestFileMeta) -> bool:
+            if not self.partition_key_predicate:
                 return True
-            return partition_predicate.test_by_simple_stats(
+            return self.partition_key_predicate.test_by_simple_stats(
                 file.partition_stats,
                 file.num_added_files + file.num_deleted_files)
 
-        return [file for file in manifest_files if test_predicate(file)]
-
-    def plan_files(self) -> List[ManifestEntry]:
-        manifest_files = self._read_manifest_files()
         deleted_entries = set()
         added_entries = []
         for manifest_file in manifest_files:
-            manifest_entries = self.manifest_file_manager.read(manifest_file.file_name,
-                                                               lambda row: self._bucket_filter(row))
+            if not filter_manifest_file(manifest_file):
+                continue
+            manifest_entries = self.manifest_file_manager.read(
+                manifest_file.file_name,
+                lambda row: self._filter_manifest_entry(row))
             for entry in manifest_entries:
                 if entry.kind == 0:
                     added_entries.append(entry)
@@ -115,8 +109,6 @@ def plan_files(self) -> List[ManifestEntry]:
             entry for entry in added_entries
             if (tuple(entry.partition.values), entry.bucket, entry.file.file_name) not in deleted_entries
         ]
-        if self.predicate:
-            file_entries = self._filter_by_predicate(file_entries)
         return file_entries
 
     def with_shard(self, idx_of_this_subtask, number_of_para_subtasks) -> 'FullStartingScanner':
@@ -203,12 +195,6 @@ def _primary_key_filter_by_shard(self, file_entries: List[ManifestEntry]) -> Lis
                 filtered_entries.append(entry)
         return filtered_entries
 
-    def _bucket_filter(self, entry: Optional[ManifestEntry]) -> bool:
-        bucket = entry.bucket
-        if self.only_read_real_buckets and bucket < 0:
-            return False
-        return True
-
     def _apply_push_down_limit(self, splits: List[Split]) -> List[Split]:
         if self.limit is None:
             return splits
@@ -224,45 +210,26 @@ def _apply_push_down_limit(self, splits: List[Split]) -> List[Split]:
 
         return limited_splits
 
-    def _filter_by_predicate(self, file_entries: List[ManifestEntry]) -> List[ManifestEntry]:
-        if not self.predicate:
-            return file_entries
-
-        filtered_files = []
-        for file_entry in file_entries:
-            if self.partition_key_predicate and not self._filter_by_partition(file_entry):
-                continue
-            if not self._filter_by_stats(file_entry):
-                continue
-            filtered_files.append(file_entry)
-
-        return filtered_files
-
-    def _filter_by_partition(self, file_entry: ManifestEntry) -> bool:
-        partition_dict = file_entry.partition.to_dict()
-        for field_name, conditions in self.partition_key_predicate.items():
-            partition_value = partition_dict[field_name]
-            for predicate in conditions:
-                if not predicate.test_by_value(partition_value):
-                    return False
-        return True
-
-    def _filter_by_stats(self, file_entry: ManifestEntry) -> bool:
-        if file_entry.kind != 0:
+    def _filter_manifest_entry(self, entry: ManifestEntry) -> bool:
+        if self.only_read_real_buckets and entry.bucket < 0:
+            return False
+        if self.partition_key_predicate and not self.partition_key_predicate.test(entry.partition):
             return False
         if self.table.is_primary_key_table:
             predicate = self.primary_key_predicate
-            stats = file_entry.file.key_stats
+            stats = entry.file.key_stats
         else:
             predicate = self.predicate
-            stats = file_entry.file.value_stats
+            stats = entry.file.value_stats
+        if not predicate:
+            return True
         return predicate.test_by_stats({
             "min_values": stats.min_values.to_dict(),
             "max_values": stats.max_values.to_dict(),
             "null_counts": {
                 stats.min_values.fields[i].name: stats.null_counts[i] for i in range(len(stats.min_values.fields))
             },
-            "row_count": file_entry.file.row_count,
+            "row_count": entry.file.row_count
         })
 
     def _create_append_only_splits(self, file_entries: List[ManifestEntry]) -> List['Split']:

paimon-python/pypaimon/read/scanner/incremental_starting_scanner.py

@@ -43,13 +43,7 @@ def plan_files(self) -> List[ManifestEntry]:
         for snapshot in snapshots_in_range:
             # Get manifest files for this snapshot
             manifest_files = self.manifest_list_manager.read_delta(snapshot)
-
-            # Read all entries from manifest files
-            for manifest_file in manifest_files:
-                entries = self.manifest_file_manager.read(manifest_file.file_name)
-                file_entries.extend(entries)
-        if self.predicate:
-            file_entries = self._filter_by_predicate(file_entries)
+            file_entries.extend(self.read_manifest_entries(manifest_files))
         return file_entries
 
     @staticmethod

paimon-python/pypaimon/read/table_read.py

@@ -22,8 +22,7 @@
 
 from pypaimon.common.core_options import CoreOptions
 from pypaimon.common.predicate import Predicate
-from pypaimon.common.predicate_builder import PredicateBuilder
-from pypaimon.read.push_down_utils import extract_predicate_to_list
+from pypaimon.read.push_down_utils import filter_predicate_by_fields
 from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader
 from pypaimon.read.split import Split
 from pypaimon.read.split_read import (MergeFileSplitRead, RawFileSplitRead,
@@ -113,14 +112,10 @@ def _push_down_predicate(self) -> Any:
         if self.predicate is None:
             return None
         elif self.table.is_primary_key_table:
-            result = []
-            extract_predicate_to_list(result, self.predicate, self.table.primary_keys)
-            if result:
-                # the field index is unused for arrow field
-                pk_predicates = (PredicateBuilder(self.table.fields).and_predicates(result)).to_arrow()
-                return pk_predicates
-            else:
+            pk_predicate = filter_predicate_by_fields(self.predicate, self.table.primary_keys)
+            if not pk_predicate:
                 return None
+            return pk_predicate.to_arrow()
         else:
             return self.predicate.to_arrow()
 

paimon-python/pypaimon/schema/table_schema.py

@@ -170,3 +170,6 @@ def get_trimmed_primary_key_fields(self) -> List[DataField]:
                 "this will result in only one record in a partition")
         field_map = {field.name: field for field in self.fields}
         return [field_map[name] for name in adjusted if name in field_map]
+
+    def get_trimmed_primary_keys(self) -> List[str]:
+        return [field.name for field in self.get_trimmed_primary_key_fields()]