polyc-query 2026.8.3

//! Shared row-serialization helper for every surface `crate::engine::QueryEngine`
//! is mounted behind.
//!
//! `output_to_json` is the ONE conversion from a `crate::engine::QueryOutput`'s
//! `RecordBatch`es to the wire JSON envelope (`{"columns": [...], "rows":
//! [[...]], "truncated": bool, "skipped_partitions": integer}`) every caller
//! returns — the control plane's `POST /api/query` handler and the
//! `polychrome_query` MCP tool both call this, never two hand-rolled versions
//! (the repo's shared-helper rule for user-facing output: see
//! `crate::humanize_tool_name`-style precedent).
//!
//! `skipped_partitions` (QRY-7) is NOT read off `QueryOutput` the way
//! `truncated` is: it is authority-level information (how many Fleet
//! partitions `ScopedQuery::replay_scoped_partitions`
//! skipped as unreadable during replay), known BEFORE
//! `QueryEngine::build` ever runs, not something the engine
//! itself computes — so `output_to_json` takes it as an explicit parameter
//! rather than reading it off `output`, and [`crate::authority::ScopedQuery::execute`]
//! is the one caller that supplies a genuinely nonzero value.
//!
//! Column-typed values are encoded through `arrow-json`'s own
//! [`arrow::json::ArrayWriter`] rather than a hand-rolled per-`DataType`
//! match — that writer already carries the full, tested `RecordBatch` →
//! JSON mapping (dates, decimals, nested lists, ...) this crate has no
//! reason to reimplement. Row order is then rebuilt from the batch's own
//! schema field order (not the parsed JSON object's key order, which
//! `serde_json::Map` does not guarantee to preserve without the
//! `preserve_order` feature this crate does not enable).

use arrow::json::ArrayWriter;
use arrow::record_batch::RecordBatch;
use serde::Serialize;

use crate::engine::QueryOutput;

