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:32 UTC
(iceberg-python) 08/10: Check for name before injecting, fix lint issues
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 f756bb50349595a5e6c55fe4f62f7f4f8baef6f3
Author: Gabriel Igliozzi <ga...@gmai.com>
AuthorDate: Sat Jan 25 16:45:02 2025 -0700
Check for name before injecting, fix lint issues
---
pyiceberg/io/pyarrow.py | 44 ++++++++++++++++++++++++--------------------
tests/io/test_pyarrow.py | 12 +++++++-----
2 files changed, 31 insertions(+), 25 deletions(-)
diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py
index 09008f4a..f0365e8f 100644
--- a/pyiceberg/io/pyarrow.py
+++ b/pyiceberg/io/pyarrow.py
@@ -125,6 +125,7 @@ from pyiceberg.manifest import (
)
from pyiceberg.partitioning import PartitionField, PartitionFieldValue, PartitionKey, PartitionSpec, partition_record_value
from pyiceberg.schema import (
+ Accessor,
PartnerAccessor,
PreOrderSchemaVisitor,
Schema,
@@ -1300,16 +1301,24 @@ class _ConvertToIcebergWithoutIDs(_ConvertToIceberg):
def _get_column_projection_values(
- file: DataFile,
- projected_schema: Schema,
- project_schema_diff: Set[int],
- partition_spec: PartitionSpec,
-) -> Dict[str, object]:
+ file: DataFile, projected_schema: Schema, partition_spec: Optional[PartitionSpec], file_project_field_ids: Set[int]
+) -> Tuple[bool, Dict[str, Any]]:
"""Apply Column Projection rules to File Schema."""
+ project_schema_diff = projected_schema.field_ids.difference(file_project_field_ids)
+ should_project_columns = len(project_schema_diff) > 0
projected_missing_fields: Dict[str, Any] = {}
- partition_schema = partition_spec.partition_type(projected_schema)
- accessors = build_position_accessors(partition_schema)
+ if not should_project_columns:
+ return False, {}
+
+ partition_schema: StructType
+ accessors: Dict[int, Accessor]
+
+ if partition_spec is not None:
+ partition_schema = partition_spec.partition_type(projected_schema)
+ accessors = build_position_accessors(partition_schema)
+ else:
+ return False, {}
for field_id in project_schema_diff:
for partition_field in partition_spec.fields_by_source_id(field_id):
@@ -1322,7 +1331,7 @@ def _get_column_projection_values(
if partition_value := accesor.get(file.partition):
projected_missing_fields[partition_field.name] = partition_value
- return projected_missing_fields
+ return True, projected_missing_fields
def _task_to_record_batches(
@@ -1354,19 +1363,12 @@ def _task_to_record_batches(
bound_file_filter = bind(file_schema, translated_row_filter, case_sensitive=case_sensitive)
pyarrow_filter = expression_to_pyarrow(bound_file_filter)
- # Apply column projection rules for missing partitions and default values
+ # Apply column projection rules
# https://iceberg.apache.org/spec/#column-projection
file_project_schema = prune_columns(file_schema, projected_field_ids, select_full_types=False)
-
- project_schema_diff = projected_field_ids.difference(file_project_schema.field_ids)
- should_project_columns = len(project_schema_diff) > 0
-
- projected_missing_fields = {}
-
- if should_project_columns and partition_spec is not None:
- projected_missing_fields = _get_column_projection_values(
- task.file, projected_schema, project_schema_diff, partition_spec
- )
+ should_project_columns, projected_missing_fields = _get_column_projection_values(
+ task.file, projected_schema, partition_spec, file_project_schema.field_ids
+ )
fragment_scanner = ds.Scanner.from_fragment(
fragment=fragment,
@@ -1417,7 +1419,9 @@ def _task_to_record_batches(
# Inject projected column values if available
if should_project_columns:
for name, value in projected_missing_fields.items():
- result_batch = result_batch.set_column(result_batch.schema.get_field_index(name), name, [value])
+ index = result_batch.schema.get_field_index(name)
+ if index != -1:
+ result_batch = result_batch.set_column(index, name, [value])
yield result_batch
diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py
index b142ebc0..b65750e5 100644
--- a/tests/io/test_pyarrow.py
+++ b/tests/io/test_pyarrow.py
@@ -81,7 +81,7 @@ from pyiceberg.schema import Schema, make_compatible_name, visit
from pyiceberg.table import FileScanTask, TableProperties
from pyiceberg.table.metadata import TableMetadataV2
from pyiceberg.table.name_mapping import create_mapping_from_schema
-from pyiceberg.transforms import IdentityTransform, VoidTransform
+from pyiceberg.transforms import IdentityTransform
from pyiceberg.typedef import UTF8, Properties, Record
from pyiceberg.types import (
BinaryType,
@@ -1132,7 +1132,7 @@ def test_projection_concat_files(schema_int: Schema, file_int: str) -> None:
assert repr(result_table.schema) == "id: int32"
-def test_projection_single_partition_value(tmp_path: str, catalog: InMemoryCatalog) -> None:
+def test_identity_transform_column_projection(tmp_path: str, catalog: InMemoryCatalog) -> None:
# Test by adding a non-partitioned data file to a partitioned table, verifying partition value projection from manifest metadata.
# TODO: Update to use a data file created by writing data to an unpartitioned table once add_files supports field IDs.
# (context: https://github.com/apache/iceberg-python/pull/1443#discussion_r1901374875)
@@ -1166,6 +1166,7 @@ def test_projection_single_partition_value(tmp_path: str, catalog: InMemoryCatal
content=DataFileContent.DATA,
file_path=file_loc,
file_format=FileFormat.PARQUET,
+ # projected value
partition=Record(partition_id=1),
file_size_in_bytes=os.path.getsize(file_loc),
sort_order_id=None,
@@ -1190,7 +1191,7 @@ partition_id: [[1]]"""
)
-def test_projection_multiple_partition_values(tmp_path: str, catalog: InMemoryCatalog) -> None:
+def test_identity_transform_columns_projection(tmp_path: str, catalog: InMemoryCatalog) -> None:
# Test by adding a non-partitioned data file to a multi-partitioned table, verifying partition value projection from manifest metadata.
# TODO: Update to use a data file created by writing data to an unpartitioned table once add_files supports field IDs.
# (context: https://github.com/apache/iceberg-python/pull/1443#discussion_r1901374875)
@@ -1199,7 +1200,7 @@ def test_projection_multiple_partition_values(tmp_path: str, catalog: InMemoryCa
)
partition_spec = PartitionSpec(
- PartitionField(2, 1000, VoidTransform(), "void_partition_id"),
+ PartitionField(2, 1000, IdentityTransform(), "void_partition_id"),
PartitionField(2, 1001, IdentityTransform(), "partition_id"),
)
@@ -1224,7 +1225,8 @@ def test_projection_multiple_partition_values(tmp_path: str, catalog: InMemoryCa
content=DataFileContent.DATA,
file_path=file_loc,
file_format=FileFormat.PARQUET,
- partition=Record(void_partition_id=None, partition_id=1),
+ # projected value
+ partition=Record(void_partition_id=12, partition_id=1),
file_size_in_bytes=os.path.getsize(file_loc),
sort_order_id=None,
spec_id=table.metadata.default_spec_id,