arrow_graph_core/

store.rs

//! ArrowGraphStore — the core partitioned graph store.
//!
//! Holds triples partitioned by namespace (string keys), with an optional layer
//! column for sub-partitioning. Supports add, query, delete (logical), and batch
//! operations.
//!
//! # Namespace Partitioning
//!
//! Namespaces are arbitrary strings — you define them when constructing the store.
//! This keeps the core generic: a knowledge graph might use `["world", "code", "self"]`,
//! while a document store might use `["drafts", "published", "archived"]`.
//!
//! ```rust
//! use arrow_graph_core::store::ArrowGraphStore;
//!
//! let mut store = ArrowGraphStore::new(&["world", "code", "self"]);
//! ```

use crate::schema::{col, triples_schema};

use arrow::array::{
    Array, BooleanArray, Float64Array, RecordBatch, StringArray, TimestampMillisecondArray,
    UInt8Array,
};
use arrow::compute;
use arrow::datatypes::SchemaRef;
use std::collections::HashMap;
use std::sync::Arc;

/// Errors from store operations.
#[derive(Debug, thiserror::Error)]
pub enum StoreError {
    #[error("Arrow error: {0}")]
    Arrow(#[from] arrow::error::ArrowError),

    #[error("Unknown namespace: {0}")]
    UnknownNamespace(String),

    #[error("Invalid layer: {0}")]
    InvalidLayer(u8),

    #[error("Triple not found: {0}")]
    TripleNotFound(String),
}

pub type Result<T> = std::result::Result<T, StoreError>;

/// A single triple to be added to the store.
#[derive(Debug, Clone)]
pub struct Triple {
    pub subject: String,
    pub predicate: String,
    pub object: String,
    pub graph: Option<String>,
    pub confidence: Option<f64>,
    pub source_document: Option<String>,
    /// FK to ChunkTable for fine-grained provenance.
    pub source_chunk_id: Option<String>,
    pub extracted_by: Option<String>,
    /// The triple_id of the triple that caused this one (causal chain).
    pub caused_by: Option<String>,
    /// The triple_id of the triple this was derived from.
    pub derived_from: Option<String>,
    /// Timestamp (ms since epoch) when this triple was consolidated.
    pub consolidated_at: Option<i64>,
}

/// Filter specification for queries.
#[derive(Debug, Default, Clone)]
pub struct QuerySpec {
    pub subject: Option<String>,
    pub predicate: Option<String>,
    pub object: Option<String>,
    pub namespace: Option<String>,
    pub layer: Option<u8>,
    pub include_deleted: bool,
}

/// A node in a causal derivation chain.
#[derive(Debug, Clone, PartialEq)]
pub struct CausalNode {
    pub triple_id: String,
    pub caused_by: Option<String>,
    pub derived_from: Option<String>,
}

/// The core Arrow-native graph store, partitioned by namespace.
///
/// Each namespace holds a vector of RecordBatches (appended over time).
/// Queries filter by namespace first, then by column predicates.
///
/// Namespaces are arbitrary strings defined at construction time.
pub struct ArrowGraphStore {
    schema: SchemaRef,
    /// Per-namespace storage: `Vec<RecordBatch>` (append-only within a partition).
    partitions: HashMap<String, Vec<RecordBatch>>,
    /// Ordered list of namespace names (for deterministic iteration).
    namespace_names: Vec<String>,
}

impl ArrowGraphStore {
    /// Create a new empty store with the given namespace partitions.
    ///
    /// # Example
    ///
    /// ```rust
    /// use arrow_graph_core::store::ArrowGraphStore;
    ///
    /// // Knowledge graph with three partitions
    /// let store = ArrowGraphStore::new(&["world", "code", "self"]);
    ///
    /// // Document store with different partitions
    /// let store = ArrowGraphStore::new(&["drafts", "published"]);
    /// ```
    pub fn new(namespaces: &[&str]) -> Self {
        let schema = Arc::new(triples_schema());
        let mut partitions = HashMap::new();
        let mut namespace_names = Vec::new();
        for ns in namespaces {
            partitions.insert(ns.to_string(), Vec::new());
            namespace_names.push(ns.to_string());
        }
        ArrowGraphStore {
            schema,
            partitions,
            namespace_names,
        }
    }

