nanoarrow 0.1.0

data/ext/nanoarrow/array_stream.c

@@ -0,0 +1,201 @@
+#include <ruby/ruby.h>
+
+#include "ext.h"
+#include "nanoarrow/nanoarrow.h"
+
+VALUE cCArrayStream;
+
+typedef struct array_stream {
+    VALUE base;
+    struct ArrowArrayStream* ptr;
+    VALUE cached_schema;
+} array_stream_t;
+
+static void array_stream_mark(void* ptr)
+{
+    array_stream_t* stream = (array_stream_t*) ptr;
+    rb_gc_mark_movable(stream->base);
+    rb_gc_mark_movable(stream->cached_schema);
+}
+
+const rb_data_type_t array_stream_data_type = {
+    .wrap_struct_name = "Nanoarrow::CArrayStream",
+    .function = {
+        .dmark = array_stream_mark,
+        .dfree = RUBY_TYPED_DEFAULT_FREE,
+    },
+    .flags = RUBY_TYPED_FREE_IMMEDIATELY
+};
+
+static VALUE array_stream_allocate(VALUE klass)
+{
+    array_stream_t* stream;
+    VALUE obj = TypedData_Make_Struct(klass, array_stream_t, &array_stream_data_type, stream);
+    stream->base = Qnil;
+    stream->ptr = NULL;
+    stream->cached_schema = Qnil;
+    return obj;
+}
+
+static VALUE array_stream_import_from_c_capsule(VALUE klass, VALUE schema_capsule)
+{
+    array_stream_t* stream;
+    VALUE obj = array_stream_allocate(klass);
+    GetArrayStream(obj, stream);
+
+    VALUE address = rb_funcall(schema_capsule, rb_intern("to_i"), 0);
+    stream->base = schema_capsule;
+    stream->ptr = (struct ArrowArrayStream*) NUM2ULL(address);
+    return obj;
+}
+
+static VALUE array_stream_from_c_arrays(VALUE klass, VALUE arrays, VALUE schema, VALUE validate)
+{
+    array_stream_t* stream;
+    VALUE out = array_stream_allocate(klass);
+    GetArrayStream(out, stream);
+
+    // TODO move to parameter
+    VALUE move = Qfalse;
+
+    VALUE out_schema;
+    VALUE validate_schema;
+    if (RTEST(validate) && !RTEST(move))
+    {
+        validate_schema = schema;
+        out_schema = rb_funcall(schema, rb_intern("deep_dup"), 0);
+    }
+    else if (RTEST(validate))
+    {
+        validate_schema = rb_funcall(schema, rb_intern("deep_dup"), 0);
+        out_schema = schema;
+    }
+    else
+        out_schema = rb_funcall(schema, rb_intern("deep_dup"), 0);
+
+    schema_t* schema_ptr;
+    GetSchema(out_schema, schema_ptr);
+
+    Check_Type(arrays, T_ARRAY);
+    long arrays_len = RARRAY_LEN(arrays);
+
+    stream->base = alloc_c_array_stream(&stream->ptr);
+
+    int code = ArrowBasicArrayStreamInit(stream->ptr, schema_ptr->ptr, arrays_len);
+    raise_error_not_ok("ArrowBasicArrayStreamInit()", code);
+
+    struct ArrowArrayStream* c_array_stream_out = stream->ptr;
+    struct ArrowArray tmp;
+    for (long i = 0; i < arrays_len; i++)
+    {
+        array_t* array;
+        GetArray(rb_ary_entry(arrays, i), array);
+
+        if (RTEST(validate))
+            assert_type_equal(array->schema, validate_schema, Qfalse);
+
+        if (!RTEST(move))
+        {
+            c_array_shallow_copy(array->base, array->ptr, &tmp);
+            ArrowBasicArrayStreamSetArray(c_array_stream_out, i, &tmp);
+        }
+        else
+            ArrowBasicArrayStreamSetArray(c_array_stream_out, i, array->ptr);
+    }
+
+    if (RTEST(validate))
+    {
+        struct ArrowError error;
+        int code = ArrowBasicArrayStreamValidate(c_array_stream_out, &error);
+        raise_message_not_ok(&error, "ArrowBasicArrayStreamValidate()", code);
+    }
+
+    return out;
+}
+
+static VALUE array_stream_assert_valid(VALUE self)
+{
+    array_stream_t* stream;
+    GetArrayStream(self, stream);
+
+    if (stream->ptr == NULL)
+        rb_raise(rb_eRuntimeError, "array stream pointer is NULL");
+    if (stream->ptr->release == NULL)
+        rb_raise(rb_eRuntimeError, "array stream is released");
+
+    return Qnil;
+}
+
+static VALUE array_stream_get_schema(VALUE self, VALUE schema)
+{
+    array_stream_t* stream;
+    GetArrayStream(self, stream);
+
+    array_stream_assert_valid(self);
+
+    struct ArrowError error;
+    schema_t* schema_ptr;
+    GetSchema(schema, schema_ptr);
+    int code = ArrowArrayStreamGetSchema(stream->ptr, schema_ptr->ptr, &error);
+    raise_message_not_ok(&error, "ArrowArrayStream::get_schema()", code);
+    return Qnil;
+}
+
+static VALUE array_stream_get_cached_schema(VALUE self)
+{
+    array_stream_t* stream;
+    GetArrayStream(self, stream);
+
+    if (NIL_P(stream->cached_schema))
+    {
+        stream->cached_schema = rb_funcall(cCSchema, rb_intern("allocate"), 0);
+        array_stream_get_schema(self, stream->cached_schema);
+    }
+    return stream->cached_schema;
+}
+
+static VALUE array_stream_get_next(VALUE self)
+{
+    array_stream_t* stream;
+    GetArrayStream(self, stream);
+
+    array_stream_assert_valid(self);
+
+    array_t* array_ptr;
+    VALUE array = rb_funcall(cCArray, rb_intern("allocate"), 0);
+    GetArray(array, array_ptr);
+    array_ptr->schema = array_stream_get_cached_schema(self);
+
+    struct ArrowError error;
+    int code = ArrowArrayStreamGetNext(stream->ptr, array_ptr->ptr, &error);
+    raise_message_not_ok(&error, "ArrowArrayStreamGetNext()", code);
+
+    if (!RTEST(rb_funcall(array, rb_intern("valid?"), 0)))
+        return Qnil;
+
+    return array;
+}
+
+static VALUE array_stream_arrow_c_stream(VALUE self)
+{
+    array_stream_t* stream;
+    GetArrayStream(self, stream);
+
+    array_stream_assert_valid(self);
+
+    struct ArrowArrayStream* c_array_stream_out;
+    VALUE array_stream_capsule = alloc_c_array_stream(&c_array_stream_out);
+    ArrowArrayStreamMove(stream->ptr, c_array_stream_out);
+    return array_stream_capsule;
+}
+
+void Init_array_stream(void)
+{
+    cCArrayStream = rb_define_class_under(mNanoarrow, "CArrayStream", rb_cObject);
+    rb_define_alloc_func(cCArrayStream, array_stream_allocate);
+    rb_define_singleton_method(cCArrayStream, "import_from_c_capsule", array_stream_import_from_c_capsule, 1);
+    rb_define_singleton_method(cCArrayStream, "from_c_arrays", array_stream_from_c_arrays, 3);
+    rb_define_method(cCArrayStream, "get_next", array_stream_get_next, 0);
+    rb_define_method(cCArrayStream, "get_cached_schema", array_stream_get_cached_schema, 0);
+    rb_define_method(cCArrayStream, "arrow_c_stream", array_stream_arrow_c_stream, 0);
+}

