nanoarrow 0.1.0

data/ext/nanoarrow/rc.c

@@ -0,0 +1,34 @@
+#include <stdatomic.h>
+#include <stdlib.h>
+
+#include "ext.h"
+
+struct rc {
+    void* ptr;
+    _Atomic int count;
+    void (*destructor)(void*);
+};
+
+rc_t* rc_new(void* ptr, void (*destructor)(void*))
+{
+    rc_t* rc = malloc(sizeof(rc_t));
+    rc->ptr = ptr;
+    rc->count = 1;
+    rc->destructor = destructor;
+    return rc;
+}
+
+void rc_increment(rc_t* rc)
+{
+    rc->count++;
+}
+
+void rc_decrement(rc_t* rc)
+{
+    int count = --rc->count;
+    if (count == 0)
+    {
+        rc->destructor(rc->ptr);
+        free(rc);
+    }
+}

data/ext/nanoarrow/schema.c

@@ -0,0 +1,346 @@
+#include <ruby/ruby.h>
+
+#include "ext.h"
+#include "nanoarrow/nanoarrow.h"
+
+VALUE cCSchema;
+
+static void schema_mark(void* ptr)
+{
+    schema_t* schema = (schema_t*) ptr;
+    rb_gc_mark_movable(schema->base);
+}
+
+const rb_data_type_t schema_data_type = {
+    .wrap_struct_name = "Nanoarrow::CSchema",
+    .function = {
+        .dmark = schema_mark,
+        .dfree = RUBY_TYPED_DEFAULT_FREE,
+    },
+    .flags = RUBY_TYPED_FREE_IMMEDIATELY
+};
+
+static VALUE schema_allocate(VALUE klass)
+{
+    schema_t* schema;
+    VALUE obj = TypedData_Make_Struct(klass, schema_t, &schema_data_type, schema);
+    schema->base = alloc_c_schema(&schema->ptr);
+    return obj;
+}
+
+static VALUE schema_import_from_c_capsule(VALUE klass, VALUE schema_capsule)
+{
+    schema_t* schema;
+    VALUE obj = TypedData_Make_Struct(klass, schema_t, &schema_data_type, schema);
+    schema->base = schema_capsule;
+    schema->ptr = (struct ArrowSchema*) NUM2ULL(rb_funcall(schema_capsule, rb_intern("to_i"), 0));
+    return obj;
+}
+
+static VALUE schema_deep_dup(VALUE self)
+{
+    schema_t* schema;
+    GetSchema(self, schema);
+
+    schema_t* out_ptr;
+    VALUE out = schema_allocate(cCSchema);
+    GetSchema(out, out_ptr);
+
+    int code = ArrowSchemaDeepCopy(schema->ptr, out_ptr->ptr);
+    raise_error_not_ok("ArrowSchemaDeepCopy()", code);
+    return out;
+}
+
+VALUE schema_assert_valid(VALUE self)
+{
+    schema_t* schema;
+    GetSchema(self, schema);
+
+    if (schema->ptr == NULL)
+        rb_raise(rb_eRuntimeError, "schema is NULL");
+    if (schema->ptr->release == NULL)
+        rb_raise(rb_eRuntimeError, "schema is released");
+
+    return Qnil;
+}
+
+static VALUE schema_type_equals(VALUE self, VALUE other, VALUE check_nullability)
+{
+    schema_t* self_ptr;
+    GetSchema(self, self_ptr);
+
+    schema_assert_valid(self);
+
+    schema_t* other_ptr;
+    GetSchema(other, other_ptr);
+
+    if (self_ptr->ptr == other_ptr->ptr)
+        return Qtrue;
+
+    // format
+
+    VALUE format = rb_funcall(self, rb_intern("format"), 0);
+    VALUE other_format = rb_funcall(other, rb_intern("format"), 0);
+
+    if (rb_equal(format, other_format) != Qtrue)
+        return Qfalse;
+
+    // flags
+
+    int64_t flags = NUM2LL(rb_funcall(self, rb_intern("flags"), 0));
+    int64_t other_flags = NUM2LL(rb_funcall(other, rb_intern("flags"), 0));
+
+    if (!RTEST(check_nullability))
+    {
+        flags &= ~ARROW_FLAG_NULLABLE;
+        other_flags &= ~ARROW_FLAG_NULLABLE;
+    }
+
+    if (flags != other_flags)
+        return Qfalse;
+
+    // children
+
+    int64_t n_children = NUM2LL(rb_funcall(self, rb_intern("n_children"), 0));
+    int64_t other_n_children = NUM2LL(rb_funcall(other, rb_intern("n_children"), 0));
+
+    if (n_children != other_n_children)
+        return Qfalse;
+
+    for (int64_t i = 0; i < n_children; i++)
+    {
+        VALUE child = rb_funcall(self, rb_intern("child"), 1, LL2NUM(i));
+        VALUE other_child = rb_funcall(other, rb_intern("child"), 1, LL2NUM(i));
+
+        if (schema_type_equals(child, other_child, check_nullability) != Qtrue)
+            return Qfalse;
+    }
+
+    // dictionary
+
+    VALUE dictionary = rb_funcall(self, rb_intern("dictionary"), 0);
+    VALUE other_dictionary = rb_funcall(other, rb_intern("dictionary"), 0);
+
+    if (NIL_P(dictionary) != NIL_P(other_dictionary))
+        return Qfalse;
+
+    if (!NIL_P(dictionary))
+    {
+        if (schema_type_equals(dictionary, other_dictionary, check_nullability) != Qtrue)
+            return Qfalse;
+    }
+
+    return Qtrue;
+}
+
+void assert_type_equal(VALUE actual, VALUE expected, VALUE check_nullability)
+{
+    if (!rb_obj_is_instance_of(actual, cCSchema))
+        rb_raise(rb_eTypeError, "actual is not CSchema");
+
+    if (!rb_obj_is_instance_of(expected, cCSchema))
+        rb_raise(rb_eTypeError, "actual is not CSchema");
+
+    if (!RTEST(schema_type_equals(actual, expected, check_nullability)))
+        rb_raise(rb_eArgError, "schema mismatch");
+}
+
+static VALUE schema_format(VALUE self)
+{
+    schema_t* schema;
+    GetSchema(self, schema);
+
+    schema_assert_valid(self);
+
+    return schema->ptr->format != NULL ? rb_utf8_str_new_cstr(schema->ptr->format) : Qnil;
+}
+
+static VALUE schema_name(VALUE self)
+{
+    schema_t* schema;
+    GetSchema(self, schema);
+
+    schema_assert_valid(self);
+
+    return schema->ptr->name != NULL ? rb_utf8_str_new_cstr(schema->ptr->name) : Qnil;
+}
+
+static VALUE schema_flags(VALUE self)
+{
+    schema_t* schema;
+    GetSchema(self, schema);
+
+    return LL2NUM(schema->ptr->flags);
+}
+
+static VALUE schema_metadata(VALUE self)
+{
+    schema_t* schema;
+    GetSchema(self, schema);
+
+    schema_assert_valid(self);
+
+    if (schema->ptr->metadata == NULL)
+        return rb_hash_new();
+
+    struct ArrowMetadataReader reader;
+    int code = ArrowMetadataReaderInit(&reader, schema->ptr->metadata);
+    raise_error_not_ok("ArrowMetadataReaderInit()", code);
+
+    VALUE items = rb_hash_new_capa(reader.remaining_keys);
+    while (reader.remaining_keys > 0)
+    {
+        struct ArrowStringView key;
+        struct ArrowStringView value;
+        code = ArrowMetadataReaderRead(&reader, &key, &value);
+        raise_error_not_ok("ArrowMetadataReaderRead()", code);
+        rb_hash_aset(items, rb_str_new(key.data, key.size_bytes), rb_str_new(value.data, value.size_bytes));
+    }
+    return items;
+}
+
+static VALUE schema_n_children(VALUE self)
+{
+    schema_t* schema;
+    GetSchema(self, schema);
+
+    schema_assert_valid(self);
+
+    return LL2NUM(schema->ptr->n_children);
+}
+
+static VALUE schema_child(VALUE self, VALUE rb_i)
+{
+    schema_t* schema;
+    GetSchema(self, schema);
+
+    schema_assert_valid(self);
+
+    int64_t i = NUM2LL(rb_i);
+    if (i < 0 || i >= schema->ptr->n_children)
+        rb_raise(rb_eIndexError, "out of range");
+
+    schema_t* schema_ptr;
+    VALUE obj = TypedData_Make_Struct(cCSchema, schema_t, &schema_data_type, schema_ptr);
+    schema_ptr->base = schema->base;
+    schema_ptr->ptr = schema->ptr->children[i];
+    return obj;
+}
+
+static VALUE schema_dictionary(VALUE self)
+{
+    schema_t* schema;
+    GetSchema(self, schema);
+
+    schema_assert_valid(self);
+
+    if (schema->ptr->dictionary != NULL)
+        raise_todo();
+
+    return Qnil;
+}
+
+static VALUE schema_to_s(VALUE self)
+{
+    schema_t* schema;
+    GetSchema(self, schema);
+
+    int64_t n_chars = 1024;
+    bool recursive = true;
+
+    char* out = ArrowMalloc(n_chars + 1);
+    if (out == NULL)
+        rb_raise(rb_eNoMemError, "out of memory");
+
+    int64_t len = ArrowSchemaToString(schema->ptr, out, n_chars + 1, recursive);
+    VALUE out_str = rb_str_new(out, len);
+    ArrowFree(out);
+    return out_str;
+}
+
+static VALUE schema_arrow_c_schema(VALUE self)
+{
+    schema_t* schema;
+    GetSchema(self, schema);
+
+    schema_assert_valid(self);
+
+    struct ArrowSchema* c_schema_out;
+    VALUE schema_capsule = alloc_c_schema(&c_schema_out);
+
+    int code = ArrowSchemaDeepCopy(schema->ptr, c_schema_out);
+    raise_error_not_ok("ArrowSchemaDeepCopy()", code);
+    return schema_capsule;
+}
+
+static VALUE schema_modify(VALUE self, VALUE format, VALUE name, VALUE flags, VALUE nullable, VALUE metadata, VALUE children, VALUE dictionary, VALUE validate)
+{
+    VALUE builder = rb_funcall(cCSchemaBuilder, rb_intern("allocate"), 0);
+
+    if (NIL_P(format))
+        rb_funcall(builder, rb_intern("set_format"), 1, rb_funcall(self, rb_intern("format"), 0));
+    else
+        rb_funcall(builder, rb_intern("set_format"), 1, format);
+
+    if (NIL_P(name))
+        rb_funcall(builder, rb_intern("set_name"), 1, rb_funcall(self, rb_intern("name"), 0));
+    else
+        rb_funcall(builder, rb_intern("set_name"), 1, name);
+
+    if (NIL_P(flags))
+        rb_funcall(builder, rb_intern("set_flags"), 1, rb_funcall(self, rb_intern("flags"), 0));
+    else
+        rb_funcall(builder, rb_intern("set_flags"), 1, flags);
+
+    if (!NIL_P(nullable))
+        rb_funcall(builder, rb_intern("set_nullable"), 1, nullable);
+
+    if (NIL_P(metadata))
+    {
+        VALUE self_metadata = rb_funcall(self, rb_intern("metadata"), 0);
+        if (!NIL_P(self_metadata))
+            rb_funcall(builder, rb_intern("append_metadata"), 1, self_metadata);
+    }
+    else
+        rb_funcall(builder, rb_intern("append_metadata"), 1, metadata);
+
+    if (NIL_P(children))
+    {
+        if (NUM2LL(rb_funcall(self, rb_intern("n_children"), 0)) > 0)
+            raise_todo();
+    }
+    else
+        raise_todo();
+
+    if (NIL_P(dictionary))
+    {
+        if (!NIL_P(rb_funcall(self, rb_intern("dictionary"), 0)))
+            raise_todo();
+    }
+    else
+        raise_todo();
+
+    if (RTEST(validate))
+        rb_funcall(builder, rb_intern("validate"), 0);
+
+    return rb_funcall(builder, rb_intern("finish"), 0);
+}
+
+void Init_schema(void)
+{
+    cCSchema = rb_define_class_under(mNanoarrow, "CSchema", rb_cObject);
+    rb_define_alloc_func(cCSchema, schema_allocate);
+    rb_define_singleton_method(cCSchema, "import_from_c_capsule", schema_import_from_c_capsule, 1);
+    rb_define_method(cCSchema, "deep_dup", schema_deep_dup, 0);
+    rb_define_method(cCSchema, "format", schema_format, 0);
+    rb_define_method(cCSchema, "name", schema_name, 0);
+    rb_define_method(cCSchema, "flags", schema_flags, 0);
+    rb_define_method(cCSchema, "metadata", schema_metadata, 0);
+    rb_define_method(cCSchema, "n_children", schema_n_children, 0);
+    rb_define_method(cCSchema, "child", schema_child, 1);
+    rb_define_method(cCSchema, "dictionary", schema_dictionary, 0);
+    rb_define_method(cCSchema, "to_s", schema_to_s, 0);
+    rb_define_method(cCSchema, "arrow_c_schema", schema_arrow_c_schema, 0);
+    rb_define_method(cCSchema, "modify", schema_modify, 8);
+    rb_define_method(cCSchema, "assert_valid", schema_assert_valid, 0);
+}