    /// Get the triples schema.
    pub fn schema(&self) -> &SchemaRef {
        &self.schema
    }

    /// Get the list of namespace names.
    pub fn namespaces(&self) -> &[String] {
        &self.namespace_names
    }

    /// Add a single triple to the specified namespace and layer.
    ///
    /// The `layer` parameter is optional — pass `None` to use layer 0 (default).
    pub fn add_triple(
        &mut self,
        triple: &Triple,
        namespace: &str,
        layer: Option<u8>,
    ) -> Result<String> {
        self.add_batch(std::slice::from_ref(triple), namespace, layer)
            .map(|ids| ids.into_iter().next().unwrap())
    }

    /// Add a batch of triples to the specified namespace and layer.
    /// Returns the generated triple IDs.
    pub fn add_batch(
        &mut self,
        triples: &[Triple],
        namespace: &str,
        layer: Option<u8>,
    ) -> Result<Vec<String>> {
        if !self.partitions.contains_key(namespace) {
            return Err(StoreError::UnknownNamespace(namespace.to_string()));
        }

        let n = triples.len();
        if n == 0 {
            return Ok(vec![]);
        }

        let now_ms = chrono::Utc::now().timestamp_millis();
        let layer_val = layer.unwrap_or(0);

        let ids: Vec<String> = (0..n).map(|_| uuid::Uuid::new_v4().to_string()).collect();

        let subjects: Vec<&str> = triples.iter().map(|t| t.subject.as_str()).collect();
        let predicates: Vec<&str> = triples.iter().map(|t| t.predicate.as_str()).collect();
        let objects: Vec<&str> = triples.iter().map(|t| t.object.as_str()).collect();
        let graphs: Vec<Option<&str>> = triples.iter().map(|t| t.graph.as_deref()).collect();
        let ns_vals: Vec<&str> = vec![namespace; n];
        let layer_vals: Vec<u8> = vec![layer_val; n];
        let confidences: Vec<Option<f64>> = triples.iter().map(|t| t.confidence).collect();
        let source_docs: Vec<Option<&str>> = triples
            .iter()
            .map(|t| t.source_document.as_deref())
            .collect();
        let source_chunks: Vec<Option<&str>> = triples
            .iter()
            .map(|t| t.source_chunk_id.as_deref())
            .collect();
        let extracted: Vec<Option<&str>> =
            triples.iter().map(|t| t.extracted_by.as_deref()).collect();
        let caused_by: Vec<Option<&str>> = triples.iter().map(|t| t.caused_by.as_deref()).collect();
        let derived_from: Vec<Option<&str>> =
            triples.iter().map(|t| t.derived_from.as_deref()).collect();
        let consolidated_at: Vec<Option<i64>> = triples.iter().map(|t| t.consolidated_at).collect();
        let timestamps: Vec<i64> = vec![now_ms; n];
        let deleted: Vec<bool> = vec![false; n];
        let id_strs: Vec<&str> = ids.iter().map(|s| s.as_str()).collect();

        let batch = RecordBatch::try_new(
            self.schema.clone(),
            vec![
                Arc::new(StringArray::from(id_strs)),
                Arc::new(StringArray::from(subjects)),
                Arc::new(StringArray::from(predicates)),
                Arc::new(StringArray::from(objects)),
                Arc::new(StringArray::from(graphs)),
                Arc::new(StringArray::from(ns_vals)),
                Arc::new(UInt8Array::from(layer_vals)),
                Arc::new(Float64Array::from(confidences)),
                Arc::new(StringArray::from(source_docs)),
                Arc::new(StringArray::from(source_chunks)),
                Arc::new(StringArray::from(extracted)),
                Arc::new(TimestampMillisecondArray::from(timestamps).with_timezone("UTC")),
                Arc::new(StringArray::from(caused_by)),
                Arc::new(StringArray::from(derived_from)),
                Arc::new(TimestampMillisecondArray::from(consolidated_at).with_timezone("UTC")),
                Arc::new(BooleanArray::from(deleted)),
            ],
        )?;

        self.partitions.get_mut(namespace).unwrap().push(batch);

        Ok(ids)
    }

