tensor_bridge.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: AGPL-3.0-or-later
// Copyright 2026 Johnson Ogundeji
 
 
#pragma once
 
// ---------------------------------------------------------------------------
// tensor_bridge.hpp — Zero-Copy Tensor Bridge for SignetStack Signet Forge
//
// Maps Parquet column data directly into ML-framework-compatible tensor views
// without copying. Provides:
//
//   TensorDataType  — enum mapping to common ML framework element types
//   TensorShape     — N-dimensional shape descriptor
//   TensorView      — non-owning, zero-copy view into contiguous memory
//   OwnedTensor     — owning tensor with heap-allocated storage
//   ColumnToTensor  — Parquet column data -> tensor conversion
//   BatchTensorBuilder — multi-column feature batch assembly
//
// Header-only. Part of the signet::forge AI module.
// ---------------------------------------------------------------------------
 
#include "signet/types.hpp"
#include "signet/error.hpp"
 
#include <algorithm>
#include <cassert>
#include <cstddef>
#include <cstdlib>
#include <stdexcept>
#include <cstdint>
#include <cstring>
#include <limits>
#include <memory>
#include <new>
#include <numeric>
#include <string>
#include <type_traits>
#include <vector>
 
#ifdef _WIN32
#include <malloc.h>
#endif
 
namespace signet::forge {
 
namespace detail {
 
template <typename T, std::size_t Alignment>
class AlignedAllocator {
public:
    using value_type = T;
    using size_type = std::size_t;
    using difference_type = std::ptrdiff_t;
    using propagate_on_container_move_assignment = std::true_type;
    using is_always_equal = std::true_type;
 
    template <typename U>
    struct rebind {
        using other = AlignedAllocator<U, Alignment>;
    };
 
    AlignedAllocator() noexcept = default;
 
    template <typename U>
    AlignedAllocator(const AlignedAllocator<U, Alignment>&) noexcept {}
 
    [[nodiscard]] T* allocate(std::size_t n) {
        static_assert(Alignment >= alignof(void*), "alignment must satisfy allocator requirements");
        static_assert((Alignment & (Alignment - 1)) == 0, "alignment must be a power of two");
        if (n == 0) return nullptr;
        if (n > (std::numeric_limits<std::size_t>::max)() / sizeof(T)) {
            throw std::bad_alloc();
        }
 
        void* ptr = nullptr;
        const std::size_t bytes = n * sizeof(T);
#ifdef _WIN32
        ptr = _aligned_malloc(bytes, Alignment);
        if (!ptr) throw std::bad_alloc();
#else
        if (::posix_memalign(&ptr, Alignment, bytes) != 0) {
            throw std::bad_alloc();
        }
#endif
        return static_cast<T*>(ptr);
    }
 
    void deallocate(T* ptr, std::size_t) noexcept {
#ifdef _WIN32
        _aligned_free(ptr);
#else
        std::free(ptr);
#endif
    }
 
    template <typename U>
    [[nodiscard]] bool operator==(const AlignedAllocator<U, Alignment>&) const noexcept {
        return true;
    }
 
