polyc-query 2026.9.0

//! `usage` typed-table decoder — `UsageEvent` → Arrow columns.
//!
//! The phase-1 typed table (see `super` module docs). Backs
//! `polychrome.events.v1.UsageEvent`'s two fields, `input_tokens` and
//! `output_tokens` (`crates/proto/proto/events.proto:118-121`), decoded
//! straight from the payload rather than dug out of an opaque blob —
//! #1311's typed-table pattern (part of the #1178 epic), matching the
//! design's "queried with an ordinary `SUM`/`GROUP BY` rather than digging
//! a numeric field out of an opaque blob" rationale. Both proto fields are `uint64`, and the generated Rust
//! type mirrors that as `u64` (`crates/control-plane/src/forensics.rs`
//! mirrors the same fields as `u64` at its own seam) — so both token Arrow
//! columns are `UInt64`, non-null.
//!
//! # Uniform keys (#1311)
//!
//! Every typed table must carry the fact model's uniform key columns so a
//! row can be joined back to the conversation and turn it belongs to — the
//! same `partition`/`position`/`turn_id` triple the wide `events`/
//! `events_raw` tables carry (`crate::decode::events_batch`). `kind_base` is
//! deliberately omitted here: every row in this table is `kind_base =
//! "usage"` by construction (that is the table's filter), so a constant
//! column would carry no information. `conversation_id` is also omitted as
//! its own column — `partition` already *is* the conversation id
//! (`conv-{id}`), so a second column would just duplicate it.
//!
//! `event_time` — a uniform attested timestamp — is deliberately **not**
//! included. No such column exists on `Event` yet; inventing one (a fake
//! decode-time clock read, or an empty always-null column) would misrepresent
//! the data. Attested event time is tracked as future work in #1327; once it
//! lands, it joins this table's key set alongside `partition`/`position`/
//! `turn_id`.
//!
//! `partition`, `position`, and `turn_id` are derived exactly the way
//! [`crate::decode::events_batch`] derives them — via
//! [`polyc_proto::kinds::parse`], the platform's single kind-grammar owner —
//! reused here, not re-implemented, so the two tables can never disagree on
//! what a row's key columns mean. [`decode_usage_events`] itself delegates
//! the filter/derive/decode loop to `crate::decode::decode_typed_kind_events`
//! — this module established the pattern; `crate::decode::model_call` is
//! the pattern's second instance, and both call through the same shared
//! loop rather than each carrying its own copy — see "Shared decode" below
//! for what `decode_typed_kind_events` decodes THROUGH.
//!
//! Decode quirk this module handles deliberately: a genuine zero-usage
//! event (e.g. a dispatch that short-circuited before any model call ran)
//! encodes to an EMPTY payload — proto3 elides all-default scalar fields,
//! and `input_tokens: 0, output_tokens: 0` is `UsageEvent::default()`, so
//! there is nothing to write. `buffa::Message::decode_from_slice` (via
//! [`polyc_proto::events_decode::try_decode_event_payload`]) decodes an
//! empty slice as the all-defaults message cleanly — no error — so a
//! genuine zero-usage row round-trips to `(0, 0)` exactly like any other
//! row, not as a hole in the batch. A NON-empty payload that fails to
//! decode, by contrast, is corruption or a schema mismatch, never a
//! legitimate zero-usage event; [`decode_usage_events`] SKIPS it rather
//! than surfacing an error or panicking, matching
//! `polyc_proto::events_decode::decode_event_payload`'s own "`None` on any
//! decode failure" contract already used everywhere else a stored payload
//! is rendered (e.g. forensics' `pretty_json`) — one malformed row
//! degrades the aggregate instead of failing the whole scan.
//!
//! # Shared decode (#1579)
//!
//! The bytes→fields step [`decode_usage_events`] passes to
//! `decode_typed_kind_events` as its `decode` argument is
//! [`polyc_facts::fold_usage_event`] — the same primitive the control
//! plane's per-turn accounting decodes through
//! (`crates/control-plane/src/grpc/mod.rs`'s `decode_usage_payload`) —
//! rather than a raw `try_decode_event_payload::<UsageEvent>` call.
//! `decode_typed_kind_events` takes `decode` as a plain
//! `Fn(&[u8]) -> Result<T, E>` (not a `T: buffa::Message` bound it resolves
//! internally) precisely so this module and
//! [`crate::decode::model_call`] can each plug in their own shared
//! `polyc_facts` fold while still sharing the loop's kind-filter/key-
//! derivation/skip-and-warn shape. What decode failure MEANS is still this
//! table's own policy, not the fold's: a corrupt non-empty payload is
//! skipped (logged) — see `decode_typed_kind_events`'s own docs.

