Python: Add limit parameter to table scan

.github/workflows/python-integration.yml

@@ -46,7 +46,7 @@ jobs:
       id: check_file_changed
       run: |
         $diff = git diff --name-only HEAD^ HEAD
-        $SourceDiff = $diff | Where-Object { $_ -match '^python/dev/Dockerfile$' }
+        $SourceDiff = $diff | Where-Object { $_ -match '^python/dev/.+$' }
         $HasDiff = $SourceDiff.Length -gt 0
         Write-Host "::set-output name=docs_changed::$HasDiff"
     - name: Restore image
@@ -84,4 +84,4 @@ jobs:
       run: make test-integration
     - name: Show debug logs
       if: ${{ failure() }}
-      run: docker-compose -f python/dev/docker-compose.yml logs
+      run: docker-compose -f python/dev/docker-compose.yml logs

python/Makefile

@@ -36,7 +36,7 @@ test-integration:
 	docker-compose -f dev/docker-compose-integration.yml kill
 	docker-compose -f dev/docker-compose-integration.yml build
 	docker-compose -f dev/docker-compose-integration.yml up -d
-	sleep 20
+	sleep 30
 	poetry run pytest tests/ -m integration ${PYTEST_ARGS}
 
 test-adlfs:

python/dev/provision.py

@@ -64,6 +64,39 @@
 """
 )
 
+spark.sql(
+    """
+  DROP TABLE IF EXISTS test_limit;
+"""
+)
+
+spark.sql(
+    """
+    CREATE TABLE test_limit
+    USING iceberg
+      AS SELECT
+          1            AS idx
+      UNION ALL SELECT
+          2            AS idx
+      UNION ALL SELECT
+          3            AS idx
+      UNION ALL SELECT
+          4            AS idx
+      UNION ALL SELECT
+          5            AS idx
+      UNION ALL SELECT
+          6            AS idx
+      UNION ALL SELECT
+          7            AS idx
+      UNION ALL SELECT
+          8            AS idx
+      UNION ALL SELECT
+          9            AS idx
+      UNION ALL SELECT
+          10           AS idx
+    """
+)
+
 spark.sql(
     """
   DROP TABLE IF EXISTS test_deletes;

python/pyiceberg/files.py

@@ -28,7 +28,7 @@ class FileContentType(Enum):
 class FileFormat(Enum):
     """An enum that includes all possible formats for an Iceberg data file"""
 
-    ORC = auto()
-    PARQUET = auto()
-    AVRO = auto()
-    METADATA = auto()
+    ORC = "ORC"
+    PARQUET = "PARQUET"
+    AVRO = "AVRO"
+    METADATA = "METADATA"

python/pyiceberg/io/pyarrow.py

