uni_store/storage/

main_edge.rs

// SPDX-License-Identifier: Apache-2.0
// Copyright 2024-2026 Dragonscale Team

//! Main edge table for unified edge storage.
//!
//! This module implements the main `edges` table as described in STORAGE_DESIGN.md.
//! The main table contains all edges in the graph with:
//! - `_eid`: Internal edge ID (primary key)
//! - `src_vid`: Source vertex ID
//! - `dst_vid`: Destination vertex ID
//! - `type`: Edge type name
//! - `props_json`: All properties as JSONB blob
//! - `_deleted`: Soft-delete flag
//! - `_version`: MVCC version
//! - `_created_at`: Creation timestamp
//! - `_updated_at`: Update timestamp

use crate::backend::StorageBackend;
use crate::backend::table_names;
use crate::backend::types::{ScalarIndexType, ScanRequest};
use crate::storage::arrow_convert::build_timestamp_column_from_eid_map;
use anyhow::{Result, anyhow};
use arrow_array::builder::{LargeBinaryBuilder, StringBuilder};
use arrow_array::{Array, ArrayRef, BooleanArray, RecordBatch, UInt64Array};
use arrow_schema::{DataType, Field, Schema as ArrowSchema, TimeUnit};
use sha3::{Digest, Sha3_256};
use std::collections::HashMap;
use std::sync::Arc;
use uni_common::Properties;
use uni_common::core::id::{Eid, UniId, Vid};

/// Which edge endpoint a pushed-down vid set constrains in
/// [`MainEdgeDataset::find_edges_by_type_names`].
///
/// `Src` for outgoing traversals, `Dst` for incoming, `Either` for
/// undirected (`Both`) traversals.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EndpointSide {
    Src,
    Dst,
    Either,
}

/// Main edge dataset for the unified `edges` table.
///
/// This table contains all edges regardless of type, providing:
/// - Fast ID-based lookups without knowing the edge type
/// - Unified traversal queries
#[derive(Debug)]
pub struct MainEdgeDataset {
    _base_uri: String,
}

impl MainEdgeDataset {
    /// Create a new MainEdgeDataset.
    pub fn new(base_uri: &str) -> Self {
        Self {
            _base_uri: base_uri.to_string(),
        }
    }

    /// Compute the content-addressed UID for an edge.
    ///
    /// Edge identity is the SHA3-256 of
    /// `(src_uid, dst_uid, edge_type, sorted_properties)` — the same
    /// content-addressed pattern as
    /// `MainVertexDataset::compute_vertex_uid` but extended with
    /// edge endpoint UIDs and the edge type. This lets fork diff and
    /// promote distinguish parallel edges between the same endpoints
    /// when their property bags differ (multi-edge support).
    ///
    /// Property iteration is sorted by key for deterministic hashing
    /// across machines and runs.
    pub fn compute_edge_uid(
        src_uid: &UniId,
        dst_uid: &UniId,
        edge_type: &str,
        props: &Properties,
    ) -> UniId {
        let mut hasher = Sha3_256::new();

        // Endpoint UIDs first — direction is significant
        // (src→dst ≠ dst→src for the same property bag).
        hasher.update(b"src:");
        hasher.update(src_uid.as_bytes());
        hasher.update(b"\0");
        hasher.update(b"dst:");
        hasher.update(dst_uid.as_bytes());
        hasher.update(b"\0");

        // Edge type.
        hasher.update(b"type:");
        hasher.update(edge_type.as_bytes());
        hasher.update(b"\0");

        // Properties sorted by key (matches compute_vertex_uid).
        let mut sorted_keys: Vec<_> = props.keys().collect();
        sorted_keys.sort();
        for key in sorted_keys {
            if let Some(val) = props.get(key) {
                hasher.update(key.as_bytes());
                hasher.update(b":");
                hasher.update(val.to_string().as_bytes());
                hasher.update(b"\0");
            }
        }

        let result = hasher.finalize();
        UniId::from_bytes(result.into())
    }

