Skip to content

Commit 7830bfb

Browse files
Gayathri Srividya Rajavarapuclaude
authored andcommitted
feat: add dictionary_columns parameter to Table.scan() for memory-efficient reads
Columns that contain large or frequently repeated strings (e.g. JSON blobs, low-cardinality categoricals) can exhaust memory when PyArrow loads them as plain string arrays. PyArrow's Parquet reader supports reading such columns as dictionary-encoded arrays, which deduplicates values and can dramatically reduce memory usage. Add a dictionary_columns: tuple[str, ...] parameter to Table.scan() (and the underlying TableScan / ArrowScan classes) that is forwarded to _get_file_format() as PyArrow's dictionary_columns kwarg. Only applies to Parquet files; silently ignored for ORC. Usage: table.scan(dictionary_columns=("payload",)).to_arrow() Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 0695f07 commit 7830bfb

3 files changed

Lines changed: 114 additions & 10 deletions

File tree

‎pyiceberg/io/pyarrow.py‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1625,8 +1625,12 @@ def _task_to_record_batches(
16251625
partition_spec: PartitionSpec | None = None,
16261626
format_version: TableVersion = TableProperties.DEFAULT_FORMAT_VERSION,
16271627
downcast_ns_timestamp_to_us: bool | None = None,
1628+
dictionary_columns: frozenset[str] = frozenset(),
16281629
) -> Iterator[pa.RecordBatch]:
1629-
arrow_format = _get_file_format(task.file.file_format, pre_buffer=True, buffer_size=(ONE_MEGABYTE * 8))
1630+
format_kwargs: dict[str, Any] = {"pre_buffer": True, "buffer_size": ONE_MEGABYTE * 8}
1631+
if dictionary_columns and task.file.file_format == FileFormat.PARQUET:
1632+
format_kwargs["dictionary_columns"] = tuple(dictionary_columns)
1633+
arrow_format = _get_file_format(task.file.file_format, **format_kwargs)
16301634
with io.new_input(task.file.file_path).open() as fin:
16311635
fragment = arrow_format.make_fragment(fin)
16321636
physical_schema = fragment.physical_schema
@@ -1729,6 +1733,7 @@ class ArrowScan:
17291733
_case_sensitive: bool
17301734
_limit: int | None
17311735
_downcast_ns_timestamp_to_us: bool | None
1736+
_dictionary_columns: frozenset[str]
17321737
"""Scan the Iceberg Table and create an Arrow construct.
17331738
17341739
Attributes:
@@ -1738,6 +1743,7 @@ class ArrowScan:
17381743
_bound_row_filter: Schema bound row expression to filter the data with
17391744
_case_sensitive: Case sensitivity when looking up column names
17401745
_limit: Limit the number of records.
1746+
_dictionary_columns: Column names to read as dictionary-encoded arrays.
17411747
"""
17421748

