genegraph_storage/

lance_storage_graph.rs

//! Lance storage backend for graph embeddings.
//!
//! Async-first implementation that matches the async `StorageBackend` trait:
//! - All I/O is async, no internal `block_on` or runtime creation.
//! - Callers (CLI, tests, services) are responsible for providing a Tokio runtime.

use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;

use arrow::array::{Array as ArrowArray, ArrayRef, Float32Array, Float64Array, UInt32Array};
use arrow::datatypes::{DataType, Field, Float64Type, Int64Type, Schema, UInt32Type};
use arrow::record_batch::RecordBatch;
use log::{debug, info};
use smartcore::linalg::basic::arrays::Array;
use smartcore::linalg::basic::matrix::DenseMatrix;
use sprs::CsMat;

use crate::graph::{GraphEdge, GraphReadOptions, GraphWriteOptions, StoredGraph};
use crate::metadata::FileInfo;
use crate::metadata::GeneMetadata;
use crate::traits::backend::StorageBackend;
use crate::traits::lance::{
    LanceStorage, graph_record_batch, stored_graph_from_batch_with_options,
    validate_vector_space_schema, vector_space_collection_batch,
};
use crate::{StorageError, StorageResult};

/// Checked `usize -> u32` conversion.
///
/// Values above `u32::MAX` must surface an error instead of being silently
/// truncated into wrong stored data (issue #51).
fn checked_u32_values(values: &[usize], what: &str) -> StorageResult<Vec<u32>> {
    values
        .iter()
        .map(|&v| {
            u32::try_from(v).map_err(|_| {
                StorageError::Overflow(format!(
                    "{} value {} exceeds u32::MAX and would be silently truncated",
                    what, v
                ))
            })
        })
        .collect()
}

/// Creates the parent directory of an explicit collection path (mirrors
/// `save_dense_to_file`): registry-free writes may target any location.
async fn ensure_parent_dir(path: &Path) -> StorageResult<()> {
    let parent = path
        .parent()
        .filter(|p| !p.as_os_str().is_empty())
        .ok_or_else(|| StorageError::Invalid(format!("path has no parent directory: {path:?}")))?;
    tokio::fs::create_dir_all(parent)
        .await
        .map_err(|e| StorageError::Io(format!("create dir {parent:?}: {e}")))
}

/// Lance-based storage backend for ArrowSpace graph embeddings.
///
/// Stores dense and sparse matrices as Lance datasets using a columnar format
/// (`row`, `col`, `value` for sparse; `col_*` for dense) schema for efficient
/// random and columnar access.
///
/// Metadata must be seeded before any `save_*` call so the storage directory
/// is initialized; the example below is executed as a doc-test on every
/// `cargo test` run so it cannot go stale.
///
/// # Examples
///
/// ```
/// use genegraph_storage::lance_storage_graph::LanceStorageGraph;
/// use genegraph_storage::metadata::GeneMetadata;
/// use genegraph_storage::traits::backend::StorageBackend;
/// use genegraph_storage::traits::metadata::Metadata;
/// use smartcore::linalg::basic::arrays::{Array, Array2};
/// use smartcore::linalg::basic::matrix::DenseMatrix;
///
/// let base = std::env::temp_dir().join(format!("genegraph_doc_{}", std::process::id()));
/// # tokio::runtime::Runtime::new().unwrap().block_on(async {
/// let storage = LanceStorageGraph::new(
///     base.to_string_lossy().to_string(),
///     "doc_example".to_string(),
/// )
/// .expect("valid instance name");
///
/// // some 2D data
/// let dense: Vec<Vec<f64>> = vec![vec![0.1, 0.4], vec![0.5, 0.2], vec![0.03, 0.8]];
/// let (nitems, nfeatures) = (dense.len(), dense[0].len());
/// let data = DenseMatrix::<f64>::from_iterator(
///     dense.iter().flatten().copied(),
///     nitems,
///     nfeatures,
///     0,
/// );
///
/// // seed metadata FIRST to initialize the storage directory
/// let md = GeneMetadata::seed_metadata("doc_example", nitems, nfeatures, &storage)
///     .await
///     .unwrap();
/// let md_path = storage.save_metadata(&md).await.unwrap();
///
/// // your data is saved in an efficient Lance format
/// storage.save_dense("my_dataset", &data, &md_path).await.unwrap();
///
/// // Loading back
/// let loaded = storage.load_dense("my_dataset").await.unwrap();
/// assert_eq!(loaded.shape(), (nitems, nfeatures));
/// # });
/// # std::fs::remove_dir_all(&base).ok();
/// ```
#[derive(Debug, Clone)]
pub struct LanceStorageGraph {
    pub(crate) base: String,
    pub(crate) name: String,
}

impl LanceStorageGraph {
    /// Creates a new Lance storage backend.
    ///
    /// This is used for on-the-fly creation. For proper setup use `Genefold<...>::seed`.
    ///
    /// # Arguments
    ///
    /// * `base` - Base directory path for all storage files
    /// * `name` - Name prefix for this storage instance
    ///
    /// The name must not carry the reserved generation suffix `__g<n>`
    /// (#117-3): only [`Self::scoped_generation`] mints generation-qualified
    /// names, otherwise generation discovery and sweep would treat the
    /// instance as a generation of another logical dataset.
    pub fn new(base: String, name: String) -> StorageResult<Self> {
        crate::generations::validate_logical_name(&name)?;
        Ok(Self::from_parts(base, name))
    }

    /// Unvalidated constructor for vetted names: stored `name_id`s read
    /// back through [`Self::spawn`], generation handles from
    /// [`Self::scoped_generation`], and internal temporary handles.
    fn from_parts(base: String, name: String) -> Self {
        info!("Creating LanceStorage at base={}, name={}", base, name);
        Self { base, name }
    }

