Query regarding Arrow Flight Server

[Eshanatnight]

Hi, I am trying to implement a basic Arrow Flight Server.

use arrow_array::RecordBatch;
use arrow_flight::flight_service_server::FlightServiceServer;
use arrow_flight::PollInfo;
use arrow_schema::{ArrowError, Schema};
use chrono::Utc;
use datafusion::common::tree_node::TreeNode;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Instant;

use futures_util::{Future, TryFutureExt};

use tonic::transport::{Identity, Server, ServerTlsConfig};
use tonic_web::GrpcWebLayer;

use crate::option::CONFIG;

use crate::handlers::livetail::cross_origin_config;

use crate::handlers::http::query::{authorize_and_set_filter_tags, into_query};
use crate::query::{TableScanVisitor, QUERY_SESSION};
use crate::storage::object_storage::commit_schema_to_storage;
use crate::utils::arrow::flight::{get_query_from_ticket, run_do_get_rpc};
use arrow_flight::{
    flight_service_server::FlightService, Action, ActionType, Criteria, Empty, FlightData,
    FlightDescriptor, FlightInfo, HandshakeRequest, HandshakeResponse, PutResult, SchemaAsIpc,
    SchemaResult, Ticket,
};
use arrow_ipc::writer::{DictionaryTracker, IpcDataGenerator, IpcWriteOptions};
use futures::stream::BoxStream;
use tonic::{Request, Response, Status, Streaming};

use crate::handlers::livetail::extract_session_key;
use crate::metadata::STREAM_INFO;
use crate::rbac::Users;

const L_CURLY: char = '{';
const R_CURLY: char = '}';

#[derive(Clone)]
pub struct AirServiceImpl {}

#[tonic::async_trait]
impl FlightService for AirServiceImpl {
    type HandshakeStream = BoxStream<'static, Result<HandshakeResponse, Status>>;
    type ListFlightsStream = BoxStream<'static, Result<FlightInfo, Status>>;
    type DoGetStream = BoxStream<'static, Result<FlightData, Status>>;
    type DoPutStream = BoxStream<'static, Result<PutResult, Status>>;
    type DoActionStream = BoxStream<'static, Result<arrow_flight::Result, Status>>;
    type ListActionsStream = BoxStream<'static, Result<ActionType, Status>>;
    type DoExchangeStream = BoxStream<'static, Result<FlightData, Status>>;

	/// other methods are unimplemented

    async fn do_get(&self, req: Request<Ticket>) -> Result<Response<Self::DoGetStream>, Status> {
        let key = extract_session_key(req.metadata())?;

        let ticket = get_query_from_ticket(req)?;

        // get the query session_state
        let session_state = QUERY_SESSION.state();

        // get the logical plan and extract the table name
        let raw_logical_plan = session_state
            .create_logical_plan(&ticket.query)
            .await
            .map_err(|err| {
                log::error!("Datafusion Error: Failed to create logical plan: {}", err);
                Status::internal("Failed to create logical plan")
            })?;

        // create a visitor to extract the table name
        let mut visitor = TableScanVisitor::default();
        let _ = raw_logical_plan.visit(&mut visitor);
        let tables = visitor.into_inner();

        // map payload to query
        let mut query = into_query(&ticket, &session_state)
            .await
            .map_err(|_| Status::internal("Failed to parse query"))?;

        // if table name is not present it is a Malformed Query
        let stream_name = query
            .first_table_name()
            .ok_or_else(|| Status::invalid_argument("Malformed Query"))?;

        let permissions = Users.get_permissions(&key);

        authorize_and_set_filter_tags(&mut query, permissions, &stream_name).map_err(|_| {
            Status::permission_denied("User Does not have permission to access this")
        })?;

        let (results, _) = query
            .execute(stream_name.clone())
            .await
            .map_err(|err| Status::internal(err.to_string()))
            .unwrap();

        let schemas = results
            .iter()
            .map(|batch| batch.schema())
            .map(|s| s.as_ref().clone())
            .collect::<Vec<_>>();

        let schema = Schema::try_merge(schemas).map_err(|err| Status::internal(err.to_string()))?;
        let options = datafusion::arrow::ipc::writer::IpcWriteOptions::default();
        let schema_flight_data = SchemaAsIpc::new(&schema, &options);

        let mut flights = vec![FlightData::from(schema_flight_data)];
        let encoder = IpcDataGenerator::default();
        let mut tracker = DictionaryTracker::new(false);
        for batch in &results {
            let (flight_dictionaries, flight_batch) = encoder
                .encoded_batch(batch, &mut tracker, &options)
                .map_err(|e| Status::internal(e.to_string()))?;
            flights.extend(flight_dictionaries.into_iter().map(Into::into));
            flights.push(flight_batch.into());
        }
        let output = futures::stream::iter(flights.into_iter().map(Ok));

        Ok(Response::new(Box::pin(output) as Self::DoGetStream))
    }
}