data/ext/nanoarrow/array_view.c

@@ -0,0 +1,177 @@
+#include <ruby/ruby.h>
+
+#include "ext.h"
+#include "nanoarrow/nanoarrow.h"
+
+VALUE cCArrayView;
+
+typedef struct array_view {
+    VALUE base;
+    VALUE array_base;
+    struct ArrowArrayView* ptr;
+} array_view_t;
+
+static void array_view_mark(void* ptr)
+{
+    array_view_t* view = (array_view_t*) ptr;
+    rb_gc_mark_movable(view->base);
+    rb_gc_mark_movable(view->array_base);
+}
+
+const rb_data_type_t array_view_data_type = {
+    .wrap_struct_name = "Nanoarrow::CArrayView",
+    .function = {
+        .dmark = array_view_mark,
+        .dfree = RUBY_TYPED_DEFAULT_FREE,
+    },
+    .flags = RUBY_TYPED_FREE_IMMEDIATELY
+};
+
+static VALUE array_view_allocate(VALUE klass)
+{
+    array_view_t* view;
+    VALUE obj = TypedData_Make_Struct(klass, array_view_t, &array_view_data_type, view);
+    view->base = Qnil;
+    view->array_base = Qnil;
+    view->ptr = NULL;
+    return obj;
+}
+
+static VALUE array_view_from_schema(VALUE klass, VALUE schema)
+{
+    struct ArrowArrayView* c_array_view;
+    VALUE base = alloc_c_array_view(&c_array_view);
+
+    schema_t* schema_ptr;
+    GetSchema(schema, schema_ptr);
+
+    struct ArrowError error;
+    int code = ArrowArrayViewInitFromSchema(c_array_view, schema_ptr->ptr, &error);
+    raise_message_not_ok(&error, "ArrowArrayViewInitFromSchema()", code);
+
+    array_view_t* view;
+    VALUE obj = TypedData_Make_Struct(klass, array_view_t, &array_view_data_type, view);
+    view->base = base;
+    view->array_base = Qnil;
+    view->ptr = c_array_view;
+    return obj;
+}
+
+static VALUE array_view_set_array(VALUE self, VALUE array)
+{
+    array_view_t* view;
+    GetArrayView(self, view);
+
+    struct ArrowError error;
+    int code;
+
+    array_t* array_ptr;
+    GetArray(array, array_ptr);
+
+    code = ArrowArrayViewSetArray(view->ptr, array_ptr->ptr, &error);
+
+    raise_message_not_ok(&error, "ArrowArrayViewSetArray()", code);
+    view->array_base = array_ptr->base;
+
+    return self;
+}
+
+static VALUE array_view_each_bool(VALUE self)
+{
+    array_view_t* view;
+    GetArrayView(self, view);
+
+    for (int64_t i = 0; i < view->ptr->length; i++)
+        rb_yield(ArrowArrayViewIsNull(view->ptr, i) ? Qnil : ArrowArrayViewGetIntUnsafe(view->ptr, i) ? Qtrue : Qfalse);
+
+    return self;
+}
+
+static VALUE array_view_each_int(VALUE self)
+{
+    array_view_t* view;
+    GetArrayView(self, view);
+
+    for (int64_t i = 0; i < view->ptr->length; i++)
+        rb_yield(ArrowArrayViewIsNull(view->ptr, i) ? Qnil : LL2NUM(ArrowArrayViewGetIntUnsafe(view->ptr, i)));
+
+    return self;
+}
+
+static VALUE array_view_each_uint(VALUE self)
+{
+    array_view_t* view;
+    GetArrayView(self, view);
+
+    for (int64_t i = 0; i < view->ptr->length; i++)
+        rb_yield(ArrowArrayViewIsNull(view->ptr, i) ? Qnil : ULL2NUM(ArrowArrayViewGetUIntUnsafe(view->ptr, i)));
+
+    return self;
+}
+
+static VALUE array_view_each_double(VALUE self)
+{
+    array_view_t* view;
+    GetArrayView(self, view);
+
+    for (int64_t i = 0; i < view->ptr->length; i++)
+        rb_yield(ArrowArrayViewIsNull(view->ptr, i) ? Qnil : DBL2NUM(ArrowArrayViewGetDoubleUnsafe(view->ptr, i)));
+
+    return self;
+}
+
+static VALUE array_view_each_string(VALUE self)
+{
+    array_view_t* view;
+    GetArrayView(self, view);
+
+    for (int64_t i = 0; i < view->ptr->length; i++)
+    {
+        VALUE v;
+        if (ArrowArrayViewIsNull(view->ptr, i))
+            v = Qnil;
+        else
+        {
+            struct ArrowStringView sv = ArrowArrayViewGetStringUnsafe(view->ptr, i);
+            v = rb_str_new(sv.data, sv.size_bytes);
+        }
+        rb_yield(v);
+    }
+
+    return self;
+}
+
+static VALUE array_view_each_binary(VALUE self)
+{
+    array_view_t* view;
+    GetArrayView(self, view);
+
+    for (int64_t i = 0; i < view->ptr->length; i++)
+    {
+        VALUE v;
+        if (ArrowArrayViewIsNull(view->ptr, i))
+            v = Qnil;
+        else
+        {
+            struct ArrowBufferView bv = ArrowArrayViewGetBytesUnsafe(view->ptr, i);
+            v = rb_str_new(bv.data.data, bv.size_bytes);
+        }
+        rb_yield(v);
+    }
+
+    return self;
+}
+
+void Init_array_view(void)
+{
+    cCArrayView = rb_define_class_under(mNanoarrow, "CArrayView", rb_cObject);
+    rb_define_alloc_func(cCArrayView, array_view_allocate);
+    rb_define_singleton_method(cCArrayView, "from_schema", array_view_from_schema, 1);
+    rb_define_method(cCArrayView, "set_array", array_view_set_array, 1);
+    rb_define_method(cCArrayView, "each_bool", array_view_each_bool, 0);
+    rb_define_method(cCArrayView, "each_int", array_view_each_int, 0);
+    rb_define_method(cCArrayView, "each_uint", array_view_each_uint, 0);
+    rb_define_method(cCArrayView, "each_double", array_view_each_double, 0);
+    rb_define_method(cCArrayView, "each_string", array_view_each_string, 0);
+    rb_define_method(cCArrayView, "each_binary", array_view_each_binary, 0);
+}

