apacheGH-38246 more improvements and tests

bookingnizer · Oct 13, 2023 · 454d196 · 454d196
1 parent ccf50a5
commit 454d196
Show file tree

Hide file tree

Showing 12 changed files with 174 additions and 11 deletions.
diff --git a/java/vector/src/main/java/org/apache/arrow/vector/complex/FixedSizeListVector.java b/java/vector/src/main/java/org/apache/arrow/vector/complex/FixedSizeListVector.java
@@ -71,7 +71,7 @@ public static FixedSizeListVector empty(String name, int size, BufferAllocator a
   private FieldVector vector;
   private ArrowBuf validityBuffer;
   private final int listSize;
-  private final Field field;
+  private Field field;
 
   private UnionFixedSizeListReader reader;
   private int valueCount;
@@ -123,8 +123,11 @@ public FixedSizeListVector(Field field,
 
   @Override
   public Field getField() {
-    return field.getChildren().isEmpty() ? new Field(field.getName(), field.getFieldType(),
-        Collections.singletonList(getDataVector().getField())) : field;
+    if(!field.getChildren().isEmpty()){
+      return field;
+    }
+    field.setChildren(Collections.singletonList(getDataVector().getField()));
+    return field;
   }
 
   @Override
@@ -152,6 +155,7 @@ public void initializeChildrenFromFields(List<Field> children) {
     checkArgument(addOrGetVector.isCreated(), "Child vector already existed: %s", addOrGetVector.getVector());
 
     addOrGetVector.getVector().initializeChildrenFromFields(field.getChildren());
+    this.field.setChildren(children);
   }
 
   @Override

diff --git a/java/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java b/java/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java
@@ -101,7 +101,7 @@ public static LargeListVector empty(String name, BufferAllocator allocator) {
   protected String defaultDataVectorName = DATA_VECTOR_NAME;
   protected ArrowBuf validityBuffer;
   protected UnionLargeListReader reader;
-  private final Field field;
+  private Field field;
   private int validityAllocationSizeInBytes;
 
   /**
@@ -158,6 +158,7 @@ public void initializeChildrenFromFields(List<Field> children) {
     checkArgument(addOrGetVector.isCreated(), "Child vector already existed: %s", addOrGetVector.getVector());
 
     addOrGetVector.getVector().initializeChildrenFromFields(field.getChildren());
+    this.field.setChildren(children);
   }
 
   @Override
@@ -814,8 +815,11 @@ public int getBufferSizeFor(int valueCount) {
 
   @Override
   public Field getField() {
-    return field.getChildren().isEmpty() ? new Field(field.getName(), field.getFieldType(),
-        Collections.singletonList(getDataVector().getField())) : field;
+    if(!field.getChildren().isEmpty()){
+      return field;
+    }
+    field.setChildren(Collections.singletonList(getDataVector().getField()));
+    return field;
   }
 
   @Override

diff --git a/java/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java b/java/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java
@@ -75,7 +75,7 @@ public static ListVector empty(String name, BufferAllocator allocator) {
   protected ArrowBuf validityBuffer;
   protected UnionListReader reader;
   private CallBack callBack;
-  protected final Field field;
+  protected Field field;
   protected int validityAllocationSizeInBytes;
 
   /**
@@ -126,6 +126,7 @@ public void initializeChildrenFromFields(List<Field> children) {
     checkArgument(addOrGetVector.isCreated(), "Child vector already existed: %s", addOrGetVector.getVector());
 
     addOrGetVector.getVector().initializeChildrenFromFields(field.getChildren());
+    this.field.setChildren(children);
   }
 
   @Override
@@ -662,8 +663,11 @@ public int getBufferSizeFor(int valueCount) {
 
   @Override
   public Field getField() {
-    return field.getChildren().isEmpty() ? new Field(field.getName(), field.getFieldType(),
-        Collections.singletonList(getDataVector().getField())) : field;
+    if(!field.getChildren().isEmpty()){
+      return field;
+    }
+    field.setChildren(Collections.singletonList(getDataVector().getField()));
+    return field;
   }
 
   @Override

diff --git a/java/vector/src/main/java/org/apache/arrow/vector/complex/MapVector.java b/java/vector/src/main/java/org/apache/arrow/vector/complex/MapVector.java
@@ -104,6 +104,7 @@ public void initializeChildrenFromFields(List<Field> children) {
     checkArgument(addOrGetVector.isCreated(), "Child vector already existed: %s", addOrGetVector.getVector());
 
     addOrGetVector.getVector().initializeChildrenFromFields(structField.getChildren());
+    this.field.setChildren(children);
   }
 
   /**

diff --git a/java/vector/src/main/java/org/apache/arrow/vector/complex/NonNullableStructVector.java b/java/vector/src/main/java/org/apache/arrow/vector/complex/NonNullableStructVector.java
@@ -474,7 +474,8 @@ public Field getField() {
     for (ValueVector child : getChildren()) {
       children.add(child.getField());
     }
-    return new Field(name, field.getFieldType(), children);
+    field.setChildren(children);
+    return field;
   }
 
   @Override

diff --git a/java/vector/src/main/java/org/apache/arrow/vector/types/pojo/Field.java b/java/vector/src/main/java/org/apache/arrow/vector/types/pojo/Field.java
@@ -70,7 +70,7 @@ public static Field notNullable(String name, ArrowType type) {
 
   private final String name;
   private final FieldType fieldType;
-  private final List<Field> children;
+  private List<Field> children;
 
   private Field(
       String name,
@@ -256,6 +256,10 @@ public List<Field> getChildren() {
     return children;
   }
 
+  public void setChildren(List<Field> children) {
+     this.children = Collections2.toImmutableList(children);
+  }
+
   @JsonIgnore
   public Map<String, String> getMetadata() {
     return fieldType.getMetadata();

diff --git a/java/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java b/java/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java
@@ -20,6 +20,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 
 import java.util.ArrayList;
@@ -992,6 +993,28 @@ public void testTotalCapacity() {
     }
   }
 
+  @Test
+  public void testGetTransferPairWithField() throws Exception {
+    try (final LargeListVector fromVector = LargeListVector.empty("list", allocator)) {
+
+      UnionLargeListWriter writer = fromVector.getWriter();
+      writer.allocate();
+
+      //set some values
+      writer.startList();
+      writer.integer().writeInt(1);
+      writer.integer().writeInt(2);
+      writer.endList();
+      fromVector.setValueCount(2);
+
+      final TransferPair transferPair = fromVector.getTransferPair(fromVector.getField(),
+          allocator);
+      final LargeListVector toVector = (LargeListVector) transferPair.getTo();
+      // Field inside a new vector created by reusing a field should be the same in memory as the original field.
+      assertSame(toVector.getField(), fromVector.getField());
+    }
+  }
+
   private void writeIntValues(UnionLargeListWriter writer, int[] values) {
     writer.startList();
     for (int v: values) {

diff --git a/java/vector/src/test/java/org/apache/arrow/vector/TestLargeVarBinaryVector.java b/java/vector/src/test/java/org/apache/arrow/vector/TestLargeVarBinaryVector.java
@@ -18,12 +18,14 @@
 package org.apache.arrow.vector;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 
 import org.apache.arrow.memory.ArrowBuf;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.vector.holders.NullableLargeVarBinaryHolder;
+import org.apache.arrow.vector.util.TransferPair;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -101,4 +103,20 @@ public void testSetNullableLargeVarBinaryHolderSafe() {
       buf.close();
     }
   }
+
+  @Test
+  public void testGetTransferPairWithField() {
+    try (BufferAllocator childAllocator1 = allocator.newChildAllocator("child1", 1000000, 1000000);
+        LargeVarBinaryVector v1 = new LargeVarBinaryVector("v1", childAllocator1)) {
+      v1.allocateNew();
+      v1.setSafe(4094, "hello world".getBytes(), 0, 11);
+      v1.setValueCount(4001);
+
+      TransferPair tp = v1.getTransferPair(v1.getField(), allocator);
+      tp.transfer();
+      LargeVarBinaryVector v2 = (LargeVarBinaryVector) tp.getTo();
+      assertSame(v1.getField(), v2.getField());
+      v2.clear();
+    }
+  }
 }
diff --git a/java/vector/src/test/java/org/apache/arrow/vector/TestLargeVarCharVector.java b/java/vector/src/test/java/org/apache/arrow/vector/TestLargeVarCharVector.java
@@ -21,6 +21,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 
 import java.nio.charset.StandardCharsets;
@@ -794,6 +795,22 @@ public void testNullableType() {
     }
   }
 
+  @Test
+  public void testGetTransferPairWithField() {
+    try (BufferAllocator childAllocator1 = allocator.newChildAllocator("child1", 1000000, 1000000);
+        LargeVarCharVector v1 = new LargeVarCharVector("v1", childAllocator1)) {
+      v1.allocateNew();
+      v1.setSafe(4094, "hello world".getBytes(), 0, 11);
+      v1.setValueCount(4001);
+
+      TransferPair tp = v1.getTransferPair(v1.getField(), allocator);
+      tp.transfer();
+      LargeVarCharVector v2 = (LargeVarCharVector) tp.getTo();
+      assertSame(v1.getField(), v2.getField());
+      v2.clear();
+    }
+  }
+
   private void populateLargeVarcharVector(final LargeVarCharVector vector, int valueCount, String[] values) {
     for (int i = 0; i < valueCount; i += 3) {
       final String s = String.format("%010d", i);

diff --git a/java/vector/src/test/java/org/apache/arrow/vector/TestListVector.java b/java/vector/src/test/java/org/apache/arrow/vector/TestListVector.java
@@ -20,6 +20,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
@@ -1135,6 +1136,26 @@ public void testTotalCapacity() {
     }
   }
 
+  @Test
+  public void testGetTransferPairWithField() {
+    try (ListVector fromVector = ListVector.empty("input", allocator)) {
+      UnionListWriter writer = fromVector.getWriter();
+      writer.allocate();
+      writer.setPosition(0); // optional
+      writer.startList();
+      writer.bigInt().writeBigInt(1);
+      writer.bigInt().writeBigInt(2);
+      writer.bigInt().writeBigInt(3);
+      writer.endList();
+      writer.setValueCount(1);
+      final TransferPair transferPair = fromVector.getTransferPair(fromVector.getField(),
+          allocator);
+      final ListVector toVector = (ListVector) transferPair.getTo();
+      // Field inside a new vector created by reusing a field should be the same in memory as the original field.
+      assertSame(toVector.getField(), fromVector.getField());
+    }
+  }
+
   private void writeIntValues(UnionListWriter writer, int[] values) {
     writer.startList();
     for (int v: values) {

diff --git a/java/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java b/java/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java
@@ -1133,4 +1133,26 @@ public void testGetTransferPair() {
       vector.clear();
     }
   }
+
+  @Test
+  public void testGetTransferPairWithField() {
+    try (MapVector mapVector = MapVector.empty("mapVector", allocator, false)) {
+
+      FieldType type = new FieldType(false, ArrowType.Struct.INSTANCE, null, null);
+      AddOrGetResult<StructVector> addResult = mapVector.addOrGetVector(type);
+      FieldType keyType = new FieldType(false, MinorType.BIGINT.getType(), null, null);
+      FieldType valueType = FieldType.nullable(MinorType.FLOAT8.getType());
+      addResult.getVector().addOrGet(MapVector.KEY_NAME, keyType, BigIntVector.class);
+      addResult.getVector().addOrGet(MapVector.VALUE_NAME, valueType, Float8Vector.class);
+      mapVector.allocateNew();
+      mapVector.setValueCount(0);
+
+      assertEquals(-1, mapVector.getLastSet());
+      TransferPair tp = mapVector.getTransferPair(mapVector.getField(), allocator);
+      tp.transfer();
+      MapVector toVector = (MapVector) tp.getTo();
+      assertSame(toVector.getField(), mapVector.getField());
+      toVector.clear();
+    }
+  }
 }
diff --git a/java/vector/src/test/java/org/apache/arrow/vector/TestStructVector.java b/java/vector/src/test/java/org/apache/arrow/vector/TestStructVector.java
@@ -29,11 +29,16 @@
 import org.apache.arrow.vector.complex.ListVector;
 import org.apache.arrow.vector.complex.StructVector;
 import org.apache.arrow.vector.complex.UnionVector;
+import org.apache.arrow.vector.complex.impl.NullableStructWriter;
+import org.apache.arrow.vector.complex.writer.Float8Writer;
+import org.apache.arrow.vector.complex.writer.IntWriter;
 import org.apache.arrow.vector.holders.ComplexHolder;
+import org.apache.arrow.vector.types.Types;
 import org.apache.arrow.vector.types.Types.MinorType;
 import org.apache.arrow.vector.types.pojo.ArrowType.Struct;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.FieldType;
+import org.apache.arrow.vector.util.TransferPair;
 import org.junit.After;
 import org.junit.Assert;
 import org.junit.Before;
@@ -290,4 +295,43 @@ public void testTypedGetters() {
       assertEquals(IntVector.class, s1.getVectorById(0, IntVector.class).getClass());
     }
   }
+
+  @Test
+  public void testGetTransferPair() {
+    try (final StructVector fromVector = simpleStructVector("s1", allocator)) {
+      TransferPair tp = fromVector.getTransferPair(fromVector.getField(), allocator);
+      final StructVector toVector = (StructVector) tp.getTo();
+      // Field inside a new vector created by reusing a field should be the same in memory as the original field.
+      assertSame(toVector.getField(), fromVector.getField());
+      toVector.clear();
+    }
+  }
+
+  private StructVector simpleStructVector(String name, BufferAllocator allocator) {
+    final String INT_COL = "struct_int_child";
+    final String FLT_COL = "struct_flt_child";
+    StructVector structVector = StructVector.empty(name, allocator);
+    final int size = 6; // number of structs
+
+    NullableStructWriter structWriter = structVector.getWriter();
+    structVector.addOrGet(
+        INT_COL, FieldType.nullable(Types.MinorType.INT.getType()), IntVector.class);
+    structVector.addOrGet(
+        FLT_COL, FieldType.nullable(Types.MinorType.INT.getType()), IntVector.class);
+    structVector.allocateNew();
+    IntWriter intWriter = structWriter.integer(INT_COL);
+    Float8Writer float8Writer = structWriter.float8(FLT_COL);
+
+    for (int i = 0; i < size; i++) {
+      structWriter.setPosition(i);
+      structWriter.start();
+      intWriter.writeInt(i);
+      float8Writer.writeFloat8(i * .1);
+      structWriter.end();
+    }
+
+    structWriter.setValueCount(size);
+
+    return structVector;
+  }
 }