You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iceberg.apache.org by ke...@apache.org on 2025/02/01 05:18:30 UTC
(iceberg-python) 06/10: Removed file_schema check and initial-default logic, separated projection logic to helper method, changed test to use high-level table scan
This is an automated email from the ASF dual-hosted git repository.
kevinjqliu pushed a commit to branch specPartitionIdentity
in repository https://gitbox.apache.org/repos/asf/iceberg-python.git
commit 2bf1ecfcfea975014020aef12063ec1de9325590
Author: Gabriel Igliozzi <gi...@netflix.com>
AuthorDate: Thu Dec 26 12:31:37 2024 -0500
Removed file_schema check and initial-default logic, separated projection logic to helper method, changed test to use high-level table scan
---
pyiceberg/io/pyarrow.py | 36 +++++++++------
tests/io/test_pyarrow.py | 115 +++++++++++++++--------------------------------
2 files changed, 59 insertions(+), 92 deletions(-)
diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py
index 63dc99da..b6e4dd50 100644
--- a/pyiceberg/io/pyarrow.py
+++ b/pyiceberg/io/pyarrow.py
@@ -1298,6 +1298,25 @@ class _ConvertToIcebergWithoutIDs(_ConvertToIceberg):
return -1
+def _get_column_projection_values(
+ file: DataFile,
+ projected_schema: Schema,
+ projected_field_ids: Set[int],
+ file_project_schema: Schema,
+ partition_spec: Optional[PartitionSpec] = None,
+) -> Dict[str, object]:
+ """Apply Column Projection rules to File Schema."""
+ projected_missing_fields = {}
+
+ for field_id in projected_field_ids.difference(file_project_schema.field_ids):
+ if partition_spec is not None:
+ for partition_field in partition_spec.fields_by_source_id(field_id):
+ if isinstance(partition_field.transform, IdentityTransform) and partition_field.name in file.partition.__dict__:
+ projected_missing_fields[partition_field.name] = file.partition.__dict__[partition_field.name]
+
+ return projected_missing_fields
+
+
def _task_to_record_batches(
fs: FileSystem,
task: FileScanTask,
@@ -1321,9 +1340,6 @@ def _task_to_record_batches(
# the table format version.
file_schema = pyarrow_to_schema(physical_schema, name_mapping, downcast_ns_timestamp_to_us=True)
- if file_schema is None:
- raise ValueError(f"Missing Iceberg schema in Metadata for file: {path}")
-
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)
@@ -1333,16 +1349,10 @@ def _task_to_record_batches(
# Apply column projection rules for missing partitions and default values
# https://iceberg.apache.org/spec/#column-projection
file_project_schema = prune_columns(file_schema, projected_field_ids, select_full_types=False)
- projected_missing_fields = {}
-
- for field_id in projected_field_ids.difference(file_project_schema.field_ids):
- if nested_field := projected_schema.find_field(field_id):
- if nested_field.initial_default is not None:
- projected_missing_fields[nested_field.name] = nested_field.initial_default
- if partition_spec is not None:
- for partition_field in partition_spec.fields_by_source_id(field_id):
- if isinstance(partition_field.transform, IdentityTransform) and task.file.partition is not None:
- projected_missing_fields[partition_field.name] = task.file.partition[0]
+
+ projected_missing_fields = _get_column_projection_values(
+ task.file, projected_schema, projected_field_ids, file_project_schema, partition_spec
+ )
fragment_scanner = ds.Scanner.from_fragment(
fragment=fragment,
diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py
index 48646759..972dad71 100644
--- a/tests/io/test_pyarrow.py
+++ b/tests/io/test_pyarrow.py
@@ -69,7 +69,10 @@ from pyiceberg.io.pyarrow import (
_read_deletes,
_to_requested_schema,
bin_pack_arrow_table,
+ compute_statistics_plan,
+ data_file_statistics_from_parquet_metadata,
expression_to_pyarrow,
+ parquet_path_to_id_mapping,
schema_to_pyarrow,
)
from pyiceberg.manifest import DataFile, DataFileContent, FileFormat
@@ -100,6 +103,7 @@ from pyiceberg.types import (
TimestamptzType,
TimeType,
)
+from tests.catalog.test_base import InMemoryCatalog
from tests.conftest import UNIFIED_AWS_SESSION_PROPERTIES
@@ -1128,110 +1132,63 @@ def test_projection_concat_files(schema_int: Schema, file_int: str) -> None:
assert repr(result_table.schema) == "id: int32"
-def test_projection_partition_inference(tmp_path: str) -> None:
+def test_projection_partition_inference(tmp_path: str, catalog: InMemoryCatalog) -> None:
schema = Schema(
- NestedField(1, "partition_field", IntegerType(), required=False),
- NestedField(2, "other_field", StringType(), required=False),
+ NestedField(1, "other_field", StringType(), required=False), NestedField(2, "partition_id", IntegerType(), required=False)
)
- partition_spec = PartitionSpec(PartitionField(1, 1000, IdentityTransform(), "partition_field"))
+ partition_spec = PartitionSpec(PartitionField(2, 1000, IdentityTransform(), "partition_id"))
- table = TableMetadataV2(
- location="file://a/b/c.json",
- last_column_id=2,
- format_version=2,
- current_schema_id=0,
- schemas=[schema],
- partition_specs=[partition_spec],
+ table = catalog.create_table(
+ "default.test_projection_partition_inference",
+ schema=schema,
+ partition_spec=partition_spec,
properties={TableProperties.DEFAULT_NAME_MAPPING: create_mapping_from_schema(schema).model_dump_json()},
)
- pa_schema = pa.schema([pa.field("other_field", pa.string())])
- pa_table = pa.table({"other_field": ["x"]}, schema=pa_schema)
- pq.write_table(pa_table, f"{tmp_path}/datafile.parquet")
+ file_data = pa.array(["foo"], type=pa.string())
+ file_loc = f"{tmp_path}/test.parquet"
+ pq.write_table(pa.table([file_data], names=["other_field"]), file_loc)
- data_file = DataFile(
- content=DataFileContent.DATA,
- file_path=f"{tmp_path}/datafile.parquet",
- file_format=FileFormat.PARQUET,
- partition=Record(partition_id=123456),
- file_size_in_bytes=os.path.getsize(f"{tmp_path}/datafile.parquet"),
- sort_order_id=None,
- spec_id=0,
- equality_ids=None,
- key_metadata=None,
- )
-
- table_result_scan = ArrowScan(
- table_metadata=table,
- io=load_file_io(),
- projected_schema=schema,
- row_filter=AlwaysTrue(),
- ).to_table(tasks=[FileScanTask(data_file=data_file)])
-
- assert (
- str(table_result_scan)
- == """pyarrow.Table
-partition_field: int64
-other_field: large_string
-----
-partition_field: [[123456]]
-other_field: [["x"]]"""
+ statistics = data_file_statistics_from_parquet_metadata(
+ parquet_metadata=pq.read_metadata(file_loc),
+ stats_columns=compute_statistics_plan(table.schema(), table.metadata.properties),
+ parquet_column_mapping=parquet_path_to_id_mapping(table.schema()),
)
-
-def test_projection_initial_default_inference(tmp_path: str) -> None:
- schema = Schema(
- NestedField(1, "other_field", StringType(), required=False),
- NestedField(2, "other_field_1", StringType(), required=False, initial_default="foo"),
- )
-
- table = TableMetadataV2(
- location="file://a/b/c.json",
- last_column_id=2,
- format_version=2,
- current_schema_id=0,
- schemas=[schema],
- partition_specs=[PartitionSpec()],
- properties={TableProperties.DEFAULT_NAME_MAPPING: create_mapping_from_schema(schema).model_dump_json()},
- )
-
- pa_schema = pa.schema([pa.field("other_field", pa.string())])
- pa_table = pa.table({"other_field": ["x"]}, schema=pa_schema)
- pq.write_table(pa_table, f"{tmp_path}/datafile.parquet")
-
- data_file = DataFile(
+ unpartitioned_file = DataFile(
content=DataFileContent.DATA,
- file_path=f"{tmp_path}/datafile.parquet",
+ file_path=file_loc,
file_format=FileFormat.PARQUET,
- partition=Record(),
- file_size_in_bytes=os.path.getsize(f"{tmp_path}/datafile.parquet"),
+ partition=Record(partition_id=1),
+ file_size_in_bytes=os.path.getsize(file_loc),
sort_order_id=None,
- spec_id=0,
+ spec_id=table.metadata.default_spec_id,
equality_ids=None,
key_metadata=None,
+ **statistics.to_serialized_dict(),
)
- table_result_scan = ArrowScan(
- table_metadata=table,
- io=load_file_io(),
- projected_schema=schema,
- row_filter=AlwaysTrue(),
- ).to_table(tasks=[FileScanTask(data_file=data_file)])
-
- print(str(table_result_scan))
+ with table.transaction() as transaction:
+ with transaction.update_snapshot().overwrite() as update:
+ update.append_data_file(unpartitioned_file)
assert (
- str(table_result_scan)
+ str(table.scan().to_arrow())
== """pyarrow.Table
other_field: large_string
-other_field_1: string
+partition_id: int64
----
-other_field: [["x"]]
-other_field_1: [["foo"]]"""
+other_field: [["foo"]]
+partition_id: [[1]]"""
)
+@pytest.fixture
+def catalog() -> InMemoryCatalog:
+ return InMemoryCatalog("test.in_memory.catalog", **{"test.key": "test.value"})
+
+
def test_projection_filter(schema_int: Schema, file_int: str) -> None:
result_table = project(schema_int, [file_int], GreaterThan("id", 4))
assert len(result_table.columns[0]) == 0