From 7ccabd4aec21cac67282519f48306a13de8bed2c Mon Sep 17 00:00:00 2001 From: Fokko Driesprong Date: Tue, 7 Mar 2023 19:18:53 +0100 Subject: [PATCH 01/11] Python: Add support for ORC Creates fragments based on the FileFormat. Blocked by: https://github.com/apache/iceberg/pull/6997 --- python/pyiceberg/files.py | 8 ++++---- python/pyiceberg/io/pyarrow.py | 36 ++++++++++++++++++++++------------ python/pyiceberg/manifest.py | 10 ++++++++-- 3 files changed, 35 insertions(+), 19 deletions(-) diff --git a/python/pyiceberg/files.py b/python/pyiceberg/files.py index b832bbbfe728..df19f37e0c40 100644 --- a/python/pyiceberg/files.py +++ b/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" diff --git a/python/pyiceberg/io/pyarrow.py b/python/pyiceberg/io/pyarrow.py index 4cf81dadd8e0..b7e9a08358e9 100644 --- a/python/pyiceberg/io/pyarrow.py +++ b/python/pyiceberg/io/pyarrow.py @@ -31,6 +31,7 @@ TYPE_CHECKING, Any, Callable, + Dict, Iterable, List, Optional, @@ -42,7 +43,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, @@ -67,6 +68,7 @@ translate_column_names, ) from pyiceberg.expressions.visitors import visit as boolean_expression_visit +from pyiceberg.files import FileFormat from pyiceberg.io import ( FileIO, InputFile, @@ -478,6 +480,16 @@ def expression_to_pyarrow(expr: BooleanExpression) -> pc.Expression: return boolean_expression_visit(expr, _ConvertToArrowExpression()) +@lru_cache +def _get_file_format(file_format: FileFormat, **kwargs: Dict[str, Any]) -> ds.FileFormat: + if file_format == FileFormat.PARQUET.value: + return ds.ParquetFileFormat(**kwargs) + elif file_format == FileFormat.ORC.value: + return ds.OrcFileFormat() + else: + raise ValueError(f"Unsupported file format: {file_format}") + + def _file_to_table( fs: FileSystem, task: FileScanTask, @@ -487,12 +499,12 @@ def _file_to_table( case_sensitive: bool, ) -> Optional[pa.Table]: _, path = PyArrowFileIO.parse_location(task.file.file_path) - - # Get the schema - with fs.open_input_file(path) as fout: - parquet_schema = pq.read_schema(fout) + arrow_format = _get_file_format(task.file.file_format, **{"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( @@ -511,14 +523,12 @@ 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, + arrow_table = ds.Scanner.from_fragment( + fragment=fragment, + schema=physical_schema, + filter=pyarrow_filter, columns=[col.name for col in file_project_schema.columns], - ) + ).to_table() # If there is no data, we don't have to go through the schema if len(arrow_table) > 0: diff --git a/python/pyiceberg/manifest.py b/python/pyiceberg/manifest.py index 757f3bd016da..110d0ae42f5f 100644 --- a/python/pyiceberg/manifest.py +++ b/python/pyiceberg/manifest.py @@ -14,7 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from enum import Enum +from enum import StrEnum from typing import ( Any, Dict, @@ -66,7 +66,7 @@ def __repr__(self) -> str: return f"ManifestEntryStatus.{self.name}" -class FileFormat(str, Enum): +class FileFormat(StrEnum): AVRO = "AVRO" PARQUET = "PARQUET" ORC = "ORC" @@ -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}) From 30ec1f1aa94256d497b4e09f0bba8843e3f4e121 Mon Sep 17 00:00:00 2001 From: Fokko Driesprong Date: Tue, 21 Mar 2023 12:50:13 +0100 Subject: [PATCH 02/11] Revert --- python/pyiceberg/manifest.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/pyiceberg/manifest.py b/python/pyiceberg/manifest.py index 110d0ae42f5f..942c582b51e2 100644 --- a/python/pyiceberg/manifest.py +++ b/python/pyiceberg/manifest.py @@ -14,7 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from enum import StrEnum +from enum import Enum from typing import ( Any, Dict, @@ -66,7 +66,7 @@ def __repr__(self) -> str: return f"ManifestEntryStatus.{self.name}" -class FileFormat(StrEnum): +class FileFormat(str, Enum): AVRO = "AVRO" PARQUET = "PARQUET" ORC = "ORC" From 82d663c335c151b4e40c847e018748ab8089ed4e Mon Sep 17 00:00:00 2001 From: polectron Date: Tue, 21 Mar 2023 21:06:44 +0100 Subject: [PATCH 03/11] TableScan add limit --- python/pyiceberg/table/__init__.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/python/pyiceberg/table/__init__.py b/python/pyiceberg/table/__init__.py index 89f3af847b5f..e4b70842f3a5 100644 --- a/python/pyiceberg/table/__init__.py +++ b/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,6 +231,7 @@ 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) @@ -235,6 +239,7 @@ def __init__( 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: From e20dfe996b5e03e0d77380e14c9b8f453045b88f Mon Sep 17 00:00:00 2001 From: polectron Date: Tue, 21 Mar 2023 21:08:58 +0100 Subject: [PATCH 04/11] pyarrow limit number of rows fetched from files if limit is set --- python/pyiceberg/io/pyarrow.py | 42 ++++++++++++++++++++++++++++------ 1 file changed, 35 insertions(+), 7 deletions(-) diff --git a/python/pyiceberg/io/pyarrow.py b/python/pyiceberg/io/pyarrow.py index cb40a00688b8..64a882ca1c08 100644 --- a/python/pyiceberg/io/pyarrow.py +++ b/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, @@ -503,7 +505,12 @@ 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]: + if limit and rows_counter.value >= limit: + return None + _, path = PyArrowFileIO.parse_location(task.file.file_path) arrow_format = _get_file_format(task.file.file_format, **{"pre_buffer": True, "buffer_size": ONE_MEGABYTE * 8}) with fs.open_input_file(path) as fin: @@ -529,12 +536,21 @@ def _file_to_table( if file_schema is None: raise ValueError(f"Missing Iceberg schema in Metadata for file: {path}") - arrow_table = ds.Scanner.from_fragment( + fragment_scanner = ds.Scanner.from_fragment( fragment=fragment, schema=physical_schema, filter=pyarrow_filter, columns=[col.name for col in file_project_schema.columns], - ).to_table() + ) + + 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: @@ -544,7 +560,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 @@ -580,23 +601,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: From 69ef3b89859f5c8e4f4bf74d670319e0d54b4df4 Mon Sep 17 00:00:00 2001 From: polectron Date: Tue, 21 Mar 2023 21:09:19 +0100 Subject: [PATCH 05/11] add tests for scan limit --- python/dev/provision.py | 33 ++++++++++++++++++++++++++++++++ python/tests/test_integration.py | 17 ++++++++++++++++ 2 files changed, 50 insertions(+) diff --git a/python/dev/provision.py b/python/dev/provision.py index ec84e87a23d4..bce11b9d206f 100644 --- a/python/dev/provision.py +++ b/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; diff --git a/python/tests/test_integration.py b/python/tests/test_integration.py index 3d498f0ec8c6..ac1759bad4d7 100644 --- a/python/tests/test_integration.py +++ b/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 From 56cf45a4dcf926c9b55699a11a1e2975afd1b426 Mon Sep 17 00:00:00 2001 From: polectron Date: Tue, 21 Mar 2023 22:12:47 +0100 Subject: [PATCH 06/11] python ci rebuild container if changes on python/dev/ --- .github/workflows/python-integration.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/python-integration.yml b/.github/workflows/python-integration.yml index 7c6269112612..64b6782a6e55 100644 --- a/.github/workflows/python-integration.yml +++ b/.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 \ No newline at end of file + run: docker-compose -f python/dev/docker-compose.yml logs From e6f5f22aee69b17c7b27eb06c706182febc91127 Mon Sep 17 00:00:00 2001 From: polectron Date: Wed, 22 Mar 2023 00:12:27 +0100 Subject: [PATCH 07/11] remove support for ORC --- python/pyiceberg/io/pyarrow.py | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/python/pyiceberg/io/pyarrow.py b/python/pyiceberg/io/pyarrow.py index 64a882ca1c08..0fcdeae40130 100644 --- a/python/pyiceberg/io/pyarrow.py +++ b/python/pyiceberg/io/pyarrow.py @@ -488,16 +488,6 @@ def expression_to_pyarrow(expr: BooleanExpression) -> pc.Expression: return boolean_expression_visit(expr, _ConvertToArrowExpression()) -@lru_cache -def _get_file_format(file_format: FileFormat, **kwargs: Dict[str, Any]) -> ds.FileFormat: - if file_format == FileFormat.PARQUET.value: - return ds.ParquetFileFormat(**kwargs) - elif file_format == FileFormat.ORC.value: - return ds.OrcFileFormat() - else: - raise ValueError(f"Unsupported file format: {file_format}") - - def _file_to_table( fs: FileSystem, task: FileScanTask, @@ -512,7 +502,7 @@ def _file_to_table( return None _, path = PyArrowFileIO.parse_location(task.file.file_path) - arrow_format = _get_file_format(task.file.file_format, **{"pre_buffer": True, "buffer_size": ONE_MEGABYTE * 8}) + 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 From 23602323771181114cb18a43cc8dd84ec3fe7297 Mon Sep 17 00:00:00 2001 From: polectron Date: Wed, 22 Mar 2023 00:21:34 +0100 Subject: [PATCH 08/11] remove unused imports --- python/pyiceberg/io/pyarrow.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/python/pyiceberg/io/pyarrow.py b/python/pyiceberg/io/pyarrow.py index 0fcdeae40130..ae817339689c 100644 --- a/python/pyiceberg/io/pyarrow.py +++ b/python/pyiceberg/io/pyarrow.py @@ -33,7 +33,6 @@ TYPE_CHECKING, Any, Callable, - Dict, Iterable, List, Optional, @@ -70,7 +69,6 @@ translate_column_names, ) from pyiceberg.expressions.visitors import visit as boolean_expression_visit -from pyiceberg.files import FileFormat from pyiceberg.io import ( S3_ACCESS_KEY_ID, S3_ENDPOINT, From d2be213430c35417d1008731faf59ef10f271905 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20R=C3=BCckert=20Garc=C3=ADa?= Date: Wed, 22 Mar 2023 11:15:14 +0100 Subject: [PATCH 09/11] increase sleep before running tests --- python/Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/Makefile b/python/Makefile index d04466dff8bb..0e3731447e42 100644 --- a/python/Makefile +++ b/python/Makefile @@ -39,7 +39,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 coverage run --source=pyiceberg/ -m pytest tests/ -m integration ${PYTEST_ARGS} test-adlfs: From 0073ef26088083cbb01952b5e4b21cb46c293287 Mon Sep 17 00:00:00 2001 From: polectron Date: Wed, 22 Mar 2023 20:28:37 +0100 Subject: [PATCH 10/11] update python docs to include limit in table query --- python/mkdocs/docs/api.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/mkdocs/docs/api.md b/python/mkdocs/docs/api.md index 302c6fd88a72..bb641aaeec59 100644 --- a/python/mkdocs/docs/api.md +++ b/python/mkdocs/docs/api.md @@ -281,7 +281,7 @@ Table( ## Query a table -To query a table, a table scan is needed. A table scan accepts a filter, columns and optionally a snapshot ID: +To query a table, a table scan is needed. A table scan accepts a filter, columns and optionally a limit and a snapshot ID: ```python from pyiceberg.catalog import load_catalog @@ -293,6 +293,7 @@ table = catalog.load_table("nyc.taxis") scan = table.scan( row_filter=GreaterThanOrEqual("trip_distance", 10.0), selected_fields=("VendorID", "tpep_pickup_datetime", "tpep_dropoff_datetime"), + limit=100 ) # Or filter using a string predicate From 77580c6c560d3cb0f2d51de808bfb036e36500d6 Mon Sep 17 00:00:00 2001 From: polectron Date: Thu, 23 Mar 2023 20:08:28 +0100 Subject: [PATCH 11/11] docs fix format --- python/mkdocs/docs/api.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/mkdocs/docs/api.md b/python/mkdocs/docs/api.md index 088b2a8b0e11..1cb26714ad86 100644 --- a/python/mkdocs/docs/api.md +++ b/python/mkdocs/docs/api.md @@ -293,7 +293,7 @@ table = catalog.load_table("nyc.taxis") scan = table.scan( row_filter=GreaterThanOrEqual("trip_distance", 10.0), selected_fields=("VendorID", "tpep_pickup_datetime", "tpep_dropoff_datetime"), - limit=100 + limit=100, ) # Or filter using a string predicate