Skip to content

Commit 75c2fe0

Browse files
committed
Use PyArrowSchemaVisitor to replace view types
1 parent 877667e commit 75c2fe0

1 file changed

Lines changed: 46 additions & 27 deletions

File tree

‎pyiceberg/io/pyarrow.py‎

Lines changed: 46 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1611,6 +1611,51 @@ def primitive(self, primitive: pa.DataType) -> pa.DataType:
16111611
return primitive
16121612

16131613

1614+
class _ConvertViewTypes(PyArrowSchemaVisitor[pa.Schema | pa.DataType]):
1615+
"""Replace string_view and binary_view, including nested occurrences, with large_string and large_binary.
1616+
1617+
Everything else is kept as is: list flavour, field names, nullability, metadata and dictionary encoding.
1618+
"""
1619+
1620+
def schema(self, schema: pa.Schema, struct_result: pa.StructType) -> pa.Schema:
1621+
return pa.schema(struct_result, metadata=schema.metadata)
1622+
1623+
def struct(self, struct: pa.StructType, field_results: builtins.list[pa.Field]) -> pa.StructType:
1624+
return pa.struct(field_results)
1625+
1626+
def field(self, field: pa.Field, field_result: pa.DataType) -> pa.Field:
1627+
return field.with_type(_keep_dictionary(field.type, field_result))
1628+
1629+
def list(self, list_type: pa.ListType | pa.LargeListType | pa.FixedSizeListType, element_result: pa.DataType) -> pa.DataType:
1630+
value_field = list_type.value_field.with_type(_keep_dictionary(list_type.value_type, element_result))
1631+
if isinstance(list_type, pa.FixedSizeListType):
1632+
return pa.list_(value_field, list_type.list_size)
1633+
elif isinstance(list_type, pa.LargeListType):
1634+
return pa.large_list(value_field)
1635+
return pa.list_(value_field)
1636+
1637+
def map(self, map_type: pa.MapType, key_result: pa.DataType, value_result: pa.DataType) -> pa.DataType:
1638+
return pa.map_(
1639+
map_type.key_field.with_type(_keep_dictionary(map_type.key_type, key_result)),
1640+
map_type.item_field.with_type(_keep_dictionary(map_type.item_type, value_result)),
1641+
keys_sorted=map_type.keys_sorted,
1642+
)
1643+
1644+
def primitive(self, primitive: pa.DataType) -> pa.DataType:
1645+
if pa.types.is_string_view(primitive):
1646+
return pa.large_string()
1647+
elif pa.types.is_binary_view(primitive):
1648+
return pa.large_binary()
1649+
return primitive
1650+
1651+
1652+
def _keep_dictionary(original: pa.DataType, result: pa.DataType) -> pa.DataType:
1653+
"""Re-apply dictionary encoding that visit_pyarrow strips before visiting the value type."""
1654+
if isinstance(original, pa.DictionaryType):
1655+
return pa.dictionary(original.index_type, result, ordered=original.ordered)
1656+
return result
1657+
1658+
16141659
class _ConvertToIcebergWithoutIDs(_ConvertToIceberg):
16151660
"""
16161661
Converts PyArrowSchema to Iceberg Schema with all -1 ids.
@@ -3091,10 +3136,7 @@ def _determine_partitions(spec: PartitionSpec, schema: Schema, arrow_table: pa.T
30913136
- We combine the chunks to create a copy to avoid GIL congestion on the original table
30923137
"""
30933138
# PyArrow has no group by, filter or take kernels for view types, so replace them with their large counterparts
3094-
arrow_schema = pa.schema(
3095-
[field.with_type(_replace_view_types(field.type)) for field in arrow_table.schema],
3096-
metadata=arrow_table.schema.metadata,
3097-
)
3139+
arrow_schema = visit_pyarrow(arrow_table.schema, _ConvertViewTypes())
30983140
if arrow_schema != arrow_table.schema:
30993141
arrow_table = arrow_table.cast(arrow_schema)
31003142

@@ -3149,29 +3191,6 @@ def _partition_value_filter(partition_field_name: str, value: Any) -> pc.Express
31493191
return pc.field(partition_field_name) == value
31503192

31513193

3152-
def _replace_view_types(arrow_type: pa.DataType) -> pa.DataType:
3153-
"""Replace string_view and binary_view, including nested occurrences, with large_string and large_binary."""
3154-
if pa.types.is_string_view(arrow_type):
3155-
return pa.large_string()
3156-
elif pa.types.is_binary_view(arrow_type):
3157-
return pa.large_binary()
3158-
elif isinstance(arrow_type, pa.StructType):
3159-
return pa.struct([field.with_type(_replace_view_types(field.type)) for field in arrow_type])
3160-
elif isinstance(arrow_type, pa.MapType):
3161-
return pa.map_(
3162-
arrow_type.key_field.with_type(_replace_view_types(arrow_type.key_type)),
3163-
arrow_type.item_field.with_type(_replace_view_types(arrow_type.item_type)),
3164-
keys_sorted=arrow_type.keys_sorted,
3165-
)
3166-
elif isinstance(arrow_type, pa.LargeListType):
3167-
return pa.large_list(arrow_type.value_field.with_type(_replace_view_types(arrow_type.value_type)))
3168-
elif isinstance(arrow_type, pa.FixedSizeListType):
3169-
return pa.list_(arrow_type.value_field.with_type(_replace_view_types(arrow_type.value_type)), arrow_type.list_size)
3170-
elif isinstance(arrow_type, pa.ListType):
3171-
return pa.list_(arrow_type.value_field.with_type(_replace_view_types(arrow_type.value_type)))
3172-
return arrow_type
3173-
3174-
31753194
def _get_field_from_arrow_table(arrow_table: pa.Table, field_path: str) -> pa.Array:
31763195
"""Get a field from an Arrow table, supporting both literal field names and nested field paths.
31773196

0 commit comments

Comments
 (0)