    template <typename U>
    [[nodiscard]] bool operator!=(const AlignedAllocator<U, Alignment>&) const noexcept {
        return false;
    }
};
 
template <typename T>
[[nodiscard]] inline bool is_pointer_aligned(const void* ptr) noexcept {
    if (ptr == nullptr) return false;
    return (reinterpret_cast<std::uintptr_t>(ptr) % alignof(T)) == 0;
}
 
template <typename T>
[[nodiscard]] inline T* aligned_ptr(void* ptr) noexcept {
    return is_pointer_aligned<T>(ptr) ? static_cast<T*>(ptr) : nullptr;
}
 
template <typename T>
[[nodiscard]] inline const T* aligned_ptr(const void* ptr) noexcept {
    return is_pointer_aligned<T>(ptr) ? static_cast<const T*>(ptr) : nullptr;
}
 
template <typename T>
[[nodiscard]] inline T* aligned_ptr_at(void* base, std::size_t offset) noexcept {
    auto* ptr = static_cast<std::uint8_t*>(base) + offset;
    return aligned_ptr<T>(ptr);
}
 
template <typename T>
[[nodiscard]] inline const T* aligned_ptr_at(const void* base, std::size_t offset) noexcept {
    auto* ptr = static_cast<const std::uint8_t*>(base) + offset;
    return aligned_ptr<T>(ptr);
}
 
} // namespace detail
 
// ===========================================================================
// TensorDataType — element data types for tensor storage
// ===========================================================================
 
enum class TensorDataType : int32_t {
    FLOAT32 = 0,  
    FLOAT64 = 1,  
    INT32   = 2,  
    INT64   = 3,  
    INT8    = 4,  
    UINT8   = 5,  
    INT16   = 6,  
    FLOAT16 = 7,  
    BOOL    = 8   
};
 
// ===========================================================================
// tensor_element_size — bytes per element for a given TensorDataType
// ===========================================================================
 
inline constexpr size_t tensor_element_size(TensorDataType dtype) noexcept {
    switch (dtype) {
        case TensorDataType::FLOAT32: return 4;
        case TensorDataType::FLOAT64: return 8;
        case TensorDataType::INT32:   return 4;
        case TensorDataType::INT64:   return 8;
        case TensorDataType::INT8:    return 1;
        case TensorDataType::UINT8:   return 1;
        case TensorDataType::INT16:   return 2;
        case TensorDataType::FLOAT16: return 2;
        case TensorDataType::BOOL:    return 1;
    }
    return 0; // unreachable
}
 
inline const char* tensor_dtype_name(TensorDataType dtype) noexcept {
    switch (dtype) {
        case TensorDataType::FLOAT32: return "float32";
        case TensorDataType::FLOAT64: return "float64";
        case TensorDataType::INT32:   return "int32";
        case TensorDataType::INT64:   return "int64";
        case TensorDataType::INT8:    return "int8";
        case TensorDataType::UINT8:   return "uint8";
        case TensorDataType::INT16:   return "int16";
        case TensorDataType::FLOAT16: return "float16";
        case TensorDataType::BOOL:    return "bool";
    }
    return "unknown";
}
 
// ===========================================================================
// TensorShape — N-dimensional shape descriptor
// ===========================================================================
 
struct TensorShape {
    std::vector<int64_t> dims;  
 
    TensorShape() = default;
 
    explicit TensorShape(std::vector<int64_t> d) : dims(std::move(d)) {}
 
    TensorShape(std::initializer_list<int64_t> il) : dims(il) {}
 
    [[nodiscard]] int64_t num_elements() const noexcept {
        if (dims.empty()) return 1;
        int64_t product = 1;
        for (auto d : dims) {
            if (d <= 0) return -1; // error sentinel: non-positive dimension
            if (product > INT64_MAX / d) return -1; // overflow sentinel
            product *= d;
        }
        return product;
    }
 
    [[nodiscard]] size_t ndim() const noexcept { return dims.size(); }
 
    [[nodiscard]] bool is_scalar() const noexcept {
        return dims.empty() || (dims.size() == 1 && dims[0] == 1);
    }
 
    [[nodiscard]] bool is_vector() const noexcept { return dims.size() == 1; }
 
    [[nodiscard]] bool is_matrix() const noexcept { return dims.size() == 2; }
 
    [[nodiscard]] bool operator==(const TensorShape& other) const {
        return dims == other.dims;
    }
 
    [[nodiscard]] bool operator!=(const TensorShape& other) const {
        return dims != other.dims;
    }
};
 
// ===========================================================================
// TensorView — non-owning, zero-copy view into contiguous tensor memory
// ===========================================================================
 
class TensorView {
public:
    TensorView() = default;
 
    TensorView(void* data, TensorShape shape, TensorDataType dtype,
               size_t byte_stride = 0) noexcept
        : data_(data)
        , shape_(std::move(shape))
        , dtype_(dtype)
        , byte_stride_(byte_stride) {}
 
    TensorView(const void* data, TensorShape shape, TensorDataType dtype,
               size_t byte_stride = 0) noexcept
        : data_(const_cast<void*>(data))
        , shape_(std::move(shape))
        , dtype_(dtype)
        , byte_stride_(byte_stride) {}
 
    // -- Raw data access ------------------------------------------------------
 
    [[nodiscard]] void* data() noexcept { return data_; }
    [[nodiscard]] const void* data() const noexcept { return data_; }
 
    template <typename T>
    [[nodiscard]] T* typed_data() noexcept {
        return detail::aligned_ptr<T>(data_);
    }
 
    template <typename T>
    [[nodiscard]] const T* typed_data() const noexcept {
        return detail::aligned_ptr<T>(data_);
    }
 
    // -- Shape and type info --------------------------------------------------
 