    /// Query triples matching the given spec.
    pub fn query(&self, spec: &QuerySpec) -> Result<Vec<RecordBatch>> {
        let namespaces: Vec<&String> = match &spec.namespace {
            Some(ns) => {
                if !self.partitions.contains_key(ns) {
                    return Err(StoreError::UnknownNamespace(ns.clone()));
                }
                vec![ns]
            }
            None => self.namespace_names.iter().collect(),
        };

        let mut results = Vec::new();

        for ns in namespaces {
            let batches = self.partitions.get(ns).unwrap();
            for batch in batches {
                let filtered = self.filter_batch(batch, spec)?;
                if filtered.num_rows() > 0 {
                    results.push(filtered);
                }
            }
        }

        Ok(results)
    }

    /// Total number of non-deleted triples across all namespaces.
    pub fn len(&self) -> usize {
        let spec = QuerySpec::default();
        self.query(&spec)
            .unwrap_or_default()
            .iter()
            .map(|b| b.num_rows())
            .sum()
    }

    /// Whether the store has no non-deleted triples.
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Total number of triples including deleted.
    pub fn len_all(&self) -> usize {
        self.partitions
            .values()
            .flat_map(|batches| batches.iter())
            .map(|b| b.num_rows())
            .sum()
    }

    /// Logically delete a triple by ID (sets deleted=true).
    pub fn delete(&mut self, triple_id: &str) -> Result<bool> {
        for batches in self.partitions.values_mut() {
            for batch in batches.iter_mut() {
                let id_col = batch
                    .column(col::TRIPLE_ID)
                    .as_any()
                    .downcast_ref::<StringArray>()
                    .expect("triple_id column must be StringArray");

                let mut found_idx = None;
                for i in 0..id_col.len() {
                    if id_col.value(i) == triple_id {
                        found_idx = Some(i);
                        break;
                    }
                }

                if let Some(idx) = found_idx {
                    let del_col = batch
                        .column(col::DELETED)
                        .as_any()
                        .downcast_ref::<BooleanArray>()
                        .expect("deleted column must be BooleanArray");

                    let mut new_del: Vec<bool> =
                        (0..del_col.len()).map(|i| del_col.value(i)).collect();
                    new_del[idx] = true;

                    let mut columns: Vec<Arc<dyn Array>> = Vec::new();
                    for c in 0..batch.num_columns() {
                        if c == col::DELETED {
                            columns.push(Arc::new(BooleanArray::from(new_del.clone())));
                        } else {
                            columns.push(batch.column(c).clone());
                        }
                    }

                    *batch = RecordBatch::try_new(self.schema.clone(), columns)?;
                    return Ok(true);
                }
            }
        }
        Ok(false)
    }

    /// Get all RecordBatches for a given namespace (including deleted triples).
    pub fn get_namespace_batches(&self, namespace: &str) -> &[RecordBatch] {
        self.partitions.get(namespace).map_or(&[], |v| v.as_slice())
    }

    /// Replace all data for a namespace (used by checkout/restore).
    pub fn set_namespace_batches(&mut self, namespace: &str, batches: Vec<RecordBatch>) {
        if let Some(existing) = self.partitions.get_mut(namespace) {
            *existing = batches;
        }
    }