    /// Spawn a LanceStorage from an existing seeded directory (with metadata.json)
    pub async fn spawn(base_path: String) -> Result<(Self, GeneMetadata), StorageError> {
        // Reuse the generic `exists` helper from the StorageBackend trait
        let (exists, md_path) = Self::exists(&base_path);

        // Replace assert! with proper error handling
        if !exists || md_path.is_none() {
            return Err(StorageError::Invalid(format!(
                "Metadata does not exist in base path: {}",
                base_path
            )));
        }

        // Load metadata from the discovered metadata.json
        let metadata = GeneMetadata::read(md_path.unwrap()).await?;

        // Construct the LanceStorage using the metadata-provided nameid
        let storage = Self::from_parts(base_path.clone(), metadata.name_id.clone());
        Ok((storage, metadata))
    }

    /// A handle over generation `gen` of this logical dataset: artifact
    /// paths resolve as `{logical}__g{gen}_{key}.lance` and the metadata
    /// path as `{logical}__g{gen}_metadata.json` — the per-generation
    /// commit pointer (#93, RFC #81-P5). The logical identity stays
    /// recoverable via [`crate::generations::logical_name`].
    pub fn scoped_generation(&self, generation: u64) -> Self {
        let logical = crate::generations::logical_name(&self.name);
        Self::from_parts(
            self.base.clone(),
            crate::generations::generation_name(logical, generation),
        )
    }
}

impl LanceStorage for LanceStorageGraph {}

// =========
// Dense file I/O helpers (#117-4): one function per format branch so each
// can be read and tested in isolation. The dispatchers above stay thin.
// =========

/// Parent directory of `path`, UTF-8 validated.
fn parent_dir_of(path: &Path) -> StorageResult<String> {
    path.parent()
        .ok_or_else(|| StorageError::Invalid(format!("Path has no parent: {:?}", path)))?
        .to_str()
        .map(str::to_string)
        .ok_or_else(|| StorageError::Invalid(format!("Non-UTF8 parent path for {:?}", path)))
}

/// Lance branch of `load_dense_from_file`: read the dataset directory
/// through a temporary backend rooted at the file's parent.
async fn load_dense_lance(path: &Path) -> StorageResult<DenseMatrix<f64>> {
    use crate::traits::backend::StorageBackend;
    use crate::traits::lance::LanceStorage;

    let parent = parent_dir_of(path)?;
    let tmp_storage = LanceStorageGraph::from_parts(parent, String::from("tmp_storage"));

    let uri = LanceStorageGraph::path_to_uri(path)?;
    let batch = tmp_storage.read_lance_all_batches_async(uri).await?;
    let matrix = tmp_storage.from_dense_record_batch(&batch)?;
    info!(
        "Loaded dense matrix from Lance: {} x {}",
        matrix.shape().0,
        matrix.shape().1
    );
    Ok(matrix)
}

/// Parquet branch of `load_dense_from_file`, reader half.
///
/// Parquet readers require a synchronous `Read` impl; run the blocking
/// open/read/concat on the dedicated blocking pool so the async executor
/// thread is not stalled (issue #52).
async fn read_parquet_batches(path: &Path) -> StorageResult<RecordBatch> {
    use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;

    let owned_path = path.to_path_buf();
    let combined = tokio::task::spawn_blocking(move || -> StorageResult<RecordBatch> {
        let file = std::fs::File::open(&owned_path)
            .map_err(|e| StorageError::Io(format!("Failed to open parquet file: {}", e)))?;

        let builder = ParquetRecordBatchReaderBuilder::try_new(file).map_err(|e| {
            StorageError::Parquet(format!("Failed to create parquet reader: {}", e))
        })?;
        let reader = builder
            .build()
            .map_err(|e| StorageError::Parquet(format!("Failed to build parquet reader: {}", e)))?;

        let batches: Vec<RecordBatch> = reader
            .collect::<Result<Vec<_>, _>>()
            .map_err(|e| StorageError::Parquet(format!("Failed to read parquet batch: {}", e)))?;

        if batches.is_empty() {
            return Err(StorageError::Invalid(format!(
                "Empty parquet dataset at {:?}",
                owned_path
            )));
        }

        let schema = batches[0].schema();
        arrow::compute::concat_batches(&schema, &batches).map_err(|e| {
            StorageError::Parquet(format!("Failed to concatenate parquet batches: {}", e))
        })
    })
    .await
    .map_err(|e| StorageError::Io(format!("Parquet reader task failed: {}", e)))??;
    Ok(combined)
}

/// Parquet branch of `load_dense_from_file`, layout half.
///
/// Supported layouts: `vector: FixedSizeList<Float64>` (decodes like the
/// Lance path) and the legacy wide layout with `Float64` columns named
/// `col_0..col_N`; anything else is rejected with `StorageError::Invalid`.
fn dense_matrix_from_parquet(
    path: &Path,
    combined: &RecordBatch,
) -> StorageResult<DenseMatrix<f64>> {
    use crate::traits::backend::StorageBackend;

    let schema = combined.schema();
    let fields = schema.fields();
    let is_vector = fields.len() == 1
        && matches!(
            fields[0].data_type(),
            DataType::FixedSizeList(inner, _)
                if matches!(inner.data_type(), DataType::Float64)
        );

    let is_wide_col = !is_vector
        && !fields.is_empty()
        && fields
            .iter()
            .all(|f| matches!(f.data_type(), DataType::Float64))
        && fields.iter().any(|f| f.name().starts_with("col_"));

    let matrix = if is_vector {
        // New format already: vector column (FixedSizeList<Float64>)
        // Reuse the same decoding as Lance.
        let parent = parent_dir_of(path)?;
        let tmp_storage = LanceStorageGraph::from_parts(parent, String::from("tmp_storage"));
        tmp_storage.from_dense_record_batch(combined)?
    } else if is_wide_col {
        // Old wide columnar: columns like col_0, col_1, ... as Float64
        let n_rows = combined.num_rows();
        let n_cols = combined.num_columns();
        if n_rows == 0 || n_cols == 0 {
            return Err(StorageError::Invalid(format!(
                "Cannot load empty wide-column parquet at {:?}",
                path
            )));
        }

        let mut data = Vec::with_capacity(n_rows * n_cols);
        for col_idx in 0..n_cols {
            let col = combined.column(col_idx);
            let arr = col.as_any().downcast_ref::<Float64Array>().ok_or_else(|| {
                StorageError::Invalid(format!(
                    "Wide-column parquet expects Float64, got {:?} in column {}",
                    col.data_type(),
                    col_idx
                ))
            })?;
            // Build column-major storage: all rows for col 0, then col 1, ...
            for row_idx in 0..n_rows {
                data.push(arr.value(row_idx));
            }
        }

        DenseMatrix::new(n_rows, n_cols, data, true)
            .map_err(|e| StorageError::Invalid(e.to_string()))?
    } else {
        return Err(StorageError::Invalid(format!(
            "Unsupported Parquet schema at {:?}: expected FixedSizeList<Float64> \
             or wide Float64 columns named col_*",
            path
        )));
    };

    info!(
        "Loaded dense matrix from Parquet: {} x {}",
        matrix.shape().0,
        matrix.shape().1
    );
    Ok(matrix)
}