use std::sync::Arc;

use arrow::array::{ArrayRef, StringBuilder, UInt64Builder};
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use arrow::error::ArrowError;
use arrow::record_batch::RecordBatch;
use polyc_eventlog::Event;
use polyc_proto::kinds;

/// One decoded `usage` row: the fact model's uniform key columns
/// (`partition`, `position`, `turn_id`) plus the decoded token counts.
///
/// Produced by [`decode_usage_events`] and consumed by
/// [`decode_usage_batch`] — kept as an explicit intermediate type (rather
/// than building the `RecordBatch` directly) so a caller can inspect or
/// filter decoded rows before materializing Arrow columns.
#[derive(Debug, Clone)]
pub(crate) struct UsageRow {
    /// The journal partition this row's event was read from (`conv-{id}`).
    pub partition: String,
    /// The journal's own monotonic append position for this event.
    pub position: u64,
    /// The `:{turn_uuid}` suffix off the event's `kind`, canonical
    /// hyphenated form, or `None` for a bare `usage` kind with no turn
    /// tagged.
    pub turn_id: Option<String>,
    /// The decoded fact ([`polyc_facts::fold_usage_event`]'s output).
    pub usage: polyc_facts::UsageFact,
}

/// The `usage` typed table's Arrow schema.
///
/// `partition` (`Utf8`, non-null), `position` (`UInt64`, non-null),
/// `turn_id` (`Utf8`, nullable), `input_tokens` (`UInt64`, non-null),
/// `output_tokens` (`UInt64`, non-null).
///
/// See the module docs for why `kind_base` and `conversation_id` are
/// deliberately omitted, and why `event_time` is deferred to #1327 rather
/// than faked.
#[must_use]
pub(crate) fn schema() -> SchemaRef {
    Arc::new(Schema::new(vec![
        Field::new("partition", DataType::Utf8, false),
        Field::new("position", DataType::UInt64, false),
        Field::new("turn_id", DataType::Utf8, true),
        Field::new("input_tokens", DataType::UInt64, false),
        Field::new("output_tokens", DataType::UInt64, false),
    ]))
}

/// Decode already-framed [`UsageRow`]s into the `usage` table's Arrow
/// `RecordBatch` (`partition`, `position`, `turn_id`, `input_tokens`,
/// `output_tokens` columns, in [`schema`] order).
///
/// # Errors
///
/// Returns [`ArrowError`] if Arrow batch construction fails.
pub(crate) fn decode_usage_batch(rows: &[UsageRow]) -> Result<RecordBatch, ArrowError> {
    let mut partition_b = StringBuilder::with_capacity(rows.len(), rows.len() * 8);
    let mut position_b = UInt64Builder::with_capacity(rows.len());
    let mut turn_id_b = StringBuilder::with_capacity(rows.len(), rows.len() * 36);
    let mut input_tokens = UInt64Builder::with_capacity(rows.len());
    let mut output_tokens = UInt64Builder::with_capacity(rows.len());

    for row in rows {
        partition_b.append_value(&row.partition);
        position_b.append_value(row.position);
        match &row.turn_id {
            Some(id) => turn_id_b.append_value(id),
            None => turn_id_b.append_null(),
        }
        input_tokens.append_value(row.usage.input_tokens);
        output_tokens.append_value(row.usage.output_tokens);
    }

    let columns: Vec<ArrayRef> = vec![
        Arc::new(partition_b.finish()),
        Arc::new(position_b.finish()),
        Arc::new(turn_id_b.finish()),
        Arc::new(input_tokens.finish()),
        Arc::new(output_tokens.finish()),
    ];
    RecordBatch::try_new(schema(), columns)
}

