Optimize reads of record batches by pushing limit to file level

Author: soumya-ghosh

pyiceberg/io/pyarrow.py

@@ -1194,6 +1194,7 @@ def _task_to_record_batches(
     case_sensitive: bool,
     name_mapping: Optional[NameMapping] = None,
     use_large_types: bool = True,
+    limit: Optional[int] = None,
 ) -> Iterator[pa.RecordBatch]:
     _, _, path = PyArrowFileIO.parse_location(task.file.file_path)
     arrow_format = ds.ParquetFileFormat(pre_buffer=True, buffer_size=(ONE_MEGABYTE * 8))
@@ -1232,8 +1233,11 @@ def _task_to_record_batches(
         )
 
         next_index = 0
+        total_rows = 0
         batches = fragment_scanner.to_batches()
         for batch in batches:
+            if limit is not None and total_rows >= limit:
+                break
             next_index = next_index + len(batch)
             current_index = next_index - len(batch)
             if positional_deletes:
@@ -1250,6 +1254,7 @@ def _task_to_record_batches(
                     if len(arrow_table) == 0:
                         continue
                     batch = arrow_table.to_batches()[0]
+            total_rows += len(batch)
             yield _to_requested_schema(
                 projected_schema, file_project_schema, batch, downcast_ns_timestamp_to_us=True, use_large_types=use_large_types
             )
@@ -1265,6 +1270,7 @@ def _task_to_table(
     case_sensitive: bool,
     name_mapping: Optional[NameMapping] = None,
     use_large_types: bool = True,
+    limit: Optional[int] = None,
 ) -> Optional[pa.Table]:
     batches = list(
         _task_to_record_batches(
@@ -1277,6 +1283,7 @@ def _task_to_table(
             case_sensitive,
             name_mapping,
             use_large_types,
+            limit,
         )
     )
 
@@ -1366,6 +1373,7 @@ def project_table(
             case_sensitive,
             table_metadata.name_mapping(),
             use_large_types,
+            limit,
         )
         for task in tasks
     ]
@@ -1464,6 +1472,7 @@ def project_batches(
             case_sensitive,
             table_metadata.name_mapping(),
             use_large_types,
+            limit,
         )
         for batch in batches:
             if limit is not None: