polyc-query 2026.9.0

//! Encodes released Arrow batches into bounded protocol frames.
//!
//! The protocol carries opaque Arrow IPC bytes in two frame kinds. This module
//! is the only place that produces them, so the framing rule is stated once:
//!
//! - a schema frame is an IPC stream that contains the schema message alone;
//! - a data frame is an IPC stream that contains the schema message and
//!   exactly one record batch.
//!
//! Each data frame therefore decodes on its own. A consumer that drops a frame
//! cannot silently mis-decode the next one against a schema it never saw.
//!
//! Slicing is deterministic. The same batch and the same byte ceiling always
//! produce the same frames, because the split points come only from the batch
//! and the ceiling. Nothing here consults a clock, a random source, or an
//! allocation address.
//!
//! Framing is incremental. One batch is held by reference with a cursor into
//! it, and one frame is produced on demand, so the encoded bytes alive here do
//! not grow with the number of frames the batch will yield. Two allocations
//! are alive at once: the frame being built, at most `ceiling_bytes`, and the
//! search candidate that overshoots it.
//!
//! The candidate is bounded in ROWS, not in bytes: the growth below tries at
//! most twice the rows of the largest fitting prefix. Rows are not uniform in
//! width, so that bound does not convert. A single row wider than the ceiling
//! is encoded in full before it can be refused, and its size is exactly what
//! `RowTooLarge` reports. The byte bound comes from the other side — every
//! candidate is a prefix of one batch, and execution holds each released batch
//! under the caller's `result_release_bytes`.
//!
//! Encoding a whole batch up front had no such bound. A small `frame_bytes`
//! multiplies the repeated schema message across many frames, so a caller
//! could make the server hold far more than the release ceiling it admitted
//! and then release only a bounded prefix of it.
//!
//! A single row that cannot fit the ceiling is refused. The encoder never
//! splits a row across frames and never releases an over-bound frame — the
//! schema frame included, which is why `start` measures it against the same
//! ceiling every data frame is measured against.
//!
//! The ceiling is the caller's. A ceiling too small for this result's schema,
//! or for one of its rows, is the caller's bound rather than this deployment's
//! fault, and [`FrameEncodeError::class`] says so. Reporting it as an internal
//! failure would put a server-fault log line under caller-chosen input.

use arrow::array::RecordBatch;
use arrow::datatypes::SchemaRef;
use arrow::ipc::writer::{IpcWriteOptions, StreamWriter};
use polyc_query_model::{DataFrame, ErrorClass, ModelError, SchemaFrame};

/// Failure encoding one released batch.
#[derive(Debug, thiserror::Error)]
pub enum FrameEncodeError {
    /// Arrow could not encode the batch or the schema.
    #[error("the result could not be encoded")]
    Encode,
    /// The encoded schema alone exceeds the frame ceiling.
    ///
    /// Reports the schema's own encoded size and the ceiling. A schema names
    /// columns the caller selected, so neither number is content: both are
    /// sizes the caller already chose to request.
    #[error("the schema needs {schema_bytes} bytes, over the {ceiling_bytes} byte frame ceiling")]
    SchemaTooLarge {
        /// The encoded size of the schema frame that did not fit.
        schema_bytes: usize,
        /// The effective frame ceiling.
        ceiling_bytes: usize,
    },
    /// One row alone exceeds the frame ceiling, so no split can help.
    ///
    /// Reports the row's own encoded size and the ceiling. Neither number is
    /// caller content: both are sizes the caller already chose to request.
    #[error("one row needs {row_bytes} bytes, over the {ceiling_bytes} byte frame ceiling")]
    RowTooLarge {
        /// The encoded size of the single row that did not fit.
        row_bytes: usize,
        /// The effective frame ceiling.
        ceiling_bytes: usize,
    },
    /// The encoded frame failed the protocol's own bound check.
    #[error("the encoded frame is outside its protocol bound")]
    Model(#[from] ModelError),
}

impl FrameEncodeError {
    /// Returns who is responsible for this failure.
    ///
    /// A ceiling too small for the schema or for one row is the caller's
    /// bound: the caller chose `frame_bytes`, and no work this deployment
    /// does can make the result fit it. Everything else is this deployment's.
    #[must_use]
    pub const fn class(&self) -> ErrorClass {
        match self {
            Self::SchemaTooLarge { .. } | Self::RowTooLarge { .. } => ErrorClass::Bounds,
            Self::Encode | Self::Model(_) => ErrorClass::Internal,
        }
    }
}

/// Produces ordered protocol frames from released batches.
///
/// `Debug` reports the ceiling and the next sequence number. A schema names
/// columns a caller selected, so the schema itself never reaches a log line
/// through this type.
pub(crate) struct FrameEncoder {
    schema: SchemaRef,
    ceiling_bytes: usize,
    next_sequence: u64,
    /// The batch being framed, and how many of its rows are already framed.
    ///
    /// Held by reference: `RecordBatch` is an `Arc` over its columns, so the
    /// cursor costs one clone of that handle and no row data.
    cursor: Option<(RecordBatch, usize)>,
    /// How many frames this encoder has built.
    ///
    /// Only a case reads it, to prove that framing stops where release stops
    /// rather than running ahead of it.
    #[cfg(test)]
    built: u64,
}

impl std::fmt::Debug for FrameEncoder {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("FrameEncoder")
            .field("ceiling_bytes", &self.ceiling_bytes)
            .field("next_sequence", &self.next_sequence)
            .finish_non_exhaustive()
    }
}

