feat: add dictionary_columns to to_arrow() / to_arrow_batch_reader() for memory-efficient reads

[GayathriSrividya]

Closes #3170

Rationale

Columns that contain large or frequently repeated string values (e.g. JSON blobs, low-cardinality categoricals) can exhaust memory when PyArrow loads them as plain string arrays. PyArrow's Parquet reader natively supports dictionary-encoded reads via its dictionary_columns kwarg, which deduplicates values and can dramatically reduce peak memory usage.

This was previously discussed in #3168 and a prior implementation (#3234) was closed as stale.

Changes

• Added dictionary_columns: tuple[str, ...] = () to Table.scan(), TableScan.__init__, and StagedTable.scan().
• Forwarded through DataScan.to_arrow() and to_arrow_batch_reader() → ArrowScan.__init__ → _task_to_record_batches → _get_file_format().
• Only applied when task.file.file_format == FileFormat.PARQUET; silently ignored for ORC (which does not support this kwarg).

Usage

# Read the "payload" column as dictionary-encoded to save memory
df = table.scan(dictionary_columns=("payload",)).to_arrow()

Verification

• Added test_dictionary_columns_produces_dict_encoded_output — confirms the requested column is dict-encoded, non-requested columns are plain, and values are identical.
• make lint ✓
• pytest tests/table/ tests/io/test_pyarrow.py ✓

[Fokko]

More of a meta question, but I don't think this helps for large JSON blobs since it would directly exceed the dictionary limit and fall back to plain encoding. I do think this helps a lot for low-cardinality strings.

Do we know how Arrow decodes the data? For example, if the Parquet column is dictionary encoded, would Arrow do something smart with the buffers to not repeat this value many times?

[Fokko]

I'm hesitant to add Arrow specific things to the public API

[GayathriSrividya]

Good point @Fokko I have moved dictionary_columns off the public scan() API and onto the Arrow-specific output methods instead:

table.scan(...).to_arrow(dictionary_columns=("payload",))
table.scan(...).to_arrow_batch_reader(dictionary_columns=("payload",))

That way it doesn't disturb the general scan interface. ArrowScan still accepts it for lower-level use. Pushed in the latest commit. Let me know if you have any further suggestions.

[Fokko]

Thanks @GayathriSrividya Moving this to the Arrow API seems a good solution. I think it is reasonable to add this. I've left some minor comments that I think would be good to clean up, but apart from that, this looks good to me

[Fokko]

Why not keep it a tuple here?

[GayathriSrividya]

I updated this to keep as a tuple.

[Fokko]

Here we convert the frozenset back to a tuple

[GayathriSrividya]

@Fokko I updated this to keep dictionary_columns as a tuple throughout the Arrow path, so we no longer convert frozenset -> tuple at read time. This is now in commit 33e4c6d.

[Fokko]

Thanks @GayathriSrividya for addressing the remaining nits 👍

[kevinjqliu]

I found one propagation gap: dictionary_columns reaches ArrowScan, but to_arrow_batch_reader() wraps the batches with the plain projected schema and casts back to it:

pa.RecordBatchReader.from_batches(target_schema, batches).cast(target_schema)

That erases the dictionary type, so to_arrow(dictionary_columns=...) preserves dictionary encoding, but to_arrow_batch_reader(dictionary_columns=...).read_all() returns the normal string type.

iceberg-python/pyiceberg/table/__init__.py

Lines 2242 to 2279 in 6da06ad

	def to_arrow_batch_reader(self, dictionary_columns: tuple[str, ...] = ()) -> pa.RecordBatchReader:
	"""Return an Arrow RecordBatchReader from this DataScan.
	For large results, using a RecordBatchReader requires less memory than
	loading an Arrow Table for the same DataScan, because a RecordBatch
	is read one at a time.
	Args:
	dictionary_columns:
	A tuple of column names that PyArrow should read as
	dictionary-encoded (``pa.DictionaryArray``). Dictionary
	encoding can substantially reduce memory usage for columns
	with low-cardinality repeated string values.
	Only applies to Parquet files; silently ignored for ORC.
	Returns:
	pa.RecordBatchReader: Arrow RecordBatchReader from the Iceberg table's DataScan
	which can be used to read a stream of record batches one by one.
	"""
	import pyarrow as pa
	from pyiceberg.io.pyarrow import ArrowScan, schema_to_pyarrow
	target_schema = schema_to_pyarrow(self.projection())
	batches = ArrowScan(
	self.table_metadata,
	self.io,
	self.projection(),
	self.row_filter,
	self.case_sensitive,
	self.limit,
	dictionary_columns=dictionary_columns,
	).to_record_batches(self.plan_files())
	return pa.RecordBatchReader.from_batches(
	target_schema,
	batches,
	).cast(target_schema)

Example public-path regression test:

def test_to_arrow_batch_reader_preserves_dictionary_columns(catalog: Catalog) -> None:
    arrow_table = pa.table(
        {
            "id": pa.array([1, 2, 3, 4], type=pa.int32()),
            "label": pa.array(["a", "b", "a", "b"], type=pa.string()),
        }
    )
    catalog.create_namespace_if_not_exists("default")
    table = catalog.create_table("default.dict_test", schema=arrow_table.schema)
    table.append(arrow_table)

    result = table.scan().to_arrow_batch_reader(dictionary_columns=("label",)).read_all()

    assert pa.types.is_dictionary(result.schema.field("label").type)
    assert result.column("label").to_pylist() == ["a", "b", "a", "b"]

I didn’t see another user-facing path where the option is accepted and dropped. Minor adjacent cleanup: update the abstract TableScan.to_arrow signature so type checkers see the new keyword.