crates/knut-flight/src/lib.rs
//! # knut-flight — an Apache **Arrow Flight** server for knut's graph output (N4)
//!
//! Where [`graphar-flight`](../graphar_flight/index.html) speaks **Flight SQL**
//! (Cypher in, Arrow out, against a live FalkorDB), this crate is a plain
//! **Arrow Flight** server over knut's *rel2graph* output: it takes a
//! [`GraphMutation`](knut_bifrost::mutation::GraphMutation) (the nodes + edges
//! the mapping fabric produced from the 12-table ERP bench) plus that bench's
//! source relational tables, and serves them as ticketed Arrow streams.
//!
//! ## Surface
//! * `list_flights` / `get_flight_info` advertise every available dataset with
//! its Arrow schema and a [`Ticket`]. The ticket is simply the dataset name,
//! UTF-8 bytes — no protobuf envelope, so any Flight client (Rust, pyarrow,
//! Go…) can round-trip it verbatim.
//! * `get_schema` returns a dataset's schema without pulling data.
//! * `do_get` streams the dataset's [`RecordBatch`]es for a ticket.
//!
//! ## The datasets ([`Catalog`])
//! | ticket | what |
//! |---|---|
//! | `nodes` | aggregate node skeleton — `(label, id)` for every node, all labels |
//! | `edges` | aggregate edge list — `(src_label, rel, dst_label, src, dst)`, resolved edges only |
//! | `node/<Label>` | one label's full **native** vertex batch (typed property columns) |
//! | `edge/<src>_<rel>_<dst>` | one relationship's edge batch (`src`/`dst` + property columns) |
//! | `table/<erp.qualified.name>` | a raw source relational table batch |
//!
//! The aggregate `nodes` + `edges` flights carry a UNIFORM schema across labels,
//! which is exactly what a `networkx.MultiDiGraph` loader wants — so a client can
//! pull them and rebuild the SAME graph the file-export demo (N3) builds,
//! proving the Flight path and the file path agree (see `demo/flight_networkx_demo.py`).
//!
//! ## Deployment note (→ N5)
//! This is a clean single-binary service: it takes its bind address and graph
//! scale from flags/env, does no interactive setup, and holds all state in
//! memory built at startup. That is the unit N5 wraps as a Skidbladnir/k8s
//! service.
use std::collections::{BTreeMap, HashSet};
use std::path::Path;
use std::pin::Pin;
use std::sync::Arc;
use arrow_array::builder::{Int64Builder, StringBuilder};
use arrow_array::{Array, Int64Array, RecordBatch};
use arrow_flight::encode::FlightDataEncoderBuilder;
use arrow_flight::error::FlightError;
use arrow_flight::flight_service_server::FlightService;
use arrow_flight::{
Action, ActionType, Criteria, Empty, FlightData, FlightDescriptor, FlightEndpoint, FlightInfo,
HandshakeRequest, HandshakeResponse, PollInfo, PutResult, SchemaAsIpc, SchemaResult, Ticket,
};
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use futures::StreamExt;
use knut_bifrost::MappingSpec;
use knut_bifrost::mutation::GraphMutation;
use knut_bifrost::rel2graph::{ArrowTables, build_mutation as engine_build_mutation};
use knut_rel2graph_bench::{DEFAULT_SEED, build_mutation, generate, mapping_spec, schema};
use tonic::{Request, Response, Status, Streaming};
/// The short git commit this crate was built from (`KNUT_GIT_HASH`, emitted by
/// `build.rs`). Defaults to `"unknown"` when git metadata is unreadable.
pub const GIT_HASH: &str = env!("KNUT_GIT_HASH");
/// `v<version> (<git-hash>)` — the `--version` line, per the knut build law.
pub const VERSION_LINE: &str = concat!(
"v",
env!("CARGO_PKG_VERSION"),
" (",
env!("KNUT_GIT_HASH"),
")"
);
/// Default bind port — `:50055`, chosen to avoid the FlightSQL server's
/// `:50051`/`:50054`.
pub const DEFAULT_PORT: u16 = 50055;
// ── the catalog ─────────────────────────────────────────────────────────────────
/// One served dataset: its ticket name, a human description (for `list_flights`
/// logging), its Arrow schema, and the batches `do_get` streams.
#[derive(Debug, Clone)]
pub struct Dataset {
pub name: String,
pub description: String,
pub schema: SchemaRef,
pub batches: Vec<RecordBatch>,
}
impl Dataset {
fn total_rows(&self) -> usize {
self.batches.iter().map(|b| b.num_rows()).sum()
}
}
/// Everything the server serves, built once at startup from a rel2graph result.
#[derive(Debug, Clone, Default)]
pub struct Catalog {
order: Vec<String>,
map: BTreeMap<String, Dataset>,
}
impl Catalog {
fn insert(&mut self, ds: Dataset) {
self.order.push(ds.name.clone());
self.map.insert(ds.name.clone(), ds);
}
/// The datasets, in a stable listing order (`nodes`, `edges`, per-label
/// vertices, per-rel edges, then source tables).
pub fn datasets(&self) -> impl Iterator<Item = &Dataset> {
self.order.iter().filter_map(move |n| self.map.get(n))
}
pub fn get(&self, name: &str) -> Option<&Dataset> {
self.map.get(name)
}
pub fn len(&self) -> usize {
self.map.len()
}
pub fn is_empty(&self) -> bool {
self.map.is_empty()
}
/// Build the catalog for the 12-table ERP bench at `scale` (deterministic;
/// uses [`DEFAULT_SEED`]). Runs the real rel2graph fabric — no live DB.
///
/// This is the synthetic entrypoint. The live-data entrypoint is
/// [`from_parquet_dir`](Self::from_parquet_dir); both funnel through the
/// single seam [`from_mutation`](Self::from_mutation), so a bench catalog and
/// a real-data catalog are byte-identical in shape.
pub fn from_bench(scale: usize) -> Self {
let spec = mapping_spec();
let db = generate(scale, DEFAULT_SEED);
let mutation = build_mutation(&spec, &db);
let source_tables: Vec<(String, Vec<RecordBatch>)> = schema()
.iter()
.map(|t| (t.name.to_string(), db.batches(t.name).to_vec()))
.filter(|(_, b)| !b.is_empty())
.collect();
Catalog::from_mutation(&mutation, source_tables)
}
/// **The live-data seam.** Serve a REAL graph read from a parquet directory —
/// an Iceberg / object-store export of the source tables — instead of the
/// synthetic bench generator.
///
/// For every table named in `spec`, reads `<dir>/<table>.parquet` (e.g.
/// `erp.hr.employees` → `dir/erp.hr.employees.parquet`) into Arrow, drives the
/// engine's canonical [`build_mutation`](knut_bifrost::rel2graph::build_mutation)
/// over the collected [`ArrowTables`], and serves the resulting graph plus the
/// source tables — exactly as [`from_bench`](Self::from_bench) does, but over
/// data that came off disk rather than a seeded PRNG. A spec table with no
/// parquet file on disk is skipped (a partial-but-consistent graph). This is
/// the shape korp's `iceberg2graph` step produces when it lands lakehouse
/// tables as parquet and hands the directory to this server.
pub fn from_parquet_dir(spec: &MappingSpec, dir: &Path) -> anyhow::Result<Self> {
use anyhow::Context;
let mut tables = ArrowTables::new();
let mut source_tables: Vec<(String, Vec<RecordBatch>)> = Vec::new();
for tm in &spec.tables {
let path = dir.join(format!("{}.parquet", tm.table));
if !path.exists() {
continue; // table not exported here — skip, stay consistent
}
let batches = read_parquet(&path)
.with_context(|| format!("read parquet for table `{}` at {path:?}", tm.table))?;
if batches.is_empty() {
continue;
}
tables.insert(tm.table.clone(), batches.clone());
source_tables.push((tm.table.clone(), batches));
}
anyhow::ensure!(
!tables.is_empty(),
"no parquet tables found in {dir:?} for any of the {} table(s) in the mapping spec",
spec.tables.len()
);
let mutation = engine_build_mutation(spec, &tables)
.context("run the rel2graph engine over the parquet source")?;
Ok(Catalog::from_mutation(&mutation, source_tables))
}
/// **The single catalog-building seam.** Materialise the served datasets from
/// a rel2graph [`GraphMutation`] plus the source tables it was built from.
///
/// Whatever produced the mutation — the bench generator or a live Iceberg /
/// parquet read — the served shape is identical: `nodes` + `edges`
/// aggregates, per-label `node/*` vertex batches, per-rel `edge/*` batches,
/// and one `table/*` per source table. `source_tables` may be empty (then no
/// `table/*` datasets are served).
pub fn from_mutation<I>(mutation: &GraphMutation, source_tables: I) -> Self
where
I: IntoIterator<Item = (String, Vec<RecordBatch>)>,
{
let mut catalog = Catalog::default();
// ── aggregate node skeleton: (label, id) for every distinct node ──────
catalog.insert(build_nodes_aggregate(mutation));
// ── aggregate edge list: (src_label, rel, dst_label, src, dst) ────────
catalog.insert(build_edges_aggregate(mutation));
// ── per-label native vertex batches (typed property columns) ──────────
// One VertexSet per label in this schema; concatenate defensively if a
// label ever spans several sets by keeping them as separate batches.
let mut by_label: BTreeMap<String, Vec<RecordBatch>> = BTreeMap::new();
let mut label_schema: BTreeMap<String, SchemaRef> = BTreeMap::new();
for vs in &mutation.vertices {
by_label
.entry(vs.label.clone())
.or_default()
.push(vs.batch.clone());
label_schema
.entry(vs.label.clone())
.or_insert_with(|| vs.batch.schema());
}
for (label, batches) in by_label {
let schema = label_schema.remove(&label).unwrap();
catalog.insert(Dataset {
name: format!("node/{label}"),
description: format!("native vertex batch for label `{label}`"),
schema,
batches,
});
}
// ── per-relationship edge batches (src/dst + property columns) ────────
for es in &mutation.edges {
let name = format!("edge/{}_{}_{}", es.src_label, es.edge_type, es.dst_label);
catalog.insert(Dataset {
name,
description: format!(
"edge batch ({})-[{}]->({})",
es.src_label, es.edge_type, es.dst_label
),
schema: es.batch.schema(),
batches: vec![es.batch.clone()],
});
}
// ── source relational tables (the rel2graph inputs) ───────────────────
for (name, batches) in source_tables {
if batches.is_empty() {
continue;
}
let table_schema = batches[0].schema();
catalog.insert(Dataset {
name: format!("table/{name}"),
description: format!("source relational table `{name}`"),
schema: table_schema,
batches,
});
}
catalog
}
}
/// Read every `RecordBatch` from a parquet file into memory (the live-data
/// source path). Uses the arrow parquet reader; no object-store round trip —
/// the caller has already localised the export to a directory.
fn read_parquet(path: &Path) -> anyhow::Result<Vec<RecordBatch>> {
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
let file = std::fs::File::open(path)?;
let reader = ParquetRecordBatchReaderBuilder::try_new(file)?.build()?;
let mut batches = Vec::new();
for b in reader {
batches.push(b?);
}
Ok(batches)
}
/// Aggregate node skeleton: `(label: Utf8, id: Int64)` for every distinct node
/// across all labels — the uniform-schema flight a networkx loader consumes.
fn build_nodes_aggregate(mutation: &GraphMutation) -> Dataset {
let schema = Arc::new(Schema::new(vec![
Field::new("label", DataType::Utf8, false),
Field::new("id", DataType::Int64, false),
]));
let mut labels = StringBuilder::new();
let mut ids = Int64Builder::new();
let mut seen: HashSet<(String, i64)> = HashSet::new();
for vs in &mutation.vertices {
let idx = vs
.batch
.schema()
.index_of(&vs.id_column)
.expect("vertex id column present");
let col = vs
.batch
.column(idx)
.as_any()
.downcast_ref::<Int64Array>()
.expect("vertex id column is Int64");
for row in 0..col.len() {
if col.is_null(row) {
continue;
}
let id = col.value(row);
if seen.insert((vs.label.clone(), id)) {
labels.append_value(&vs.label);
ids.append_value(id);
}
}
}
let batch = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(labels.finish()), Arc::new(ids.finish())],
)
.expect("nodes-aggregate batch");
Dataset {
name: "nodes".to_string(),
description: "aggregate node skeleton (label, id) across all labels".to_string(),
schema,
batches: vec![batch],
}
}
/// Aggregate edge list: `(src_label, rel, dst_label, src, dst)` for every
/// RESOLVED edge (both endpoints non-null) — the uniform-schema flight.
fn build_edges_aggregate(mutation: &GraphMutation) -> Dataset {
let schema = Arc::new(Schema::new(vec![
Field::new("src_label", DataType::Utf8, false),
Field::new("rel", DataType::Utf8, false),
Field::new("dst_label", DataType::Utf8, false),
Field::new("src", DataType::Int64, false),
Field::new("dst", DataType::Int64, false),
]));
let mut src_l = StringBuilder::new();
let mut rels = StringBuilder::new();
let mut dst_l = StringBuilder::new();
let mut srcs = Int64Builder::new();
let mut dsts = Int64Builder::new();
for es in &mutation.edges {
let s_idx = es.batch.schema().index_of("src").expect("edge `src`");
let d_idx = es.batch.schema().index_of("dst").expect("edge `dst`");
let s_col = es
.batch
.column(s_idx)
.as_any()
.downcast_ref::<Int64Array>()
.expect("edge src is Int64");
let d_col = es
.batch
.column(d_idx)
.as_any()
.downcast_ref::<Int64Array>()
.expect("edge dst is Int64");
for row in 0..s_col.len() {
if s_col.is_null(row) || d_col.is_null(row) {
continue; // null endpoint = "no edge"
}
src_l.append_value(&es.src_label);
rels.append_value(&es.edge_type);
dst_l.append_value(&es.dst_label);
srcs.append_value(s_col.value(row));
dsts.append_value(d_col.value(row));
}
}
let batch = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(src_l.finish()),
Arc::new(rels.finish()),
Arc::new(dst_l.finish()),
Arc::new(srcs.finish()),
Arc::new(dsts.finish()),
],
)
.expect("edges-aggregate batch");
Dataset {
name: "edges".to_string(),
description: "aggregate edge list (src_label, rel, dst_label, src, dst)".to_string(),
schema,
batches: vec![batch],
}
}
// ── the Flight service ────────────────────────────────────────────────────────
type FlightStream<T> = Pin<Box<dyn futures::Stream<Item = Result<T, Status>> + Send + 'static>>;
/// The Arrow Flight [`FlightService`] over a [`Catalog`]. Read-only: it
/// advertises datasets and streams their batches; the mutating verbs (`do_put`,
/// `do_exchange`, `do_action`) are unimplemented.
#[derive(Clone)]
pub struct GraphFlightService {
catalog: Arc<Catalog>,
}
impl GraphFlightService {
pub fn new(catalog: Catalog) -> Self {
Self {
catalog: Arc::new(catalog),
}
}
pub fn catalog(&self) -> &Catalog {
&self.catalog
}
/// Build a [`FlightInfo`] for one dataset: schema, a `Cmd` descriptor and a
/// `Ticket`, both the dataset name's UTF-8 bytes.
fn flight_info(name: &str, ds: &Dataset) -> Result<FlightInfo, Status> {
let info = FlightInfo::new()
.try_with_schema(ds.schema.as_ref())
.map_err(|e| Status::internal(format!("encode schema for {name}: {e}")))?
.with_descriptor(FlightDescriptor::new_cmd(name.as_bytes().to_vec()))
.with_endpoint(FlightEndpoint::new().with_ticket(Ticket::new(name.as_bytes().to_vec())))
.with_total_records(ds.total_rows() as i64)
.with_total_bytes(-1);
Ok(info)
}
/// The dataset name a descriptor addresses: its `cmd` bytes if present, else
/// its path joined by `/`.
fn descriptor_name(desc: &FlightDescriptor) -> Result<String, Status> {
if !desc.cmd.is_empty() {
return String::from_utf8(desc.cmd.to_vec())
.map_err(|_| Status::invalid_argument("descriptor cmd is not UTF-8"));
}
if !desc.path.is_empty() {
return Ok(desc.path.join("/"));
}
Err(Status::invalid_argument(
"descriptor carries neither cmd nor path",
))
}
}
#[tonic::async_trait]
impl FlightService for GraphFlightService {
type HandshakeStream = FlightStream<HandshakeResponse>;
type ListFlightsStream = FlightStream<FlightInfo>;
type DoGetStream = FlightStream<FlightData>;
type DoPutStream = FlightStream<PutResult>;
type DoExchangeStream = FlightStream<FlightData>;
type DoActionStream = FlightStream<arrow_flight::Result>;
type ListActionsStream = FlightStream<ActionType>;
async fn list_flights(
&self,
_request: Request<Criteria>,
) -> Result<Response<Self::ListFlightsStream>, Status> {
let infos: Vec<Result<FlightInfo, Status>> = self
.catalog
.datasets()
.map(|ds| Self::flight_info(&ds.name, ds))
.collect();
let stream = futures::stream::iter(infos);
Ok(Response::new(Box::pin(stream)))
}
async fn get_flight_info(
&self,
request: Request<FlightDescriptor>,
) -> Result<Response<FlightInfo>, Status> {
let desc = request.into_inner();
let name = Self::descriptor_name(&desc)?;
let ds = self
.catalog
.get(&name)
.ok_or_else(|| Status::not_found(format!("no such dataset: {name}")))?;
Ok(Response::new(Self::flight_info(&name, ds)?))
}
async fn get_schema(
&self,
request: Request<FlightDescriptor>,
) -> Result<Response<SchemaResult>, Status> {
let desc = request.into_inner();
let name = Self::descriptor_name(&desc)?;
let ds = self
.catalog
.get(&name)
.ok_or_else(|| Status::not_found(format!("no such dataset: {name}")))?;
let options = arrow::ipc::writer::IpcWriteOptions::default();
let result = SchemaResult::try_from(SchemaAsIpc::new(ds.schema.as_ref(), &options))
.map_err(|e| Status::internal(format!("encode schema for {name}: {e}")))?;
Ok(Response::new(result))
}
async fn do_get(
&self,
request: Request<Ticket>,
) -> Result<Response<Self::DoGetStream>, Status> {
let ticket = request.into_inner();
let name = String::from_utf8(ticket.ticket.to_vec())
.map_err(|_| Status::invalid_argument("ticket is not UTF-8"))?;
let ds = self
.catalog
.get(&name)
.ok_or_else(|| Status::not_found(format!("no such dataset: {name}")))?;
let batches = ds.batches.clone();
let stream = FlightDataEncoderBuilder::new()
.with_schema(ds.schema.clone())
.build(futures::stream::iter(batches.into_iter().map(Ok)))
.map(|r| r.map_err(|e: FlightError| Status::internal(e.to_string())));
Ok(Response::new(Box::pin(stream)))
}
// ── unimplemented surface — this is a read-only graph server ──────────────
async fn handshake(
&self,
_request: Request<Streaming<HandshakeRequest>>,
) -> Result<Response<Self::HandshakeStream>, Status> {
Err(Status::unimplemented("handshake"))
}
async fn poll_flight_info(
&self,
_request: Request<FlightDescriptor>,
) -> Result<Response<PollInfo>, Status> {
Err(Status::unimplemented("poll_flight_info"))
}
async fn do_put(
&self,
_request: Request<Streaming<FlightData>>,
) -> Result<Response<Self::DoPutStream>, Status> {
Err(Status::unimplemented("do_put: read-only graph server"))
}
async fn do_exchange(
&self,
_request: Request<Streaming<FlightData>>,
) -> Result<Response<Self::DoExchangeStream>, Status> {
Err(Status::unimplemented("do_exchange"))
}
async fn do_action(
&self,
_request: Request<Action>,
) -> Result<Response<Self::DoActionStream>, Status> {
Err(Status::unimplemented("do_action"))
}
async fn list_actions(
&self,
_request: Request<Empty>,
) -> Result<Response<Self::ListActionsStream>, Status> {
let stream = futures::stream::iter(Vec::<Result<ActionType, Status>>::new());
Ok(Response::new(Box::pin(stream)))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn catalog_advertises_nodes_edges_and_tables() {
let cat = Catalog::from_bench(1);
assert!(cat.get("nodes").is_some(), "nodes flight present");
assert!(cat.get("edges").is_some(), "edges flight present");
assert!(
cat.get("edge/Employee_REPORTS_TO_Employee").is_some(),
"self-ref edge flight present"
);
assert!(
cat.get("table/erp.hr.employees").is_some(),
"source table flight present"
);
// Aggregate node skeleton: 1106 distinct nodes at scale 1.
let nodes = cat.get("nodes").unwrap();
assert_eq!(nodes.total_rows(), 1106, "scale-1 node total");
// Aggregate edges: 2955 resolved edges at scale 1 (null endpoints dropped).
let edges = cat.get("edges").unwrap();
assert_eq!(edges.total_rows(), 2955, "scale-1 resolved-edge total");
}
#[test]
fn edges_aggregate_has_no_self_loops() {
// The a7208dc self-ref fix: no REPORTS_TO edge points a node at itself.
let cat = Catalog::from_bench(1);
let edges = cat.get("edges").unwrap();
let b = &edges.batches[0];
let src_l = b
.column(b.schema().index_of("src_label").unwrap())
.as_any()
.downcast_ref::<arrow_array::StringArray>()
.unwrap();
let dst_l = b
.column(b.schema().index_of("dst_label").unwrap())
.as_any()
.downcast_ref::<arrow_array::StringArray>()
.unwrap();
let src = b
.column(b.schema().index_of("src").unwrap())
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
let dst = b
.column(b.schema().index_of("dst").unwrap())
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
let mut reports_to = 0;
for i in 0..b.num_rows() {
if src_l.value(i) == dst_l.value(i) {
assert_ne!(src.value(i), dst.value(i), "no self-loop edge");
}
if b.column(b.schema().index_of("rel").unwrap())
.as_any()
.downcast_ref::<arrow_array::StringArray>()
.unwrap()
.value(i)
== "REPORTS_TO"
{
reports_to += 1;
}
}
assert_eq!(reports_to, 21, "REPORTS_TO edges (24 employees - 3 roots)");
}
}