Allow dictionary output for selected string columns in ArrowReader #3614

Is your feature request related to a problem or challenge?

iceberg-rust's Arrow reader returns string columns as plain strings, even when the Parquet file stores them with a dictionary. A caller that hashes or groups by those columns then hashes every row's value, or rebuilds the dictionary that the file already had.

I reproduced this with ArrowReader at main 896f9d6c50ff52fa9e43aec170d3f14bfb40b512. ArrowReaderBuilder does not let callers request dictionaries. The reader creates parquet's ArrowReaderOptions internally, without a string type hint, so arrow-rs decodes these columns as Utf8. Even if the Parquet reader returned dictionaries, the output transformer would cast them back to the snapshot schema's ordinary string type.

Minimal reproduction

Run this as src/main.rs of a Rust 2024 crate with cargo run. I used Rust 1.97.0 on linux/amd64. The pinned arrow-array, arrow-schema and parquet versions match upstream's Cargo.lock at the tested revision.

[dependencies]
iceberg = { git = "https://github.com/apache/iceberg-rust", rev = "896f9d6c50ff52fa9e43aec170d3f14bfb40b512" }
arrow-array = "=59.3.2"
arrow-schema = "=59.3.2"
futures = "0.3"
parquet = { version = "=59.3.0", default-features = false, features = ["arrow"] }
tokio = { version = "1", features = ["macros", "rt"] }

use std::{fs::File, sync::Arc};

use arrow_array::{RecordBatch, StringArray};
use arrow_schema::{DataType, Schema};
use futures::{StreamExt, TryStreamExt, stream};
use iceberg::arrow::{ArrowReaderBuilder, schema_to_arrow_schema};
use iceberg::spec::{DataFileFormat, NestedField, PrimitiveType, Schema as IcebergSchema};
use iceberg::{Runtime, io::FileIO, scan::FileScanTask};
use parquet::arrow::ArrowWriter;
use parquet::arrow::arrow_reader::{ArrowReaderOptions, ParquetRecordBatchReaderBuilder};
use parquet::basic::Encoding;

#[tokio::main(flavor = "current_thread")]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let path = std::env::temp_dir().join("strings.parquet");
    let fields = [NestedField::required(1, "s", PrimitiveType::String).into()];
    let iceberg_schema = Arc::new(IcebergSchema::builder().with_fields(fields).build()?);
    // Adds the `PARQUET:field_id` metadata that ArrowReader uses to match columns.
    let arrow_schema = Arc::new(schema_to_arrow_schema(&iceberg_schema)?);

    let values = StringArray::from(vec!["a", "b", "a", "b"]);
    let batch = RecordBatch::try_new(arrow_schema.clone(), vec![Arc::new(values)])?;
    let mut writer = ArrowWriter::try_new(File::create(&path)?, arrow_schema.clone(), None)?;
    writer.write(&batch)?;
    writer.close()?;

    let size = std::fs::metadata(&path)?.len();
    let task = FileScanTask::builder()
        .with_data_file_path(path.display().to_string())
        .with_data_file_format(DataFileFormat::Parquet)
        .with_file_size_in_bytes(size)
        .with_start(0)
        .with_length(size)
        .with_schema(iceberg_schema)
        .with_project_field_ids(vec![1])
        .with_case_sensitive(true)
        .build()?;
    let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
    let mut batches = reader.read(stream::iter([Ok(task)]).boxed())?.stream();
    let batch = batches.try_next().await?.unwrap();
    println!("Iceberg: {:?}", batch.column(0).data_type());

    let dictionary = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
    let dictionary_field = arrow_schema.field(0).clone().with_data_type(dictionary);
    let hint = Arc::new(Schema::new(vec![dictionary_field]));
    let options = ArrowReaderOptions::new().with_schema(hint);
    let file = File::open(&path)?;
    let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options)?;
    let chunk = builder.metadata().row_group(0).column(0);
    assert!(chunk.encodings().any(|e| e == Encoding::RLE_DICTIONARY));
    let batch = builder.build()?.next().unwrap()?;
    println!("Parquet: {:?}", batch.column(0).data_type());
    Ok(())
}

