From 864c25089fea56f0a7acdc158477473ab6c1bcaf Mon Sep 17 00:00:00 2001 From: Xin Sun Date: Sun, 20 Sep 2026 11:10:58 +0800 Subject: [PATCH 1/3] GH-1302: Add support for the arrow.json extension type Change-Id: I94ff9448ea6a2ac4089489c49d2bc61f1312a16f --- .../org/apache/arrow/vector/VectorLoader.java | 6 +- .../apache/arrow/vector/VectorUnloader.java | 8 +- .../arrow/vector/extension/JsonType.java | 132 +++++++++ .../arrow/vector/extension/JsonVector.java | 118 ++++++++ .../org/apache/arrow/vector/TestJsonType.java | 255 ++++++++++++++++++ 5 files changed, 516 insertions(+), 3 deletions(-) create mode 100644 vector/src/main/java/org/apache/arrow/vector/extension/JsonType.java create mode 100644 vector/src/main/java/org/apache/arrow/vector/extension/JsonVector.java create mode 100644 vector/src/test/java/org/apache/arrow/vector/TestJsonType.java diff --git a/vector/src/main/java/org/apache/arrow/vector/VectorLoader.java b/vector/src/main/java/org/apache/arrow/vector/VectorLoader.java index 9b9a890346..9c3942ec20 100644 --- a/vector/src/main/java/org/apache/arrow/vector/VectorLoader.java +++ b/vector/src/main/java/org/apache/arrow/vector/VectorLoader.java @@ -106,11 +106,15 @@ private void loadBuffers( Iterator nodes, CompressionCodec codec, Iterator variadicBufferCounts) { + FieldVector storageVector = vector; + while (storageVector instanceof ExtensionTypeVector) { + storageVector = ((ExtensionTypeVector) storageVector).getUnderlyingVector(); + } checkArgument(nodes.hasNext(), "no more field nodes for field %s and vector %s", field, vector); ArrowFieldNode fieldNode = nodes.next(); // variadicBufferLayoutCount will be 0 for vectors of a type except BaseVariableWidthViewVector long variadicBufferLayoutCount = 0; - if (vector instanceof BaseVariableWidthViewVector) { + if (storageVector instanceof BaseVariableWidthViewVector) { if (variadicBufferCounts.hasNext()) { variadicBufferLayoutCount = variadicBufferCounts.next(); } else { diff --git a/vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java b/vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java index 342f210b82..66d0debff0 100644 --- a/vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java +++ b/vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java @@ -104,14 +104,18 @@ private void appendNodes( List nodes, List buffers, List variadicBufferCounts) { + FieldVector storageVector = vector; + while (storageVector instanceof ExtensionTypeVector) { + storageVector = ((ExtensionTypeVector) storageVector).getUnderlyingVector(); + } nodes.add( new ArrowFieldNode(vector.getValueCount(), includeNullCount ? vector.getNullCount() : -1)); List fieldBuffers = vector.getFieldBuffers(); - long variadicBufferCount = getVariadicBufferCount(vector); + long variadicBufferCount = getVariadicBufferCount(storageVector); int expectedBufferCount = (int) (TypeLayout.getTypeBufferCount(vector.getField().getType()) + variadicBufferCount); // only update variadicBufferCounts for vectors that have variadic buffers - if (vector instanceof BaseVariableWidthViewVector) { + if (storageVector instanceof BaseVariableWidthViewVector) { variadicBufferCounts.add(variadicBufferCount); } if (fieldBuffers.size() != expectedBufferCount) { diff --git a/vector/src/main/java/org/apache/arrow/vector/extension/JsonType.java b/vector/src/main/java/org/apache/arrow/vector/extension/JsonType.java new file mode 100644 index 0000000000..be3ebb4f7b --- /dev/null +++ b/vector/src/main/java/org/apache/arrow/vector/extension/JsonType.java @@ -0,0 +1,132 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.vector.extension; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.DeserializationFeature; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import java.util.Collections; +import java.util.Objects; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.ExtensionTypeRegistry; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; + +/** + * Canonical extension type for UTF-8 encoded RFC 8259 JSON values. + * + *

The storage type is {@link ArrowType.Utf8}, {@link ArrowType.LargeUtf8}, or {@link + * ArrowType.Utf8View}. Values use the corresponding string vector; this type does not parse or + * validate individual JSON values. + * + *

Register the type before reading schemas containing {@code arrow.json}: + * + * 

{@code
+ * JsonType.ensureRegistered();
+ * Field field = Field.nullable("json", new JsonType(ArrowType.Utf8.INSTANCE));
+ * try (JsonVector vector = (JsonVector) field.createVector(allocator)) {
+ *   VarCharVector storage = (VarCharVector) vector.getUnderlyingVector();
+ *   storage.setSafe(0, "{}".getBytes(java.nio.charset.StandardCharsets.UTF_8));
+ *   vector.setValueCount(1);
+ *   Text value = vector.getObject(0);
+ * }
+ * }
+ */ +public class JsonType extends ArrowType.ExtensionType { + public static final String EXTENSION_NAME = "arrow.json"; + private static final ObjectMapper MAPPER = + new ObjectMapper().enable(DeserializationFeature.FAIL_ON_TRAILING_TOKENS); + private final ArrowType storageType; + + /** Register a prototype that can deserialize all supported JSON storage types. */ + public static void ensureRegistered() { + ExtensionTypeRegistry.register(new JsonType(ArrowType.Utf8.INSTANCE)); + } + + /** + * Create a JSON type backed by the specified string type. + * + * @param storageType Utf8, LargeUtf8, or Utf8View + * @throws IllegalArgumentException if the storage type is not a supported string type + */ + public JsonType(ArrowType storageType) { + Objects.requireNonNull(storageType, "storageType"); + if (!(storageType instanceof ArrowType.Utf8) + && !(storageType instanceof ArrowType.LargeUtf8) + && !(storageType instanceof ArrowType.Utf8View)) { + throw new IllegalArgumentException( + "arrow.json requires Utf8, LargeUtf8, or Utf8View storage, got " + storageType); + } + this.storageType = storageType; + } + + @Override + public ArrowType storageType() { + return storageType; + } + + @Override + public String extensionName() { + return EXTENSION_NAME; + } + + @Override + public boolean extensionEquals(ExtensionType other) { + return other instanceof JsonType && storageType.equals(other.storageType()); + } + + @Override + public String serialize() { + return ""; + } + + @Override + public ArrowType deserialize(ArrowType storageType, String serializedData) { + JsonType type = new JsonType(storageType); + if (serializedData == null) { + throw new InvalidExtensionMetadataException("arrow.json metadata must not be null"); + } + if (!serializedData.isEmpty()) { + try { + JsonNode metadata = MAPPER.readTree(serializedData); + if (metadata == null || !metadata.isObject()) { + throw new InvalidExtensionMetadataException("arrow.json metadata must be a JSON object"); + } + } catch (JsonProcessingException e) { + throw new InvalidExtensionMetadataException("arrow.json metadata is invalid", e); + } + } + return type; + } + + @Override + public boolean isComplex() { + return false; + } + + @Override + public FieldVector getNewVector(String name, FieldType fieldType, BufferAllocator allocator) { + Field field = new Field(name, fieldType, Collections.emptyList()); + FieldType storageFieldType = + new FieldType(fieldType.isNullable(), storageType, fieldType.getDictionary(), null); + FieldVector storage = storageFieldType.createNewSingleVector(name, allocator, null); + return new JsonVector(field, allocator, storage); + } +} diff --git a/vector/src/main/java/org/apache/arrow/vector/extension/JsonVector.java b/vector/src/main/java/org/apache/arrow/vector/extension/JsonVector.java new file mode 100644 index 0000000000..549b1a2190 --- /dev/null +++ b/vector/src/main/java/org/apache/arrow/vector/extension/JsonVector.java @@ -0,0 +1,118 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.vector.extension; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.util.hash.ArrowBufHasher; +import org.apache.arrow.vector.ExtensionTypeVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.ValueIterableVector; +import org.apache.arrow.vector.ValueVector; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.util.CallBack; +import org.apache.arrow.vector.util.Text; +import org.apache.arrow.vector.util.TransferPair; + +/** + * A JSON extension vector backed by a string vector. + * + *

Use {@link Field#createVector(BufferAllocator)} with a {@link JsonType} field to create an + * instance. Write UTF-8 JSON through {@link #getUnderlyingVector()}; values are not parsed or + * validated. + */ +public class JsonVector extends ExtensionTypeVector + implements ValueIterableVector { + private final Field field; + + JsonVector(Field field, BufferAllocator allocator, FieldVector underlyingVector) { + super(field, allocator, underlyingVector); + this.field = field; + } + + @Override + public Field getField() { + return field; + } + + @Override + public Text getObject(int index) { + return (Text) getUnderlyingVector().getObject(index); + } + + @Override + public TransferPair getTransferPair(BufferAllocator allocator) { + return getTransferPair(field, allocator); + } + + @Override + public TransferPair getTransferPair(String name, BufferAllocator allocator) { + return getTransferPair(new Field(name, field.getFieldType(), field.getChildren()), allocator); + } + + @Override + public TransferPair getTransferPair(String name, BufferAllocator allocator, CallBack callBack) { + return getTransferPair(name, allocator); + } + + @Override + public TransferPair getTransferPair(Field targetField, BufferAllocator allocator) { + return makeTransferPair(targetField.createVector(allocator)); + } + + @Override + public TransferPair getTransferPair( + Field targetField, BufferAllocator allocator, CallBack callBack) { + return getTransferPair(targetField, allocator); + } + + @Override + public TransferPair makeTransferPair(ValueVector target) { + JsonVector to = (JsonVector) target; + TransferPair storagePair = getUnderlyingVector().makeTransferPair(to.getUnderlyingVector()); + return new TransferPair() { + @Override + public void transfer() { + storagePair.transfer(); + } + + @Override + public void splitAndTransfer(int startIndex, int length) { + storagePair.splitAndTransfer(startIndex, length); + } + + @Override + public JsonVector getTo() { + return to; + } + + @Override + public void copyValueSafe(int fromIndex, int toIndex) { + storagePair.copyValueSafe(fromIndex, toIndex); + } + }; + } + + @Override + public int hashCode(int index) { + return hashCode(index, null); + } + + @Override + public int hashCode(int index, ArrowBufHasher hasher) { + return getUnderlyingVector().hashCode(index, hasher); + } +} diff --git a/vector/src/test/java/org/apache/arrow/vector/TestJsonType.java b/vector/src/test/java/org/apache/arrow/vector/TestJsonType.java new file mode 100644 index 0000000000..17d09c6a05 --- /dev/null +++ b/vector/src/test/java/org/apache/arrow/vector/TestJsonType.java @@ -0,0 +1,255 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.vector; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +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.RootAllocator; +import org.apache.arrow.vector.extension.InvalidExtensionMetadataException; +import org.apache.arrow.vector.extension.JsonType; +import org.apache.arrow.vector.extension.JsonVector; +import org.apache.arrow.vector.ipc.ArrowStreamReader; +import org.apache.arrow.vector.ipc.ArrowStreamWriter; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.ExtensionTypeRegistry; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +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.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 { + static Stream storageTypes() { + return Stream.of( + ArrowType.Utf8.INSTANCE, ArrowType.LargeUtf8.INSTANCE, ArrowType.Utf8View.INSTANCE); + } + + @ParameterizedTest + @MethodSource("storageTypes") + void testType(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); + } + + @ParameterizedTest + @NullSource + @ValueSource(strings = {" ", "null", "[]", "1", "true", "\"json\"", "{", "{} {}", "{} trailing"}) + void testInvalidMetadata(String metadata) { + assertThrows( + InvalidExtensionMetadataException.class, + () -> new JsonType(ArrowType.Utf8.INSTANCE).deserialize(ArrowType.Utf8.INSTANCE, metadata)); + } + + @Test + void testInvalidStorage() { + for (ArrowType storage : + new ArrowType[] { + ArrowType.Binary.INSTANCE, + ArrowType.LargeBinary.INSTANCE, + ArrowType.BinaryView.INSTANCE, + ArrowType.Null.INSTANCE, + new ArrowType.Int(32, true) + }) { + assertThrows(IllegalArgumentException.class, () -> new JsonType(storage)); + assertThrows( + IllegalArgumentException.class, + () -> new JsonType(ArrowType.Utf8.INSTANCE).deserialize(storage, "")); + } + } + + @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)); + assertEquals(schema, Schema.deserializeMessage(ByteBuffer.wrap(schema.serializeAsMessage()))); + } + } + + // Generated with PyArrow 24.0.0 using pa.schema([pa.field("json", pa.json_(t), + // 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))); + ArrowType[] storage = storageTypes().toArray(ArrowType[]::new); + for (int i = 0; i < storage.length; i++) { + assertEquals(field(storage[i], false), schema.getFields().get(i)); + } + } + + private static Field field(ArrowType storage, boolean nullable) { + return new Field( + "json", + new FieldType( + nullable, new JsonType(storage), null, Collections.singletonMap("custom", "preserved")), + Collections.emptyList()); + } + + @ParameterizedTest + @MethodSource("storageTypes") + void testTransfer(ArrowType storage) { + Field field = field(storage, true); + try (RootAllocator allocator = new RootAllocator(); + 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); + TransferPair split = source.getTransferPair("copy", allocator); + try (JsonVector target = assertInstanceOf(JsonVector.class, split.getTo())) { + split.splitAndTransfer(0, 2); + assertEquals("copy", target.getName()); + assertEquals(field.getFieldType(), target.getField().getFieldType()); + assertEquals(source.getObject(0), 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); + target.setValueCount(2); + assertEquals(source.getObject(0), 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()); + 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); + } + } + } +} From afc2ca1ac05d0cd45a18f9bd6684446f4d006dec Mon Sep 17 00:00:00 2001 From: Xin Sun Date: Sun, 20 Sep 2026 11:17:49 +0800 Subject: [PATCH 2/3] GH-1302: Simplify JSON extension tests Change-Id: I2f961fbc983bd29bffd2b4280021f7b696e3ad50 --- .../org/apache/arrow/vector/TestJsonType.java | 215 +++++++++--------- .../test/resources/pyarrow_json_schema.arrow | Bin 0 -> 728 bytes 2 files changed, 110 insertions(+), 105 deletions(-) create mode 100644 vector/src/test/resources/pyarrow_json_schema.arrow diff --git a/vector/src/test/java/org/apache/arrow/vector/TestJsonType.java b/vector/src/test/java/org/apache/arrow/vector/TestJsonType.java index 17d09c6a05..ebad0d148f 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestJsonType.java +++ b/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,6 +46,8 @@ 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; @@ -50,6 +55,25 @@ 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 storageTypes() { return Stream.of( ArrowType.Utf8.INSTANCE, ArrowType.LargeUtf8.INSTANCE, ArrowType.Utf8View.INSTANCE); @@ -57,21 +81,26 @@ static Stream storageTypes() { @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,43 +166,34 @@ 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()); } } @@ -192,64 +201,60 @@ void testTransfer(ArrowType storage) { @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)); + } } diff --git a/vector/src/test/resources/pyarrow_json_schema.arrow b/vector/src/test/resources/pyarrow_json_schema.arrow new file mode 100644 index 0000000000000000000000000000000000000000..6c44bff645566eb27323659c317fd097ce9f219a GIT binary patch literal 728 zcmbVKJ8r^25FICLFi{H2a^W*j|HHsF#&aEu%x1tkR~1!b<0 zEbkeR5JhA&(p$}p=kwmKt+jUdAmR({NuH%6u6XcZjzTPmJe=x$iaa9PI3oKxEAswl zE!-%?@eQz=K1~6iMV8WMUj+;S<`2KtMrWV Date: Sun, 20 Sep 2026 11:21:23 +0800 Subject: [PATCH 3/3] GH-1302: Follow vector transfer implementation style Change-Id: If49bdef2c665414b786ea0f2b60d0336b3839b9b --- .../org/apache/arrow/vector/VectorLoader.java | 2 +- .../arrow/vector/extension/JsonVector.java | 54 +++++++++++-------- 2 files changed, 32 insertions(+), 24 deletions(-) diff --git a/vector/src/main/java/org/apache/arrow/vector/VectorLoader.java b/vector/src/main/java/org/apache/arrow/vector/VectorLoader.java index 9c3942ec20..9e2992761f 100644 --- a/vector/src/main/java/org/apache/arrow/vector/VectorLoader.java +++ b/vector/src/main/java/org/apache/arrow/vector/VectorLoader.java @@ -112,7 +112,7 @@ private void loadBuffers( } checkArgument(nodes.hasNext(), "no more field nodes for field %s and vector %s", field, vector); ArrowFieldNode fieldNode = nodes.next(); - // variadicBufferLayoutCount will be 0 for vectors of a type except BaseVariableWidthViewVector + // Only view storage has variadic buffers. long variadicBufferLayoutCount = 0; if (storageVector instanceof BaseVariableWidthViewVector) { if (variadicBufferCounts.hasNext()) { diff --git a/vector/src/main/java/org/apache/arrow/vector/extension/JsonVector.java b/vector/src/main/java/org/apache/arrow/vector/extension/JsonVector.java index 549b1a2190..76c8579d06 100644 --- a/vector/src/main/java/org/apache/arrow/vector/extension/JsonVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/extension/JsonVector.java @@ -81,29 +81,7 @@ public TransferPair getTransferPair( @Override public TransferPair makeTransferPair(ValueVector target) { - JsonVector to = (JsonVector) target; - TransferPair storagePair = getUnderlyingVector().makeTransferPair(to.getUnderlyingVector()); - return new TransferPair() { - @Override - public void transfer() { - storagePair.transfer(); - } - - @Override - public void splitAndTransfer(int startIndex, int length) { - storagePair.splitAndTransfer(startIndex, length); - } - - @Override - public JsonVector getTo() { - return to; - } - - @Override - public void copyValueSafe(int fromIndex, int toIndex) { - storagePair.copyValueSafe(fromIndex, toIndex); - } - }; + return new TransferImpl((JsonVector) target); } @Override @@ -115,4 +93,34 @@ public int hashCode(int index) { public int hashCode(int index, ArrowBufHasher hasher) { return getUnderlyingVector().hashCode(index, hasher); } + + private class TransferImpl implements TransferPair { + private final JsonVector to; + private final TransferPair storagePair; + + TransferImpl(JsonVector to) { + this.to = to; + this.storagePair = getUnderlyingVector().makeTransferPair(to.getUnderlyingVector()); + } + + @Override + public void transfer() { + storagePair.transfer(); + } + + @Override + public void splitAndTransfer(int startIndex, int length) { + storagePair.splitAndTransfer(startIndex, length); + } + + @Override + public JsonVector getTo() { + return to; + } + + @Override + public void copyValueSafe(int fromIndex, int toIndex) { + storagePair.copyValueSafe(fromIndex, toIndex); + } + } }