red-arrow 0.14.1 → 0.15.0

data/ext/arrow/raw-records.cpp

@@ -17,622 +17,16 @@
  * under the License.
  */
 
-#include "red-arrow.hpp"
-
-#include <ruby.hpp>
-#include <ruby/encoding.h>
-
-#include <arrow-glib/error.hpp>
-
-#include <arrow/util/logging.h>
+#include "converters.hpp"
 
 namespace red_arrow {
   namespace {
-    using Status = arrow::Status;
-
-    void check_status(const Status&& status, const char* context) {
-      GError* error = nullptr;
-      if (!garrow_error_check(&error, status, context)) {
-        RG_RAISE_ERROR(error);
-      }
-    }
-
-    class ListArrayValueConverter;
-    class StructArrayValueConverter;
-    class UnionArrayValueConverter;
-    class DictionaryArrayValueConverter;
-
-    class ArrayValueConverter {
-    public:
-      ArrayValueConverter()
-        : decimal_buffer_(),
-          list_array_value_converter_(nullptr),
-          struct_array_value_converter_(nullptr),
-          union_array_value_converter_(nullptr),
-          dictionary_array_value_converter_(nullptr) {
-      }
-
-      void set_sub_value_converters(ListArrayValueConverter* list_array_value_converter,
-                                    StructArrayValueConverter* struct_array_value_converter,
-                                    UnionArrayValueConverter* union_array_value_converter,
-                                    DictionaryArrayValueConverter* dictionary_array_value_converter) {
-        list_array_value_converter_ = list_array_value_converter;
-        struct_array_value_converter_ = struct_array_value_converter;
-        union_array_value_converter_ = union_array_value_converter;
-        dictionary_array_value_converter_ = dictionary_array_value_converter;
-      }
-
-      inline VALUE convert(const arrow::NullArray& array,
-                           const int64_t i) {
-        return Qnil;
-      }
-
-      inline VALUE convert(const arrow::BooleanArray& array,
-                           const int64_t i) {
-        return array.Value(i) ? Qtrue : Qfalse;
-      }
-
-      inline VALUE convert(const arrow::Int8Array& array,
-                           const int64_t i) {
-        return INT2NUM(array.Value(i));
-      }
-
-      inline VALUE convert(const arrow::Int16Array& array,
-                           const int64_t i) {
-        return INT2NUM(array.Value(i));
-      }
-
-      inline VALUE convert(const arrow::Int32Array& array,
-                           const int64_t i) {
-        return INT2NUM(array.Value(i));
-      }
-
-      inline VALUE convert(const arrow::Int64Array& array,
-                           const int64_t i) {
-        return LL2NUM(array.Value(i));
-      }
-
-      inline VALUE convert(const arrow::UInt8Array& array,
-                           const int64_t i) {
-        return UINT2NUM(array.Value(i));
-      }
-
-      inline VALUE convert(const arrow::UInt16Array& array,
-                           const int64_t i) {
-        return UINT2NUM(array.Value(i));
-      }
-
-      inline VALUE convert(const arrow::UInt32Array& array,
-                           const int64_t i) {
-        return UINT2NUM(array.Value(i));
-      }
-
-      inline VALUE convert(const arrow::UInt64Array& array,
-                           const int64_t i) {
-        return ULL2NUM(array.Value(i));
-      }
-
-      // TODO
-      // inline VALUE convert(const arrow::HalfFloatArray& array,
-      //                      const int64_t i) {
-      // }
-
-      inline VALUE convert(const arrow::FloatArray& array,
-                           const int64_t i) {
-        return DBL2NUM(array.Value(i));
-      }
-
-      inline VALUE convert(const arrow::DoubleArray& array,
-                           const int64_t i) {
-        return DBL2NUM(array.Value(i));
-      }
-
-      inline VALUE convert(const arrow::BinaryArray& array,
-                           const int64_t i) {
-        int32_t length;
-        const auto value = array.GetValue(i, &length);
-        // TODO: encoding support
-        return rb_enc_str_new(reinterpret_cast<const char*>(value),
-                              length,
-                              rb_ascii8bit_encoding());
-      }
-
-      inline VALUE convert(const arrow::StringArray& array,
-                           const int64_t i) {
-        int32_t length;
-        const auto value = array.GetValue(i, &length);
-        return rb_utf8_str_new(reinterpret_cast<const char*>(value),
-                               length);
-      }
-
-      inline VALUE convert(const arrow::FixedSizeBinaryArray& array,
-                           const int64_t i) {
-        return rb_enc_str_new(reinterpret_cast<const char*>(array.Value(i)),
-                              array.byte_width(),
-                              rb_ascii8bit_encoding());
-      }
-
-      constexpr static int32_t JULIAN_DATE_UNIX_EPOCH = 2440588;
-      inline VALUE convert(const arrow::Date32Array& array,
-                           const int64_t i) {
-        const auto value = array.Value(i);
-        const auto days_in_julian = value + JULIAN_DATE_UNIX_EPOCH;
-        return rb_funcall(cDate, id_jd, 1, LONG2NUM(days_in_julian));
-      }
-
-      inline VALUE convert(const arrow::Date64Array& array,
-                           const int64_t i) {
-        const auto value = array.Value(i);
-        auto msec = LL2NUM(value);
-        auto sec = rb_rational_new(msec, INT2NUM(1000));
-        auto time_value = rb_time_num_new(sec, Qnil);
-        return rb_funcall(time_value, id_to_datetime, 0, 0);
-      }
-
-      inline VALUE convert(const arrow::Time32Array& array,
-                           const int64_t i) {
-        // TODO: unit treatment
-        const auto value = array.Value(i);
-        return INT2NUM(value);
-      }
-
-      inline VALUE convert(const arrow::Time64Array& array,
-                           const int64_t i) {
-        // TODO: unit treatment
-        const auto value = array.Value(i);
-        return LL2NUM(value);
-      }
-
-      inline VALUE convert(const arrow::TimestampArray& array,
-                           const int64_t i) {
-        const auto type =
-          arrow::internal::checked_cast<const arrow::TimestampType*>(array.type().get());
-        auto scale = time_unit_to_scale(type->unit());
-        if (NIL_P(scale)) {
-          rb_raise(rb_eArgError, "Invalid TimeUnit");
-        }
-        auto value = array.Value(i);
-        auto sec = rb_rational_new(LL2NUM(value), scale);
-        return rb_time_num_new(sec, Qnil);
-      }
-
-      // TODO
-      // inline VALUE convert(const arrow::IntervalArray& array,
-      //                      const int64_t i) {
-      // };
-
-      VALUE convert(const arrow::ListArray& array,
-                    const int64_t i);
-
-      VALUE convert(const arrow::StructArray& array,
-                    const int64_t i);
-
-      VALUE convert(const arrow::UnionArray& array,
-                    const int64_t i);
-
-      VALUE convert(const arrow::DictionaryArray& array,
-                    const int64_t i);
-
-      inline VALUE convert(const arrow::Decimal128Array& array,
-                           const int64_t i) {
-        decimal_buffer_ = array.FormatValue(i);
-        return rb_funcall(rb_cObject,
-                          id_BigDecimal,
-                          1,
-                          rb_enc_str_new(decimal_buffer_.data(),
-                                         decimal_buffer_.length(),
-                                         rb_ascii8bit_encoding()));
-      }
-
-    private:
-      std::string decimal_buffer_;
-      ListArrayValueConverter* list_array_value_converter_;
-      StructArrayValueConverter* struct_array_value_converter_;
-      UnionArrayValueConverter* union_array_value_converter_;
-      DictionaryArrayValueConverter* dictionary_array_value_converter_;
-    };
-
-    class ListArrayValueConverter : public arrow::ArrayVisitor {
-    public:
-      explicit ListArrayValueConverter(ArrayValueConverter* converter)
-        : array_value_converter_(converter),
-          offset_(0),
-          length_(0),
-          result_(Qnil) {}
-
-      VALUE convert(const arrow::ListArray& array, const int64_t index) {
-        auto values = array.values().get();
-        auto offset_keep = offset_;
-        auto length_keep = length_;
-        offset_ = array.value_offset(index);
-        length_ = array.value_length(index);
-        auto result_keep = result_;
-        result_ = rb_ary_new_capa(length_);
-        check_status(values->Accept(this),
-                     "[raw-records][list-array]");
-        offset_ = offset_keep;
-        length_ = length_keep;
-        auto result_return = result_;
-        result_ = result_keep;
-        return result_return;
-      }
-
-#define VISIT(TYPE)                                                     \
-      Status Visit(const arrow::TYPE ## Array& array) override {        \
-        return visit_value(array);                                      \
-      }
-
-      VISIT(Null)
-      VISIT(Boolean)
-      VISIT(Int8)
-      VISIT(Int16)
-      VISIT(Int32)
-      VISIT(Int64)
-      VISIT(UInt8)
-      VISIT(UInt16)
-      VISIT(UInt32)
-      VISIT(UInt64)
-      // TODO
-      // VISIT(HalfFloat)
-      VISIT(Float)
-      VISIT(Double)
-      VISIT(Binary)
-      VISIT(String)
-      VISIT(FixedSizeBinary)
-      VISIT(Date32)
-      VISIT(Date64)
-      VISIT(Time32)
-      VISIT(Time64)
-      VISIT(Timestamp)
-      // TODO
-      // VISIT(Interval)
-      VISIT(List)
-      VISIT(Struct)
-      VISIT(Union)
-      VISIT(Dictionary)
-      VISIT(Decimal128)
-      // TODO
-      // VISIT(Extension)
-
-#undef VISIT
-
-    private:
-      template <typename ArrayType>
-      inline VALUE convert_value(const ArrayType& array,
-                                 const int64_t i) {
-        return array_value_converter_->convert(array, i);
-      }
-
-      template <typename ArrayType>
-      Status visit_value(const ArrayType& array) {
-        if (array.null_count() > 0) {
-          for (int64_t i = 0; i < length_; ++i) {
-            auto value = Qnil;
-            if (!array.IsNull(i + offset_)) {
-              value = convert_value(array, i + offset_);
-            }
-            rb_ary_push(result_, value);
-          }
-        } else {
-          for (int64_t i = 0; i < length_; ++i) {
-            rb_ary_push(result_, convert_value(array, i + offset_));
-          }
-        }
-        return Status::OK();
-      }
-
-      ArrayValueConverter* array_value_converter_;
-      int32_t offset_;
-      int32_t length_;
-      VALUE result_;
-    };
-
-    class StructArrayValueConverter : public arrow::ArrayVisitor {
-    public:
-      explicit StructArrayValueConverter(ArrayValueConverter* converter)
-        : array_value_converter_(converter),
-          key_(Qnil),
-          index_(0),
-          result_(Qnil) {}
-
-      VALUE convert(const arrow::StructArray& array,
-                    const int64_t index) {
-        auto index_keep = index_;
-        auto result_keep = result_;
-        index_ = index;
-        result_ = rb_hash_new();
-        const auto struct_type = array.struct_type();
-        const auto n = struct_type->num_children();
-        for (int i = 0; i < n; ++i) {
-          const auto field_type = struct_type->child(i).get();
-          const auto& field_name = field_type->name();
-          auto key_keep = key_;
-          key_ = rb_utf8_str_new(field_name.data(), field_name.length());
-          const auto field_array = array.field(i).get();
-          check_status(field_array->Accept(this),
-                       "[raw-records][struct-array]");
-          key_ = key_keep;
-        }
-        auto result_return = result_;
-        result_ = result_keep;
-        index_ = index_keep;
-        return result_return;
-      }
-
-#define VISIT(TYPE)                                                     \
-      Status Visit(const arrow::TYPE ## Array& array) override {        \
-        fill_field(array);                                              \
-        return Status::OK();                                            \
-      }
-
-      VISIT(Null)
-      VISIT(Boolean)
-      VISIT(Int8)
-      VISIT(Int16)
-      VISIT(Int32)
-      VISIT(Int64)
-      VISIT(UInt8)
-      VISIT(UInt16)
-      VISIT(UInt32)
-      VISIT(UInt64)
-      // TODO
-      // VISIT(HalfFloat)
-      VISIT(Float)
-      VISIT(Double)
-      VISIT(Binary)
-      VISIT(String)
-      VISIT(FixedSizeBinary)
-      VISIT(Date32)
-      VISIT(Date64)
-      VISIT(Time32)
-      VISIT(Time64)
-      VISIT(Timestamp)
-      // TODO
-      // VISIT(Interval)
-      VISIT(List)
-      VISIT(Struct)
-      VISIT(Union)
-      VISIT(Dictionary)
-      VISIT(Decimal128)
-      // TODO
-      // VISIT(Extension)
-
-#undef VISIT
-
-    private:
-      template <typename ArrayType>
-      inline VALUE convert_value(const ArrayType& array,
-                                 const int64_t i) {
-        return array_value_converter_->convert(array, i);
-      }
-
-      template <typename ArrayType>
-      void fill_field(const ArrayType& array) {
-        if (array.IsNull(index_)) {
-          rb_hash_aset(result_, key_, Qnil);
-        } else {
-          rb_hash_aset(result_, key_, convert_value(array, index_));
-        }
-      }
-
-      ArrayValueConverter* array_value_converter_;
-      VALUE key_;
-      int64_t index_;
-      VALUE result_;
-    };
-
-    class UnionArrayValueConverter : public arrow::ArrayVisitor {
-    public:
-      explicit UnionArrayValueConverter(ArrayValueConverter* converter)
-        : array_value_converter_(converter),
-          index_(0),
-          result_(Qnil) {}
-
-      VALUE convert(const arrow::UnionArray& array,
-                    const int64_t index) {
-        const auto index_keep = index_;
-        const auto result_keep = result_;
-        index_ = index;
-        switch (array.mode()) {
-        case arrow::UnionMode::SPARSE:
-          convert_sparse(array);
-          break;
-        case arrow::UnionMode::DENSE:
-          convert_dense(array);
-          break;
-        default:
-          rb_raise(rb_eArgError, "Invalid union mode");
-          break;
-        }
-        auto result_return = result_;
-        index_ = index_keep;
-        result_ = result_keep;
-        return result_return;
-      }
-
-#define VISIT(TYPE)                                                     \
-      Status Visit(const arrow::TYPE ## Array& array) override {        \
-        convert_value(array);                                           \
-        return Status::OK();                                            \
-      }
-
-      VISIT(Null)
-      VISIT(Boolean)
-      VISIT(Int8)
-      VISIT(Int16)
-      VISIT(Int32)
-      VISIT(Int64)
-      VISIT(UInt8)
-      VISIT(UInt16)
-      VISIT(UInt32)
-      VISIT(UInt64)
-      // TODO
-      // VISIT(HalfFloat)
-      VISIT(Float)
-      VISIT(Double)
-      VISIT(Binary)
-      VISIT(String)
-      VISIT(FixedSizeBinary)
-      VISIT(Date32)
-      VISIT(Date64)
-      VISIT(Time32)
-      VISIT(Time64)
-      VISIT(Timestamp)
-      // TODO
-      // VISIT(Interval)
-      VISIT(List)
-      VISIT(Struct)
-      VISIT(Union)
-      VISIT(Dictionary)
-      VISIT(Decimal128)
-      // TODO
-      // VISIT(Extension)
-
-#undef VISIT
-    private:
-      template <typename ArrayType>
-      inline void convert_value(const ArrayType& array) {
-        auto result = rb_hash_new();
-        if (array.IsNull(index_)) {
-          rb_hash_aset(result, field_name_, Qnil);
-        } else {
-          rb_hash_aset(result,
-                       field_name_,
-                       array_value_converter_->convert(array, index_));
-        }
-        result_ = result;
-      }
-
-      uint8_t compute_child_index(const arrow::UnionArray& array,
-                                  arrow::UnionType* type,
-                                  const char* tag) {
-        const auto type_id = array.raw_type_ids()[index_];
-        const auto& type_codes = type->type_codes();
-        for (uint8_t i = 0; i < type_codes.size(); ++i) {
-          if (type_codes[i] == type_id) {
-            return i;
-          }
-        }
-        check_status(Status::Invalid("Unknown type ID: ", type_id),
-                     tag);
-        return 0;
-      }
-
-      void convert_sparse(const arrow::UnionArray& array) {
-        const auto type =
-          std::static_pointer_cast<arrow::UnionType>(array.type()).get();
-        const auto tag = "[raw-records][union-sparse-array]";
-        const auto child_index = compute_child_index(array, type, tag);
-        const auto child_field = type->child(child_index).get();
-        const auto& field_name = child_field->name();
-        const auto field_name_keep = field_name_;
-        field_name_ = rb_utf8_str_new(field_name.data(), field_name.length());
-        const auto child_array = array.child(child_index).get();
-        check_status(child_array->Accept(this), tag);
-        field_name_ = field_name_keep;
-      }
-
-      void convert_dense(const arrow::UnionArray& array) {
-        const auto type =
-          std::static_pointer_cast<arrow::UnionType>(array.type()).get();
-        const auto tag = "[raw-records][union-dense-array]";
-        const auto child_index = compute_child_index(array, type, tag);
-        const auto child_field = type->child(child_index).get();
-        const auto& field_name = child_field->name();
-        const auto field_name_keep = field_name_;
-        field_name_ = rb_utf8_str_new(field_name.data(), field_name.length());
-        const auto child_array = array.child(child_index);
-        const auto index_keep = index_;
-        index_ = array.value_offset(index_);
-        check_status(child_array->Accept(this), tag);
-        index_ = index_keep;
-        field_name_ = field_name_keep;
-      }
-
-      ArrayValueConverter* array_value_converter_;
-      int64_t index_;
-      VALUE field_name_;
-      VALUE result_;
-    };
-
-    class DictionaryArrayValueConverter : public arrow::ArrayVisitor {
-    public:
-      explicit DictionaryArrayValueConverter(ArrayValueConverter* converter)
-        : array_value_converter_(converter),
-          index_(0),
-          result_(Qnil) {
-      }
-
-      VALUE convert(const arrow::DictionaryArray& array,
-                    const int64_t index) {
-        index_ = index;
-        auto indices = array.indices().get();
-        check_status(indices->Accept(this),
-                     "[raw-records][dictionary-array]");
-        return result_;
-      }
-
-      // TODO: Convert to real value.
-#define VISIT(TYPE)                                                     \
-      Status Visit(const arrow::TYPE ## Array& array) override {        \
-        result_ = convert_value(array, index_);                         \
-        return Status::OK();                                            \
-      }
-
-      VISIT(Int8)
-      VISIT(Int16)
-      VISIT(Int32)
-      VISIT(Int64)
-
-#undef VISIT
-
-    private:
-      template <typename ArrayType>
-      inline VALUE convert_value(const ArrayType& array,
-                                 const int64_t i) {
-        return array_value_converter_->convert(array, i);
-      }
-
-      ArrayValueConverter* array_value_converter_;
-      int64_t index_;
-      VALUE result_;
-    };
-
-    VALUE ArrayValueConverter::convert(const arrow::ListArray& array,
-                                       const int64_t i) {
-      return list_array_value_converter_->convert(array, i);
-    }
-
-    VALUE ArrayValueConverter::convert(const arrow::StructArray& array,
-                                       const int64_t i) {
-      return struct_array_value_converter_->convert(array, i);
-    }
-
-    VALUE ArrayValueConverter::convert(const arrow::UnionArray& array,
-                                       const int64_t i) {
-      return union_array_value_converter_->convert(array, i);
-    }
-
-    VALUE ArrayValueConverter::convert(const arrow::DictionaryArray& array,
-                                       const int64_t i) {
-      return dictionary_array_value_converter_->convert(array, i);
-    }
-
-    class RawRecordsBuilder : public arrow::ArrayVisitor {
+    class RawRecordsBuilder : private Converter, public arrow::ArrayVisitor {
     public:
       explicit RawRecordsBuilder(VALUE records, int n_columns)
-        : array_value_converter_(),
-          list_array_value_converter_(&array_value_converter_),
-          struct_array_value_converter_(&array_value_converter_),
-          union_array_value_converter_(&array_value_converter_),
-          dictionary_array_value_converter_(&array_value_converter_),
+        : Converter(),
           records_(records),
           n_columns_(n_columns) {
-        array_value_converter_.
-          set_sub_value_converters(&list_array_value_converter_,
-                                   &struct_array_value_converter_,
-                                   &union_array_value_converter_,
-                                   &dictionary_array_value_converter_);
       }
 
       void build(const arrow::RecordBatch& record_batch) {
@@ -661,8 +55,7 @@ namespace red_arrow {
             rb_ary_push(records_, record);
           }
           for (int i = 0; i < n_columns_; ++i) {
-            const auto column = table.column(i).get();
-            const auto chunked_array = column->data();
+            const auto& chunked_array = table.column(i).get();
             column_index_ = i;
             row_offset_ = 0;
             for (const auto array : chunked_array->chunks()) {
@@ -676,9 +69,9 @@ namespace red_arrow {
       }
 
 #define VISIT(TYPE)                                                     \
-      Status Visit(const arrow::TYPE ## Array& array) override {        \
+      arrow::Status Visit(const arrow::TYPE ## Array& array) override { \
         convert(array);                                                 \
-        return Status::OK();                                            \
+        return arrow::Status::OK();                                     \
       }
 
       VISIT(Null)
@@ -716,12 +109,6 @@ namespace red_arrow {
 #undef VISIT
 
     private:
-      template <typename ArrayType>
-      inline VALUE convert_value(const ArrayType& array,
-                                 const int64_t i) {
-        return array_value_converter_.convert(array, i);
-      }
-
       template <typename ArrayType>
       void convert(const ArrayType& array) {
         const auto n = array.length();
@@ -742,12 +129,6 @@ namespace red_arrow {
         }
       }
 
-      ArrayValueConverter array_value_converter_;
-      ListArrayValueConverter list_array_value_converter_;
-      StructArrayValueConverter struct_array_value_converter_;
-      UnionArrayValueConverter union_array_value_converter_;
-      DictionaryArrayValueConverter dictionary_array_value_converter_;
-
       // Destination for converted records.
       VALUE records_;