Skip to content

Support Arrow type Dictionary with value FixedSizeBinary in Parquet #7445

@albertlockett

Description

@albertlockett

Is your feature request related to a problem or challenge? Please describe what you are trying to do.

I would like to be able to use an Arrow type of Dictionary(UInt8, FixedSizeBinary(16)) and be able to write this to parquet.

use std::sync::Arc;

use arrow_array::{FixedSizeBinaryArray, RecordBatch, UInt8Array, UInt8DictionaryArray};
use arrow_schema::{DataType, Field, Schema};
use object_store::{local::LocalFileSystem, path::Path};
use parquet::{
    arrow::async_writer::{AsyncArrowWriter, ParquetObjectWriter},
    file::properties::WriterProperties,
};

#[tokio::main]
async fn main() {
    let schema = Arc::new(Schema::new(vec![Field::new(
        "trace_id",
        DataType::Dictionary(
            Box::new(DataType::UInt8),
            Box::new(DataType::FixedSizeBinary(16)),
        ),
        false,
    )]));
    let keys = UInt8Array::from_iter_values(vec![0, 0, 0, 1, 1, 1, 2, 2]);
    let values = FixedSizeBinaryArray::try_from_iter(
        vec![
            vec![0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3],
            vec![2, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3],
            vec![2, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3],
        ]
        .into_iter(),
    )
    .unwrap();
    let arr = UInt8DictionaryArray::new(keys, Arc::new(values));

    let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr)]).unwrap();

    let object_store = Arc::new(LocalFileSystem::new_with_prefix("/tmp/").unwrap());
    let parquet_object_writer =
        ParquetObjectWriter::new(object_store.clone(), Path::from("albert.parquet"));
    let mut parquet_writer = AsyncArrowWriter::try_new(
        parquet_object_writer,
        schema.clone(),
        Some(WriterProperties::default()),
    )
    .unwrap();
    parquet_writer.write(&batch).await.unwrap();
    parquet_writer.close().await.unwrap();
}

Currently when I do this, I get an error:

called `Result::unwrap()` on an `Err` value: NYI("Attempting to write an Arrow type that is not yet implemented")

Describe the solution you'd like

I would like it if the type was supported

Describe alternatives you've considered

Decoding the dictionary into an array of the value type before sending the record batch to write. This seemed inefficient because I'd need to copy the values

Additional context

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementAny new improvement worthy of a entry in the changelog

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions