candding 0.1.0

//! Reading a NumPy `.npy` file (`Qdrant/minicoil-v1`'s head matrix) into a candle tensor, per numpy's format spec.
//! Only `f4`, `f8`, `i2`, `i4`, `i8`, `u1` and `u4` are read; `f2` would need the `half` crate and is refused.

mod header;

use std::fs::File;
use std::io::Read;
use std::path::Path;

use candle_core::DType;
use candle_core::Device;
use candle_core::Tensor;
use header::Header;

use crate::error::CanddingError;
use crate::error::Result;

const MAGIC: [u8; 6] = *b"\x93NUMPY";

/// Reads `path` as a CPU tensor in the array's dtype; a file that disagrees with its own header is refused.
pub fn read(path: &Path) -> Result<Tensor> {
  let mut file = File::open(path)?;
  let header = read_header(path, &mut file)?;
  let elements: usize = header.shape.iter().product();
  let expected_bytes = elements * header.itemsize;
  let mut data = Vec::new();
  file.read_to_end(&mut data)?;
  if data.len() != expected_bytes {
    return Err(CanddingError::NpyShapeMismatch {
      file: path.to_path_buf(),
      shape: header.shape,
      expected_bytes,
      actual_bytes: data.len(),
    });
  }
  tensor_from_bytes(&header, &data)
}

/// The dtype `path` stores its array in, read from the header alone.
pub(crate) fn stored_dtype(path: &Path) -> Result<DType> {
  let mut file = File::open(path)?;
  Ok(read_header(path, &mut file)?.dtype)
}

/// Magic, version, header length and dictionary; read piece by piece so a truncated file names the short piece.
fn read_header(path: &Path, file: &mut File) -> Result<Header> {
  let mut prefix = [0u8; 8];
  read_exact_named(
    path,
    file,
    &mut prefix,
    "the magic bytes and format version",
  )?;
  if prefix[..6] != MAGIC {
    return Err(header_error(path, "missing the `\\x93NUMPY` magic bytes"));
  }
  let (major, minor) = (prefix[6], prefix[7]);
  let header_len_size = match major {
    1 => 2,
    2 => 4,
    _ => {
      return Err(header_error(
        path,
        &format!(
          "format version {major}.{minor} is not supported; only versions 1.x and 2.x are read"
        ),
      ));
    }
  };
  let mut header_len_bytes = [0u8; 4];
  read_exact_named(
    path,
    file,
    &mut header_len_bytes[..header_len_size],
    "the header-length field",
  )?;
  let header_len = if header_len_size == 2 {
    u16::from_le_bytes([header_len_bytes[0], header_len_bytes[1]]) as usize
  } else {
    u32::from_le_bytes(header_len_bytes) as usize
  };
  let mut dict_bytes = vec![0u8; header_len];
  read_exact_named(path, file, &mut dict_bytes, "the header dictionary")?;
  let dict = String::from_utf8(dict_bytes)
    .map_err(|_| header_error(path, "the header dictionary is not valid ASCII/UTF-8"))?;
  header::parse_header(path, &dict)
}

fn read_exact_named(path: &Path, file: &mut File, buf: &mut [u8], what: &str) -> Result<()> {
  file
    .read_exact(buf)
    .map_err(|_| header_error(path, &format!("truncated before {what} could be read")))
}

fn header_error(path: &Path, reason: &str) -> CanddingError {
  CanddingError::NpyHeader {
    file: path.to_path_buf(),
    reason: reason.to_string(),
  }
}

/// Decodes `data` as little-endian `header.dtype` values; `read` has already checked its length against the shape.
fn tensor_from_bytes(header: &Header, data: &[u8]) -> Result<Tensor> {
  let device = Device::Cpu;
  let shape = header.shape.as_slice();
  Ok(match header.dtype {
    DType::F32 => Tensor::from_vec(
      decode_le::<4, f32>(data, f32::from_le_bytes),
      shape,
      &device,
    )?,
    DType::F64 => Tensor::from_vec(
      decode_le::<8, f64>(data, f64::from_le_bytes),
      shape,
      &device,
    )?,
    DType::I16 => Tensor::from_vec(
      decode_le::<2, i16>(data, i16::from_le_bytes),
      shape,
      &device,
    )?,
    DType::I32 => Tensor::from_vec(
      decode_le::<4, i32>(data, i32::from_le_bytes),
      shape,
      &device,
    )?,
    DType::I64 => Tensor::from_vec(
      decode_le::<8, i64>(data, i64::from_le_bytes),
      shape,
      &device,
    )?,
    DType::U8 => Tensor::from_vec(data.to_vec(), shape, &device)?,
    DType::U32 => Tensor::from_vec(
      decode_le::<4, u32>(data, u32::from_le_bytes),
      shape,
      &device,
    )?,
    other => {
      unreachable!(
        "header::resolve_dtype only returns dtypes SUPPORTED_DTYPES lists, never {other:?}"
      )
    }
  })
}

fn decode_le<const N: usize, T>(data: &[u8], from_bytes: fn([u8; N]) -> T) -> Vec<T> {
  data
    .as_chunks::<N>()
    .0
    .iter()
    .copied()
    .map(from_bytes)
    .collect()
}

#[cfg(test)]
mod tests;