    [[nodiscard]] const TensorShape& shape() const noexcept { return shape_; }
    [[nodiscard]] TensorDataType dtype() const noexcept { return dtype_; }
 
    [[nodiscard]] size_t element_size() const noexcept {
        return tensor_element_size(dtype_);
    }
 
    [[nodiscard]] int64_t num_elements() const noexcept {
        return shape_.num_elements();
    }
 
    [[nodiscard]] size_t byte_size() const noexcept {
        const int64_t n = num_elements();
        if (n <= 0) return 0;
        return static_cast<size_t>(n) * element_size();
    }
 
    [[nodiscard]] size_t effective_byte_stride() const noexcept {
        if (byte_stride_ != 0) return byte_stride_;
        // Dense stride: product of dims[1..] * element_size
        if (shape_.ndim() <= 1) return element_size();
        size_t inner_size = element_size();
        for (size_t i = 1; i < shape_.ndim(); ++i) {
            inner_size *= static_cast<size_t>(shape_.dims[i]);
        }
        return inner_size;
    }
 
    // -- Typed element accessors ----------------------------------------------
 
    template <typename T>
    [[nodiscard]] T& at(int64_t i) {
        if (data_ == nullptr || i < 0 || i >= num_elements())
            throw std::out_of_range("TensorView::at(i): index out of range");
        if (byte_stride_ != 0) {
            auto* elem = detail::aligned_ptr_at<T>(data_, static_cast<size_t>(i) * byte_stride_);
            if (elem == nullptr)
                throw std::runtime_error("TensorView::at(i): misaligned tensor access");
            return *elem;
        }
        auto* ptr = typed_data<T>();
        if (ptr == nullptr)
            throw std::runtime_error("TensorView::at(i): misaligned tensor access");
        return ptr[i];
    }
 
    template <typename T>
    [[nodiscard]] const T& at(int64_t i) const {
        if (data_ == nullptr || i < 0 || i >= num_elements())
            throw std::out_of_range("TensorView::at(i): index out of range");
        if (byte_stride_ != 0) {
            const auto* elem = detail::aligned_ptr_at<T>(data_, static_cast<size_t>(i) * byte_stride_);
            if (elem == nullptr)
                throw std::runtime_error("TensorView::at(i): misaligned tensor access");
            return *elem;
        }
        const auto* ptr = typed_data<T>();
        if (ptr == nullptr)
            throw std::runtime_error("TensorView::at(i): misaligned tensor access");
        return ptr[i];
    }
 
    template <typename T>
    [[nodiscard]] T& at(int64_t row, int64_t col) {
        if (data_ == nullptr || shape_.ndim() != 2 ||
            row < 0 || row >= shape_.dims[0] ||
            col < 0 || col >= shape_.dims[1])
            throw std::out_of_range("TensorView::at(row,col): index out of range");
        const int64_t cols = shape_.dims[1];
        if (byte_stride_ != 0) {
            auto* row_ptr = detail::aligned_ptr_at<T>(data_, static_cast<size_t>(row) * byte_stride_);
            if (row_ptr == nullptr)
                throw std::runtime_error("TensorView::at(row,col): misaligned tensor access");
            return row_ptr[col];
        }
        auto* ptr = typed_data<T>();
        if (ptr == nullptr)
            throw std::runtime_error("TensorView::at(row,col): misaligned tensor access");
        return ptr[row * cols + col];
    }
 
    template <typename T>
    [[nodiscard]] const T& at(int64_t row, int64_t col) const {
        if (data_ == nullptr || shape_.ndim() != 2 ||
            row < 0 || row >= shape_.dims[0] ||
            col < 0 || col >= shape_.dims[1])
            throw std::out_of_range("TensorView::at(row,col): index out of range");
        const int64_t cols = shape_.dims[1];
        if (byte_stride_ != 0) {
            const auto* row_ptr = detail::aligned_ptr_at<T>(data_, static_cast<size_t>(row) * byte_stride_);
            if (row_ptr == nullptr)
                throw std::runtime_error("TensorView::at(row,col): misaligned tensor access");
            return row_ptr[col];
        }
        const auto* ptr = typed_data<T>();
        if (ptr == nullptr)
            throw std::runtime_error("TensorView::at(row,col): misaligned tensor access");
        return ptr[row * cols + col];
    }
 
    // -- Predicates -----------------------------------------------------------
 
    [[nodiscard]] bool is_contiguous() const noexcept {
        return byte_stride_ == 0;
    }
 
