arrow_bridge.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: AGPL-3.0-or-later
// Copyright 2026 Johnson Ogundeji
#pragma once
 
 
#include "signet/ai/tensor_bridge.hpp"
#include "signet/types.hpp"
#include "signet/error.hpp"
 
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <memory>
#include <string>
 
 
#ifndef SIGNET_ARROW_C_DATA_DEFINED
#define SIGNET_ARROW_C_DATA_DEFINED
 
extern "C" {
 
struct ArrowSchema {
    const char*    format;       
    const char*    name;         
    const char*    metadata;     
    int64_t        flags;        
    int64_t        n_children;   
    ArrowSchema**  children;     
    ArrowSchema*   dictionary;   
    void (*release)(ArrowSchema*); 
    void*          private_data; 
};
 
struct ArrowArray {
    int64_t        length;       
    int64_t        null_count;   
    int64_t        offset;       
    int64_t        n_buffers;    
    int64_t        n_children;   
    const void**   buffers;      
    ArrowArray**   children;     
    ArrowArray*    dictionary;   
    void (*release)(ArrowArray*); 
    void*          private_data; 
};
 
} // extern "C"
 
#endif // SIGNET_ARROW_C_DATA_DEFINED
 
namespace signet::forge {
 
namespace detail {
 
struct ArrowSchemaPrivate {
    std::string format_storage;   
    std::string name_storage;     
};
 
struct ArrowArrayPrivate {
    bool         owns_data   = false;   
    const void*  buffer_ptrs[2] = {nullptr, nullptr}; 
};
 
inline void release_arrow_schema(ArrowSchema* schema) {
    if (schema == nullptr) return;
    if (schema->private_data != nullptr) {
        delete static_cast<ArrowSchemaPrivate*>(schema->private_data);
        schema->private_data = nullptr;
    }
    schema->format  = nullptr;
    schema->name    = nullptr;
    schema->release = nullptr;
}
 
inline void release_arrow_array(ArrowArray* array) {
    if (array == nullptr) return;
    if (array->private_data != nullptr) {
        auto* ctx = static_cast<ArrowArrayPrivate*>(array->private_data);
        if (ctx->owns_data && ctx->buffer_ptrs[1] != nullptr) {
            std::free(const_cast<void*>(ctx->buffer_ptrs[1]));
        }
        delete ctx;
        array->private_data = nullptr;
    }
    array->buffers = nullptr;
    array->release = nullptr;
}
 
} // namespace detail
 
 
inline const char* parquet_to_arrow_format(PhysicalType pt) {
    switch (pt) {
        case PhysicalType::BOOLEAN: return "b";
        case PhysicalType::INT32:   return "i";
        case PhysicalType::INT64:   return "l";
        case PhysicalType::FLOAT:   return "f";
        case PhysicalType::DOUBLE:  return "g";
        default:                    return nullptr;
    }
}
 
inline const char* tensor_dtype_to_arrow_format(TensorDataType dtype) {
    switch (dtype) {
        case TensorDataType::FLOAT32: return "f";
        case TensorDataType::FLOAT64: return "g";
        case TensorDataType::INT32:   return "i";
        case TensorDataType::INT64:   return "l";
        case TensorDataType::INT8:    return "c";
        case TensorDataType::UINT8:   return "C";
        case TensorDataType::INT16:   return "s";
        case TensorDataType::FLOAT16: return "e";
        case TensorDataType::BOOL:    return "b";
    }
    return nullptr; // unreachable
}
 
inline expected<TensorDataType> arrow_format_to_tensor_dtype(const char* format) {
    if (format == nullptr || format[0] == '\0') {
        return Error{ErrorCode::UNSUPPORTED_TYPE,
                     "null or empty Arrow format string"};
    }
 
    // Only single-character format codes are handled for primitive types.
    if (format[1] != '\0') {
        return Error{ErrorCode::UNSUPPORTED_TYPE,
                     std::string("unsupported Arrow format: ") + format};
    }
 
    switch (format[0]) {
        case 'f': return TensorDataType::FLOAT32;
        case 'g': return TensorDataType::FLOAT64;
        case 'i': return TensorDataType::INT32;
        case 'l': return TensorDataType::INT64;
        case 'c': return TensorDataType::INT8;
        case 'C': return TensorDataType::UINT8;
        case 's': return TensorDataType::INT16;
        case 'e': return TensorDataType::FLOAT16;
        case 'b': return TensorDataType::BOOL;
        default:
            return Error{ErrorCode::UNSUPPORTED_TYPE,
                         std::string("unsupported Arrow format character: '")
                         + format[0] + "'"};
    }
}
 
inline expected<TensorDataType> physical_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;
        default:
            return Error{ErrorCode::UNSUPPORTED_TYPE,
                         "PhysicalType has no TensorDataType mapping"};
    }
}
 
