chore: bump dlpark to 0.8

Cargo.toml

@@ -41,7 +41,10 @@ arrow-buffer = "59"
 arrow-schema = "59"
 async-trait = { version = "0.1", optional = true }
 bytes = "1.12.0"
-dlpark = { git = "https://github.com/kylebarron/dlpark", rev = "31c6f49c064e634326c97172d39a00acecd854b6", features = [
+# Fork of dlpark 0.8.0 that gates bindgen's 64-bit layout tests so the crate
+# compiles for 32-bit targets (wasm32-unknown-emscripten). Upstreamed to
+# SunDoge/dlpark; revert to the crates.io release once a fixed version ships.
+dlpark = { git = "https://github.com/kylebarron/dlpark", rev = "f6887bf8b367d98d079710be9659c4177ba899f8", features = [
     "pyo3",
 ] }
 icechunk = { version = "2.0.0", default-features = false, optional = true }

src/data/tensor.rs

@@ -2,9 +2,9 @@ use std::ffi::{c_int, c_void};
 use std::sync::Arc;
 
 use bytes::Bytes;
-use dlpark::SafeManagedTensor;
-use dlpark::ffi::{Device, DeviceType};
-use dlpark::traits::{RowMajorCompactLayout, TensorLike};
+use dlpark::ffi::{DLDataType, DLDataTypeCode, DLDevice, DLDeviceType};
+use dlpark::metadata::CopiedSlice;
+use dlpark::{Builder, legacy};
 use pyo3::exceptions::{PyNotImplementedError, PyValueError};
 use pyo3::ffi;
 use pyo3::prelude::*;
@@ -14,7 +14,7 @@ use zarrs::array::DataType;
 
 use crate::data::buffer_protocol::PyTensorBuffer;
 use crate::dtype::PyDataType;
-use crate::error::{ZarristaError, ZarristaResult};
+use crate::error::ZarristaResult;
 
 /// Fixed-width, dense decoded data.
 ///
@@ -118,13 +118,39 @@ impl PyTensor {
     fn __dlpack__<'py>(
         &self,
         _kwargs: Option<Bound<'py, PyDict>>,
-    ) -> ZarristaResult<SafeManagedTensor> {
-        SafeManagedTensor::new(self.clone())
+    ) -> ZarristaResult<legacy::Dlpack> {
+        let dtype = self.dlpack_data_type()?;
+        let shape = self
+            .shape
+            .iter()
+            .map(|s| i64::try_from(*s).expect("overflow converting shape to i64"))
+            .collect::<Vec<_>>();
+        let strides = row_major_compact_strides(&shape);
+
+        // The boxed `Bytes` handed to the builder is what keeps the buffer alive: dlpark stores it
+        // as the managed tensor's `manager_ctx` and drops it from the deleter, which a consumer may
+        // run on any thread and long after this `PyTensor` is gone.
+        let data = self.bytes.as_ptr().cast::<c_void>().cast_mut();
+        let builder = Builder::new(
+            Box::new(self.bytes.clone()),
+            CopiedSlice::new(shape, strides),
+        );
+        // SAFETY: `data` points at the start of the `Bytes` allocation moved into the context, so
+        // it stays valid until the deleter runs. The shape is the tensor's own shape, the strides
+        // are row-major compact, and `dtype` is the element type those bytes were decoded as, so
+        // together they describe exactly the initialized elements of the buffer.
+        let builder = unsafe { builder.data(data) };
+
+        Ok(builder
+            .device(DLDevice::CPU)
+            .dtype(dtype)
+            .try_build()
+            .map_err(|err| PyValueError::new_err(err.to_string()))?)
     }
 
     /// The DLPack device this data lives on: `(device_type, device_id)`. Always CPU.
-    fn __dlpack_device__(&self) -> (i32, i32) {
-        (DeviceType::Cpu as i32, 0)
+    fn __dlpack_device__(&self) -> (u32, i32) {
+        (DLDeviceType::CPU.0, 0)
     }
 
     /// Export as a PEP 3118 buffer: an N-dimensional, typed, read-only,
@@ -148,64 +174,52 @@ impl PyTensor {
     unsafe fn __releasebuffer__(&self, _view: *mut ffi::Py_buffer) {}
 }
 
-impl TensorLike<RowMajorCompactLayout> for PyTensor {
-    type Error = ZarristaError;
-
-    fn data_ptr(&self) -> *mut c_void {
-        self.bytes.as_ptr().cast::<c_void>().cast_mut()
-    }
-
-    fn memory_layout(&self) -> RowMajorCompactLayout {
-        let shape = self
-            .shape()
-            .iter()
-            .map(|s| i64::try_from(*s).expect("overflow converting shape to i64"))
-            .collect();
-        RowMajorCompactLayout::new(shape)
-    }
-
-    fn byte_offset(&self) -> u64 {
-        0
-    }
-
-    fn device(&self) -> Result<Device, Self::Error> {
-        Ok(Device::CPU)
-    }
-
-    fn data_type(&self) -> Result<dlpark::ffi::DataType, Self::Error> {
+impl PyTensor {
+    /// The DLPack element descriptor for this tensor's zarr data type.
+    fn dlpack_data_type(&self) -> ZarristaResult<DLDataType> {
         use zarrs::array::data_type::*;
 
         let dtype = &self.data_type;
-        if dtype.is::<BoolDataType>() {
-            Ok(dlpark::ffi::DataType::BOOL)
+        let (code, bits) = if dtype.is::<BoolDataType>() {
+            (DLDataTypeCode::BOOL, 8)
         } else if dtype.is::<Int8DataType>() {
-            Ok(dlpark::ffi::DataType::I8)
+            (DLDataTypeCode::INT, 8)
         } else if dtype.is::<Int16DataType>() {
-            Ok(dlpark::ffi::DataType::I16)
+            (DLDataTypeCode::INT, 16)
         } else if dtype.is::<Int32DataType>() {
-            Ok(dlpark::ffi::DataType::I32)
+            (DLDataTypeCode::INT, 32)
         } else if dtype.is::<Int64DataType>() {
-            Ok(dlpark::ffi::DataType::I64)
+            (DLDataTypeCode::INT, 64)
         } else if dtype.is::<UInt8DataType>() {
-            Ok(dlpark::ffi::DataType::U8)
+            (DLDataTypeCode::UINT, 8)
         } else if dtype.is::<UInt16DataType>() {
-            Ok(dlpark::ffi::DataType::U16)
+            (DLDataTypeCode::UINT, 16)
         } else if dtype.is::<UInt32DataType>() {
-            Ok(dlpark::ffi::DataType::U32)
+            (DLDataTypeCode::UINT, 32)
         } else if dtype.is::<UInt64DataType>() {
-            Ok(dlpark::ffi::DataType::U64)
+            (DLDataTypeCode::UINT, 64)
         } else if dtype.is::<Float16DataType>() {
-            Ok(dlpark::ffi::DataType::F16)
+            (DLDataTypeCode::FLOAT, 16)
         } else if dtype.is::<Float32DataType>() {
-            Ok(dlpark::ffi::DataType::F32)
+            (DLDataTypeCode::FLOAT, 32)
         } else if dtype.is::<Float64DataType>() {
-            Ok(dlpark::ffi::DataType::F64)
+            (DLDataTypeCode::FLOAT, 64)
         } else if dtype.is::<BFloat16DataType>() {
-            Ok(dlpark::ffi::DataType::BF16)
+            (DLDataTypeCode::BFLOAT, 16)
         } else {
-            Err(PyValueError::new_err("Unsupported data type in dlpack").into())
-        }
+            return Err(PyValueError::new_err("Unsupported data type in dlpack").into());
+        };
+        Ok(DLDataType::scalar(code, bits))
+    }
+}
+
+/// Strides, in elements, for a row-major contiguous buffer of the given shape.
+fn row_major_compact_strides(shape: &[i64]) -> Vec<i64> {
+    let mut strides = vec![1; shape.len()];
+    for axis in (0..shape.len().saturating_sub(1)).rev() {
+        strides[axis] = strides[axis + 1] * shape[axis + 1];
     }
+    strides
 }
 
 /// Fixed-width data with a validity mask. Skeleton.