    /// Follow the caused_by/derived_from chain from a triple to build a derivation graph.
    ///
    /// Returns a list of causal nodes representing the full ancestry of the given
    /// triple. The first element is always the queried triple itself.
    /// Traversal is breadth-first.
    pub fn causal_chain(&self, triple_id: &str) -> Vec<CausalNode> {
        let mut result = Vec::new();
        let mut visited = std::collections::HashSet::new();
        let mut queue = std::collections::VecDeque::new();
        queue.push_back(triple_id.to_string());

        // Build an index for efficient lookup
        let mut index: HashMap<String, (Option<String>, Option<String>)> = HashMap::new();
        for batches in self.partitions.values() {
            for batch in batches {
                let id_col = batch
                    .column(col::TRIPLE_ID)
                    .as_any()
                    .downcast_ref::<StringArray>()
                    .expect("triple_id column");
                let caused_col = batch
                    .column(col::CAUSED_BY)
                    .as_any()
                    .downcast_ref::<StringArray>()
                    .expect("caused_by column");
                let derived_col = batch
                    .column(col::DERIVED_FROM)
                    .as_any()
                    .downcast_ref::<StringArray>()
                    .expect("derived_from column");
                let del_col = batch
                    .column(col::DELETED)
                    .as_any()
                    .downcast_ref::<BooleanArray>()
                    .expect("deleted column");

                for i in 0..batch.num_rows() {
                    if del_col.value(i) {
                        continue;
                    }
                    let id = id_col.value(i).to_string();
                    let caused = if caused_col.is_null(i) {
                        None
                    } else {
                        Some(caused_col.value(i).to_string())
                    };
                    let derived = if derived_col.is_null(i) {
                        None
                    } else {
                        Some(derived_col.value(i).to_string())
                    };
                    index.insert(id, (caused, derived));
                }
            }
        }

        while let Some(tid) = queue.pop_front() {
            if !visited.insert(tid.clone()) {
                continue;
            }
            if let Some((caused, derived)) = index.get(&tid) {
                result.push(CausalNode {
                    triple_id: tid.clone(),
                    caused_by: caused.clone(),
                    derived_from: derived.clone(),
                });
                if let Some(cb) = caused
                    && !visited.contains(cb)
                {
                    queue.push_back(cb.clone());
                }
                if let Some(df) = derived
                    && !visited.contains(df)
                {
                    queue.push_back(df.clone());
                }
            }
        }

        result
    }

    /// Clear all data from the store.
    pub fn clear(&mut self) {
        for batches in self.partitions.values_mut() {
            batches.clear();
        }
    }

    /// Filter a RecordBatch by the QuerySpec predicates.
    fn filter_batch(&self, batch: &RecordBatch, spec: &QuerySpec) -> Result<RecordBatch> {
        let n = batch.num_rows();
        let mut mask = BooleanArray::from(vec![true; n]);

        // Filter out deleted unless include_deleted
        if !spec.include_deleted {
            let del_col = batch
                .column(col::DELETED)
                .as_any()
                .downcast_ref::<BooleanArray>()
                .expect("deleted column must be BooleanArray");
            let not_deleted = compute::not(del_col)?;
            mask = compute::and(&mask, &not_deleted)?;
        }

        if let Some(ref subj) = spec.subject {
            let c = batch
                .column(col::SUBJECT)
                .as_any()
                .downcast_ref::<StringArray>()
                .expect("subject column must be StringArray");
            let eq = string_eq_scalar(c, subj);
            mask = compute::and(&mask, &eq)?;
        }

        if let Some(ref pred) = spec.predicate {
            let c = batch
                .column(col::PREDICATE)
                .as_any()
                .downcast_ref::<StringArray>()
                .expect("predicate column must be StringArray");
            let eq = string_eq_scalar(c, pred);
            mask = compute::and(&mask, &eq)?;
        }

        if let Some(ref obj) = spec.object {
            let c = batch
                .column(col::OBJECT)
                .as_any()
                .downcast_ref::<StringArray>()
                .expect("object column must be StringArray");
            let eq = string_eq_scalar(c, obj);
            mask = compute::and(&mask, &eq)?;
        }

        if let Some(layer) = spec.layer {
            let c = batch
                .column(col::LAYER)
                .as_any()
                .downcast_ref::<UInt8Array>()
                .expect("layer column must be UInt8Array");
            let eq = u8_eq_scalar(c, layer);
            mask = compute::and(&mask, &eq)?;
        }

        let filtered = compute::filter_record_batch(batch, &mask)?;
        Ok(filtered)
    }
}

/// Scalar string equality: returns BooleanArray where each element == value.
fn string_eq_scalar(array: &StringArray, value: &str) -> BooleanArray {
    let bools: Vec<bool> = (0..array.len()).map(|i| array.value(i) == value).collect();
    BooleanArray::from(bools)
}

/// Scalar u8 equality.
fn u8_eq_scalar(array: &UInt8Array, value: u8) -> BooleanArray {
    let bools: Vec<bool> = (0..array.len()).map(|i| array.value(i) == value).collect();
    BooleanArray::from(bools)
}

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

