nordisk/skade

skade-flight/src/encode.rs

//! Parallel Flight encoding.
//!
//! `DataFrame::execute_stream` coalesces every partition into ONE stream and
//! arrow-flight's encoder then IPC-encodes every batch on the request task —
//! one core, however many the scan used. MEASURED 2026-09-09 (2M rows × 6
//! cols, 16 partitions): the engine alone drains at 61M rows/s, over Flight
//! 23M rows/s. The serial encode was the ceiling.
//!
//! Here each partition's stream is encoded on its own task and the encoded
//! frames are merged through a bounded channel. A `FlightData` is a pair of
//! `Bytes` (header, body), so the merge moves refcounts, not bytes: zero
//! copy between the encoder and the wire. The schema frame is emitted once,
//! up front; the per-partition encoders' own schema frames are dropped.
//!
//! LAW: no serial CPU wait. Tasks here are the SAME tokio runtime DataFusion
//! already fans its partitions across — no second thread pool.

use std::pin::Pin;
use std::sync::Arc;

use arrow_array::RecordBatch;
use arrow_flight::encode::FlightDataEncoderBuilder;
use arrow_flight::error::FlightError;
use arrow_flight::{FlightData, SchemaAsIpc};
use arrow_schema::SchemaRef;
use datafusion::execution::SendableRecordBatchStream;
use futures::{Stream, StreamExt, TryStreamExt};
use tonic::Status;

pub type FrameStream = Pin<Box<dyn Stream<Item = Result<FlightData, Status>> + Send + 'static>>;

/// Frames in flight between the encoders and the wire, per partition.
const CHANNEL_DEPTH_PER_PARTITION: usize = 2;

fn is_schema_frame(fd: &FlightData) -> bool {
    arrow_ipc::root_as_message(&fd.data_header)
        .map(|m| m.header_type() == arrow_ipc::MessageHeader::Schema)
        .unwrap_or(false)
}

/// The IPC *Message* carrying the schema — what the wire wants. (`IpcMessage`
/// is the bare schema flatbuffer `FlightInfo.schema` carries; a decoder fed
/// that as a frame fails with "unaligned".)
fn schema_frame(schema: &SchemaRef) -> FlightData {
    let opts = arrow_ipc::writer::IpcWriteOptions::default();
    FlightData::from(SchemaAsIpc::new(schema.as_ref(), &opts))
}

fn to_flight_err(e: datafusion::error::DataFusionError) -> FlightError {
    FlightError::ExternalError(Box::new(e))
}

fn to_status(e: FlightError) -> Status {
    match e {
        FlightError::ExternalError(inner) => Status::internal(inner.to_string()),
        other => Status::internal(other.to_string()),
    }
}

/// Encode `partitions` into one Flight stream. One partition → the plain
/// serial encoder (no task, no channel). Many → parallel.
pub fn encode_partitions(schema: SchemaRef, mut partitions: Vec<SendableRecordBatchStream>) -> Result<FrameStream, Status> {
    if partitions.len() <= 1 {
        let batches: Pin<Box<dyn Stream<Item = Result<RecordBatch, FlightError>> + Send>> = match partitions.pop() {
            Some(s) => Box::pin(s.map_err(to_flight_err)),
            None => Box::pin(futures::stream::empty()),
        };
        let out = FlightDataEncoderBuilder::new().with_schema(schema).build(batches).map_err(to_status);
        return Ok(Box::pin(out));
    }

    let n = partitions.len();
    let (tx, rx) = tokio::sync::mpsc::channel::<Result<FlightData, Status>>(n * CHANNEL_DEPTH_PER_PARTITION);
    let schema = Arc::clone(&schema);
    for part in partitions {
        let tx = tx.clone();
        let schema = Arc::clone(&schema);
        tokio::spawn(async move {
            let mut enc = FlightDataEncoderBuilder::new()
                .with_schema(schema)
                .build(part.map_err(to_flight_err))
                .map_err(to_status);
            while let Some(item) = enc.next().await {
                let item = match item {
                    Ok(fd) if is_schema_frame(&fd) => continue,
                    other => other,
                };
                let failed = item.is_err();
                if tx.send(item).await.is_err() || failed {
                    // Receiver gone (client hung up) or we reported an error:
                    // stop this partition; dropping `part` cancels its scan.
                    break;
                }
            }
        });
    }
    drop(tx);
    let first = schema_frame(&schema);
    let merged = tokio_stream_receiver(rx);
    Ok(Box::pin(futures::stream::once(async move { Ok(first) }).chain(merged)))
}

fn tokio_stream_receiver(rx: tokio::sync::mpsc::Receiver<Result<FlightData, Status>>) -> impl Stream<Item = Result<FlightData, Status>> + Send {
    futures::stream::unfold(rx, |mut rx| async move { rx.recv().await.map(|item| (item, rx)) })
}