Flight SQL - DictId Schema Serialization Issue - FlatBuf ParseError

[amunra]

TL;DR:

Failed deserializing batch: ParseError("Unable to convert flight info to a message: Type `i64` at position 1380 is unaligned.\n\twhile verifying table field `id` at position 1380\n\twhile verifying table field `dictionary` at position 1356\n\twhile verifying vector element 1 at position 48\n\twhile verifying table field `fields` at position 36\n\t while verifying union variant `MessageHeader::Schema` at position 24\n\twhile verifying table field `header` at position 24\n\n")

Printing that string:

Unable to convert flight info to a message: Type `i64` at position 1380 is unaligned.
	while verifying table field `id` at position 1380
	while verifying table field `dictionary` at position 1356
	while verifying vector element 1 at position 48
	while verifying table field `fields` at position 36
	 while verifying union variant `MessageHeader::Schema` at position 24
	while verifying table field `header` at position 24

When querying from a Rust client to a Rust server.

Using: arrow-rs 0.45 (both client and server)

Details

I'm struggling with sending dictionaries back in Flight SQL from my Rust implementation.

I'm prototyping a Flight SQL server implementation which works fine so long as I don't query dictionary types.
I've had a few stumbling blocks, but the current one that I'm stuck on is serializing a schema with multiple dictionary columns. These columns are all Dictionary<Int32, Utf8>: They are effectively interned strings and a core piece of functionality for what I do.

My async fn do_action_create_prepared_statement runs successfully and looks like so:

    async fn do_action_create_prepared_statement(
        &self,
        query: ActionCreatePreparedStatementRequest,
        request: Request<Action>,
    ) -> Result<ActionCreatePreparedStatementResult, Status> {
        self.check_token(&request)?;
        let sql_query = query.query.as_str();
        let cursor = Query::execute(&self.engine, sql_query).map_err(to_status_err)?;
        let schema = cursor.arrow_schema();
        info!("do_action_create_prepared_statement schema: {:?}", *schema);
        let message = SchemaAsIpc::new(&schema, &IpcWriteOptions::default())
            .try_into()
            .map_err(|e| status!("Unable to serialize schema", e))?;
        let IpcMessage(schema_bytes) = message;
        let uuid = Uuid::new_v4();
        let uuid_bytes = prost::bytes::Bytes::copy_from_slice(uuid.as_bytes());
        let res = ActionCreatePreparedStatementResult {
            prepared_statement_handle: uuid_bytes,
            dataset_schema: schema_bytes,
            parameter_schema: Default::default(), // TODO: parameters
        };
        self.handle_queries.insert(uuid, cursor);
        info!("prepared statement handle uuid: {}", uuid);
        Ok(res)
    }

The logged schema looks a bit like this:

