This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-python.git
The following commit(s) were added to refs/heads/main by this push:
new bd98ff1de PyArrow: Support view types when writing to a partitioned
table (#4121)
bd98ff1de is described below
commit bd98ff1de88febaca645d54004a83de37a2d5ce9
Author: Dhruv <[email protected]>
AuthorDate: Sat Oct 24 18:17:38 2026 +0000
PyArrow: Support view types when writing to a partitioned table (#4121)
* PyArrow: Support view types when writing to a partitioned table
Appending an Arrow table with string_view or binary_view columns, which
pa.table() produces from a Polars DataFrame, worked for unpartitioned
tables but failed for partitioned ones: PyArrow has no group by, filter
or take kernels for view types, and pyiceberg-core's bucket and truncate
transforms reject them. Cast view types, including nested ones, to their
large counterparts before splitting the table into partitions.
* Use PyArrowSchemaVisitor to replace view types
* Test appending Polars categorical columns to a partitioned table
Polars exports Categorical as dictionary<string_view> and List(Categorical)
as large_list<dictionary<string_view>>, which exercise the _keep_dictionary
paths in _ConvertViewTypes.
Co-authored-by: Copilot App <[email protected]>
---------
Co-authored-by: Kevin Liu <[email protected]>
Co-authored-by: Copilot App <[email protected]>
---
pyiceberg/io/pyarrow.py | 50 ++++++++++++++++++++++++++++++
tests/io/test_pyarrow.py | 79 +++++++++++++++++++++++++++++++++++++++++++++++-
2 files changed, 128 insertions(+), 1 deletion(-)
diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py
index 874646b71..843dc0449 100644
--- a/pyiceberg/io/pyarrow.py
+++ b/pyiceberg/io/pyarrow.py
@@ -1607,6 +1607,51 @@ class
_ConvertToSmallTypes(PyArrowSchemaVisitor[IcebergType | pa.Schema]):
return primitive
+class _ConvertViewTypes(PyArrowSchemaVisitor[pa.Schema | pa.DataType]):
+ """Replace string_view and binary_view, including nested occurrences, with
large_string and large_binary.
+
+ Everything else is kept as is: list flavour, field names, nullability,
metadata and dictionary encoding.
+ """
+
+ def schema(self, schema: pa.Schema, struct_result: pa.StructType) ->
pa.Schema:
+ return pa.schema(struct_result, metadata=schema.metadata)
+
+ def struct(self, struct: pa.StructType, field_results:
builtins.list[pa.Field]) -> pa.StructType:
+ return pa.struct(field_results)
+
+ def field(self, field: pa.Field, field_result: pa.DataType) -> pa.Field:
+ return field.with_type(_keep_dictionary(field.type, field_result))
+
+ def list(self, list_type: pa.ListType | pa.LargeListType |
pa.FixedSizeListType, element_result: pa.DataType) -> pa.DataType:
+ value_field =
list_type.value_field.with_type(_keep_dictionary(list_type.value_type,
element_result))
+ if isinstance(list_type, pa.FixedSizeListType):
+ return pa.list_(value_field, list_type.list_size)
+ elif isinstance(list_type, pa.LargeListType):
+ return pa.large_list(value_field)
+ return pa.list_(value_field)
+
+ def map(self, map_type: pa.MapType, key_result: pa.DataType, value_result:
pa.DataType) -> pa.DataType:
+ return pa.map_(
+ map_type.key_field.with_type(_keep_dictionary(map_type.key_type,
key_result)),
+ map_type.item_field.with_type(_keep_dictionary(map_type.item_type,
value_result)),
+ keys_sorted=map_type.keys_sorted,
+ )
+
+ def primitive(self, primitive: pa.DataType) -> pa.DataType:
+ if pa.types.is_string_view(primitive):
+ return pa.large_string()
+ elif pa.types.is_binary_view(primitive):
+ return pa.large_binary()
+ return primitive
+
+
+def _keep_dictionary(original: pa.DataType, result: pa.DataType) ->
pa.DataType:
+ """Re-apply dictionary encoding that visit_pyarrow strips before visiting
the value type."""
+ if isinstance(original, pa.DictionaryType):
+ return pa.dictionary(original.index_type, result,
ordered=original.ordered)
+ return result
+
+
class _ConvertToIcebergWithoutIDs(_ConvertToIceberg):
"""
Converts PyArrowSchema to Iceberg Schema with all -1 ids.
@@ -3116,6 +3161,11 @@ def _determine_partitions(spec: PartitionSpec, schema:
Schema, arrow_table: pa.T
- Then we produce a set of partitions by filtering on each of the
combinations
- We combine the chunks to create a copy to avoid GIL congestion on the
original table
"""
+ # PyArrow has no group by, filter or take kernels for view types, so
replace them with their large counterparts
+ arrow_schema = visit_pyarrow(arrow_table.schema, _ConvertViewTypes())
+ if arrow_schema != arrow_table.schema:
+ arrow_table = arrow_table.cast(arrow_schema)
+
# Assign unique names to columns where the partition transform has been
applied
# to avoid conflicts
partition_fields = [f"_partition_{field.name}" for field in spec.fields]
diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py
index e97b7813d..e6a085b8d 100644
--- a/tests/io/test_pyarrow.py
+++ b/tests/io/test_pyarrow.py
@@ -106,7 +106,7 @@ from pyiceberg.table.deletion_vector import (
from pyiceberg.table.metadata import TableMetadataV2, TableMetadataV3
from pyiceberg.table.name_mapping import create_mapping_from_schema
from pyiceberg.table.puffin import MAGIC_BYTES, PuffinFile
-from pyiceberg.transforms import HourTransform, IdentityTransform
+from pyiceberg.transforms import BucketTransform, HourTransform,
IdentityTransform, Transform, TruncateTransform
from pyiceberg.typedef import UTF8, Properties, Record, TableVersion
from pyiceberg.types import (
BinaryType,
@@ -3194,6 +3194,83 @@ def test_determine_partitions_identity_nan() -> None:
assert rows_by_partition == {"1.0": [1], "nan": [2, 4], "None": [3]}
[email protected]("transform", [IdentityTransform(),
BucketTransform(4), TruncateTransform(1)])
+def test_determine_partitions_with_view_types(transform: Transform[Any, Any])
-> None:
+ schema = Schema(
+ NestedField(id=1, name="region", field_type=StringType(),
required=False),
+ NestedField(id=2, name="payload", field_type=BinaryType(),
required=False),
+ NestedField(
+ id=3,
+ name="info",
+ field_type=StructType(
+ NestedField(id=4, name="tags", field_type=ListType(5,
StringType(), element_required=False), required=False),
+ NestedField(id=6, name="attrs", field_type=MapType(7,
StringType(), 8, BinaryType()), required=False),
+ ),
+ required=False,
+ ),
+ )
+ spec = PartitionSpec(PartitionField(source_id=1, field_id=1000,
transform=transform, name="region_part"))
+ arrow_schema = pa.schema(
+ [
+ pa.field("region", pa.string_view()),
+ pa.field("payload", pa.binary_view()),
+ pa.field(
+ "info",
+ pa.struct(
+ [
+ pa.field("tags", pa.list_(pa.string_view())),
+ pa.field("attrs", pa.map_(pa.string_view(),
pa.binary_view())),
+ ]
+ ),
+ ),
+ ]
+ )
+ rows = [
+ {"region": "us", "payload": b"a", "info": {"tags": ["x"], "attrs":
[("k", b"v")]}},
+ {"region": "eu", "payload": None, "info": None},
+ {"region": "us", "payload": b"b", "info": {"tags": None, "attrs":
None}},
+ {"region": None, "payload": b"c", "info": {"tags": ["y", None],
"attrs": []}},
+ ]
+ arrow_table = pa.Table.from_pylist(rows, schema=arrow_schema)
+
+ partitions = list(_determine_partitions(spec, schema, arrow_table))
+
+ assert {p.partition_key.partition[0] for p in partitions} == {
+ transform.transform(StringType())(region) for region in ("us", "eu",
None)
+ }
+ written = pa.concat_tables([p.arrow_table_partition for p in partitions])
+ assert sorted(written.to_pylist(), key=str) == sorted(rows, key=str)
+
+
+def test_append_view_types_to_partitioned_table(catalog: InMemoryCatalog) ->
None:
+ schema = Schema(
+ NestedField(id=1, name="id", field_type=LongType(), required=False),
+ NestedField(id=2, name="region", field_type=StringType(),
required=False),
+ NestedField(id=3, name="payload", field_type=BinaryType(),
required=False),
+ NestedField(id=4, name="category", field_type=StringType(),
required=False),
+ NestedField(id=5, name="labels", field_type=ListType(6, StringType(),
element_required=False), required=False),
+ )
+ spec = PartitionSpec(PartitionField(source_id=2, field_id=1000,
transform=IdentityTransform(), name="region"))
+ catalog.create_namespace("default")
+ table = catalog.create_table("default.test_append_view_types",
schema=schema, partition_spec=spec)
+
+ categories = pa.array(["a", "b", None],
type=pa.string_view()).dictionary_encode()
+ arrow_table = pa.table(
+ {
+ "id": pa.array([1, 2, 3], type=pa.int64()),
+ "region": pa.array(["us", "eu", "us"], type=pa.string_view()),
+ "payload": pa.array([b"a", None, b"c"], type=pa.binary_view()),
+ # Polars exports Categorical as dictionary<string_view> and
List(Categorical) as large_list<dictionary<string_view>>
+ "category": categories,
+ "labels": pa.LargeListArray.from_arrays([0, 1, 1, 3], categories),
+ }
+ )
+ table.append(arrow_table)
+
+ assert len(table.inspect.partitions()) == 2
+ assert sorted(table.scan().to_arrow().to_pylist(), key=lambda row:
row["id"]) == arrow_table.to_pylist()
+
+
def test_inspect_partition_for_nested_field(catalog: InMemoryCatalog) -> None:
schema = Schema(
NestedField(id=1, name="foo", field_type=StringType(), required=True),