/// Builds the vector-layout RecordBatch for `data` and verifies the row
/// count survived the conversion.
fn checked_dense_batch(
    storage: &LanceStorageGraph,
    data: &DenseMatrix<f64>,
) -> StorageResult<RecordBatch> {
    use crate::traits::backend::StorageBackend;

    let batch = storage.to_dense_record_batch(data)?;
    let (n_rows, _) = data.shape();
    if batch.num_rows() != n_rows {
        return Err(StorageError::Invalid(format!(
            "RecordBatch has {} rows but matrix has {} rows",
            batch.num_rows(),
            n_rows
        )));
    }
    Ok(batch)
}

/// Lance branch of `save_dense_to_file`.
async fn save_dense_lance(
    parent_dir: String,
    data: &DenseMatrix<f64>,
    path: &Path,
) -> StorageResult<()> {
    use crate::traits::backend::StorageBackend;
    use crate::traits::lance::LanceStorage;

    let tmp_storage = LanceStorageGraph::from_parts(parent_dir, String::from("tmp_storage"));
    let batch = checked_dense_batch(&tmp_storage, data)?;
    debug!(
        "Created RecordBatch with {} rows for Lance",
        batch.num_rows()
    );

    let uri = LanceStorageGraph::path_to_uri(path)?;
    tmp_storage.write_lance_batch_async(uri, batch).await?;
    let (n_rows, n_cols) = data.shape();
    info!("Saved dense matrix to Lance: {} x {}", n_rows, n_cols);
    Ok(())
}

/// Parquet branch of `save_dense_to_file`.
///
/// The parquet writer is synchronous: run it on the blocking pool so the
/// async executor thread is not stalled (see #52 for the matching read
/// path).
async fn save_dense_parquet(
    parent_dir: String,
    data: &DenseMatrix<f64>,
    path: &Path,
) -> StorageResult<()> {
    use parquet::arrow::ArrowWriter;
    use parquet::file::properties::WriterProperties;
    use std::fs::File;

    let tmp_storage = LanceStorageGraph::from_parts(parent_dir, String::from("tmp_storage"));
    let batch = checked_dense_batch(&tmp_storage, data)?;
    debug!(
        "Created RecordBatch with {} rows for Parquet",
        batch.num_rows()
    );

    let owned_path = path.to_path_buf();
    tokio::task::spawn_blocking(move || -> StorageResult<()> {
        let file = File::create(&owned_path)
            .map_err(|e| StorageError::Io(format!("Failed to create parquet file: {}", e)))?;

        let props = WriterProperties::builder()
            .set_compression(parquet::basic::Compression::SNAPPY)
            .build();

        let mut writer = ArrowWriter::try_new(file, batch.schema(), Some(props)).map_err(|e| {
            StorageError::Parquet(format!("Failed to create parquet writer: {}", e))
        })?;

        writer
            .write(&batch)
            .map_err(|e| StorageError::Parquet(format!("Failed to write batch: {}", e)))?;

        writer
            .close()
            .map_err(|e| StorageError::Parquet(format!("Failed to close writer: {}", e)))?;

        Ok(())
    })
    .await
    .map_err(|e| StorageError::Io(format!("parquet writer task failed: {}", e)))??;

    let (n_rows, n_cols) = data.shape();
    info!("Saved dense matrix to Parquet: {} x {}", n_rows, n_cols);
    Ok(())
}

impl StorageBackend for LanceStorageGraph {
    fn get_base(&self) -> String {
        self.base.clone()
    }

    fn get_name(&self) -> String {
        self.name.clone()
    }

    fn base_path(&self) -> PathBuf {
        PathBuf::from(&self.base)
    }

    fn metadata_path(&self) -> PathBuf {
        crate::generations::metadata_file_path(&self.base_path(), &self.name)
    }

    /// Converts the base path for the store to a `file://` URI for Lance.
    fn basepath_to_uri(&self) -> StorageResult<String> {
        Self::path_to_uri(PathBuf::from(self.base.clone()).as_path())
    }