data/ext/nanoarrow/buffer.c

@@ -0,0 +1,46 @@
+#include <ruby/ruby.h>
+
+#include "ext.h"
+#include "nanoarrow/nanoarrow.h"
+
+VALUE cCBuffer;
+
+static void buffer_mark(void* ptr)
+{
+    buffer_t* buffer = (buffer_t*) ptr;
+    rb_gc_mark_movable(buffer->base);
+}
+
+const rb_data_type_t buffer_data_type = {
+    .wrap_struct_name = "Nanoarrow::CArrayView",
+    .function = {
+        .dmark = buffer_mark,
+        .dfree = RUBY_TYPED_DEFAULT_FREE,
+    },
+    .flags = RUBY_TYPED_FREE_IMMEDIATELY
+};
+
+static VALUE buffer_allocate(VALUE klass)
+{
+    buffer_t* buffer;
+    VALUE obj = TypedData_Make_Struct(klass, buffer_t, &buffer_data_type, buffer);
+    buffer->base = Qnil;
+    buffer->ptr = NULL;
+    return obj;
+}
+
+static VALUE buffer_empty(VALUE klass)
+{
+    buffer_t* buffer;
+    VALUE out = buffer_allocate(klass);
+    GetBuffer(out, buffer);
+    buffer->base = alloc_c_buffer(&buffer->ptr);
+    return out;
+}
+
+void Init_buffer(void)
+{
+    cCBuffer = rb_define_class_under(mNanoarrow, "CBuffer", rb_cObject);
+    rb_define_alloc_func(cCBuffer, buffer_allocate);
+    rb_define_singleton_method(cCBuffer, "empty", buffer_empty, 0);
+}