    /// Get the Arrow schema for the main edges table.
    pub fn get_arrow_schema() -> Arc<ArrowSchema> {
        Arc::new(ArrowSchema::new(vec![
            Field::new("_eid", DataType::UInt64, false),
            Field::new("src_vid", DataType::UInt64, false),
            Field::new("dst_vid", DataType::UInt64, false),
            Field::new("type", DataType::Utf8, false),
            Field::new("props_json", DataType::LargeBinary, true),
            Field::new("_deleted", DataType::Boolean, false),
            Field::new("_version", DataType::UInt64, false),
            Field::new(
                "_created_at",
                DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())),
                true,
            ),
            Field::new(
                "_updated_at",
                DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())),
                true,
            ),
        ]))
    }

    /// Get the table name for the main edges table.
    pub fn table_name() -> &'static str {
        "edges"
    }

    /// Build a record batch for the main edges table.
    ///
    /// # Arguments
    /// * `edges` - List of (eid, src_vid, dst_vid, edge_type, properties, deleted, version) tuples
    /// * `created_at` - Optional map of Eid -> nanoseconds since epoch
    /// * `updated_at` - Optional map of Eid -> nanoseconds since epoch
    pub fn build_record_batch(
        edges: &[(Eid, Vid, Vid, String, Properties, bool, u64)],
        created_at: Option<&HashMap<Eid, i64>>,
        updated_at: Option<&HashMap<Eid, i64>>,
    ) -> Result<RecordBatch> {
        let arrow_schema = Self::get_arrow_schema();
        let mut columns: Vec<ArrayRef> = Vec::with_capacity(arrow_schema.fields().len());

        // _eid column
        let eids: Vec<u64> = edges
            .iter()
            .map(|(e, _, _, _, _, _, _)| e.as_u64())
            .collect();
        columns.push(Arc::new(UInt64Array::from(eids)));

        // src_vid column
        let src_vids: Vec<u64> = edges
            .iter()
            .map(|(_, s, _, _, _, _, _)| s.as_u64())
            .collect();
        columns.push(Arc::new(UInt64Array::from(src_vids)));

        // dst_vid column
        let dst_vids: Vec<u64> = edges
            .iter()
            .map(|(_, _, d, _, _, _, _)| d.as_u64())
            .collect();
        columns.push(Arc::new(UInt64Array::from(dst_vids)));

        // type column
        let mut type_builder = StringBuilder::new();
        for (_, _, _, edge_type, _, _, _) in edges.iter() {
            type_builder.append_value(edge_type);
        }
        columns.push(Arc::new(type_builder.finish()));

        // props_json column (JSONB binary encoding)
        let mut props_json_builder = LargeBinaryBuilder::new();
        for (_, _, _, _, props, _, _) in edges.iter() {
            let jsonb_bytes = {
                let json_val = serde_json::to_value(props).unwrap_or(serde_json::json!({}));
                let uni_val: uni_common::Value = json_val.into();
                uni_common::cypher_value_codec::encode(&uni_val)
            };
            props_json_builder.append_value(&jsonb_bytes);
        }
        columns.push(Arc::new(props_json_builder.finish()));

        // _deleted column
        let deleted: Vec<bool> = edges.iter().map(|(_, _, _, _, _, d, _)| *d).collect();
        columns.push(Arc::new(BooleanArray::from(deleted)));

        // _version column
        let versions: Vec<u64> = edges.iter().map(|(_, _, _, _, _, _, v)| *v).collect();
        columns.push(Arc::new(UInt64Array::from(versions)));

        // _created_at and _updated_at columns using shared builder
        let eids = edges.iter().map(|(e, _, _, _, _, _, _)| *e);
        columns.push(build_timestamp_column_from_eid_map(
            eids.clone(),
            created_at,
        ));
        columns.push(build_timestamp_column_from_eid_map(eids, updated_at));

        RecordBatch::try_new(arrow_schema, columns).map_err(|e| anyhow!(e))
    }

    /// Write a batch to the main edges table.
    ///
    /// Creates the table if it doesn't exist, otherwise appends to it.
    /// Race-safe under async-flush — see
    /// `crate::storage::manager::write_batch_with_lance_conflict_retry`.
    pub async fn write_batch(backend: &dyn StorageBackend, batch: RecordBatch) -> Result<()> {
        let table_name = table_names::main_edge_table_name();
        crate::storage::manager::write_batch_with_lance_conflict_retry(backend, table_name, batch)
            .await
    }

    /// Ensure default indexes exist on the main edges table.
    ///
    /// Checks for existing indexes before creating to avoid expensive
    /// full-table rebuilds on every flush (LanceDB replaces indexes on create).
    pub async fn ensure_default_indexes(backend: &dyn StorageBackend) -> Result<()> {
        let table_name = table_names::main_edge_table_name();
        let indices = backend.list_indexes(table_name).await?;

        let has_index = |col: &str| {
            indices
                .iter()
                .any(|idx| idx.columns.contains(&col.to_string()))
        };

        for (column, idx_type) in [
            ("_eid", ScalarIndexType::BTree),
            ("src_vid", ScalarIndexType::BTree),
            ("dst_vid", ScalarIndexType::BTree),
            ("type", ScalarIndexType::BTree),
        ] {
            if has_index(column) {
                continue;
            }
            log::info!("Creating {} index on main_edges", column);
            if let Err(e) = backend
                .create_scalar_index(table_name, column, idx_type)
                .await
            {
                log::warn!("Failed to create {} index on main_edges: {}", column, e);
            }
        }

        Ok(())
    }

    /// Query the main edges table for an edge by eid.
    pub async fn find_by_eid(
        backend: &dyn StorageBackend,
        eid: Eid,
    ) -> Result<Option<(Vid, Vid, String, Properties)>> {
        let filter = format!("_eid = {}", eid.as_u64());
        let results = Self::execute_query(backend, &filter, None).await?;

        for batch in results {
            if batch.num_rows() > 0 {
                let src_vid_col = batch.column_by_name("src_vid");
                let dst_vid_col = batch.column_by_name("dst_vid");
                let type_col = batch.column_by_name("type");
                let props_col = batch.column_by_name("props_json");

                if let (Some(src), Some(dst), Some(typ), Some(props)) =
                    (src_vid_col, dst_vid_col, type_col, props_col)
                    && let (Some(src_arr), Some(dst_arr), Some(type_arr), Some(props_arr)) = (
                        src.as_any().downcast_ref::<UInt64Array>(),
                        dst.as_any().downcast_ref::<UInt64Array>(),
                        typ.as_any().downcast_ref::<arrow_array::StringArray>(),
                        props
                            .as_any()
                            .downcast_ref::<arrow_array::LargeBinaryArray>(),
                    )
                {
                    let src_vid = Vid::from(src_arr.value(0));
                    let dst_vid = Vid::from(dst_arr.value(0));
                    let edge_type = type_arr.value(0).to_string();
                    let properties: Properties = if props_arr.is_null(0)
                        || props_arr.value(0).is_empty()
                    {
                        Properties::new()
                    } else {
                        let uni_val = uni_common::cypher_value_codec::decode(props_arr.value(0))
                            .unwrap_or(uni_common::Value::Null);
                        let json_val: serde_json::Value = uni_val.into();
                        serde_json::from_value(json_val).unwrap_or_default()
                    };

                    return Ok(Some((src_vid, dst_vid, edge_type, properties)));
                }
            }
        }

        Ok(None)
    }

    /// Check whether an edge exists by EID, regardless of deletion status.
    ///
    /// Unlike `find_props_by_eid`, this does NOT filter by `_deleted = false`,
    /// so it returns true for both active and soft-deleted edges. Used by the
    /// compaction invariant check to verify dual-writes occurred.
    pub async fn exists_by_eid(backend: &dyn StorageBackend, eid: Eid) -> Result<bool> {
        let filter = format!("_eid = {}", eid.as_u64());
        let batches = Self::execute_query(backend, &filter, Some(vec!["_eid"])).await?;
        Ok(!batches.is_empty() && batches.iter().any(|b| b.num_rows() > 0))
    }

    /// Execute a query on the main edges table.
    ///
    /// Returns empty vec if table doesn't exist.
    async fn execute_query(
        backend: &dyn StorageBackend,
        filter: &str,
        columns: Option<Vec<&str>>,
    ) -> Result<Vec<RecordBatch>> {
        let table_name = table_names::main_edge_table_name();

        if !backend.table_exists(table_name).await? {
            return Ok(Vec::new());
        }

        let mut request = ScanRequest::all(table_name).with_filter(filter);
        if let Some(cols) = columns {
            request = request.with_columns(cols.into_iter().map(String::from).collect());
        }

        backend.scan(request).await
    }

    /// Extract EIDs from record batches.
    fn extract_eids(batches: &[RecordBatch]) -> Vec<Eid> {
        let mut eids = Vec::new();
        for batch in batches {
            if let Some(eid_col) = batch.column_by_name("_eid")
                && let Some(eid_arr) = eid_col.as_any().downcast_ref::<UInt64Array>()
            {
                for i in 0..eid_arr.len() {
                    if !eid_arr.is_null(i) {
                        eids.push(Eid::new(eid_arr.value(i)));
                    }
                }
            }
        }
        eids
    }

    /// Find all non-deleted EIDs from the main edges table.
    pub async fn find_all_eids(backend: &dyn StorageBackend) -> Result<Vec<Eid>> {
        let batches = Self::execute_query(backend, "_deleted = false", Some(vec!["_eid"])).await?;
        Ok(Self::extract_eids(&batches))
    }

    /// Find EIDs by type name in the main edges table.
    pub async fn find_eids_by_type_name(
        backend: &dyn StorageBackend,
        type_name: &str,
    ) -> Result<Vec<Eid>> {
        let filter = format!(
            "_deleted = false AND type = '{}'",
            type_name.replace('\'', "''")
        );
        let batches = Self::execute_query(backend, &filter, Some(vec!["_eid"])).await?;
        Ok(Self::extract_eids(&batches))
    }

    /// Find properties for an edge by EID in the main edges table.
    ///
    /// Returns the props_json parsed into a Properties HashMap if found.
    /// This is used as a fallback for unknown/schemaless edge types.
    pub async fn find_props_by_eid(
        backend: &dyn StorageBackend,
        eid: Eid,
    ) -> Result<Option<Properties>> {
        // MVCC (review C2): the scan must see deletion tombstones — the
        // highest-version row wins, and a deleted winner yields `None`.
        // Filtering `_deleted = false` here would let an OLDER live version
        // resurrect an edge whose tombstone is the true (highest-version)
        // winner.
        let filter = format!("_eid = {}", eid.as_u64());
        let batches = Self::execute_query(
            backend,
            &filter,
            Some(vec!["props_json", "_version", "_deleted"]),
        )
        .await?;

        if batches.is_empty() {
            return Ok(None);
        }

        // Find the row with highest version (latest), tombstones included.
        let mut best_props: Option<Properties> = None;
        let mut best_version: u64 = 0;
        let mut best_deleted = false;

        for batch in &batches {
            let props_col = batch.column_by_name("props_json");
            let version_col = batch.column_by_name("_version");
            let deleted_col = batch
                .column_by_name("_deleted")
                .and_then(|c| c.as_any().downcast_ref::<arrow_array::BooleanArray>());

            if let (Some(props_arr), Some(ver_arr)) = (
                props_col.and_then(|c| c.as_any().downcast_ref::<arrow_array::LargeBinaryArray>()),
                version_col.and_then(|c| c.as_any().downcast_ref::<UInt64Array>()),
            ) {
                for i in 0..batch.num_rows() {
                    let version = if ver_arr.is_null(i) {
                        0
                    } else {
                        ver_arr.value(i)
                    };

                    if version >= best_version {
                        best_version = version;
                        best_deleted = deleted_col.is_some_and(|d| d.value(i));
                        best_props = if best_deleted {
                            Some(Properties::new())
                        } else {
                            Some(Self::parse_props_json(props_arr, i)?)
                        };
                    }
                }
            }
        }

        if best_deleted {
            return Ok(None);
        }
        Ok(best_props)
    }

    /// Parse props_json from a LargeBinaryArray (JSONB) at the given index.
    fn parse_props_json(arr: &arrow_array::LargeBinaryArray, idx: usize) -> Result<Properties> {
        if arr.is_null(idx) || arr.value(idx).is_empty() {
            return Ok(Properties::new());
        }
        let bytes = arr.value(idx);
        let uni_val = uni_common::cypher_value_codec::decode(bytes)
            .map_err(|e| anyhow!("Failed to decode CypherValue: {}", e))?;
        let json_val: serde_json::Value = uni_val.into();
        serde_json::from_value(json_val).map_err(|e| anyhow!("Failed to parse props_json: {}", e))
    }

    /// Find edge type name by EID in the main edges table.
    pub async fn find_type_by_eid(
        backend: &dyn StorageBackend,
        eid: Eid,
    ) -> Result<Option<String>> {
        // MVCC (review C2): take the highest-version row — including tombstones —
        // and return `None` if that winner is deleted. The old code filtered
        // `_deleted = false` and took `value(0)` with no version comparison, so
        // it could return an arbitrary (and possibly stale) edge type.
        let filter = format!("_eid = {}", eid.as_u64());
        let batches =
            Self::execute_query(backend, &filter, Some(vec!["type", "_version", "_deleted"]))
                .await?;

        let mut best_type: Option<String> = None;
        let mut best_version: u64 = 0;
        let mut best_deleted = false;

        for batch in &batches {
            let type_col = batch
                .column_by_name("type")
                .and_then(|c| c.as_any().downcast_ref::<arrow_array::StringArray>());
            let ver_col = batch
                .column_by_name("_version")
                .and_then(|c| c.as_any().downcast_ref::<UInt64Array>());
            let deleted_col = batch
                .column_by_name("_deleted")
                .and_then(|c| c.as_any().downcast_ref::<arrow_array::BooleanArray>());

            if let (Some(type_arr), Some(ver_arr)) = (type_col, ver_col) {
                for i in 0..batch.num_rows() {
                    let version = if ver_arr.is_null(i) {
                        0
                    } else {
                        ver_arr.value(i)
                    };

                    if version >= best_version {
                        best_version = version;
                        best_deleted = deleted_col.is_some_and(|d| d.value(i));
                        best_type = if best_deleted || type_arr.is_null(i) {
                            None
                        } else {
                            Some(type_arr.value(i).to_string())
                        };
                    }
                }
            }
        }

        if best_deleted {
            return Ok(None);
        }
        Ok(best_type)
    }

    /// Find edge data (eid, src_vid, dst_vid, props) by type name in the main edges table.
    ///
    /// Returns all non-deleted edges with the given type name.
    pub async fn find_edges_by_type_name(
        backend: &dyn StorageBackend,
        type_name: &str,
    ) -> Result<Vec<(Eid, Vid, Vid, Properties)>> {
        let filter = format!(
            "_deleted = false AND type = '{}'",
            type_name.replace('\'', "''")
        );
        // Fetch all columns for edge data
        let batches = Self::execute_query(backend, &filter, None).await?;

        let mut edges = Vec::new();
        for batch in &batches {
            Self::extract_edges_from_batch(batch, &mut edges)?;
        }

        Ok(edges)
    }

    /// Find edge data (eid, src_vid, dst_vid, edge_type, props) by multiple type names in the main edges table.
    ///
    /// Returns all non-deleted edges with any of the given type names.
    /// This is used for OR relationship type queries like `[:KNOWS|HATES]`.
    ///
    /// `endpoint_filter` pushes a bounded endpoint set into the scan (review
    /// perf #5: a 1-source schemaless traversal used to materialize the whole
    /// edge type). `None` keeps the full-type scan.
    pub async fn find_edges_by_type_names(
        backend: &dyn StorageBackend,
        type_names: &[&str],
        endpoint_filter: Option<(EndpointSide, &[Vid])>,
    ) -> Result<Vec<(Eid, Vid, Vid, String, Properties)>> {
        if type_names.is_empty() {
            return Ok(Vec::new());
        }

        // Build IN clause: type IN ('T1', 'T2', ...)
        let escaped_types: Vec<String> = type_names
            .iter()
            .map(|t| format!("'{}'", t.replace('\'', "''")))
            .collect();
        let base_filter = format!(
            "_deleted = false AND type IN ({})",
            escaped_types.join(", ")
        );

        let mut edges = Vec::new();
        match endpoint_filter {
            None => {
                // Fetch all columns for edge data
                let batches = Self::execute_query(backend, &base_filter, None).await?;
                for batch in &batches {
                    Self::extract_edges_with_type_from_batch(batch, &mut edges)?;
                }
            }
            Some((_, [])) => {}
            Some((side, vids)) => {
                // Chunked so the filter string stays parseable for large sets.
                const VID_CHUNK: usize = 8192;
                for chunk in vids.chunks(VID_CHUNK) {
                    let list = chunk
                        .iter()
                        .map(|v| v.as_u64().to_string())
                        .collect::<Vec<_>>()
                        .join(", ");
                    let endpoint_clause = match side {
                        EndpointSide::Src => format!("src_vid IN ({list})"),
                        EndpointSide::Dst => format!("dst_vid IN ({list})"),
                        EndpointSide::Either => {
                            format!("(src_vid IN ({list}) OR dst_vid IN ({list}))")
                        }
                    };
                    let filter = format!("{base_filter} AND {endpoint_clause}");
                    let batches = Self::execute_query(backend, &filter, None).await?;
                    for batch in &batches {
                        Self::extract_edges_with_type_from_batch(batch, &mut edges)?;
                    }
                }
            }
        }

        Ok(edges)
    }

    /// Extract edge data from a record batch (without edge type).
    fn extract_edges_from_batch(
        batch: &RecordBatch,
        edges: &mut Vec<(Eid, Vid, Vid, Properties)>,
    ) -> Result<()> {
        // Reuse the with-type extraction and discard the edge type
        let mut edges_with_type = Vec::new();
        Self::extract_edges_with_type_from_batch(batch, &mut edges_with_type)?;
        edges.extend(
            edges_with_type
                .into_iter()
                .map(|(eid, src, dst, _type, props)| (eid, src, dst, props)),
        );
        Ok(())
    }

    /// Extract edge data with type from a record batch.
    fn extract_edges_with_type_from_batch(
        batch: &RecordBatch,
        edges: &mut Vec<(Eid, Vid, Vid, String, Properties)>,
    ) -> Result<()> {
        let Some(eid_arr) = batch
            .column_by_name("_eid")
            .and_then(|c| c.as_any().downcast_ref::<UInt64Array>())
        else {
            return Ok(());
        };
        let Some(src_arr) = batch
            .column_by_name("src_vid")
            .and_then(|c| c.as_any().downcast_ref::<UInt64Array>())
        else {
            return Ok(());
        };
        let Some(dst_arr) = batch
            .column_by_name("dst_vid")
            .and_then(|c| c.as_any().downcast_ref::<UInt64Array>())
        else {
            return Ok(());
        };
        let type_arr = batch
            .column_by_name("type")
            .and_then(|c| c.as_any().downcast_ref::<arrow_array::StringArray>());
        let props_arr = batch
            .column_by_name("props_json")
            .and_then(|c| c.as_any().downcast_ref::<arrow_array::LargeBinaryArray>());

        for i in 0..batch.num_rows() {
            if eid_arr.is_null(i) || src_arr.is_null(i) || dst_arr.is_null(i) {
                continue;
            }

            let eid = Eid::new(eid_arr.value(i));
            let src_vid = Vid::new(src_arr.value(i));
            let dst_vid = Vid::new(dst_arr.value(i));
            let edge_type = type_arr
                .filter(|arr| !arr.is_null(i))
                .map(|arr| arr.value(i).to_string())
                .unwrap_or_default();
            let props = props_arr
                .map(|arr| Self::parse_props_json(arr, i))
                .transpose()?
                .unwrap_or_default();

            edges.push((eid, src_vid, dst_vid, edge_type, props));
        }

        Ok(())
    }
}

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

    #[test]
    fn test_main_edge_schema() {
        let schema = MainEdgeDataset::get_arrow_schema();
        assert_eq!(schema.fields().len(), 9);
        assert!(schema.field_with_name("_eid").is_ok());
        assert!(schema.field_with_name("src_vid").is_ok());
        assert!(schema.field_with_name("dst_vid").is_ok());
        assert!(schema.field_with_name("type").is_ok());
        assert!(schema.field_with_name("props_json").is_ok());
        assert!(schema.field_with_name("_deleted").is_ok());
        assert!(schema.field_with_name("_version").is_ok());
        assert!(schema.field_with_name("_created_at").is_ok());
        assert!(schema.field_with_name("_updated_at").is_ok());
    }

    #[test]
    fn test_build_record_batch() {
        use uni_common::Value;
        let mut props = HashMap::new();
        props.insert("weight".to_string(), Value::Float(0.5));

        let edges = vec![(
            Eid::new(1),
            Vid::new(1),
            Vid::new(2),
            "KNOWS".to_string(),
            props,
            false,
            1u64,
        )];

        let batch = MainEdgeDataset::build_record_batch(&edges, None, None).unwrap();
        assert_eq!(batch.num_rows(), 1);
        assert_eq!(batch.num_columns(), 9);
    }

    #[test]
    fn test_build_record_batch_multiple_edges() {
        use uni_common::Value;

        let edges = vec![
            (
                Eid::new(1),
                Vid::new(1),
                Vid::new(2),
                "KNOWS".to_string(),
                HashMap::from([("since".to_string(), Value::Int(2020))]),
                false,
                1u64,
            ),
            (
                Eid::new(2),
                Vid::new(2),
                Vid::new(3),
                "WORKS_AT".to_string(),
                HashMap::new(),
                false,
                2u64,
            ),
            (
                Eid::new(3),
                Vid::new(1),
                Vid::new(3),
                "KNOWS".to_string(),
                HashMap::new(),
                true, // deleted
                3u64,
            ),
        ];

        let batch = MainEdgeDataset::build_record_batch(&edges, None, None).unwrap();
        assert_eq!(batch.num_rows(), 3);
        assert_eq!(batch.num_columns(), 9);

        // Verify type column has correct values
        let type_col = batch
            .column_by_name("type")
            .unwrap()
            .as_any()
            .downcast_ref::<arrow_array::StringArray>()
            .unwrap();
        assert_eq!(type_col.value(0), "KNOWS");
        assert_eq!(type_col.value(1), "WORKS_AT");
        assert_eq!(type_col.value(2), "KNOWS");
    }

    #[test]
    fn test_build_record_batch_with_timestamps() {
        let edges = vec![(
            Eid::new(1),
            Vid::new(1),
            Vid::new(2),
            "KNOWS".to_string(),
            HashMap::new(),
            false,
            1u64,
        )];

        let mut created_at: HashMap<Eid, i64> = HashMap::new();
        created_at.insert(Eid::new(1), 1_000_000_000);

        let mut updated_at: HashMap<Eid, i64> = HashMap::new();
        updated_at.insert(Eid::new(1), 2_000_000_000);

        let batch =
            MainEdgeDataset::build_record_batch(&edges, Some(&created_at), Some(&updated_at))
                .unwrap();
        assert_eq!(batch.num_rows(), 1);

        // Timestamp columns should exist and not be all null
        let created_col = batch.column_by_name("_created_at").unwrap();
        assert!(!created_col.is_null(0), "created_at should be populated");
    }

    /// MVCC regression (review C2): a deletion tombstone written at a higher
    /// version must win over the older live row. `find_props_by_eid` filtered
    /// `_deleted = false` before version-ranking (so an older live version
    /// resurrected a deleted edge), and `find_type_by_eid` took `value(0)` with
    /// no version comparison.
    #[tokio::test]
    async fn test_edge_key_reads_respect_tombstone_winner() {
        use crate::backend::lance::LanceDbBackend;
        use uni_common::Value;

        let dir = tempfile::TempDir::new().unwrap();
        let be = LanceDbBackend::connect(dir.path().to_str().unwrap(), None)
            .await
            .unwrap();
        let backend: &dyn StorageBackend = &be;

        let mut props = HashMap::new();
        props.insert("weight".to_string(), Value::Float(0.5));

        // v1: live edge.
        let live = MainEdgeDataset::build_record_batch(
            &[(
                Eid::new(1),
                Vid::new(1),
                Vid::new(2),
                "KNOWS".to_string(),
                props.clone(),
                false,
                1u64,
            )],
            None,
            None,
        )
        .unwrap();
        MainEdgeDataset::write_batch(backend, live).await.unwrap();

        // Sanity: visible while live.
        assert!(
            MainEdgeDataset::find_props_by_eid(backend, Eid::new(1))
                .await
                .unwrap()
                .is_some()
        );
        assert_eq!(
            MainEdgeDataset::find_type_by_eid(backend, Eid::new(1))
                .await
                .unwrap(),
            Some("KNOWS".to_string())
        );

        // v2: deletion tombstone at a higher version — the winning row.
        let dead = MainEdgeDataset::build_record_batch(
            &[(
                Eid::new(1),
                Vid::new(1),
                Vid::new(2),
                "KNOWS".to_string(),
                props,
                true,
                2u64,
            )],
            None,
            None,
        )
        .unwrap();
        MainEdgeDataset::write_batch(backend, dead).await.unwrap();

        assert_eq!(
            MainEdgeDataset::find_props_by_eid(backend, Eid::new(1))
                .await
                .unwrap(),
            None,
            "deleted (highest-version) winner must not resurrect edge props"
        );
        assert_eq!(
            MainEdgeDataset::find_type_by_eid(backend, Eid::new(1))
                .await
                .unwrap(),
            None,
            "deleted winner must not return an edge type"
        );
    }
}