ARROW-590: Add integration tests for Union types

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

@@ -269,10 +269,10 @@ class SchemaWriter {
     writer_->Key("mode");
     switch (type.mode()) {
       case UnionMode::SPARSE:
-        writer_->String("SPARSE");
+        writer_->String("Sparse");
         break;
       case UnionMode::DENSE:
-        writer_->String("DENSE");
+        writer_->String("Dense");
         break;
     }
 
@@ -569,7 +569,7 @@ class ArrayWriter {
     WriteValidityField(array);
     const auto& type = static_cast<const UnionType&>(*array.type());
 
-    WriteIntegerField("TYPE_ID", array.raw_type_ids(), array.length());
+    WriteIntegerField("TYPE", array.raw_type_ids(), array.length());
     if (type.mode() == UnionMode::DENSE) {
       WriteIntegerField("OFFSET", array.raw_value_offsets(), array.length());
     }
@@ -763,9 +763,9 @@ static Status GetUnion(const RjObject& json_type,
   std::string mode_str = it_mode->value.GetString();
   UnionMode mode;
 
-  if (mode_str == "SPARSE") {
+  if (mode_str == "Sparse") {
     mode = UnionMode::SPARSE;
-  } else if (mode_str == "DENSE") {
+  } else if (mode_str == "Dense") {
     mode = UnionMode::DENSE;
   } else {
     std::stringstream ss;
@@ -774,13 +774,25 @@ static Status GetUnion(const RjObject& json_type,
   }
 
   const auto& it_type_codes = json_type.FindMember("typeIds");
-  RETURN_NOT_ARRAY("typeIds", it_type_codes, json_type);
 
   std::vector<uint8_t> type_codes;
-  const auto& id_array = it_type_codes->value.GetArray();
-  for (const rj::Value& val : id_array) {
-    DCHECK(val.IsUint());
-    type_codes.push_back(static_cast<uint8_t>(val.GetUint()));
+  if (it_type_codes == json_type.MemberEnd()) {
+    for (uint8_t code = 0; code < static_cast<uint8_t>(children.size()); ++code) {
+      type_codes.push_back(code);
+    }
+  } else {
+    RETURN_NOT_ARRAY("typeIds", it_type_codes, json_type);
+    const auto& id_array = it_type_codes->value.GetArray();
+    if (id_array.Size() == 0) {
+      for (uint8_t code = 0; code < static_cast<uint8_t>(children.size()); ++code) {
+        type_codes.push_back(code);
+      }
+    } else {
+      for (const rj::Value& val : id_array) {
+        DCHECK(val.IsUint());
+        type_codes.push_back(static_cast<uint8_t>(val.GetUint()));
+      }
+    }
   }
 
   *type = union_(children, type_codes, mode);
@@ -1142,8 +1154,8 @@ class ArrayReader {
 
     RETURN_NOT_OK(GetValidityBuffer(is_valid_, &null_count, &validity_buffer));
 
-    const auto& json_type_ids = obj_->FindMember("TYPE_ID");
-    RETURN_NOT_ARRAY("TYPE_ID", json_type_ids, *obj_);
+    const auto& json_type_ids = obj_->FindMember("TYPE");
+    RETURN_NOT_ARRAY("TYPE", json_type_ids, *obj_);
     RETURN_NOT_OK(
         GetIntArray<uint8_t>(json_type_ids->value.GetArray(), length_, &type_id_buffer));
 

integration/data/union.json

@@ -0,0 +1,82 @@
+{
+  "schema" : {
+    "fields" : [{
+      "name" : "union",
+      "nullable" : true,
+      "type" : {
+        "name" : "union",
+        "mode" : "Sparse",
+        "typeIds" : [4,5]
+      },
+      "children" : [{
+        "name" : "int",
+        "nullable" : true,
+        "type" : {
+          "name" : "int",
+          "bitWidth" : 32,
+          "isSigned" : true
+        },
+        "children" : [ ],
+        "typeLayout" : {
+          "vectors" : [{
+            "type" : "VALIDITY",
+            "typeBitWidth" : 1
+          },{
+            "type" : "DATA",
+            "typeBitWidth" : 32
+          }]
+        }
+      },{
+        "name" : "bigint",
+        "nullable" : true,
+        "type" : {
+          "name" : "int",
+          "bitWidth" : 64,
+          "isSigned" : true
+        },
+        "children" : [ ],
+        "typeLayout" : {
+          "vectors" : [{
+            "type" : "VALIDITY",
+            "typeBitWidth" : 1
+          },{
+            "type" : "DATA",
+            "typeBitWidth" : 64
+          }]
+        }
+      }],
+      "typeLayout" : {
+          "vectors" : [
+              {
+                  "type": "VALIDITY",
+                  "typeBitWidth": 1
+              },
+              {
+                  "type" : "TYPE",
+                  "typeBitWidth" : 8
+              }
+          ]
+      }
+    }]
+  },
+  "batches" : [{
+    "count" : 10,
+    "columns" : [{
+        "name" : "union",
+        "count" : 10,
+        "VALIDITY" : [1,1,1,1,1,1,1,1,1,1],
+        "TYPE" : [4,5,4,5,4,5,4,5,4,5],
+        "children" : [{
+            "name" : "int",
+            "count" : 10,
+            "VALIDITY" : [1,0,1,0,1,0,1,0,1,0],
+            "DATA" : [0,0,2,0,4,0,6,0,8,0]
+        },{
+            "name" : "bigint",
+            "count" : 10,
+            "VALIDITY" : [0,1,0,1,0,1,0,1,0,1],
+            "DATA" : [0,1,0,3,0,5,0,7,0,9]
+        }]
+    }]
+  }]
+}

integration/integration_test.py

@@ -533,6 +533,55 @@ def generate_column(self, size, name=None):
         return StructColumn(name, size, is_valid, field_values)
 
 
+class UnionType(DataType):
+
+    def __init__(self, name, mode, type_ids, field_types, nullable=True):
+        DataType.__init__(self, name, nullable=nullable)
+        self.mode = mode
+        self.type_ids = type_ids
+        self.field_types = field_types
+
+    def _get_type(self):
+        type_ids = self.type_ids if self.type_ids is not None else []
+
+        attrs = [
+            ('name', 'union'),
+            ('mode', self.mode),
+            ('typeIds', type_ids)
+        ]
+
+        return OrderedDict(attrs)
+
+    def _get_children(self):
+        return [type_.get_json() for type_ in self.field_types]
+
+    def _get_type_layout(self):
+        return OrderedDict([
+            ('vectors',
+             [OrderedDict([('type', 'VALIDITY'),
+                           ('typeBitWidth', 1)]),
+              OrderedDict([('type', 'TYPE'),
+                           ('typeBitWidth', 8)])])])
+
+    def _make_type(self, size):
+        if self.type_ids is not None:
+            type_ids = self.type_ids
+        else:
+            type_ids = np.arange(len(self.field_types))
+
+        return np.random.choice(type_ids, size)
+
+    def generate_column(self, size, name=None):
+        is_valid = self._make_is_valid(size)
+        types = self._make_type(size)
+
+        field_values = [type_.generate_column(size)
+                        for type_ in self.field_types]
+        if name is None:
+            name = self.name
+        return UnionColumn(name, size, is_valid, types, field_values)
+
+
 class Dictionary(object):
 
     def __init__(self, id_, field, values, ordered=False):
@@ -603,6 +652,23 @@ def _get_children(self):
         return [field.get_json() for field in self.field_values]
 
 
+class UnionColumn(Column):
+    def __init__(self, name, count, is_valid, types, field_values):
+        Column.__init__(self, name, count)
+        self.is_valid = is_valid
+        self.types = types
+        self.field_values = field_values
+
+    def _get_buffers(self):
+        return [
+            ('VALIDITY', [int(v) for v in self.is_valid]),
+            ('TYPE', [int(v) for v in self.types])
+        ]
+
+    def _get_children(self):
+        return [field.get_json() for field in self.field_values]
+
+
 class JsonRecordBatch(object):
 
     def __init__(self, count, columns):
@@ -747,6 +813,22 @@ def generate_dictionary_case():
                           dictionaries=[dict1, dict2])
 
 
+def _generate_union_field(type_ids=None):
+    return UnionType('union_nullable', "Sparse", type_ids,
+                     [get_field('f1', 'int64'),
+                      get_field('f2', 'float64'),
+                      get_field('f3', 'utf8'),
+                      get_field('f4', 'binary'),
+                      StructType('f5', [get_field('f1', 'int32'),
+                                        get_field('f2', 'utf8')]),])
+
+
+def generate_union_case():
+    type_ids = np.random.choice(range(128), 5, replace=False).tolist()
+    fields = [_generate_union_field(type_ids)]
+    return _generate_file("union", fields, [10])
+
+
 def get_generated_json_files():
     temp_dir = tempfile.mkdtemp()
 
@@ -758,7 +840,8 @@ def _temp_path():
         generate_primitive_case([0, 0, 0]),
         generate_datetime_case(),
         generate_nested_case(),
-        generate_dictionary_case()
+        generate_dictionary_case(),
+        generate_union_case(),
     ]
 
     generated_paths = []
@@ -948,6 +1031,7 @@ def get_static_json_files():
 
 def run_all_tests(debug=False):
     testers = [CPPTester(debug=debug), JavaTester(debug=debug)]
+
     static_json_files = get_static_json_files()
     generated_json_files = get_generated_json_files()
     json_files = static_json_files + generated_json_files

java/tools/src/main/java/org/apache/arrow/tools/Integration.java

@@ -187,7 +187,7 @@ public void execute(File arrowFile, File jsonFile) throws IOException {
           LOGGER.debug("ARROW schema: " + arrowSchema);
           LOGGER.debug("JSON Input file size: " + jsonFile.length());
           LOGGER.debug("JSON schema: " + jsonSchema);
-          Validator.compareSchemas(jsonSchema, arrowSchema);
+          Validator.compareSchemas(arrowSchema, jsonSchema);
 
           List<ArrowBlock> recordBatches = arrowReader.getRecordBlocks();
           Iterator<ArrowBlock> iterator = recordBatches.iterator();