/// Filter `partition`'s framed `events` to kind-base `usage` rows, decode
/// each payload, and pair it with that row's key columns.
///
/// `partition` and each event's own `position` are stored verbatim;
/// `turn_id` is derived via [`polyc_proto::kinds::parse`] — the platform's
/// single kind-grammar owner, reused here exactly as
/// [`crate::decode::events_batch`] reuses it, not re-derived — so a
/// `usage:{turn_uuid}` kind yields `turn_id = Some(turn_uuid)` and a bare
/// `usage` kind yields `turn_id = None`.
///
/// A non-empty payload that fails to decode is silently skipped (logged) —
/// see the module docs for why that (rather than an empty payload, which
/// decodes cleanly to an all-zero fact) is the corruption signal.
///
/// Delegates the filter/derive/decode loop to
/// `crate::decode::decode_typed_kind_events` (see that function's docs and
/// the module docs' "Shared decode" section) — only the
/// [`polyc_facts::fold_usage_event`] `decode` argument and the
/// [`polyc_facts::UsageFact`] → [`UsageRow`] mapping are unique to this
/// module.
#[must_use]
pub(crate) fn decode_usage_events(partition: &str, events: &[(u64, Event)]) -> Vec<UsageRow> {
    crate::decode::decode_typed_kind_events(
        partition,
        events,
        &[kinds::USAGE],
        "usage",
        polyc_facts::fold_usage_event,
        |partition, position, turn_id, usage| UsageRow {
            partition,
            position,
            turn_id,
            usage,
        },
    )
}

#[cfg(test)]
mod tests {
    use arrow::array::Array as _;
    use buffa::Message as _;
    use polyc_proto::proto::polychrome::events::v1::UsageEvent;
    use uuid::Uuid;

    use super::*;

    #[test]
    fn schema_shape() {
        let schema = schema();
        let names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
        assert_eq!(
            names,
            vec![
                "partition",
                "position",
                "turn_id",
                "input_tokens",
                "output_tokens",
            ]
        );

        let expect = [
            ("partition", DataType::Utf8, false),
            ("position", DataType::UInt64, false),
            ("turn_id", DataType::Utf8, true),
            ("input_tokens", DataType::UInt64, false),
            ("output_tokens", DataType::UInt64, false),
        ];
        for (field, (name, ty, nullable)) in schema.fields().iter().zip(expect) {
            assert_eq!(field.name(), name);
            assert_eq!(field.data_type(), &ty);
            assert_eq!(field.is_nullable(), nullable);
        }
    }

    #[test]
    fn decode_usage_batch_round_trips_and_sums() {
        let rows = vec![
            UsageRow {
                partition: "conv-a".to_string(),
                position: 5,
                turn_id: None,
                usage: polyc_facts::UsageFact {
                    input_tokens: 42,
                    output_tokens: 7,
                },
            },
            UsageRow {
                partition: "conv-a".to_string(),
                position: 9,
                turn_id: Some("turn-xyz".to_string()),
                usage: polyc_facts::UsageFact {
                    input_tokens: 100,
                    output_tokens: 58,
                },
            },
        ];
        let batch = decode_usage_batch(&rows).expect("batch build");
        assert_eq!(batch.num_rows(), 2);
        assert_eq!(batch.schema(), schema());

        let partition = batch
            .column(0)
            .as_any()
            .downcast_ref::<arrow::array::StringArray>()
            .unwrap();
        assert_eq!(partition.value(0), "conv-a");
        assert_eq!(partition.value(1), "conv-a");

        let position = batch
            .column(1)
            .as_any()
            .downcast_ref::<arrow::array::UInt64Array>()
            .unwrap();
        assert_eq!(position.values(), &[5, 9]);

        let turn_id = batch
            .column(2)
            .as_any()
            .downcast_ref::<arrow::array::StringArray>()
            .unwrap();
        assert!(turn_id.is_null(0));
        assert_eq!(turn_id.value(1), "turn-xyz");

        let input = batch
            .column(3)
            .as_any()
            .downcast_ref::<arrow::array::UInt64Array>()
            .unwrap();
        let output = batch
            .column(4)
            .as_any()
            .downcast_ref::<arrow::array::UInt64Array>()
            .unwrap();
        assert_eq!(input.values(), &[42, 100]);
        assert_eq!(output.values(), &[7, 58]);
        assert_eq!(input.values().iter().sum::<u64>(), 142);
        assert_eq!(output.values().iter().sum::<u64>(), 65);
    }