    /// Save dense matrix using Lance-optimized vector format.
    ///
    /// Each row of the matrix becomes a FixedSizeList entry for efficient vector operations.
    /// This format is optimized for vector search and enables Lance's full-zip encoding.
    ///
    /// # Arguments
    /// * `filename` - any name
    /// * `matrix` - Dense matrix to save (N rows × F cols)
    /// * `md_path` - Metadata file path for validation
    async fn save_dense(
        &self,
        key: &str,
        matrix: &DenseMatrix<f64>,
        md_path: &Path,
    ) -> StorageResult<()> {
        self.validate_initialized(md_path)?;
        let path = self.file_path(key);
        let (n_rows, n_cols) = matrix.shape();

        info!(
            "Saving dense {} matrix: {} x {} at {:?}",
            key, n_rows, n_cols, path
        );

        // Convert to Lance-optimized RecordBatch (FixedSizeList format)
        let batch = self.to_dense_record_batch(matrix)?;

        // Verify batch has correct number of rows
        if batch.num_rows() != n_rows {
            return Err(StorageError::Invalid(format!(
                "RecordBatch has {} rows but matrix has {} rows",
                batch.num_rows(),
                n_rows
            )));
        }

        {
            // Write to Lance
            let uri = Self::path_to_uri(&path)?;
            self.write_lance_batch_async(uri, batch).await?;
            // Commit the registry entry through the instance's commit actor:
            // concurrent save_* calls serialize their read-modify-write
            // cycles instead of losing updates (duva actor model).
            crate::commit::with_commit_actor(&self.metadata_path(), || async {
                let mut md = self.load_metadata().await?;
                md = md.add_file(
                    key,
                    FileInfo::new(
                        crate::generations::artifact_file_name(self.get_name().as_str(), key),
                        "dense",
                        matrix.shape(),
                        None,
                        None,
                    )?,
                );
                self.save_metadata(&md).await
            })
            .await?;
            info!("Dense {} matrix saved successfully", key);
        }
        Ok(())
    }

    /// Load dense matrix from Lance-optimized vector format.
    ///
    /// Reads FixedSizeList vectors and reconstructs a column-major DenseMatrix.
    ///
    /// # Arguments
    /// * `filename` - any name previously assigned
    ///
    /// # Returns
    /// Column-major DenseMatrix matching smartcore conventions
    async fn load_dense(&self, key: &str) -> StorageResult<DenseMatrix<f64>> {
        let path = self.file_path(key);
        info!("Loading dense {} matrix from {:?}", key, path);

        // Read all batches from Lance (may span multiple batches for large datasets)
        let uri = Self::path_to_uri(&path)?;
        let batch = self.read_lance_all_batches_async(uri).await?;

        // Convert from FixedSizeList format to DenseMatrix
        let matrix = self.from_dense_record_batch(&batch)?;

        let (n_rows, n_cols) = matrix.shape();
        info!("Loaded dense {} matrix: {} x {}", key, n_rows, n_cols);

        Ok(matrix)
    }

    /// Load initial data using columnar format from a file path.
    ///
    /// Async test helper that avoids any internal blocking runtimes.
    /// Format dispatch only (#117-4): each branch lives in its own helper.
    async fn load_dense_from_file(&self, path: &Path) -> StorageResult<DenseMatrix<f64>> {
        info!("Loading dense matrix from file (async): {:?}", path);

        if !path.exists() {
            return Err(StorageError::Invalid(format!(
                "Dense file does not exist: {:?}",
                path
            )));
        }

        let extension = path
            .extension()
            .and_then(|e| e.to_str())
            .ok_or_else(|| StorageError::Invalid(format!("Invalid file path: {:?}", path)))?;

        match extension {
            "lance" => load_dense_lance(path).await,
            "parquet" => {
                let combined = read_parquet_batches(path).await?;
                dense_matrix_from_parquet(path, &combined)
            }
            _ => Err(StorageError::Invalid(format!(
                "Unsupported file format: {}. Only .lance and .parquet are supported",
                extension
            ))),
        }
    }

    fn file_path(&self, key: &str) -> PathBuf {
        crate::generations::artifact_file_path(&self.base_path(), &self.name, key)
    }

    // =========
    // ASYNC API (matches StorageBackend)
    // =========

    async fn save_sparse(
        &self,
        key: &str,
        matrix: &CsMat<f64>,
        md_path: &Path,
    ) -> StorageResult<()> {
        self.validate_initialized(md_path)?;
        // #102, fail early and typed: an nnz=0 sparse matrix (a fully
        // disconnected adjacency — a legitimate data state, e.g. from a
        // too-small eps) cannot be persisted as a sparse artifact. Reject
        // here, before the metadata commit cycle and before any artifact
        // write, so no partial directory is left behind.
        if matrix.nnz() == 0 {
            return Err(StorageError::Invalid(format!(
                "sparse '{key}' is fully disconnected: adjacency nnz=0 ({}x{}) — \
                 raise --eps/--k or gate the input before persisting",
                matrix.rows(),
                matrix.cols()
            )));
        }
        let path = self.file_path(key);
        info!(
            "Saving sparse {} matrix: {} x {}, nnz={} at {:?}",
            key,
            matrix.rows(),
            matrix.cols(),
            matrix.nnz(),
            path
        );

        let filetype = FileInfo::which_filetype(key)?;
        {
            crate::commit::with_commit_actor(&self.metadata_path(), || async {
                let mut metadata = self.load_metadata().await?;
                metadata = metadata.add_file(
                    key,
                    FileInfo::new(
                        crate::generations::artifact_file_name(self.get_name().as_str(), key),
                        filetype.as_str(),
                        (matrix.rows(), matrix.cols()),
                        Some(matrix.nnz()),
                        None,
                    )?,
                );
                self.save_metadata(&metadata).await
            })
            .await?;

            let batch = self.to_sparse_record_batch(matrix)?;
            let uri = Self::path_to_uri(&path)?;
            self.write_lance_batch_async(uri, batch).await?;
        }
        info!("Sparse matrix {} saved successfully", filetype);
        Ok(())
    }