17431749
def __init__(
@@ -1748,6 +1754,8 @@ def __init__(
17481754
row_filter: BooleanExpression,
17491755
case_sensitive: bool = True,
17501756
limit: int | None = None,
1757+
*,
1758+
dictionary_columns: tuple[str, ...] = (),
17511759
) -> None:
17521760
self._table_metadata = table_metadata
17531761
self._io = io
@@ -1756,6 +1764,7 @@ def __init__(
17561764
self._case_sensitive = case_sensitive
17571765
self._limit = limit
17581766
self._downcast_ns_timestamp_to_us = Config().get_bool(DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE)
1767+
self._dictionary_columns = frozenset(dictionary_columns)
17591768

17601769
@property
17611770
def _projected_field_ids(self) -> set[int]:
@@ -1866,6 +1875,7 @@ def _record_batches_from_scan_tasks_and_deletes(
18661875
self._table_metadata.specs().get(task.file.spec_id),
18671876
self._table_metadata.format_version,
18681877
self._downcast_ns_timestamp_to_us,
1878+
self._dictionary_columns,
18691879
)
18701880
for batch in batches:
18711881
if self._limit is not None:

‎pyiceberg/table/__init__.py‎

Lines changed: 33 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1219,6 +1219,7 @@ def scan(
12191219
snapshot_id: int | None = None,
12201220
options: Properties = EMPTY_DICT,
12211221
limit: int | None = None,
1222+
dictionary_columns: tuple[str, ...] = (),
12221223
) -> DataScan:
12231224
"""Fetch a DataScan based on the table's current metadata.
12241225
@@ -1245,6 +1246,14 @@ def scan(
12451246
An integer representing the number of rows to
12461247
return in the scan result. If None, fetches all
12471248
matching rows.
1249+
dictionary_columns:
1250+
A tuple of column names that PyArrow should read as
1251+
dictionary-encoded (``pa.DictionaryArray``). Dictionary
1252+
encoding can substantially reduce memory usage for columns
1253+
that contain large or frequently repeated string values
1254+
(e.g. large JSON blobs or low-cardinality categoricals).
1255+
Only applies to Parquet files; silently ignored for ORC.
1256+
Columns absent from the file are silently skipped.
12481257
12491258
Returns:
12501259
A DataScan based on the table's current metadata.
@@ -1260,6 +1269,7 @@ def scan(
12601269
limit=limit,
12611270
catalog=self.catalog,
12621271
table_identifier=self._identifier,
1272+
dictionary_columns=dictionary_columns,
12631273
)
12641274

12651275
@property
@@ -1775,6 +1785,7 @@ def scan(
17751785
snapshot_id: int | None = None,
17761786
options: Properties = EMPTY_DICT,
17771787
limit: int | None = None,
1788+
dictionary_columns: tuple[str, ...] = (),
17781789
) -> DataScan:
17791790
raise ValueError("Cannot scan a staged table")
17801791

@@ -1809,6 +1820,7 @@ class TableScan(ABC):
18091820
limit: int | None
18101821
catalog: Catalog | None
18111822
table_identifier: Identifier | None
1823+
dictionary_columns: tuple[str, ...]
18121824

18131825
def __init__(
18141826
self,
@@ -1822,6 +1834,7 @@ def __init__(
18221834
limit: int | None = None,
18231835
catalog: Catalog | None = None,
18241836
table_identifier: Identifier | None = None,
1837+
dictionary_columns: tuple[str, ...] = (),
18251838
):
18261839
self.table_metadata = table_metadata
18271840
self.io = io
@@ -1833,6 +1846,7 @@ def __init__(
18331846
self.limit = limit
18341847
self.catalog = catalog
18351848
self.table_identifier = table_identifier
1849+
self.dictionary_columns = dictionary_columns
18361850

18371851
def snapshot(self) -> Snapshot | None:
18381852
if self.snapshot_id:
@@ -2072,13 +2086,11 @@ def _build_residual_evaluator(self, spec_id: int) -> Callable[[DataFile], Residu
20722086
# The lambda created here is run in multiple threads.
20732087
# So we avoid creating _EvaluatorExpression methods bound to a single
20742088
# shared instance across multiple threads.
2075-
return lambda datafile: (
2076-
residual_evaluator_of(
2077-
spec=spec,
2078-
expr=self.row_filter,
2079-
case_sensitive=self.case_sensitive,
2080-
schema=self.table_metadata.schema(),
2081-
)
2089+
return lambda datafile: residual_evaluator_of(
2090+
spec=spec,
2091+
expr=self.row_filter,
2092+
case_sensitive=self.case_sensitive,
2093+
schema=self.table_metadata.schema(),
20822094
)
20832095

20842096
@staticmethod
@@ -2224,7 +2236,13 @@ def to_arrow(self) -> pa.Table:
22242236
from pyiceberg.io.pyarrow import ArrowScan
22252237

22262238
return ArrowScan(
2227-
self.table_metadata, self.io, self.projection(), self.row_filter, self.case_sensitive, self.limit
2239+
self.table_metadata,
2240+
self.io,
2241+
self.projection(),
2242+
self.row_filter,
2243+
self.case_sensitive,
2244+
self.limit,
2245+
dictionary_columns=self.dictionary_columns,
22282246
).to_table(self.plan_files())
22292247

22302248
def to_arrow_batch_reader(self) -> pa.RecordBatchReader:
@@ -2244,7 +2262,13 @@ def to_arrow_batch_reader(self) -> pa.RecordBatchReader:
22442262

22452263
target_schema = schema_to_pyarrow(self.projection())
22462264
batches = ArrowScan(
2247-
self.table_metadata, self.io, self.projection(), self.row_filter, self.case_sensitive, self.limit
2265+
self.table_metadata,
2266+
self.io,
2267+
self.projection(),
2268+
self.row_filter,
2269+
self.case_sensitive,
2270+
self.limit,
2271+
dictionary_columns=self.dictionary_columns,
22482272
).to_record_batches(self.plan_files())
22492273

22502274
return pa.RecordBatchReader.from_batches(

‎tests/io/test_pyarrow.py‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5103,3 +5103,73 @@ def test_partition_column_projection_with_schema_evolution(catalog: InMemoryCata
51035103
result_sorted = result.sort_by("name")
51045104
assert result_sorted["name"].to_pylist() == ["Alice", "Bob", "Charlie", "David"]
51055105
assert result_sorted["new_column"].to_pylist() == [None, None, "new1", "new2"]
5106+
5107+
5108+
def test_dictionary_columns_produces_dict_encoded_output(tmpdir: str) -> None:
5109+
"""dictionary_columns passed to ArrowScan must yield dictionary-encoded arrays.
5110+
5111+
Verifies that:
5112+
1. The requested column is returned as a pa.DictionaryArray.
5113+
2. Values are identical to a plain (non-dict) scan.
5114+
3. A column NOT in dictionary_columns is still returned as a plain array.
5115+
"""
5116+
from pyiceberg.expressions import AlwaysTrue
5117+
from pyiceberg.io.pyarrow import ArrowScan, PyArrowFileIO
5118+
from pyiceberg.partitioning import PartitionSpec
5119+
from pyiceberg.table import FileScanTask
5120+
from pyiceberg.table.metadata import TableMetadataV2
5121+
5122+
arrow_schema = pa.schema(
5123+
[
5124+
pa.field("id", pa.int32(), nullable=True, metadata={PYARROW_PARQUET_FIELD_ID_KEY: "1"}),
5125+
pa.field("label", pa.string(), nullable=True, metadata={PYARROW_PARQUET_FIELD_ID_KEY: "2"}),
5126+
]
5127+
)
5128+
arrow_table = pa.table(
5129+
[pa.array([1, 2, 3, 4], type=pa.int32()), pa.array(["a", "b", "a", "b"], type=pa.string())],
5130+
schema=arrow_schema,
5131+
)
5132+
data_file = _write_table_to_data_file(f"{tmpdir}/test_dict_cols.parquet", arrow_schema, arrow_table)
5133+
data_file.spec_id = 0
5134+
5135+
iceberg_schema = Schema(
5136+
NestedField(1, "id", IntegerType(), required=False),
5137+
NestedField(2, "label", StringType(), required=False),
5138+
)
5139+
table_metadata = TableMetadataV2(
5140+
location=f"file://{tmpdir}",
5141+
last_column_id=2,
5142+
format_version=2,
5143+
schemas=[iceberg_schema],
5144+
partition_specs=[PartitionSpec()],
5145+
)
5146+
io = PyArrowFileIO()
5147+
task = FileScanTask(data_file)
5148+
5149+
scan_plain = ArrowScan(
5150+
table_metadata=table_metadata,
5151+
io=io,
5152+
projected_schema=iceberg_schema,
5153+
row_filter=AlwaysTrue(),
5154+
)
5155+
scan_dict = ArrowScan(
5156+
table_metadata=table_metadata,
5157+
io=io,
5158+
projected_schema=iceberg_schema,
5159+
row_filter=AlwaysTrue(),
5160+
dictionary_columns=("label",),
5161+
)
5162+
5163+
result_plain = scan_plain.to_table([task])
5164+
result_dict = scan_dict.to_table([task])
5165+
5166+
# id column is not in dictionary_columns — both scans should return int32
5167+
assert result_plain.schema.field("id").type == pa.int32()
5168+
assert result_dict.schema.field("id").type == pa.int32()
5169+
5170+
# label column: plain scan → string, dict scan → dictionary<values=string, indices=int32>
5171+
assert result_plain.schema.field("label").type == pa.string()
5172+
assert pa.types.is_dictionary(result_dict.schema.field("label").type)
5173+
5174+
# Values must be identical
5175+
assert result_plain.column("label").to_pylist() == result_dict.column("label").to_pylist()

0 commit comments

Comments
 (0)