Apply residuals when reading a table

Author: Fokko

pyiceberg/io/pyarrow.py

@@ -1342,9 +1342,8 @@ def _get_column_projection_values(
 def _task_to_record_batches(
     fs: FileSystem,
     task: FileScanTask,
-    bound_row_filter: BooleanExpression,
+    schema: Schema,
     projected_schema: Schema,
-    projected_field_ids: Set[int],
     positional_deletes: Optional[List[ChunkedArray]],
     case_sensitive: bool,
     name_mapping: Optional[NameMapping] = None,
@@ -1363,8 +1362,8 @@ def _task_to_record_batches(
         file_schema = pyarrow_to_schema(physical_schema, name_mapping, downcast_ns_timestamp_to_us=True)
 
         pyarrow_filter = None
-        if bound_row_filter is not AlwaysTrue():
-            translated_row_filter = translate_column_names(bound_row_filter, file_schema, case_sensitive=case_sensitive)
+        if task.residual is not AlwaysTrue():
+            translated_row_filter = translate_column_names(task.residual, file_schema, case_sensitive=case_sensitive)
             bound_file_filter = bind(file_schema, translated_row_filter, case_sensitive=case_sensitive)
             pyarrow_filter = expression_to_pyarrow(bound_file_filter)
 
@@ -1374,7 +1373,13 @@ def _task_to_record_batches(
             task.file, projected_schema, partition_spec, file_schema.field_ids
         )
 
-        file_project_schema = prune_columns(file_schema, projected_field_ids, select_full_types=False)
+        file_project_schema = prune_columns(
+            file_schema,
+            {
+                id for id in projected_schema.field_ids if not isinstance(projected_schema.find_type(id), (MapType, ListType))
+            }.union(extract_field_ids(task.residual)),
+            select_full_types=False,
+        )
 
         fragment_scanner = ds.Scanner.from_fragment(
             fragment=fragment,
@@ -1474,7 +1479,7 @@ class ArrowScan:
     _table_metadata: TableMetadata
     _io: FileIO
     _projected_schema: Schema
-    _bound_row_filter: BooleanExpression
+    _bound_row_filter: Optional[BooleanExpression]
     _case_sensitive: bool
     _limit: Optional[int]
     """Scan the Iceberg Table and create an Arrow construct.
@@ -1493,14 +1498,18 @@ def __init__(
         table_metadata: TableMetadata,
         io: FileIO,
         projected_schema: Schema,
-        row_filter: BooleanExpression,
+        row_filter: Optional[BooleanExpression] = None,
         case_sensitive: bool = True,
         limit: Optional[int] = None,
     ) -> None:
         self._table_metadata = table_metadata
         self._io = io
         self._projected_schema = projected_schema
-        self._bound_row_filter = bind(table_metadata.schema(), row_filter, case_sensitive=case_sensitive)
+        # TBD: Should we deprecate the `row_filter` argument?
+        if row_filter is not None:
+            self._bound_row_filter = bind(table_metadata.schema(), row_filter, case_sensitive=case_sensitive)
+        else:
+            self._bound_row_filter = None
         self._case_sensitive = case_sensitive
         self._limit = limit
 
@@ -1512,15 +1521,6 @@ def _use_large_types(self) -> bool:
         """
         return property_as_bool(self._io.properties, PYARROW_USE_LARGE_TYPES_ON_READ, True)
 
-    @property
-    def _projected_field_ids(self) -> Set[int]:
-        """Set of field IDs that should be projected from the data files."""
-        return {
-            id
-            for id in self._projected_schema.field_ids
-            if not isinstance(self._projected_schema.find_type(id), (MapType, ListType))
-        }.union(extract_field_ids(self._bound_row_filter))
-
     def to_table(self, tasks: Iterable[FileScanTask]) -> pa.Table:
         """Scan the Iceberg table and return a pa.Table.
 
@@ -1541,7 +1541,10 @@ def to_table(self, tasks: Iterable[FileScanTask]) -> pa.Table:
         deletes_per_file = _read_all_delete_files(self._io, tasks)
         executor = ExecutorFactory.get_or_create()
 
-        def _table_from_scan_task(task: FileScanTask) -> pa.Table:
+        if self._bound_row_filter is not None:
+            tasks = [task.set_residual(expr=self._bound_row_filter) for task in tasks]
+
+        def _table_from_scan_task(task: FileScanTask) -> Optional[pa.Table]:
             batches = list(self._record_batches_from_scan_tasks_and_deletes([task], deletes_per_file))
             if len(batches) > 0:
                 return pa.Table.from_batches(batches)
@@ -1601,6 +1604,9 @@ def to_record_batches(self, tasks: Iterable[FileScanTask]) -> Iterator[pa.Record
             ResolveError: When a required field cannot be found in the file
             ValueError: When a field type in the file cannot be projected to the schema type
         """
+        if self._bound_row_filter is not None:
+            tasks = [task.set_residual(expr=self._bound_row_filter) for task in tasks]
+
         deletes_per_file = _read_all_delete_files(self._io, tasks)
         return self._record_batches_from_scan_tasks_and_deletes(tasks, deletes_per_file)
 
@@ -1614,9 +1620,8 @@ def _record_batches_from_scan_tasks_and_deletes(
             batches = _task_to_record_batches(
                 _fs_from_file_path(self._io, task.file.file_path),
                 task,
-                self._bound_row_filter,
+                self._table_metadata.schema(),
                 self._projected_schema,
-                self._projected_field_ids,
                 deletes_per_file.get(task.file.file_path),
                 self._case_sensitive,
                 self._table_metadata.name_mapping(),

pyiceberg/table/__init__.py

@@ -1394,6 +1394,10 @@ def __init__(
         self.length = length or data_file.file_size_in_bytes
         self.residual = residual
 
+    def set_residual(self, expr: BooleanExpression) -> "FileScanTask":
+        self.residual = expr
+        return self
+
 
 def _open_manifest(
     io: FileIO,
@@ -1584,8 +1588,12 @@ def plan_files(self) -> Iterable[FileScanTask]:
                     data_entry,
                     positional_delete_entries,
                 ),
-                residual=residual_evaluators[data_entry.data_file.spec_id](data_entry.data_file).residual_for(
-                    data_entry.data_file.partition
+                residual=bind(
+                    self.table_metadata.schema(),
+                    residual_evaluators[data_entry.data_file.spec_id](data_entry.data_file).residual_for(
+                        data_entry.data_file.partition
+                    ),
+                    case_sensitive=self.case_sensitive,
                 ),
             )
             for data_entry in data_entries

tests/io/test_pyarrow.py

@@ -984,7 +984,8 @@ def project(
                     partition={},
                     record_count=3,
                     file_size_in_bytes=3,
-                )
+                ),
+                residual=expr or AlwaysTrue(),
             )
             for file in files
         ]