    async fn load_sparse(&self, key: &str) -> StorageResult<CsMat<f64>> {
        info!("Loading sparse {} matrix", key);

        let metadata = self.load_metadata().await?;
        let filetype = FileInfo::which_filetype(key)?;
        let file_info = metadata
            .files
            .get(key)
            .ok_or_else(|| StorageError::Invalid(format!("{key} not found in metadata")))?;

        let expected_rows = file_info.rows;
        let expected_cols = file_info.cols;
        debug!(
            "Expected dimensions from storage metadata: {} x {}",
            expected_rows, expected_cols
        );

        let path = self.file_path(key);
        let uri = Self::path_to_uri(&path)?;
        let batch = self.read_lance_all_batches_async(uri).await?;
        let matrix = self.from_sparse_record_batch(batch, expected_rows, expected_cols)?;
        info!(
            "Sparse {} matrix loaded: {} x {}, nnz={}",
            filetype,
            matrix.rows(),
            matrix.cols(),
            matrix.nnz()
        );
        Ok(matrix)
    }

    async fn save_lambdas(&self, lambdas: &[f64], md_path: &Path) -> StorageResult<()> {
        info!("Saving {} lambda values", lambdas.len());
        self.save_primitive_column::<Float64Type>("lambdas", "lambda", lambdas.to_vec(), md_path)
            .await
    }

    async fn load_lambdas(&self) -> StorageResult<Vec<f64>> {
        let path = self.file_path("lambdas");
        info!("Loading lambda values from {:?}", path);

        let uri = Self::path_to_uri(&path)?;
        let batch = self.read_lance_all_batches_async(uri).await?;
        let arr = batch
            .column(0)
            .as_any()
            .downcast_ref::<Float64Array>()
            .ok_or_else(|| StorageError::Invalid("lambda column type mismatch".into()))?;

        let lambdas: Vec<f64> = (0..arr.len()).map(|i| arr.value(i)).collect();
        info!("Loaded {} lambda values", lambdas.len());
        Ok(lambdas)
    }

    async fn save_vector(&self, key: &str, vector: &[f64], md_path: &Path) -> StorageResult<()> {
        info!("Saving {} values for vector {}", vector.len(), key);
        self.save_primitive_column::<Float64Type>(key, "element", vector.to_vec(), md_path)
            .await
    }

    async fn save_index(&self, key: &str, vector: &[usize], md_path: &Path) -> StorageResult<()> {
        info!("Saving {} values for index {}", vector.len(), key);
        // Checked cast: usize values above u32::MAX must not be silently
        // truncated (issue #51).
        let values = checked_u32_values(vector, "index")?;
        self.save_primitive_column::<UInt32Type>(key, "id", values, md_path)
            .await
    }

    async fn load_vector(&self, filename: &str) -> StorageResult<Vec<f64>> {
        let path = self.file_path(filename);
        info!("Loading vector {} from {:?}", filename, path);

        let uri = Self::path_to_uri(&path)?;
        let batch = self.read_lance_all_batches_async(uri).await?;
        let arr = batch
            .column(0)
            .as_any()
            .downcast_ref::<Float64Array>()
            .ok_or_else(|| StorageError::Invalid("column type mismatch".into()))?;

        let vector: Vec<f64> = (0..arr.len()).map(|i| arr.value(i)).collect();
        info!("Loaded {} vector values for {}", vector.len(), filename);
        Ok(vector)
    }

    async fn load_index(&self, filename: &str) -> StorageResult<Vec<usize>> {
        let path = self.file_path(filename);
        info!("Loading vector {} from {:?}", filename, path);

        let uri = Self::path_to_uri(&path)?;
        let batch = self.read_lance_all_batches_async(uri).await?;
        let arr = batch
            .column(0)
            .as_any()
            .downcast_ref::<UInt32Array>()
            .ok_or_else(|| StorageError::Invalid("column type mismatch".into()))?;

        let vector: Vec<usize> = (0..arr.len()).map(|i| arr.value(i) as usize).collect();
        info!("Loaded {} vector values for {}", vector.len(), filename);
        Ok(vector)
    }

    /// Save dense matrix to file in columnar format (col_0, col_1, ..., col_N)
    ///
    /// Async test helper that avoids any internal blocking runtimes.
    async fn save_dense_to_file(data: &DenseMatrix<f64>, path: &Path) -> StorageResult<()> {
        info!("Saving dense matrix to file (async): {:?}", path);

        // Create missing parent directories instead of failing late in the
        // format-specific writers.
        let parent = path
            .parent()
            .filter(|p| !p.as_os_str().is_empty())
            .ok_or_else(|| {
                StorageError::Invalid(format!("path has no parent directory: {:?}", path))
            })?;
        tokio::fs::create_dir_all(parent)
            .await
            .map_err(|e| StorageError::Io(format!("create dir {:?}: {}", parent, e)))?;
        let parent_str = parent
            .to_str()
            .ok_or_else(|| StorageError::Invalid(format!("non-UTF8 parent path for {:?}", path)))?
            .to_string();

        let extension = path
            .extension()
            .and_then(|e| e.to_str())
            .ok_or_else(|| StorageError::Invalid(format!("Invalid file path: {:?}", path)))?;

        let (n_rows, n_cols) = data.shape();
        info!("Saving matrix: {} rows x {} cols", n_rows, n_cols);

        match extension {
            "lance" => save_dense_lance(parent_str, data, path).await,
            "parquet" => save_dense_parquet(parent_str, data, path).await,
            _ => Err(StorageError::Invalid(format!(
                "Unsupported file format: {}. Only .lance and .parquet are supported",
                extension
            ))),
        }
    }

    // =========
    // NAMED COLLECTIONS (RFC #81)
    // =========