    fn sample_triple(subj: &str, pred: &str, obj: &str) -> Triple {
        Triple {
            subject: subj.to_string(),
            predicate: pred.to_string(),
            object: obj.to_string(),
            graph: None,
            confidence: Some(0.9),
            source_document: None,
            source_chunk_id: None,
            extracted_by: Some("test".to_string()),
            caused_by: None,
            derived_from: None,
            consolidated_at: None,
        }
    }

    #[test]
    fn test_add_and_query_single() {
        let mut store = ArrowGraphStore::new(&["world", "work"]);
        let id = store
            .add_triple(&sample_triple("s1", "p1", "o1"), "world", Some(1))
            .unwrap();

        assert!(!id.is_empty());
        assert_eq!(store.len(), 1);

        let results = store
            .query(&QuerySpec {
                subject: Some("s1".to_string()),
                ..Default::default()
            })
            .unwrap();
        let total: usize = results.iter().map(|b| b.num_rows()).sum();
        assert_eq!(total, 1);
    }

    #[test]
    fn test_namespace_isolation() {
        let mut store = ArrowGraphStore::new(&["world", "work"]);

        let world_triples: Vec<Triple> = (0..100)
            .map(|i| sample_triple(&format!("w{i}"), "rdf:type", "Thing"))
            .collect();
        store.add_batch(&world_triples, "world", Some(1)).unwrap();

        let work_triples: Vec<Triple> = (0..100)
            .map(|i| sample_triple(&format!("k{i}"), "rdf:type", "Task"))
            .collect();
        store.add_batch(&work_triples, "work", Some(1)).unwrap();

        let world_results = store
            .query(&QuerySpec {
                namespace: Some("world".to_string()),
                ..Default::default()
            })
            .unwrap();
        let world_count: usize = world_results.iter().map(|b| b.num_rows()).sum();
        assert_eq!(world_count, 100);

        let work_results = store
            .query(&QuerySpec {
                namespace: Some("work".to_string()),
                ..Default::default()
            })
            .unwrap();
        let work_count: usize = work_results.iter().map(|b| b.num_rows()).sum();
        assert_eq!(work_count, 100);

        assert_eq!(store.len(), 200);
    }

    #[test]
    fn test_layer_query() {
        let mut store = ArrowGraphStore::new(&["default"]);

        store
            .add_triple(&sample_triple("s1", "p1", "o1"), "default", Some(0))
            .unwrap();
        store
            .add_triple(&sample_triple("s2", "p2", "o2"), "default", Some(1))
            .unwrap();

        let l0_results = store
            .query(&QuerySpec {
                layer: Some(0),
                ..Default::default()
            })
            .unwrap();
        let l0_count: usize = l0_results.iter().map(|b| b.num_rows()).sum();
        assert_eq!(l0_count, 1);
    }