    #[test]
    fn decode_usage_events_filters_and_decodes_real_buffa_bytes() {
        let turn = Uuid::from_u128(0x0195_abcd_ef01_2345_6789_abcd_ef01_2345);
        let usage = UsageEvent {
            input_tokens: 9,
            output_tokens: 3,
            ..Default::default()
        };
        let bytes = usage.encode_to_vec();

        let events = vec![
            (1, Event::new(kinds::TURN_START, Vec::new())),
            (2, Event::new(kinds::tagged(kinds::USAGE, &turn), bytes)),
            (3, Event::new(kinds::USER_MSG, b"not usage".to_vec())),
        ];

        let decoded = decode_usage_events("conv-real", &events);
        assert_eq!(decoded.len(), 1);
        assert_eq!(decoded[0].partition, "conv-real");
        assert_eq!(decoded[0].position, 2);
        assert_eq!(decoded[0].turn_id, Some(turn.to_string()));
        assert_eq!(decoded[0].usage.input_tokens, 9);
        assert_eq!(decoded[0].usage.output_tokens, 3);
    }

    #[test]
    fn decode_usage_events_bare_kind_has_no_turn_id() {
        let events = vec![(7, Event::new(kinds::USAGE, Vec::new()))];
        let decoded = decode_usage_events("conv-bare", &events);
        assert_eq!(decoded.len(), 1);
        assert_eq!(decoded[0].partition, "conv-bare");
        assert_eq!(decoded[0].position, 7);
        assert_eq!(decoded[0].turn_id, None);
    }

    #[test]
    fn empty_payload_usage_event_decodes_to_zero_zero() {
        let events = vec![(1, Event::new(kinds::USAGE, Vec::new()))];
        let decoded = decode_usage_events("conv-empty", &events);
        assert_eq!(decoded.len(), 1, "an empty payload must decode, not skip");
        assert_eq!(decoded[0].usage.input_tokens, 0);
        assert_eq!(decoded[0].usage.output_tokens, 0);
    }

    #[test]
    fn undecodable_non_empty_payload_is_skipped() {
        let events = vec![
            (1, Event::new(kinds::USAGE, vec![0xFF, 0xFE, 0xFD])),
            (2, Event::new(kinds::USAGE, Vec::new())),
        ];
        let decoded = decode_usage_events("conv-corrupt", &events);
        assert_eq!(decoded.len(), 1, "only the empty (valid) payload decodes");
        assert_eq!(decoded[0].position, 2);
        assert_eq!(decoded[0].usage.input_tokens, 0);
        assert_eq!(decoded[0].usage.output_tokens, 0);
    }

    #[test]
    fn usage_event_round_trip_through_decode_and_batch() {
        let turn = Uuid::from_u128(0x0195_abcd_ef01_2345_6789_abcd_ef01_9999);
        let usage = UsageEvent {
            input_tokens: 1234,
            output_tokens: 5678,
            ..Default::default()
        };
        let bytes = usage.encode_to_vec();
        let events = vec![(1, Event::new(kinds::tagged(kinds::USAGE, &turn), bytes))];

        let decoded = decode_usage_events("conv-rt", &events);
        let batch = decode_usage_batch(&decoded).expect("batch build");

        let partition = batch
            .column(0)
            .as_any()
            .downcast_ref::<arrow::array::StringArray>()
            .unwrap();
        assert_eq!(partition.value(0), "conv-rt");

        let position = batch
            .column(1)
            .as_any()
            .downcast_ref::<arrow::array::UInt64Array>()
            .unwrap();
        assert_eq!(position.value(0), 1);

        let turn_id = batch
            .column(2)
            .as_any()
            .downcast_ref::<arrow::array::StringArray>()
            .unwrap();
        assert_eq!(turn_id.value(0), turn.to_string());

        let input = batch
            .column(3)
            .as_any()
            .downcast_ref::<arrow::array::UInt64Array>()
            .unwrap();
        let output = batch
            .column(4)
            .as_any()
            .downcast_ref::<arrow::array::UInt64Array>()
            .unwrap();
        assert_eq!(input.value(0), 1234);
        assert_eq!(output.value(0), 5678);
    }
}