pub fn server() -> impl Future<Output = Result<(), Box<dyn std::error::Error + Send>>> + Send {
    let mut addr: SocketAddr = CONFIG
        .parseable
        .address
        .parse()
        .expect("valid socket address");
    addr.set_port(CONFIG.parseable.flight_port);

    let service = AirServiceImpl {};

    let svc = FlightServiceServer::new(service);

    let cors = cross_origin_config();

    let identity = match (
        &CONFIG.parseable.tls_cert_path,
        &CONFIG.parseable.tls_key_path,
    ) {
        (Some(cert), Some(key)) => {
            match (std::fs::read_to_string(cert), std::fs::read_to_string(key)) {
                (Ok(cert_file), Ok(key_file)) => {
                    let identity = Identity::from_pem(cert_file, key_file);
                    Some(identity)
                }
                _ => None,
            }
        }
        (_, _) => None,
    };

    let config = identity.map(|id| ServerTlsConfig::new().identity(id));

    // rust is treating closures as different types
    let err_map_fn = |err| Box::new(err) as Box<dyn std::error::Error + Send>;

    // match on config to decide if we want to use tls or not
   Server::builder()
            .accept_http1(true)
            .max_frame_size((16 * 1024 * 1024) - 2) // 6MB ish
            .layer(cors)
            .layer(GrpcWebLayer::new())
            .add_service(svc)
            .serve(addr)
            .map_err(err_map_fn),

}

I am using a load testing tool to ingest data and query it with the arrow flight server.
When I run a query I get this error
Error, message length too large: found 5605265 bytes, the limit is: 4194304 bytes

Any help would be much appreciated.

[Eshanatnight]

My guess is I need to configure the tonic server to use a bigger message length but I could not find any resource for help

[Eshanatnight]

Thanks A lot for the help!!

[Eshanatnight]

I added the .max_decoding_message_size(32*1024*1024).max_encoding_message_size(32*1024*1024); but nothing changed

[Eshanatnight]

It still says size is about 7mb and limit is 4mb

[mbrobbel]

Added to let svc = FlightServiceServer::new(service);?

[tustvold]

https://docs.rs/arrow-flight/latest/arrow_flight/encode/struct.FlightDataEncoderBuilder.html#method.with_max_flight_data_size

The limit may be being enforced in the load testing tool

[Eshanatnight]

Added to let svc = FlightServiceServer::new(service);?

Yes I did it in there. I ever tried using it compression

    let svc = FlightServiceServer::new(service)
        .max_encoding_message_size(usize::MAX)
        .max_decoding_message_size(usize::MAX)
        .send_compressed(CompressionEncoding::Gzip);

But the result is the same.

if I dbg!(&svc); I am seeing the a the inner encoding size and decoding size are set as the values I passed.

I found this issue in tonic. But it seems to be closed.

[Eshanatnight]

docs.rs/arrow-flight/latest/arrow_flight/encode/struct.FlightDataEncoderBuilder.html#method.with_max_flight_data_size

The limit may be being enforced in the load testing tool

I dont think the load testing tool is the issue, because the ingestion takes place over http, and If I ingest the some data that is relatively small

example

[
	     {
            "id": "8080",
            "datetime": "24/Jun/2022:14:12:15 +0000",
            "host": "153.10.110.81"
        }
]

there is no problem with the do_get request

[Eshanatnight]

I am invoking the request with Postman, and a simple client I wrote

pub fn record_batches_to_json(records: &[&RecordBatch]) -> Vec<Map<String, Value>> {
    let buf = vec![];

    let mut writer = arrow_json::ArrayWriter::new(buf);
    writer.write_batches(records).unwrap();
    writer.finish().unwrap();

    let buf = writer.into_inner();

    let json_rows: Vec<Map<String, Value>> = match serde_json::from_reader(buf.as_slice()) {
        Ok(json) => json,
        Err(_) => vec![],
    };

    json_rows
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let mut args = args().collect::<Vec<_>>();
    args.remove(0);

    let end_time = args.pop().unwrap_or_else(|| "now".to_owned());
    let start_time = args.pop().unwrap_or_else(|| "1min".to_owned());
    let query = args
        .pop()
        .unwrap_or_else(|| "select * from teststream limit 1".to_owned());
    println!("{}:{}:{}", query, start_time, end_time);
    let channel = Channel::from_static("http://localhost:8002")
        .connect()
        .await?;

    let client = FlightClient::new(channel);
    let inn = client.into_inner().accept_compressed(tonic::codec::CompressionEncoding::Gzip).max_decoding_message_size(usize::MAX).max_encoding_message_size(4*1024*1024);
    let mut client = FlightClient::new_from_inner(inn);

    client.add_header("authorization", "Basic YWRtaW46YWRtaW4=")?;


    let td = format!(
        "{}\"query\":\"{}\", \"startTime\": \"{}\", \"endTime\": \"{}\"{}",
        '{', query, start_time, end_time, '}'
    );
    let ticket_data = serde_json::from_str::<Value>(&td)?;
    let mut ticket: Vec<u8> = vec![];
    serde_json::to_writer(&mut ticket, &ticket_data)?;

    let response = client
        .do_get(Ticket {
            ticket: ticket.into(),
        })
        .await?;

    let batches: Vec<RecordBatch> = response.try_collect().await?;
    dbg!(&batches);
    let q = batches.iter().collect::<Vec<&RecordBatch>>();
    let s = record_batches_to_json(&q);
    let s = s.into_iter().map(Value::Object).collect::<Vec<_>>();
    let q = serde_json::to_string_pretty(&s)?;
    println!("{}", q);
    Ok(())
}