Arrow schema issues for partition columns

[sebbegg]

Hi,

hope this is an okay place to mention this...
I forked off this repo to play a bit with ballista+delta myself.

I built a small web api in front of the scheduler usign axum+axum-streams.
When yielding results as arrow ipc streams I came across one issue:
abdolence/axum-streams-rs#80

When yielding results from delta tables there seem to be cases where the schema as reported by the dataframe and the schema of individual batches don't match. E.g. in:

let df = ctx.sql("SELECT * FROM example").await.unwrap();

let batches = df.collect().await.unwrap();
let batch_schema = batches.get(0).unwrap().schema();

println!("df.schema = {:?}", df.schema().clone());
println!("batch_schema = {:?}", batch_schema);

In datafusion alone I cannot reproduce the error, only in combination with ballista. I guess this diff might occur due to ballista writing partition results to arrow & reading them again?

Glad about any insights !

[milenkovicm]

Is the schema completely wrong or some columns are missing?

[sebbegg]

The schema seems correct except for some columns that are Dictionary(UInt16, Utf8) in the dataframe schema and plain Utf8 in the record batch.
Creating the IPC stream with one of the schemas and writing batches with a different schema, then of course leads to errors.

Digged a bit into delta and it seems that delta-rs wraps partition columns into dictionary arrays, which matches my observations:

https://github.com/delta-io/delta-rs/blob/3d7ab94617ed0be7806177c2f0ae6ec49f9aee0b/crates/core/src/delta_datafusion/table_provider/next/scan/plan.rs#L233-L235

All other columns, also complex/nested types, seem fine.

[milenkovicm]

I'm not sure if i can wrap my head around it at the moment
so delta has dictionary schema but record batch at the end has utf8 as result column, is it ?
are values correct apart from type ?

[sebbegg]

Yes, values in the record batches are correct.

Essentially the code looks like this:

async fn datafusion_sql_with_batches() -> impl IntoResponse {
    let ctx = SessionContext::new();
    let delta_table =
        open_table(Url::parse("file:///path/to/my_delta_data").unwrap())
            .await
            .unwrap();
    ctx.register_table("example", Arc::new(delta_table))
        .unwrap();

    let df = ctx.sql("SELECT * FROM example").await.unwrap();
    let df_schema = df.schema().clone().as_arrow().clone();
    println!("df.schema = {:?}", df_schema);

    let batches = df.collect().await.unwrap();
    let batch_schema = batches.get(0).unwrap().schema();

    StreamBodyAs::arrow_ipc(
        batches.get(0).unwrap().schema(), // okay
        // Arc::new(df_schema), // bad - because it lists partition colums as Dictionary types
        stream::iter(batches),
    )
}

Using collect and picking the schema of the first batch works. The IPC stream can be read e.g. in python and the data looks correct.

Using the df_schema or e.g the one returned as part of the stream from df.execute_stream() (which would of course be better choice than the collect) contains Dictionary types.
The actual batches then don't, leading to a corrupt IPC stream and the error messages shown in the ticket I linked.

[milenkovicm]

it looks like an issue due to stage writes. maybe we could add schema adapter ?

[sebbegg]

Hm, I think you lost me there - what's a schema adapter? 😉
Any existing examples in ballista I could have a look at?

Besides:
In general I think it would be smart if ballista keep/re-used the existing dictionary encoding. As it obviously saves bytes during transfer/fetch.

[milenkovicm]

Hm, I think you lost me there - what's a schema adapter? 😉
Any existing examples in ballista I could have a look at?

Not sure if tit will help

• https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/custom_data_source/custom_file_casts.rs

• https://docs.rs/datafusion/latest/datafusion/physical_expr_adapter/trait.PhysicalExprAdapterFactory.html

Besides:
In general I think it would be smart if ballista keep/re-used the existing dictionary encoding. As it obviously saves bytes during transfer/fetch.

If it can be done I have no problem, have a look at shuffle writer in ballista,

