crates/knut-bifrost/src/flight.rs
//! Flight transform plugins (#6): a [`Transform`] whose logic lives in an
//! **external Arrow-Flight service**, reached over the bidirectional
//! [`DoExchange`] stream — **Arrow record batches in, Arrow record batches
//! out** — instead of in-process Rust or a sandboxed wasm module. The service
//! can be written in any language with an Arrow Flight server (Python, Go,
//! C++, Java, …); bifröst just proxies the transform to it and keeps the same
//! [`GraphMutation`] contract.
//!
//! This is the sibling of [`crate::wasm::WasmTransform`]: same shape (rows →
//! transformed rows → one [`VertexSet`] per output batch under a configured
//! label + id column), but **out-of-process and polyglot** rather than an
//! in-process wasm interpreter.
//!
//! ## Wire protocol
//!
//! `apply` opens a `DoExchange` call to the configured endpoint, encodes the
//! input batches as a Flight data stream ([`FlightDataEncoderBuilder`]), and
//! reads the server's response stream back as record batches
//! ([`FlightRecordBatchStream`], via the client). The server is free to
//! filter / derive / reshape the rows; an identity server (echo the input)
//! yields the input rows as vertices.
//!
//! [`DoExchange`]: arrow_flight::flight_service_client::FlightServiceClient::do_exchange
//!
//! ```no_run
//! # #[cfg(feature = "flight")]
//! # async fn demo() -> knut_bifrost::Result<()> {
//! use knut_bifrost::flight::FlightTransform;
//! let t = FlightTransform::connect("http://transform-svc:50055", "Person", "_gar_id").await?;
//! # let _ = t; Ok(()) }
//! ```
use std::sync::Arc;
use arrow_array::RecordBatch;
use arrow_flight::{
FlightClient, decode::FlightRecordBatchStream, encode::FlightDataEncoderBuilder,
};
use futures::TryStreamExt;
use tokio::sync::Mutex;
use tonic::transport::Channel;
use crate::error::{BifrostError, Result};
use crate::mutation::{GraphMutation, VertexSet};
use crate::transform::{TableContext, Transform};
/// A [`Transform`] that proxies to an external Arrow-Flight `DoExchange`
/// service (Arrow record batches in/out).
///
/// The connected [`FlightClient`] is shared behind an `Arc<Mutex<…>>` so the
/// transform is `Send + Sync` and a single TCP channel is reused across
/// `apply` calls (a `DoExchange` call borrows the client mutably).
pub struct FlightTransform {
client: Arc<Mutex<FlightClient>>,
label: String,
id_column: String,
}
impl FlightTransform {
/// Connect to an Arrow-Flight endpoint (e.g. `http://host:50055`) and bind
/// the transform's output to `label` / `id_column`.
///
/// The connection is established once here; each `apply` opens a fresh
/// `DoExchange` call over the same channel.
pub async fn connect(
endpoint: impl Into<String>,
label: impl Into<String>,
id_column: impl Into<String>,
) -> Result<Self> {
let endpoint = endpoint.into();
let channel = Channel::from_shared(endpoint.clone())
.map_err(|e| flight_err(format!("invalid endpoint {endpoint:?}: {e}")))?
.connect()
.await
.map_err(|e| flight_err(format!("connect {endpoint:?}: {e}")))?;
Ok(Self::from_client(
FlightClient::new(channel),
label,
id_column,
))
}
/// Build from an already-connected [`FlightClient`] (e.g. over a test
/// channel). `connect` is the usual entry point.
pub fn from_client(
client: FlightClient,
label: impl Into<String>,
id_column: impl Into<String>,
) -> Self {
Self {
client: Arc::new(Mutex::new(client)),
label: label.into(),
id_column: id_column.into(),
}
}
/// Run one `DoExchange` round: send `batches`, collect the transformed
/// batches the service streams back.
async fn exchange(&self, batches: &[RecordBatch]) -> Result<Vec<RecordBatch>> {
// Encode the input batches as a Flight data stream. The encoder emits
// the schema message followed by one dictionary/record-batch message
// set per batch — exactly what the server's `FlightRecordBatchStream`
// decoder expects.
let owned: Vec<RecordBatch> = batches.to_vec();
let input =
FlightDataEncoderBuilder::new().build(futures::stream::iter(owned.into_iter().map(Ok)));
let mut client = self.client.lock().await;
let response: FlightRecordBatchStream = client
.do_exchange(input)
.await
.map_err(|e| flight_err(format!("do_exchange: {e}")))?;
response
.try_collect::<Vec<_>>()
.await
.map_err(|e| flight_err(format!("decode response: {e}")))
}
}
impl Transform for FlightTransform {
fn apply(&self, _ctx: &TableContext, batches: &[RecordBatch]) -> Result<GraphMutation> {
if batches.is_empty() {
return Ok(GraphMutation::default());
}
// `Transform::apply` is sync; bridge to the async Flight client the
// same way the rest of the workspace does — drive the future on the
// ambient Tokio runtime if there is one (the listener runs inside
// `#[tokio::main]`), else spin up a throwaway current-thread runtime.
let out_batches = block_on(self.exchange(batches))?;
let vertices = out_batches
.into_iter()
.map(|batch| VertexSet {
label: self.label.clone(),
id_column: self.id_column.clone(),
batch,
})
.collect();
Ok(GraphMutation {
edge_deletes: Vec::new(),
vertices,
..Default::default()
})
}
}
/// Map a transport/Flight error string into a [`BifrostError`], prefixed so
/// it's recognisable without touching `error.rs`.
fn flight_err(msg: String) -> BifrostError {
BifrostError::Other(format!("flight: {msg}"))
}
/// Drive an async transform to completion from the sync [`Transform::apply`].
///
/// `apply` may be called from inside a running Tokio runtime (the listener's
/// `#[tokio::main]`) or from a plain thread. Blocking a worker thread on a
/// future would deadlock the current-thread flavor, so when a runtime handle
/// is present we hand the future to a **dedicated blocking thread** with its
/// own current-thread runtime (`block_in_place` needs the multi-thread
/// flavor, which we can't assume). With no ambient runtime we just build a
/// throwaway current-thread runtime here.
fn block_on<F>(fut: F) -> F::Output
where
F: std::future::Future + Send,
F::Output: Send,
{
match tokio::runtime::Handle::try_current() {
Ok(_) => std::thread::scope(|scope| {
scope
.spawn(|| {
tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("build current-thread runtime")
.block_on(fut)
})
.join()
.expect("flight transform worker thread panicked")
}),
Err(_) => tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("build current-thread runtime")
.block_on(fut),
}
}
#[cfg(test)]
mod tests {
// Shadow the crate's single-arg `Result<T>` alias (pulled in via `super::*`)
// with the std two-arg `Result` the `FlightService` trait methods need.
use std::result::Result;
use super::{
Channel, FlightClient, FlightDataEncoderBuilder, FlightRecordBatchStream, FlightTransform,
RecordBatch, TableContext, Transform,
};
use futures::TryStreamExt;
use std::net::SocketAddr;
use std::sync::Arc;
use arrow_array::{Int64Array, StringArray};
use arrow_flight::{
FlightData, PutResult, SchemaResult, Ticket,
error::FlightError,
flight_service_server::{FlightService, FlightServiceServer},
};
use arrow_schema::{DataType, Field, Schema};
use futures::stream::BoxStream;
use tonic::transport::server::TcpIncoming;
use tonic::{Request, Response, Status, Streaming};
/// A minimal in-process `DoExchange` server: it decodes the incoming Flight
/// stream into record batches and re-encodes them back unchanged (an
/// **echo** transform). This exercises the full client round-trip — encode
/// → DoExchange → decode → VertexSet — without a separate process. A real
/// service would filter/derive/reshape between decode and re-encode.
#[derive(Clone)]
struct EchoExchange;
type ExchangeStream = BoxStream<'static, Result<FlightData, Status>>;
#[tonic::async_trait]
impl FlightService for EchoExchange {
type HandshakeStream = BoxStream<'static, Result<arrow_flight::HandshakeResponse, Status>>;
type ListFlightsStream = BoxStream<'static, Result<arrow_flight::FlightInfo, Status>>;
type DoGetStream = ExchangeStream;
type DoPutStream = BoxStream<'static, Result<PutResult, Status>>;
type DoExchangeStream = ExchangeStream;
type DoActionStream = BoxStream<'static, Result<arrow_flight::Result, Status>>;
type ListActionsStream = BoxStream<'static, Result<arrow_flight::ActionType, Status>>;
async fn do_exchange(
&self,
request: Request<Streaming<FlightData>>,
) -> Result<Response<Self::DoExchangeStream>, Status> {
// Decode the client's input stream into RecordBatches …
let inbound = request.into_inner();
let batches = FlightRecordBatchStream::new_from_flight_data(
inbound.map_err(|s| FlightError::Tonic(Box::new(s))),
)
.try_collect::<Vec<_>>()
.await
.map_err(|e| Status::internal(e.to_string()))?;
// … and echo them straight back out as a Flight data stream.
let out: ExchangeStream = Box::pin(
FlightDataEncoderBuilder::new()
.build(futures::stream::iter(batches.into_iter().map(Ok)))
.map_err(|e: FlightError| Status::internal(e.to_string())),
);
Ok(Response::new(out))
}
// ── Unused surface — this test server only speaks DoExchange ──────────
async fn handshake(
&self,
_r: Request<Streaming<arrow_flight::HandshakeRequest>>,
) -> Result<Response<Self::HandshakeStream>, Status> {
Err(Status::unimplemented("handshake"))
}
async fn list_flights(
&self,
_r: Request<arrow_flight::Criteria>,
) -> Result<Response<Self::ListFlightsStream>, Status> {
Err(Status::unimplemented("list_flights"))
}
async fn get_flight_info(
&self,
_r: Request<arrow_flight::FlightDescriptor>,
) -> Result<Response<arrow_flight::FlightInfo>, Status> {
Err(Status::unimplemented("get_flight_info"))
}
async fn poll_flight_info(
&self,
_r: Request<arrow_flight::FlightDescriptor>,
) -> Result<Response<arrow_flight::PollInfo>, Status> {
Err(Status::unimplemented("poll_flight_info"))
}
async fn get_schema(
&self,
_r: Request<arrow_flight::FlightDescriptor>,
) -> Result<Response<SchemaResult>, Status> {
Err(Status::unimplemented("get_schema"))
}
async fn do_get(&self, _r: Request<Ticket>) -> Result<Response<Self::DoGetStream>, Status> {
Err(Status::unimplemented("do_get"))
}
async fn do_put(
&self,
_r: Request<Streaming<FlightData>>,
) -> Result<Response<Self::DoPutStream>, Status> {
Err(Status::unimplemented("do_put"))
}
async fn do_action(
&self,
_r: Request<arrow_flight::Action>,
) -> Result<Response<Self::DoActionStream>, Status> {
Err(Status::unimplemented("do_action"))
}
async fn list_actions(
&self,
_r: Request<arrow_flight::Empty>,
) -> Result<Response<Self::ListActionsStream>, Status> {
Err(Status::unimplemented("list_actions"))
}
}
fn sample() -> RecordBatch {
RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("_gar_id", DataType::Int64, false),
Field::new("name", DataType::Utf8, false),
])),
vec![
Arc::new(Int64Array::from(vec![1, 2, 3])),
Arc::new(StringArray::from(vec!["a", "b", "c"])),
],
)
.unwrap()
}
/// Stand up the echo `DoExchange` server on an ephemeral port (background
/// task), point a `FlightTransform` at it, run `apply` over a sample batch,
/// and assert the resulting `GraphMutation.vertices` carry the rows back.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn flight_exchange_round_trips_to_vertices() {
// Bind on an ephemeral port first so the address is known before the
// server task starts (no fixed-port race). `TcpIncoming` lets the
// server consume the already-bound listener.
let incoming = TcpIncoming::bind("127.0.0.1:0".parse::<SocketAddr>().unwrap())
.expect("bind ephemeral port");
let addr: SocketAddr = incoming.local_addr().unwrap();
tokio::spawn(async move {
tonic::transport::Server::builder()
.add_service(FlightServiceServer::new(EchoExchange))
.serve_with_incoming(incoming)
.await
.unwrap();
});
let endpoint = format!("http://{addr}");
let transform = FlightTransform::connect(&endpoint, "Person", "_gar_id")
.await
.expect("connect to echo flight service");
let batch = sample();
// `apply` is sync and bridges to async internally; run it off the
// async test thread so the inner runtime isn't nested on this one.
let mutation = tokio::task::spawn_blocking(move || {
transform.apply(&TableContext { table: "t".into() }, &[batch])
})
.await
.unwrap()
.expect("apply succeeds");
assert_eq!(
mutation.vertices.len(),
1,
"one output batch → one vertex set"
);
let vs = &mutation.vertices[0];
assert_eq!(vs.label, "Person");
assert_eq!(vs.id_column, "_gar_id");
assert_eq!(vs.batch.num_rows(), 3);
assert_eq!(vs.batch, sample(), "echo service round-trips the rows");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn empty_input_is_a_no_op() {
// No server needed: an empty batch slice short-circuits before any
// network call.
let client = FlightClient::new(Channel::from_static("http://127.0.0.1:1").connect_lazy());
let transform = FlightTransform::from_client(client, "X", "_gar_id");
let mutation = tokio::task::spawn_blocking(move || {
transform.apply(&TableContext { table: "t".into() }, &[])
})
.await
.unwrap()
.unwrap();
assert!(mutation.is_empty());
}
}