    [[nodiscard]] bool is_valid() const noexcept {
        return data_ != nullptr;
    }
 
    // -- Subview and reshape --------------------------------------------------
 
    [[nodiscard]] TensorView slice(int64_t start, int64_t count) const {
        if (data_ == nullptr || shape_.ndim() < 1 ||
            start < 0 || count < 0 || start + count > shape_.dims[0])
            return TensorView{}; // return invalid view
 
        // Compute the byte offset to the start of the slice
        const size_t stride = effective_byte_stride();
        auto* base = static_cast<uint8_t*>(const_cast<void*>(data_));
        void* slice_data = base + static_cast<size_t>(start) * stride;
 
        // Build the new shape: replace dims[0] with count, keep the rest
        TensorShape new_shape;
        new_shape.dims = shape_.dims;
        new_shape.dims[0] = count;
 
        return TensorView(slice_data, std::move(new_shape), dtype_, byte_stride_);
    }
 
    [[nodiscard]] expected<TensorView> reshape(TensorShape new_shape) const {
        if (!is_contiguous()) {
            return Error{ErrorCode::INTERNAL_ERROR,
                         "cannot reshape a non-contiguous tensor view"};
        }
        if (new_shape.num_elements() != shape_.num_elements()) {
            return Error{ErrorCode::SCHEMA_MISMATCH,
                         "reshape: total elements mismatch ("
                         + std::to_string(shape_.num_elements()) + " vs "
                         + std::to_string(new_shape.num_elements()) + ")"};
        }
        return TensorView(data_, std::move(new_shape), dtype_, 0);
    }
 
private:
    void*          data_        = nullptr;
    TensorShape    shape_;
    TensorDataType dtype_       = TensorDataType::FLOAT32;
    size_t         byte_stride_ = 0;  // 0 = contiguous (densely packed)
};
 
// ===========================================================================
// OwnedTensor — heap-allocated, owning tensor
// ===========================================================================
 
class OwnedTensor {
public:
    OwnedTensor() = default;
 
    OwnedTensor(TensorShape shape, TensorDataType dtype)
        : shape_(std::move(shape))
        , dtype_(dtype) {
        const auto num_elements = shape_.num_elements();
        const auto element_size = tensor_element_size(dtype_);
        if (num_elements <= 0 || static_cast<size_t>(num_elements) > SIZE_MAX / element_size) {
            // Overflow or invalid shape — leave buffer empty (invalid tensor)
            return;
        }
        const size_t sz = static_cast<size_t>(num_elements) * element_size;
        buffer_.resize(sz, 0);
    }
 
    OwnedTensor(const void* data, TensorShape shape, TensorDataType dtype)
        : shape_(std::move(shape))
        , dtype_(dtype) {
        const auto num_elements = shape_.num_elements();
        const auto element_size = tensor_element_size(dtype_);
        if (num_elements <= 0 || static_cast<size_t>(num_elements) > SIZE_MAX / element_size) {
            return; // Overflow or invalid shape — leave buffer empty
        }
        const size_t sz = static_cast<size_t>(num_elements) * element_size;
        buffer_.resize(sz);
        if (data && sz > 0) {
            std::memcpy(buffer_.data(), data, sz);
        }
    }
 
    // Move semantics
    OwnedTensor(OwnedTensor&&) noexcept = default;
    OwnedTensor& operator=(OwnedTensor&&) noexcept = default;
 
    // No implicit copy — use clone()
    OwnedTensor(const OwnedTensor&) = delete;
    OwnedTensor& operator=(const OwnedTensor&) = delete;
 
    [[nodiscard]] OwnedTensor clone() const {
        OwnedTensor copy;
        copy.buffer_ = buffer_;
        copy.shape_  = shape_;
        copy.dtype_  = dtype_;
        return copy;
    }
 
    // -- View access ----------------------------------------------------------
 
    [[nodiscard]] TensorView view() {
        return TensorView(buffer_.data(), shape_, dtype_, 0);
    }
 
    [[nodiscard]] TensorView view() const {
        return TensorView(
            const_cast<uint8_t*>(buffer_.data()), shape_, dtype_, 0);
    }
 
    // -- Data access (forwarded from view) ------------------------------------
 
    [[nodiscard]] void* data() noexcept { return buffer_.data(); }
    [[nodiscard]] const void* data() const noexcept { return buffer_.data(); }
 