@@ -24,9 +24,11 @@
 """
 from __future__ import annotations
 
+import multiprocessing
 import os
 from functools import lru_cache
 from multiprocessing.pool import ThreadPool
+from multiprocessing.sharedctypes import Synchronized
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -42,7 +44,7 @@
 
 import pyarrow as pa
 import pyarrow.compute as pc
-import pyarrow.parquet as pq
+import pyarrow.dataset as ds
 from pyarrow.fs import (
     FileInfo,
     FileSystem,
@@ -491,14 +493,19 @@ def _file_to_table(
     projected_schema: Schema,
     projected_field_ids: Set[int],
     case_sensitive: bool,
+    rows_counter: Synchronized[int],
+    limit: Optional[int] = None,
 ) -> Optional[pa.Table]:
-    _, path = PyArrowFileIO.parse_location(task.file.file_path)
+    if limit and rows_counter.value >= limit:
+        return None
 
-    # Get the schema
-    with fs.open_input_file(path) as fout:
-        parquet_schema = pq.read_schema(fout)
+    _, path = PyArrowFileIO.parse_location(task.file.file_path)
+    arrow_format = ds.ParquetFileFormat(pre_buffer=True, buffer_size=(ONE_MEGABYTE * 8))
+    with fs.open_input_file(path) as fin:
+        fragment = arrow_format.make_fragment(fin)
+        physical_schema = fragment.physical_schema
         schema_raw = None
-        if metadata := parquet_schema.metadata:
+        if metadata := physical_schema.metadata:
             schema_raw = metadata.get(ICEBERG_SCHEMA)
         if schema_raw is None:
             raise ValueError(
@@ -517,15 +524,22 @@ def _file_to_table(
         if file_schema is None:
             raise ValueError(f"Missing Iceberg schema in Metadata for file: {path}")
 
-        arrow_table = pq.read_table(
-            source=fout,
-            schema=parquet_schema,
-            pre_buffer=True,
-            buffer_size=8 * ONE_MEGABYTE,
-            filters=pyarrow_filter,
+        fragment_scanner = ds.Scanner.from_fragment(
+            fragment=fragment,
+            schema=physical_schema,
+            filter=pyarrow_filter,
             columns=[col.name for col in file_project_schema.columns],
         )
 
+        if limit:
+            arrow_table = fragment_scanner.head(limit)
+            with rows_counter.get_lock():
+                if rows_counter.value >= limit:
+                    return None
+                rows_counter.value += len(arrow_table)
+        else:
+            arrow_table = fragment_scanner.to_table()
+
         # If there is no data, we don't have to go through the schema
         if len(arrow_table) > 0:
             return to_requested_schema(projected_schema, file_project_schema, arrow_table)
@@ -534,7 +548,12 @@ def _file_to_table(
 
 
 def project_table(
-    tasks: Iterable[FileScanTask], table: Table, row_filter: BooleanExpression, projected_schema: Schema, case_sensitive: bool
+    tasks: Iterable[FileScanTask],
+    table: Table,
+    row_filter: BooleanExpression,
+    projected_schema: Schema,
+    case_sensitive: bool,
+    limit: Optional[int] = None,
 ) -> pa.Table:
     """Resolves the right columns based on the identifier
 
@@ -570,23 +589,30 @@ def project_table(
         id for id in projected_schema.field_ids if not isinstance(projected_schema.find_type(id), (MapType, ListType))
     }.union(extract_field_ids(bound_row_filter))
 
+    rows_counter = multiprocessing.Value("i", 0)
+
     with ThreadPool() as pool:
         tables = [
             table
             for table in pool.starmap(
                 func=_file_to_table,
-                iterable=[(fs, task, bound_row_filter, projected_schema, projected_field_ids, case_sensitive) for task in tasks],
+                iterable=[
+                    (fs, task, bound_row_filter, projected_schema, projected_field_ids, case_sensitive, rows_counter, limit)
+                    for task in tasks
+                ],
                 chunksize=None,  # we could use this to control how to materialize the generator of tasks (we should also make the expression above lazy)
             )
             if table is not None
         ]
 
     if len(tables) > 1:
-        return pa.concat_tables(tables)
+        final_table = pa.concat_tables(tables)
     elif len(tables) == 1:
-        return tables[0]
+        final_table = tables[0]
     else:
-        return pa.Table.from_batches([], schema=schema_to_pyarrow(projected_schema))
+        final_table = pa.Table.from_batches([], schema=schema_to_pyarrow(projected_schema))
+
+    return final_table.slice(0, limit)
 
 
 def to_requested_schema(requested_schema: Schema, file_schema: Schema, table: pa.Table) -> pa.Table:

python/pyiceberg/manifest.py

@@ -177,6 +177,12 @@ class DataFile(Record):
     sort_order_id: Optional[int]
     spec_id: Optional[int]
 
+    def __setattr__(self, name: str, value: Any) -> None:
+        # The file_format is written as a string, so we need to cast it to the Enum
+        if name == "file_format":
+            value = FileFormat[value]
+        super().__setattr__(name, value)
+
     def __init__(self, *data: Any, **named_data: Any) -> None:
         super().__init__(*data, **{"struct": DATA_FILE_TYPE, **named_data})
 

