WIP JSON array reader code path

Change-Id: I7c52a441edcef57a7a868c80f0caa2f4ac734f22

Author: wesm

cpp/src/arrow/ipc/json-internal.cc

@@ -33,6 +33,7 @@
 #include "arrow/types/primitive.h"
 #include "arrow/types/string.h"
 #include "arrow/types/struct.h"
+#include "arrow/util/memory-pool.h"
 #include "arrow/util/status.h"
 
 namespace arrow {
@@ -847,23 +848,163 @@ class JsonSchemaReader {
   const rj::Value& json_schema_;
 };
 
-// class JsonArrayReader {
-//  public:
-//   explicit JsonArrayReader(const rj::Value& json_array, const Schema& schema)
-//       : json_array_(json_array), schema_(schema) {}
+class JsonArrayReader {
+ public:
+  explicit JsonArrayReader(
+      MemoryPool* pool, const rj::Value& json_array, const Schema& schema)
+      : pool_(pool), json_array_(json_array), schema_(schema) {}
+
+  Status GetResult(std::shared_ptr<Array>* array) {
+    if (!json_array_.IsObject()) {
+      return Status::Invalid("Array was not a JSON object");
+    }
+    const auto& json_array = json_array_.GetObject();
+
+    const auto& json_name = json_array.FindMember("name");
+    RETURN_NOT_STRING("name", json_name, json_array);
+
+    return GetArrayFromStruct(
+        json_array_, json_name.GetString(), schema_.fields(), array);
+  }
+
+  Status GetArrayFromStruct(const rj::Value& obj, const std::string& name,
+      const std::vector<std::shared_ptr<Field>>& fields, std::shared_ptr<Array>* array) {
+    std::shared_ptr<Field> result = nullptr;
+
+    for (const std::shared_ptr<Field>& field : fields) {
+      if (field->name == name) {
+        result = field;
+        break;
+      }
+    }
+
+    if (result == nullptr) {
+      std::stringstream ss;
+      ss << "Field named " << name << " not found in struct/schema";
+      return Status::KeyError(ss.str());
+    }
+
+    return GetArray(obj, result->type, array);
+  }
+
+  template <typename T>
+  Status ReadArray(const rj::Value& obj, const std::vector<bool>& is_valid,
+      const std::shared_ptr<DataType>& type, std::shared_ptr<Array>* array) {
+    typename TypeTraits<T>::BuilderType builder(pool_, type);
+    const auto& json_array = obj.GetObject();
 
-//   Status GetArray(std::shared_ptr<Array>* array) {
-//     if (!json_array_.IsObject()) {
-//       return Status::Invalid("Array was not a JSON object");
-//     }
+    const auto& json_data = json_array.FindMember("DATA");
+    RETURN_NOT_ARRAY("DATA", json_data, json_array);
 
-//     return Status::OK();
-//   }
+    const auto& json_data_arr = json_type_ids->value.GetArray();
 
-//  private:
-//   const rj::Value& json_array_;
-//   const Schema& schema_;
-// };
+    for (auto i = 0; i < json_data_arr.Size(); ++i) {
+      if (!is_valid[i]) {
+        builder.AppendNull();
+        continue;
+      }
+
+      const rj::Value& val = json_data_arr[i];
+      if (IsSignedInt<T>::value) {
+        builder.Append(val.GetInt64());
+      } else if (IsUnsignedInt<T>::value) {
+        builder.Append(val.GetUint64());
+      } else if (IsFloatingPoint<T>::value) {
+        builder.Append(val.GetFloat());
+      } else if (std::is_base_of<BooleanType, T>::value) {
+        builder.Append(val.GetBool());
+      } else {
+        // We are in the wrong function
+        return Status::Invalid(type->ToString());
+      }
+    }
+
+    return builder.Finish(array);
+  }
+
+  template <typename T>
+  typename std::enable_if<std::is_base_of<BinaryType, T>::value, Status>::type
+  ReadArray(const rj::Value& obj, const std::shared_ptr<DataType>& type,
+      std::shared_ptr<Array>* array) {}
+
+  template <typename T>
+  typename std::enable_if<std::is_base_of<ListType, T>::value, Status>::type
+  ReadArray(const rj::Value& obj, const std::shared_ptr<DataType>& type,
+      std::shared_ptr<Array>* array) {}
+
+  template <typename T>
+  typename std::enable_if<std::is_base_of<StructType, T>::value, Status>::type
+  ReadArray(const rj::Value& obj, const std::shared_ptr<DataType>& type,
+      std::shared_ptr<Array>* array) {}
+
+  template <typename T>
+  typename std::enable_if<std::is_base_of<NullType, T>::value, Status>::type
+  ReadArray(const rj::Value& obj, const std::shared_ptr<DataType>& type,
+      std::shared_ptr<Array>* array) {}
+
+  Status GetArray(const rj::Value& obj, const std::shared_ptr<DataType>& type,
+      std::shared_ptr<Array>* array) {
+    if (!obj.IsObject()) { return Status::Invalid("Array was not a JSON object"); }
+    const auto& json_array = obj.GetObject();
+
+    const auto& json_length = json_array.FindMember("count");
+    RETURN_NOT_INT("count", json_length, json_array);
+
+    const auto& json_validity = json_array.FindMember("VALIDITY");
+    RETURN_NOT_ARRAY("VALIDITY", json_validity, json_array);
+
+    std::vector<bool> is_valid(count);
+
+#define TYPE_CASE(TYPE)                         \
+    case TYPE::type_enum:                       \
+      return ReadArray<TYPE>(obj, type, array);
+
+#define NOT_IMPLEMENTED_CASE(TYPE_ENUM)         \
+    case Type::TYPE_ENUM:                       \
+      std::stringstream ss;                     \
+      ss << type->ToString();                   \
+      return Status::NotImplemented(ss.str());
+
+    switch (type->type) {
+      TYPE_CASE(NullType);
+      TYPE_CASE(BooleanType);
+      TYPE_CASE(UInt8Type);
+      TYPE_CASE(Int8Type);
+      TYPE_CASE(UInt16Type);
+      TYPE_CASE(Int16Type);
+      TYPE_CASE(UInt32Type);
+      TYPE_CASE(Int32Type);
+      TYPE_CASE(UInt64Type);
+      TYPE_CASE(Int64Type);
+      TYPE_CASE(HalfFloatType);
+      TYPE_CASE(FloatType);
+      TYPE_CASE(DoubleType);
+      TYPE_CASE(StringType);
+      TYPE_CASE(BinaryType);
+      NOT_IMPLEMENTED_CASE(DATE);
+      NOT_IMPLEMENTED_CASE(TIMESTAMP);
+      NOT_IMPLEMENTED_CASE(TIME);
+      NOT_IMPLEMENTED_CASE(INTERVAL);
+      TYPE_CASE(ListType);
+      TYPE_CASE(StructType);
+      NOT_IMPLEMENTED_CASE(UNION);
+      default:
+        std::stringstream ss;
+        ss << type->ToString();
+        return Status::NotImplemented(ss.str());
+    };
+
+#undef TYPE_CASE
+#undef NOT_IMPLEMENTED_CASE
+
+    return Status::OK();
+  }
+
+ private:
+  MemoryPool* pool;
+  const rj::Value& json_array_;
+  const Schema& schema_;
+};
 
 Status WriteJsonSchema(const Schema& schema, RjWriter* json_writer) {
   JsonSchemaWriter converter(schema, json_writer);
@@ -875,17 +1016,17 @@ Status ReadJsonSchema(const rj::Value& json_schema, std::shared_ptr<Schema>* sch
   return converter.GetSchema(schema);
 }
 
-// Status WriteJsonArray(
-//     const std::string& name, const Array& array, RjWriter* json_writer) {
-//   JsonArrayWriter converter(name, array, json_writer);
-//   converter.Write();
-// }
-
-// Status ReadJsonArray(
-//     const rj::Value& json_array, const Schema& schema, std::shared_ptr<Array>* array) {
-//   JsonArrayReader converter(json_array, schema);
-//   return converter.GetArray(array);
-// }
+Status WriteJsonArray(
+    const std::string& name, const Array& array, RjWriter* json_writer) {
+  JsonArrayWriter converter(name, array, json_writer);
+  return converter.Write();
+}
+
+Status ReadJsonArray(MemoryPool* pool, const rj::Value& json_array, const Schema& schema,
+    std::shared_ptr<Array>* array) {
+  JsonArrayReader converter(pool, json_array, schema);
+  return converter.GetArray(array);
+}
 
 }  // namespace ipc
 }  // namespace arrow

