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::events_decode::try_decode_event_payload;
use polyc_proto::kinds;
use polyc_proto::proto::polychrome::events::v1::SummaryEvent;
#[derive(Debug, Clone)]
pub(crate) struct SummaryRow {
pub partition: String,
pub position: u64,
pub turn_id: Option<String>,
pub summary: SummaryEvent,
}
#[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("text", DataType::Utf8, false),
Field::new("covers_through_position", DataType::UInt64, false),
]))
}
pub(crate) fn decode_summary_batch(rows: &[SummaryRow]) -> 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 text_b = StringBuilder::with_capacity(rows.len(), rows.len() * 64);
let mut covers_through_b = 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(),
}
text_b.append_value(&row.summary.text);
covers_through_b.append_value(row.summary.covers_through_position);
}
let columns: Vec<ArrayRef> = vec![
Arc::new(partition_b.finish()),
Arc::new(position_b.finish()),
Arc::new(turn_id_b.finish()),
Arc::new(text_b.finish()),
Arc::new(covers_through_b.finish()),
];
RecordBatch::try_new(schema(), columns)
}
#[must_use]
pub(crate) fn decode_summary_events(partition: &str, events: &[(u64, Event)]) -> Vec<SummaryRow> {
crate::decode::decode_typed_kind_events(
partition,
events,
&[kinds::SUMMARY],
"summary",
try_decode_event_payload::<SummaryEvent>,
|partition, position, turn_id, summary: SummaryEvent| SummaryRow {
partition,
position,
turn_id,
summary,
},
)
}
#[cfg(test)]
mod tests {
use arrow::array::Array as _;
use buffa::Message as _;
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",
"text",
"covers_through_position",
]
);
let expect = [
("partition", DataType::Utf8, false),
("position", DataType::UInt64, false),
("turn_id", DataType::Utf8, true),
("text", DataType::Utf8, false),
("covers_through_position", 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_summary_events_filters_and_decodes_real_buffa_bytes() {
let summary_id = Uuid::from_u128(0x0195_abcd_ef01_2345_6789_abcd_ef01_2345);
let summary = SummaryEvent {
text: "the older prefix, condensed".to_string(),
covers_through_position: 12,
..Default::default()
};
let bytes = summary.encode_to_vec();
let events = vec![
(1, Event::new(kinds::TURN_START, Vec::new())),
(
2,
Event::new(kinds::tagged(kinds::SUMMARY, &summary_id), bytes),
),
(3, Event::new(kinds::USER_MSG, b"not summary".to_vec())),
];
let decoded = decode_summary_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(summary_id.to_string()));
assert_eq!(decoded[0].summary.text, "the older prefix, condensed");
assert_eq!(decoded[0].summary.covers_through_position, 12);
}
#[test]
fn decode_summary_events_bare_kind_has_no_turn_id() {
let summary = SummaryEvent {
text: "bare".to_string(),
covers_through_position: 3,
..Default::default()
};
let events = vec![(7, Event::new(kinds::SUMMARY, summary.encode_to_vec()))];
let decoded = decode_summary_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 undecodable_non_empty_payload_is_skipped() {
let good = SummaryEvent {
text: "ok".to_string(),
covers_through_position: 1,
..Default::default()
};
let events = vec![
(1, Event::new(kinds::SUMMARY, vec![0xFF, 0xFE, 0xFD])),
(2, Event::new(kinds::SUMMARY, good.encode_to_vec())),
];
let decoded = decode_summary_events("conv-corrupt", &events);
assert_eq!(decoded.len(), 1, "only the valid payload decodes");
assert_eq!(decoded[0].position, 2);
assert_eq!(decoded[0].summary.text, "ok");
}
#[test]
fn decode_summary_batch_round_trips_multiple_rows() {
let rows = vec![
SummaryRow {
partition: "conv-a".to_string(),
position: 5,
turn_id: None,
summary: SummaryEvent {
text: "first summary".to_string(),
covers_through_position: 4,
..Default::default()
},
},
SummaryRow {
partition: "conv-a".to_string(),
position: 40,
turn_id: Some("summary-xyz".to_string()),
summary: SummaryEvent {
text: "second, later summary".to_string(),
covers_through_position: 30,
..Default::default()
},
},
];
let batch = decode_summary_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, 40]);
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), "summary-xyz");
let text = batch
.column(3)
.as_any()
.downcast_ref::<arrow::array::StringArray>()
.unwrap();
assert_eq!(text.value(0), "first summary");
assert_eq!(text.value(1), "second, later summary");
let covers_through = batch
.column(4)
.as_any()
.downcast_ref::<arrow::array::UInt64Array>()
.unwrap();
assert_eq!(covers_through.values(), &[4, 30]);
}
#[test]
fn empty_payload_summary_event_decodes_to_defaults() {
let events = vec![(1, Event::new(kinds::SUMMARY, Vec::new()))];
let decoded = decode_summary_events("conv-empty", &events);
assert_eq!(decoded.len(), 1, "an empty payload must decode, not skip");
assert_eq!(decoded[0].summary.text, String::new());
assert_eq!(decoded[0].summary.covers_through_position, 0);
}
}