    template <typename T>
    [[nodiscard]] T* typed_data() noexcept {
        return detail::aligned_ptr<T>(buffer_.data());
    }
 
    template <typename T>
    [[nodiscard]] const T* typed_data() const noexcept {
        return detail::aligned_ptr<T>(buffer_.data());
    }
 
    [[nodiscard]] const TensorShape& shape() const noexcept { return shape_; }
    [[nodiscard]] TensorDataType dtype() const noexcept { return dtype_; }
 
    [[nodiscard]] size_t byte_size() const noexcept { return buffer_.size(); }
 
    [[nodiscard]] int64_t num_elements() const noexcept {
        return shape_.num_elements();
    }
 
    [[nodiscard]] bool is_valid() const noexcept { return !buffer_.empty(); }
 
private:
    using Buffer = std::vector<uint8_t,
        detail::AlignedAllocator<uint8_t, alignof(std::max_align_t)>>;
 
    Buffer           buffer_;
    TensorShape      shape_;
    TensorDataType   dtype_ = TensorDataType::FLOAT32;
};
 
// ===========================================================================
// ColumnToTensor — map Parquet column data to tensor representations
// ===========================================================================
 
class ColumnToTensor {
public:
    // -----------------------------------------------------------------------
    // Zero-copy path: wrap existing column data as a TensorView
    // -----------------------------------------------------------------------
 
    static inline expected<TensorView> wrap_column(
            const void* column_data,
            int64_t num_values,
            PhysicalType physical_type,
            int32_t type_length = -1) {
        if (!column_data || num_values <= 0) {
            return Error{ErrorCode::INTERNAL_ERROR,
                         "wrap_column: null data or non-positive count"};
        }
 
        switch (physical_type) {
            case PhysicalType::INT32:
                return TensorView(column_data,
                                  TensorShape{num_values},
                                  TensorDataType::INT32);
 
            case PhysicalType::INT64:
                return TensorView(column_data,
                                  TensorShape{num_values},
                                  TensorDataType::INT64);
 
            case PhysicalType::FLOAT:
                return TensorView(column_data,
                                  TensorShape{num_values},
                                  TensorDataType::FLOAT32);
 
            case PhysicalType::DOUBLE:
                return TensorView(column_data,
                                  TensorShape{num_values},
                                  TensorDataType::FLOAT64);
 
            case PhysicalType::FIXED_LEN_BYTE_ARRAY: {
                if (type_length <= 0) {
                    return Error{ErrorCode::SCHEMA_MISMATCH,
                                 "wrap_column: FIXED_LEN_BYTE_ARRAY requires "
                                 "positive type_length"};
                }
                // Expose as a 2D {num_values, type_length} uint8 view
                return TensorView(column_data,
                                  TensorShape{num_values,
                                              static_cast<int64_t>(type_length)},
                                  TensorDataType::UINT8);
            }
 
            case PhysicalType::BOOLEAN:
                // Parquet booleans are bit-packed; cannot zero-copy as a
                // byte-addressable tensor without unpacking.
                return Error{ErrorCode::UNSUPPORTED_TYPE,
                             "wrap_column: BOOLEAN columns require copy "
                             "(bit-packed, not byte-addressable)"};
 
            case PhysicalType::BYTE_ARRAY:
                return Error{ErrorCode::UNSUPPORTED_TYPE,
                             "wrap_column: BYTE_ARRAY (variable-length) "
                             "cannot be zero-copy wrapped as a tensor"};
 
            case PhysicalType::INT96:
                return Error{ErrorCode::UNSUPPORTED_TYPE,
                             "wrap_column: INT96 is deprecated and "
                             "not supported for tensor wrapping"};
        }
 
        return Error{ErrorCode::UNSUPPORTED_TYPE,
                     "wrap_column: unknown physical type"};
    }
 
    // -----------------------------------------------------------------------
    // Zero-copy path: wrap vector column data
    // -----------------------------------------------------------------------
 
    static inline expected<TensorView> wrap_vectors(
            const void* column_data,
            int64_t num_vectors,
            uint32_t dimension) {
        if (!column_data || num_vectors <= 0) {
            return Error{ErrorCode::INTERNAL_ERROR,
                         "wrap_vectors: null data or non-positive count"};
        }
        if (dimension == 0) {
            return Error{ErrorCode::SCHEMA_MISMATCH,
                         "wrap_vectors: dimension must be > 0"};
        }
 
        return TensorView(column_data,
                          TensorShape{num_vectors,
                                      static_cast<int64_t>(dimension)},
                          TensorDataType::FLOAT32);
    }
 
