Skip to content

Commit 5ea0617

Browse files
Accept pyarrow LargeListType and FixedSizeListType (#458)
* Accept pyarrow LargeListType and FixedSizeListType * Combine the 3 register in one
1 parent 0c7c21c commit 5ea0617

File tree

2 files changed

+23
-1
lines changed

2 files changed

+23
-1
lines changed

pyiceberg/io/pyarrow.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -690,7 +690,9 @@ def _(obj: pa.StructType, visitor: PyArrowSchemaVisitor[T]) -> T:
690690

691691

692692
@visit_pyarrow.register(pa.ListType)
693-
def _(obj: pa.ListType, visitor: PyArrowSchemaVisitor[T]) -> T:
693+
@visit_pyarrow.register(pa.FixedSizeListType)
694+
@visit_pyarrow.register(pa.LargeListType)
695+
def _(obj: Union[pa.ListType, pa.LargeListType, pa.FixedSizeListType], visitor: PyArrowSchemaVisitor[T]) -> T:
694696
visitor.before_list_element(obj.value_field)
695697
result = visit_pyarrow(obj.value_type, visitor)
696698
visitor.after_list_element(obj.value_field)

tests/io/test_pyarrow_visitor.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -245,6 +245,26 @@ def test_pyarrow_list_to_iceberg() -> None:
245245
assert visit_pyarrow(pyarrow_list, _ConvertToIceberg()) == expected
246246

247247

248+
def test_pyarrow_large_list_to_iceberg() -> None:
249+
pyarrow_list = pa.large_list(pa.field("element", pa.int32(), nullable=False, metadata={"PARQUET:field_id": "1"}))
250+
expected = ListType(
251+
element_id=1,
252+
element_type=IntegerType(),
253+
element_required=True,
254+
)
255+
assert visit_pyarrow(pyarrow_list, _ConvertToIceberg()) == expected
256+
257+
258+
def test_pyarrow_fixed_size_list_to_iceberg() -> None:
259+
pyarrow_list = pa.list_(pa.field("element", pa.int32(), nullable=False, metadata={"PARQUET:field_id": "1"}), 1)
260+
expected = ListType(
261+
element_id=1,
262+
element_type=IntegerType(),
263+
element_required=True,
264+
)
265+
assert visit_pyarrow(pyarrow_list, _ConvertToIceberg()) == expected
266+
267+
248268
def test_pyarrow_map_to_iceberg() -> None:
249269
pyarrow_map = pa.map_(
250270
pa.field("key", pa.int32(), nullable=False, metadata={"PARQUET:field_id": "1"}),

0 commit comments

Comments
 (0)