    async fn save_vectors_with(
        &self,
        name: &str,
        batch: &RecordBatch,
        properties: &BTreeMap<String, String>,
        md_path: &Path,
    ) -> StorageResult<()> {
        self.validate_initialized(md_path)?;
        let (batch, dim) = vector_space_collection_batch(batch, properties)?;
        let rows = batch.num_rows();

        let path = self.file_path(name);
        info!(
            "Saving vector-space collection '{}' ({} rows, vector dim {}) at {:?}",
            name, rows, dim, path
        );
        let uri = Self::path_to_uri(&path)?;
        // Artifact first, registry publish second: the registry update
        // (through the commit actor) is the single commit point, so a
        // live entry only exists once the artifact is durable at its
        // final location (#106 review). A failed artifact write leaves
        // no registry entry; a failed publish leaves the artifact as
        // unreferenced residue that a later sweep reclaims — never a
        // live entry pointing at a missing artifact. Crash recovery for
        // the residue case is a sweep/gc, not a compensating rewrite.
        self.write_lance_batch_async(uri, batch).await?;
        crate::commit::with_commit_actor(&self.metadata_path(), || async {
            let mut md = self.load_metadata().await?;
            let mut info = FileInfo::new(
                crate::generations::artifact_file_name(self.get_name().as_str(), name),
                "vectors",
                (rows, dim as usize),
                None,
                None,
            )?;
            info.properties = properties.clone();
            md = md.add_file(name, info);
            self.save_metadata(&md).await
        })
        .await?;
        info!("Vector-space collection '{}' saved successfully", name);
        Ok(())
    }

    async fn load_vectors(&self, name: &str) -> StorageResult<RecordBatch> {
        info!("Loading vector-space collection '{}'", name);
        self.load_vectors_from_path(&self.file_path(name)).await
    }

    async fn save_graph_with(
        &self,
        name: &str,
        edges: &[GraphEdge],
        options: &GraphWriteOptions,
        md_path: &Path,
    ) -> StorageResult<()> {
        self.validate_initialized(md_path)?;
        let (batch, num_nodes) = graph_record_batch(edges, options)?;
        let num_nodes_usize = usize::try_from(num_nodes)
            .map_err(|_| StorageError::Overflow(format!("node count {num_nodes} exceeds usize")))?;
        let weighted = edges[0].weight.is_some();
        let width = options.node_id_width;

        let path = self.file_path(name);
        info!(
            "Saving graph collection '{}' ({} edges, {} nodes, width {:?}, weight width {:?}, weighted {}) at {:?}",
            name,
            edges.len(),
            num_nodes,
            width,
            options.weight_type,
            weighted,
            path
        );
        let uri = Self::path_to_uri(&path)?;
        // Artifact first, registry publish second: the registry update
        // (through the commit actor) is the single commit point, so a
        // live entry only exists once the artifact is durable at its
        // final location (#106 review). A failed artifact write leaves
        // no registry entry; a failed publish leaves the artifact as
        // unreferenced residue that a later sweep reclaims — never a
        // live entry pointing at a missing artifact.
        self.write_lance_batch_async(uri, batch).await?;
        crate::commit::with_commit_actor(&self.metadata_path(), || async {
            let mut md = self.load_metadata().await?;
            let mut info = FileInfo::new(
                crate::generations::artifact_file_name(self.get_name().as_str(), name),
                "graph",
                (num_nodes_usize, num_nodes_usize),
                Some(edges.len()),
                None,
            )?;
            info.properties = options.properties.clone();
            info.properties
                .insert("node_id_width".to_string(), width.as_str().to_string());
            info.properties
                .insert("weighted".to_string(), weighted.to_string());
            info.properties
                .insert("num_nodes".to_string(), num_nodes.to_string());
            info.properties.insert(
                "weight_type".to_string(),
                options.weight_type.as_str().to_string(),
            );
            md = md.add_file(name, info);
            self.save_metadata(&md).await
        })
        .await?;
        info!("Graph collection '{}' saved successfully", name);
        Ok(())
    }

    async fn load_graph(&self, name: &str) -> StorageResult<StoredGraph> {
        info!("Loading graph collection '{}'", name);
        self.load_graph_from_path(&self.file_path(name)).await
    }

    // =========
    // REGISTRY-FREE COLLECTION I/O (#106)
    // =========

    async fn save_vectors_to_path(
        &self,
        path: &Path,
        batch: &RecordBatch,
        properties: &BTreeMap<String, String>,
    ) -> StorageResult<()> {
        let (batch, dim) = vector_space_collection_batch(batch, properties)?;
        ensure_parent_dir(path).await?;
        info!(
            "Saving vector-space collection ({} rows, vector dim {}) at {:?} (registry-free)",
            batch.num_rows(),
            dim,
            path
        );
        let uri = Self::path_to_uri(path)?;
        self.write_lance_batch_async(uri, batch).await?;
        info!("Vector-space collection saved successfully (registry-free)");
        Ok(())
    }

    async fn load_vectors_from_path(&self, path: &Path) -> StorageResult<RecordBatch> {
        info!(
            "Loading vector-space collection from {:?} (registry-free)",
            path
        );
        let uri = Self::path_to_uri(path)?;
        let batch = self.read_lance_all_batches_async(uri).await?;
        validate_vector_space_schema(batch.schema().as_ref())?;
        info!(
            "Loaded vector-space collection: {} rows (registry-free)",
            batch.num_rows()
        );
        Ok(batch)
    }

    async fn save_graph_to_path(
        &self,
        path: &Path,
        edges: &[GraphEdge],
        options: &GraphWriteOptions,
    ) -> StorageResult<()> {
        let (batch, num_nodes) = graph_record_batch(edges, options)?;
        ensure_parent_dir(path).await?;
        info!(
            "Saving graph collection ({} edges, {} nodes, width {:?}, weight width {:?}, weighted {}) at {:?} (registry-free)",
            edges.len(),
            num_nodes,
            options.node_id_width,
            options.weight_type,
            edges.first().map(|e| e.weight.is_some()).unwrap_or(false),
            path
        );
        let uri = Self::path_to_uri(path)?;
        self.write_lance_batch_async(uri, batch).await?;
        info!("Graph collection saved successfully (registry-free)");
        Ok(())
    }

    async fn load_graph_from_path(&self, path: &Path) -> StorageResult<StoredGraph> {
        self.load_graph_from_path_with_options(path, &GraphReadOptions::default())
            .await
    }