python/pyiceberg/table/__init__.py

@@ -99,6 +99,7 @@ def scan(
         case_sensitive: bool = True,
         snapshot_id: Optional[int] = None,
         options: Properties = EMPTY_DICT,
+        limit: Optional[int] = None,
     ) -> DataScan:
         return DataScan(
             table=self,
@@ -107,6 +108,7 @@ def scan(
             case_sensitive=case_sensitive,
             snapshot_id=snapshot_id,
             options=options,
+            limit=limit,
         )
 
     def schema(self) -> Schema:
@@ -219,6 +221,7 @@ class TableScan(ABC):
     case_sensitive: bool
     snapshot_id: Optional[int]
     options: Properties
+    limit: Optional[int]
 
     def __init__(
         self,
@@ -228,13 +231,15 @@ def __init__(
         case_sensitive: bool = True,
         snapshot_id: Optional[int] = None,
         options: Properties = EMPTY_DICT,
+        limit: Optional[int] = None,
     ):
         self.table = table
         self.row_filter = _parse_row_filter(row_filter)
         self.selected_fields = selected_fields
         self.case_sensitive = case_sensitive
         self.snapshot_id = snapshot_id
         self.options = options
+        self.limit = limit
 
     def snapshot(self) -> Optional[Snapshot]:
         if self.snapshot_id:
@@ -335,8 +340,9 @@ def __init__(
         case_sensitive: bool = True,
         snapshot_id: Optional[int] = None,
         options: Properties = EMPTY_DICT,
+        limit: Optional[int] = None,
     ):
-        super().__init__(table, row_filter, selected_fields, case_sensitive, snapshot_id, options)
+        super().__init__(table, row_filter, selected_fields, case_sensitive, snapshot_id, options, limit)
 
     def _build_partition_projection(self, spec_id: int) -> BooleanExpression:
         project = inclusive_projection(self.table.schema(), self.table.specs()[spec_id])
@@ -402,7 +408,12 @@ def to_arrow(self) -> pa.Table:
         from pyiceberg.io.pyarrow import project_table
 
         return project_table(
-            self.plan_files(), self.table, self.row_filter, self.projection(), case_sensitive=self.case_sensitive
+            self.plan_files(),
+            self.table,
+            self.row_filter,
+            self.projection(),
+            case_sensitive=self.case_sensitive,
+            limit=self.limit,
         )
 
     def to_pandas(self, **kwargs: Any) -> pd.DataFrame:

python/tests/test_integration.py

@@ -49,6 +49,11 @@ def table_test_null_nan_rewritten(catalog: Catalog) -> Table:
     return catalog.load_table("default.test_null_nan_rewritten")
 
 
+@pytest.fixture()
+def table_test_limit(catalog: Catalog) -> Table:
+    return catalog.load_table("default.test_limit")
+
+
 @pytest.mark.integration
 def test_pyarrow_nan(table_test_null_nan: Table) -> None:
     arrow_table = table_test_null_nan.scan(row_filter=IsNaN("col_numeric"), selected_fields=("idx", "col_numeric")).to_arrow()
@@ -80,3 +85,15 @@ def test_duckdb_nan(table_test_null_nan_rewritten: Table) -> None:
     result = con.query("SELECT idx, col_numeric FROM table_test_null_nan WHERE isnan(col_numeric)").fetchone()
     assert result[0] == 1
     assert math.isnan(result[1])
+
+
+@pytest.mark.integration
+def test_pyarrow_limit(table_test_limit: Table) -> None:
+    limited_result = table_test_limit.scan(selected_fields=("idx",), limit=1).to_arrow()
+    assert len(limited_result) == 1
+
+    empty_result = table_test_limit.scan(selected_fields=("idx",), limit=0).to_arrow()
+    assert len(empty_result) == 0
+
+    full_result = table_test_limit.scan(selected_fields=("idx",), limit=999).to_arrow()
+    assert len(full_result) == 10