opendeviationbar-streaming 13.80.0

//! Arrow export for CompletedBar vectors (v7.2 #313)
//!
//! Extends the core `opendeviationbar_vec_to_record_batch()` schema with
//! `_symbol` (Utf8) and `_threshold_decimal_bps` (UInt32) metadata columns.
//! This enables zero-copy Polars grouping in Python's sync_flush path,
//! eliminating ~200K Python dict allocations per drain.

use arrow_array::RecordBatch;
use arrow_array::builder::{StringBuilder, UInt32Builder};
use arrow_schema::{DataType, Field, Schema};
use std::sync::Arc;

use opendeviationbar_core::arrow_export::opendeviationbar_vec_to_record_batch;

use crate::live_engine::CompletedBar;

/// Convert a slice of `CompletedBar`s to an Arrow RecordBatch.
///
/// The output RecordBatch has all columns from `opendeviationbar_vec_to_record_batch()`
/// plus two metadata columns appended:
/// - `_symbol` (Utf8): Trading pair (e.g., "BTCUSDT")
/// - `_threshold_decimal_bps` (UInt32): Threshold in decimal basis points
///
/// This is a two-step process:
/// 1. Extract `OpenDeviationBar` refs and build the core RecordBatch
/// 2. Append `_symbol` and `_threshold_decimal_bps` columns
///
/// The resulting RecordBatch can be zero-copy converted to a Polars DataFrame
/// in Python for efficient grouping by (symbol, threshold).
pub fn completed_bars_to_record_batch(bars: &[CompletedBar]) -> RecordBatch {
    if bars.is_empty() {
        // Return empty RecordBatch with extended schema
        let core_batch = opendeviationbar_vec_to_record_batch(&[]);
        return append_metadata_columns(core_batch, &[], &[]);
    }

    // Step 1: Extract OpenDeviationBar refs for core conversion
    let odb_bars: Vec<&opendeviationbar_core::types::OpenDeviationBar> =
        bars.iter().map(|cb| &cb.bar).collect();

    // Build core RecordBatch from borrowed bars (avoids clone)
    let core_batch = opendeviationbar_vec_to_record_batch_borrowed(&odb_bars);

    // Step 2: Collect metadata
    let symbols: Vec<&str> = bars.iter().map(|cb| &*cb.symbol).collect();
    let thresholds: Vec<u32> = bars.iter().map(|cb| cb.threshold_decimal_bps).collect();

    append_metadata_columns(core_batch, &symbols, &thresholds)
}

/// Append `_symbol` and `_threshold_decimal_bps` columns to an existing RecordBatch.
fn append_metadata_columns(
    core_batch: RecordBatch,
    symbols: &[&str],
    thresholds: &[u32],
) -> RecordBatch {
    let n = core_batch.num_rows();

    // Build metadata columns
    let mut symbol_builder = StringBuilder::with_capacity(n, n * 10);
    let mut threshold_builder = UInt32Builder::with_capacity(n);

    for &sym in symbols {
        symbol_builder.append_value(sym);
    }
    for &thr in thresholds {
        threshold_builder.append_value(thr);
    }

    // Extend schema with metadata fields
    let mut fields: Vec<Arc<Field>> = core_batch.schema().fields().iter().cloned().collect();
    fields.push(Arc::new(Field::new("_symbol", DataType::Utf8, false)));
    fields.push(Arc::new(Field::new(
        "_threshold_decimal_bps",
        DataType::UInt32,
        false,
    )));
    let extended_schema = Arc::new(Schema::new(fields));

    // Extend columns
    let mut columns: Vec<Arc<dyn arrow_array::Array>> = core_batch.columns().to_vec();
    columns.push(Arc::new(symbol_builder.finish()));
    columns.push(Arc::new(threshold_builder.finish()));

    RecordBatch::try_new(extended_schema, columns)
        .expect("Failed to create extended RecordBatch with metadata columns")
}

/// Like `opendeviationbar_vec_to_record_batch` but takes borrowed refs to avoid cloning.
///
/// This is a thin wrapper that dereferences the borrowed slice for the core function.
fn opendeviationbar_vec_to_record_batch_borrowed(
    bars: &[&opendeviationbar_core::types::OpenDeviationBar],
) -> RecordBatch {
    // The core function takes &[OpenDeviationBar]. We need to convert
    // &[&OpenDeviationBar] to a contiguous slice. Since OpenDeviationBar
    // is not trivially small, we pass through the existing function by
    // collecting into a temporary Vec. The Arrow builder allocation
    // dominates cost for large batches, so this clone is negligible.
    //
    // Future optimization: refactor core to accept Iterator<Item = &OpenDeviationBar>.
    let owned: Vec<opendeviationbar_core::types::OpenDeviationBar> =
        bars.iter().map(|b| (*b).clone()).collect();
    opendeviationbar_vec_to_record_batch(&owned)
}

#[cfg(test)]
mod tests {
    use super::*;
    use opendeviationbar_core::types::OpenDeviationBar;

    fn make_test_bar() -> OpenDeviationBar {
        OpenDeviationBar::default()
    }

    #[test]
    fn test_completed_bars_to_record_batch_empty() {
        let batch = completed_bars_to_record_batch(&[]);
        assert_eq!(batch.num_rows(), 0);
        // Should have core columns + _symbol + _threshold_decimal_bps
        assert!(batch.schema().field_with_name("_symbol").is_ok());
        assert!(
            batch
                .schema()
                .field_with_name("_threshold_decimal_bps")
                .is_ok()
        );
    }

    #[test]
    fn test_completed_bars_to_record_batch_single() {
        let bars = vec![CompletedBar {
            symbol: Arc::from("BTCUSDT"),
            threshold_decimal_bps: 250,
            bar: make_test_bar(),
        }];
        let batch = completed_bars_to_record_batch(&bars);
        assert_eq!(batch.num_rows(), 1);

        // Verify metadata columns
        let sym_col = batch
            .column_by_name("_symbol")
            .unwrap()
            .as_any()
            .downcast_ref::<arrow_array::StringArray>()
            .unwrap();
        assert_eq!(sym_col.value(0), "BTCUSDT");

        let thr_col = batch
            .column_by_name("_threshold_decimal_bps")
            .unwrap()
            .as_any()
            .downcast_ref::<arrow_array::UInt32Array>()
            .unwrap();
        assert_eq!(thr_col.value(0), 250);
    }

    #[test]
    fn test_completed_bars_to_record_batch_mixed() {
        let bars = vec![
            CompletedBar {
                symbol: Arc::from("BTCUSDT"),
                threshold_decimal_bps: 250,
                bar: make_test_bar(),
            },
            CompletedBar {
                symbol: Arc::from("ETHUSDT"),
                threshold_decimal_bps: 500,
                bar: make_test_bar(),
            },
            CompletedBar {
                symbol: Arc::from("BTCUSDT"),
                threshold_decimal_bps: 500,
                bar: make_test_bar(),
            },
        ];
        let batch = completed_bars_to_record_batch(&bars);
        assert_eq!(batch.num_rows(), 3);

        let sym_col = batch
            .column_by_name("_symbol")
            .unwrap()
            .as_any()
            .downcast_ref::<arrow_array::StringArray>()
            .unwrap();
        assert_eq!(sym_col.value(0), "BTCUSDT");
        assert_eq!(sym_col.value(1), "ETHUSDT");
        assert_eq!(sym_col.value(2), "BTCUSDT");
    }
}