GH-558: ComplexCopier supports for Duration columns.

apache · Jan 24, 2025 · d6a20f6 · d6a20f6
1 parent 7cfc89d
commit d6a20f6
Show file tree

Hide file tree

Showing 2 changed files with 140 additions and 0 deletions.
diff --git a/vector/src/main/codegen/templates/ComplexCopier.java b/vector/src/main/codegen/templates/ComplexCopier.java
@@ -162,6 +162,8 @@ private static FieldWriter getStructWriterForReader(FieldReader reader, StructWr
       return (FieldWriter) writer.map(name);
     case LISTVIEW:
       return (FieldWriter) writer.listView(name);
+    case DURATION:
+      return (FieldWriter) writer.duration(name);
     default:
       throw new UnsupportedOperationException(reader.getMinorType().toString());
     }
@@ -186,6 +188,8 @@ private static FieldWriter getListWriterForReader(FieldReader reader, ListWriter
       return (FieldWriter) writer.list();
     case LISTVIEW:
       return (FieldWriter) writer.listView();
+    case DURATION:
+      return (FieldWriter) writer.duration();
     default:
       throw new UnsupportedOperationException(reader.getMinorType().toString());
     }
@@ -211,6 +215,8 @@ private static FieldWriter getMapWriterForReader(FieldReader reader, MapWriter w
         return (FieldWriter) writer.listView();
       case MAP:
         return (FieldWriter) writer.map(false);
+      case DURATION:
+        return (FieldWriter) writer.duration();
       default:
         throw new UnsupportedOperationException(reader.getMinorType().toString());
     }

diff --git a/vector/src/test/java/org/apache/arrow/vector/complex/impl/TestComplexCopier.java b/vector/src/test/java/org/apache/arrow/vector/complex/impl/TestComplexCopier.java
@@ -33,6 +33,7 @@
 import org.apache.arrow.vector.complex.writer.BaseWriter.StructWriter;
 import org.apache.arrow.vector.complex.writer.FieldWriter;
 import org.apache.arrow.vector.holders.DecimalHolder;
+import org.apache.arrow.vector.holders.DurationHolder;
 import org.apache.arrow.vector.types.Types;
 import org.apache.arrow.vector.types.pojo.ArrowType;
 import org.apache.arrow.vector.types.pojo.FieldType;
@@ -845,4 +846,137 @@ public void testCopyMapVectorWithMapValue() {
       assertTrue(VectorEqualsVisitor.vectorEquals(from, to));
     }
   }
+
+  @Test
+  public void testCopyDurationVector() {
+    try (ListVector from = ListVector.empty("v", allocator);
+        ListVector to = ListVector.empty("v", allocator)) {
+
+      UnionListWriter listWriter = from.getWriter();
+      listWriter.allocate();
+
+      DurationHolder durationHolder = new DurationHolder();
+      for (int i = 0; i < COUNT; i++) {
+        listWriter.setPosition(i);
+        listWriter.startList();
+        durationHolder.value = 123456789L;
+        listWriter.duration().write(durationHolder);
+        listWriter.endList();
+      }
+      from.setValueCount(COUNT);
+
+      FieldReader in = from.getReader();
+      FieldWriter out = to.getWriter();
+      for (int i = 0; i < COUNT; i++) {
+        in.setPosition(i);
+        out.setPosition(i);
+        ComplexCopier.copy(in, out);
+      }
+
+      to.setValueCount(COUNT);
+
+      assertTrue(VectorEqualsVisitor.vectorEquals(from, to));
+    }
+  }
+
+  @Test
+  public void testCopyListVectorWithDuration() {
+    try (ListVector from = ListVector.empty("v", allocator);
+        ListVector to = ListVector.empty("v", allocator)) {
+
+      UnionListWriter listWriter = from.getWriter();
+      listWriter.allocate();
+
+      DurationHolder durationHolder = new DurationHolder();
+      for (int i = 0; i < COUNT; i++) {
+        listWriter.setPosition(i);
+        listWriter.startList();
+        durationHolder.value = 123456789L;
+        listWriter.duration().write(durationHolder);
+        listWriter.endList();
+      }
+      from.setValueCount(COUNT);
+
+      // Copy values
+      FieldReader in = from.getReader();
+      FieldWriter out = to.getWriter();
+      for (int i = 0; i < COUNT; i++) {
+        in.setPosition(i);
+        out.setPosition(i);
+        ComplexCopier.copy(in, out);
+      }
+
+      to.setValueCount(COUNT);
+
+      assertTrue(VectorEqualsVisitor.vectorEquals(from, to));
+    }
+  }
+
+  @Test
+  public void testCopyMapVectorWithDurationValue() {
+    try (MapVector from = MapVector.empty("v", allocator, false);
+        MapVector to = MapVector.empty("v", allocator, false)) {
+
+      UnionMapWriter mapWriter = from.getWriter();
+      mapWriter.allocate();
+
+      DurationHolder durationHolder = new DurationHolder();
+      for (int i = 0; i < COUNT; i++) {
+        mapWriter.setPosition(i);
+        mapWriter.startMap();
+        mapWriter.startEntry();
+        mapWriter.key().integer().writeInt(i);
+        durationHolder.value = 123456789L;
+        mapWriter.value().duration().write(durationHolder);
+        mapWriter.endEntry();
+        mapWriter.endMap();
+      }
+      from.setValueCount(COUNT);
+
+      FieldReader in = from.getReader();
+      FieldWriter out = to.getWriter();
+      for (int i = 0; i < COUNT; i++) {
+        in.setPosition(i);
+        out.setPosition(i);
+        ComplexCopier.copy(in, out);
+      }
+
+      to.setValueCount(COUNT);
+
+      assertTrue(VectorEqualsVisitor.vectorEquals(from, to));
+    }
+  }
+
+  @Test
+  public void testCopyStructVectorWithDurationValue() {
+    try (StructVector from = StructVector.empty("v", allocator);
+        StructVector to = StructVector.empty("v", allocator)) {
+
+      from.allocateNew();
+
+      NullableStructWriter structWriter = from.getWriter();
+      DurationHolder durationHolder = new DurationHolder();
+
+      for (int i = 0; i < COUNT; i++) {
+        structWriter.setPosition(i);
+        structWriter.start();
+        durationHolder.value = 123456789L;
+        structWriter.duration("durationField").write(durationHolder); // 使用 DurationHolder
+        structWriter.end();
+      }
+
+      from.setValueCount(COUNT);
+
+      FieldReader in = from.getReader();
+      FieldWriter out = to.getWriter();
+      for (int i = 0; i < COUNT; i++) {
+        in.setPosition(i);
+        out.setPosition(i);
+        ComplexCopier.copy(in, out);
+      }
+      to.setValueCount(COUNT);
+
+      assertTrue(VectorEqualsVisitor.vectorEquals(from, to));
+    }
+  }
 }