inline size_t physical_type_byte_size(PhysicalType pt) {
    switch (pt) {
        case PhysicalType::BOOLEAN: return 1; // Arrow uses 1 bit, but data buffer is byte-aligned
        case PhysicalType::INT32:   return 4;
        case PhysicalType::INT64:   return 8;
        case PhysicalType::FLOAT:   return 4;
        case PhysicalType::DOUBLE:  return 8;
        default:                    return 0;
    }
}
 
 
class ArrowExporter {
public:
    static inline expected<void> export_tensor(
        const TensorView& tensor,
        const std::string& name,
        ArrowArray* out_array,
        ArrowSchema* out_schema)
    {
        // CWE-457: Use of Uninitialized Variable — zero-init prevents double-free on partial init failure
        // M27: Zero-initialize outputs so callers see release=nullptr on early error
        std::memset(out_schema, 0, sizeof(ArrowSchema));
        std::memset(out_array, 0, sizeof(ArrowArray));
 
        if (!tensor.is_valid()) {
            return Error{ErrorCode::INTERNAL_ERROR,
                         "cannot export invalid tensor to Arrow"};
        }
 
        if (!tensor.is_contiguous()) {
            return Error{ErrorCode::UNSUPPORTED_TYPE,
                         "non-contiguous tensor cannot be exported to Arrow"};
        }
 
        const char* fmt = tensor_dtype_to_arrow_format(tensor.dtype());
        if (fmt == nullptr) {
            return Error{ErrorCode::UNSUPPORTED_TYPE,
                         "tensor dtype has no Arrow format mapping"};
        }
 
        // -- Fill schema (RAII: unique_ptr until setup complete) --
        auto schema_owner = std::make_unique<detail::ArrowSchemaPrivate>();
        schema_owner->format_storage = fmt;
        schema_owner->name_storage   = name;
 
        out_schema->format       = schema_owner->format_storage.c_str();
        out_schema->name         = schema_owner->name_storage.c_str();
        out_schema->metadata     = nullptr;
        out_schema->flags        = 0;   // non-nullable
        out_schema->n_children   = 0;
        out_schema->children     = nullptr;
        out_schema->dictionary   = nullptr;
        out_schema->release      = detail::release_arrow_schema;
        out_schema->private_data = schema_owner.release(); // transfer ownership
 
        // -- Fill array (RAII: unique_ptr until setup complete) --
        auto array_owner = std::make_unique<detail::ArrowArrayPrivate>();
        array_owner->owns_data      = false; // zero-copy
        array_owner->buffer_ptrs[0] = nullptr;  // no validity bitmap (non-nullable)
        array_owner->buffer_ptrs[1] = tensor.data();
 
        const int64_t num_elements = static_cast<int64_t>(tensor.num_elements());
 
        out_array->length       = num_elements;
        out_array->null_count   = 0;
        out_array->offset       = 0;
        out_array->n_buffers    = 2;
        out_array->n_children   = 0;
        out_array->buffers      = array_owner->buffer_ptrs;
        out_array->children     = nullptr;
        out_array->dictionary   = nullptr;
        out_array->release      = detail::release_arrow_array;
        out_array->private_data = array_owner.release(); // transfer ownership
 
        return expected<void>{};
    }
 
    static inline expected<void> export_column(
        const void* data,
        int64_t num_values,
        PhysicalType physical_type,
        const std::string& name,
        ArrowArray* out_array,
        ArrowSchema* out_schema)
    {
        // CWE-457: Use of Uninitialized Variable — zero-init prevents double-free on partial init failure
        // M27: Zero-initialize outputs so callers see release=nullptr on early error
        std::memset(out_schema, 0, sizeof(ArrowSchema));
        std::memset(out_array, 0, sizeof(ArrowArray));
 
        if (data == nullptr || num_values <= 0) {
            return Error{ErrorCode::INTERNAL_ERROR,
                         "cannot export null/empty column to Arrow"};
        }
 
        const char* fmt = parquet_to_arrow_format(physical_type);
        if (fmt == nullptr) {
            return Error{ErrorCode::UNSUPPORTED_TYPE,
                         "PhysicalType has no direct Arrow format mapping"};
        }
 
        // -- Fill schema (RAII: unique_ptr until setup complete) --
        auto schema_owner = std::make_unique<detail::ArrowSchemaPrivate>();
        schema_owner->format_storage = fmt;
        schema_owner->name_storage   = name;
 
        out_schema->format       = schema_owner->format_storage.c_str();
        out_schema->name         = schema_owner->name_storage.c_str();
        out_schema->metadata     = nullptr;
        out_schema->flags        = 0;
        out_schema->n_children   = 0;
        out_schema->children     = nullptr;
        out_schema->dictionary   = nullptr;
        out_schema->release      = detail::release_arrow_schema;
        out_schema->private_data = schema_owner.release(); // transfer ownership
 
        // -- Fill array (RAII: unique_ptr until setup complete) --
        auto array_owner = std::make_unique<detail::ArrowArrayPrivate>();
        array_owner->owns_data      = false;
        array_owner->buffer_ptrs[0] = nullptr; // no validity bitmap
        array_owner->buffer_ptrs[1] = data;
 
        out_array->length       = num_values;
        out_array->null_count   = 0;
        out_array->offset       = 0;
        out_array->n_buffers    = 2;
        out_array->n_children   = 0;
        out_array->buffers      = array_owner->buffer_ptrs;
        out_array->children     = nullptr;
        out_array->dictionary   = nullptr;
        out_array->release      = detail::release_arrow_array;
        out_array->private_data = array_owner.release(); // transfer ownership
 
        return expected<void>{};
    }
};
 