    // -----------------------------------------------------------------------
    // Copy path: read + convert column data into an OwnedTensor
    // -----------------------------------------------------------------------
 
    static inline expected<OwnedTensor> copy_column(
            const void* column_data,
            int64_t num_values,
            PhysicalType physical_type,
            TensorDataType target_dtype,
            int32_t type_length = -1) {
        if (!column_data || num_values <= 0) {
            return Error{ErrorCode::INTERNAL_ERROR,
                         "copy_column: null data or non-positive count"};
        }
 
        // For BYTE_ARRAY we cannot produce a dense tensor
        if (physical_type == PhysicalType::BYTE_ARRAY) {
            return Error{ErrorCode::UNSUPPORTED_TYPE,
                         "copy_column: BYTE_ARRAY (strings) cannot be "
                         "converted to a dense tensor"};
        }
        if (physical_type == PhysicalType::INT96) {
            return Error{ErrorCode::UNSUPPORTED_TYPE,
                         "copy_column: INT96 is deprecated and not supported"};
        }
 
        // First, try the zero-copy wrap to get a typed view of the source
        // For FIXED_LEN_BYTE_ARRAY, we handle specially below
        if (physical_type == PhysicalType::FIXED_LEN_BYTE_ARRAY) {
            if (type_length <= 0) {
                return Error{ErrorCode::SCHEMA_MISMATCH,
                             "copy_column: FIXED_LEN_BYTE_ARRAY requires "
                             "positive type_length"};
            }
            // Treat as flat bytes, then cast into the target dtype
            TensorView src(column_data,
                           TensorShape{num_values * static_cast<int64_t>(type_length)},
                           TensorDataType::UINT8);
            // If target is UINT8, just copy directly
            if (target_dtype == TensorDataType::UINT8) {
                OwnedTensor out(TensorShape{num_values,
                                            static_cast<int64_t>(type_length)},
                                target_dtype);
                std::memcpy(out.data(), column_data,
                            static_cast<size_t>(num_values) *
                            static_cast<size_t>(type_length));
                return out;
            }
            // Otherwise interpret as float32 vectors if type_length is a
            // multiple of sizeof(float)
            if (type_length % static_cast<int32_t>(sizeof(float)) == 0) {
                int64_t dim = type_length / static_cast<int32_t>(sizeof(float));
                TensorView float_src(column_data,
                                     TensorShape{num_values, dim},
                                     TensorDataType::FLOAT32);
                return cast(float_src, target_dtype);
            }
            return Error{ErrorCode::UNSUPPORTED_TYPE,
                         "copy_column: FIXED_LEN_BYTE_ARRAY with type_length "
                         "not a multiple of 4 can only be copied as UINT8"};
        }
 
        // For standard numeric types, wrap then cast
        auto src_dtype_result = parquet_to_tensor_dtype(physical_type);
        if (!src_dtype_result) {
            return Error{src_dtype_result.error().code,
                         src_dtype_result.error().message};
        }
 
        TensorDataType src_dtype = src_dtype_result.value();
        TensorView src(column_data, TensorShape{num_values}, src_dtype);
 
        // If source dtype matches target, just copy the bytes
        if (src_dtype == target_dtype) {
            return OwnedTensor(column_data,
                               TensorShape{num_values}, target_dtype);
        }
 
        return cast(src, target_dtype);
    }
 
    // -----------------------------------------------------------------------
    // Type casting
    // -----------------------------------------------------------------------
 
    static inline expected<OwnedTensor> cast(
            const TensorView& src,
            TensorDataType target_dtype) {
        if (!src.is_valid()) {
            return Error{ErrorCode::INTERNAL_ERROR,
                         "cast: source tensor is null"};
        }
        if (!src.is_contiguous()) {
            return Error{ErrorCode::INTERNAL_ERROR,
                         "cast: source tensor must be contiguous"};
        }
 
        // Same type — just copy
        if (src.dtype() == target_dtype) {
            return OwnedTensor(src.data(), src.shape(), target_dtype);
        }
 
        const int64_t n = src.num_elements();
        OwnedTensor out(src.shape(), target_dtype);
 
        // Dispatch on (src_dtype, target_dtype) using a helper
        bool ok = dispatch_cast(src.data(), src.dtype(),
                                out.data(), target_dtype, n);
        if (!ok) {
            return Error{ErrorCode::UNSUPPORTED_TYPE,
                         std::string("cast: unsupported conversion from ")
                         + tensor_dtype_name(src.dtype()) + " to "
                         + tensor_dtype_name(target_dtype)};
        }
 
        return out;
    }
 
