feat: adding rest of the test cases: snapshot 1: requires transferPai…

…r to proceed
vibhatha · Jun 7, 2024 · b89cbe6 · b89cbe6
1 parent d0fdb61
commit b89cbe6
Show file tree

Hide file tree

Showing 5 changed files with 217 additions and 81 deletions.
diff --git a/java/vector/src/main/codegen/templates/AbstractFieldWriter.java b/java/vector/src/main/codegen/templates/AbstractFieldWriter.java
@@ -184,6 +184,12 @@ public ListWriter list() {
     return null;
   }
 
+  @Override
+  public ListWriter listView() {
+    fail("ListView");
+    return null;
+  }
+
   @Override
   public MapWriter map() {
     fail("Map");
@@ -202,6 +208,12 @@ public ListWriter list(String name) {
     return null;
   }
 
+  @Override
+  public ListWriter listView(String name) {
+    fail("ListView");
+    return null;
+  }
+
   @Override
   public MapWriter map(String name) {
     fail("Map");

diff --git a/java/vector/src/main/codegen/templates/BaseWriter.java b/java/vector/src/main/codegen/templates/BaseWriter.java
@@ -62,6 +62,7 @@ public interface StructWriter extends BaseWriter {
     void copyReaderToField(String name, FieldReader reader);
     StructWriter struct(String name);
     ListWriter list(String name);
+    ListWriter listView(String name);
     MapWriter map(String name);
     MapWriter map(String name, boolean keysSorted);
     void start();
@@ -73,6 +74,7 @@ public interface ListWriter extends BaseWriter {
     void endList();
     StructWriter struct();
     ListWriter list();
+    ListWriter listView();
     MapWriter map();
     MapWriter map(boolean keysSorted);
     void copyReader(FieldReader reader);

diff --git a/java/vector/src/main/codegen/templates/StructWriters.java b/java/vector/src/main/codegen/templates/StructWriters.java
@@ -200,6 +200,28 @@ public ListWriter list(String name) {
     return writer;
   }
 
+  @Override
+  public ListWriter listView(String name) {
+    String finalName = handleCase(name);
+    FieldWriter writer = fields.get(finalName);
+    int vectorCount = container.size();
+    if(writer == null) {
+      FieldType fieldType = new FieldType(addVectorAsNullable, MinorType.LISTVIEW.getType(), null, null);
+      writer = new PromotableWriter(container.addOrGet(name, fieldType, ListViewVector.class), container, getNullableStructWriterFactory());
+      if (container.size() > vectorCount) {
+        writer.allocate();
+      }
+      writer.setPosition(idx());
+      fields.put(finalName, writer);
+    } else {
+      if (writer instanceof PromotableWriter) {
+        // ensure writers are initialized
+        ((PromotableWriter)writer).getWriter(MinorType.LISTVIEW);
+      }
+    }
+    return writer;
+  }
+
   @Override
   public MapWriter map(String name) {
     return map(name, false);

diff --git a/java/vector/src/main/java/org/apache/arrow/vector/complex/impl/PromotableWriter.java b/java/vector/src/main/java/org/apache/arrow/vector/complex/impl/PromotableWriter.java
@@ -244,6 +244,9 @@ private void setWriter(ValueVector v) {
       case LIST:
         writer = new UnionListWriter((ListVector) vector, nullableStructWriterFactory);
         break;
+      case LISTVIEW:
+        writer = new UnionListViewWriter((ListViewVector) vector, nullableStructWriterFactory);
+        break;
       case MAP:
         writer = new UnionMapWriter((MapVector) vector);
         break;