Test JSON array roundtrip for numeric types, strings, lists, structs

Change-Id: I9403a253307d304d0dc5a71e5d8b7e623fbfa69f

Author: wesm

cpp/src/arrow/array.cc

@@ -18,13 +18,24 @@
 #include "arrow/array.h"
 
 #include <cstdint>
+#include <cstring>
 
 #include "arrow/util/bit-util.h"
 #include "arrow/util/buffer.h"
 #include "arrow/util/status.h"
 
 namespace arrow {
 
+Status GetEmptyBitmap(
+    MemoryPool* pool, int32_t length, std::shared_ptr<MutableBuffer>* result) {
+  auto buffer = std::make_shared<PoolBuffer>(pool);
+  RETURN_NOT_OK(buffer->Resize(BitUtil::BytesForBits(length)));
+  memset(buffer->mutable_data(), 0, buffer->size());
+
+  *result = buffer;
+  return Status::OK();
+}
+
 // ----------------------------------------------------------------------
 // Base array class
 

cpp/src/arrow/array.h

@@ -29,6 +29,8 @@
 namespace arrow {
 
 class Buffer;
+class MemoryPool;
+class MutableBuffer;
 class Status;
 
 // Immutable data array with some logical type and some length. Any memory is
@@ -103,6 +105,10 @@ class ARROW_EXPORT NullArray : public Array {
 };
 
 typedef std::shared_ptr<Array> ArrayPtr;
+
+Status ARROW_EXPORT GetEmptyBitmap(
+    MemoryPool* pool, int32_t length, std::shared_ptr<MutableBuffer>* result);
+
 }  // namespace arrow
 
 #endif

cpp/src/arrow/column-test.cc

@@ -22,6 +22,7 @@
 
 #include "gtest/gtest.h"
 
+#include "arrow/array.h"
 #include "arrow/column.h"
 #include "arrow/schema.h"
 #include "arrow/test-util.h"

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

@@ -35,6 +35,8 @@
 #include "arrow/type.h"
 #include "arrow/type_traits.h"
 #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"
 
@@ -64,15 +66,20 @@ void TestArrayRoundTrip(const Array& array) {
 
   ASSERT_OK(WriteJsonArray(name, array, &writer));
 
+  std::string array_as_json = sb.GetString();
+
   rj::Document d;
-  d.Parse(sb.GetString());
+  d.Parse(array_as_json);
+
+  if (d.HasParseError()) { FAIL() << "JSON parsing failed"; }
 
   std::shared_ptr<Array> out;
   ASSERT_OK(ReadJsonArray(default_memory_pool(), d, array.type(), &out));
 
-  ASSERT_TRUE(array.Equals(out));
-}
+  std::cout << array_as_json << std::endl;
 
+  ASSERT_TRUE(array.Equals(out)) << array_as_json;
+}
 
 template <typename T, typename ValueType>
 void CheckPrimitive(const std::shared_ptr<DataType>& type,
@@ -109,12 +116,70 @@ TEST(TestJsonSchemaWriter, FlatTypes) {
   TestSchemaRoundTrip(schema);
 }
 
+template <typename T>
+void PrimitiveTypesCheckOne() {
+  using c_type = typename T::c_type;
+
+  std::vector<bool> is_valid = {true, false, true, true, true, false, true, true};
+  std::vector<c_type> values = {0, 1, 2, 3, 4, 5, 6, 7};
+  CheckPrimitive<T, c_type>(std::make_shared<T>(), is_valid, values);
+}
 
 TEST(TestJsonArrayWriter, PrimitiveTypes) {
+  PrimitiveTypesCheckOne<Int8Type>();
+  PrimitiveTypesCheckOne<Int16Type>();
+  PrimitiveTypesCheckOne<Int32Type>();
+  PrimitiveTypesCheckOne<Int64Type>();
+  PrimitiveTypesCheckOne<UInt8Type>();
+  PrimitiveTypesCheckOne<UInt16Type>();
+  PrimitiveTypesCheckOne<UInt32Type>();
+  PrimitiveTypesCheckOne<UInt64Type>();
+  PrimitiveTypesCheckOne<FloatType>();
+  PrimitiveTypesCheckOne<DoubleType>();
+
   std::vector<bool> is_valid = {true, false, true, true, true, false, true, true};
+  std::vector<std::string> values = {"foo", "bar", "", "baz", "qux", "foo", "a", "1"};
+
+  CheckPrimitive<StringType, std::string>(utf8(), is_valid, values);
+  CheckPrimitive<BinaryType, std::string>(binary(), is_valid, values);
+}
+
+TEST(TestJsonArrayWriter, NestedTypes) {
+  auto value_type = int32();
+
+  std::vector<bool> values_is_valid = {true, false, true, true, false, true, true};
+  std::vector<int32_t> values = {0, 1, 2, 3, 4, 5, 6};
+
+  std::shared_ptr<Buffer> values_buffer = test::GetBufferFromVector(values);
+  std::shared_ptr<Buffer> values_bitmap;
+  ASSERT_OK(test::GetBitmapFromBoolVector(values_is_valid, &values_bitmap));
+  auto values_array = std::make_shared<Int32Array>(
+      value_type, static_cast<int32_t>(values.size()), values_buffer, 2, values_bitmap);
+
+  // List
+  std::vector<bool> list_is_valid = {true, false, true, true, true};
+  std::vector<int32_t> offsets = {0, 0, 0, 1, 4, 7};
+
+  std::shared_ptr<Buffer> list_bitmap;
+  ASSERT_OK(test::GetBitmapFromBoolVector(list_is_valid, &list_bitmap));
+  std::shared_ptr<Buffer> offsets_buffer = test::GetBufferFromVector(offsets);
+
+  ListArray list_array(list(value_type), 5, offsets_buffer, values_array, 1, list_bitmap);
+
+  TestArrayRoundTrip(list_array);
+
+  // Struct
+  std::vector<bool> struct_is_valid = {true, false, true, true, true, false, true};
+  std::shared_ptr<Buffer> struct_bitmap;
+  ASSERT_OK(test::GetBitmapFromBoolVector(struct_is_valid, &struct_bitmap));
+
+  auto struct_type =
+      struct_({field("f1", int32()), field("f2", int32()), field("f3", int32())});
 
-  std::vector<uint8_t> u1 = {0, 1, 2, 3, 4, 5, 6, 7};
-  CheckPrimitive<UInt8Type, uint8_t>(uint8(), is_valid, u1);
+  std::vector<std::shared_ptr<Array>> fields = {values_array, values_array, values_array};
+  StructArray struct_array(
+      struct_type, static_cast<int>(struct_is_valid.size()), fields, 2, struct_bitmap);
+  TestArrayRoundTrip(struct_array);
 }
 
 }  // namespace ipc

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