data/ext/nanoarrow/capsule.c

@@ -0,0 +1,89 @@
+#include <ruby/ruby.h>
+
+#include "ext.h"
+#include "nanoarrow/nanoarrow.h"
+
+VALUE cCapsule;
+
+struct capsule {
+    void* ptr;
+    const char* name;
+    void (*destructor)(capsule_t*);
+};
+
+static void capsule_free(void* ptr)
+{
+    capsule_t* capsule = (capsule_t*) ptr;
+
+    if (capsule->destructor)
+        capsule->destructor(capsule);
+
+    xfree(ptr);
+}
+
+const rb_data_type_t capsule_data_type = {
+    .wrap_struct_name = "Nanoarrow::Capsule",
+    .function = {
+        .dfree = capsule_free,
+    },
+    .flags = RUBY_TYPED_FREE_IMMEDIATELY
+};
+
+VALUE capsule_new(void* ptr, const char* name, void (*destructor)(capsule_t*))
+{
+    capsule_t* capsule;
+    VALUE obj = TypedData_Make_Struct(cCapsule, capsule_t, &capsule_data_type, capsule);
+    capsule->ptr = ptr;
+    capsule->name = name;
+    capsule->destructor = destructor;
+    return obj;
+}
+
+void* capsule_pointer(capsule_t* capsule, const char* name)
+{
+    bool matches;
+    if (name == NULL || capsule->name == NULL)
+        matches = name == capsule->name;
+    else
+        matches = strcmp(name, capsule->name) == 0;
+
+    // cannot raise since may be called in dfree function
+    if (!matches)
+        return NULL;
+
+    return capsule->ptr;
+}
+
+static VALUE capsule_allocate(VALUE klass)
+{
+    capsule_t* capsule;
+    VALUE obj = TypedData_Make_Struct(klass, capsule_t, &capsule_data_type, capsule);
+    capsule->ptr = NULL;
+    capsule->name = NULL;
+    capsule->destructor = NULL;
+    return obj;
+}
+
+static VALUE capsule_to_i(VALUE self)
+{
+    capsule_t* capsule;
+    GetCapsule(self, capsule);
+
+    return ULL2NUM((uintptr_t) capsule->ptr);
+}
+
+static VALUE capsule_name(VALUE self)
+{
+    capsule_t* capsule;
+    GetCapsule(self, capsule);
+
+    return rb_utf8_str_new_cstr(capsule->name);
+}
+
+void Init_capsule(void)
+{
+    cCapsule = rb_define_class_under(mNanoarrow, "Capsule", rb_cObject);
+    rb_define_alloc_func(cCapsule, capsule_allocate);
+    rb_define_method(cCapsule, "to_i", capsule_to_i, 0);
+    rb_define_method(cCapsule, "name", capsule_name, 0);
+}