    // -----------------------------------------------------------------------
    // Parquet → Tensor dtype mapping
    // -----------------------------------------------------------------------
 
    static inline expected<TensorDataType> parquet_to_tensor_dtype(
            PhysicalType pt) {
        switch (pt) {
            case PhysicalType::BOOLEAN:  return TensorDataType::BOOL;
            case PhysicalType::INT32:    return TensorDataType::INT32;
            case PhysicalType::INT64:    return TensorDataType::INT64;
            case PhysicalType::FLOAT:    return TensorDataType::FLOAT32;
            case PhysicalType::DOUBLE:   return TensorDataType::FLOAT64;
            case PhysicalType::FIXED_LEN_BYTE_ARRAY:
                return TensorDataType::UINT8;
            case PhysicalType::BYTE_ARRAY:
                return Error{ErrorCode::UNSUPPORTED_TYPE,
                             "BYTE_ARRAY has no fixed tensor type mapping"};
            case PhysicalType::INT96:
                return Error{ErrorCode::UNSUPPORTED_TYPE,
                             "INT96 has no tensor type mapping"};
        }
        return Error{ErrorCode::UNSUPPORTED_TYPE,
                     "unknown PhysicalType"};
    }
 
private:
    // -- Cast dispatch --------------------------------------------------------
 
    template <typename T>
    static inline T read_element(const void* data, int64_t idx) {
        return static_cast<const T*>(data)[idx];
    }
 
    template <typename T>
    static inline void write_element(void* data, int64_t idx, T val) {
        static_cast<T*>(data)[idx] = val;
    }
 
    template <typename Src, typename Dst>
    static inline void convert_loop(const void* src, void* dst, int64_t n) {
        const auto* s = static_cast<const Src*>(src);
        auto* d = static_cast<Dst*>(dst);
        for (int64_t i = 0; i < n; ++i) {
            d[i] = static_cast<Dst>(s[i]);
        }
    }
 
    template <typename Src>
    static inline bool dispatch_target(const void* src, void* dst,
                                       TensorDataType target, int64_t n) {
        switch (target) {
            case TensorDataType::FLOAT32: convert_loop<Src, float>(src, dst, n); return true;
            case TensorDataType::FLOAT64: convert_loop<Src, double>(src, dst, n); return true;
            case TensorDataType::INT32:   convert_loop<Src, int32_t>(src, dst, n); return true;
            case TensorDataType::INT64:   convert_loop<Src, int64_t>(src, dst, n); return true;
            case TensorDataType::INT8:    convert_loop<Src, int8_t>(src, dst, n); return true;
            case TensorDataType::UINT8:   convert_loop<Src, uint8_t>(src, dst, n); return true;
            case TensorDataType::INT16:   convert_loop<Src, int16_t>(src, dst, n); return true;
            case TensorDataType::BOOL:    convert_loop<Src, bool>(src, dst, n); return true;
            default: return false;
        }
    }
 
    static inline bool dispatch_cast(const void* src, TensorDataType src_dtype,
                                     void* dst, TensorDataType target, int64_t n) {
        switch (src_dtype) {
            case TensorDataType::FLOAT32: return dispatch_target<float>(src, dst, target, n);
            case TensorDataType::FLOAT64: return dispatch_target<double>(src, dst, target, n);
            case TensorDataType::INT32:   return dispatch_target<int32_t>(src, dst, target, n);
            case TensorDataType::INT64:   return dispatch_target<int64_t>(src, dst, target, n);
            case TensorDataType::INT8:    return dispatch_target<int8_t>(src, dst, target, n);
            case TensorDataType::UINT8:   return dispatch_target<uint8_t>(src, dst, target, n);
            case TensorDataType::INT16:   return dispatch_target<int16_t>(src, dst, target, n);
            case TensorDataType::BOOL:    return dispatch_target<bool>(src, dst, target, n);
            default: return false;
        }
    }
};
 
// ===========================================================================
// BatchTensorBuilder — assemble multi-column feature batches for ML inference
// ===========================================================================
 
class BatchTensorBuilder {
public:
    BatchTensorBuilder() = default;
 