    #[test]
    fn test_default_layer() {
        let mut store = ArrowGraphStore::new(&["default"]);

        // None layer defaults to 0
        store
            .add_triple(&sample_triple("s1", "p1", "o1"), "default", None)
            .unwrap();

        let results = store
            .query(&QuerySpec {
                layer: Some(0),
                ..Default::default()
            })
            .unwrap();
        let count: usize = results.iter().map(|b| b.num_rows()).sum();
        assert_eq!(count, 1);
    }

    #[test]
    fn test_unknown_namespace_error() {
        let mut store = ArrowGraphStore::new(&["world"]);
        let result = store.add_triple(&sample_triple("s", "p", "o"), "unknown", None);
        assert!(matches!(result, Err(StoreError::UnknownNamespace(_))));
    }

    #[test]
    fn test_logical_delete() {
        let mut store = ArrowGraphStore::new(&["world"]);
        let id = store
            .add_triple(&sample_triple("s1", "p1", "o1"), "world", Some(1))
            .unwrap();

        assert_eq!(store.len(), 1);
        assert!(store.delete(&id).unwrap());
        assert_eq!(store.len(), 0);
        assert_eq!(store.len_all(), 1); // Still physically present
    }

    #[test]
    fn test_batch_add_performance() {
        let mut store = ArrowGraphStore::new(&["world"]);

        let triples: Vec<Triple> = (0..10_000)
            .map(|i| sample_triple(&format!("s{i}"), "rdf:type", "Entity"))
            .collect();

        let start = std::time::Instant::now();
        store.add_batch(&triples, "world", Some(1)).unwrap();
        let elapsed = start.elapsed();

        assert_eq!(store.len(), 10_000);
        assert!(
            elapsed.as_millis() < 100,
            "Batch add took too long: {:?}",
            elapsed
        );
    }

    #[test]
    fn test_custom_namespaces() {
        let mut store = ArrowGraphStore::new(&["drafts", "published", "archived"]);
        assert_eq!(store.namespaces().len(), 3);

        store
            .add_triple(&sample_triple("doc1", "status", "draft"), "drafts", None)
            .unwrap();
        store
            .add_triple(&sample_triple("doc2", "status", "live"), "published", None)
            .unwrap();

        assert_eq!(store.len(), 2);

        let drafts = store
            .query(&QuerySpec {
                namespace: Some("drafts".to_string()),
                ..Default::default()
            })
            .unwrap();
        assert_eq!(drafts.iter().map(|b| b.num_rows()).sum::<usize>(), 1);
    }

    #[test]
    fn test_causal_chain_linear() {
        let mut store = ArrowGraphStore::new(&["world"]);

        let t0 = Triple {
            caused_by: None,
            derived_from: None,
            ..sample_triple("s0", "p", "o0")
        };
        let id0 = store.add_triple(&t0, "world", Some(1)).unwrap();

        let t1 = Triple {
            caused_by: Some(id0.clone()),
            derived_from: None,
            ..sample_triple("s1", "p", "o1")
        };
        let id1 = store.add_triple(&t1, "world", Some(1)).unwrap();

        let t2 = Triple {
            caused_by: Some(id1.clone()),
            derived_from: None,
            ..sample_triple("s2", "p", "o2")
        };
        let id2 = store.add_triple(&t2, "world", Some(1)).unwrap();

        let chain = store.causal_chain(&id2);
        assert_eq!(chain.len(), 3);
        assert_eq!(chain[0].triple_id, id2);
        assert_eq!(chain[0].caused_by, Some(id1.clone()));
        assert_eq!(chain[1].triple_id, id1);
        assert_eq!(chain[1].caused_by, Some(id0.clone()));
        assert_eq!(chain[2].triple_id, id0);
        assert_eq!(chain[2].caused_by, None);
    }

    #[test]
    fn test_causal_chain_nonexistent_triple() {
        let store = ArrowGraphStore::new(&["world"]);
        let chain = store.causal_chain("nonexistent");
        assert!(chain.is_empty());
    }
}