GH-1302: Simplify JSON extension tests

Change-Id: I2f961fbc983bd29bffd2b4280021f7b696e3ad50

Author: ddupg

vector/src/test/java/org/apache/arrow/vector/TestJsonType.java

@@ -16,6 +16,7 @@
  */
 package org.apache.arrow.vector;
 
+import static org.apache.arrow.vector.testing.ValueVectorDataPopulator.setVector;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertInstanceOf;
@@ -25,11 +26,13 @@
 
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
-import java.util.Base64;
 import java.util.Collections;
 import java.util.stream.Stream;
+import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.vector.extension.InvalidExtensionMetadataException;
 import org.apache.arrow.vector.extension.JsonType;
@@ -43,35 +46,61 @@
 import org.apache.arrow.vector.types.pojo.Schema;
 import org.apache.arrow.vector.util.Text;
 import org.apache.arrow.vector.util.TransferPair;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.MethodSource;
 import org.junit.jupiter.params.provider.NullSource;
 import org.junit.jupiter.params.provider.ValueSource;
 
 class TestJsonType {
+  BufferAllocator allocator;
+  ArrowType.ExtensionType previousType;
+
+  @BeforeEach
+  void beforeEach() {
+    allocator = new RootAllocator();
+    previousType = ExtensionTypeRegistry.lookup(JsonType.EXTENSION_NAME);
+    JsonType.ensureRegistered();
+  }
+
+  @AfterEach
+  void afterEach() {
+    ExtensionTypeRegistry.unregister(new JsonType(ArrowType.Utf8.INSTANCE));
+    if (previousType != null) {
+      ExtensionTypeRegistry.register(previousType);
+    }
+    allocator.close();
+  }
+
   static Stream<ArrowType> storageTypes() {
     return Stream.of(
         ArrowType.Utf8.INSTANCE, ArrowType.LargeUtf8.INSTANCE, ArrowType.Utf8View.INSTANCE);
   }
 
   @ParameterizedTest
   @MethodSource("storageTypes")
-  void testType(ArrowType storage) {
+  void testRoundTrip(ArrowType storage) {
     JsonType type = new JsonType(storage);
     assertEquals("arrow.json", type.extensionName());
     assertEquals(storage, type.storageType());
     assertFalse(type.isComplex());
     assertEquals("", type.serialize());
-    for (String metadata : new String[] {"", "{}", " { } ", "{\"future\": {\"value\": 1}}"}) {
-      ArrowType restored = type.deserialize(storage, metadata);
-      assertEquals(type, restored);
-      assertEquals(type.hashCode(), restored.hashCode());
-    }
-    storageTypes()
-        .filter(other -> !storage.equals(other))
-        .forEach(other -> assertNotEquals(type, new JsonType(other)));
-    assertNotEquals(type, storage);
+    assertEquals(type, type.deserialize(storage, type.serialize()));
+    assertNotEquals(
+        type,
+        new JsonType(
+            storage instanceof ArrowType.Utf8
+                ? ArrowType.LargeUtf8.INSTANCE
+                : ArrowType.Utf8.INSTANCE));
+  }
+
+  @ParameterizedTest
+  @ValueSource(strings = {"", "{}", " { } ", "{\"future\": 1}"})
+  void testDeserializeValid(String metadata) {
+    JsonType type = new JsonType(ArrowType.Utf8.INSTANCE);
+    assertEquals(type, type.deserialize(type.storageType(), metadata));
   }
 
   @ParameterizedTest
@@ -103,7 +132,6 @@ void testInvalidStorage() {
   @ParameterizedTest
   @MethodSource("storageTypes")
   void testSchemaRoundTrip(ArrowType storage) {
-    JsonType.ensureRegistered();
     for (boolean nullable : new boolean[] {false, true}) {
       Field field = field(storage, nullable);
       Schema schema = new Schema(Collections.singletonList(field));
@@ -115,21 +143,11 @@ void testSchemaRoundTrip(ArrowType storage) {
   // nullable=False, metadata={"custom": "preserved"}) for t in
   // [pa.string(), pa.large_string(), pa.string_view()]]).serialize().
   @Test
-  void testPyArrowSchema() {
-    JsonType.ensureRegistered();
-    String encoded =
-        "/////9ACAAAQAAAAAAAKAAwABgAFAAgACgAAAAABBAAMAAAACAAIAAAABAAIAAAABAAAAAMAAADEAQAA2AAAAAQAAABa/v//"
-            + "AAAAGBQAAADAAAAACAAAABQAAAAAAAAABAAAAGpzb24AAAAAAwAAAHgAAABAAAAABAAAANz9//8YAAAABAAAAAoAAABhcnJv"
-            + "dy5qc29uAAAUAAAAQVJST1c6ZXh0ZW5zaW9uOm5hbWUAAAAAFP7//xAAAAAEAAAAAAAAAAAAAAAYAAAAQVJST1c6ZXh0ZW5z"
-            + "aW9uOm1ldGFkYXRhAAAAAEj+//8YAAAABAAAAAkAAABwcmVzZXJ2ZWQAAAAGAAAAY3VzdG9tAABA/v//Kv///wAAABQUAAAA"
-            + "xAAAAAgAAAAUAAAAAAAAAAQAAABqc29uAAAAAAMAAAB4AAAAQAAAAAQAAACs/v//GAAAAAQAAAAKAAAAYXJyb3cuanNvbgAA"
-            + "FAAAAEFSUk9XOmV4dGVuc2lvbjpuYW1lAAAAAOT+//8QAAAABAAAAAAAAAAAAAAAGAAAAEFSUk9XOmV4dGVuc2lvbjptZXRh"
-            + "ZGF0YQAAAAAY////GAAAAAQAAAAJAAAAcHJlc2VydmVkAAAABgAAAGN1c3RvbQAABAAGAAQAAAAAABIAGAAIAAAABwAMAAAA"
-            + "EAAUABIAAAAAAAAFFAAAAMwAAAAIAAAAFAAAAAAAAAAEAAAAanNvbgAAAAADAAAAgAAAAEAAAAAEAAAAlP///xgAAAAEAAAA"
-            + "CgAAAGFycm93Lmpzb24AABQAAABBUlJPVzpleHRlbnNpb246bmFtZQAAAADM////EAAAAAQAAAAAAAAAAAAAABgAAABBUlJP"
-            + "VzpleHRlbnNpb246bWV0YWRhdGEAAAAACAAMAAQACAAIAAAAGAAAAAQAAAAJAAAAcHJlc2VydmVkAAAABgAAAGN1c3RvbQAA"
-            + "BAAEAAQAAAA=";
-    Schema schema = Schema.deserializeMessage(ByteBuffer.wrap(Base64.getDecoder().decode(encoded)));
+  void testPyArrowSchema() throws IOException {
+    Schema schema;
+    try (InputStream input = getClass().getResourceAsStream("/pyarrow_json_schema.arrow")) {
+      schema = Schema.deserializeMessage(ByteBuffer.wrap(input.readAllBytes()));
+    }
     ArrowType[] storage = storageTypes().toArray(ArrowType[]::new);
     for (int i = 0; i < storage.length; i++) {
       assertEquals(field(storage[i], false), schema.getFields().get(i));
@@ -148,108 +166,95 @@ nullable, new JsonType(storage), null, Collections.singletonMap("custom", "prese
   @MethodSource("storageTypes")
   void testTransfer(ArrowType storage) {
     Field field = field(storage, true);
-    try (RootAllocator allocator = new RootAllocator();
-        JsonVector source = (JsonVector) field.createVector(allocator)) {
+    try (JsonVector source = (JsonVector) field.createVector(allocator)) {
       byte[] bytes =
           "{\"key\":\"value longer than twelve bytes\"}".getBytes(StandardCharsets.UTF_8);
-      FieldVector underlying = source.getUnderlyingVector();
-      if (underlying instanceof VarCharVector) {
-        ((VarCharVector) underlying).setSafe(0, bytes);
-      } else if (underlying instanceof LargeVarCharVector) {
-        ((LargeVarCharVector) underlying).setSafe(0, bytes);
-      } else {
-        ((ViewVarCharVector) underlying).setSafe(0, bytes);
-      }
-      source.setNull(1);
-      source.setValueCount(2);
+      setVector((VariableWidthFieldVector) source.getUnderlyingVector(), null, bytes, null);
       TransferPair split = source.getTransferPair("copy", allocator);
       try (JsonVector target = assertInstanceOf(JsonVector.class, split.getTo())) {
-        split.splitAndTransfer(0, 2);
+        split.splitAndTransfer(1, 2);
         assertEquals("copy", target.getName());
         assertEquals(field.getFieldType(), target.getField().getFieldType());
-        assertEquals(source.getObject(0), target.getObject(0));
+        assertEquals(new Text(bytes), target.getObject(0));
         assertTrue(target.isNull(1));
       }
       try (JsonVector target = (JsonVector) field.createVector(allocator)) {
         TransferPair copy = source.makeTransferPair(target);
-        copy.copyValueSafe(0, 0);
-        copy.copyValueSafe(1, 1);
+        copy.copyValueSafe(1, 0);
+        copy.copyValueSafe(2, 1);
         target.setValueCount(2);
-        assertEquals(source.getObject(0), target.getObject(0));
+        assertEquals(new Text(bytes), target.getObject(0));
         assertTrue(target.isNull(1));
       }
       TransferPair transfer = source.getTransferPair(allocator);
       try (JsonVector target = assertInstanceOf(JsonVector.class, transfer.getTo())) {
         transfer.transfer();
         assertEquals(field, target.getField());
-        assertEquals(new Text(bytes), target.getObject(0));
-        assertTrue(target.isNull(1));
-        assertEquals(2, target.getValueCount());
+        assertTrue(target.isNull(0));
+        assertEquals(new Text(bytes), target.getObject(1));
+        assertTrue(target.isNull(2));
+        assertEquals(3, target.getValueCount());
         assertEquals(0, source.getValueCount());
       }
     }
   }
 
   @ParameterizedTest
   @MethodSource("storageTypes")
-  void testIpcAndUnregisteredFallback(ArrowType storage) throws Exception {
-    ArrowType.ExtensionType previous = ExtensionTypeRegistry.lookup(JsonType.EXTENSION_NAME);
-    JsonType type = new JsonType(storage);
-    JsonType.ensureRegistered();
-    try (RootAllocator allocator = new RootAllocator()) {
-      Field field = field(storage, true);
-      Schema schema = new Schema(Collections.singletonList(field));
-      ByteArrayOutputStream out = new ByteArrayOutputStream();
-      String[] values = {
-        "{\"message\":\"你好, a JSON value longer than twelve bytes\"}", null, "null", "42", "[]"
-      };
-      try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator);
-          ArrowStreamWriter writer = new ArrowStreamWriter(root, null, out)) {
-        JsonVector vector = assertInstanceOf(JsonVector.class, root.getVector(0));
-        FieldVector underlying = vector.getUnderlyingVector();
-        for (int i = 0; i < values.length; i++) {
-          if (values[i] == null) {
-            vector.setNull(i);
-          } else {
-            byte[] bytes = values[i].getBytes(StandardCharsets.UTF_8);
-            if (storage instanceof ArrowType.Utf8) {
-              assertInstanceOf(VarCharVector.class, underlying).setSafe(i, bytes);
-            } else if (storage instanceof ArrowType.LargeUtf8) {
-              assertInstanceOf(LargeVarCharVector.class, underlying).setSafe(i, bytes);
-            } else {
-              assertInstanceOf(ViewVarCharVector.class, underlying).setSafe(i, bytes);
-            }
-          }
-        }
-        root.setRowCount(values.length);
-        assertEquals(underlying.hashCode(0), vector.hashCode(0));
-        writer.start();
-        writer.writeBatch();
-        writer.end();
-      }
-      for (boolean registered : new boolean[] {true, false}) {
-        if (!registered) {
-          ExtensionTypeRegistry.unregister(type);
-        }
-        try (ArrowStreamReader reader =
-            new ArrowStreamReader(new ByteArrayInputStream(out.toByteArray()), allocator)) {
-          assertTrue(reader.loadNextBatch());
-          FieldVector vector = reader.getVectorSchemaRoot().getVector(0);
-          assertEquals(registered ? type : storage, vector.getField().getType());
-          assertEquals(field.getMetadata(), vector.getField().getMetadata());
-          assertTrue(vector.getField().isNullable());
-          assertEquals(values.length, vector.getValueCount());
-          for (int i = 0; i < values.length; i++) {
-            assertEquals(values[i] == null ? null : new Text(values[i]), vector.getObject(i));
-          }
-          assertFalse(reader.loadNextBatch());
-        }
-      }
-    } finally {
-      ExtensionTypeRegistry.unregister(type);
-      if (previous != null) {
-        ExtensionTypeRegistry.register(previous);
-      }
+  void testVectorIpcRoundTrip(ArrowType storage) throws IOException {
+    Field field = field(storage, true);
+    byte[] serialized = writeStream(field);
+    try (ArrowStreamReader reader =
+        new ArrowStreamReader(new ByteArrayInputStream(serialized), allocator)) {
+      assertTrue(reader.loadNextBatch());
+      JsonVector vector =
+          assertInstanceOf(JsonVector.class, reader.getVectorSchemaRoot().getVector(0));
+      assertEquals(field, vector.getField());
+      assertValues(vector);
     }
   }
+
+  @ParameterizedTest
+  @MethodSource("storageTypes")
+  void testReadUnderlyingType(ArrowType storage) throws IOException {
+    Field field = field(storage, true);
+    byte[] serialized = writeStream(field);
+    ExtensionTypeRegistry.unregister((JsonType) field.getType());
+    try (ArrowStreamReader reader =
+        new ArrowStreamReader(new ByteArrayInputStream(serialized), allocator)) {
+      assertTrue(reader.loadNextBatch());
+      FieldVector vector = reader.getVectorSchemaRoot().getVector(0);
+      assertEquals(storage, vector.getField().getType());
+      assertEquals(field.getMetadata(), vector.getField().getMetadata());
+      assertValues(vector);
+    }
+  }
+
+  private static final String JSON = "{\"message\":\"你好, a JSON value longer than twelve bytes\"}";
+
+  private byte[] writeStream(Field field) throws IOException {
+    ByteArrayOutputStream out = new ByteArrayOutputStream();
+    try (VectorSchemaRoot root =
+            VectorSchemaRoot.create(new Schema(Collections.singletonList(field)), allocator);
+        ArrowStreamWriter writer = new ArrowStreamWriter(root, null, out)) {
+      JsonVector vector = (JsonVector) root.getVector(0);
+      setVector(
+          (VariableWidthFieldVector) vector.getUnderlyingVector(),
+          JSON.getBytes(StandardCharsets.UTF_8),
+          null,
+          "null".getBytes(StandardCharsets.UTF_8));
+      root.setRowCount(3);
+      writer.start();
+      writer.writeBatch();
+      writer.end();
+    }
+    return out.toByteArray();
+  }
+
+  private static void assertValues(FieldVector vector) {
+    assertEquals(3, vector.getValueCount());
+    assertEquals(new Text(JSON), vector.getObject(0));
+    assertTrue(vector.isNull(1));
+    assertEquals(new Text("null"), vector.getObject(2));
+  }
 }