Output. The first line is what ArrowReader returns. The second is what arrow-rs returns for the same file when given a dictionary hint, which is the expected type.

Iceberg: Utf8
Parquet: Dictionary(Int32, Utf8)

Measurements

I wrote a synthetic, unpartitioned Iceberg table: 30 files of 1,000,000 rows, ZSTD(1), and 131,072-row row groups, totaling 585 MiB. Four string columns have 4,096, 200, 64 and 3 distinct values. A fifth has up to 1,048,576 distinct values per file and falls back to plain pages. The table also has a sorted timestamp and two numeric columns.

I wrote the files with ParquetWriter and DataFileWriter and committed them with fast_append through a MemoryCatalog. I planned with TableScan::plan_files and read with 8,192-row batches and data-file concurrency 1. The consumer hashes strings and sums numbers. For dictionaries, it hashes each dictionary value once per batch and adds the corresponding hash for each row.

I ran on local disk in a linux/amd64 VM with 8 AMD EPYC vCPUs and 16 GB RAM, using Rust 1.97.0, a current-thread Tokio runtime, CPU affinity 2-5, and Arrow 59.3 and Parquet 59.3.0 from upstream's Cargo.lock.

I also built a local prototype. It passes a dictionary type hint for the listed columns to arrow-rs and keeps those types in RecordBatchTransformer. Time to read and hash the table, in seconds (median of five fresh processes, no predicate):

In the prototype build, dictionaries for all five string columns took 4.42 s, against 2.65 s for the four dictionary-encoded columns. The plain-encoded column removed the whole gain. Percentages compare against stock main, not the prototype with its option off. In this run, differences under about 10% between builds were within noise. With a consumer that only counts rows, reading took 2.50 s with the option off and 2.13 s with four dictionaries, in the same build. Peak memory stayed around 25-26 MiB. Row counts and per-column hashing checksums matched.

The parquet crate alone took 4.07 s with Utf8 and 2.54 s with four dictionaries, close to the prototype's 2.65 s. Its Utf8View result was 3.51 s. These references need no prototype and bypass the Iceberg transformer. Object storage, other writers, and nullable or partitioned tables were not measured at benchmark scale.

Describe the solution you'd like

Allow callers to request selected string columns as Arrow dictionaries. Make the request per column: including the plain-encoded column removed the gain in this workload. Leave unlisted columns and defaults unchanged, and preserve the requested representation through the output transformer. A scan must return consistent types across files, including files with plain pages or missing columns. Choosing the output type independently per file from its page encodings would give one scan different types in different files.

Callers of the scan API must be able to make this request, as well as callers of ArrowReaderBuilder. The request must survive planning on one host and reading on another. At its current head, the open #3514 moves ArrowReaderBuilder settings into serializable reader options and removes its per-field setters. Under that API, the dictionary request belongs in those options.

A DataFusion table provider would need to declare these types too. That work belongs in https://github.com/apache/datafusion-iceberg and depends on the reader supporting the request.

Open questions: option naming, column names versus field IDs, dictionary key type, and whether to support nested strings such as list<string>.

Related work:

• feat: add dictionary_columns to to_arrow() / to_arrow_batch_reader() for memory-efficient reads iceberg-python#3461 is merged and adds dictionary_columns to reduce memory use for large or repeated strings.
• fix: preserve dictionary encoding in to_arrow_batch_reader iceberg-python#3595 is merged and preserves dictionary_columns through the streaming output conversion.
• Support DataType::Dictionary for arrow schema conversion #1277 (closed) and feat: support arrow dictionary in schema conversion #1293 (merged) concern Arrow-to-Iceberg schema conversion, not reader output.
• fix(reader): keep matching rows when filtering a promoted column #3540 is open and adds a RecordBatchTransformer for row-filter columns, another place where the requested representation would need to be preserved.

Willingness to contribute

I would be willing to contribute to this feature with guidance from the Iceberg Rust community