nodedb-columnar 0.6.1

// SPDX-License-Identifier: Apache-2.0

//! Row append, truncate, drain, and dictionary-encoding mutations.

use nodedb_types::columnar::ColumnarSchema;
use nodedb_types::value::Value;

use crate::error::ColumnarError;

use super::column_data::{ColumnData, DICT_ENCODE_MAX_CARDINALITY};
use super::core::ColumnarMemtable;
use super::ingest_value::IngestValue;

impl ColumnarMemtable {
    /// Append a row of values. Validates types and nullability.
    ///
    /// The append is all-or-nothing. On an error, every column is cut back
    /// to `row_count`, so the memtable holds exactly the rows it held before
    /// the call. A rolled-back push to a `DictEncoded` column can leave one
    /// dictionary entry that no row references. Every id stays valid.
    pub fn append_row(&mut self, values: &[Value]) -> Result<(), ColumnarError> {
        if values.len() != self.schema.columns.len() {
            return Err(ColumnarError::SchemaMismatch {
                expected: self.schema.columns.len(),
                got: values.len(),
            });
        }

        let pushed = self
            .schema
            .columns
            .iter()
            .zip(values.iter())
            .zip(self.columns.iter_mut())
            .try_for_each(|((col_def, value), column)| {
                if matches!(value, Value::Null) && !col_def.nullable {
                    return Err(ColumnarError::NullViolation(col_def.name.clone()));
                }
                column.push(value, &col_def.name, &col_def.column_type)
            });
        if let Err(e) = pushed {
            self.rollback_partial_row();
            return Err(e);
        }

        self.row_count += 1;
        debug_assert!(
            self.columns.iter().all(|c| c.len() == self.row_count),
            "column lengths must stay aligned with row_count"
        );
        Ok(())
    }

    /// Cut every column back to `row_count` after a failed row push.
    ///
    /// The columns before the failed one hold one extra value. The failed
    /// column and the columns after it hold `row_count` values already.
    fn rollback_partial_row(&mut self) {
        let n = self.row_count;
        for col in &mut self.columns {
            col.truncate(n);
        }
        debug_assert!(
            self.columns.iter().all(|c| c.len() == self.row_count),
            "column lengths must stay aligned with row_count after rollback"
        );
    }

    /// Convert low-cardinality `String` columns to `DictEncoded` in-place.
    pub fn try_dict_encode_columns(&mut self, max_cardinality: u32) {
        for col in &mut self.columns {
            if let ColumnData::String { .. } = col
                && let Some(encoded) = ColumnData::try_dict_encode(col, max_cardinality)
            {
                *col = encoded;
            }
        }
    }

    /// Truncate the memtable to `n` rows, discarding all rows beyond that point.
    ///
    /// Used by transaction rollback to restore the memtable to its pre-write state.
    /// No-op if `n >= self.row_count`.
    pub fn truncate_to(&mut self, n: usize) {
        if n >= self.row_count {
            return;
        }
        for col in &mut self.columns {
            col.truncate(n);
        }
        self.row_count = n;
        debug_assert!(
            self.columns.iter().all(|c| c.len() == self.row_count),
            "column lengths misaligned after truncate_to"
        );
    }

    /// Drain the memtable: return all column data and reset to empty.
    pub fn drain(&mut self) -> (ColumnarSchema, Vec<ColumnData>, usize) {
        let columns = std::mem::replace(
            &mut self.columns,
            self.schema
                .columns
                .iter()
                .map(|col| ColumnData::new(&col.column_type, col.nullable))
                .collect(),
        );
        let row_count = self.row_count;
        self.row_count = 0;
        (self.schema.clone(), columns, row_count)
    }

    /// Drain with automatic dictionary encoding for low-cardinality String columns.
    pub fn drain_optimized(&mut self) -> (ColumnarSchema, Vec<ColumnData>, usize) {
        self.try_dict_encode_columns(DICT_ENCODE_MAX_CARDINALITY);
        self.drain()
    }