Schema { fields: [
  Field { 
    name: "hostname",
    data_type: Dictionary(Int32, Utf8),
    nullable: false,
    dict_id: 0,
    dict_is_ordered: false,
    metadata: {}
  },
  Field {
    name: "region",
    data_type: Dictionary(Int32, Utf8),
    nullable: false,
    dict_id: 1,
    dict_is_ordered: false, metadata: {}
  },
  Field {
    name: "datacenter",
    data_type: Dictionary(Int32, Utf8),
    nullable: false,
    dict_id: 2,
    dict_is_ordered: false, metadata: {}
  }
  .... // other non-dict fields
}

This causes the client to reject the schema when processing the prepare action's response.

Note how it contains multiple dict_ids: I previously (incorrectly) had left all the dict fields with dict_id == 0 and the schema was deserialized just fine by the client. It failed (rather cryptically) later on, but I was able to debug that issue.

When ran in a debugger, this is the backtrace:

arrow_ipc::convert::try_schema_from_ipc_buffer convert.rs:190
arrow_flight::{impl#15}::try_from lib.rs:366
arrow_flight::sql::client::{impl#0}::prepare::{async_fn#0} client.rs:368
flight_sql_demo::async_main::{async_block#0} main.rs:20
tokio::runtime::park::{impl#4}::block_on::{closure#0}<flight_sql_demo::async_main::{async_block_env#0}> park.rs:282
tokio::runtime::coop::with_budget<core::task::poll::Poll<core::result::Result<(), alloc::boxed::Box<dyn core::error::Error, alloc::alloc::Global>>>, tokio::runtime::park::{impl#4}::block_on::{closure_env#0}<flight_sql_demo::async_main::{async_block_env#0}>> coop.rs:107
tokio::runtime::coop::budget<core::task::poll::Poll<core::result::Result<(), alloc::boxed::Box<dyn core::error::Error, alloc::alloc::Global>>>, tokio::runtime::park::{impl#4}::block_on::{closure_env#0}<flight_sql_demo::async_main::{async_block_env#0}>> coop.rs:73
tokio::runtime::park::CachedParkThread::block_on<flight_sql_demo::async_main::{async_block_env#0}> park.rs:282
tokio::runtime::context::blocking::BlockingRegionGuard::block_on<flight_sql_demo::async_main::{async_block_env#0}> blocking.rs:66
tokio::runtime::scheduler::multi_thread::{impl#0}::block_on::{closure#0}<flight_sql_demo::async_main::{async_block_env#0}> mod.rs:87
tokio::runtime::context::runtime::enter_runtime<tokio::runtime::scheduler::multi_thread::{impl#0}::block_on::{closure_env#0}<flight_sql_demo::async_main::{async_block_env#0}>, core::result::Result<(), alloc::boxed::Box<dyn core::error::Error, alloc::alloc::Global>>> runtime.rs:65
tokio::runtime::scheduler::multi_thread::MultiThread::block_on<flight_sql_demo::async_main::{async_block_env#0}> mod.rs:86
tokio::runtime::runtime::Runtime::block_on<flight_sql_demo::async_main::{async_block_env#0}> runtime.rs:349
flight_sql_demo::async_main main.rs:50
flight_sql_demo::main main.rs:66
core::ops::function::FnOnce::call_once<fn(), ()> function.rs:250
std::sys_common::backtrace::__rust_begin_short_backtrace<fn(), ()> backtrace.rs:134
std::rt::lang_start::{closure#0}<()> rt.rs:166
core::ops::function::impls::{impl#2}::call_once<(), (dyn core::ops::function::Fn<(), Output=i32> + core::marker::Sync + core::panic::unwind_safe::RefUnwindSafe)> function.rs:287
std::panicking::try::do_call<&(dyn core::ops::function::Fn<(), Output=i32> + core::marker::Sync + core::panic::unwind_safe::RefUnwindSafe), i32> panicking.rs:485
std::panicking::try<i32, &(dyn core::ops::function::Fn<(), Output=i32> + core::marker::Sync + core::panic::unwind_safe::RefUnwindSafe)> panicking.rs:449
std::panic::catch_unwind<&(dyn core::ops::function::Fn<(), Output=i32> + core::marker::Sync + core::panic::unwind_safe::RefUnwindSafe), i32> panic.rs:140
std::rt::lang_start_internal::{closure#2} rt.rs:148
std::panicking::try::do_call<std::rt::lang_start_internal::{closure_env#2}, isize> panicking.rs:485
std::panicking::try<isize, std::rt::lang_start_internal::{closure_env#2}> panicking.rs:449
std::panic::catch_unwind<std::rt::lang_start_internal::{closure_env#2}, isize> panic.rs:140
std::rt::lang_start_internal rt.rs:148
std::rt::lang_start<()> rt.rs:165
main 0x0000555555654e1e
__libc_start_call_main 0x00007ffff7c23a90
__libc_start_main_impl 0x00007ffff7c23b49
_start 0x000055555564a6b5

Here's my test Rust client:

use arrow_flight::sql::client::FlightSqlServiceClient;
use futures::{StreamExt, TryStreamExt};
use log::{error, info};
use std::error::Error;
use arrow_flight::decode::FlightRecordBatchStream;
use arrow_flight::error::FlightError;
use tonic::transport::Channel;

#[tokio::main]
async fn async_main(query: String) -> Result<(), Box<dyn Error>> {
    let uri = "grpc+tcp://localhost:50051".parse()?;
    let token = "uuid_token";
    let start = std::time::Instant::now();
    let channel = Channel::builder(uri).connect().await?;
    info!("Connected to gRPC channel: {channel:?}, about to prepare query: {query:?}");
    let mut client = FlightSqlServiceClient::new(channel);
    client.set_token(token.to_string());
    let mut prepared = client
        .prepare(query, None)
        .await?;
    let info = prepared.execute().await?;
    let ticket = info
        .endpoint
        .get(0)
        .expect("endpoint defined")
        .ticket
        .as_ref()
        .expect("ticked defined")
        .clone();
    info!("Got ticket: {ticket:?}");

    // This API I don't understand. What am I supposed to do with `FlightData`?
    let response_stream = client.do_get(ticket).await?;
    let mut response_stream = FlightRecordBatchStream::new_from_flight_data(
        response_stream.map_err(FlightError::Tonic));
    let mut batch_count = 0usize;
    let mut n_rows = 0usize;
    while let Some(batch) = response_stream.next().await {
        if batch.is_err() {
            error!("Batch {batch_count} errored");
        }
        let batch = batch?;
        n_rows += batch.num_rows();
        batch_count += 1;
    }
    let end = std::time::Instant::now();
    let elapsed = end - start;
    let rows_sec = n_rows as f64 / elapsed.as_secs_f64();
    info!("Got {n_rows} rows back as {batch_count} batches, took {elapsed:?} ({rows_sec} rows/sec)");
    Ok(())
}

fn main() {
    std::env::set_var("RUST_LOG", "trace");
    env_logger::init();

    let query = std::env::args().nth(1);
    let query = match query {
        None => {
            error!("No query provided");
            return;
        },
        Some(query) => query,
    };

    if let Err(err) = async_main(query) {
        error!("Failed deserializing batch: {err:?}");
    }
}

Arrow Flight SQL is pretty cool, but I keep encountering issues with dictionaries. For more context on what I'm up to:

• Java - Flight SQL - StackOverflowException caused by org.apache.arrow.flight.FlightStream.next arrow#37094
• Python adbc_driver_flightsql segfault in Go code arrow-adbc#975

Any help would be very welcome!

[tustvold]

The error you are getting suggests you are decoding something as the wrong flatbuffer type, that is an internal error from the flatbuffer verification not arrow. IIRC DictionaryBatch are encoded differently from regular RecordBatch

There is also #3389 which may be related

[amunra]

I'm quite sure that the issue I'm encountering doesn't have much to do with RecordBatch as it's not involved in the response being sent from the server to the client.

• The client fails to parse the server's ActionCreatePreparedStatementResult response as produced by the do_action_create_prepared_statement message.
• The server encodes the schema (see code provided earlier) via SchemaAsIpc::new(&schema, &IpcWriteOptions::default()).try_into() which in turn calls schema_to_ipc_format -> flight_schema_as_encoded_data -> IpcDataGenerator::default().schema_to_bytes().

The schema_to_bytes impl is as follows:

impl IpcDataGenerator {
    pub fn schema_to_bytes(
        &self,
        schema: &Schema,
        write_options: &IpcWriteOptions,
    ) -> EncodedData {
        let mut fbb = FlatBufferBuilder::new();
        let schema = {
            let fb = crate::convert::schema_to_fb_offset(&mut fbb, schema);
            fb.as_union_value()
        };

        let mut message = crate::MessageBuilder::new(&mut fbb);
        message.add_version(write_options.metadata_version);
        message.add_header_type(crate::MessageHeader::Schema);
        message.add_bodyLength(0);
        message.add_header(schema);
        // TODO: custom metadata
        let data = message.finish();
        fbb.finish(data, None);

        let data = fbb.finished_data();
        EncodedData {
            ipc_message: data.to_vec(),
            arrow_data: vec![],
        }
    }

I.e., it's all flatbuffer code, and at this stage I'm a little lost.

What does "Type i64 at position 1380 is unaligned." mean exactly? I'm assuming this means that the layout of the fields sent doesn't match the layout of the data being received, but they're both the same lib, so I can only assume the client Rust code is decoding with different definitions than the server is sending.

P.S.: I didn't realise that dictionaries weren't fully supported yet: Thanks for pointing me to that ticket!

[tustvold]

it's all flatbuffer code, and at this stage I'm a little lost.

If only it were that simple 😅 Flight itself is gRPC using protobuf as the message encoding (as arrow IPC probably should have done). FlightSQL is then a layer on top of this.

Unfortunately the response message for Flight::do_action is opaque

rpc DoAction(Action) returns (stream Result) {}

 /*
  * An opaque action specific for the service.
  */
 message Action {
   string type = 1;
   bytes body = 2;
 }

 /*
  * An opaque result returned after executing an action.
  */
 message Result {
   bytes body = 1;
 }

And additionally, as FlightSQL multiplexes multiple actions onto this DoAction endpoint, it needs a way to return multiple different payloads.

To achieve this it wraps both the ActionCreatePreparedStatementRequest and the ActionCreatePreparedStatementResult in protobuf::Any.

This can be seen in the FlightSqlServiceClient here - https://github.com/apache/arrow-rs/blob/master/arrow-flight/src/sql/client.rs#L350.

And in the server here - https://github.com/apache/arrow-rs/blob/master/arrow-flight/src/sql/server.rs#L830

I don't know exactly where this is breaking down, but the error you are getting is basically it saying that what you are feeding the flatbuffer decoder is not a valid flatbuffer. I'm afraid without your actual code I can't suggest why this might be occuring, but hopefully the above provides some color that might help...

[amunra]

Here's a repro:
https://github.com/amunra/flight_sql_repro/tree/main

[amunra]

Oops!
I had just noticed that I had marked the repo as private by mistake.
Sorry!
It should be accessible now.

As mentioned before, changing the dict_id from 1 to 0 for the last column hides the bug.