diff --git a/Cargo.toml b/Cargo.toml index 6661823..34f1da5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,9 +11,11 @@ b-table = { path = "../b-table", features = ["stream"] } destream = "0.10" freqfs = "0.12" ha-ndarray = "0.5" +futures = "0.3" rand = "0.8" safecast = "0.2" smallvec = "1" [dev-dependencies] tokio = { version = "1", features = ["macros", "rt-multi-thread", "fs"] } +tbon = "0.9" diff --git a/README.md b/README.md index d428912..c755072 100644 --- a/README.md +++ b/README.md @@ -8,3 +8,10 @@ A filesystem-backed `Tensor` data structure featuring support for dense and spar - `fensor` does not attempt to repair, recover, or auto-heal corrupted metadata or tensor data. - If metadata or data is malformed, inconsistent, or unreadable, operations must return a structured error with a clear message. - Recovery workflows (restore/rebuild/migration) are external operational concerns, not `fensor` runtime behavior. + +## Serializing a tensor + +`Tensor` has two, independent wire-format surfaces: + +- **Schema-only (`Tensor: ToStream`/`FromStream`/`IntoStream`)**: encodes just the `TensorSchema` (dtype/shape/layout). Only a base (identity-view) tensor can be encoded this way; encoding a transformed (sliced/transposed/reshaped) view is rejected, since views are metadata-only and never persisted. Decoding always builds a fresh, empty base tensor at the given directory via `Tensor::create` — no element data is carried. +- **View + data streaming (`Tensor::view_encoder` / `TensorViewDecoder`)**: `tensor.view_encoder()` returns a `TensorViewEncoder<'_, FE, T>` that streams the tensor's *current* view — identity or transformed, dense or sparse — directly to the wire via `destream`'s `ToStream`/`IntoStream` contract. The encoder lazily reads from the tensor's filesystem-backed storage and emits values one at a time, with no full in-memory buffering; only non-default (nonzero) values are transmitted, reducing network traffic for sparse-heavy or mostly-empty tensors. On the receiving end, `TensorViewDecoder` implements `destream`'s `FromStream` and writes each arriving value directly to a fresh, independent, identity base tensor's filesystem storage as it arrives off the wire — also with no full in-memory buffering. There is no trailer or checksum on this wire format: a successful transfer is signaled by natural exhaustion of the pairs sequence, and end-to-end transfer completeness/integrity is left to the transport/caller layer rather than re-implemented here. A read failure on the sending side propagates as a bounded error-code sentinel — a closed classification of which failure shape occurred, deliberately excluding free-text detail (filesystem paths, coordinates, or other sender-machine specifics) since this wire format is meant to cross a machine boundary. On that sentinel, a malformed/truncated stream, or any other decode-time failure after the destination storage is created, the directory is truncated and deleted before a fail-closed error is returned. Call `.into_inner()` on the decoder to extract the reconstructed tensor. Like the design it replaces, this produces a fresh, independent identity base tensor with no link back to the source storage; it remains a distinct, additive wire surface alongside the existing schema-only `Tensor: ToStream/FromStream/IntoStream` contract, which is completely unrelated and still only carries dtype/shape/layout metadata without element data. diff --git a/ROADMAP.md b/ROADMAP.md index 93fcf3d..05791e2 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -87,6 +87,7 @@ Transactional orchestration belongs to `tc-collection`, which composes `fensor` - Validate metadata and block persistence across real `freqfs` reload boundaries. - Corruption and malformed metadata must fail closed with structured errors. - Gate: integration tests cover create->write->reload->read and corruption->error flows. + - Implemented: `Tensor::view_encoder()` and `TensorViewDecoder` provide a genuinely streaming serialization for a tensor's current view (identity or transformed, dense or sparse), with no full in-memory buffering on either side. The encoder lazily reads values from filesystem storage and emits only non-default payloads to the wire; the decoder writes arriving values directly to fresh independent base tensor storage. There is no trailer or checksum: success is natural exhaustion of the pairs sequence, and end-to-end transfer completeness/integrity is a transport/caller concern rather than something this wire format re-verifies itself. A sender-side read failure propagates as a bounded error-code sentinel (a closed classification of the failure shape, with no free-text/sender-internal detail such as filesystem paths); that sentinel, a malformed/truncated stream, or any other decode-time error triggers fail-closed cleanup (directory truncation and deletion). Reconstruction always produces a fresh, independent identity base tensor with no link back to the source storage. This is a distinct, additive wire surface alongside the existing schema-only `Tensor: ToStream/FromStream` contract, not a replacement for it; it does not by itself satisfy the "Phase 5: Conversion and materialization" `Dense <-> Sparse` conversion exit criteria below, since it always preserves the source's layout category. 4. **Base/view write-through semantics.** - Define which transformed tensors are writable and which are read-only. @@ -155,6 +156,7 @@ Transactional orchestration belongs to `tc-collection`, which composes `fensor` - **Adaptive block sizing.** Evaluate configurable or data-driven block sizing after baseline persistence semantics are stable. - **Typed tensor families.** Expand beyond `f32` once core lifecycle and sparse-index behavior are validated. - **Cross-host sharding hooks.** Keep routing/sharding orchestration in client libraries while exposing reusable shard-local primitives in `fensor`. +- **Encode-side enumeration optimization.** Encode-side view streaming currently walks every logical coordinate and filters defaults at the codec layer. A future optimization could use sparse-index row enumeration and filesystem directory listing to skip known-empty storage regions entirely, reducing I/O cost for extremely sparse large tensors — but this requires solving how to invert arbitrary view transforms (some non-closed-form) back to base storage coordinates, which is out of scope for now. ## ha-ndarray execution dependency acknowledgement diff --git a/src/error.rs b/src/error.rs index fcf1867..6f6c321 100644 --- a/src/error.rs +++ b/src/error.rs @@ -15,6 +15,7 @@ pub enum Error { hint: String, }, Unsupported(String), + DataMismatch(String), } pub type Result = std::result::Result; @@ -40,7 +41,8 @@ impl fmt::Display for Error { | Self::InvalidCoord(cause) | Self::InvalidLayout(cause) | Self::SparseIndex(cause) - | Self::Unsupported(cause) => f.write_str(cause), + | Self::Unsupported(cause) + | Self::DataMismatch(cause) => f.write_str(cause), Self::UnsupportedSparseIterationOrder { requested_order, base_order, diff --git a/src/lib.rs b/src/lib.rs index 58fcb41..d471b24 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,9 +1,9 @@ -use std::io; use std::sync::Arc; use b_table::{TableLock, collate::Collator}; use destream::{de, en}; use freqfs::{DirLock, FileLoad}; +use futures::StreamExt as _; use ha_ndarray::{Axes, Range, Shape}; use safecast::AsType; @@ -17,14 +17,15 @@ mod wire_tags; pub use error::{Error, Result}; pub use schema::{ - AxisContribSchema, DType, Layout, SparseIndexSchema, SparseTableSchema, TensorSchema, - TensorShape, ViewSchema, contiguous_strides, + DType, Layout, SparseIndexSchema, SparseTableSchema, TensorSchema, TensorShape, + contiguous_strides, }; +pub use stream::{TensorViewDecoder, TensorViewEncoder}; pub use traits::{ - BoxFuture, SparseZeroPolicy, TensorArray, TensorBlockStore, TensorMatMul, TensorMath, - TensorMathScalar, TensorRead, TensorReadBulk, TensorReduce, TensorReduceAll, - TensorReduceBoolean, TensorSparseIndex, TensorSparseLifecycle, TensorTransform, TensorUnary, - TensorViewSemantics, TensorWrite, TensorWriteBulk, + BoxFuture, TensorArray, TensorBlockStore, TensorMatMul, TensorMath, TensorMathScalar, + TensorRead, TensorReadBulk, TensorReduce, TensorReduceAll, TensorReduceBoolean, + TensorSparseIndex, TensorTransform, TensorUnary, TensorViewSemantics, TensorWrite, + TensorWriteBulk, }; use view::{TensorView, default_permutation}; @@ -32,7 +33,7 @@ use view::{TensorView, default_permutation}; const BLOCKS: &str = "blocks"; const INDEX: &str = "index"; const METADATA: &str = "metadata"; -const METADATA_VERSION: u32 = 1; +const METADATA_VERSION: u32 = 2; pub trait TensorElement: Copy @@ -43,6 +44,7 @@ pub trait TensorElement: + 'static + de::FromStream + for<'en> en::ToStream<'en> + + for<'en> en::IntoStream<'en> { const DTYPE: DType; } @@ -67,24 +69,34 @@ where { } -struct TensorStorage { +// --------------------------------------------------------------------------- +// Private storage structs +// --------------------------------------------------------------------------- + +struct DenseStorage { + blocks: DirLock, + schema: TensorSchema, +} + +struct SparseStorage { blocks: DirLock, - index: Option, FE>>, + index: TableLock, FE>, schema: TensorSchema, } +// --------------------------------------------------------------------------- +// DenseTensor +// --------------------------------------------------------------------------- + #[derive(Clone)] -pub struct Tensor { - storage: Arc>, +pub struct DenseTensor { + storage: Arc>, schema: TensorSchema, view: TensorView, _dtype: std::marker::PhantomData, } -pub type TensorF32 = Tensor; -pub type TensorF64 = Tensor; - -impl Tensor +impl DenseTensor where FE: TensorFileEntry, T: TensorElement, @@ -96,10 +108,162 @@ where validate_tensor_dtype::(&schema)?; let mut dir_guard = dir.try_write()?; - let blocks_dir = dir_guard.create_dir(BLOCKS.to_string())?; - let index = create_sparse_index(&schema, &mut dir_guard)?; + let tensor = Self::new_storage(blocks_dir, schema); + tensor.persist_metadata().await?; + Ok(tensor) + } + + pub async fn load(dir: DirLock) -> Result + where + FE: AsType + From, + { + let mut dir_guard = dir.try_write()?; + let blocks_dir = dir_guard.get_or_create_dir(BLOCKS.to_string())?; + let schema = load_metadata_inner(&blocks_dir).await?; + validate_tensor_dtype::(&schema)?; + Ok(Self::new_storage(blocks_dir, schema)) + } + pub async fn load_with_schema(dir: DirLock, expected: &TensorSchema) -> Result + where + FE: AsType + From, + { + validate_tensor_dtype::(expected)?; + let tensor = Self::load(dir).await?; + if tensor.schema() != expected { + return Err(Error::InvalidSchema(format!( + "persisted metadata mismatch: expected {:?}, found {:?}", + expected, + tensor.schema() + ))); + } + Ok(tensor) + } + + fn new_storage(blocks: DirLock, schema: TensorSchema) -> Self { + let view = TensorView::identity(&schema); + let storage = Arc::new(DenseStorage { + blocks, + schema: schema.clone(), + }); + Self { + storage, + schema, + view, + _dtype: std::marker::PhantomData, + } + } + + pub(crate) fn block_len(&self) -> usize { + self.storage.schema.block_len().max(1) + } + + pub(crate) fn resolve_base_coord(&self, coord: &[u64]) -> Result> { + self.schema.validate_coord(coord)?; + let k = self.view.flat_offset(coord)?; + if k < 0 { + return Err(Error::InvalidCoord("negative linear offset".to_string())); + } + let k = k as u64; + let base_coord: Vec = self + .storage + .schema + .strides() + .iter() + .zip(self.storage.schema.shape().iter()) + .map(|(stride, dim)| (k / *stride as u64) % *dim as u64) + .collect(); + self.storage.schema.validate_coord(&base_coord)?; + Ok(base_coord) + } + + pub(crate) fn block_position_from_base_coord(&self, base_coord: &[u64]) -> (u64, usize) { + let offset: u64 = base_coord + .iter() + .zip(self.storage.schema.strides().iter()) + .map(|(coord, stride)| *coord * (*stride as u64)) + .sum(); + let block_len = self.block_len() as u64; + let block_offset = offset / block_len; + let offset_in_block = (offset % block_len) as usize; + (block_offset, offset_in_block) + } + + fn default_block(&self) -> Vec { + vec![T::default(); self.block_len()] + } + + async fn write_value_to_block( + &self, + block_id: u64, + offset_in_block: usize, + value: T, + create_if_missing: bool, + ) -> Result<()> { + let mut block = match self.read_block(block_id).await? { + Some(block) => block, + None if create_if_missing => self.default_block(), + None => { + return Err(Error::SparseIndex( + "sparse index points to missing block".to_string(), + )); + } + }; + validate::ensure_offset_in_bounds(offset_in_block, block.len())?; + block[offset_in_block] = value; + self.write_block(block_id, block).await + } + + async fn persist_metadata(&self) -> Result<()> + where + FE: AsType + From, + { + let payload = encode_schema(&self.storage.schema); + let _ = decode_schema(&payload)?; + write_metadata_file(&self.storage.blocks, METADATA, &payload).await + } +} + +// --------------------------------------------------------------------------- +// SparseTensor +// --------------------------------------------------------------------------- + +enum SparseWriteAction { + Write(u64), + DeleteRow(u64), + NoOp, +} + +#[derive(Clone)] +pub struct SparseTensor { + storage: Arc>, + schema: TensorSchema, + view: TensorView, + _dtype: std::marker::PhantomData, +} + +impl SparseTensor +where + FE: TensorFileEntry, + T: TensorElement, +{ + pub async fn create(dir: DirLock, schema: TensorSchema) -> Result + where + FE: AsType + From, + { + validate_tensor_dtype::(&schema)?; + if !matches!(schema.layout(), Layout::Sparse { .. }) { + return Err(Error::InvalidLayout( + "SparseTensor requires a sparse schema layout".to_string(), + )); + } + + let mut dir_guard = dir.try_write()?; + let blocks_dir = dir_guard.create_dir(BLOCKS.to_string())?; + let index_dir = dir_guard.create_dir(INDEX.to_string())?; + let index = + TableLock::create(SparseTableSchema::default(), Collator::default(), index_dir)?; let tensor = Self::new_storage(blocks_dir, index, schema); tensor.persist_metadata().await?; Ok(tensor) @@ -110,12 +274,13 @@ where FE: AsType + From, { let mut dir_guard = dir.try_write()?; - let blocks_dir = dir_guard.get_or_create_dir(BLOCKS.to_string())?; - let schema = Self::load_metadata(&blocks_dir).await?; + let schema = load_metadata_inner(&blocks_dir).await?; validate_tensor_dtype::(&schema)?; - let index = load_sparse_index(&schema, &mut dir_guard)?; - + let index_dir = dir_guard.get_dir(INDEX).cloned().ok_or_else(|| { + Error::InvalidSchema("sparse tensor missing index directory".to_string()) + })?; + let index = TableLock::load(SparseTableSchema::default(), Collator::default(), index_dir)?; Ok(Self::new_storage(blocks_dir, index, schema)) } @@ -124,31 +289,28 @@ where FE: AsType + From, { validate_tensor_dtype::(expected)?; - let tensor = Self::load(dir).await?; - - if &tensor.schema != expected { + if tensor.schema() != expected { return Err(Error::InvalidSchema(format!( "persisted metadata mismatch: expected {:?}, found {:?}", - expected, tensor.schema + expected, + tensor.schema() ))); } - Ok(tensor) } fn new_storage( blocks: DirLock, - index: Option, FE>>, + index: TableLock, FE>, schema: TensorSchema, ) -> Self { let view = TensorView::identity(&schema); - let storage = Arc::new(TensorStorage { + let storage = Arc::new(SparseStorage { blocks, index, schema: schema.clone(), }); - Self { storage, schema, @@ -157,35 +319,17 @@ where } } - pub fn view_schema(&self) -> Result { - self.view.to_schema() - } - - pub fn with_view_schema(mut self, view_schema: &ViewSchema) -> Result { - let view = TensorView::from_schema(view_schema)?; - if view.rank() != self.schema.rank() { - return Err(Error::InvalidSchema( - "view rank must match tensor rank".to_string(), - )); - } - - self.view = view; - Ok(self) - } - pub(crate) fn block_len(&self) -> usize { self.storage.schema.block_len().max(1) } pub(crate) fn resolve_base_coord(&self, coord: &[u64]) -> Result> { self.schema.validate_coord(coord)?; - let k = self.view.flat_offset(coord)?; if k < 0 { return Err(Error::InvalidCoord("negative linear offset".to_string())); } let k = k as u64; - let base_coord: Vec = self .storage .schema @@ -194,7 +338,6 @@ where .zip(self.storage.schema.shape().iter()) .map(|(stride, dim)| (k / *stride as u64) % *dim as u64) .collect(); - self.storage.schema.validate_coord(&base_coord)?; Ok(base_coord) } @@ -205,11 +348,9 @@ where .zip(self.storage.schema.strides().iter()) .map(|(coord, stride)| *coord * (*stride as u64)) .sum(); - let block_len = self.block_len() as u64; let block_offset = offset / block_len; let offset_in_block = (offset % block_len) as usize; - (block_offset, offset_in_block) } @@ -218,15 +359,10 @@ where Layout::Sparse { axis } => axis.unwrap_or(0), Layout::Dense => 0, }; - let axis = sparse_axis.min(coords.len().saturating_sub(1)); vec![coords[axis], block_offset] } - pub(crate) fn has_sparse_index(&self) -> bool { - self.storage.index.is_some() - } - async fn lookup_sparse_block_for_coord( &self, base_coord: &[u64], @@ -245,25 +381,27 @@ where base_coord: &[u64], block_offset: u64, value: T, - ) -> Result> { + ) -> Result { let key = self.sparse_key(base_coord, block_offset); if let Some(block_id) = self .lookup_sparse_block_for_coord(base_coord, block_offset) .await? { - return Ok(Some(block_id)); + if value == T::default() { + return Ok(SparseWriteAction::DeleteRow(block_id)); + } + return Ok(SparseWriteAction::Write(block_id)); } if value == T::default() { - return Ok(None); + return Ok(SparseWriteAction::NoOp); } let block_id: u64 = rand::random(); self.write_block(block_id, self.default_block()).await?; self.upsert_block_id(key, block_id).await?; - - Ok(Some(block_id)) + Ok(SparseWriteAction::Write(block_id)) } async fn write_value_to_block( @@ -282,67 +420,153 @@ where )); } }; - validate::ensure_offset_in_bounds(offset_in_block, block.len())?; block[offset_in_block] = value; self.write_block(block_id, block).await } + async fn delete_block(&self, block_id: u64) { + let mut blocks = self.storage.blocks.write().await; + blocks.delete(&block_id.to_string()).await; + } + async fn persist_metadata(&self) -> Result<()> where FE: AsType + From, { let payload = encode_schema(&self.storage.schema); - - // Validate metadata payload before writing it. let _ = decode_schema(&payload)?; + write_metadata_file(&self.storage.blocks, METADATA, &payload).await + } + + pub async fn compact_sparse(&self) -> Result<()> { + let all_rows = { + let guard = self.storage.index.read().await; + let mut rows = guard.into_rows().await.map_err(Error::from)?; + let mut collected: Vec> = Vec::new(); + while let Some(row) = rows.next().await { + let row = row.map_err(Error::from)?; + collected.push(row.to_vec()); + } + collected + }; - self.write_metadata_file(METADATA, &payload).await?; + let mut to_delete: Vec<(Vec, u64)> = Vec::new(); + for row in &all_rows { + let key = vec![row[0], row[1]]; + let block_id = row[2]; + let all_zero = match self.read_block(block_id).await? { + Some(block) => block.iter().all(|v| *v == T::default()), + None => true, + }; + if all_zero { + to_delete.push((key, block_id)); + } + } + + for (key, block_id) in to_delete { + self.delete_row(key).await?; + self.delete_block(block_id).await; + } Ok(()) } +} + +// --------------------------------------------------------------------------- +// Tensor enum +// --------------------------------------------------------------------------- + +#[derive(Clone)] +pub enum Tensor { + Dense(DenseTensor), + Sparse(SparseTensor), +} + +pub type TensorF32 = Tensor; +pub type TensorF64 = Tensor; +pub type DenseTensorF32 = DenseTensor; +pub type DenseTensorF64 = DenseTensor; +pub type SparseTensorF32 = SparseTensor; +pub type SparseTensorF64 = SparseTensor; - async fn load_metadata(blocks: &DirLock) -> Result +impl Tensor +where + FE: TensorFileEntry, + T: TensorElement, +{ + pub async fn create(dir: DirLock, schema: TensorSchema) -> Result where FE: AsType + From, { - let file = { - let dir = blocks.read().await; - dir.get_file(METADATA) - .cloned() - .ok_or_else(|| Error::InvalidSchema("missing tensor metadata file".to_string()))? - }; + match schema.layout() { + Layout::Dense => DenseTensor::create(dir, schema).await.map(Self::Dense), + Layout::Sparse { .. } => SparseTensor::create(dir, schema).await.map(Self::Sparse), + } + } - let payload = { - let guard = file.read::().await?; - guard.clone() + pub async fn load(dir: DirLock) -> Result + where + FE: AsType + From, + { + let layout = { + let mut dir_guard = dir.try_write()?; + let blocks_dir = dir_guard.get_or_create_dir(BLOCKS.to_string())?; + let schema = load_metadata_inner(&blocks_dir).await?; + schema.layout().clone() }; - - decode_schema(&payload) + match layout { + Layout::Dense => DenseTensor::load(dir).await.map(Self::Dense), + Layout::Sparse { .. } => SparseTensor::load(dir).await.map(Self::Sparse), + } } - async fn write_metadata_file(&self, name: &str, payload: &str) -> Result<()> + pub async fn load_with_schema(dir: DirLock, expected: &TensorSchema) -> Result where FE: AsType + From, { - let existing = { - let dir = self.storage.blocks.read().await; - dir.get_file(name).cloned() - }; + validate_tensor_dtype::(expected)?; + let tensor = Self::load(dir).await?; + if tensor.schema() != expected { + return Err(Error::InvalidSchema(format!( + "persisted metadata mismatch: expected {:?}, found {:?}", + expected, + tensor.schema() + ))); + } + Ok(tensor) + } - if let Some(file) = existing { - let mut guard = file.write::().await?; - *guard = payload.to_string(); - } else { - let mut dir = self.storage.blocks.write().await; - dir.create_file(name.to_string(), payload.to_string(), payload.len())?; + /// Build a lazily-streamed encoder for this tensor's current view + /// (identity or transformed, dense or sparse): schema followed by a + /// nested sequence of non-default `(coord, value)` pairs and a trailing + /// verification record. No full in-memory buffering -- each value is + /// read from storage only as the returned value is actually driven by + /// a destream encoder (e.g. `tbon::en::encode(tensor.view_encoder())`). + pub fn view_encoder(&self) -> stream::TensorViewEncoder<'_, FE, T> { + stream::TensorViewEncoder::new(self) + } + + pub fn as_sparse(&self) -> Option<&SparseTensor> { + match self { + Self::Sparse(inner) => Some(inner), + Self::Dense(_) => None, } + } - Ok(()) + pub fn into_sparse(self) -> Option> { + match self { + Self::Sparse(inner) => Some(inner), + Self::Dense(_) => None, + } } } -impl TensorArray for Tensor +// --------------------------------------------------------------------------- +// TensorArray impls +// --------------------------------------------------------------------------- + +impl TensorArray for DenseTensor where FE: TensorFileEntry, T: TensorElement, @@ -362,7 +586,54 @@ where } } -impl TensorRead for Tensor +impl TensorArray for SparseTensor +where + FE: TensorFileEntry, + T: TensorElement, +{ + type DType = T; + + fn schema(&self) -> &TensorSchema { + &self.schema + } + + fn dtype(&self) -> Self::DType { + T::default() + } + + fn schema_dtype(&self) -> DType { + self.schema.dtype() + } +} + +impl TensorArray for Tensor +where + FE: TensorFileEntry, + T: TensorElement, +{ + type DType = T; + + fn schema(&self) -> &TensorSchema { + match self { + Self::Dense(inner) => inner.schema(), + Self::Sparse(inner) => inner.schema(), + } + } + + fn dtype(&self) -> Self::DType { + T::default() + } + + fn schema_dtype(&self) -> DType { + self.schema().dtype() + } +} + +// --------------------------------------------------------------------------- +// TensorRead impls +// --------------------------------------------------------------------------- + +impl TensorRead for DenseTensor where FE: TensorFileEntry, T: TensorElement, @@ -372,60 +643,182 @@ where let base_coord = self.resolve_base_coord(coord)?; let (block_offset, offset_in_block) = self.block_position_from_base_coord(&base_coord); - let block_id = if self.has_sparse_index() { - self.lookup_sparse_block_for_coord(&base_coord, block_offset) - .await? + if let Some(block) = self.read_block(block_offset).await? { + validate::ensure_offset_in_bounds(offset_in_block, block.len())?; + Ok(block[offset_in_block]) } else { - Some(block_offset) - }; + Ok(T::default()) + } + }) + } +} + +impl TensorRead for SparseTensor +where + FE: TensorFileEntry, + T: TensorElement, +{ + fn read_value<'a>(&'a self, coord: &'a [u64]) -> BoxFuture<'a, Result> { + Box::pin(async move { + let base_coord = self.resolve_base_coord(coord)?; + let (block_offset, offset_in_block) = self.block_position_from_base_coord(&base_coord); - if let Some(block_id) = block_id { + if let Some(block_id) = self + .lookup_sparse_block_for_coord(&base_coord, block_offset) + .await? + { if let Some(block) = self.read_block(block_id).await? { validate::ensure_offset_in_bounds(offset_in_block, block.len())?; Ok(block[offset_in_block]) } else { - Ok(T::default()) + Err(Error::SparseIndex("Block is missing".to_string())) } } else { Ok(T::default()) } }) } -} +} + +impl TensorRead for Tensor +where + FE: TensorFileEntry, + T: TensorElement, +{ + fn read_value<'a>(&'a self, coord: &'a [u64]) -> BoxFuture<'a, Result> { + match self { + Self::Dense(inner) => inner.read_value(coord), + Self::Sparse(inner) => inner.read_value(coord), + } + } +} + +// --------------------------------------------------------------------------- +// TensorWrite impls +// --------------------------------------------------------------------------- + +impl TensorWrite for DenseTensor +where + FE: TensorFileEntry, + T: TensorElement, +{ + fn write_value<'a>( + &'a self, + coord: &'a [u64], + value: Self::DType, + ) -> BoxFuture<'a, Result<()>> { + Box::pin(async move { + let base_coord = self.resolve_base_coord(coord)?; + let (block_offset, offset_in_block) = self.block_position_from_base_coord(&base_coord); + self.write_value_to_block(block_offset, offset_in_block, value, true) + .await + }) + } +} + +impl TensorWrite for SparseTensor +where + FE: TensorFileEntry, + T: TensorElement, +{ + fn write_value<'a>( + &'a self, + coord: &'a [u64], + value: Self::DType, + ) -> BoxFuture<'a, Result<()>> { + Box::pin(async move { + let base_coord = self.resolve_base_coord(coord)?; + let (block_offset, offset_in_block) = self.block_position_from_base_coord(&base_coord); + + match self + .resolve_sparse_block_for_write(&base_coord, block_offset, value) + .await? + { + SparseWriteAction::Write(block_id) => { + self.write_value_to_block(block_id, offset_in_block, value, false) + .await + } + SparseWriteAction::DeleteRow(block_id) => { + let key = self.sparse_key(&base_coord, block_offset); + self.delete_row(key).await?; + self.delete_block(block_id).await; + Ok(()) + } + SparseWriteAction::NoOp => Ok(()), + } + }) + } +} + +impl TensorWrite for Tensor +where + FE: TensorFileEntry, + T: TensorElement, +{ + fn write_value<'a>( + &'a self, + coord: &'a [u64], + value: Self::DType, + ) -> BoxFuture<'a, Result<()>> { + match self { + Self::Dense(inner) => inner.write_value(coord, value), + Self::Sparse(inner) => inner.write_value(coord, value), + } + } +} + +// --------------------------------------------------------------------------- +// TensorTransform impls +// --------------------------------------------------------------------------- + +impl TensorTransform for DenseTensor +where + FE: TensorFileEntry, + T: TensorElement, +{ + fn reshape(mut self, shape: Shape) -> Result { + let old_shape = self.schema.shape(); + let old_size: usize = old_shape.iter().product(); + let new_size: usize = shape.iter().product(); + if old_size != new_size { + return Err(Error::InvalidLayout( + "reshape requires an equal number of elements".to_string(), + )); + } + self.view = self.view.reshape(self.schema.shape(), &shape)?; + self.schema.set_shape(shape)?; + self.schema + .set_strides(contiguous_strides(self.schema.shape()))?; + Ok(self) + } -impl TensorWrite for Tensor -where - FE: TensorFileEntry, - T: TensorElement, -{ - fn write_value<'a>( - &'a self, - coord: &'a [u64], - value: Self::DType, - ) -> BoxFuture<'a, Result<()>> { - Box::pin(async move { - let base_coord = self.resolve_base_coord(coord)?; - let (block_offset, offset_in_block) = self.block_position_from_base_coord(&base_coord); + fn slice(mut self, range: Range) -> Result { + let current_shape = self.schema.shape(); + let (view, shape, strides) = self.view.slice(current_shape, &range)?; + self.view = view; + self.schema.set_shape(shape)?; + self.schema.set_strides(strides)?; + Ok(self) + } - if self.has_sparse_index() { - if let Some(block_id) = self - .resolve_sparse_block_for_write(&base_coord, block_offset, value) - .await? - { - self.write_value_to_block(block_id, offset_in_block, value, false) - .await - } else { - Ok(()) - } - } else { - self.write_value_to_block(block_offset, offset_in_block, value, true) - .await - } - }) + fn transpose(mut self, permutation: Option) -> Result { + let current_shape = self.schema.shape(); + let current_strides = self.schema.strides().clone(); + let permutation = default_permutation(current_shape.len(), permutation)?; + self.view = self.view.transpose(&permutation)?; + let mut shape = Shape::with_capacity(current_shape.len()); + let mut strides = Vec::with_capacity(current_shape.len()); + for axis in permutation { + shape.push(current_shape[axis]); + strides.push(current_strides[axis]); + } + self.schema.set_shape(shape)?; + self.schema.set_strides(strides.into())?; + Ok(self) } } -impl TensorTransform for Tensor +impl TensorTransform for SparseTensor where FE: TensorFileEntry, T: TensorElement, @@ -434,18 +827,15 @@ where let old_shape = self.schema.shape(); let old_size: usize = old_shape.iter().product(); let new_size: usize = shape.iter().product(); - if old_size != new_size { return Err(Error::InvalidLayout( "reshape requires an equal number of elements".to_string(), )); } - self.view = self.view.reshape(self.schema.shape(), &shape)?; self.schema.set_shape(shape)?; self.schema .set_strides(contiguous_strides(self.schema.shape()))?; - Ok(self) } @@ -455,33 +845,58 @@ where self.view = view; self.schema.set_shape(shape)?; self.schema.set_strides(strides)?; - Ok(self) } fn transpose(mut self, permutation: Option) -> Result { let current_shape = self.schema.shape(); let current_strides = self.schema.strides().clone(); - let permutation = default_permutation(current_shape.len(), permutation)?; self.view = self.view.transpose(&permutation)?; - let mut shape = Shape::with_capacity(current_shape.len()); let mut strides = Vec::with_capacity(current_shape.len()); - for axis in permutation { shape.push(current_shape[axis]); strides.push(current_strides[axis]); } - self.schema.set_shape(shape)?; self.schema.set_strides(strides.into())?; - Ok(self) } } -impl TensorBlockStore for Tensor +impl TensorTransform for Tensor +where + FE: TensorFileEntry, + T: TensorElement, +{ + fn reshape(self, shape: Shape) -> Result { + match self { + Self::Dense(inner) => inner.reshape(shape).map(Self::Dense), + Self::Sparse(inner) => inner.reshape(shape).map(Self::Sparse), + } + } + + fn slice(self, range: Range) -> Result { + match self { + Self::Dense(inner) => inner.slice(range).map(Self::Dense), + Self::Sparse(inner) => inner.slice(range).map(Self::Sparse), + } + } + + fn transpose(self, permutation: Option) -> Result { + match self { + Self::Dense(inner) => inner.transpose(permutation).map(Self::Dense), + Self::Sparse(inner) => inner.transpose(permutation).map(Self::Sparse), + } + } +} + +// --------------------------------------------------------------------------- +// TensorBlockStore impls +// --------------------------------------------------------------------------- + +impl TensorBlockStore for DenseTensor where FE: TensorFileEntry, T: TensorElement, @@ -506,36 +921,111 @@ where let blocks = self.storage.blocks.read().await; blocks.get_file(&block_id.to_string()).cloned() }; - if let Some(file) = file { let mut guard = file.write::>().await?; *guard = block; - Ok(()) } else { let mut blocks = self.storage.blocks.write().await; blocks.create_file(block_id.to_string(), block, 0)?; + Ok(()) + } + }) + } +} + +impl TensorBlockStore for SparseTensor +where + FE: TensorFileEntry, + T: TensorElement, +{ + type Block = Vec; + + fn read_block<'a>(&'a self, block_id: u64) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let blocks = self.storage.blocks.read().await; + if let Some(file) = blocks.get_file(&block_id.to_string()) { + let guard = file.read::>().await?; + Ok(Some(guard.clone())) + } else { + Ok(None) + } + }) + } + fn write_block<'a>(&'a self, block_id: u64, block: Self::Block) -> BoxFuture<'a, Result<()>> { + Box::pin(async move { + let file = { + let blocks = self.storage.blocks.read().await; + blocks.get_file(&block_id.to_string()).cloned() + }; + if let Some(file) = file { + let mut guard = file.write::>().await?; + *guard = block; + Ok(()) + } else { + let mut blocks = self.storage.blocks.write().await; + blocks.create_file(block_id.to_string(), block, 0)?; Ok(()) } }) } } -impl TensorSparseIndex for Tensor +impl TensorBlockStore for Tensor +where + FE: TensorFileEntry, + T: TensorElement, +{ + type Block = Vec; + + fn read_block<'a>(&'a self, block_id: u64) -> BoxFuture<'a, Result>> { + match self { + Self::Dense(inner) => inner.read_block(block_id), + Self::Sparse(inner) => inner.read_block(block_id), + } + } + + fn write_block<'a>(&'a self, block_id: u64, block: Self::Block) -> BoxFuture<'a, Result<()>> { + match self { + Self::Dense(inner) => inner.write_block(block_id, block), + Self::Sparse(inner) => inner.write_block(block_id, block), + } + } +} + +// --------------------------------------------------------------------------- +// TensorSparseIndex impls +// --------------------------------------------------------------------------- + +impl TensorSparseIndex for DenseTensor +where + FE: TensorFileEntry, + T: TensorElement, +{ + fn lookup_block_id<'a>(&'a self, _key: &'a [u64]) -> BoxFuture<'a, Result>> { + Box::pin(async move { Ok(None) }) + } + + fn upsert_block_id<'a>(&'a self, _key: Vec, _block_id: u64) -> BoxFuture<'a, Result<()>> { + Box::pin(async move { + Err(Error::Unsupported( + "sparse index is not available for dense layout".to_string(), + )) + }) + } +} + +impl TensorSparseIndex for SparseTensor where FE: TensorFileEntry, T: TensorElement, { fn lookup_block_id<'a>(&'a self, key: &'a [u64]) -> BoxFuture<'a, Result>> { Box::pin(async move { - if let Some(index) = &self.storage.index { - let index_lock = index.read().await; - if let Some(row) = index_lock.get_row(key).await? { - Ok(row.get(2).copied()) - } else { - Ok(None) - } + let index_lock = self.storage.index.read().await; + if let Some(row) = index_lock.get_row(key).await? { + Ok(row.get(2).copied()) } else { Ok(None) } @@ -544,59 +1034,156 @@ where fn upsert_block_id<'a>(&'a self, key: Vec, block_id: u64) -> BoxFuture<'a, Result<()>> { Box::pin(async move { - if let Some(index) = &self.storage.index { - let mut index_lock = index.write().await; - index_lock - .upsert(key, vec![block_id]) - .await - .map(|_| ()) - .map_err(Error::from) - } else { - Err(Error::Unsupported( - "sparse index is not available for dense layout".to_string(), - )) - } + let mut index_lock = self.storage.index.write().await; + index_lock + .upsert(key, vec![block_id]) + .await + .map(|_| ()) + .map_err(Error::from) + }) + } + + fn delete_row<'a>(&'a self, key: Vec) -> BoxFuture<'a, Result> { + Box::pin(async move { + let mut index_lock = self.storage.index.write().await; + index_lock.delete_row(&key).await.map_err(Error::from) }) } } -// Trait-surface wiring only: every method inherits the default `Unsupported` -// behavior from `traits.rs`. These exist so the test matrix can call into the -// trait surface and observe failures in the right places. They are NOT a -// feature implementation — see Requirements_1.md. -impl TensorReadBulk for Tensor +impl TensorSparseIndex for Tensor +where + FE: TensorFileEntry, + T: TensorElement, +{ + fn lookup_block_id<'a>(&'a self, key: &'a [u64]) -> BoxFuture<'a, Result>> { + match self { + Self::Dense(inner) => inner.lookup_block_id(key), + Self::Sparse(inner) => inner.lookup_block_id(key), + } + } + + fn upsert_block_id<'a>(&'a self, key: Vec, block_id: u64) -> BoxFuture<'a, Result<()>> { + match self { + Self::Dense(inner) => inner.upsert_block_id(key, block_id), + Self::Sparse(inner) => inner.upsert_block_id(key, block_id), + } + } + + fn delete_row<'a>(&'a self, key: Vec) -> BoxFuture<'a, Result> { + match self { + Self::Dense(inner) => inner.delete_row(key), + Self::Sparse(inner) => inner.delete_row(key), + } + } +} + +// --------------------------------------------------------------------------- +// TensorViewSemantics impls +// --------------------------------------------------------------------------- + +impl TensorViewSemantics for DenseTensor where FE: TensorFileEntry, T: TensorElement, { + fn is_base_tensor(&self) -> bool { + self.view.is_identity(&self.storage.schema) + } + + fn supports_write_through(&self) -> bool { + !self.view.has_gather_axes() + } } -impl TensorWriteBulk for Tensor +impl TensorViewSemantics for SparseTensor where FE: TensorFileEntry, T: TensorElement, { + fn is_base_tensor(&self) -> bool { + self.view.is_identity(&self.storage.schema) + } + + fn supports_write_through(&self) -> bool { + !self.view.has_gather_axes() + } } impl TensorViewSemantics for Tensor +where + FE: TensorFileEntry, + T: TensorElement, +{ + fn is_base_tensor(&self) -> bool { + match self { + Self::Dense(inner) => inner.is_base_tensor(), + Self::Sparse(inner) => inner.is_base_tensor(), + } + } + + fn supports_write_through(&self) -> bool { + match self { + Self::Dense(inner) => inner.supports_write_through(), + Self::Sparse(inner) => inner.supports_write_through(), + } + } +} + +// --------------------------------------------------------------------------- +// Unsupported bulk traits (trait-surface wiring only) +// --------------------------------------------------------------------------- + +impl TensorReadBulk for DenseTensor +where + FE: TensorFileEntry, + T: TensorElement, +{ +} + +impl TensorReadBulk for SparseTensor where FE: TensorFileEntry, T: TensorElement, { } -impl TensorSparseLifecycle for Tensor +impl TensorReadBulk for Tensor +where + FE: TensorFileEntry, + T: TensorElement, +{ +} + +impl TensorWriteBulk for DenseTensor +where + FE: TensorFileEntry, + T: TensorElement, +{ +} + +impl TensorWriteBulk for SparseTensor +where + FE: TensorFileEntry, + T: TensorElement, +{ +} + +impl TensorWriteBulk for Tensor where FE: TensorFileEntry, T: TensorElement, { } -fn default_block_shape(shape: &Shape) -> Shape { +// --------------------------------------------------------------------------- +// Free functions +// --------------------------------------------------------------------------- + +pub fn default_block_shape(shape: &Shape) -> Shape { if shape.is_empty() { return Shape::new(); } - let mut block_shape = Shape::with_capacity(shape.len()); for (i, dim) in shape.iter().enumerate() { if i + 1 == shape.len() { @@ -605,22 +1192,54 @@ fn default_block_shape(shape: &Shape) -> Shape { block_shape.push(1); } } - block_shape } +async fn load_metadata_inner(blocks: &DirLock) -> Result +where + FE: AsType + FileLoad + Send + Sync + 'static, +{ + let file = { + let dir = blocks.read().await; + dir.get_file(METADATA) + .cloned() + .ok_or_else(|| Error::InvalidSchema("missing tensor metadata file".to_string()))? + }; + let payload = { + let guard = file.read::().await?; + guard.clone() + }; + decode_schema(&payload) +} + +async fn write_metadata_file(blocks: &DirLock, name: &str, payload: &str) -> Result<()> +where + FE: AsType + From + FileLoad + Send + Sync + 'static, +{ + let existing = { + let dir = blocks.read().await; + dir.get_file(name).cloned() + }; + if let Some(file) = existing { + let mut guard = file.write::().await?; + *guard = payload.to_string(); + } else { + let mut dir = blocks.write().await; + dir.create_file(name.to_string(), payload.to_string(), payload.len())?; + } + Ok(()) +} + fn encode_schema(schema: &TensorSchema) -> String { let dtype = schema.dtype().as_str(); - let layout = match schema.layout() { Layout::Dense => "dense".to_string(), Layout::Sparse { axis } => format!( "sparse:{}", - axis.map(|value| value.to_string()) + axis.map(|a| a.to_string()) .unwrap_or_else(|| "none".to_string()) ), }; - let shape = schema .shape() .iter() @@ -639,10 +1258,10 @@ fn encode_schema(schema: &TensorSchema) -> String { .map(|dim| dim.to_string()) .collect::>() .join(","); - - format!( + let s = format!( "version={METADATA_VERSION}\ndtype={dtype}\nlayout={layout}\nshape={shape}\nblock_shape={block_shape}\nstrides={strides}\n" - ) + ); + s } fn decode_schema(payload: &str) -> Result { @@ -695,20 +1314,21 @@ fn decode_schema(payload: &str) -> Result { .ok_or_else(|| Error::InvalidSchema("missing strides in metadata".to_string()))?, )?; - TensorSchema::new( + let schema = TensorSchema::new( dtype, shape.into(), layout, block_shape.into(), strides.into(), - ) + )?; + + Ok(schema) } fn parse_layout(layout: &str) -> Result { if layout == "dense" { return Ok(Layout::Dense); } - if let Some(axis_hint) = layout.strip_prefix("sparse:") { let axis = if axis_hint == "none" { None @@ -717,10 +1337,8 @@ fn parse_layout(layout: &str) -> Result { Error::InvalidSchema(format!("invalid sparse axis hint: {cause}")) })?) }; - return Ok(Layout::Sparse { axis }); } - Err(Error::InvalidSchema(format!( "invalid layout in metadata: {layout}" ))) @@ -730,7 +1348,6 @@ fn parse_usize_vec(value: &str) -> Result> { if value.is_empty() { return Ok(vec![]); } - value .split(',') .map(|dim| { @@ -748,49 +1365,12 @@ fn validate_tensor_dtype(schema: &TensorSchema) -> Result<()> T::DTYPE, ))); } - Ok(()) } -type SparseIndex = TableLock, FE>; - -fn create_sparse_index( - schema: &TensorSchema, - dir_guard: &mut freqfs::DirWriteGuard<'_, FE>, -) -> io::Result>> -where - FE: FileLoad + AsType> + Send + Sync + 'static, -{ - if matches!(schema.layout(), Layout::Sparse { .. }) { - let index_dir = dir_guard.create_dir(INDEX.to_string())?; - - let table_schema = SparseTableSchema::default(); - let collator = Collator::default(); - - TableLock::create(table_schema, collator, index_dir).map(Some) - } else { - Ok(None) - } -} - -fn load_sparse_index( - schema: &TensorSchema, - dir_guard: &mut freqfs::DirWriteGuard<'_, FE>, -) -> io::Result>> -where - FE: FileLoad + AsType> + Send + Sync + 'static, -{ - if matches!(schema.layout(), Layout::Sparse { .. }) - && let Some(index_dir) = dir_guard.get_dir(INDEX).cloned() - { - let table_schema = SparseTableSchema::default(); - let collator = Collator::default(); - - TableLock::load(table_schema, collator, index_dir).map(Some) - } else { - Ok(None) - } -} +// --------------------------------------------------------------------------- +// Unit tests +// --------------------------------------------------------------------------- #[cfg(test)] mod metadata_tests { @@ -842,14 +1422,14 @@ mod metadata_tests { #[test] fn metadata_rejects_invalid_layout() { let payload = - "version=1\ndtype=f32\nlayout=weird\nshape=2,3\nblock_shape=1,3\nstrides=3,1\n"; + "version=2\ndtype=f32\nlayout=weird\nshape=2,3\nblock_shape=1,3\nstrides=3,1\n"; let err = decode_schema(payload).expect_err("should reject"); assert!(matches!(err, Error::InvalidSchema(_))); } #[test] fn metadata_rejects_missing_fields() { - let payload = "version=1\ndtype=f32\nlayout=dense\nshape=2,3\n"; + let payload = "version=2\ndtype=f32\nlayout=dense\nshape=2,3\n"; let err = decode_schema(payload).expect_err("should reject"); assert!(matches!(err, Error::InvalidSchema(_))); } diff --git a/src/schema.rs b/src/schema.rs index b528885..a6ab137 100644 --- a/src/schema.rs +++ b/src/schema.rs @@ -9,19 +9,6 @@ use crate::{Error, Result as FResult}; pub const PORTABLE_INLINE_RANK: usize = 8; pub type TensorShape = SmallVec<[u64; PORTABLE_INLINE_RANK]>; -#[derive(Clone, Eq, PartialEq, Debug)] -pub enum AxisContribSchema { - Stride(i64), - Gather(SmallVec<[i64; PORTABLE_INLINE_RANK]>), -} - -#[derive(Clone, Eq, PartialEq, Debug)] -pub struct ViewSchema { - pub base_rank: usize, - pub base_offset: i64, - pub axes: SmallVec<[AxisContribSchema; PORTABLE_INLINE_RANK]>, -} - #[derive(Clone, Eq, PartialEq, Debug)] struct InternalLayoutMetadata { layout: Layout, @@ -234,6 +221,13 @@ impl TensorSchema { pub fn rank(&self) -> usize { self.shape.len() } + + /// Total number of elements described by this schema's shape, computed + /// with overflow checking so a corrupted/adversarial shape fails closed + /// instead of silently wrapping. + pub fn element_count(&self) -> FResult { + checked_product(self.shape.as_slice()) + } } fn validate_shape_dims(shape: &[T]) -> FResult<()> @@ -268,6 +262,89 @@ pub fn contiguous_strides(shape: &[usize]) -> Strides { strides.into() } +fn checked_product(shape: &[usize]) -> FResult { + shape + .iter() + .try_fold(1usize, |acc, &dim| acc.checked_mul(dim)) + .ok_or_else(|| Error::InvalidSchema("shape element count overflows usize".to_string())) +} + +/// Row-major (C-order) coordinate walk over `shape`, used to materialize or +/// reconstruct a tensor view's data as a flat sequence for wire transfer. +pub(crate) struct RowMajorCoords { + shape: Vec, + coord: Vec, + remaining: usize, +} + +pub(crate) fn row_major_coords(shape: &[usize]) -> FResult { + let remaining = if shape.is_empty() { + 0 + } else { + checked_product(shape)? + }; + + Ok(RowMajorCoords { + shape: shape.to_vec(), + coord: vec![0u64; shape.len()], + remaining, + }) +} + +impl Iterator for RowMajorCoords { + type Item = Vec; + + fn next(&mut self) -> Option { + if self.remaining == 0 { + return None; + } + + let out = self.coord.clone(); + self.remaining -= 1; + + if self.remaining > 0 { + let mut axis = self.shape.len() - 1; + loop { + self.coord[axis] += 1; + if (self.coord[axis] as usize) < self.shape[axis] { + break; + } + self.coord[axis] = 0; + if axis == 0 { + break; + } + axis -= 1; + } + } + + Some(out) + } +} + +/// Build the destination schema for a materialized view snapshot: same +/// dtype/shape as `source`, with layout normalized for a fresh, independent +/// destination tensor and freshly computed `block_shape`/`strides`. +/// +/// A `Sparse { axis }` hint is preserved only when `source` is currently an +/// identity/base view (`is_identity`) -- `transpose`/`slice`/`reshape` never +/// update `layout` when they mutate `shape`/`strides`, so a stale axis hint +/// surviving a transform could silently denote the wrong logical axis (or, +/// after a rank-reducing slice, be out of bounds). Resetting it to `None` on +/// any non-identity view sidesteps that silent-corruption risk; `None` is +/// always valid and defaults to axis 0 downstream. +pub(crate) fn snapshot_schema(source: &TensorSchema, is_identity: bool) -> FResult { + let shape = source.shape().clone(); + let layout = match source.layout() { + Layout::Dense => Layout::Dense, + Layout::Sparse { axis } => Layout::Sparse { + axis: if is_identity { *axis } else { None }, + }, + }; + let block_shape = crate::default_block_shape(&shape); + let strides = contiguous_strides(&shape); + TensorSchema::new(source.dtype(), shape, layout, block_shape, strides) +} + #[derive(Clone, Eq, PartialEq, Debug)] pub struct SparseIndexSchema { columns: Vec, diff --git a/src/stream.rs b/src/stream.rs index 2f488fa..948d818 100644 --- a/src/stream.rs +++ b/src/stream.rs @@ -1,11 +1,11 @@ +use std::fmt; + use destream::{de, en}; -use crate::wire_tags::{ - AXIS_CONTRIB_TAG_GATHER, AXIS_CONTRIB_TAG_STRIDE, LAYOUT_TAG_DENSE, LAYOUT_TAG_SPARSE, -}; +use crate::wire_tags::{ELEM_TAG_ERROR, ELEM_TAG_PAIR, LAYOUT_TAG_DENSE, LAYOUT_TAG_SPARSE}; use crate::{ - AxisContribSchema, DType, Layout, Tensor, TensorElement, TensorFileEntry, TensorSchema, - ViewSchema, contiguous_strides, + DType, Error, Layout, Tensor, TensorArray, TensorElement, TensorFileEntry, TensorRead, + TensorSchema, TensorViewSemantics, TensorWrite, contiguous_strides, schema, }; fn encode_sparse_axis(axis: Option) -> Result, E> { @@ -20,20 +20,6 @@ fn encode_layout(layout: Layout) -> Result<(u8, Option), E> { } } -fn encode_axis_contrib_ref(a: &AxisContribSchema) -> (u8, Vec) { - match a { - AxisContribSchema::Stride(s) => (AXIS_CONTRIB_TAG_STRIDE, vec![*s]), - AxisContribSchema::Gather(g) => (AXIS_CONTRIB_TAG_GATHER, g.iter().copied().collect()), - } -} - -fn encode_axis_contrib(a: AxisContribSchema) -> (u8, Vec) { - match a { - AxisContribSchema::Stride(s) => (AXIS_CONTRIB_TAG_STRIDE, vec![s]), - AxisContribSchema::Gather(g) => (AXIS_CONTRIB_TAG_GATHER, g.into_iter().collect()), - } -} - fn encode_tensor_schema_ref(schema: &TensorSchema) -> Result<(DType, Vec, Layout), String> { let shape = schema .shape_u64() @@ -47,42 +33,28 @@ fn encode_tensor_schema(schema: TensorSchema) -> Result<(DType, Vec, Layout encode_tensor_schema_ref(&schema) } -type EncodedViewSchema = (u64, i64, Vec); - -fn encode_view_schema_ref(schema: &ViewSchema) -> Result { - let base_rank = - u64::try_from(schema.base_rank).map_err(|_| "base rank overflow".to_string())?; - let axes = schema.axes.iter().cloned().collect(); - Ok((base_rank, schema.base_offset, axes)) -} - -fn encode_view_schema(schema: ViewSchema) -> Result { - let base_rank = - u64::try_from(schema.base_rank).map_err(|_| "base rank overflow".to_string())?; - let axes = schema.axes.into_iter().collect(); - Ok((base_rank, schema.base_offset, axes)) -} - -fn encode_tensor_ref(tensor: &Tensor) -> Result<(TensorSchema, ViewSchema), String> +fn encode_tensor_ref(tensor: &Tensor) -> Result where FE: TensorFileEntry, T: TensorElement, { - let schema = tensor.schema.clone(); - let view = tensor.view_schema().map_err(|err| err.to_string())?; + if !tensor.is_base_tensor() { + return Err( + "cannot serialize a tensor with a non-identity view; views are metadata-only and \ + are never persisted" + .to_string(), + ); + } - Ok((schema, view)) + Ok(tensor.schema().clone()) } -fn encode_tensor(tensor: Tensor) -> Result<(TensorSchema, ViewSchema), String> +fn encode_tensor(tensor: Tensor) -> Result where FE: TensorFileEntry, T: TensorElement, { - let view = tensor.view_schema().map_err(|err| err.to_string())?; - let schema = tensor.schema; - - Ok((schema, view)) + encode_tensor_ref(&tensor) } impl de::FromStream for DType { @@ -183,76 +155,304 @@ impl<'en> en::IntoStream<'en> for TensorSchema { } } -impl de::FromStream for AxisContribSchema { +impl de::FromStream for Tensor +where + FE: TensorFileEntry + safecast::AsType + From, + T: TensorElement, +{ + type Context = freqfs::DirLock; + + async fn from_stream( + dir: Self::Context, + decoder: &mut D, + ) -> Result { + let schema = TensorSchema::from_stream((), decoder).await?; + + Tensor::create(dir, schema).await.map_err(de::Error::custom) + } +} + +impl<'en, FE, T> en::ToStream<'en> for Tensor +where + FE: TensorFileEntry, + T: TensorElement, +{ + fn to_stream>(&'en self, encoder: E) -> Result { + en::IntoStream::into_stream(encode_tensor_ref(self).map_err(en::Error::custom)?, encoder) + } +} + +impl<'en, FE, T> en::IntoStream<'en> for Tensor +where + FE: TensorFileEntry, + T: TensorElement, +{ + fn into_stream>(self, encoder: E) -> Result { + en::IntoStream::into_stream(encode_tensor(self).map_err(en::Error::custom)?, encoder) + } +} + +// --------------------------------------------------------------------------- +// Tensor view streaming: lazy encode + streaming decode +// --------------------------------------------------------------------------- + +/// A closed, wire-safe classification of a `crate::Error` that occurred while +/// the encoder was lazily reading values from storage. Carries no payload +/// data (no paths, no coordinates, no underlying `io::Error` text) since this +/// crosses a machine boundary; only which failure shape occurred survives. +#[derive(Debug, Clone, Copy)] +#[repr(u8)] +enum ElemErrorCode { + Io = 0, + Nd = 1, + InvalidSchema = 2, + InvalidCoord = 3, + InvalidLayout = 4, + SparseIndex = 5, + UnsupportedSparseIterationOrder = 6, + Unsupported = 7, + DataMismatch = 8, +} + +impl ElemErrorCode { + fn from_error(err: &Error) -> Self { + match err { + Error::Io(_) => Self::Io, + Error::Nd(_) => Self::Nd, + Error::InvalidSchema(_) => Self::InvalidSchema, + Error::InvalidCoord(_) => Self::InvalidCoord, + Error::InvalidLayout(_) => Self::InvalidLayout, + Error::SparseIndex(_) => Self::SparseIndex, + Error::UnsupportedSparseIterationOrder { .. } => Self::UnsupportedSparseIterationOrder, + Error::Unsupported(_) => Self::Unsupported, + Error::DataMismatch(_) => Self::DataMismatch, + } + } + + fn try_from_u8(code: u8) -> Option { + match code { + 0 => Some(Self::Io), + 1 => Some(Self::Nd), + 2 => Some(Self::InvalidSchema), + 3 => Some(Self::InvalidCoord), + 4 => Some(Self::InvalidLayout), + 5 => Some(Self::SparseIndex), + 6 => Some(Self::UnsupportedSparseIterationOrder), + 7 => Some(Self::Unsupported), + 8 => Some(Self::DataMismatch), + _ => None, + } + } +} + +impl fmt::Display for ElemErrorCode { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::Io => "I/O error", + Self::Nd => "array computation error", + Self::InvalidSchema => "invalid schema", + Self::InvalidCoord => "invalid coordinate", + Self::InvalidLayout => "invalid layout", + Self::SparseIndex => "sparse index error", + Self::UnsupportedSparseIterationOrder => "unsupported sparse iteration order", + Self::Unsupported => "unsupported operation", + Self::DataMismatch => "data mismatch", + }) + } +} + +enum Elem { + Pair(Vec, T), + Error(ElemErrorCode), +} + +impl de::FromStream for Elem +where + T: TensorElement, +{ type Context = (); async fn from_stream(_: (), decoder: &mut D) -> Result { - let (tag, data): (u8, Vec) = <(u8, Vec)>::from_stream((), decoder).await?; + let (tag, coord, value, error_code): (u8, Vec, Option, Option) = + <(u8, Vec, Option, Option)>::from_stream((), decoder).await?; match tag { - AXIS_CONTRIB_TAG_STRIDE => { - if data.len() != 1 { - return Err(de::Error::custom( - "stride axis contrib expects one i64 payload", - )); - } - Ok(Self::Stride(data[0])) + ELEM_TAG_PAIR => { + let value = + value.ok_or_else(|| de::Error::custom("missing tensor view pair value"))?; + Ok(Elem::Pair(coord, value)) + } + ELEM_TAG_ERROR => { + let code = error_code + .ok_or_else(|| de::Error::custom("missing tensor view error code"))?; + let code = ElemErrorCode::try_from_u8(code).ok_or_else(|| { + de::Error::custom(format!("unknown tensor view error code {code}")) + })?; + Ok(Elem::Error(code)) } - AXIS_CONTRIB_TAG_GATHER => Ok(Self::Gather(data.into())), - _ => Err(de::Error::custom(format!("unknown axis contrib tag {tag}"))), + _ => Err(de::Error::custom(format!( + "unknown tensor view stream element tag {tag}" + ))), } } } -impl<'en> en::ToStream<'en> for AxisContribSchema { +impl<'en, T> en::ToStream<'en> for Elem +where + T: TensorElement, +{ fn to_stream>(&'en self, encoder: E) -> Result { - en::IntoStream::into_stream(encode_axis_contrib_ref(self), encoder) + match self { + Elem::Pair(coord, value) => en::IntoStream::into_stream( + (ELEM_TAG_PAIR, coord.clone(), Some(*value), None::), + encoder, + ), + Elem::Error(code) => en::IntoStream::into_stream( + ( + ELEM_TAG_ERROR, + Vec::::new(), + None::, + Some(*code as u8), + ), + encoder, + ), + } } } -impl<'en> en::IntoStream<'en> for AxisContribSchema { +impl<'en, T> en::IntoStream<'en> for Elem +where + T: TensorElement, +{ fn into_stream>(self, encoder: E) -> Result { - en::IntoStream::into_stream(encode_axis_contrib(self), encoder) + match self { + Elem::Pair(coord, value) => en::IntoStream::into_stream( + (ELEM_TAG_PAIR, coord, Some(value), None::), + encoder, + ), + Elem::Error(code) => en::IntoStream::into_stream( + ( + ELEM_TAG_ERROR, + Vec::::new(), + None::, + Some(code as u8), + ), + encoder, + ), + } } } -impl de::FromStream for ViewSchema { - type Context = (); +pub struct TensorViewEncoder<'a, FE, T> { + tensor: &'a Tensor, +} - async fn from_stream(_: (), decoder: &mut D) -> Result { - let (base_rank, base_offset, axes): (u64, i64, Vec) = - <(u64, i64, Vec)>::from_stream((), decoder).await?; +impl<'a, FE, T> TensorViewEncoder<'a, FE, T> { + pub(crate) fn new(tensor: &'a Tensor) -> Self { + Self { tensor } + } +} - let base_rank = - usize::try_from(base_rank).map_err(|_| de::Error::custom("base rank overflow"))?; +fn encode_view<'en, E, FE, T>(tensor: &'en Tensor, encoder: E) -> Result +where + E: en::Encoder<'en>, + FE: TensorFileEntry, + T: TensorElement, +{ + let is_identity = tensor.is_base_tensor(); + let schema = + schema::snapshot_schema(tensor.schema(), is_identity).map_err(en::Error::custom)?; + let shape: Vec = tensor.schema().shape().iter().copied().collect(); + let pairs = TensorPairs { tensor, shape }; + en::IntoStream::into_stream((schema, pairs), encoder) +} - Ok(Self { - base_rank, - base_offset, - axes: axes.into(), - }) +impl<'en, FE, T> en::ToStream<'en> for TensorViewEncoder<'_, FE, T> +where + FE: TensorFileEntry, + T: TensorElement, +{ + fn to_stream>(&'en self, encoder: E) -> Result { + encode_view(self.tensor, encoder) } } -impl<'en> en::ToStream<'en> for ViewSchema { - fn to_stream>(&'en self, encoder: E) -> Result { - en::IntoStream::into_stream( - encode_view_schema_ref(self).map_err(en::Error::custom)?, - encoder, - ) +impl<'a, 'en, FE, T> en::IntoStream<'en> for TensorViewEncoder<'a, FE, T> +where + 'a: 'en, + FE: TensorFileEntry, + T: TensorElement, +{ + fn into_stream>(self, encoder: E) -> Result { + encode_view(self.tensor, encoder) } } -impl<'en> en::IntoStream<'en> for ViewSchema { +struct TensorPairs<'a, FE, T> { + tensor: &'a Tensor, + shape: Vec, +} + +enum PairState<'a, FE, T> { + Walking { + tensor: &'a Tensor, + coords: crate::schema::RowMajorCoords, + }, + Done, +} + +impl<'a, 'en, FE, T> en::IntoStream<'en> for TensorPairs<'a, FE, T> +where + 'a: 'en, + FE: TensorFileEntry, + T: TensorElement, +{ fn into_stream>(self, encoder: E) -> Result { - en::IntoStream::into_stream( - encode_view_schema(self).map_err(en::Error::custom)?, - encoder, - ) + let coords = crate::schema::row_major_coords(&self.shape).map_err(en::Error::custom)?; + let initial = PairState::Walking { + tensor: self.tensor, + coords, + }; + let stream = Box::pin(futures::stream::unfold(initial, |state| async move { + match state { + PairState::Walking { tensor, mut coords } => loop { + match coords.next() { + Some(coord) => match tensor.read_value(&coord).await { + Ok(value) if value == T::default() => continue, + Ok(value) => { + break Some(( + Elem::Pair(coord, value), + PairState::Walking { tensor, coords }, + )); + } + Err(err) => { + break Some(( + Elem::Error(ElemErrorCode::from_error(&err)), + PairState::Done, + )); + } + }, + None => break None, + } + }, + PairState::Done => None, + } + })); + encoder.encode_seq_stream(stream) } } -impl de::FromStream for Tensor +pub struct TensorViewDecoder { + tensor: Tensor, +} + +impl TensorViewDecoder { + pub fn into_inner(self) -> Tensor { + self.tensor + } +} + +impl de::FromStream for TensorViewDecoder where FE: TensorFileEntry + safecast::AsType + From, T: TensorElement, @@ -263,34 +463,130 @@ where dir: Self::Context, decoder: &mut D, ) -> Result { - let (schema, view): (TensorSchema, ViewSchema) = - <(TensorSchema, ViewSchema)>::from_stream((), decoder).await?; + decoder + .decode_seq(TensorViewVisitor { + dir, + _marker: std::marker::PhantomData, + }) + .await + } +} + +struct TensorViewVisitor { + dir: freqfs::DirLock, + _marker: std::marker::PhantomData, +} + +impl de::Visitor for TensorViewVisitor +where + FE: TensorFileEntry + safecast::AsType + From, + T: TensorElement, +{ + type Value = TensorViewDecoder; + + fn expecting() -> &'static str { + "a tensor view (schema followed by pairs sequence)" + } + + async fn visit_seq(self, mut seq: A) -> Result { + let schema: TensorSchema = seq + .next_element(()) + .await? + .ok_or_else(|| de::Error::custom("missing tensor view schema"))?; + + if schema.dtype() != T::DTYPE { + return Err(de::Error::custom(format!( + "tensor dtype mismatch: schema {:?} != tensor {:?}", + schema.dtype(), + T::DTYPE, + ))); + } - let tensor = Tensor::create(dir, schema) + let tensor = Tensor::::create(self.dir.clone(), schema) .await .map_err(de::Error::custom)?; - tensor.with_view_schema(&view).map_err(de::Error::custom) + // Decode the nested pairs sequence, moving `tensor` in by value + // (Context) and getting it back out as the decoded Value on success + // -- each nonzero value is written to storage as soon as it arrives + // off the wire, with no in-memory buffering. + match seq.next_element::>(tensor).await { + Ok(Some(DecodedPairs { tensor })) => Ok(TensorViewDecoder { tensor }), + Ok(None) => { + let mut dir_guard = self.dir.write().await; + let _ = dir_guard.truncate_and_sync().await; + Err(de::Error::custom("missing tensor view data")) + } + Err(err) => { + let mut dir_guard = self.dir.write().await; + let _ = dir_guard.truncate_and_sync().await; + Err(err) + } + } } } -impl<'en, FE, T> en::ToStream<'en> for Tensor +struct DecodedPairs { + tensor: Tensor, +} + +impl de::FromStream for DecodedPairs where FE: TensorFileEntry, T: TensorElement, { - fn to_stream>(&'en self, encoder: E) -> Result { - en::IntoStream::into_stream(encode_tensor_ref(self).map_err(en::Error::custom)?, encoder) + type Context = Tensor; + + async fn from_stream( + tensor: Tensor, + decoder: &mut D, + ) -> Result { + let tensor = decoder.decode_seq(PairsVisitor { tensor }).await?; + Ok(Self { tensor }) } } -impl<'en, FE, T> en::IntoStream<'en> for Tensor +struct PairsVisitor { + tensor: Tensor, +} + +impl de::Visitor for PairsVisitor where FE: TensorFileEntry, T: TensorElement, { - fn into_stream>(self, encoder: E) -> Result { - en::IntoStream::into_stream(encode_tensor(self).map_err(en::Error::custom)?, encoder) + type Value = Tensor; + + fn expecting() -> &'static str { + "a sequence of tensor view element pairs (coord, value)" + } + + async fn visit_seq(self, mut seq: A) -> Result { + while let Some(elem) = seq.next_element::>(()).await? { + match elem { + Elem::Pair(coord, value) => { + if value == T::default() { + // defensive: encode never emits defaults, but don't + // trust the wire. + continue; + } + self.tensor + .write_value(&coord, value) + .await + .map_err(de::Error::custom)?; + } + Elem::Error(code) => { + return Err(de::Error::custom(format!( + "tensor view encode-side failure: {code}" + ))); + } + } + } + // Natural exhaustion of the sequence is itself the success signal -- + // `tbon`'s SeqAccess only returns `None` after consuming a real + // `LIST_END` delimiter, so a genuinely truncated stream surfaces as + // a decode error before this point is ever reached. + Ok(self.tensor) } } @@ -310,7 +606,5 @@ mod tests { assert_stream::(); assert_stream::(); assert_stream::(); - assert_stream::(); - assert_stream::(); } } diff --git a/src/traits.rs b/src/traits.rs index 08b47b5..d56c955 100644 --- a/src/traits.rs +++ b/src/traits.rs @@ -171,6 +171,14 @@ pub trait TensorSparseIndex: Send + Sync { fn lookup_block_id<'a>(&'a self, key: &'a [u64]) -> BoxFuture<'a, Result>>; fn upsert_block_id<'a>(&'a self, key: Vec, block_id: u64) -> BoxFuture<'a, Result<()>>; + + fn delete_row<'a>(&'a self, _key: Vec) -> BoxFuture<'a, Result> { + Box::pin(async move { + Err(Error::Unsupported( + "delete_row is not implemented for this tensor backend".to_string(), + )) + }) + } } /// Base/view capability contract, aligned with v1 writeability semantics. @@ -184,28 +192,6 @@ pub trait TensorViewSemantics: TensorArray { } } -/// Sparse lifecycle policy for zero-write handling and index cleanup. -#[derive(Copy, Clone, Debug, Eq, PartialEq)] -pub enum SparseZeroPolicy { - RemoveRow, - Tombstone, - RetainZero, -} - -pub trait TensorSparseLifecycle: TensorArray { - fn sparse_zero_policy(&self) -> SparseZeroPolicy { - SparseZeroPolicy::RemoveRow - } - - fn compact_sparse<'a>(&'a self) -> BoxFuture<'a, Result<()>> { - Box::pin(async move { - Err(Error::Unsupported( - "sparse compaction is not implemented for this tensor backend".to_string(), - )) - }) - } -} - /// Unary tensor math operations. pub trait TensorUnary: TensorArray + Sized { fn exp<'a>(&'a self) -> BoxFuture<'a, Result> { diff --git a/src/view.rs b/src/view.rs index 77d6726..4b58eeb 100644 --- a/src/view.rs +++ b/src/view.rs @@ -2,11 +2,10 @@ use std::sync::Arc; use ha_ndarray::{Axes, AxisRange, Range, Shape, Strides}; -use crate::{AxisContribSchema, Error, Result, TensorSchema, ViewSchema, schema}; +use crate::{Error, Result, TensorSchema, schema}; #[derive(Clone)] pub struct TensorView { - base_rank: usize, base_offset: i64, axes: Vec, } @@ -18,13 +17,8 @@ pub(crate) enum AxisContrib { } impl TensorView { - pub fn rank(&self) -> usize { - self.axes.len() - } - pub fn identity(schema: &TensorSchema) -> Self { Self { - base_rank: schema.rank(), base_offset: 0, axes: schema .strides() @@ -56,7 +50,6 @@ impl TensorView { } Ok(Self { - base_rank: self.base_rank, base_offset: self.base_offset, axes, }) @@ -159,7 +152,6 @@ impl TensorView { Ok(( Self { - base_rank: self.base_rank, base_offset: new_base_offset, axes: new_axes, }, @@ -168,42 +160,6 @@ impl TensorView { )) } - pub fn to_schema(&self) -> Result { - Ok(ViewSchema { - base_rank: self.base_rank, - base_offset: self.base_offset, - axes: self - .axes - .iter() - .map(|a| match a { - AxisContrib::Stride(s) => AxisContribSchema::Stride(*s), - AxisContrib::Gather(g) => { - AxisContribSchema::Gather(g.iter().copied().collect()) - } - }) - .collect(), - }) - } - - pub fn from_schema(schema: &ViewSchema) -> Result { - let axes = schema - .axes - .iter() - .map(|a| match a { - AxisContribSchema::Stride(s) => Ok(AxisContrib::Stride(*s)), - AxisContribSchema::Gather(g) => Ok(AxisContrib::Gather( - g.iter().copied().collect::>().into(), - )), - }) - .collect::>>()?; - - Ok(Self { - base_rank: schema.base_rank, - base_offset: schema.base_offset, - axes, - }) - } - pub fn reshape(&self, current_shape: &[usize], new_shape: &[usize]) -> Result { if !self.is_c_contiguous(current_shape) { return Err(Error::Unsupported( @@ -213,7 +169,6 @@ impl TensorView { )); } Ok(Self { - base_rank: self.base_rank, base_offset: self.base_offset, axes: schema::contiguous_strides(new_shape) .iter() @@ -246,6 +201,16 @@ impl TensorView { Ok(k) } + pub fn is_identity(&self, base_schema: &TensorSchema) -> bool { + self.base_offset == 0 && self.is_c_contiguous(base_schema.shape()) + } + + pub fn has_gather_axes(&self) -> bool { + self.axes + .iter() + .any(|a| matches!(a, AxisContrib::Gather(_))) + } + fn is_c_contiguous(&self, shape: &[usize]) -> bool { if self.axes.len() != shape.len() { return false; @@ -378,32 +343,6 @@ mod tests { assert_eq!(sliced.flat_offset(&[1, 1]).expect("offset"), 53); } - // -- schema roundtrip --------------------------------------------------------- - - #[test] - fn view_schema_roundtrip() { - // [5,4,3], strides [12,3,1], perm [2,0,1] → axes [Stride(1),Stride(12),Stride(3)] - // flat_offset([1,2,3]) = 1*1 + 2*12 + 3*3 = 34 - let schema = schema(&[5, 4, 3]); - let identity = TensorView::identity(&schema); - let transposed = identity.transpose(&[2, 0, 1]).expect("transpose"); - let view_schema = transposed.to_schema().expect("schema"); - let restored = TensorView::from_schema(&view_schema).expect("restore"); - assert_eq!(restored.flat_offset(&[1, 2, 3]).expect("coord"), 34); - } - - #[test] - fn view_schema_roundtrip_with_reshape() { - // ViewSchema carries only base_rank, base_offset, axes — no shape/strides duplication - // [6,4]→[2,3,4]: axes [Stride(12),Stride(4),Stride(1)], coord [1,0,2] → k=14 - let schema = schema(&[6, 4]); - let view = TensorView::identity(&schema); - let reshaped = view.reshape(schema.shape(), &[2, 3, 4]).expect("reshape"); - let view_schema = reshaped.to_schema().expect("to_schema"); - let restored = TensorView::from_schema(&view_schema).expect("from_schema"); - assert_eq!(restored.flat_offset(&[1, 0, 2]).expect("offset"), 14); - } - // -- invalid reshape chain guards --------------------------------------------- #[test] diff --git a/src/wire_tags.rs b/src/wire_tags.rs index 060160a..99aec22 100644 --- a/src/wire_tags.rs +++ b/src/wire_tags.rs @@ -1,5 +1,5 @@ pub(crate) const LAYOUT_TAG_DENSE: u8 = 0; pub(crate) const LAYOUT_TAG_SPARSE: u8 = 1; -pub(crate) const AXIS_CONTRIB_TAG_STRIDE: u8 = 0; -pub(crate) const AXIS_CONTRIB_TAG_GATHER: u8 = 1; +pub(crate) const ELEM_TAG_PAIR: u8 = 0; +pub(crate) const ELEM_TAG_ERROR: u8 = 2; diff --git a/tests/access_matrix.rs b/tests/access_matrix.rs index d449f17..03ed3de 100644 --- a/tests/access_matrix.rs +++ b/tests/access_matrix.rs @@ -17,9 +17,9 @@ use std::collections::HashMap; use std::path::PathBuf; use fensor::{ - BoxFuture, DType, Error, Layout, SparseZeroPolicy, Tensor, TensorArray, TensorBlockStore, - TensorRead, TensorReadBulk, TensorSchema, TensorSparseIndex, TensorSparseLifecycle, - TensorTransform, TensorViewSemantics, TensorWrite, TensorWriteBulk, contiguous_strides, + BoxFuture, DType, Error, Layout, SparseTensor, Tensor, TensorArray, TensorBlockStore, + TensorRead, TensorReadBulk, TensorSchema, TensorSparseIndex, TensorTransform, + TensorViewSemantics, TensorWrite, TensorWriteBulk, contiguous_strides, }; use ha_ndarray::{Axes, AxisRange, Range, Shape, axes, range, shape}; @@ -62,10 +62,10 @@ async fn create_sparse( shape: Shape, block_shape: Shape, axis: Option, -) -> (PathBuf, Tensor, TensorSchema) { +) -> (PathBuf, SparseTensor, TensorSchema) { let (root, dir) = new_dir(name).await; let schema = sparse_schema_f32(shape, block_shape, axis); - let tensor = Tensor::::create(dir, schema.clone()) + let tensor = SparseTensor::::create(dir, schema.clone()) .await .expect("create sparse"); (root, tensor, schema) @@ -75,7 +75,10 @@ fn encode_value(coord: &[u64]) -> f32 { (coord[0] as f32) * 100.0 + (coord[1] as f32) * 10.0 + (coord[2] as f32) } -async fn seed_values(tensor: &Tensor) { +async fn seed_values(tensor: &T) +where + T: TensorWrite, +{ for coord in iter_coords(tensor.shape()) { tensor .write_value(&coord, encode_value(&coord)) @@ -1239,30 +1242,6 @@ mod section_e_sparse_lifecycle { cleanup(&root).await; } - #[tokio::test] - #[ignore = "requires Tensor::with_sparse_zero_policy / create_with_policy"] - async fn sparse_nonzero_to_zero_retain_zero_policy() { - // Intended shape once the policy setter lands: - // - // let tensor = Tensor::::create_with_policy( - // dir, schema, SparseZeroPolicy::RetainZero, - // ).await.expect("create"); - // tensor.write_value(&[0, 1, 2], 5.0).await.expect("write nz"); - // tensor.write_value(&[0, 1, 2], 0.0).await.expect("write z"); - // assert!(tensor.lookup_block_id(&[1, 1]).await.unwrap().is_some()); - // assert_eq!(tensor.read_value(&[0, 1, 2]).await.unwrap(), 0.0); - panic!("requires Tensor::with_sparse_zero_policy / create_with_policy"); - } - - #[tokio::test] - #[ignore = "requires Tensor::with_sparse_zero_policy / create_with_policy"] - async fn sparse_nonzero_to_zero_tombstone_policy() { - // Same shape as above, but with `SparseZeroPolicy::Tombstone`. - // Exact contract (sentinel block_id vs zero-filled block) to be - // specified during implementation. - panic!("requires Tensor::with_sparse_zero_policy / create_with_policy"); - } - #[tokio::test] async fn sparse_overwrite_nonzero_preserves_row_id() { let (root, tensor, schema) = @@ -1308,13 +1287,90 @@ mod section_e_sparse_lifecycle { } #[tokio::test] - #[ignore = "requires policy persistence + setter"] - async fn sparse_zero_policy_persists_across_reload() { - // Once the policy setter lands and is persisted in metadata: - // 1. Create tensor with non-default policy - // 2. Write values, sync, drop - // 3. Reload, assert sparse_zero_policy() matches what was set - panic!("requires SparseZeroPolicy persistence"); + async fn compact_sparse_preserves_nonzero_rows() { + let (root, tensor, schema) = create_sparse( + "e_compact_preserve", + shape![2, 3, 4], + shape![1, 1, 4], + Some(1), + ) + .await; + + tensor.write_value(&[0, 1, 2], 5.0).await.expect("write nz"); + + tensor.compact_sparse().await.expect("compact"); + + assert!( + block_id_for_coord(&tensor, &schema, &[0, 1, 2]) + .await + .is_some(), + "compact must not remove rows with nonzero values" + ); + assert_eq!(tensor.read_value(&[0, 1, 2]).await.expect("read"), 5.0); + + cleanup(&root).await; + } + + #[tokio::test] + async fn sparse_write_zero_to_new_coord_is_noop() { + let (root, tensor, schema) = + create_sparse("e_noop", shape![2, 3, 4], shape![1, 1, 4], Some(1)).await; + + tensor + .write_value(&[0, 1, 2], 0.0) + .await + .expect("write zero"); + + assert!( + block_id_for_coord(&tensor, &schema, &[0, 1, 2]) + .await + .is_none(), + "zero write to new coord must not create a row" + ); + + cleanup(&root).await; + } + + #[tokio::test] + async fn compact_sparse_idempotent() { + let (root, tensor, schema) = + create_sparse("e_compact_idem", shape![2, 3, 4], shape![1, 1, 4], Some(1)).await; + + tensor.write_value(&[0, 1, 2], 5.0).await.expect("nz"); + tensor.write_value(&[1, 2, 3], 3.0).await.expect("nz2"); + tensor.write_value(&[0, 1, 2], 0.0).await.expect("zero"); + + tensor.compact_sparse().await.expect("first compact"); + + assert!( + block_id_for_coord(&tensor, &schema, &[0, 1, 2]) + .await + .is_none(), + "zero row removed after first compact" + ); + assert!( + block_id_for_coord(&tensor, &schema, &[1, 2, 3]) + .await + .is_some(), + "nonzero row preserved after first compact" + ); + + tensor.compact_sparse().await.expect("second compact"); + + assert!( + block_id_for_coord(&tensor, &schema, &[0, 1, 2]) + .await + .is_none(), + "still absent after second compact" + ); + assert!( + block_id_for_coord(&tensor, &schema, &[1, 2, 3]) + .await + .is_some(), + "nonzero row still present after second compact" + ); + + cleanup(&root).await; } } @@ -1634,38 +1690,6 @@ mod section_h_persistence { cleanup(&root).await; } - #[tokio::test] - async fn view_schema_persists_via_with_view_schema() { - let root = common::unique_tmp_dir("h_view_persist"); - tokio::fs::create_dir(&root).await.expect("mkdir"); - let schema = dense_schema_f32(shape![2, 3, 4], shape![1, 1, 4]); - - let view_schema = { - let dir = open_dir(&root).expect("open"); - let tensor = Tensor::::create(dir.clone(), schema.clone()) - .await - .expect("create"); - seed_values(&tensor).await; - let transposed = tensor.clone().transpose(Some(axes![2, 0, 1])).expect("tx"); - let vs = transposed.view_schema().expect("view schema"); - dir.sync().await.expect("sync"); - vs - }; - - let dir2 = open_dir(&root).expect("reopen"); - let loaded = Tensor::::load_with_schema(dir2, &schema) - .await - .expect("reload"); - let rehydrated = loaded - .with_view_schema(&view_schema) - .expect("rehydrate view"); - assert_eq!(rehydrated.shape(), &[4, 2, 3]); - // sample one coord to verify the view still resolves correctly - let _ = rehydrated.read_value(&[0, 0, 0]).await.expect("read"); - - cleanup(&root).await; - } - #[tokio::test] async fn metadata_file_missing_fails_closed() { let root = common::unique_tmp_dir("h_meta_missing"); diff --git a/tests/common.rs b/tests/common.rs index 5403023..0fae2d4 100644 --- a/tests/common.rs +++ b/tests/common.rs @@ -10,7 +10,7 @@ use std::path::{Path, PathBuf}; use b_table::Node; use destream::{de, en}; -use fensor::{Layout, Tensor, TensorSchema, TensorSparseIndex}; +use fensor::{Layout, TensorSchema, TensorSparseIndex}; use freqfs::{Cache, DirLock}; use safecast::as_type; @@ -22,108 +22,29 @@ pub enum FsEntry { Text(String), } -#[derive(Clone, Copy)] -#[repr(u8)] -pub enum FsEntryTag { - Node = 0, - F32 = 1, - F64 = 2, - Text = 3, -} - -impl FsEntryTag { - fn to_u8(self) -> u8 { - self as u8 - } - - fn from_u8(tag: u8) -> Option { - match tag { - x if x == Self::Node as u8 => Some(Self::Node), - x if x == Self::F32 as u8 => Some(Self::F32), - x if x == Self::F64 as u8 => Some(Self::F64), - x if x == Self::Text as u8 => Some(Self::Text), - _ => None, - } - } -} - impl<'en> en::ToStream<'en> for FsEntry { fn to_stream>(&'en self, encoder: E) -> Result { match self { - Self::Node(node) => en::IntoStream::into_stream( - ( - FsEntryTag::Node.to_u8(), - Some(node.clone()), - None::>, - None::>, - None::, - ), - encoder, - ), - Self::F32(values) => en::IntoStream::into_stream( - ( - FsEntryTag::F32.to_u8(), - None::>, - Some(values.clone()), - None::>, - None::, - ), - encoder, - ), - Self::F64(values) => en::IntoStream::into_stream( - ( - FsEntryTag::F64.to_u8(), - None::>, - None::>, - Some(values.clone()), - None::, - ), - encoder, - ), - Self::Text(text) => en::IntoStream::into_stream( - ( - FsEntryTag::Text.to_u8(), - None::>, - None::>, - None::>, - Some(text.clone()), - ), - encoder, - ), + Self::Node(node) => node.to_stream(encoder), + Self::F32(values) => values.to_stream(encoder), + Self::F64(values) => values.to_stream(encoder), + Self::Text(text) => text.to_stream(encoder), } } } -type DecodedFsEntry = ( - u8, - Option>, - Option>, - Option>, - Option, -); - +// `TensorFileEntry: FileLoad` is only satisfiable via the blanket +// `impl FileLoad for T`, so `FsEntry` needs a `FromStream` impl to +// type-check. Every read in this codebase goes through a concrete `AsType` target +// (`String`/`Vec`/`Vec`/`Node`), never through `FsEntry` itself, so +// this is never actually invoked at runtime. impl de::FromStream for FsEntry { type Context = (); - async fn from_stream(_: (), decoder: &mut D) -> Result { - let (tag, node, f32_values, f64_values, text): DecodedFsEntry = - ::from_stream((), decoder).await?; - - match FsEntryTag::from_u8(tag) { - Some(FsEntryTag::Node) => node - .map(Self::Node) - .ok_or_else(|| de::Error::custom("missing node payload")), - Some(FsEntryTag::F32) => f32_values - .map(Self::F32) - .ok_or_else(|| de::Error::custom("missing f32 payload")), - Some(FsEntryTag::F64) => f64_values - .map(Self::F64) - .ok_or_else(|| de::Error::custom("missing f64 payload")), - Some(FsEntryTag::Text) => text - .map(Self::Text) - .ok_or_else(|| de::Error::custom("missing string payload")), - None => Err(de::Error::custom(format!("unknown fs entry tag {tag}"))), - } + async fn from_stream(_: (), _decoder: &mut D) -> Result { + Err(de::Error::custom( + "FsEntry does not support generic decoding; read via a concrete AsType target", + )) } } @@ -195,8 +116,8 @@ pub fn block_key_for_coord(schema: &TensorSchema, coord: &[u64]) -> Vec { /// Test helper: probe the sparse index for the block backing `coord`. /// Wraps `block_key_for_coord` + `TensorSparseIndex::lookup_block_id` so test /// bodies can talk in coord-space instead of hand-spelling block keys. -pub async fn block_id_for_coord( - tensor: &Tensor, +pub async fn block_id_for_coord( + tensor: &T, schema: &TensorSchema, coord: &[u64], ) -> Option { diff --git a/tests/stream_format.rs b/tests/stream_format.rs new file mode 100644 index 0000000..b0c0bbe --- /dev/null +++ b/tests/stream_format.rs @@ -0,0 +1,64 @@ +//! Exercises `Tensor`'s actual destream wire format via a real `tbon` encoder/decoder. + +mod common; + +use fensor::{ + DType, Layout, Tensor, TensorArray, TensorSchema, TensorTransform, contiguous_strides, +}; +use ha_ndarray::{axes, shape}; + +use common::{FsEntry, cleanup, open_dir, unique_tmp_dir}; + +fn dense_schema_f32(shape: ha_ndarray::Shape, block_shape: ha_ndarray::Shape) -> TensorSchema { + let strides = contiguous_strides(&shape); + TensorSchema::new(DType::F32, shape, Layout::Dense, block_shape, strides).expect("schema") +} + +#[tokio::test] +async fn base_tensor_round_trips_through_tbon() { + let root = unique_tmp_dir("stream_base_roundtrip"); + tokio::fs::create_dir(&root).await.expect("mkdir"); + let schema = dense_schema_f32(shape![2, 3], shape![1, 3]); + + let dir = open_dir(&root).expect("open"); + let tensor = Tensor::::create(dir, schema.clone()) + .await + .expect("create"); + + let encoded = tbon::en::encode(tensor).expect("encode base tensor"); + + let decode_root = root.join("decoded"); + tokio::fs::create_dir(&decode_root) + .await + .expect("mkdir decoded"); + let dir2 = open_dir(&decode_root).expect("open decode target"); + let decoded: Tensor = tbon::de::try_decode(dir2, encoded) + .await + .expect("decode base tensor"); + + assert_eq!(decoded.schema(), &schema); + + cleanup(&root).await; +} + +#[tokio::test] +async fn non_identity_view_is_rejected_on_encode() { + let root = unique_tmp_dir("stream_non_identity_rejected"); + tokio::fs::create_dir(&root).await.expect("mkdir"); + let schema = dense_schema_f32(shape![2, 3], shape![1, 3]); + + let dir = open_dir(&root).expect("open"); + let tensor = Tensor::::create(dir, schema) + .await + .expect("create"); + + let transposed = tensor.transpose(Some(axes![1, 0])).expect("transpose"); + + let result = tbon::en::encode(transposed); + assert!( + result.is_err(), + "encoding a non-identity view must be rejected, not silently drop the transform" + ); + + cleanup(&root).await; +} diff --git a/tests/view_snapshot_stream.rs b/tests/view_snapshot_stream.rs new file mode 100644 index 0000000..f739b4c --- /dev/null +++ b/tests/view_snapshot_stream.rs @@ -0,0 +1,433 @@ +//! Exercises `Tensor::view_encoder` / `TensorViewEncoder` / `TensorViewDecoder`: +//! streaming a tensor's current view (identity or transformed, dense or +//! sparse) together with only its non-default element data, and reconstructing +//! a fresh, independent base tensor from it via a real `tbon` round trip. + +mod common; + +use fensor::{ + DType, Layout, Tensor, TensorArray, TensorRead, TensorSchema, TensorTransform, + TensorViewDecoder, TensorViewSemantics, TensorWrite, contiguous_strides, +}; +use futures::stream::TryStreamExt; +use ha_ndarray::{AxisRange, Shape, axes, range, shape}; + +use common::{FsEntry, cleanup, iter_coords, open_dir, unique_tmp_dir}; + +fn dense_schema_f32(shape: Shape, block_shape: Shape) -> TensorSchema { + let strides = contiguous_strides(&shape); + TensorSchema::new(DType::F32, shape, Layout::Dense, block_shape, strides).expect("schema") +} + +fn sparse_schema_f32(shape: Shape, block_shape: Shape, axis: Option) -> TensorSchema { + let strides = contiguous_strides(&shape); + TensorSchema::new( + DType::F32, + shape, + Layout::Sparse { axis }, + block_shape, + strides, + ) + .expect("schema") +} + +#[tokio::test] +async fn identity_dense_view_round_trips_with_data() { + let root = unique_tmp_dir("view_snapshot_identity_dense"); + tokio::fs::create_dir(&root).await.expect("mkdir"); + let schema = dense_schema_f32(shape![2, 3], shape![1, 3]); + + let dir = open_dir(&root).expect("open"); + let tensor = Tensor::::create(dir, schema.clone()) + .await + .expect("create"); + + let mut expected = Vec::new(); + for coord in iter_coords(schema.shape()) { + let value = (coord[0] * 10 + coord[1]) as f32 + 1.0; + tensor.write_value(&coord, value).await.expect("write"); + expected.push((coord, value)); + } + + let encoded = tbon::en::encode(tensor.view_encoder()).expect("encode view"); + + let decode_root = root.join("decoded"); + tokio::fs::create_dir(&decode_root) + .await + .expect("mkdir decoded"); + let dir2 = open_dir(&decode_root).expect("open decode target"); + let decoded: TensorViewDecoder = tbon::de::try_decode(dir2, encoded) + .await + .expect("decode view"); + + let decoded = decoded.into_inner(); + + assert!(decoded.is_base_tensor()); + assert_eq!(decoded.schema().shape(), schema.shape()); + assert_eq!(decoded.schema().layout(), &Layout::Dense); + + for (coord, value) in expected { + let read = decoded.read_value(&coord).await.expect("read"); + assert_eq!(read, value); + } + + cleanup(&root).await; +} + +#[tokio::test] +async fn identity_sparse_view_round_trips_with_data_and_axis() { + let root = unique_tmp_dir("view_snapshot_identity_sparse"); + tokio::fs::create_dir(&root).await.expect("mkdir"); + let schema = sparse_schema_f32(shape![2, 3], shape![1, 3], Some(0)); + + let dir = open_dir(&root).expect("open"); + let tensor = Tensor::::create(dir, schema.clone()) + .await + .expect("create"); + + tensor.write_value(&[0, 1], 5.0f32).await.expect("write"); + tensor.write_value(&[1, 2], 9.0f32).await.expect("write"); + + let encoded = tbon::en::encode(tensor.view_encoder()).expect("encode view"); + + let decode_root = root.join("decoded"); + tokio::fs::create_dir(&decode_root) + .await + .expect("mkdir decoded"); + let dir2 = open_dir(&decode_root).expect("open decode target"); + let decoded: TensorViewDecoder = tbon::de::try_decode(dir2, encoded) + .await + .expect("decode view"); + + let decoded = decoded.into_inner(); + + assert!(decoded.is_base_tensor()); + assert_eq!(decoded.schema().layout(), &Layout::Sparse { axis: Some(0) }); + + for coord in iter_coords(schema.shape()) { + let expected = tensor.read_value(&coord).await.expect("read source"); + let actual = decoded.read_value(&coord).await.expect("read decoded"); + assert_eq!(actual, expected); + } + + cleanup(&root).await; +} + +#[tokio::test] +async fn non_identity_dense_view_round_trips_with_data() { + let root = unique_tmp_dir("view_snapshot_non_identity_dense"); + tokio::fs::create_dir(&root).await.expect("mkdir"); + let schema = dense_schema_f32(shape![3, 4], shape![1, 4]); + + let dir = open_dir(&root).expect("open"); + let tensor = Tensor::::create(dir, schema.clone()) + .await + .expect("create"); + + for coord in iter_coords(schema.shape()) { + let value = (coord[0] * 10 + coord[1]) as f32 + 1.0; + tensor.write_value(&coord, value).await.expect("write"); + } + + let sliced = tensor + .clone() + .slice(range![AxisRange::In(1, 3, 1), AxisRange::In(0, 4, 2)]) + .expect("slice") + .transpose(Some(axes![1, 0])) + .expect("transpose"); + + assert!(!sliced.is_base_tensor()); + + let mut expected = Vec::new(); + for coord in iter_coords(sliced.schema().shape()) { + let value = sliced.read_value(&coord).await.expect("read source view"); + expected.push((coord, value)); + } + + let encoded = tbon::en::encode(sliced.view_encoder()).expect("encode view"); + + let decode_root = root.join("decoded"); + tokio::fs::create_dir(&decode_root) + .await + .expect("mkdir decoded"); + let dir2 = open_dir(&decode_root).expect("open decode target"); + let decoded: TensorViewDecoder = tbon::de::try_decode(dir2, encoded) + .await + .expect("decode view"); + + let decoded = decoded.into_inner(); + + assert!(decoded.is_base_tensor()); + assert_eq!(decoded.schema().shape(), sliced.schema().shape()); + + for (coord, value) in expected { + let read = decoded.read_value(&coord).await.expect("read"); + assert_eq!(read, value); + } + + cleanup(&root).await; +} + +#[tokio::test] +async fn non_identity_sparse_view_round_trips_with_data_resets_axis() { + let root = unique_tmp_dir("view_snapshot_non_identity_sparse"); + tokio::fs::create_dir(&root).await.expect("mkdir"); + let schema = sparse_schema_f32(shape![2, 3], shape![1, 3], Some(0)); + + let dir = open_dir(&root).expect("open"); + let tensor = Tensor::::create(dir, schema.clone()) + .await + .expect("create"); + + tensor.write_value(&[0, 1], 7.0f32).await.expect("write"); + tensor.write_value(&[1, 2], 3.0f32).await.expect("write"); + + let transposed = tensor + .clone() + .transpose(Some(axes![1, 0])) + .expect("transpose"); + assert!(!transposed.is_base_tensor()); + + let mut expected = Vec::new(); + for coord in iter_coords(transposed.schema().shape()) { + let value = transposed + .read_value(&coord) + .await + .expect("read source view"); + expected.push((coord, value)); + } + + let encoded = tbon::en::encode(transposed.view_encoder()).expect("encode view"); + + let decode_root = root.join("decoded"); + tokio::fs::create_dir(&decode_root) + .await + .expect("mkdir decoded"); + let dir2 = open_dir(&decode_root).expect("open decode target"); + let decoded: TensorViewDecoder = tbon::de::try_decode(dir2, encoded) + .await + .expect("decode view"); + + let decoded = decoded.into_inner(); + + assert!(decoded.is_base_tensor()); + assert_eq!(decoded.schema().layout(), &Layout::Sparse { axis: None }); + + for (coord, value) in expected { + let read = decoded.read_value(&coord).await.expect("read"); + assert_eq!(read, value); + } + + cleanup(&root).await; +} + +#[tokio::test] +async fn only_nonzero_values_are_transmitted() { + let root = unique_tmp_dir("view_snapshot_sparse_transmission"); + tokio::fs::create_dir(&root).await.expect("mkdir"); + + // Create a large dense tensor (20x20 = 400 elements), but write only ONE nonzero value. + // This demonstrates that the encoder only transmits non-default values, not all 400 elements. + let schema = dense_schema_f32(shape![20, 20], shape![5, 5]); + + let dir = open_dir(&root).expect("open"); + let tensor = Tensor::::create(dir, schema.clone()) + .await + .expect("create"); + + // Write exactly one nonzero value; everything else remains at default (0.0) + tensor.write_value(&[5, 7], 3.0f32).await.expect("write"); + + let encoded_stream = tbon::en::encode(tensor.view_encoder()).expect("encode view"); + // Collect the stream to measure size and prepare for decoding + let encoded_parts = encoded_stream + .try_collect::>() + .await + .expect("collect encoded stream"); + let mut total_size = 0; + for part in &encoded_parts { + total_size += part.len(); + } + + // Assert that the encoded size is small. A single f32 value (4 bytes) plus schema + // overhead should be far smaller than 400 f32 values (1600 bytes). + // Being generous, we allow up to 200 bytes to account for wire framing, tags, and metadata. + // This verifies that only the one nonzero value was transmitted, not all 400 elements. + assert!( + total_size < 200, + "encoded size {} should be small for sparse transmission (proves only nonzero values transmitted)", + total_size + ); + + // Round-trip: decode and verify correctness + let decode_root = root.join("decoded"); + tokio::fs::create_dir(&decode_root) + .await + .expect("mkdir decoded"); + let dir2 = open_dir(&decode_root).expect("open decode target"); + // Create a stream from the collected chunks for decoding + let encoded_stream_for_decode = + futures::stream::iter(encoded_parts.into_iter().map(Ok::<_, tbon::de::Error>)); + let decoded: TensorViewDecoder = + tbon::de::try_decode(dir2, encoded_stream_for_decode) + .await + .expect("decode view"); + + let decoded = decoded.into_inner(); + + assert!(decoded.is_base_tensor()); + assert_eq!(decoded.schema().shape(), schema.shape()); + + // Verify the single written value is present + let value = decoded.read_value(&[5, 7]).await.expect("read"); + assert_eq!(value, 3.0f32); + + // Verify all other coordinates are zero (default) + for coord in iter_coords(schema.shape()) { + if coord[0] == 5 && coord[1] == 7 { + // Skip the one written value, we already checked it + continue; + } + let value = decoded.read_value(&coord).await.expect("read"); + assert_eq!(value, 0.0f32, "expected zero at {:?}", coord); + } + + cleanup(&root).await; +} + +#[tokio::test] +async fn truncated_stream_fails_closed_and_removes_partial_storage() { + let root = unique_tmp_dir("view_snapshot_verification_mismatch"); + tokio::fs::create_dir(&root).await.expect("mkdir"); + let schema = dense_schema_f32(shape![2, 3], shape![1, 3]); + + let dir = open_dir(&root).expect("open"); + let tensor = Tensor::::create(dir, schema) + .await + .expect("create"); + + // Write a couple of nonzero values + tensor.write_value(&[0, 1], 5.0f32).await.expect("write"); + tensor.write_value(&[1, 2], 9.0f32).await.expect("write"); + + let encoded_stream = tbon::en::encode(tensor.view_encoder()).expect("encode view"); + // Collect the stream so we can create a truncated version + let encoded_parts = encoded_stream + .try_collect::>() + .await + .expect("collect encoded stream"); + + // Create a corrupted stream that is incomplete: keep only the first chunk (schema). + // The decoder will create the tensor, then try to read the pairs sequence. Since the + // stream ends immediately, the underlying `tbon` decoder itself fails with an + // "unexpected end of stream" error (it never sees the pairs sequence's closing + // delimiter), before the decoder-side visitor ever gets a chance to return. + let corrupted_parts: Vec<_> = if !encoded_parts.is_empty() { + vec![encoded_parts[0].clone()] + } else { + vec![] + }; + + let decode_root = root.join("decoded"); + tokio::fs::create_dir(&decode_root) + .await + .expect("mkdir decoded"); + let dir2 = open_dir(&decode_root).expect("open decode target"); + + // Create a stream from the truncated chunks for decoding attempt + let encoded_stream_corrupted = + futures::stream::iter(corrupted_parts.into_iter().map(Ok::<_, tbon::de::Error>)); + + // Attempt to decode with incomplete stream; should fail due to the truncated pairs sequence + let result: std::result::Result, _> = + tbon::de::try_decode(dir2.clone(), encoded_stream_corrupted).await; + + assert!( + result.is_err(), + "corrupted verification (incomplete stream) must fail" + ); + + // Verify that the decode target directory was cleaned up (no tensor storage left behind). + // Try to read the directory contents; it should be empty or sparse. + let mut dir_entries = tokio::fs::read_dir(&decode_root) + .await + .expect("read decode_root"); + + let mut found_tensor_content = false; + while let Some(entry) = dir_entries.next_entry().await.expect("read next entry") { + let path = entry.path(); + // If there's a "blocks" directory with actual block files, that's tensor content that should have been cleaned up. + if path.ends_with("blocks") + && let Ok(mut block_entries) = tokio::fs::read_dir(&path).await + && block_entries.next_entry().await.ok().flatten().is_some() + { + found_tensor_content = true; + } + } + + assert!( + !found_tensor_content, + "partial tensor storage should have been cleaned up after decode failure" + ); + + cleanup(&root).await; +} + +#[tokio::test] +async fn corrupted_dtype_mismatch_fails_closed() { + let root = unique_tmp_dir("view_snapshot_dtype_mismatch"); + tokio::fs::create_dir(&root).await.expect("mkdir"); + + // Create an f64 tensor + let shape_f64 = shape![2, 2]; + let block_shape_f64 = shape![1, 2]; + let f64_schema = TensorSchema::new( + DType::F64, + shape_f64.clone(), + Layout::Dense, + block_shape_f64, + contiguous_strides(&shape_f64), + ) + .expect("f64 schema"); + + let dir = open_dir(&root).expect("open"); + let tensor_f64 = Tensor::::create(dir, f64_schema) + .await + .expect("create f64 tensor"); + + tensor_f64.write_value(&[0, 0], 1.5).await.expect("write"); + + let encoded = tbon::en::encode(tensor_f64.view_encoder()).expect("encode f64 view"); + + let decode_root = root.join("decoded"); + tokio::fs::create_dir(&decode_root) + .await + .expect("mkdir decoded"); + let dir2 = open_dir(&decode_root).expect("open decode target"); + + // Try to decode f64-encoded data as f32; should fail with dtype mismatch + let result: std::result::Result, _> = + tbon::de::try_decode(dir2.clone(), encoded).await; + + assert!(result.is_err(), "dtype mismatch must fail closed"); + + // Verify no storage was created in the decode target (dtype check is early, before Tensor::create) + let mut dir_entries = tokio::fs::read_dir(&decode_root) + .await + .expect("read decode_root"); + + let mut found_storage = false; + while let Some(entry) = dir_entries.next_entry().await.expect("read next entry") { + let path = entry.path(); + if path.ends_with("blocks") || path.ends_with("metadata") { + found_storage = true; + } + } + + assert!( + !found_storage, + "no tensor storage should be created on dtype mismatch" + ); + + cleanup(&root).await; +}