data/ext/nanoarrow/schema_builder.c

@@ -0,0 +1,251 @@
+#include <ruby/ruby.h>
+
+#include "ext.h"
+#include "nanoarrow/nanoarrow.h"
+
+VALUE cCSchemaBuilder;
+
+typedef struct schema_builder {
+    VALUE c_schema;
+    struct ArrowSchema* ptr;
+} schema_builder_t;
+
+static void schema_builder_mark(void* ptr)
+{
+    schema_builder_t* builder = (schema_builder_t*) ptr;
+    rb_gc_mark_movable(builder->c_schema);
+}
+
+const rb_data_type_t schema_builder_data_type = {
+    .wrap_struct_name = "Nanoarrow::CSchemaBuilder",
+    .function = {
+        .dmark = schema_builder_mark,
+        .dfree = RUBY_TYPED_DEFAULT_FREE,
+    },
+    .flags = RUBY_TYPED_FREE_IMMEDIATELY
+};
+
+VALUE schema_builder_allocate(VALUE klass)
+{
+    schema_builder_t* builder;
+    VALUE obj = TypedData_Make_Struct(klass, schema_builder_t, &schema_builder_data_type, builder);
+
+    schema_t* schema_ptr;
+    VALUE schema = rb_funcall(cCSchema, rb_intern("allocate"), 0);
+    GetSchema(schema, schema_ptr);
+
+    builder->c_schema = schema;
+    builder->ptr = schema_ptr->ptr;
+    if (builder->ptr->release == NULL)
+        ArrowSchemaInit(builder->ptr);
+
+    return obj;
+}
+
+static int foreach_metadata(st_data_t key_data, st_data_t value_data, st_data_t buffer_data)
+{
+    VALUE k = (VALUE) key_data;
+    VALUE v = (VALUE) value_data;
+    buffer_t* buffer;
+    GetBuffer((VALUE) buffer_data, buffer);
+
+    struct ArrowStringView key;
+    struct ArrowStringView value;
+
+    Check_Type(k, T_STRING);
+    key.data = StringValuePtr(k);
+    key.size_bytes = RSTRING_LEN(k);
+
+    Check_Type(v, T_STRING);
+    value.data = StringValuePtr(v);
+    value.size_bytes = RSTRING_LEN(v);
+
+    int code = ArrowMetadataBuilderAppend(buffer->ptr, key, value);
+    raise_error_not_ok("ArrowMetadataBuilderAppend()", code);
+
+    return ST_CONTINUE;
+}
+
+static VALUE schema_builder_append_metadata(VALUE self, VALUE metadata)
+{
+    schema_builder_t* builder;
+    GetSchemaBuilder(self, builder);
+
+    Check_Type(metadata, T_HASH);
+
+    if (RHASH_SIZE(metadata) > 0)
+    {
+        VALUE buffer = rb_funcall(cCBuffer, rb_intern("empty"), 0);
+        buffer_t* buffer_ptr;
+        GetBuffer(buffer, buffer_ptr);
+
+        schema_t* c_schema;
+        GetSchema(builder->c_schema, c_schema);
+
+        const char* existing_metadata = c_schema->ptr->metadata;
+        int code = ArrowMetadataBuilderInit(buffer_ptr->ptr, existing_metadata);
+        raise_error_not_ok("ArrowMetadataBuilderInit()", code);
+
+        rb_hash_foreach(metadata, foreach_metadata, buffer);
+
+        code = ArrowSchemaSetMetadata(c_schema->ptr, (const char*) buffer_ptr->ptr->data);
+        raise_error_not_ok("ArrowSchemaSetMetadata()", code);
+    }
+
+    return self;
+}
+
+static VALUE schema_builder_set_type(VALUE self, VALUE type_id)
+{
+    schema_builder_t* builder;
+    GetSchemaBuilder(self, builder);
+
+    schema_assert_valid(builder->c_schema);
+
+    int code = ArrowSchemaSetType(builder->ptr, NUM2INT(type_id));
+    raise_error_not_ok("ArrowSchemaSetType()", code);
+
+    return self;
+}
+
+static VALUE schema_builder_set_type_decimal(VALUE self, VALUE type_id, VALUE precision, VALUE scale)
+{
+    schema_builder_t* builder;
+    GetSchemaBuilder(self, builder);
+
+    schema_assert_valid(builder->c_schema);
+
+    int code = ArrowSchemaSetTypeDecimal(builder->ptr, NUM2INT(type_id), NUM2INT(precision), NUM2INT(scale));
+    raise_error_not_ok("ArrowSchemaSetTypeDecimal()", code);
+
+    return self;
+}
+
+static VALUE schema_builder_set_type_fixed_size(VALUE self, VALUE type_id, VALUE fixed_size)
+{
+    schema_builder_t* builder;
+    GetSchemaBuilder(self, builder);
+
+    schema_assert_valid(builder->c_schema);
+
+    int code = ArrowSchemaSetTypeFixedSize(builder->ptr, NUM2INT(type_id), NUM2INT(fixed_size));
+    raise_error_not_ok("ArrowSchemaSetTypeFixedSize()", code);
+
+    return self;
+}
+
+static VALUE schema_builder_set_type_date_time(VALUE self, VALUE type_id, VALUE time_unit, VALUE timezone)
+{
+    schema_builder_t* builder;
+    GetSchemaBuilder(self, builder);
+
+    schema_assert_valid(builder->c_schema);
+
+    int code;
+    if (NIL_P(timezone))
+        code = ArrowSchemaSetTypeDateTime(builder->ptr, NUM2INT(type_id), NUM2INT(time_unit), NULL);
+    else
+        code = ArrowSchemaSetTypeDateTime(builder->ptr, NUM2INT(type_id), NUM2INT(time_unit), StringValueCStr(timezone));
+
+    raise_error_not_ok("ArrowSchemaSetTypeDateTime()", code);
+
+    return self;
+}
+
+static VALUE schema_builder_set_format(VALUE self, VALUE format)
+{
+    schema_builder_t* builder;
+    GetSchemaBuilder(self, builder);
+
+    schema_assert_valid(builder->c_schema);
+
+    int code = ArrowSchemaSetFormat(builder->ptr, StringValueCStr(format));
+    raise_error_not_ok("ArrowSchemaSetFormat()", code);
+
+    return self;
+}
+
+static VALUE schema_builder_set_name(VALUE self, VALUE name)
+{
+    schema_builder_t* builder;
+    GetSchemaBuilder(self, builder);
+
+    schema_assert_valid(builder->c_schema);
+
+    int code;
+    if (NIL_P(name))
+        code = ArrowSchemaSetName(builder->ptr, NULL);
+    else
+        code = ArrowSchemaSetName(builder->ptr, StringValueCStr(name));
+
+    raise_error_not_ok("ArrowSchemaSetName()", code);
+
+    return self;
+}
+
+static VALUE schema_builder_set_flags(VALUE self, VALUE flags)
+{
+    schema_builder_t* builder;
+    GetSchemaBuilder(self, builder);
+
+    builder->ptr->flags = NUM2LL(flags);
+    return self;
+}
+
+static VALUE schema_builder_set_nullable(VALUE self, VALUE nullable)
+{
+    schema_builder_t* builder;
+    GetSchemaBuilder(self, builder);
+
+    if (RTEST(nullable))
+        builder->ptr->flags |= ARROW_FLAG_NULLABLE;
+    else
+        builder->ptr->flags &= ~ARROW_FLAG_NULLABLE;
+
+    return self;
+}
+
+static VALUE schema_builder_validate(VALUE self)
+{
+    schema_builder_t* builder;
+    GetSchemaBuilder(self, builder);
+
+    return rb_funcall(cCSchemaView, rb_intern("new"), 1, builder->c_schema);
+}
+
+static VALUE schema_builder_finish(VALUE self)
+{
+    schema_builder_t* builder;
+    GetSchemaBuilder(self, builder);
+
+    schema_assert_valid(builder->c_schema);
+
+    VALUE out = rb_funcall(cCSchema, rb_intern("allocate"), 0);
+
+    schema_t* c_schema_ptr;
+    GetSchema(builder->c_schema, c_schema_ptr);
+
+    schema_t* out_ptr;
+    GetSchema(out, out_ptr);
+
+    ArrowSchemaMove(c_schema_ptr->ptr, out_ptr->ptr);
+    ArrowSchemaInit(c_schema_ptr->ptr);
+    return out;
+}
+
+void Init_schema_builder(void)
+{
+    cCSchemaBuilder = rb_define_class_under(mNanoarrow, "CSchemaBuilder", rb_cObject);
+    rb_define_alloc_func(cCSchemaBuilder, schema_builder_allocate);
+    rb_define_method(cCSchemaBuilder, "append_metadata", schema_builder_append_metadata, 1);
+    rb_define_method(cCSchemaBuilder, "set_type", schema_builder_set_type, 1);
+    rb_define_method(cCSchemaBuilder, "set_type_decimal", schema_builder_set_type_decimal, 3);
+    rb_define_method(cCSchemaBuilder, "set_type_fixed_size", schema_builder_set_type_fixed_size, 2);
+    rb_define_method(cCSchemaBuilder, "set_type_date_time", schema_builder_set_type_date_time, 3);
+    rb_define_method(cCSchemaBuilder, "set_format", schema_builder_set_format, 1);
+    rb_define_method(cCSchemaBuilder, "set_name", schema_builder_set_name, 1);
+    rb_define_method(cCSchemaBuilder, "set_flags", schema_builder_set_flags, 1);
+    rb_define_method(cCSchemaBuilder, "set_nullable", schema_builder_set_nullable, 1);
+    rb_define_method(cCSchemaBuilder, "validate", schema_builder_validate, 0);
+    rb_define_method(cCSchemaBuilder, "finish", schema_builder_finish, 0);
+}