Skip to content

Commit 747ebee

Browse files
feat: Support passing kwargs to polars.scan_iceberg in Table.to_polars (#3868)
feat: Support passing kwargs to polars.scan_iceberg in Table.to_polars (#3128) Co-authored-by: Kevin Liu <kevinjqliu@users.noreply.github.com>
1 parent aa0a911 commit 747ebee

2 files changed

Lines changed: 33 additions & 3 deletions

File tree

‎pyiceberg/table/__init__.py‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1900,15 +1900,19 @@ def to_bodo(self) -> bd.DataFrame:
19001900

19011901
return bd.read_iceberg_table(self)
19021902

1903-
def to_polars(self) -> pl.LazyFrame:
1903+
def to_polars(self, **kwargs: Any) -> pl.LazyFrame:
19041904
"""Lazily read from this Apache Iceberg table.
19051905
1906+
Args:
1907+
**kwargs: Additional keyword arguments to forward to :func:`polars.scan_iceberg`
1908+
(e.g. ``storage_options``, ``snapshot_id``, ``reader_override``).
1909+
19061910
Returns:
1907-
pl.LazyFrame: Unmaterialized Polars LazyFrame created from the Iceberg table
1911+
pl.LazyFrame: Unmaterialized Polars LazyFrame created from the Iceberg table.
19081912
"""
19091913
import polars as pl
19101914

1911-
return pl.scan_iceberg(self)
1915+
return pl.scan_iceberg(self, **kwargs)
19121916

19131917
def __datafusion_table_provider__(self, session: Any | None = None) -> IcebergDataFusionTable:
19141918
"""Return the DataFusion table provider PyCapsule interface.

‎tests/table/test_init.py‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2067,3 +2067,29 @@ def _spy(*args: Any, **kwargs: Any) -> FileIO:
20672067

20682068
assert seen_locations, "expected at least one load_file_io call"
20692069
assert all(loc is not None for loc in seen_locations), f"load_file_io called without a location: {seen_locations}"
2070+
2071+
2072+
def test_table_to_polars_forwards_kwargs(table_v2: Table, monkeypatch: pytest.MonkeyPatch) -> None:
2073+
import polars as pl
2074+
2075+
recorded_calls: list[tuple[Any, dict[str, Any]]] = []
2076+
2077+
def mock_scan_iceberg(source: Any, **kwargs: Any) -> Any:
2078+
recorded_calls.append((source, kwargs))
2079+
return "mock_lazy_frame"
2080+
2081+
monkeypatch.setattr(pl, "scan_iceberg", mock_scan_iceberg)
2082+
2083+
# Test without kwargs
2084+
res1 = table_v2.to_polars()
2085+
assert res1 == "mock_lazy_frame"
2086+
assert recorded_calls[0] == (table_v2, {})
2087+
2088+
# Test with kwargs
2089+
storage_options = {"s3": {"endpoint": "https://s3.custom.endpoint"}}
2090+
res2 = table_v2.to_polars(storage_options=storage_options, snapshot_id=12345, reader_override="native")
2091+
assert res2 == "mock_lazy_frame"
2092+
assert recorded_calls[1] == (
2093+
table_v2,
2094+
{"storage_options": storage_options, "snapshot_id": 12345, "reader_override": "native"},
2095+
)

0 commit comments

Comments
 (0)