cpp/src/arrow/schema-test.cc

@@ -98,4 +98,25 @@ f3: list<item: int16>)";
   ASSERT_EQ(expected, result);
 }
 
+TEST_F(TestSchema, GetFieldByName) {
+  auto f0 = field("f0", int32());
+  auto f1 = field("f1", uint8(), false);
+  auto f2 = field("f2", utf8());
+  auto f3 = field("f3", list(int16()));
+
+  vector<shared_ptr<Field>> fields = {f0, f1, f2, f3};
+  auto schema = std::make_shared<Schema>(fields);
+
+  std::shared_ptr<Field> result;
+
+  result = schema->GetFieldByName("f1");
+  ASSERT_TRUE(f1->Equals(result));
+
+  result = schema->GetFieldByName("f3");
+  ASSERT_TRUE(f3->Equals(result));
+
+  result = schema->GetFieldByName("not-found");
+  ASSERT_TRUE(result == nullptr);
+}
+
 }  // namespace arrow

cpp/src/arrow/schema.cc

@@ -42,6 +42,21 @@ bool Schema::Equals(const std::shared_ptr<Schema>& other) const {
   return Equals(*other.get());
 }
 
+std::shared_ptr<Field> Schema::GetFieldByName(const std::string& name) {
+  if (fields_.size() > 0 && name_to_index_.size() == 0) {
+    for (size_t i = 0; i < fields_.size(); ++i) {
+      name_to_index_[fields_[i]->name] = i;
+    }
+  }
+
+  auto it = name_to_index_.find(name);
+  if (it == name_to_index_.end()) {
+    return nullptr;
+  } else {
+    return fields_[it->second];
+  }
+}
+
 std::string Schema::ToString() const {
   std::stringstream buffer;
 

cpp/src/arrow/schema.h

@@ -20,14 +20,14 @@
 
 #include <memory>
 #include <string>
+#include <unordered_map>
 #include <vector>
 
+#include "arrow/type.h"
 #include "arrow/util/visibility.h"
 
 namespace arrow {
 
-struct Field;
-
 class ARROW_EXPORT Schema {
  public:
   explicit Schema(const std::vector<std::shared_ptr<Field>>& fields);
@@ -37,7 +37,10 @@ class ARROW_EXPORT Schema {
   bool Equals(const std::shared_ptr<Schema>& other) const;
 
   // Return the ith schema element. Does not boundscheck
-  const std::shared_ptr<Field>& field(int i) const { return fields_[i]; }
+  std::shared_ptr<Field> field(int i) const { return fields_[i]; }
+
+  // Returns nullptr if name not found
+  std::shared_ptr<Field> GetFieldByName(const std::string& name);
 
   const std::vector<std::shared_ptr<Field>>& fields() const { return fields_; }
 
@@ -48,6 +51,7 @@ class ARROW_EXPORT Schema {
 
  private:
   std::vector<std::shared_ptr<Field>> fields_;
+  std::unordered_map<std::string, int> name_to_index_;
 };
 
 }  // namespace arrow

cpp/src/arrow/type.h

@@ -74,21 +74,15 @@ struct Type {
     // Default unit millisecond
     TIMESTAMP = 17,
 
-    // Timestamp as double seconds since the UNIX epoch
-    TIMESTAMP_DOUBLE = 18,
-
     // Exact time encoded with int64, default unit millisecond
-    TIME = 19,
+    TIME = 18,
 
     // YEAR_MONTH or DAY_TIME interval in SQL style
-    INTERVAL = 20,
+    INTERVAL = 19,
 
     // Precision- and scale-based decimal type. Storage type depends on the
     // parameters.
-    DECIMAL = 21,
-
-    // Decimal value encoded as a text string
-    DECIMAL_TEXT = 22,
+    DECIMAL = 20,
 
     // A list of some logical data type
     LIST = 30,
@@ -98,6 +92,12 @@ struct Type {
 
     // Unions of logical types
     UNION = 32,
+
+    // Timestamp as double seconds since the UNIX epoch
+    TIMESTAMP_DOUBLE = 33,
+
+    // Decimal value encoded as a text string
+    DECIMAL_TEXT = 34,
   };
 };
 
@@ -155,7 +155,7 @@ struct ARROW_EXPORT Field {
   std::string name;
 
   // The field's data type
-  TypePtr type;
+  std::shared_ptr<DataType> type;
 
   // Fields can be nullable
   bool nullable;
@@ -164,8 +164,8 @@ struct ARROW_EXPORT Field {
   // 0 means it's not dictionary encoded
   int64_t dictionary;
 
-  Field(const std::string& name, const TypePtr& type, bool nullable = true,
-      int64_t dictionary = 0)
+  Field(const std::string& name, const std::shared_ptr<DataType>& type,
+      bool nullable = true, int64_t dictionary = 0)
       : name(name), type(type), nullable(nullable), dictionary(dictionary) {}
 
   bool operator==(const Field& other) const { return this->Equals(other); }
@@ -459,8 +459,8 @@ std::shared_ptr<DataType> ARROW_EXPORT union_(
     const std::vector<std::shared_ptr<Field>>& child_fields,
     const std::vector<uint8_t>& type_ids, UnionMode mode = UnionMode::SPARSE);
 
-std::shared_ptr<Field> ARROW_EXPORT field(const std::string& name, const TypePtr& type,
-    bool nullable = true, int64_t dictionary = 0);
+std::shared_ptr<Field> ARROW_EXPORT field(const std::string& name,
+    const std::shared_ptr<DataType>& type, bool nullable = true, int64_t dictionary = 0);
 
 }  // namespace arrow