    async fn load_graph_from_path_with_options(
        &self,
        path: &Path,
        options: &GraphReadOptions,
    ) -> StorageResult<StoredGraph> {
        info!(
            "Loading graph collection from {:?} (registry-free, strict={})",
            path, options.strict
        );
        let uri = Self::path_to_uri(path)?;
        let batch = self.read_lance_all_batches_async(uri).await?;
        let graph = stored_graph_from_batch_with_options(batch, options)?;
        info!(
            "Loaded graph collection: {} edges, {} nodes, width {:?}, weight width {:?} (registry-free)",
            graph.edges.len(),
            graph.num_nodes,
            graph.node_id_width,
            graph.weight_type
        );
        Ok(graph)
    }

    async fn load_scalars(&self, name: &str) -> StorageResult<Vec<f64>> {
        self.load_scalars_from_path(&self.file_path(name)).await
    }

    async fn save_scalars_to_path(&self, path: &Path, values: &[f64]) -> StorageResult<()> {
        if values.is_empty() {
            return Err(StorageError::Invalid(
                "empty scalar collections are not supported".into(),
            ));
        }
        ensure_parent_dir(path).await?;
        info!(
            "Saving scalar collection ({} values) at {:?} (registry-free)",
            values.len(),
            path
        );
        // Dataset-level kind (RFC #81-P1): scalar artifacts are
        // vector-space collections, mirroring the registry-level
        // `CollectionKind::for_filetype("vector")` shim, so logical-kind
        // readers (`load_scalars`) can enforce the kind (#106 review). The
        // column name is the neutral default `lambda`.
        let schema =
            Schema::new(vec![Field::new("lambda", DataType::Float64, false)]).with_metadata(
                std::collections::HashMap::from([("kind".to_string(), "vector-space".to_string())]),
            );
        let batch = RecordBatch::try_new(
            Arc::new(schema),
            vec![Arc::new(Float64Array::from(values.to_vec())) as ArrayRef],
        )
        .map_err(|e| StorageError::Lance(e.to_string()))?;
        let uri = Self::path_to_uri(path)?;
        self.write_lance_batch_async(uri, batch).await?;
        info!("Scalar collection saved successfully (registry-free)");
        Ok(())
    }

    async fn load_scalars_from_path(&self, path: &Path) -> StorageResult<Vec<f64>> {
        info!("Loading scalar collection from {:?} (registry-free)", path);
        let uri = Self::path_to_uri(path)?;
        let batch = self.read_lance_all_batches_async(uri).await?;

        let schema = batch.schema();
        // Logical kind gate (#106 review): physical shape determines
        // decodability, but the stamped dataset kind determines whether
        // the collection is semantically a scalar vector-space
        // collection. Missing or mismatched kinds are rejected.
        match schema.metadata().get("kind").map(String::as_str) {
            Some("vector-space") => {}
            other => {
                return Err(StorageError::Invalid(format!(
                    "scalar collection at {path:?} has dataset kind {other:?}, \
                     expected 'vector-space'"
                )));
            }
        }
        if schema.fields().len() != 1 {
            return Err(StorageError::Invalid(format!(
                "scalar collection expects exactly one column, found {}",
                schema.fields().len()
            )));
        }
        let field = schema.field(0);
        if field.is_nullable() {
            return Err(StorageError::Invalid(format!(
                "scalar collection column '{}' is nullable",
                field.name()
            )));
        }
        let column = batch.column(0);
        if column.null_count() != 0 {
            return Err(StorageError::Invalid(
                "scalar collection column contains nulls".into(),
            ));
        }
        let values: Vec<f64> = match field.data_type() {
            DataType::Float64 => {
                let a = column
                    .as_any()
                    .downcast_ref::<Float64Array>()
                    .ok_or_else(|| StorageError::Invalid("scalar column type mismatch".into()))?;
                (0..a.len()).map(|i| a.value(i)).collect()
            }
            DataType::Float32 => {
                // lossless upcast at the API boundary
                let a = column
                    .as_any()
                    .downcast_ref::<Float32Array>()
                    .ok_or_else(|| StorageError::Invalid("scalar column type mismatch".into()))?;
                (0..a.len()).map(|i| f64::from(a.value(i))).collect()
            }
            other => {
                return Err(StorageError::Invalid(format!(
                    "scalar collection column must be Float64|Float32, found {other:?}"
                )));
            }
        };
        info!("Loaded {} scalar values (registry-free)", values.len());
        Ok(values)
    }

    async fn collection_schema_from_path(&self, path: &Path) -> StorageResult<Schema> {
        info!("Reading collection schema from {:?} (registry-free)", path);
        let path = path.to_path_buf();
        tokio::task::spawn_blocking(move || crate::lancefmt::read_schema(&path))
            .await
            .map_err(|e| StorageError::Io(format!("lancefmt read_schema task failed: {e}")))?
    }
}

impl crate::traits::clustering::ClusteringArtifacts for LanceStorageGraph {
    /// Save centroid_map (item-to-centroid assignments)
    async fn save_centroid_map(&self, map: &[usize], md_path: &Path) -> StorageResult<()> {
        info!("Saving {} centroid map entries", map.len());
        // Checked cast: usize values above u32::MAX must not be silently
        // truncated (issue #51).
        let values = checked_u32_values(map, "centroid map")?;
        self.save_primitive_column::<UInt32Type>("centroid_map", "centroid_id", values, md_path)
            .await
    }

    /// Load centroid_map
    async fn load_centroid_map(&self) -> StorageResult<Vec<usize>> {
        let path = self.file_path("centroid_map");
        info!("Loading centroid map from {:?}", path);

        let uri = Self::path_to_uri(&path)?;
        let batch = self.read_lance_all_batches_async(uri).await?;
        let arr = batch
            .column(0)
            .as_any()
            .downcast_ref::<UInt32Array>()
            .ok_or_else(|| StorageError::Invalid("centroid_id column type mismatch".into()))?;

        let map: Vec<usize> = (0..arr.len()).map(|i| arr.value(i) as usize).collect();
        info!("Loaded {} centroid map entries", map.len());
        Ok(map)
    }

