[iceberg] branch master updated: Python: Add partition_type to partition_spec (#5929)

[bl...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

blue pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iceberg.git


The following commit(s) were added to refs/heads/master by this push:
     new bc62cf8dab Python: Add partition_type to partition_spec (#5929)
bc62cf8dab is described below

commit bc62cf8dab42df2b113f9e56f846425daf22123b
Author: Fokko Driesprong <fo...@apache.org>
AuthorDate: Fri Oct 7 22:53:43 2022 +0200

    Python: Add partition_type to partition_spec (#5929)
---
 python/pyiceberg/table/partitioning.py  | 27 +++++++++++++++++++++++++--
 python/tests/table/test_partitioning.py | 20 ++++++++++++++++++++
 2 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/python/pyiceberg/table/partitioning.py b/python/pyiceberg/table/partitioning.py
index 8dcbf68454..1445a3f327 100644
--- a/python/pyiceberg/table/partitioning.py
+++ b/python/pyiceberg/table/partitioning.py
@@ -27,6 +27,7 @@ from pydantic import Field
 
 from pyiceberg.schema import Schema
 from pyiceberg.transforms import Transform
+from pyiceberg.types import NestedField, StructType
 from pyiceberg.utils.iceberg_base_model import IcebergBaseModel
 
 INITIAL_PARTITION_SPEC_ID = 0
@@ -77,9 +78,8 @@ class PartitionSpec(IcebergBaseModel):
     PartitionSpec captures the transformation from table data to partition values
 
     Attributes:
-        schema(Schema): the schema of data table
         spec_id(int): any change to PartitionSpec will produce a new specId
-        fields(List[PartitionField): list of partition fields to produce partition values
+        fields(Tuple[PartitionField): list of partition fields to produce partition values
     """
 
     spec_id: int = Field(alias="spec-id", default=INITIAL_PARTITION_SPEC_ID)
@@ -158,6 +158,29 @@ class PartitionSpec(IcebergBaseModel):
             for this_field, that_field in zip(self.fields, other.fields)
         )
 
+    def partition_type(self, schema: Schema) -> StructType:
+        """Produces a struct of the PartitionSpec
+
+        The partition fields should be optional:
+
+        - All partition transforms are required to produce null if the input value is null, so it can
+          happen when the source column is optional
+        - Partition fields may be added later, in which case not all files would have the result field,
+          and it may be null.
+
+        There is a case where we can guarantee that a partition field in the first and only partition spec
+        that uses a required source column will never be null, but it doesn't seem worth tracking this case.
+
+        :param schema: The schema to bind to
+        :return: A StructType that represents the PartitionSpec, with a NestedField for each PartitionField
+        """
+        nested_fields = []
+        for field in self.fields:
+            source_type = schema.find_type(field.source_id)
+            result_type = field.transform.result_type(source_type)
+            nested_fields.append(NestedField(field.field_id, field.name, result_type, required=False))
+        return StructType(*nested_fields)
+
 
 UNPARTITIONED_PARTITION_SPEC = PartitionSpec(spec_id=0)
 
diff --git a/python/tests/table/test_partitioning.py b/python/tests/table/test_partitioning.py
index 9aa0f73891..b3cbf30f8f 100644
--- a/python/tests/table/test_partitioning.py
+++ b/python/tests/table/test_partitioning.py
@@ -14,8 +14,15 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+from pyiceberg.schema import Schema
 from pyiceberg.table.partitioning import UNPARTITIONED_PARTITION_SPEC, PartitionField, PartitionSpec
 from pyiceberg.transforms import BucketTransform, TruncateTransform
+from pyiceberg.types import (
+    IntegerType,
+    NestedField,
+    StringType,
+    StructType,
+)
 
 
 def test_partition_field_init():
@@ -102,3 +109,16 @@ def test_deserialize_partition_spec():
             PartitionField(source_id=2, field_id=1001, transform=BucketTransform(num_buckets=25), name="int_bucket"),
         ),
     )
+
+
+def test_partition_type(table_schema_simple: Schema) -> None:
+    spec = PartitionSpec(
+        PartitionField(source_id=1, field_id=1000, transform=TruncateTransform(width=19), name="str_truncate"),
+        PartitionField(source_id=2, field_id=1001, transform=BucketTransform(num_buckets=25), name="int_bucket"),
+        spec_id=3,
+    )
+
+    assert spec.partition_type(table_schema_simple) == StructType(
+        NestedField(field_id=1000, name="str_truncate", field_type=StringType(), required=False),
+        NestedField(field_id=1001, name="int_bucket", field_type=IntegerType(), required=False),
+    )