/// Failure converting a [`QueryOutput`]'s Arrow batches to the wire JSON
/// envelope.
#[derive(Debug, thiserror::Error)]
pub(crate) enum OutputJsonError {
    /// `arrow-json`'s writer failed to encode a batch (e.g. an unsupported
    /// `DataType` for JSON output).
    #[error("failed to encode result rows as JSON: {0}")]
    Arrow(#[from] arrow::error::ArrowError),
    /// The bytes `arrow-json` wrote did not parse back as JSON — should be
    /// unreachable in practice (the writer only ever emits valid JSON), but
    /// surfaced rather than unwrapped so a caller sees a typed error instead
    /// of a panic.
    #[error("failed to parse the encoded JSON rows: {0}")]
    Parse(#[from] serde_json::Error),
}

/// The wire JSON envelope every query surface returns.
///
/// Column names in schema order, one row per result row (each cell in the
/// same column order), whether [`crate::QueryLimits::row_cap`] truncated the
/// real result, and — for a `QueryScope::Fleet` query —
/// whether every partition the caller expected actually replayed.
#[derive(Debug, Clone, Serialize)]
pub struct QueryResultJson {
    /// Column names, in the result schema's own order.
    pub columns: Vec<String>,
    /// One entry per result row; each row's values are ordered to match
    /// [`QueryResultJson::columns`].
    pub rows: Vec<Vec<serde_json::Value>>,
    /// Mirrors `QueryOutput::truncated`.
    pub truncated: bool,
    /// How many partitions `ScopedQuery::replay_scoped_partitions`
    /// skipped as unreadable during replay (QRY-7) — always `0` for a
    /// non-`Fleet` scope, since a `Conversations` replay failure is a hard
    /// error there, never a skip-and-continue. A nonzero value here is the
    /// signal a Fleet caller needs to tell "this result covers every
    /// conversation partition" apart from "this result is missing whatever
    /// this many partitions couldn't be read" — the identical completeness
    /// gap `truncated` closes for the ROW cap, applied here to the
    /// PARTITION set instead. Before this field existed, an unreadable
    /// partition was only `tracing::warn!`-logged server-side, with no
    /// signal at all reaching the caller.
    pub skipped_partitions: usize,
}

/// Convert `output`'s collected `RecordBatch`es into the shared wire JSON
/// envelope, stamping `skipped_partitions` (see [`QueryResultJson::skipped_partitions`]'s
/// doc for why that count is a caller-supplied parameter rather than
/// something read off `output` itself).
///
/// `columns` always comes from `output.schema` — the query's PLANNED schema,
/// read off the `DataFrame` before `.collect()` ran (see
/// [`crate::engine::QueryOutput::schema`]) — never from `batches.first()`.
/// A query whose plan eliminates every row (an inner join with no matches,
/// for example) still has a real, known column list; deriving `columns` from
/// the first batch instead made that legitimate empty result
/// indistinguishable from "no schema could be determined" (issue #1916). An
/// `output` with zero batches therefore yields `columns` populated from the
/// schema and an empty `rows`, not empty `columns` too.
///
/// # Errors
///
/// Returns [`OutputJsonError`] if `arrow-json`'s writer fails to encode a
/// batch, or if (unreachably in practice) the bytes it wrote fail to parse
/// back as JSON.
pub(crate) fn output_to_json(
    output: &QueryOutput,
    skipped_partitions: usize,
) -> Result<QueryResultJson, OutputJsonError> {
    let columns: Vec<String> = output
        .schema
        .fields()
        .iter()
        .map(|field| field.name().clone())
        .collect();

    if output.batches.is_empty() {
        return Ok(QueryResultJson {
            columns,
            rows: Vec::new(),
            truncated: output.truncated,
            skipped_partitions,
        });
    }

    let refs: Vec<&RecordBatch> = output.batches.iter().collect();
    let mut buf = Vec::new();
    {
        let mut writer = ArrayWriter::new(&mut buf);
        writer.write_batches(&refs)?;
        writer.finish()?;
    }
    let parsed: Vec<serde_json::Value> = serde_json::from_slice(&buf)?;

    let rows = parsed
        .into_iter()
        .map(|row| {
            let mut obj = match row {
                serde_json::Value::Object(map) => map,
                // Every row `arrow-json` writes is a JSON object keyed by
                // column name; anything else would be a bug in that writer,
                // not a shape this crate needs to preserve. Fall back to an
                // all-null row rather than panicking on it.
                _ => serde_json::Map::new(),
            };
            columns
                .iter()
                .map(|column| obj.remove(column).unwrap_or(serde_json::Value::Null))
                .collect()
        })
        .collect();

    Ok(QueryResultJson {
        columns,
        rows,
        truncated: output.truncated,
        skipped_partitions,
    })
}

#[cfg(test)]
mod tests {
    use arrow::array::{Int64Array, StringArray};
    use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
    use std::sync::Arc;

    use super::*;

    fn schema() -> SchemaRef {
        Arc::new(Schema::new(vec![
            Field::new("id", DataType::Int64, false),
            Field::new("name", DataType::Utf8, false),
        ]))
    }

    fn batch(ids: &[i64], names: &[&str]) -> RecordBatch {
        RecordBatch::try_new(
            schema(),
            vec![
                Arc::new(Int64Array::from(ids.to_vec())),
                Arc::new(StringArray::from(names.to_vec())),
            ],
        )
        .expect("build fixture batch")
    }

    #[test]
    fn converts_columns_and_rows_in_schema_order() {
        let output = QueryOutput {
            batches: vec![batch(&[1, 2], &["a", "b"])],
            truncated: false,
            schema: schema(),
        };
        let json = output_to_json(&output, 0).expect("convert");
        assert_eq!(json.columns, vec!["id", "name"]);
        assert_eq!(
            json.rows,
            vec![
                vec![serde_json::json!(1), serde_json::json!("a")],
                vec![serde_json::json!(2), serde_json::json!("b")],
            ]
        );
        assert!(!json.truncated);
        assert_eq!(json.skipped_partitions, 0);
    }

    #[test]
    fn multiple_batches_concatenate_rows_in_order() {
        let output = QueryOutput {
            batches: vec![batch(&[1], &["a"]), batch(&[2], &["b"])],
            truncated: true,
            schema: schema(),
        };
        let json = output_to_json(&output, 0).expect("convert");
        assert_eq!(json.rows.len(), 2);
        assert_eq!(json.rows[0][0], serde_json::json!(1));
        assert_eq!(json.rows[1][0], serde_json::json!(2));
        assert!(json.truncated);
    }

    /// Pins the fix for #1916: a query whose plan produced zero batches (an
    /// inner join that eliminates every row, for example) still has a real,
    /// known column list — `columns` comes from the planned schema, not from
    /// `batches.first()`, so it must NOT come back empty just because there
    /// were no rows to report.
    #[test]
    fn zero_batches_still_yields_columns_from_the_planned_schema() {
        let output = QueryOutput {
            batches: Vec::new(),
            truncated: false,
            schema: schema(),
        };
        let json = output_to_json(&output, 0).expect("convert");
        assert_eq!(json.columns, vec!["id", "name"]);
        assert!(json.rows.is_empty());
        assert_eq!(json.skipped_partitions, 0);
    }

    /// QRY-7: `skipped_partitions` is a caller-supplied stamp, not derived
    /// from `output` itself — pins that `output_to_json` threads whatever
    /// value the caller passes straight onto the envelope, for both the
    /// zero-batches early return and the normal path.
    #[test]
    fn skipped_partitions_is_stamped_from_the_caller_supplied_count() {
        let with_rows = QueryOutput {
            batches: vec![batch(&[1], &["a"])],
            truncated: false,
            schema: schema(),
        };
        let json = output_to_json(&with_rows, 3).expect("convert");
        assert_eq!(json.skipped_partitions, 3);

        let zero_batches = QueryOutput {
            batches: Vec::new(),
            truncated: false,
            schema: schema(),
        };
        let json = output_to_json(&zero_batches, 3).expect("convert");
        assert_eq!(
            json.skipped_partitions, 3,
            "the zero-batches early return must still stamp the caller's count"
        );
    }
}