fix: adding json reader/writer

vibhatha · Jul 22, 2024 · f4c18e7 · f4c18e7
1 parent 0ca4446
commit f4c18e7
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 3 deletions.
diff --git a/java/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java b/java/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java
@@ -22,6 +22,7 @@
 import static com.fasterxml.jackson.core.JsonToken.START_OBJECT;
 import static org.apache.arrow.vector.BufferLayout.BufferType.DATA;
 import static org.apache.arrow.vector.BufferLayout.BufferType.OFFSET;
+import static org.apache.arrow.vector.BufferLayout.BufferType.SIZE;
 import static org.apache.arrow.vector.BufferLayout.BufferType.TYPE;
 import static org.apache.arrow.vector.BufferLayout.BufferType.VALIDITY;
 import static org.apache.arrow.vector.BufferLayout.BufferType.VARIADIC_DATA_BUFFERS;
@@ -72,6 +73,7 @@
 import org.apache.arrow.vector.ipc.message.ArrowFieldNode;
 import org.apache.arrow.vector.types.Types.MinorType;
 import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.ArrowType.ListView;
 import org.apache.arrow.vector.types.pojo.ArrowType.Union;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.Schema;
@@ -724,7 +726,7 @@ private List<ArrowBuf> readIntoBuffer(
 
     if (bufferType.equals(VALIDITY)) {
       reader = helper.BIT;
-    } else if (bufferType.equals(OFFSET)) {
+    } else if (bufferType.equals(OFFSET) || bufferType.equals(SIZE)) {
       if (type == MinorType.LARGELIST
           || type == MinorType.LARGEVARCHAR
           || type == MinorType.LARGEVARBINARY) {
@@ -888,7 +890,7 @@ private void readFromJsonIntoVector(Field field, FieldVector vector) throws IOEx
         BufferType bufferType = vectorTypes.get(v);
         nextFieldIs(bufferType.getName());
         int innerBufferValueCount = valueCount;
-        if (bufferType.equals(OFFSET) && !(type instanceof Union)) {
+        if (bufferType.equals(OFFSET) && !(type instanceof Union) && !(type instanceof ListView)) {
           /* offset buffer has 1 additional value capacity except for dense unions */
           innerBufferValueCount = valueCount + 1;
         }

diff --git a/java/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileWriter.java b/java/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileWriter.java
@@ -73,6 +73,7 @@
 import org.apache.arrow.vector.UInt4Vector;
 import org.apache.arrow.vector.UInt8Vector;
 import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.complex.BaseRepeatedValueViewVector;
 import org.apache.arrow.vector.dictionary.Dictionary;
 import org.apache.arrow.vector.dictionary.DictionaryProvider;
 import org.apache.arrow.vector.types.Types.MinorType;
@@ -229,7 +230,9 @@ private void writeFromVectorIntoJson(Field field, FieldVector vector) throws IOE
         // thus the values are only written to a single entity.
         generator.writeArrayFieldStart(bufferType.getName());
         final int bufferValueCount =
-            (bufferType.equals(OFFSET) && vector.getMinorType() != MinorType.DENSEUNION)
+            (bufferType.equals(OFFSET)
+                    && vector.getMinorType() != MinorType.DENSEUNION
+                    && vector.getMinorType() != MinorType.LISTVIEW)
                 ? valueCount + 1
                 : valueCount;
         for (int i = 0; i < bufferValueCount; i++) {
@@ -259,6 +262,7 @@ private void writeFromVectorIntoJson(Field field, FieldVector vector) throws IOE
           } else if (bufferType.equals(OFFSET)
               && vector.getValueCount() == 0
               && (vector.getMinorType() == MinorType.LIST
+                  || vector.getMinorType() == MinorType.LISTVIEW
                   || vector.getMinorType() == MinorType.MAP
                   || vector.getMinorType() == MinorType.VARBINARY
                   || vector.getMinorType() == MinorType.VARCHAR)) {
@@ -277,6 +281,14 @@ private void writeFromVectorIntoJson(Field field, FieldVector vector) throws IOE
               vectorBufferTmp.setLong(0, 0);
               writeValueToGenerator(bufferType, vectorBufferTmp, null, vector, i);
             }
+          } else if (bufferType.equals(SIZE)
+              && vector.getValueCount() == 0
+              && vector.getMinorType() == MinorType.LISTVIEW) {
+            // Empty vectors may not have allocated a sizes buffer
+            try (ArrowBuf vectorBufferTmp = vector.getAllocator().buffer(4)) {
+              vectorBufferTmp.setInt(0, 0);
+              writeValueToGenerator(bufferType, vectorBufferTmp, null, vector, i);
+            }
           } else {
             writeValueToGenerator(bufferType, vectorBuffer, null, vector, i);
           }
@@ -419,6 +431,10 @@ private void writeValueToGenerator(
         case MAP:
           generator.writeNumber(buffer.getInt((long) index * BaseVariableWidthVector.OFFSET_WIDTH));
           break;
+        case LISTVIEW:
+          generator.writeNumber(
+              buffer.getInt((long) index * BaseRepeatedValueViewVector.OFFSET_WIDTH));
+          break;
         case LARGELIST:
         case LARGEVARBINARY:
         case LARGEVARCHAR:
@@ -573,6 +589,8 @@ private void writeValueToGenerator(
         default:
           throw new UnsupportedOperationException("minor type: " + vector.getMinorType());
       }
+    } else if (bufferType.equals(SIZE)) {
+      generator.writeNumber(buffer.getInt((long) index * BaseRepeatedValueViewVector.SIZE_WIDTH));
     }
   }