impl FrameEncoder {
    /// Builds the encoder and the schema frame that must precede every batch.
    ///
    /// # Errors
    ///
    /// Returns [`FrameEncodeError::Encode`] if Arrow cannot encode the schema,
    /// [`FrameEncodeError::Model`] if the encoded schema is outside the
    /// protocol's own frame bound, and [`FrameEncodeError::SchemaTooLarge`] if
    /// it is outside the caller's.
    pub(crate) fn start(
        schema: SchemaRef,
        ceiling_bytes: usize,
    ) -> Result<(Self, SchemaFrame), FrameEncodeError> {
        let encoded = encode(&schema, None)?;
        // Against the caller's ceiling, not only the protocol's. The pump
        // refuses any frame over the caller's bound, so releasing this one
        // would break the stream contract from inside the server and report
        // an internal fault for a bound the caller chose.
        if encoded.len() > ceiling_bytes {
            return Err(FrameEncodeError::SchemaTooLarge {
                schema_bytes: encoded.len(),
                ceiling_bytes,
            });
        }
        let frame = SchemaFrame::try_new(encoded)?;
        Ok((
            Self {
                schema,
                ceiling_bytes,
                next_sequence: 0,
                cursor: None,
                #[cfg(test)]
                built: 0,
            },
            frame,
        ))
    }

    /// Takes one released batch to frame.
    ///
    /// An empty batch is ignored: a frame carrying zero rows would consume a
    /// sequence number and tell a consumer nothing. Any batch still being
    /// framed is replaced, which cannot happen — the caller drains one batch
    /// before it asks the stream for another.
    pub(crate) fn begin(&mut self, batch: RecordBatch) {
        if batch.num_rows() > 0 {
            self.cursor = Some((batch, 0));
        }
    }

    /// Builds the next frame of the batch in hand, or `None` when it is done.
    ///
    /// Exactly one frame is encoded per call, so a caller that stops early
    /// pays for the frames it took and one search, never for the whole batch.
    ///
    /// # Errors
    ///
    /// Returns [`FrameEncodeError::RowTooLarge`] when one row alone exceeds
    /// the ceiling, [`FrameEncodeError::Encode`] on an Arrow failure, and
    /// [`FrameEncodeError::Model`] when an encoded frame is outside the
    /// protocol's own bound.
    pub(crate) fn next(&mut self) -> Result<Option<DataFrame>, FrameEncodeError> {
        let Some((batch, offset)) = self.cursor.take() else {
            return Ok(None);
        };
        let remaining = batch.num_rows() - offset;
        let (encoded, rows) = self.fit(&batch, offset, remaining)?;
        let sequence = self.next_sequence;
        self.next_sequence = self.next_sequence.saturating_add(1);
        #[cfg(test)]
        {
            self.built = self.built.saturating_add(1);
        }
        let advanced = offset + rows;
        if advanced < batch.num_rows() {
            self.cursor = Some((batch, advanced));
        }
        let row_count = u64::try_from(rows).unwrap_or(u64::MAX);
        DataFrame::try_new(sequence, row_count, encoded)
            .map(Some)
            .map_err(FrameEncodeError::from)
    }

    /// Returns whether a batch is still being framed.
    pub(crate) const fn has_rows(&self) -> bool {
        self.cursor.is_some()
    }

    /// Drops the batch in hand without framing the rest of it.
    pub(crate) fn discard(&mut self) {
        self.cursor = None;
    }

    /// Returns how many frames this encoder has built.
    #[cfg(test)]
    pub(crate) const fn built(&self) -> u64 {
        self.built
    }

    /// Finds a prefix of `remaining` rows that fits the ceiling.
    ///
    /// Grows from one row, doubling while the encoding still fits, and takes
    /// the last length that did. Growing rather than halving is what keeps the
    /// buffer from following the batch: starting at `remaining` encodes the
    /// whole remainder just to discover it does not fit. The candidate that
    /// overshoots here carries at most twice the rows of the one that fits —
    /// exactly twice whenever the remainder allows it.
    ///
    /// The price is encodes. A batch that fits one frame took one encode when
    /// the search started at the whole remainder; it now takes about
    /// `log2(rows)` of them, summing to roughly twice the batch's bytes. That
    /// is the cost of not holding the whole remainder at once.
    ///
    /// The lengths tried depend only on `remaining` and the data, so the split
    /// points are the same on every run.
    fn fit(
        &self,
        batch: &RecordBatch,
        offset: usize,
        remaining: usize,
    ) -> Result<(Vec<u8>, usize), FrameEncodeError> {
        let mut length = 1;
        let mut fitting: Option<(Vec<u8>, usize)> = None;
        loop {
            let candidate = batch.slice(offset, length);
            let encoded = encode(&self.schema, Some(&candidate))?;
            if encoded.len() > self.ceiling_bytes {
                return fitting.ok_or(FrameEncodeError::RowTooLarge {
                    row_bytes: encoded.len(),
                    ceiling_bytes: self.ceiling_bytes,
                });
            }
            if length == remaining {
                return Ok((encoded, length));
            }
            fitting = Some((encoded, length));
            length = length.saturating_mul(2).min(remaining);
        }
    }
}

/// Encodes one self-contained IPC stream.
///
/// `batch` of `None` produces the schema-only stream the protocol's first
/// frame carries.
fn encode(schema: &SchemaRef, batch: Option<&RecordBatch>) -> Result<Vec<u8>, FrameEncodeError> {
    let options = IpcWriteOptions::default();
    let mut writer = StreamWriter::try_new_with_options(Vec::new(), schema, options)
        .map_err(|_arrow| FrameEncodeError::Encode)?;
    if let Some(batch) = batch {
        writer
            .write(batch)
            .map_err(|_arrow| FrameEncodeError::Encode)?;
    }
    writer.finish().map_err(|_arrow| FrameEncodeError::Encode)?;
    writer
        .into_inner()
        .map_err(|_arrow| FrameEncodeError::Encode)
}