kevinjqliu commented on code in PR #1443:
URL: https://github.com/apache/iceberg-python/pull/1443#discussion_r1938141310
##########
tests/io/test_pyarrow.py:
##########
@@ -1122,6 +1127,129 @@ def test_projection_concat_files(schema_int: Schema,
file_int: str) -> None:
assert repr(result_table.schema) == "id: int32"
+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)
+
+ schema = Schema(
+ NestedField(1, "other_field", StringType(), required=False),
NestedField(2, "partition_id", IntegerType(), required=False)
+ )
+
+ partition_spec = PartitionSpec(
+ PartitionField(2, 1000, IdentityTransform(), "partition_id"),
+ )
+
+ table = catalog.create_table(
+ "default.test_projection_partition",
+ schema=schema,
+ partition_spec=partition_spec,
+ properties={TableProperties.DEFAULT_NAME_MAPPING:
create_mapping_from_schema(schema).model_dump_json()},
+ )
+
+ 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)
+
+ 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()),
+ )
+
+ unpartitioned_file = DataFile(
+ 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,
+ spec_id=table.metadata.default_spec_id,
+ equality_ids=None,
+ key_metadata=None,
+ **statistics.to_serialized_dict(),
+ )
+
+ with table.transaction() as transaction:
+ with transaction.update_snapshot().overwrite() as update:
+ update.append_data_file(unpartitioned_file)
+
+ assert (
+ str(table.scan().to_arrow())
+ == """pyarrow.Table
+other_field: large_string
+partition_id: int64
+----
+other_field: [["foo"]]
+partition_id: [[1]]"""
+ )
+
+
+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)
+ schema = Schema(
+ NestedField(1, "other_field", StringType(), required=False),
NestedField(2, "partition_id", IntegerType(), required=False)
+ )
+
+ partition_spec = PartitionSpec(
+ PartitionField(2, 1000, IdentityTransform(), "void_partition_id"),
Review Comment:
nit: avoid using `void` since its a type of transform
https://iceberg.apache.org/spec/#partition-transforms
##########
tests/io/test_pyarrow.py:
##########
@@ -1122,6 +1127,129 @@ def test_projection_concat_files(schema_int: Schema,
file_int: str) -> None:
assert repr(result_table.schema) == "id: int32"
+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)
+
+ schema = Schema(
+ NestedField(1, "other_field", StringType(), required=False),
NestedField(2, "partition_id", IntegerType(), required=False)
+ )
+
+ partition_spec = PartitionSpec(
+ PartitionField(2, 1000, IdentityTransform(), "partition_id"),
+ )
+
+ table = catalog.create_table(
+ "default.test_projection_partition",
+ schema=schema,
+ partition_spec=partition_spec,
+ properties={TableProperties.DEFAULT_NAME_MAPPING:
create_mapping_from_schema(schema).model_dump_json()},
+ )
+
+ 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)
+
+ 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()),
+ )
+
+ unpartitioned_file = DataFile(
+ 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,
+ spec_id=table.metadata.default_spec_id,
+ equality_ids=None,
+ key_metadata=None,
+ **statistics.to_serialized_dict(),
+ )
+
+ with table.transaction() as transaction:
+ with transaction.update_snapshot().overwrite() as update:
+ update.append_data_file(unpartitioned_file)
+
+ assert (
+ str(table.scan().to_arrow())
+ == """pyarrow.Table
+other_field: large_string
+partition_id: int64
+----
+other_field: [["foo"]]
+partition_id: [[1]]"""
+ )
+
+
+def test_identity_transform_columns_projection(tmp_path: str, catalog:
InMemoryCatalog) -> None:
Review Comment:
i was thinking we can test something like 3 fields where 2 are identity
partitions.
to check the scenario for multi-level hive partition, for example
`s3://foo/year=2025/month=06/blah.parquet`
##########
tests/io/test_pyarrow.py:
##########
@@ -1122,6 +1127,129 @@ def test_projection_concat_files(schema_int: Schema,
file_int: str) -> None:
assert repr(result_table.schema) == "id: int32"
+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)
+
+ schema = Schema(
+ NestedField(1, "other_field", StringType(), required=False),
NestedField(2, "partition_id", IntegerType(), required=False)
+ )
+
+ partition_spec = PartitionSpec(
+ PartitionField(2, 1000, IdentityTransform(), "partition_id"),
+ )
+
+ table = catalog.create_table(
+ "default.test_projection_partition",
+ schema=schema,
+ partition_spec=partition_spec,
+ properties={TableProperties.DEFAULT_NAME_MAPPING:
create_mapping_from_schema(schema).model_dump_json()},
+ )
+
+ 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)
+
+ 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()),
+ )
+
+ unpartitioned_file = DataFile(
+ 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,
+ spec_id=table.metadata.default_spec_id,
+ equality_ids=None,
+ key_metadata=None,
+ **statistics.to_serialized_dict(),
+ )
+
+ with table.transaction() as transaction:
+ with transaction.update_snapshot().overwrite() as update:
+ update.append_data_file(unpartitioned_file)
+
+ assert (
+ str(table.scan().to_arrow())
+ == """pyarrow.Table
+other_field: large_string
+partition_id: int64
+----
+other_field: [["foo"]]
+partition_id: [[1]]"""
+ )
+
+
+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)
+ schema = Schema(
+ NestedField(1, "other_field", StringType(), required=False),
NestedField(2, "partition_id", IntegerType(), required=False)
+ )
+
+ partition_spec = PartitionSpec(
+ PartitionField(2, 1000, IdentityTransform(), "void_partition_id"),
+ PartitionField(2, 1001, IdentityTransform(), "partition_id"),
+ )
+
+ table = catalog.create_table(
+ "default.test_projection_partitions",
+ schema=schema,
+ partition_spec=partition_spec,
+ properties={TableProperties.DEFAULT_NAME_MAPPING:
create_mapping_from_schema(schema).model_dump_json()},
+ )
+
+ 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)
+
+ 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()),
+ )
+
+ unpartitioned_file = DataFile(
+ content=DataFileContent.DATA,
+ file_path=file_loc,
+ file_format=FileFormat.PARQUET,
+ # 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,
+ equality_ids=None,
+ key_metadata=None,
+ **statistics.to_serialized_dict(),
+ )
+
+ with table.transaction() as transaction:
+ with transaction.update_snapshot().overwrite() as update:
+ update.append_data_file(unpartitioned_file)
+
+ assert (
+ str(table.scan().to_arrow())
+ == """pyarrow.Table
+other_field: large_string
+partition_id: int64
+----
+other_field: [["foo"]]
+partition_id: [[1]]"""
Review Comment:
shouldnt this project `void_partition_id=12` as well?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscr...@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscr...@iceberg.apache.org
For additional commands, e-mail: issues-h...@iceberg.apache.org