    /// Zero-copy row ingest for timeseries and high-throughput paths.
    ///
    /// Accepts borrowed values via `IngestValue<'_>`, avoiding string cloning
    /// for tag columns that are already interned in the `DictEncoded` dictionary.
    /// The ingest is all-or-nothing, as in [`Self::append_row`].
    pub fn ingest_row_refs(&mut self, values: &[IngestValue<'_>]) -> Result<(), ColumnarError> {
        if values.len() != self.schema.columns.len() {
            return Err(ColumnarError::SchemaMismatch {
                expected: self.schema.columns.len(),
                got: values.len(),
            });
        }

        let pushed = self
            .schema
            .columns
            .iter()
            .zip(values.iter())
            .zip(self.columns.iter_mut())
            .try_for_each(|((col_def, value), column)| {
                if matches!(value, IngestValue::Null) && !col_def.nullable {
                    return Err(ColumnarError::NullViolation(col_def.name.clone()));
                }
                column.push_ref(value, &col_def.name)
            });
        if let Err(e) = pushed {
            self.rollback_partial_row();
            return Err(e);
        }

        self.row_count += 1;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use nodedb_types::columnar::{ColumnDef, ColumnType, ColumnarSchema};

    use super::*;

    fn test_schema() -> ColumnarSchema {
        ColumnarSchema::new(vec![
            ColumnDef::required("id", ColumnType::Int64).with_primary_key(),
            ColumnDef::required("name", ColumnType::String),
            ColumnDef::nullable("score", ColumnType::Float64),
        ])
        .expect("valid schema")
    }

    #[test]
    fn append_and_count() {
        let schema = test_schema();
        let mut mt = ColumnarMemtable::new(&schema);

        mt.append_row(&[
            Value::Integer(1),
            Value::String("Alice".into()),
            Value::Float(0.75),
        ])
        .expect("append");

        mt.append_row(&[Value::Integer(2), Value::String("Bob".into()), Value::Null])
            .expect("append");

        assert_eq!(mt.row_count(), 2);
        assert!(!mt.is_empty());
    }

    #[test]
    fn null_violation_rejected() {
        let schema = test_schema();
        let mut mt = ColumnarMemtable::new(&schema);

        let err = mt
            .append_row(&[Value::Null, Value::String("x".into()), Value::Null])
            .unwrap_err();
        assert!(matches!(err, ColumnarError::NullViolation(ref s) if s == "id"));
    }

    #[test]
    fn null_violation_mid_row_leaves_columns_aligned() {
        let schema = test_schema();
        let mut mt = ColumnarMemtable::new(&schema);

        let err = mt
            .append_row(&[Value::Integer(1), Value::Null, Value::Null])
            .unwrap_err();
        assert!(matches!(err, ColumnarError::NullViolation(ref s) if s == "name"));
        assert_eq!(mt.row_count(), 0);
        assert!(mt.columns().iter().all(|c| c.len() == 0));

        mt.append_row(&[Value::Integer(2), Value::String("b".into()), Value::Null])
            .expect("append after rejected row");
        assert_eq!(mt.row_count(), 1);
        assert!(mt.columns().iter().all(|c| c.len() == 1));
        assert_eq!(
            mt.get_row(0).expect("read"),
            Some(vec![
                Value::Integer(2),
                Value::String("b".into()),
                Value::Null
            ])
        );
    }

    #[test]
    fn type_mismatch_mid_row_leaves_columns_aligned() {
        let schema = test_schema();
        let mut mt = ColumnarMemtable::new(&schema);
        mt.append_row(&[
            Value::Integer(1),
            Value::String("a".into()),
            Value::Float(0.5),
        ])
        .expect("append");

        let err = mt
            .append_row(&[
                Value::Integer(2),
                Value::String("b".into()),
                Value::Bool(true),
            ])
            .unwrap_err();
        assert!(matches!(err, ColumnarError::TypeMismatch { ref column, .. } if column == "score"));
        assert_eq!(mt.row_count(), 1);
        assert!(mt.columns().iter().all(|c| c.len() == 1));

        mt.append_row(&[
            Value::Integer(3),
            Value::String("c".into()),
            Value::Float(0.25),
        ])
        .expect("append after rejected row");
        assert_eq!(
            mt.get_row(1).expect("read"),
            Some(vec![
                Value::Integer(3),
                Value::String("c".into()),
                Value::Float(0.25)
            ])
        );
    }

    #[test]
    fn rejected_ref_ingest_leaves_columns_aligned() {
        let schema = test_schema();
        let mut mt = ColumnarMemtable::new(&schema);

        let err = mt
            .ingest_row_refs(&[
                IngestValue::Int64(1),
                IngestValue::Str("a"),
                IngestValue::Bool(true),
            ])
            .unwrap_err();
        assert!(matches!(err, ColumnarError::TypeMismatch { .. }));
        assert_eq!(mt.row_count(), 0);
        assert!(mt.columns().iter().all(|c| c.len() == 0));
    }

    #[test]
    fn schema_mismatch_rejected() {
        let schema = test_schema();
        let mut mt = ColumnarMemtable::new(&schema);

        let err = mt.append_row(&[Value::Integer(1)]).unwrap_err();
        assert!(matches!(err, ColumnarError::SchemaMismatch { .. }));
    }

    #[test]
    fn flush_threshold() {
        let schema = test_schema();
        let mut mt = ColumnarMemtable::with_threshold(&schema, 3);

        for i in 0..2 {
            mt.append_row(&[
                Value::Integer(i),
                Value::String(format!("u{i}")),
                Value::Null,
            ])
            .expect("append");
        }
        assert!(!mt.should_flush());

        mt.append_row(&[Value::Integer(2), Value::String("u2".into()), Value::Null])
            .expect("append");
        assert!(mt.should_flush());
    }

    #[test]
    fn drain_resets() {
        let schema = test_schema();
        let mut mt = ColumnarMemtable::new(&schema);

        mt.append_row(&[
            Value::Integer(1),
            Value::String("x".into()),
            Value::Float(0.5),
        ])
        .expect("append");

        let (_schema, columns, row_count) = mt.drain();
        assert_eq!(row_count, 1);
        assert_eq!(columns.len(), 3);
        assert_eq!(mt.row_count(), 0);
        assert!(mt.is_empty());

        match &columns[0] {
            ColumnData::Int64 { values, valid } => {
                assert_eq!(values, &[1]);
                assert!(valid.is_none());
            }
            _ => panic!("expected Int64"),
        }
        match &columns[1] {
            ColumnData::String {
                data,
                offsets,
                valid,
            } => {
                assert_eq!(std::str::from_utf8(data).unwrap(), "x");
                assert_eq!(offsets, &[0, 1]);
                assert!(valid.is_none());
            }
            _ => panic!("expected String"),
        }
    }

    #[test]
    fn all_types() {
        let schema = ColumnarSchema::new(vec![
            ColumnDef::required("i", ColumnType::Int64),
            ColumnDef::required("f", ColumnType::Float64),
            ColumnDef::required("b", ColumnType::Bool),
            ColumnDef::required("ts", ColumnType::Timestamp),
            ColumnDef::required("s", ColumnType::String),
            ColumnDef::required("raw", ColumnType::Bytes),
            ColumnDef::required("vec", ColumnType::Vector(3)),
        ])
        .expect("valid");

        let mut mt = ColumnarMemtable::new(&schema);
        mt.append_row(&[
            Value::Integer(42),
            Value::Float(0.25),
            Value::Bool(true),
            Value::Integer(1_700_000_000),
            Value::String("hello".into()),
            Value::Bytes(vec![0xDE, 0xAD]),
            Value::Array(vec![
                Value::Float(1.0),
                Value::Float(2.0),
                Value::Float(3.0),
            ]),
        ])
        .expect("append all types");

        assert_eq!(mt.row_count(), 1);
    }

    #[test]
    fn dict_encode_low_cardinality() {
        let schema = ColumnarSchema::new(vec![ColumnDef::required("qtype", ColumnType::String)])
            .expect("valid");

        let mut mt = ColumnarMemtable::new(&schema);
        let qtypes = ["A", "B", "AAAA", "NS", "MX", "SOA", "CNAME", "PTR"];
        for _ in 0..10 {
            for &q in &qtypes {
                mt.append_row(&[Value::String(q.into())]).expect("append");
            }
        }
        assert_eq!(mt.row_count(), 80);

        mt.try_dict_encode_columns(DICT_ENCODE_MAX_CARDINALITY);

        let (_schema, columns, _row_count) = mt.drain();
        match &columns[0] {
            ColumnData::DictEncoded {
                ids,
                dictionary,
                valid,
                ..
            } => {
                assert_eq!(ids.len(), 80);
                assert!(valid.is_none());
                assert_eq!(dictionary.len(), 8);
                for &id in ids {
                    assert!((id as usize) < dictionary.len());
                }
                for (i, &q) in qtypes.iter().enumerate().take(8) {
                    let expected_id = dictionary.iter().position(|s| s == q).expect("in dict");
                    assert_eq!(ids[i], expected_id as u32);
                }
            }
            _ => panic!("expected DictEncoded after try_dict_encode_columns"),
        }
    }

    #[test]
    fn dict_encode_exceeds_cardinality_stays_string() {
        let schema = ColumnarSchema::new(vec![ColumnDef::required("name", ColumnType::String)])
            .expect("valid");

        let mut mt = ColumnarMemtable::new(&schema);
        let max: u32 = 4;
        for i in 0..=max {
            mt.append_row(&[Value::String(format!("val_{i}"))])
                .expect("append");
        }

        mt.try_dict_encode_columns(max);

        let (_schema, columns, _row_count) = mt.drain();
        assert!(matches!(columns[0], ColumnData::String { .. }));
    }

    #[test]
    fn dict_encode_with_nulls() {
        let schema = ColumnarSchema::new(vec![ColumnDef::nullable("tag", ColumnType::String)])
            .expect("valid");

        let mut mt = ColumnarMemtable::new(&schema);
        mt.append_row(&[Value::String("foo".into())])
            .expect("append");
        mt.append_row(&[Value::Null]).expect("append null");
        mt.append_row(&[Value::String("bar".into())])
            .expect("append");
        mt.append_row(&[Value::Null]).expect("append null");

        mt.try_dict_encode_columns(DICT_ENCODE_MAX_CARDINALITY);

        let (_schema, columns, _row_count) = mt.drain();
        match &columns[0] {
            ColumnData::DictEncoded {
                ids,
                valid,
                dictionary,
                ..
            } => {
                assert_eq!(ids.len(), 4);
                let v = valid.as_ref().expect("nullable column has validity bitmap");
                assert_eq!(v.len(), 4);
                assert!(v[0]);
                assert!(!v[1]);
                assert!(v[2]);
                assert!(!v[3]);
                assert_eq!(dictionary.len(), 2);
            }
            _ => panic!("expected DictEncoded"),
        }
    }
}