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            excluded: false,
642            metadata: BTreeMap::new(),
643            tags: Vec::new(),
644        }
645    }
646
647    // [3, 2, 2] u8 RGB-ish tensor: 3 rows (obs.s1..s3), each a 2x2 frame.
648    fn rgb_input() -> NdTensorInput {
649        NdTensorInput {
650            tensor_id: "rgb".to_string(),
651            representation_id: RepresentationId::new("rgb_image").unwrap(),
652            container: "pil_image_batch".to_string(),
653            dtype: NdTensorDType::U8,
654            shape: vec![3, 2, 2],
655            observation_ids: vec![
656                ObservationId::new("obs.s1").unwrap(),
657                ObservationId::new("obs.s2").unwrap(),
658                ObservationId::new("obs.s3").unwrap(),
659            ],
660            sample_ids: None,
661            data: (0u8..12).collect(),
662            row_presence: None,
663        }
664    }
665
666    #[test]
667    fn from_input_validates_and_projects_in_relation_order() {
668        let store = NdTensorStore::from_inputs(vec![rgb_input()]).unwrap();
669        let relations = CoordinatorRelationSet {
670            records: vec![
671                relation("obs.s3", "s3", "cam"),
672                relation("obs.s1", "s1", "cam"),
673            ],
674        };
675        let block = store.project_relations("rgb", &relations, None).unwrap();
676        assert_eq!(block.shape, vec![2, 2, 2]);
677        assert_eq!(block.dtype, NdTensorDType::U8);
678        // Rows gathered in relation order: s3 (bytes 8..12) then s1 (0..4).
679        assert_eq!(block.data, vec![8, 9, 10, 11, 0, 1, 2, 3]);
680        assert_eq!(
681            block.sample_ids,
682            vec![SampleId::new("s3").unwrap(), SampleId::new("s1").unwrap()]
683        );
684    }
685
686    #[test]
687    fn rejects_wrong_data_len() {
688        let mut input = rgb_input();
689        input.data.pop();
690        let error = NdTensor::from_input(input).unwrap_err();
691        assert!(format!("{error}").contains("data bytes"));
692    }
693
694    #[test]
695    fn rejects_observation_count_mismatch() {
696        let mut input = rgb_input();
697        input.observation_ids.pop();
698        let error = NdTensor::from_input(input).unwrap_err();
699        assert!(format!("{error}").contains("observation ids"));
700    }
701
702    #[test]
703    fn rejects_rank_zero_and_over_max() {
704        let mut zero = rgb_input();
705        zero.shape = vec![];
706        assert!(NdTensor::from_input(zero).is_err());
707        let mut huge = rgb_input();
708        huge.shape = vec![3; ND_TENSOR_MAX_RANK + 1];
709        assert!(NdTensor::from_input(huge).is_err());
710    }
711
712    #[test]
713    fn rejects_non_binary_bool_payload() {
714        let input = NdTensorInput {
715            tensor_id: "mask".to_string(),
716            representation_id: RepresentationId::new("mask_image").unwrap(),
717            container: "ndarray".to_string(),
718            dtype: NdTensorDType::Bool,
719            shape: vec![2, 2],
720            observation_ids: vec![
721                ObservationId::new("obs.s1").unwrap(),
722                ObservationId::new("obs.s2").unwrap(),
723            ],
724            sample_ids: None,
725            data: vec![1, 0, 2, 1],
726            row_presence: None,
727        };
728        let error = NdTensor::from_input(input).unwrap_err();
729        assert!(format!("{error}").contains("not 0 or 1"));
730    }
731
732    #[test]
733    fn arena_binds_and_refuses_unbound_or_wrong_source() {
734        let mut arena = NdTensorArena::new(NdTensorStore::from_inputs(vec![rgb_input()]).unwrap());
735        let relations = CoordinatorRelationSet {
736            records: vec![
737                relation("obs.s1", "s1", "cam"),
738                relation("obs.s2", "s2", "cam"),
739                relation("obs.s3", "s3", "cam"),
740            ],
741        };
742        let representation = RepresentationId::new("rgb_image").unwrap();
743        let bindings = arena
744            .bind_data_handle(1, &relations, &representation)
745            .unwrap();
746        assert_eq!(bindings.len(), 1);
747        assert_eq!(bindings[0].source_ids, vec![SourceId::new("cam").unwrap()]);
748
749        // Bound source projects.
750        let block = arena
751            .project_bound_relations(1, "rgb", &relations, Some(&SourceId::new("cam").unwrap()))
752            .unwrap();
753        assert_eq!(block.shape, vec![3, 2, 2]);
754
755        // Unbound source / handle are refused.
756        assert!(arena
757            .project_bound_relations(1, "rgb", &relations, Some(&SourceId::new("nope").unwrap()))
758            .is_err());
759        assert!(arena
760            .project_bound_relations(2, "rgb", &relations, None)
761            .is_err());
762        assert!(arena.release_data_handle(1));
763        assert!(arena.bindings_for_data_handle(1).is_err());
764    }
765
766    #[test]
767    fn rejects_empty_tensor_id() {
768        let mut input = rgb_input();
769        input.tensor_id = "  ".to_string();
770        let error = NdTensor::from_input(input).unwrap_err();
771        assert!(format!("{error}").contains("empty tensor id"));
772    }
773
774    #[test]
775    fn rejects_zero_dimension() {
776        let mut input = rgb_input();
777        input.shape = vec![3, 0];
778        input.data = Vec::new();
779        let error = NdTensor::from_input(input).unwrap_err();
780        assert!(format!("{error}").contains("zero dimension"));
781    }
782
783    #[test]
784    fn arena_refuses_unscoped_export_when_a_view_source_is_unbound() {
785        // Tensor only contains source `a` observations.
786        let input = NdTensorInput {
787            tensor_id: "multi".to_string(),
788            representation_id: RepresentationId::new("rgb_image").unwrap(),
789            container: "ndarray".to_string(),
790            dtype: NdTensorDType::U8,
791            shape: vec![1, 2],
792            observation_ids: vec![ObservationId::new("obs.a1").unwrap()],
793            sample_ids: None,
794            data: vec![1, 2],
795            row_presence: None,
796        };
797        let mut arena = NdTensorArena::new(NdTensorStore::from_inputs(vec![input]).unwrap());
798        // The view spans sources `a` and `b`.
799        let relations = CoordinatorRelationSet {
800            records: vec![relation("obs.a1", "a1", "a"), relation("obs.b1", "b1", "b")],
801        };
802        let representation = RepresentationId::new("rgb_image").unwrap();
803        let bindings = arena
804            .bind_data_handle(1, &relations, &representation)
805            .unwrap();
806        // Only source `a` is covered.
807        assert_eq!(bindings[0].source_ids, vec![SourceId::new("a").unwrap()]);
808
809        // Unscoped export over the a+b view is refused (b is unbound).
810        assert!(arena
811            .project_bound_relations(1, "multi", &relations, None)
812            .is_err());
813        // Scoped to the uncovered source `b` is refused.
814        assert!(arena
815            .project_bound_relations(1, "multi", &relations, Some(&SourceId::new("b").unwrap()))
816            .is_err());
817        // Scoped to the covered source `a` over an a-only view succeeds.
818        let a_only = CoordinatorRelationSet {
819            records: vec![relation("obs.a1", "a1", "a")],
820        };
821        let block = arena
822            .project_bound_relations(1, "multi", &a_only, Some(&SourceId::new("a").unwrap()))
823            .unwrap();
824        assert_eq!(block.shape, vec![1, 2]);
825    }
826
827    #[test]
828    fn manifest_carries_shape_and_fingerprint() {
829        let store = NdTensorStore::from_inputs(vec![rgb_input()]).unwrap();
830        let manifests = store.manifests().unwrap();
831        assert_eq!(manifests.len(), 1);
832        assert_eq!(manifests[0].shape, vec![3, 2, 2]);
833        assert_eq!(manifests[0].data_bytes, 12);
834        assert_eq!(manifests[0].element_bytes, 1);
835        assert_eq!(manifests[0].tensor_fingerprint.len(), 64);
836    }
837
838    fn fp(input: NdTensorInput) -> String {
839        NdTensor::from_input(input).unwrap().fingerprint().unwrap()
840    }
841
842    #[test]
843    fn tensor_fingerprint_is_64_lowercase_hex() {
844        let fingerprint = fp(rgb_input());
845        assert_eq!(fingerprint.len(), 64);
846        assert!(fingerprint
847            .chars()
848            .all(|c| c.is_ascii_hexdigit() && !c.is_ascii_uppercase()));
849    }
850
851    #[test]
852    fn tensor_fingerprint_is_deterministic_across_calls_and_clone() {
853        let tensor = NdTensor::from_input(rgb_input()).unwrap();
854        let once = tensor.fingerprint().unwrap();
855        assert_eq!(once, tensor.fingerprint().unwrap());
856        assert_eq!(once, tensor.clone().fingerprint().unwrap());
857    }
858
859    #[test]
860    fn tensor_fingerprint_changes_when_a_single_data_byte_flips() {
861        let baseline = fp(rgb_input());
862        let mut flipped = rgb_input();
863        flipped.data[0] ^= 0xFF;
864        assert_ne!(baseline, fp(flipped));
865    }
866
867    #[test]
868    fn tensor_fingerprint_changes_when_tensor_id_is_renamed() {
869        let baseline = fp(rgb_input());
870        let mut renamed = rgb_input();
871        renamed.tensor_id = "rgb_renamed".to_string();
872        assert_ne!(baseline, fp(renamed));
873    }
874
875    #[test]
876    fn tensor_fingerprint_changes_when_observation_ids_are_reordered() {
877        // Permute both the rows' bytes and their ids together so the tensor
878        // stays a valid permutation of the same logical rows.
879        let baseline = fp(rgb_input());
880        let mut reordered = rgb_input();
881        reordered.observation_ids.swap(0, 2);
882        // Each row is 4 bytes ([_,2,2] u8); swap row 0 and row 2.
883        let (head, tail) = reordered.data.split_at_mut(8);
884        head[0..4].swap_with_slice(&mut tail[0..4]);
885        assert_ne!(baseline, fp(reordered));
886    }
887
888    #[test]
889    fn tensor_fingerprint_distinguishes_transposed_shapes_with_identical_bytes() {
890        // Same 12 bytes, axis-0 size 3 either way; [3,2,2] vs [3,4] differ only
891        // in the trailing shape, which the rank+dims framing must separate.
892        let base = NdTensorInput {
893            tensor_id: "t".to_string(),
894            representation_id: RepresentationId::new("rgb_image").unwrap(),
895            container: "ndarray".to_string(),
896            dtype: NdTensorDType::U8,
897            shape: vec![3, 2, 2],
898            observation_ids: vec![
899                ObservationId::new("obs.s1").unwrap(),
900                ObservationId::new("obs.s2").unwrap(),
901                ObservationId::new("obs.s3").unwrap(),
902            ],
903            sample_ids: None,
904            data: (0u8..12).collect(),
905            row_presence: None,
906        };
907        let mut reshaped = base.clone();
908        reshaped.shape = vec![3, 4];
909        assert_ne!(fp(base), fp(reshaped));
910    }
911
912    #[test]
913    fn tensor_fingerprint_distinguishes_dtype_with_identical_bytes() {
914        // Four bytes, axis-0 size 1: one I32 element vs four U8 elements share
915        // the same payload bytes; the dtype tag must separate them.
916        let as_i32 = NdTensorInput {
917            tensor_id: "t".to_string(),
918            representation_id: RepresentationId::new("rgb_image").unwrap(),
919            container: "ndarray".to_string(),
920            dtype: NdTensorDType::I32,
921            shape: vec![1, 1],
922            observation_ids: vec![ObservationId::new("obs.s1").unwrap()],
923            sample_ids: None,
924            data: vec![1, 2, 3, 4],
925            row_presence: None,
926        };
927        let mut as_u8 = as_i32.clone();
928        as_u8.dtype = NdTensorDType::U8;
929        as_u8.shape = vec![1, 4];
930        assert_ne!(fp(as_i32), fp(as_u8));
931    }
932
933    #[test]
934    fn tensor_fingerprint_hashes_data_bytes_verbatim_in_declared_order() {
935        // Conformance: the layer hashes element bytes as-is and never
936        // byte-swaps, so the little-endian input contract is what makes the
937        // fingerprint platform-independent. An f32 = 1.0 is `00 00 80 3f` LE;
938        // feeding the byte-reversed (big-endian) arrangement of the same
939        // logical value must change the fingerprint, proving order is honored.
940        let f32_one_le: [u8; 4] = 1.0f32.to_le_bytes(); // 00 00 80 3f
941        let le = NdTensorInput {
942            tensor_id: "t".to_string(),
943            representation_id: RepresentationId::new("hyperspectral").unwrap(),
944            container: "ndarray".to_string(),
945            dtype: NdTensorDType::F32,
946            shape: vec![1, 1],
947            observation_ids: vec![ObservationId::new("obs.s1").unwrap()],
948            sample_ids: None,
949            data: f32_one_le.to_vec(),
950            row_presence: None,
951        };
952        let mut be = le.clone();
953        let mut reversed = f32_one_le;
954        reversed.reverse(); // 3f 80 00 00 — the BE arrangement
955        be.data = reversed.to_vec();
956        // Two LE builds match; the BE arrangement differs.
957        assert_eq!(fp(le.clone()), fp(le.clone()));
958        assert_ne!(fp(le), fp(be));
959    }
960
961    #[test]
962    fn tensor_fingerprint_distinguishes_row_presence_states() {
963        // No presence vs all-present-present mask are different logical states
964        // and must fingerprint differently (the absent-tag vs the framed mask).
965        let baseline = fp(rgb_input());
966        let mut with_presence = rgb_input();
967        with_presence.row_presence = Some(vec![true, true, true]);
968        let present_fp = fp(with_presence);
969        assert_ne!(baseline, present_fp);
970
971        let mut one_absent = rgb_input();
972        one_absent.row_presence = Some(vec![true, false, true]);
973        assert_ne!(present_fp, fp(one_absent));
974    }
975
976    #[test]
977    #[ignore = "perf sanity probe; run with --release --ignored --nocapture"]
978    fn tensor_fingerprint_large_payload_under_500ms() {
979        // Stream a ~12 MB u8-equivalent / ~12.7 M-element f32 payload; the
980        // streamed hash must stay well under the 500 ms budget (the old
981        // serde_json path expanded every byte to a JSON integer first). The
982        // budget is asserted only in optimized builds — an unoptimized
983        // `cargo test` runs sha2 with overflow checks and no SIMD, so its wall
984        // time is not "native" performance. The number is always printed.
985        let rows = 3021usize;
986        let cols = 1050usize;
987        let element_size = NdTensorDType::F32.element_size();
988        let data = vec![0x3Cu8; rows * cols * element_size];
989        let input = NdTensorInput {
990            tensor_id: "big".to_string(),
991            representation_id: RepresentationId::new("hyperspectral").unwrap(),
992            container: "ndarray".to_string(),
993            dtype: NdTensorDType::F32,
994            shape: vec![rows, cols],
995            observation_ids: (0..rows)
996                .map(|r| ObservationId::new(format!("obs.{r}")).unwrap())
997                .collect(),
998            sample_ids: None,
999            data,
1000            row_presence: None,
1001        };
1002        let tensor = NdTensor::from_input(input).unwrap();
1003        let start = std::time::Instant::now();
1004        let fingerprint = tensor.fingerprint().unwrap();
1005        let elapsed = start.elapsed();
1006        println!(
1007            "nd tensor fingerprint({rows}x{cols} f32) = {:.3} ms (fp={fingerprint})",
1008            elapsed.as_secs_f64() * 1e3
1009        );
1010        assert_eq!(fingerprint.len(), 64);
1011        if !cfg!(debug_assertions) {
1012            assert!(
1013                elapsed.as_millis() < 500,
1014                "tensor fingerprint took {} ms (>= 500 ms budget)",
1015                elapsed.as_millis()
1016            );
1017        }
1018    }
1019}