data/ext/nanoarrow/ext.c

@@ -0,0 +1,23 @@
+#include <ruby/ruby.h>
+
+#include "ext.h"
+#include "nanoarrow/nanoarrow.h"
+
+VALUE mNanoarrow;
+
+void Init_ext(void)
+{
+    mNanoarrow = rb_define_module("Nanoarrow");
+
+    Init_array();
+    Init_array_builder();
+    Init_array_stream();
+    Init_array_view();
+    Init_buffer();
+    Init_capsule();
+    Init_materialized();
+    Init_schema();
+    Init_schema_builder();
+    Init_schema_view();
+    Init_types();
+}

data/ext/nanoarrow/ext.h

@@ -0,0 +1,89 @@
+#pragma once
+
+#include <ruby/ruby.h>
+
+#include "nanoarrow/nanoarrow.h"
+
+#define GetArray(obj, ptr) TypedData_Get_Struct((obj), array_t, &array_data_type, (ptr));
+#define GetArrayBuilder(obj, ptr) TypedData_Get_Struct((obj), array_builder_t, &array_builder_data_type, (ptr));
+#define GetArrayStream(obj, ptr) TypedData_Get_Struct((obj), array_stream_t, &array_stream_data_type, (ptr));
+#define GetArrayView(obj, ptr) TypedData_Get_Struct((obj), array_view_t, &array_view_data_type, (ptr));
+#define GetBuffer(obj, ptr) TypedData_Get_Struct((obj), buffer_t, &buffer_data_type, (ptr));
+#define GetCapsule(obj, ptr) TypedData_Get_Struct((obj), capsule_t, &capsule_data_type, (ptr));
+#define GetMatArrayStream(obj, ptr) TypedData_Get_Struct((obj), mat_array_stream_t, &mat_array_stream_data_type, (ptr));
+#define GetSchema(obj, ptr) TypedData_Get_Struct((obj), schema_t, &schema_data_type, (ptr));
+#define GetSchemaBuilder(obj, ptr) TypedData_Get_Struct((obj), schema_builder_t, &schema_builder_data_type, (ptr));
+#define GetSchemaView(obj, ptr) TypedData_Get_Struct((obj), schema_view_t, &schema_view_data_type, (ptr));
+
+typedef struct array {
+    VALUE base;
+    struct ArrowArray* ptr;
+    VALUE schema;
+} array_t;
+
+extern const rb_data_type_t array_data_type;
+
+typedef struct buffer {
+    VALUE base;
+    struct ArrowBuffer* ptr;
+} buffer_t;
+
+extern const rb_data_type_t buffer_data_type;
+
+typedef struct schema {
+    VALUE base;
+    struct ArrowSchema* ptr;
+} schema_t;
+
+extern const rb_data_type_t schema_data_type;
+
+extern VALUE mNanoarrow;
+extern VALUE cCArray;
+extern VALUE cCArrayBuilder;
+extern VALUE cCArrayStream;
+extern VALUE cCArrayView;
+extern VALUE cCBuffer;
+extern VALUE cCapsule;
+extern VALUE cCMatArrayStream;
+extern VALUE cCSchema;
+extern VALUE cCSchemaBuilder;
+extern VALUE cCSchemaView;
+
+void Init_array(void);
+void Init_array_builder(void);
+void Init_array_stream(void);
+void Init_array_view(void);
+void Init_buffer(void);
+void Init_capsule(void);
+void Init_materialized(void);
+void Init_schema(void);
+void Init_schema_builder(void);
+void Init_schema_view(void);
+void Init_types(void);
+
+typedef struct capsule capsule_t;
+
+extern const rb_data_type_t capsule_data_type;
+
+VALUE capsule_new(void* ptr, const char* name, void (*destructor)(capsule_t*));
+void* capsule_pointer(capsule_t* capsule, const char* name);
+
+typedef struct rc rc_t;
+
+rc_t* rc_new(void* ptr, void (*destructor)(void*));
+void rc_increment(rc_t* rc);
+void rc_decrement(rc_t* rc);
+
+VALUE schema_assert_valid(VALUE self);
+void assert_type_equal(VALUE actual, VALUE expected, VALUE check_nullability);
+
+void raise_message_not_ok(struct ArrowError* error, const char* message, int code);
+void raise_error_not_ok(const char* message, int code);
+void raise_todo(void);
+
+VALUE alloc_c_array(struct ArrowArray** c_array);
+VALUE alloc_c_array_stream(struct ArrowArrayStream** c_stream);
+VALUE alloc_c_array_view(struct ArrowArrayView** c_array_view);
+VALUE alloc_c_buffer(struct ArrowBuffer** c_buffer);
+VALUE alloc_c_schema(struct ArrowSchema** c_schema);
+void c_array_shallow_copy(VALUE base, const struct ArrowArray* src, struct ArrowArray* dst);

data/ext/nanoarrow/extconf.rb

@@ -0,0 +1,7 @@
+require "mkmf"
+
+append_cflags("-Wall -Wextra -Wpedantic")
+append_cflags("-Wno-suggest-attribute=format")
+append_cflags("-Wno-suggest-attribute=noreturn")
+
+create_makefile("nanoarrow/ext")