[milenkovicm]

SELECT * FROM example

Should not repartition, it would wort to check how ipc files are written and if there is something ballista does to flatten dictionaries

[milenkovicm]

Which ballista version you use?

[milenkovicm]

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

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

use arrow::array::{
    Array, DictionaryArray, Int32Array, Int32Builder, StringArray,
    StringDictionaryBuilder,
};
use arrow::datatypes::{DataType, Field, Int32Type, Schema};
use arrow::ipc::reader::StreamReader;
use arrow::ipc::writer::StreamWriter;
use arrow::record_batch::RecordBatch;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // --- Build schema with a dictionary column ---
    let schema = Arc::new(Schema::new(vec![
        Field::new("id", DataType::Int32, false),
        Field::new(
            "category",
            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
            true,
        ),
    ]));

    // --- Build arrays ---
    let ids = Int32Array::from(vec![1, 2, 3, 4, 5]);

    let mut dict_builder = StringDictionaryBuilder::<Int32Type>::new();
    dict_builder.append_value("alpha");
    dict_builder.append_value("beta");
    dict_builder.append_value("alpha");
    dict_builder.append_null();
    dict_builder.append_value("gamma");
    let categories: DictionaryArray<Int32Type> = dict_builder.finish();

    let batch =
        RecordBatch::try_new(schema.clone(), vec![Arc::new(ids), Arc::new(categories)])?;

    println!("Original schema:\n{}", batch.schema());
    println!("Rows: {}", batch.num_rows());

    // --- Write to target/ipc.arrows ---
    let path = "target/ipc.arrows";
    {
        let file = File::create(path)?;
        let mut writer = StreamWriter::try_new(file, &schema)?;
        writer.write(&batch)?;
        writer.finish()?;
    }
    println!("\nWrote IPC stream to {path}");

    // --- Read back ---
    let file = File::open(path)?;
    let reader = StreamReader::try_new(file, None)?;
    let read_schema = reader.schema();

    println!("\nRead-back schema:\n{read_schema}");

    // --- Compare schemas ---
    if schema == read_schema {
        println!("\nSchemas match ✓");
    } else {
        eprintln!("\nSchema mismatch!");
        eprintln!("  original : {schema:?}");
        eprintln!("  read-back: {read_schema:?}");
        std::process::exit(1);
    }

    // --- Verify data round-trips correctly ---
    let mut batches = vec![];
    for result in reader {
        batches.push(result?);
    }
    // Re-open since StreamReader consumes the file handle above
    let file = File::open(path)?;
    let mut reader2 = StreamReader::try_new(file, None)?;
    let read_batch = reader2.next().unwrap()?;

    assert_eq!(batch.num_rows(), read_batch.num_rows());
    println!("Row count matches: {} rows ✓", read_batch.num_rows());

    // Check dictionary column values
    let orig_col = batch
        .column_by_name("category")
        .unwrap()
        .as_any()
        .downcast_ref::<DictionaryArray<Int32Type>>()
        .unwrap();
    let read_col = read_batch
        .column_by_name("category")
        .unwrap()
        .as_any()
        .downcast_ref::<DictionaryArray<Int32Type>>()
        .unwrap();

    let orig_values = orig_col
        .values()
        .as_any()
        .downcast_ref::<StringArray>()
        .unwrap();
    let read_values = read_col
        .values()
        .as_any()
        .downcast_ref::<StringArray>()
        .unwrap();

    assert_eq!(
        orig_values.len(),
        read_values.len(),
        "dictionary sizes differ"
    );
    for i in 0..orig_values.len() {
        assert_eq!(orig_values.value(i), read_values.value(i));
    }
    println!("Dictionary values match ✓");

    Ok(())
}

writing and reading dictionary columns works as expected

this issue might be related to spiceai/datafusion-ballista#24

[sebbegg]

I'm on ballista 51 at the moment.
For the time being I just use the schema of the first returned batch - which seems fine.