    /// Save subcentroid_lambdas (tau values for subcentroids)
    async fn save_subcentroid_lambdas(&self, lambdas: &[f64], md_path: &Path) -> StorageResult<()> {
        info!("Saving {} subcentroid lambda values", lambdas.len());
        self.save_primitive_column::<Float64Type>(
            "subcentroid_lambdas",
            "subcentroid_lambda",
            lambdas.to_vec(),
            md_path,
        )
        .await
    }

    /// Load subcentroid_lambdas
    async fn load_subcentroid_lambdas(&self) -> StorageResult<Vec<f64>> {
        let path = self.file_path("subcentroid_lambdas");
        info!("Loading subcentroid lambda values from {:?}", path);

        let uri = Self::path_to_uri(&path)?;
        let batch = self.read_lance_all_batches_async(uri).await?;
        let arr = batch
            .column(0)
            .as_any()
            .downcast_ref::<Float64Array>()
            .ok_or_else(|| {
                StorageError::Invalid("subcentroid_lambda column type mismatch".into())
            })?;

        let lambdas: Vec<f64> = (0..arr.len()).map(|i| arr.value(i)).collect();
        info!("Loaded {} subcentroid lambda values", lambdas.len());
        Ok(lambdas)
    }

    /// Save subcentroids (dense matrix)
    async fn save_subcentroids(
        &self,
        subcentroids: &DenseMatrix<f64>,
        md_path: &Path,
    ) -> StorageResult<()> {
        self.validate_initialized(md_path)?;
        let key = "sub_centroids";
        let path = self.file_path(key);
        let (n_rows, n_cols) = subcentroids.shape();
        info!(
            "Saving subcentroids matrix {} x {} at {:?}",
            n_rows, n_cols, path
        );

        let batch = self.to_dense_record_batch(subcentroids)?;
        {
            crate::commit::with_commit_actor(&self.metadata_path(), || async {
                let mut metadata = self.load_metadata().await?;
                metadata = metadata.add_file(
                    key,
                    FileInfo::new(
                        crate::generations::artifact_file_name(self.get_name().as_str(), key),
                        "vector",
                        subcentroids.shape(),
                        None,
                        None,
                    )?,
                );
                self.save_metadata(&metadata).await
            })
            .await?;

            let uri = Self::path_to_uri(&path)?;
            self.write_lance_batch_async(uri, batch).await?;
        }
        debug!("Subcentroids matrix saved successfully");
        Ok(())
    }

    /// Load subcentroids as Vec<Vec<f64>>
    async fn load_subcentroids(&self) -> StorageResult<Vec<Vec<f64>>> {
        let path = self.file_path("sub_centroids");
        info!("Loading sub_centroids from {:?}", path);

        let uri = Self::path_to_uri(&path)?;
        let batch = self.read_lance_all_batches_async(uri).await?;
        let matrix = self.from_dense_record_batch(&batch)?;

        // Convert DenseMatrix to Vec<Vec<f64>>
        let (n_rows, n_cols) = matrix.shape();
        let mut result = Vec::with_capacity(n_rows);

        for row_idx in 0..n_rows {
            let row: Vec<f64> = (0..n_cols)
                .map(|col_idx| *matrix.get((row_idx, col_idx)))
                .collect();
            result.push(row);
        }

        info!(
            "Loaded sub_centroids: {} x {} as Vec<Vec<f64>>",
            n_rows, n_cols
        );
        Ok(result)
    }

    /// Save item norms vector
    async fn save_item_norms(&self, item_norms: &[f64], md_path: &Path) -> StorageResult<()> {
        info!("Saving {} item norm values", item_norms.len());
        self.save_primitive_column::<Float64Type>(
            "item_norms",
            "norm",
            item_norms.to_vec(),
            md_path,
        )
        .await
    }

    /// Load item norms vector
    async fn load_item_norms(&self) -> StorageResult<Vec<f64>> {
        let path = self.file_path("item_norms");
        info!("Loading item norms from {:?}", path);

        let uri = Self::path_to_uri(&path)?;
        let batch = self.read_lance_all_batches_async(uri).await?;
        let arr = batch
            .column(0)
            .as_any()
            .downcast_ref::<Float64Array>()
            .ok_or_else(|| StorageError::Invalid("norm column type mismatch".into()))?;

        let norms: Vec<f64> = (0..arr.len()).map(|i| arr.value(i)).collect();
        info!("Loaded {} item norm values", norms.len());
        Ok(norms)
    }

    async fn save_cluster_assignments(
        &self,
        assignments: &[Option<usize>],
        md_path: &Path,
    ) -> StorageResult<()> {
        info!("Saving {} cluster assignments", assignments.len());

        // Convert Option<usize> to i64 (-1 for None)
        let values: Vec<i64> = assignments
            .iter()
            .map(|opt| opt.map(|v| v as i64).unwrap_or(-1))
            .collect();

        self.save_primitive_column::<Int64Type>(
            "cluster_assignments",
            "cluster_id",
            values,
            md_path,
        )
        .await
    }

    async fn load_cluster_assignments(&self) -> StorageResult<Vec<Option<usize>>> {
        use arrow::array::Int64Array;
        let path = self.file_path("cluster_assignments");
        info!("Loading cluster assignments from {:?}", path);

        let uri = Self::path_to_uri(&path)?;
        let batch = self.read_lance_all_batches_async(uri).await?;
        let arr = batch
            .column(0)
            .as_any()
            .downcast_ref::<Int64Array>()
            .ok_or_else(|| StorageError::Invalid("cluster_id column type mismatch".into()))?;

        let assignments: Vec<Option<usize>> = (0..arr.len())
            .map(|i| {
                let v = arr.value(i);
                if v < 0 { None } else { Some(v as usize) }
            })
            .collect();
        info!("Loaded {} cluster assignments", assignments.len());
        Ok(assignments)
    }
}