[ARROW-6738][Java] Fix problems with current union comparison logic

DataDog · Jan 21, 2020 · 4d8b570 · 4d8b570
1 parent 3dc2813
commit 4d8b570
Show file tree

Hide file tree

Showing 4 changed files with 66 additions and 14 deletions.
diff --git a/java/vector/src/main/codegen/templates/UnionVector.java b/java/vector/src/main/codegen/templates/UnionVector.java
@@ -550,7 +550,7 @@ public Iterator<ValueVector> iterator() {
     return vectors.iterator();
   }
 
-    private ValueVector getVector(int index) {
+    public ValueVector getVector(int index) {
       int type = typeBuffer.getByte(index * TYPE_WIDTH);
       switch (MinorType.values()[type]) {
         case NULL:

diff --git a/java/vector/src/main/java/org/apache/arrow/vector/compare/RangeEqualsVisitor.java b/java/vector/src/main/java/org/apache/arrow/vector/compare/RangeEqualsVisitor.java
@@ -24,7 +24,6 @@
 import org.apache.arrow.util.Preconditions;
 import org.apache.arrow.vector.BaseFixedWidthVector;
 import org.apache.arrow.vector.BaseVariableWidthVector;
-import org.apache.arrow.vector.FieldVector;
 import org.apache.arrow.vector.NullVector;
 import org.apache.arrow.vector.ValueVector;
 import org.apache.arrow.vector.complex.BaseRepeatedValueVector;
@@ -189,26 +188,37 @@ public Boolean visit(NullVector left, Range range) {
   /**
    * Creates a visitor to visit child vectors.
    * It is used for complex vector types.
-   * @return the visitor for child vecors.
+   * @return the visitor for child vectors.
    */
   protected RangeEqualsVisitor createInnerVisitor(ValueVector leftInner, ValueVector rightInner) {
+    return this.createInnerVisitor(leftInner, rightInner, this.typeComparator);
+  }
+
+  protected RangeEqualsVisitor createInnerVisitor(
+          ValueVector leftInner, ValueVector rightInner,
+          BiFunction<ValueVector, ValueVector, Boolean> typeComparator) {
     return new RangeEqualsVisitor(leftInner, rightInner, typeComparator);
   }
 
   protected boolean compareUnionVectors(Range range) {
     UnionVector leftVector = (UnionVector) left;
     UnionVector rightVector = (UnionVector) right;
 
-    List<FieldVector> leftChildren = leftVector.getChildrenFromFields();
-    List<FieldVector> rightChildren = rightVector.getChildrenFromFields();
-
-    if (leftChildren.size() != rightChildren.size()) {
-      return false;
-    }
-
-    for (int k = 0; k < leftChildren.size(); k++) {
-      RangeEqualsVisitor visitor = createInnerVisitor(leftChildren.get(k), rightChildren.get(k));
-      if (!visitor.rangeEquals(range)) {
+    Range subRange = new Range(0, 0, 1);
+    for (int i = 0; i < range.getLength(); i++) {
+      subRange.setLeftStart(range.getLeftStart() + i).setRightStart(range.getRightStart() + i);
+      ValueVector leftSubVector = leftVector.getVector(range.getLeftStart() + i);
+      ValueVector rightSubVector = rightVector.getVector(range.getRightStart() + i);
+
+      if (leftSubVector == null || rightSubVector == null) {
+        if (leftSubVector == rightSubVector) {
+          continue;
+        } else {
+          return false;
+        }
+      }
+      RangeEqualsVisitor visitor = createInnerVisitor(leftSubVector, rightSubVector, true);
+      if (!visitor.rangeEquals(subRange)) {
         return false;
       }
     }

diff --git a/java/vector/src/test/java/org/apache/arrow/vector/TestValueVector.java b/java/vector/src/test/java/org/apache/arrow/vector/TestValueVector.java
@@ -2320,7 +2320,7 @@ public void testDecimalVectorEquals() {
       VectorEqualsVisitor visitor2 = new VectorEqualsVisitor();
 
       assertTrue(visitor1.vectorEquals(vector1, vector2));
-      assertFalse(visitor2.vectorEquals(vector1, vector3));
+      assertTrue(visitor2.vectorEquals(vector1, vector3));
     }
   }
 

diff --git a/java/vector/src/test/java/org/apache/arrow/vector/compare/TestRangeEqualsVisitor.java b/java/vector/src/test/java/org/apache/arrow/vector/compare/TestRangeEqualsVisitor.java
@@ -282,6 +282,48 @@ public void testUnionVectorRangeEquals() {
     }
   }
 
+  /**
+   * Test comparing two union vectors.
+   * The two vectors are different in total, but have a range with equal values.
+   */
+  @Test
+  public void testUnionVectorSubRangeEquals() {
+    try (final UnionVector vector1 = new UnionVector("union", allocator, null, null);
+         final UnionVector vector2 = new UnionVector("union", allocator, null, null);) {
+
+      final NullableUInt4Holder uInt4Holder = new NullableUInt4Holder();
+      uInt4Holder.value = 10;
+      uInt4Holder.isSet = 1;
+
+      final NullableIntHolder intHolder = new NullableIntHolder();
+      uInt4Holder.value = 20;
+      uInt4Holder.isSet = 1;
+
+      vector1.setType(0, Types.MinorType.UINT4);
+      vector1.setSafe(0, uInt4Holder);
+
+      vector1.setType(1, Types.MinorType.INT);
+      vector1.setSafe(1, intHolder);
+
+      vector1.setType(2, Types.MinorType.INT);
+      vector1.setSafe(2, intHolder);
+      vector1.setValueCount(3);
+
+      vector2.setType(0, Types.MinorType.SMALLINT);
+      vector2.setSafe(0, intHolder);
+
+      vector2.setType(1, Types.MinorType.INT);
+      vector2.setSafe(1, intHolder);
+
+      vector2.setType(2, Types.MinorType.INT);
+      vector2.setSafe(2, intHolder);
+      vector2.setValueCount(3);
+
+      RangeEqualsVisitor visitor = new RangeEqualsVisitor(vector1, vector2);
+      assertTrue(visitor.rangeEquals(new Range(1, 1, 2)));
+    }
+  }
+
   @Ignore
   @Test
   public void testEqualsWithOutTypeCheck() {