@@ -458,6 +458,7 @@ class JsonArrayWriter : public ArrayVisitor {
 
   template <typename T>
   void WriteDataField(const T& arr) {
+    writer_->Key("DATA");
     writer_->StartArray();
     WriteDataValues(arr);
     writer_->EndArray();
@@ -858,10 +859,9 @@ class JsonArrayReader {
       std::shared_ptr<Buffer>* validity_buffer) {
     int length = static_cast<int>(is_valid.size());
 
-    auto out_buffer = std::make_shared<PoolBuffer>(pool_);
-    RETURN_NOT_OK(out_buffer->Resize(BitUtil::BytesForBits(length)));
-    uint8_t* bitmap = reinterpret_cast<uint8_t*>(out_buffer->mutable_data());
-    memset(bitmap, 0, out_buffer->size());
+    std::shared_ptr<MutableBuffer> out_buffer;
+    RETURN_NOT_OK(GetEmptyBitmap(pool_, length, &out_buffer));
+    uint8_t* bitmap = out_buffer->mutable_data();
 
     *null_count = 0;
     for (int i = 0; i < length; ++i) {
@@ -1015,9 +1015,17 @@ class JsonArrayReader {
     }
 
     for (int i = 0; i < static_cast<int>(json_children_arr.Size()); ++i) {
-      DCHECK(json_children_arr[i].IsObject());
+      const rj::Value& json_child = json_children_arr[i];
+      DCHECK(json_child.IsObject());
+
+      std::shared_ptr<Field> child_field = type->child(i);
+
+      auto it = json_child.FindMember("name");
+      RETURN_NOT_STRING("name", it, json_child);
+
+      DCHECK_EQ(it->value.GetString(), child_field->name);
       std::shared_ptr<Array> child;
-      RETURN_NOT_OK(GetArray(json_children_arr[i], type->child(i)->type, &child));
+      RETURN_NOT_OK(GetArray(json_children_arr[i], child_field->type, &child));
       array->emplace_back(child);
     }
 
@@ -1042,7 +1050,7 @@ class JsonArrayReader {
 
     DCHECK_EQ(static_cast<int>(json_validity.Size()), length);
 
-    std::vector<bool> is_valid(length);
+    std::vector<bool> is_valid;
     for (const rj::Value& val : json_validity) {
       DCHECK(val.IsInt());
       is_valid.push_back(static_cast<bool>(val.GetInt()));
@@ -1120,9 +1128,8 @@ Status ReadJsonArray(MemoryPool* pool, const rj::Value& json_array,
   return converter.GetArray(json_array, type, array);
 }
 
-
-Status ReadJsonArray(MemoryPool* pool, const rj::Value& json_array,
-    const Schema& schema, std::shared_ptr<Array>* array) {
+Status ReadJsonArray(MemoryPool* pool, const rj::Value& json_array, const Schema& schema,
+    std::shared_ptr<Array>* array) {
   if (!json_array.IsObject()) { return Status::Invalid("Element was not a JSON object"); }
 
   const auto& json_obj = json_array.GetObject();

cpp/src/arrow/test-util.h

@@ -27,6 +27,7 @@
 
 #include "gtest/gtest.h"
 
+#include "arrow/array.h"
 #include "arrow/column.h"
 #include "arrow/schema.h"
 #include "arrow/table.h"
@@ -102,11 +103,27 @@ void random_real(int n, uint32_t seed, T min_value, T max_value, std::vector<T>*
 }
 
 template <typename T>
-std::shared_ptr<Buffer> to_buffer(const std::vector<T>& values) {
+std::shared_ptr<Buffer> GetBufferFromVector(const std::vector<T>& values) {
   return std::make_shared<Buffer>(
       reinterpret_cast<const uint8_t*>(values.data()), values.size() * sizeof(T));
 }
 
+static inline Status GetBitmapFromBoolVector(
+    const std::vector<bool>& is_valid, std::shared_ptr<Buffer>* result) {
+  int length = static_cast<int>(is_valid.size());
+
+  std::shared_ptr<MutableBuffer> buffer;
+  RETURN_NOT_OK(GetEmptyBitmap(default_memory_pool(), length, &buffer));
+
+  uint8_t* bitmap = buffer->mutable_data();
+  for (int i = 0; i < length; ++i) {
+    if (is_valid[i]) { BitUtil::SetBit(bitmap, i); }
+  }
+
+  *result = buffer;
+  return Status::OK();
+}
+
 // Sets approximately pct_null of the first n bytes in null_bytes to zero
 // and the rest to non-zero (true) values.
 void random_null_bytes(int64_t n, double pct_null, uint8_t* null_bytes) {

cpp/src/arrow/types/string-test.cc

@@ -66,8 +66,8 @@ class TestStringContainer : public ::testing::Test {
 
   void MakeArray() {
     length_ = offsets_.size() - 1;
-    value_buf_ = test::to_buffer(chars_);
-    offsets_buf_ = test::to_buffer(offsets_);
+    value_buf_ = test::GetBufferFromVector(chars_);
+    offsets_buf_ = test::GetBufferFromVector(offsets_);
     null_bitmap_ = test::bytes_to_null_buffer(valid_bytes_);
     null_count_ = test::null_count(valid_bytes_);
 
@@ -131,7 +131,7 @@ TEST_F(TestStringContainer, TestGetString) {
 
 TEST_F(TestStringContainer, TestEmptyStringComparison) {
   offsets_ = {0, 0, 0, 0, 0, 0};
-  offsets_buf_ = test::to_buffer(offsets_);
+  offsets_buf_ = test::GetBufferFromVector(offsets_);
   length_ = offsets_.size() - 1;
 
   auto strings_a = std::make_shared<StringArray>(
@@ -227,8 +227,8 @@ class TestBinaryContainer : public ::testing::Test {
 
   void MakeArray() {
     length_ = offsets_.size() - 1;
-    value_buf_ = test::to_buffer(chars_);
-    offsets_buf_ = test::to_buffer(offsets_);
+    value_buf_ = test::GetBufferFromVector(chars_);
+    offsets_buf_ = test::GetBufferFromVector(offsets_);
 
     null_bitmap_ = test::bytes_to_null_buffer(valid_bytes_);
     null_count_ = test::null_count(valid_bytes_);