GitHub user Eshanatnight closed a discussion: Query regarding Arrow Flight
Server
Hi, I am trying to implement a basic Arrow Flight Server.
```rust
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.
GitHub link: https://github.com/apache/arrow-rs/discussions/5706
----
This is an automatically sent email for [email protected].
To unsubscribe, please send an email to: [email protected]
