Skip to main content

dag_ml_data_core/
nd_tensor.rs

1//! Borrowed N-D tensor transport: host-owned dense tensors (e.g. uint8
2//! `[N,H,W,C]` RGB, float32 `[N,H,W,Wavelength]` hyperspectral) copied into
3//! provider-owned, canonical row-major storage, keyed and gathered by
4//! `observation_id`.
5//!
6//! Axis 0 is ALWAYS the sample/observation axis; the remaining axes are opaque
7//! to this layer (no conv / decode / preprocessing / interpolation). Strides are
8//! a transport detail of the *input*: borrowed strided input is copied into
9//! canonical contiguous bytes at construction and strides are then discarded, so
10//! everything stored here is contiguous row-major and fingerprintable.
11//!
12//! This is the ND parallel of [`crate::buffer`]: a [`NdTensorStore`] of immutable
13//! tensors, an [`NdTensorArena`] that binds them to materialized data handles by
14//! output representation, and a relation-ordered [`NdTensorBlock`] export.
15
16use std::collections::{BTreeMap, BTreeSet};
17
18use serde::{Deserialize, Serialize};
19
20use crate::content_hash::StreamingHasher;
21use crate::coordinator::CoordinatorRelationSet;
22use crate::error::{DataError, Result};
23use crate::ids::{ObservationId, RepresentationId, SampleId, SourceId};
24
25/// Schema version of [`NdTensorManifest`] / [`NdTensorBinding`].
26pub const ND_TENSOR_MANIFEST_SCHEMA_VERSION: u32 = 1;
27
28/// Maximum supported tensor rank. Bounds the shape-product arithmetic and keeps
29/// the contract finite.
30pub const ND_TENSOR_MAX_RANK: usize = 16;
31
32/// Element dtype of an N-D tensor. The C ABI exposes a matching `#[repr(C)]`
33/// `DagMlDataTensorDType` enum and maps to/from this type explicitly; this core
34/// enum carries no ABI guarantee of its own.
35#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
36#[serde(rename_all = "snake_case")]
37pub enum NdTensorDType {
38    F64,
39    F32,
40    U8,
41    I32,
42    Bool,
43}
44
45impl NdTensorDType {
46    /// Size in bytes of one element of this dtype.
47    pub fn element_size(self) -> usize {
48        match self {
49            NdTensorDType::F64 => 8,
50            NdTensorDType::F32 => 4,
51            NdTensorDType::U8 => 1,
52            NdTensorDType::I32 => 4,
53            NdTensorDType::Bool => 1,
54        }
55    }
56
57    /// Stable, explicit tag absorbed into the content fingerprint. Hard-coded
58    /// (not the enum's memory discriminant or its serde name) so the dtype is
59    /// part of the preimage and reordering variants here can never silently
60    /// change a fingerprint. Two tensors whose `data` bytes happen to coincide
61    /// but whose dtypes differ (e.g. `[N]` `U8` vs the low byte of an `I32`)
62    /// fingerprint differently.
63    fn fingerprint_tag(self) -> u64 {
64        match self {
65            NdTensorDType::F64 => 1,
66            NdTensorDType::F32 => 2,
67            NdTensorDType::U8 => 3,
68            NdTensorDType::I32 => 4,
69            NdTensorDType::Bool => 5,
70        }
71    }
72}
73
74/// A canonical (contiguous row-major) N-D tensor handed to the store. The C ABI
75/// layer copies any strided borrowed input into this shape before calling in.
76#[derive(Clone, Debug, PartialEq)]
77pub struct NdTensorInput {
78    pub tensor_id: String,
79    pub representation_id: RepresentationId,
80    pub container: String,
81    pub dtype: NdTensorDType,
82    /// Full shape; `shape[0]` is the sample/observation axis.
83    pub shape: Vec<usize>,
84    /// One id per axis-0 row; `len == shape[0]`.
85    pub observation_ids: Vec<ObservationId>,
86    /// Optional advisory sample ids (validated for length only; the authoritative
87    /// sample mapping comes from the coordinator relations at export time).
88    pub sample_ids: Option<Vec<SampleId>>,
89    /// Contiguous row-major bytes; `len == product(shape) * dtype.element_size()`.
90    ///
91    /// Multibyte elements (`F64`/`F32`/`I32`) MUST be encoded **little-endian**.
92    /// This layer copies the bytes verbatim and never reinterprets them as
93    /// native-endian values, so the byte order is part of the input contract,
94    /// not something the layer can fix up. The content fingerprint hashes these
95    /// bytes as-is; pinning little-endian here is what makes the fingerprint
96    /// platform-independent (a big-endian host must canonicalize to LE before
97    /// constructing the input).
98    pub data: Vec<u8>,
99    /// Optional per-axis-0-row presence mask; `len == shape[0]`.
100    pub row_presence: Option<Vec<bool>>,
101}
102
103/// A stored, immutable, contiguous row-major N-D tensor.
104#[derive(Clone, Debug, PartialEq)]
105pub struct NdTensor {
106    tensor_id: String,
107    representation_id: RepresentationId,
108    container: String,
109    dtype: NdTensorDType,
110    shape: Vec<usize>,
111    observation_ids: Vec<ObservationId>,
112    data: Vec<u8>,
113    row_presence: Option<Vec<bool>>,
114    row_index_by_observation: BTreeMap<ObservationId, usize>,
115    /// Bytes per axis-0 row: `product(shape[1..]) * element_size`.
116    row_stride_bytes: usize,
117}
118
119/// Provider-wide manifest of one stored tensor (no payload bytes).
120#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
121pub struct NdTensorManifest {
122    pub schema_version: u32,
123    pub tensor_id: String,
124    pub representation_id: RepresentationId,
125    pub container: String,
126    pub dtype: NdTensorDType,
127    pub shape: Vec<usize>,
128    pub observation_ids: Vec<ObservationId>,
129    pub row_count: usize,
130    pub element_bytes: usize,
131    pub data_bytes: usize,
132    pub tensor_fingerprint: String,
133}
134
135/// Binding of one stored tensor to a materialized data handle.
136#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
137pub struct NdTensorBinding {
138    pub tensor_id: String,
139    pub representation_id: RepresentationId,
140    pub container: String,
141    pub dtype: NdTensorDType,
142    pub source_ids: Vec<SourceId>,
143    pub shape: Vec<usize>,
144    pub row_count: usize,
145    pub tensor_fingerprint: String,
146}
147
148/// A relation-ordered, axis-0-filtered export of a stored tensor. `shape[0]` is
149/// the number of selected rows; `data` is contiguous row-major bytes.
150#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
151pub struct NdTensorBlock {
152    pub tensor_id: String,
153    pub representation_id: RepresentationId,
154    pub container: String,
155    pub dtype: NdTensorDType,
156    pub shape: Vec<usize>,
157    pub observation_ids: Vec<ObservationId>,
158    pub sample_ids: Vec<SampleId>,
159    pub data: Vec<u8>,
160    #[serde(default, skip_serializing_if = "Option::is_none")]
161    pub row_presence: Option<Vec<bool>>,
162}
163
164/// An immutable collection of N-D tensors keyed by `tensor_id`.
165#[derive(Clone, Debug, Default, PartialEq)]
166pub struct NdTensorStore {
167    tensors: BTreeMap<String, NdTensor>,
168}
169
170/// A store plus per-data-handle bindings, mirroring `NumericFeatureBufferArena`.
171#[derive(Clone, Debug, Default, PartialEq)]
172pub struct NdTensorArena {
173    store: NdTensorStore,
174    data_bindings: BTreeMap<u64, BTreeMap<String, NdTensorBinding>>,
175}
176
177fn checked_shape_product(tensor_id: &str, shape: &[usize]) -> Result<usize> {
178    let mut product: usize = 1;
179    for dim in shape {
180        product = product.checked_mul(*dim).ok_or_else(|| {
181            DataError::Validation(format!(
182                "nd tensor `{tensor_id}` shape product overflows usize"
183            ))
184        })?;
185    }
186    Ok(product)
187}
188
189impl NdTensor {
190    /// Builds and validates a stored tensor from canonical contiguous input.
191    pub fn from_input(input: NdTensorInput) -> Result<Self> {
192        let NdTensorInput {
193            tensor_id,
194            representation_id,
195            container,
196            dtype,
197            shape,
198            observation_ids,
199            sample_ids,
200            data,
201            row_presence,
202        } = input;
203
204        if tensor_id.trim().is_empty() {
205            return Err(DataError::Validation(
206                "nd tensor has an empty tensor id".to_string(),
207            ));
208        }
209        if container.trim().is_empty() {
210            return Err(DataError::Validation(format!(
211                "nd tensor `{tensor_id}` has an empty container"
212            )));
213        }
214        if shape.is_empty() || shape.len() > ND_TENSOR_MAX_RANK {
215            return Err(DataError::Validation(format!(
216                "nd tensor `{tensor_id}` rank {} is not in 1..={ND_TENSOR_MAX_RANK}",
217                shape.len()
218            )));
219        }
220        if shape.contains(&0) {
221            return Err(DataError::Validation(format!(
222                "nd tensor `{tensor_id}` has a zero dimension in shape {shape:?}"
223            )));
224        }
225        let row_count = shape[0];
226        if observation_ids.len() != row_count {
227            return Err(DataError::Validation(format!(
228                "nd tensor `{tensor_id}` has {} observation ids for axis-0 size {row_count}",
229                observation_ids.len()
230            )));
231        }
232        if observation_ids.is_empty() {
233            return Err(DataError::Validation(format!(
234                "nd tensor `{tensor_id}` has no observations"
235            )));
236        }
237        let mut row_index_by_observation = BTreeMap::new();
238        for (idx, observation_id) in observation_ids.iter().enumerate() {
239            if row_index_by_observation
240                .insert(observation_id.clone(), idx)
241                .is_some()
242            {
243                return Err(DataError::Validation(format!(
244                    "nd tensor `{tensor_id}` has duplicate observation `{observation_id}`"
245                )));
246            }
247        }
248        if let Some(sample_ids) = &sample_ids {
249            if sample_ids.len() != row_count {
250                return Err(DataError::Validation(format!(
251                    "nd tensor `{tensor_id}` has {} sample ids for axis-0 size {row_count}",
252                    sample_ids.len()
253                )));
254            }
255        }
256
257        let element_size = dtype.element_size();
258        let total_elements = checked_shape_product(&tensor_id, &shape)?;
259        let expected_bytes = total_elements.checked_mul(element_size).ok_or_else(|| {
260            DataError::Validation(format!("nd tensor `{tensor_id}` byte size overflows usize"))
261        })?;
262        if data.len() != expected_bytes {
263            return Err(DataError::Validation(format!(
264                "nd tensor `{tensor_id}` has {} data bytes for shape {shape:?} dtype {dtype:?} ({expected_bytes} expected)",
265                data.len()
266            )));
267        }
268        if dtype == NdTensorDType::Bool && data.iter().any(|byte| *byte > 1) {
269            return Err(DataError::Validation(format!(
270                "nd tensor `{tensor_id}` bool payload contains a byte that is not 0 or 1"
271            )));
272        }
273        if let Some(row_presence) = &row_presence {
274            if row_presence.len() != row_count {
275                return Err(DataError::Validation(format!(
276                    "nd tensor `{tensor_id}` row presence has {} flags for axis-0 size {row_count}",
277                    row_presence.len()
278                )));
279            }
280        }
281
282        // product(shape[1..]) * element_size; rank-1 rows are one element each.
283        // All dims are > 0 (rejected above), so row_count > 0.
284        let row_stride_bytes = expected_bytes / row_count;
285
286        Ok(Self {
287            tensor_id,
288            representation_id,
289            container,
290            dtype,
291            shape,
292            observation_ids,
293            data,
294            row_presence,
295            row_index_by_observation,
296            row_stride_bytes,
297        })
298    }
299
300    fn contains_observation(&self, observation_id: &ObservationId) -> bool {
301        self.row_index_by_observation.contains_key(observation_id)
302    }
303
304    /// Gathers axis-0 rows in relation order (optionally scoped to one source),
305    /// keyed by `observation_id`. Output `sample_ids` come from the relations,
306    /// not the stored tensor.
307    pub fn project_relations(
308        &self,
309        relations: &CoordinatorRelationSet,
310        source_id: Option<&SourceId>,
311    ) -> Result<NdTensorBlock> {
312        relations.validate()?;
313        let selected = relations.records.iter().filter(|relation| {
314            source_id
315                .map(|source_id| relation.source_id.as_ref() == Some(source_id))
316                .unwrap_or(true)
317        });
318
319        let mut observation_ids = Vec::new();
320        let mut sample_ids = Vec::new();
321        let mut data = Vec::new();
322        let mut presence = Vec::new();
323        for relation in selected {
324            let row_idx = *self
325                .row_index_by_observation
326                .get(&relation.observation_id)
327                .ok_or_else(|| {
328                    DataError::Validation(format!(
329                        "nd tensor `{}` has no row for observation `{}`",
330                        self.tensor_id, relation.observation_id
331                    ))
332                })?;
333            let start = row_idx * self.row_stride_bytes;
334            data.extend_from_slice(&self.data[start..start + self.row_stride_bytes]);
335            if let Some(row_presence) = &self.row_presence {
336                presence.push(row_presence[row_idx]);
337            }
338            observation_ids.push(relation.observation_id.clone());
339            sample_ids.push(relation.sample_id.clone());
340        }
341
342        let mut shape = self.shape.clone();
343        shape[0] = observation_ids.len();
344
345        Ok(NdTensorBlock {
346            tensor_id: self.tensor_id.clone(),
347            representation_id: self.representation_id.clone(),
348            container: self.container.clone(),
349            dtype: self.dtype,
350            shape,
351            observation_ids,
352            sample_ids,
353            data,
354            row_presence: self.row_presence.as_ref().map(|_| presence),
355        })
356    }
357
358    /// Reproducibility fingerprint of the tensor's logical content as a 64-char
359    /// lowercase hex SHA-256.
360    ///
361    /// Streamed, never serialized to an intermediate byte vector: the bulk
362    /// `data` payload is fed straight through the hasher (the old
363    /// `serde_json::to_vec` path materialized the whole `Vec<u8>` as a JSON
364    /// array of integers, multiplying both time and transient memory).
365    ///
366    /// Preimage layout (every multi-byte field little-endian; every
367    /// variable-length field length-prefixed):
368    ///
369    /// ```text
370    /// domain tag  b"dag-ml-data.nd-tensor.v2\0"   (fixed literal)
371    /// tensor_id         : str   (u64 byte-len, then UTF-8)
372    /// representation_id : str
373    /// container         : str
374    /// dtype             : u64   (stable fingerprint tag, see fingerprint_tag)
375    /// shape             : u64 rank, then each dim as u64
376    /// observation_ids   : str[] (u64 count, then each str)
377    /// data              : u64 byte-len, then the contiguous row-major bytes
378    /// row_presence      : u64 tag (0 = absent), or (1, u64 len, one 0/1 byte each)
379    /// ```
380    ///
381    /// `data` is the host's **canonical contiguous row-major** payload: the
382    /// element bytes are already in the platform-independent little-endian
383    /// encoding the C-ABI copies in verbatim (the layer never re-interprets
384    /// them as native-endian values), so feeding the raw bytes IS the
385    /// per-dtype LE encoding. The dtype tag and the byte-length prefix mean
386    /// identical bytes under a different dtype or a different shape never
387    /// collide.
388    fn fingerprint(&self) -> Result<String> {
389        let mut hasher = StreamingHasher::new(b"dag-ml-data.nd-tensor.v2\0");
390        hasher.absorb_str(&self.tensor_id);
391        hasher.absorb_str(self.representation_id.as_str());
392        hasher.absorb_str(&self.container);
393        hasher.absorb_u64(self.dtype.fingerprint_tag());
394        hasher.absorb_len(self.shape.len());
395        for dim in &self.shape {
396            hasher.absorb_len(*dim);
397        }
398        hasher.absorb_str_collection(self.observation_ids.iter().map(ObservationId::as_str));
399        hasher.absorb_len(self.data.len());
400        hasher.absorb_raw(&self.data);
401        match &self.row_presence {
402            None => hasher.absorb_u64(0),
403            Some(presence) => {
404                hasher.absorb_u64(1);
405                hasher.absorb_len(presence.len());
406                for present in presence {
407                    hasher.absorb_raw(&[u8::from(*present)]);
408                }
409            }
410        }
411        Ok(hasher.finalize_hex())
412    }
413
414    fn manifest(&self) -> Result<NdTensorManifest> {
415        Ok(NdTensorManifest {
416            schema_version: ND_TENSOR_MANIFEST_SCHEMA_VERSION,
417            tensor_id: self.tensor_id.clone(),
418            representation_id: self.representation_id.clone(),
419            container: self.container.clone(),
420            dtype: self.dtype,
421            shape: self.shape.clone(),
422            observation_ids: self.observation_ids.clone(),
423            row_count: self.shape[0],
424            element_bytes: self.dtype.element_size(),
425            data_bytes: self.data.len(),
426            tensor_fingerprint: self.fingerprint()?,
427        })
428    }
429
430    fn binding_for_sources(&self, source_ids: Vec<SourceId>) -> Result<NdTensorBinding> {
431        Ok(NdTensorBinding {
432            tensor_id: self.tensor_id.clone(),
433            representation_id: self.representation_id.clone(),
434            container: self.container.clone(),
435            dtype: self.dtype,
436            source_ids,
437            shape: self.shape.clone(),
438            row_count: self.shape[0],
439            tensor_fingerprint: self.fingerprint()?,
440        })
441    }
442}
443
444impl NdTensorStore {
445    /// Builds a store from canonical inputs, rejecting duplicate tensor ids.
446    pub fn from_inputs(inputs: Vec<NdTensorInput>) -> Result<Self> {
447        let mut tensors = BTreeMap::new();
448        for input in inputs {
449            let tensor = NdTensor::from_input(input)?;
450            let tensor_id = tensor.tensor_id.clone();
451            if tensors.insert(tensor_id.clone(), tensor).is_some() {
452                return Err(DataError::Validation(format!(
453                    "duplicate nd tensor `{tensor_id}`"
454                )));
455            }
456        }
457        Ok(Self { tensors })
458    }
459
460    pub fn is_empty(&self) -> bool {
461        self.tensors.is_empty()
462    }
463
464    /// Provider-wide manifests for every stored tensor.
465    pub fn manifests(&self) -> Result<Vec<NdTensorManifest>> {
466        self.tensors.values().map(NdTensor::manifest).collect()
467    }
468
469    /// Bindings for tensors whose representation matches `representation_id` and
470    /// whose observations cover the relation set (per source).
471    pub fn bindings_for_relations(
472        &self,
473        relations: &CoordinatorRelationSet,
474        representation_id: &RepresentationId,
475    ) -> Result<Vec<NdTensorBinding>> {
476        relations.validate()?;
477        let source_ids = relations
478            .records
479            .iter()
480            .filter_map(|relation| relation.source_id.as_ref())
481            .collect::<BTreeSet<_>>();
482
483        let mut bindings = Vec::new();
484        for tensor in self.tensors.values() {
485            if &tensor.representation_id != representation_id {
486                continue;
487            }
488            if source_ids.is_empty() {
489                if relations
490                    .records
491                    .iter()
492                    .all(|relation| tensor.contains_observation(&relation.observation_id))
493                {
494                    bindings.push(tensor.binding_for_sources(Vec::new())?);
495                }
496                continue;
497            }
498            let mut covered_sources = Vec::new();
499            for source_id in &source_ids {
500                let covers_source = relations
501                    .records
502                    .iter()
503                    .filter(|relation| relation.source_id.as_ref() == Some(*source_id))
504                    .all(|relation| tensor.contains_observation(&relation.observation_id));
505                if covers_source {
506                    covered_sources.push((*source_id).clone());
507                }
508            }
509            if !covered_sources.is_empty() {
510                bindings.push(tensor.binding_for_sources(covered_sources)?);
511            }
512        }
513        Ok(bindings)
514    }
515
516    /// Projects a single tensor over the given relations.
517    pub fn project_relations(
518        &self,
519        tensor_id: &str,
520        relations: &CoordinatorRelationSet,
521        source_id: Option<&SourceId>,
522    ) -> Result<NdTensorBlock> {
523        let tensor = self.tensors.get(tensor_id).ok_or_else(|| {
524            DataError::Validation(format!("nd tensor `{tensor_id}` is not present"))
525        })?;
526        tensor.project_relations(relations, source_id)
527    }
528}
529
530impl NdTensorArena {
531    pub fn new(store: NdTensorStore) -> Self {
532        Self {
533            store,
534            data_bindings: BTreeMap::new(),
535        }
536    }
537
538    /// Binds tensors matching `representation_id` to `data_handle` and returns
539    /// the bindings.
540    pub fn bind_data_handle(
541        &mut self,
542        data_handle: u64,
543        relations: &CoordinatorRelationSet,
544        representation_id: &RepresentationId,
545    ) -> Result<Vec<NdTensorBinding>> {
546        let bindings = self
547            .store
548            .bindings_for_relations(relations, representation_id)?;
549        self.data_bindings.insert(
550            data_handle,
551            bindings
552                .iter()
553                .cloned()
554                .map(|binding| (binding.tensor_id.clone(), binding))
555                .collect(),
556        );
557        Ok(bindings)
558    }
559
560    pub fn release_data_handle(&mut self, data_handle: u64) -> bool {
561        self.data_bindings.remove(&data_handle).is_some()
562    }
563
564    pub fn manifests(&self) -> Result<Vec<NdTensorManifest>> {
565        self.store.manifests()
566    }
567
568    pub fn bindings_for_data_handle(&self, data_handle: u64) -> Result<Vec<NdTensorBinding>> {
569        let bindings = self.data_bindings.get(&data_handle).ok_or_else(|| {
570            DataError::Validation(format!(
571                "data handle `{data_handle}` has no nd tensor bindings"
572            ))
573        })?;
574        Ok(bindings.values().cloned().collect())
575    }
576
577    /// Projects a bound tensor over `relations`. Mirrors the feature-buffer
578    /// bound-source check: the tensor must be bound to `data_handle`; an explicit
579    /// `source_id` must be present in the view AND bound; an unscoped export
580    /// requires EVERY source present in the view to be bound.
581    pub fn project_bound_relations(
582        &self,
583        data_handle: u64,
584        tensor_id: &str,
585        relations: &CoordinatorRelationSet,
586        source_id: Option<&SourceId>,
587    ) -> Result<NdTensorBlock> {
588        relations.validate()?;
589        let binding = self
590            .data_bindings
591            .get(&data_handle)
592            .and_then(|bindings| bindings.get(tensor_id))
593            .ok_or_else(|| {
594                DataError::Validation(format!(
595                    "nd tensor `{tensor_id}` is not bound to data handle `{data_handle}`"
596                ))
597            })?;
598        let view_source_ids = relations
599            .records
600            .iter()
601            .filter_map(|relation| relation.source_id.as_ref())
602            .cloned()
603            .collect::<BTreeSet<_>>();
604        let required_source_ids: Vec<SourceId> = if let Some(source_id) = source_id {
605            if view_source_ids.is_empty() || !view_source_ids.contains(source_id) {
606                return Err(DataError::Validation(format!(
607                    "nd tensor `{tensor_id}` source `{source_id}` is not present in view for data handle `{data_handle}`"
608                )));
609            }
610            vec![source_id.clone()]
611        } else {
612            view_source_ids.into_iter().collect()
613        };
614        for required in &required_source_ids {
615            if !binding.source_ids.contains(required) {
616                return Err(DataError::Validation(format!(
617                    "nd tensor `{tensor_id}` is not bound to source `{required}` for data handle `{data_handle}`"
618                )));
619            }
620        }
621        self.store
622            .project_relations(tensor_id, relations, source_id)
623    }
624}
625
626#[cfg(test)]
627mod tests {
628    use super::*;
629    use crate::coordinator::CoordinatorRelation;
630    use crate::ids::SampleId;
631
632    fn relation(observation: &str, sample: &str, source: &str) -> CoordinatorRelation {
633        CoordinatorRelation {
634            observation_id: ObservationId::new(observation).unwrap(),
635            sample_id: SampleId::new(sample).unwrap(),
636            target_id: None,
637            group_id: None,
638            origin_sample_id: None,
639            source_id: Some(SourceId::new(source).unwrap()),
640            is_augmented: false,
641        }
642    }
643
644    // [3, 2, 2] u8 RGB-ish tensor: 3 rows (obs.s1..s3), each a 2x2 frame.
645    fn rgb_input() -> NdTensorInput {
646        NdTensorInput {
647            tensor_id: "rgb".to_string(),
648            representation_id: RepresentationId::new("rgb_image").unwrap(),
649            container: "pil_image_batch".to_string(),
650            dtype: NdTensorDType::U8,
651            shape: vec![3, 2, 2],
652            observation_ids: vec![
653                ObservationId::new("obs.s1").unwrap(),
654                ObservationId::new("obs.s2").unwrap(),
655                ObservationId::new("obs.s3").unwrap(),
656            ],
657            sample_ids: None,
658            data: (0u8..12).collect(),
659            row_presence: None,
660        }
661    }
662
663    #[test]
664    fn from_input_validates_and_projects_in_relation_order() {
665        let store = NdTensorStore::from_inputs(vec![rgb_input()]).unwrap();
666        let relations = CoordinatorRelationSet {
667            records: vec![
668                relation("obs.s3", "s3", "cam"),
669                relation("obs.s1", "s1", "cam"),
670            ],
671        };
672        let block = store.project_relations("rgb", &relations, None).unwrap();
673        assert_eq!(block.shape, vec![2, 2, 2]);
674        assert_eq!(block.dtype, NdTensorDType::U8);
675        // Rows gathered in relation order: s3 (bytes 8..12) then s1 (0..4).
676        assert_eq!(block.data, vec![8, 9, 10, 11, 0, 1, 2, 3]);
677        assert_eq!(
678            block.sample_ids,
679            vec![SampleId::new("s3").unwrap(), SampleId::new("s1").unwrap()]
680        );
681    }
682
683    #[test]
684    fn rejects_wrong_data_len() {
685        let mut input = rgb_input();
686        input.data.pop();
687        let error = NdTensor::from_input(input).unwrap_err();
688        assert!(format!("{error}").contains("data bytes"));
689    }
690
691    #[test]
692    fn rejects_observation_count_mismatch() {
693        let mut input = rgb_input();
694        input.observation_ids.pop();
695        let error = NdTensor::from_input(input).unwrap_err();
696        assert!(format!("{error}").contains("observation ids"));
697    }
698
699    #[test]
700    fn rejects_rank_zero_and_over_max() {
701        let mut zero = rgb_input();
702        zero.shape = vec![];
703        assert!(NdTensor::from_input(zero).is_err());
704        let mut huge = rgb_input();
705        huge.shape = vec![3; ND_TENSOR_MAX_RANK + 1];
706        assert!(NdTensor::from_input(huge).is_err());
707    }
708
709    #[test]
710    fn rejects_non_binary_bool_payload() {
711        let input = NdTensorInput {
712            tensor_id: "mask".to_string(),
713            representation_id: RepresentationId::new("mask_image").unwrap(),
714            container: "ndarray".to_string(),
715            dtype: NdTensorDType::Bool,
716            shape: vec![2, 2],
717            observation_ids: vec![
718                ObservationId::new("obs.s1").unwrap(),
719                ObservationId::new("obs.s2").unwrap(),
720            ],
721            sample_ids: None,
722            data: vec![1, 0, 2, 1],
723            row_presence: None,
724        };
725        let error = NdTensor::from_input(input).unwrap_err();
726        assert!(format!("{error}").contains("not 0 or 1"));
727    }
728
729    #[test]
730    fn arena_binds_and_refuses_unbound_or_wrong_source() {
731        let mut arena = NdTensorArena::new(NdTensorStore::from_inputs(vec![rgb_input()]).unwrap());
732        let relations = CoordinatorRelationSet {
733            records: vec![
734                relation("obs.s1", "s1", "cam"),
735                relation("obs.s2", "s2", "cam"),
736                relation("obs.s3", "s3", "cam"),
737            ],
738        };
739        let representation = RepresentationId::new("rgb_image").unwrap();
740        let bindings = arena
741            .bind_data_handle(1, &relations, &representation)
742            .unwrap();
743        assert_eq!(bindings.len(), 1);
744        assert_eq!(bindings[0].source_ids, vec![SourceId::new("cam").unwrap()]);
745
746        // Bound source projects.
747        let block = arena
748            .project_bound_relations(1, "rgb", &relations, Some(&SourceId::new("cam").unwrap()))
749            .unwrap();
750        assert_eq!(block.shape, vec![3, 2, 2]);
751
752        // Unbound source / handle are refused.
753        assert!(arena
754            .project_bound_relations(1, "rgb", &relations, Some(&SourceId::new("nope").unwrap()))
755            .is_err());
756        assert!(arena
757            .project_bound_relations(2, "rgb", &relations, None)
758            .is_err());
759        assert!(arena.release_data_handle(1));
760        assert!(arena.bindings_for_data_handle(1).is_err());
761    }
762
763    #[test]
764    fn rejects_empty_tensor_id() {
765        let mut input = rgb_input();
766        input.tensor_id = "  ".to_string();
767        let error = NdTensor::from_input(input).unwrap_err();
768        assert!(format!("{error}").contains("empty tensor id"));
769    }
770
771    #[test]
772    fn rejects_zero_dimension() {
773        let mut input = rgb_input();
774        input.shape = vec![3, 0];
775        input.data = Vec::new();
776        let error = NdTensor::from_input(input).unwrap_err();
777        assert!(format!("{error}").contains("zero dimension"));
778    }
779
780    #[test]
781    fn arena_refuses_unscoped_export_when_a_view_source_is_unbound() {
782        // Tensor only contains source `a` observations.
783        let input = NdTensorInput {
784            tensor_id: "multi".to_string(),
785            representation_id: RepresentationId::new("rgb_image").unwrap(),
786            container: "ndarray".to_string(),
787            dtype: NdTensorDType::U8,
788            shape: vec![1, 2],
789            observation_ids: vec![ObservationId::new("obs.a1").unwrap()],
790            sample_ids: None,
791            data: vec![1, 2],
792            row_presence: None,
793        };
794        let mut arena = NdTensorArena::new(NdTensorStore::from_inputs(vec![input]).unwrap());
795        // The view spans sources `a` and `b`.
796        let relations = CoordinatorRelationSet {
797            records: vec![relation("obs.a1", "a1", "a"), relation("obs.b1", "b1", "b")],
798        };
799        let representation = RepresentationId::new("rgb_image").unwrap();
800        let bindings = arena
801            .bind_data_handle(1, &relations, &representation)
802            .unwrap();
803        // Only source `a` is covered.
804        assert_eq!(bindings[0].source_ids, vec![SourceId::new("a").unwrap()]);
805
806        // Unscoped export over the a+b view is refused (b is unbound).
807        assert!(arena
808            .project_bound_relations(1, "multi", &relations, None)
809            .is_err());
810        // Scoped to the uncovered source `b` is refused.
811        assert!(arena
812            .project_bound_relations(1, "multi", &relations, Some(&SourceId::new("b").unwrap()))
813            .is_err());
814        // Scoped to the covered source `a` over an a-only view succeeds.
815        let a_only = CoordinatorRelationSet {
816            records: vec![relation("obs.a1", "a1", "a")],
817        };
818        let block = arena
819            .project_bound_relations(1, "multi", &a_only, Some(&SourceId::new("a").unwrap()))
820            .unwrap();
821        assert_eq!(block.shape, vec![1, 2]);
822    }
823
824    #[test]
825    fn manifest_carries_shape_and_fingerprint() {
826        let store = NdTensorStore::from_inputs(vec![rgb_input()]).unwrap();
827        let manifests = store.manifests().unwrap();
828        assert_eq!(manifests.len(), 1);
829        assert_eq!(manifests[0].shape, vec![3, 2, 2]);
830        assert_eq!(manifests[0].data_bytes, 12);
831        assert_eq!(manifests[0].element_bytes, 1);
832        assert_eq!(manifests[0].tensor_fingerprint.len(), 64);
833    }
834
835    fn fp(input: NdTensorInput) -> String {
836        NdTensor::from_input(input).unwrap().fingerprint().unwrap()
837    }
838
839    #[test]
840    fn tensor_fingerprint_is_64_lowercase_hex() {
841        let fingerprint = fp(rgb_input());
842        assert_eq!(fingerprint.len(), 64);
843        assert!(fingerprint
844            .chars()
845            .all(|c| c.is_ascii_hexdigit() && !c.is_ascii_uppercase()));
846    }
847
848    #[test]
849    fn tensor_fingerprint_is_deterministic_across_calls_and_clone() {
850        let tensor = NdTensor::from_input(rgb_input()).unwrap();
851        let once = tensor.fingerprint().unwrap();
852        assert_eq!(once, tensor.fingerprint().unwrap());
853        assert_eq!(once, tensor.clone().fingerprint().unwrap());
854    }
855
856    #[test]
857    fn tensor_fingerprint_changes_when_a_single_data_byte_flips() {
858        let baseline = fp(rgb_input());
859        let mut flipped = rgb_input();
860        flipped.data[0] ^= 0xFF;
861        assert_ne!(baseline, fp(flipped));
862    }
863
864    #[test]
865    fn tensor_fingerprint_changes_when_tensor_id_is_renamed() {
866        let baseline = fp(rgb_input());
867        let mut renamed = rgb_input();
868        renamed.tensor_id = "rgb_renamed".to_string();
869        assert_ne!(baseline, fp(renamed));
870    }
871
872    #[test]
873    fn tensor_fingerprint_changes_when_observation_ids_are_reordered() {
874        // Permute both the rows' bytes and their ids together so the tensor
875        // stays a valid permutation of the same logical rows.
876        let baseline = fp(rgb_input());
877        let mut reordered = rgb_input();
878        reordered.observation_ids.swap(0, 2);
879        // Each row is 4 bytes ([_,2,2] u8); swap row 0 and row 2.
880        let (head, tail) = reordered.data.split_at_mut(8);
881        head[0..4].swap_with_slice(&mut tail[0..4]);
882        assert_ne!(baseline, fp(reordered));
883    }
884
885    #[test]
886    fn tensor_fingerprint_distinguishes_transposed_shapes_with_identical_bytes() {
887        // Same 12 bytes, axis-0 size 3 either way; [3,2,2] vs [3,4] differ only
888        // in the trailing shape, which the rank+dims framing must separate.
889        let base = NdTensorInput {
890            tensor_id: "t".to_string(),
891            representation_id: RepresentationId::new("rgb_image").unwrap(),
892            container: "ndarray".to_string(),
893            dtype: NdTensorDType::U8,
894            shape: vec![3, 2, 2],
895            observation_ids: vec![
896                ObservationId::new("obs.s1").unwrap(),
897                ObservationId::new("obs.s2").unwrap(),
898                ObservationId::new("obs.s3").unwrap(),
899            ],
900            sample_ids: None,
901            data: (0u8..12).collect(),
902            row_presence: None,
903        };
904        let mut reshaped = base.clone();
905        reshaped.shape = vec![3, 4];
906        assert_ne!(fp(base), fp(reshaped));
907    }
908
909    #[test]
910    fn tensor_fingerprint_distinguishes_dtype_with_identical_bytes() {
911        // Four bytes, axis-0 size 1: one I32 element vs four U8 elements share
912        // the same payload bytes; the dtype tag must separate them.
913        let as_i32 = NdTensorInput {
914            tensor_id: "t".to_string(),
915            representation_id: RepresentationId::new("rgb_image").unwrap(),
916            container: "ndarray".to_string(),
917            dtype: NdTensorDType::I32,
918            shape: vec![1, 1],
919            observation_ids: vec![ObservationId::new("obs.s1").unwrap()],
920            sample_ids: None,
921            data: vec![1, 2, 3, 4],
922            row_presence: None,
923        };
924        let mut as_u8 = as_i32.clone();
925        as_u8.dtype = NdTensorDType::U8;
926        as_u8.shape = vec![1, 4];
927        assert_ne!(fp(as_i32), fp(as_u8));
928    }
929
930    #[test]
931    fn tensor_fingerprint_hashes_data_bytes_verbatim_in_declared_order() {
932        // Conformance: the layer hashes element bytes as-is and never
933        // byte-swaps, so the little-endian input contract is what makes the
934        // fingerprint platform-independent. An f32 = 1.0 is `00 00 80 3f` LE;
935        // feeding the byte-reversed (big-endian) arrangement of the same
936        // logical value must change the fingerprint, proving order is honored.
937        let f32_one_le: [u8; 4] = 1.0f32.to_le_bytes(); // 00 00 80 3f
938        let le = NdTensorInput {
939            tensor_id: "t".to_string(),
940            representation_id: RepresentationId::new("hyperspectral").unwrap(),
941            container: "ndarray".to_string(),
942            dtype: NdTensorDType::F32,
943            shape: vec![1, 1],
944            observation_ids: vec![ObservationId::new("obs.s1").unwrap()],
945            sample_ids: None,
946            data: f32_one_le.to_vec(),
947            row_presence: None,
948        };
949        let mut be = le.clone();
950        let mut reversed = f32_one_le;
951        reversed.reverse(); // 3f 80 00 00 — the BE arrangement
952        be.data = reversed.to_vec();
953        // Two LE builds match; the BE arrangement differs.
954        assert_eq!(fp(le.clone()), fp(le.clone()));
955        assert_ne!(fp(le), fp(be));
956    }
957
958    #[test]
959    fn tensor_fingerprint_distinguishes_row_presence_states() {
960        // No presence vs all-present-present mask are different logical states
961        // and must fingerprint differently (the absent-tag vs the framed mask).
962        let baseline = fp(rgb_input());
963        let mut with_presence = rgb_input();
964        with_presence.row_presence = Some(vec![true, true, true]);
965        let present_fp = fp(with_presence);
966        assert_ne!(baseline, present_fp);
967
968        let mut one_absent = rgb_input();
969        one_absent.row_presence = Some(vec![true, false, true]);
970        assert_ne!(present_fp, fp(one_absent));
971    }
972
973    #[test]
974    #[ignore = "perf sanity probe; run with --release --ignored --nocapture"]
975    fn tensor_fingerprint_large_payload_under_500ms() {
976        // Stream a ~12 MB u8-equivalent / ~12.7 M-element f32 payload; the
977        // streamed hash must stay well under the 500 ms budget (the old
978        // serde_json path expanded every byte to a JSON integer first). The
979        // budget is asserted only in optimized builds — an unoptimized
980        // `cargo test` runs sha2 with overflow checks and no SIMD, so its wall
981        // time is not "native" performance. The number is always printed.
982        let rows = 3021usize;
983        let cols = 1050usize;
984        let element_size = NdTensorDType::F32.element_size();
985        let data = vec![0x3Cu8; rows * cols * element_size];
986        let input = NdTensorInput {
987            tensor_id: "big".to_string(),
988            representation_id: RepresentationId::new("hyperspectral").unwrap(),
989            container: "ndarray".to_string(),
990            dtype: NdTensorDType::F32,
991            shape: vec![rows, cols],
992            observation_ids: (0..rows)
993                .map(|r| ObservationId::new(format!("obs.{r}")).unwrap())
994                .collect(),
995            sample_ids: None,
996            data,
997            row_presence: None,
998        };
999        let tensor = NdTensor::from_input(input).unwrap();
1000        let start = std::time::Instant::now();
1001        let fingerprint = tensor.fingerprint().unwrap();
1002        let elapsed = start.elapsed();
1003        println!(
1004            "nd tensor fingerprint({rows}x{cols} f32) = {:.3} ms (fp={fingerprint})",
1005            elapsed.as_secs_f64() * 1e3
1006        );
1007        assert_eq!(fingerprint.len(), 64);
1008        if !cfg!(debug_assertions) {
1009            assert!(
1010                elapsed.as_millis() < 500,
1011                "tensor fingerprint took {} ms (>= 500 ms budget)",
1012                elapsed.as_millis()
1013            );
1014        }
1015    }
1016}