class ArrowImporter {
public:
    static inline expected<TensorView> import_array(
        const ArrowArray* array,
        const ArrowSchema* schema)
    {
        if (array == nullptr || schema == nullptr) {
            return Error{ErrorCode::INTERNAL_ERROR,
                         "null ArrowArray or ArrowSchema pointer"};
        }
 
        if (array->release == nullptr) {
            return Error{ErrorCode::INTERNAL_ERROR,
                         "ArrowArray has already been released"};
        }
 
        if (array->null_count != 0) {
            return Error{ErrorCode::UNSUPPORTED_TYPE,
                         "ArrowArray with nulls cannot be imported as a dense tensor"};
        }
 
        auto dtype_result = arrow_format_to_tensor_dtype(schema->format);
        if (!dtype_result) {
            return dtype_result.error();
        }
 
        TensorDataType dtype = *dtype_result;
        const int64_t length = array->length;
 
        if (length <= 0) {
            return Error{ErrorCode::INTERNAL_ERROR,
                         "ArrowArray has zero or negative length"};
        }
 
        // Reject unreasonably large arrays to prevent OOB reads in copy paths.
        // Arrow C Data Interface doesn't carry buffer sizes, so this is a
        // defense-in-depth cap. 1 billion elements ≈ 8 GB for float64.
        static constexpr int64_t MAX_IMPORT_ELEMENTS = 1'000'000'000;
        static constexpr int64_t MAX_ARROW_OFFSET = 1'000'000'000LL;
        static constexpr int64_t MAX_ARROW_LENGTH = 1'000'000'000LL;
 
        if (array->offset > MAX_ARROW_OFFSET) {
            return Error{ErrorCode::INVALID_ARGUMENT, "ArrowArray offset exceeds 1B cap (CWE-190)"};
        }
        if (array->length > MAX_ARROW_LENGTH) {
            return Error{ErrorCode::INVALID_ARGUMENT, "ArrowArray length exceeds 1B cap (CWE-190)"};
        }
 
        if (length > MAX_IMPORT_ELEMENTS) {
            return Error{ErrorCode::IO_ERROR,
                         "ArrowArray length exceeds import limit (1 billion elements)"};
        }
 
        if (array->offset < 0) {
            return Error{ErrorCode::IO_ERROR,
                         "ArrowArray has negative offset"};
        }
 
        if (array->n_buffers < 2 || array->buffers == nullptr) {
            return Error{ErrorCode::INTERNAL_ERROR,
                         "ArrowArray does not have expected buffer layout"};
        }
 
        const void* data_buf = array->buffers[1];
        if (data_buf == nullptr) {
            return Error{ErrorCode::INTERNAL_ERROR,
                         "ArrowArray data buffer (buffers[1]) is null"};
        }
 
        if (static_cast<int64_t>(length) > INT64_MAX - array->offset) {
            return Error{ErrorCode::IO_ERROR,
                         "ArrowArray: offset + length overflows int64"};
        }
 
        // Account for offset: data starts at offset * element_size bytes
        const size_t elem_size = tensor_element_size(dtype);
        const size_t offset_val = static_cast<size_t>(array->offset);
        if (offset_val > static_cast<size_t>(MAX_IMPORT_ELEMENTS)) {
            return Error{ErrorCode::IO_ERROR,
                         "ArrowArray offset exceeds import limit"};
        }
 
        // Validate (offset + length) * elem_size doesn't overflow size_t.
        // The Arrow C Data Interface doesn't carry buffer sizes, so we
        // cannot verify the buffer is large enough.  This overflow guard
        // prevents wild pointer arithmetic from crafted metadata.
        const size_t total_elems = offset_val + static_cast<size_t>(length);
        if (total_elems > static_cast<size_t>(MAX_IMPORT_ELEMENTS)) {
            return Error{ErrorCode::IO_ERROR,
                         "ArrowArray: offset + length exceeds import limit"};
        }
        if (elem_size > 0 && total_elems > SIZE_MAX / elem_size) {
            return Error{ErrorCode::IO_ERROR,
                         "ArrowArray: (offset + length) * elem_size overflows size_t"};
        }
        const uint8_t* base = static_cast<const uint8_t*>(data_buf)
                            + offset_val * elem_size;
 
        TensorShape shape;
        shape.dims = {length};
 
        return TensorView(const_cast<void*>(static_cast<const void*>(base)),
                          shape, dtype);
    }
 
    static inline expected<OwnedTensor> import_array_copy(
        const ArrowArray* array,
        const ArrowSchema* schema)
    {
        auto view_result = import_array(array, schema);
        if (!view_result) {
            return view_result.error();
        }
 
        // Create an owned copy from the view
        const auto& v = *view_result;
        return OwnedTensor(v.data(), v.shape(), v.dtype());
    }
};
 
} // namespace signet::forge