    BatchTensorBuilder& add_column(const std::string& name,
                                   const TensorView& tensor) {
        columns_.push_back(ColumnEntry{name, tensor});
        return *this;
    }
 
    [[nodiscard]] TensorShape expected_shape() const {
        if (columns_.empty()) return TensorShape{0, 0};
 
        const int64_t rows = column_rows(columns_[0].tensor);
        int64_t total_cols = 0;
        for (const auto& entry : columns_) {
            total_cols += column_width(entry.tensor);
        }
        return TensorShape{rows, total_cols};
    }
 
    [[nodiscard]] size_t num_features() const noexcept {
        return columns_.size();
    }
 
    [[nodiscard]] expected<OwnedTensor> build(
            TensorDataType output_dtype = TensorDataType::FLOAT32) {
        if (columns_.empty()) {
            return Error{ErrorCode::INTERNAL_ERROR,
                         "BatchTensorBuilder: no columns added"};
        }
 
        // Determine row count from first column
        const int64_t rows = column_rows(columns_[0].tensor);
        if (rows <= 0) {
            return Error{ErrorCode::INTERNAL_ERROR,
                         "BatchTensorBuilder: first column has no rows"};
        }
 
        // Validate all columns have the same row count
        for (size_t i = 1; i < columns_.size(); ++i) {
            const int64_t col_rows = column_rows(columns_[i].tensor);
            if (col_rows != rows) {
                return Error{ErrorCode::SCHEMA_MISMATCH,
                             "BatchTensorBuilder: column '"
                             + columns_[i].name + "' has "
                             + std::to_string(col_rows) + " rows, expected "
                             + std::to_string(rows)};
            }
        }
 
        // Compute total output columns
        int64_t total_cols = 0;
        for (const auto& entry : columns_) {
            total_cols += column_width(entry.tensor);
        }
 
        // Allocate output tensor
        TensorShape out_shape{rows, total_cols};
        OwnedTensor output(out_shape, output_dtype);
        const size_t out_elem_size = tensor_element_size(output_dtype);
 
        // Fill column by column
        int64_t col_offset = 0;
        for (const auto& entry : columns_) {
            const TensorView& src = entry.tensor;
            const int64_t width = column_width(src);
 
            // Get or cast the source data to the output dtype
            // We need a contiguous float32 (or target type) source
            if (src.dtype() == output_dtype && src.is_contiguous()) {
                // Direct copy path
                copy_column_into(output, rows, total_cols, col_offset,
                                 width, src.data(), out_elem_size);
            } else {
                // Need to cast first
                auto cast_result = ColumnToTensor::cast(src, output_dtype);
                if (!cast_result) {
                    return Error{cast_result.error().code,
                                 "BatchTensorBuilder: failed to cast column '"
                                 + entry.name + "': "
                                 + cast_result.error().message};
                }
                copy_column_into(output, rows, total_cols, col_offset,
                                 width, cast_result.value().data(),
                                 out_elem_size);
            }
 
            col_offset += width;
        }
 
        return output;
    }
 
private:
    struct ColumnEntry {
        std::string name;
        TensorView  tensor;
    };
 
    std::vector<ColumnEntry> columns_;
 
    static int64_t column_rows(const TensorView& t) noexcept {
        if (t.shape().ndim() == 0) return 1;
        return t.shape().dims[0];
    }
 
    static int64_t column_width(const TensorView& t) noexcept {
        if (t.shape().ndim() <= 1) return 1;
        return t.shape().dims[1];
    }
 
    static void copy_column_into(
            OwnedTensor& output,
            int64_t rows,
            int64_t total_cols,
            int64_t col_offset,
            int64_t width,
            const void* src_data,
            size_t elem_size) {
        auto* dst_base = static_cast<uint8_t*>(output.data());
        const auto* src_base = static_cast<const uint8_t*>(src_data);
 
        const size_t row_byte_stride = static_cast<size_t>(total_cols) * elem_size;
        const size_t src_row_bytes   = static_cast<size_t>(width) * elem_size;
        const size_t col_byte_offset = static_cast<size_t>(col_offset) * elem_size;
 
        for (int64_t r = 0; r < rows; ++r) {
            const size_t dst_offset = static_cast<size_t>(r) * row_byte_stride
                                      + col_byte_offset;
            const size_t src_offset = static_cast<size_t>(r) * src_row_bytes;
            std::memcpy(dst_base + dst_offset,
                        src_base + src_offset,
                        src_row_bytes);
        }
    }
};
 
} // namespace signet::forge