ITADN

Add support for structured dtypes to `zarr3` driver, open zarr 2 and 3 structs as void

#264Pull RequestBrianMichell 创建于 2025-12-04
B
BrianMichellcommented
Resolves #241 This PR implements support for structured data in the `zarr3` driver as well as the ability to interact with the full array as raw bytes to avoid needing to manage multiple Stores in client code. ~~I am also happy to revisit the question of #125 for a more general Zarr support for this behavior.~~ I have also implemented the `open_as_void` feature in the `zarr` driver for feature parity between the formats. Below is an example workflow that this usecase would support. The Python snippet will generate a structured array of several different elements and types. A real world example of this is creating an array to store the [SEG-Y 240-byte Trace Headers](https://seg.org/wp-content/uploads/2025/11/seg_y_rev1.pdf#page=15). <details> <summary>Click here to expand Python example.</summary> ```python import numpy as np import zarr store = "foo.zarr" np_dtype = np.dtype( [ ("field_1", "<i4"), ("field_2", ">i4"), # More theoretical fields here... ("field_n", "<f4"), ] ) z = zarr.create_array( store=store, shape=(128, 128), dtype=np_dtype, chunks=(32, 32), ) arr = np.zeros((128, 128), dtype=np_dtype) f1 = np.arange(128, dtype="<i4") # NOTE: We populate f2 as little endian for the demonstration f2 = np.arange(128, dtype="<i4") fn = np.arange(128, dtype="<f4") / 10 arr["field_1"][:] = f1[:, None] arr["field_2"][:] = f2[:, None] arr["field_n"][:] = fn[:, None] z[:] = arr ``` </details> The usecase continues with an HPC C++ application consuming either a single field from the structured dtype, or the entirety of the dtype to perform its intended workload. To keep this example code simple, I only demonstrate consuming the same field out of the raw data, however this could reasonably be extended to 70+ fields per element if writing a trivial data loader. <details> <summary>Click here to expand C++ example.</summary> ```cpp // Minimal example demonstrating structured Zarr v3 data reading. // This example reads the 'foo.zarr' output generated by generate_struct.py // // Usage: // ./read_foo_zarr --zarr_path=/path/to/foo.zarr #include <stdint.h> #include <iostream> #include <string> #include <cstring> #include "absl/flags/flag.h" #include "absl/flags/parse.h" #include "absl/status/status.h" #include <nlohmann/json.hpp> #include "tensorstore/array.h" #include "tensorstore/context.h" #include "tensorstore/data_type.h" #include "tensorstore/index.h" #include "tensorstore/open.h" #include "tensorstore/open_mode.h" #include "tensorstore/spec.h" #include "tensorstore/tensorstore.h" #include "tensorstore/util/result.h" #include "tensorstore/util/status.h" ABSL_FLAG(std::string, zarr_path, "/foo.zarr", "Path to the foo.zarr directory"); namespace { using ::tensorstore::Index; // Helper function to read and display data from a tensorstore // T: the logical value type you want to interpret the data as // offset_bytes: // -1 (default): interpret array.data() as a T* over the logical elements // >=0 : treat array.data() as a byte buffer and start reading T // values at `offset_bytes`, stepping by sizeof(T). template <typename T> absl::Status ReadAndDisplayData(const tensorstore::TensorStore<>& store, const std::string& description, Index offset_bytes = -1) { std::cout << "\n=== " << description << " ===" << std::endl; // Get array information auto domain = store.domain(); std::cout << "Domain: " << domain << std::endl; std::cout << "Data type: " << store.dtype() << std::endl; // Read all data TENSORSTORE_ASSIGN_OR_RETURN( auto array, tensorstore::Read<tensorstore::zero_origin>(store).result()); std::cout << "Successfully read array with " << array.num_elements() << " elements" << std::endl; auto shape = domain.shape(); if (shape.size() < 2) { std::cout << "Rank < 2, skipping pretty-print of first 2D block." << std::endl; return absl::OkStatus(); } Index rows = std::min(shape[0], Index{5}); Index cols = std::min(shape[1], Index{5}); std::cout << "First " << rows << "x" << cols << " elements"; if (offset_bytes >= 0) { std::cout << " starting at byte offset " << offset_bytes; } std::cout << " interpreted as " << sizeof(T) * 8 << "-bit values:" << std::endl; // No offset: “normal” interpretation as T[] if (offset_bytes < 0) { const T* data = reinterpret_cast<const T*>(array.data()); for (Index i = 0; i < rows; ++i) { for (Index j = 0; j < cols; ++j) { Index idx = i * shape[1] + j; if (idx >= array.num_elements()) break; std::cout << data[idx] << "\t"; } std::cout << std::endl; } return absl::OkStatus(); } // Offset mode: interpret as raw bytes and then per-record struct. const auto* bytes = reinterpret_cast<const uint8_t*>(array.data()); // Total bytes in buffer (works for any dtype). std::size_t dtype_size = store.dtype().size(); Index total_bytes = static_cast<Index>(array.num_elements()) * static_cast<Index>(dtype_size == 0 ? 1 : dtype_size); if (total_bytes == 0) { std::cout << "[empty buffer]" << std::endl; return absl::OkStatus(); } // Determine "record" size in bytes. // // For open_as_void on a structured array, we get something like: // shape = [N0, N1, record_bytes] // dtype = byte // // In that case, treat the last dimension as the record layout. Index record_bytes = static_cast<Index>(store.dtype().size()); if (store.dtype() == tensorstore::dtype_v<tensorstore::dtypes::byte_t> && shape.size() >= 3) { record_bytes = shape.back(); // e.g. 12 bytes in your example } if (offset_bytes < 0 || offset_bytes >= record_bytes) { std::cout << "[offset outside record size (" << record_bytes << " bytes); nothing to display]" << std::endl; return absl::OkStatus(); } for (Index i = 0; i < rows; ++i) { for (Index j = 0; j < cols; ++j) { Index record_index_2d = i * shape[1] + j; Index base_offset = record_index_2d * record_bytes + offset_bytes; if (base_offset + static_cast<Index>(sizeof(T)) > total_bytes) { std::cout << "[out]\t"; continue; } T value{}; std::memcpy(&value, bytes + base_offset, sizeof(T)); std::cout << value << "\t"; } std::cout << std::endl; } return absl::OkStatus(); } absl::Status Run(const std::string& zarr_path) { std::cout << "=== Structured Zarr v3 Example ===" << std::endl; std::cout << "Reading from: " << zarr_path << std::endl; auto context = tensorstore::Context::Default(); // Open the structured Zarr v3 array with field access ::nlohmann::json spec = ::nlohmann::json::object(); spec["driver"] = "zarr3"; spec["kvstore"] = ::nlohmann::json::object(); spec["kvstore"]["driver"] = "file"; spec["kvstore"]["path"] = zarr_path + "/"; spec["field"] = "field_2"; auto open_result = tensorstore::Open(spec, context, tensorstore::OpenMode::open, tensorstore::ReadWriteMode::read).result(); if (!open_result.ok()) { std::cout << "Failed to open structured array: " << open_result.status() << std::endl; return open_result.status(); } auto store = std::move(open_result).value(); TENSORSTORE_RETURN_IF_ERROR(ReadAndDisplayData<int32_t>(store, "Structured Array")); ::nlohmann::json void_spec = spec; void_spec["open_as_void"] = true; void_spec.erase("field"); auto void_open_result = tensorstore::Open(void_spec, context, tensorstore::OpenMode::open, tensorstore::ReadWriteMode::read).result(); if (!void_open_result.ok()) { std::cout << "Failed to open with open_as_void: " << void_open_result.status() << std::endl; return void_open_result.status(); } auto void_store = std::move(void_open_result).value(); constexpr Index kFieldOffsetBytes = 4; TENSORSTORE_RETURN_IF_ERROR(ReadAndDisplayData<int32_t>(void_store, "Raw Bytes (open_as_void)", kFieldOffsetBytes)); return absl::OkStatus(); } } // namespace int main(int argc, char** argv) { absl::ParseCommandLine(argc, argv); std::string zarr_path = absl::GetFlag(FLAGS_zarr_path); if (zarr_path.empty()) { std::cerr << "Error: --zarr_path is required" << std::endl; return 1; } auto status = Run(zarr_path); if (!status.ok()) { std::cerr << "\nFinal status: " << status << std::endl; return 1; } return 0; } ``` </details>
合并状态:未合并 关闭于 2026-01-05 4 条评论