|
33 | 33 | import org.apache.arrow.vector.complex.writer.BaseWriter.StructWriter; |
34 | 34 | import org.apache.arrow.vector.complex.writer.FieldWriter; |
35 | 35 | import org.apache.arrow.vector.holders.DecimalHolder; |
| 36 | +import org.apache.arrow.vector.holders.DurationHolder; |
36 | 37 | import org.apache.arrow.vector.types.Types; |
37 | 38 | import org.apache.arrow.vector.types.pojo.ArrowType; |
38 | 39 | import org.apache.arrow.vector.types.pojo.FieldType; |
@@ -845,4 +846,137 @@ public void testCopyMapVectorWithMapValue() { |
845 | 846 | assertTrue(VectorEqualsVisitor.vectorEquals(from, to)); |
846 | 847 | } |
847 | 848 | } |
| 849 | + |
| 850 | + @Test |
| 851 | + public void testCopyDurationVector() { |
| 852 | + try (ListVector from = ListVector.empty("v", allocator); |
| 853 | + ListVector to = ListVector.empty("v", allocator)) { |
| 854 | + |
| 855 | + UnionListWriter listWriter = from.getWriter(); |
| 856 | + listWriter.allocate(); |
| 857 | + |
| 858 | + DurationHolder durationHolder = new DurationHolder(); |
| 859 | + for (int i = 0; i < COUNT; i++) { |
| 860 | + listWriter.setPosition(i); |
| 861 | + listWriter.startList(); |
| 862 | + durationHolder.value = 123456789L; |
| 863 | + listWriter.duration().write(durationHolder); |
| 864 | + listWriter.endList(); |
| 865 | + } |
| 866 | + from.setValueCount(COUNT); |
| 867 | + |
| 868 | + FieldReader in = from.getReader(); |
| 869 | + FieldWriter out = to.getWriter(); |
| 870 | + for (int i = 0; i < COUNT; i++) { |
| 871 | + in.setPosition(i); |
| 872 | + out.setPosition(i); |
| 873 | + ComplexCopier.copy(in, out); |
| 874 | + } |
| 875 | + |
| 876 | + to.setValueCount(COUNT); |
| 877 | + |
| 878 | + assertTrue(VectorEqualsVisitor.vectorEquals(from, to)); |
| 879 | + } |
| 880 | + } |
| 881 | + |
| 882 | + @Test |
| 883 | + public void testCopyListVectorWithDuration() { |
| 884 | + try (ListVector from = ListVector.empty("v", allocator); |
| 885 | + ListVector to = ListVector.empty("v", allocator)) { |
| 886 | + |
| 887 | + UnionListWriter listWriter = from.getWriter(); |
| 888 | + listWriter.allocate(); |
| 889 | + |
| 890 | + DurationHolder durationHolder = new DurationHolder(); |
| 891 | + for (int i = 0; i < COUNT; i++) { |
| 892 | + listWriter.setPosition(i); |
| 893 | + listWriter.startList(); |
| 894 | + durationHolder.value = 123456789L; |
| 895 | + listWriter.duration().write(durationHolder); |
| 896 | + listWriter.endList(); |
| 897 | + } |
| 898 | + from.setValueCount(COUNT); |
| 899 | + |
| 900 | + // Copy values |
| 901 | + FieldReader in = from.getReader(); |
| 902 | + FieldWriter out = to.getWriter(); |
| 903 | + for (int i = 0; i < COUNT; i++) { |
| 904 | + in.setPosition(i); |
| 905 | + out.setPosition(i); |
| 906 | + ComplexCopier.copy(in, out); |
| 907 | + } |
| 908 | + |
| 909 | + to.setValueCount(COUNT); |
| 910 | + |
| 911 | + assertTrue(VectorEqualsVisitor.vectorEquals(from, to)); |
| 912 | + } |
| 913 | + } |
| 914 | + |
| 915 | + @Test |
| 916 | + public void testCopyMapVectorWithDurationValue() { |
| 917 | + try (MapVector from = MapVector.empty("v", allocator, false); |
| 918 | + MapVector to = MapVector.empty("v", allocator, false)) { |
| 919 | + |
| 920 | + UnionMapWriter mapWriter = from.getWriter(); |
| 921 | + mapWriter.allocate(); |
| 922 | + |
| 923 | + DurationHolder durationHolder = new DurationHolder(); |
| 924 | + for (int i = 0; i < COUNT; i++) { |
| 925 | + mapWriter.setPosition(i); |
| 926 | + mapWriter.startMap(); |
| 927 | + mapWriter.startEntry(); |
| 928 | + mapWriter.key().integer().writeInt(i); |
| 929 | + durationHolder.value = 123456789L; |
| 930 | + mapWriter.value().duration().write(durationHolder); |
| 931 | + mapWriter.endEntry(); |
| 932 | + mapWriter.endMap(); |
| 933 | + } |
| 934 | + from.setValueCount(COUNT); |
| 935 | + |
| 936 | + FieldReader in = from.getReader(); |
| 937 | + FieldWriter out = to.getWriter(); |
| 938 | + for (int i = 0; i < COUNT; i++) { |
| 939 | + in.setPosition(i); |
| 940 | + out.setPosition(i); |
| 941 | + ComplexCopier.copy(in, out); |
| 942 | + } |
| 943 | + |
| 944 | + to.setValueCount(COUNT); |
| 945 | + |
| 946 | + assertTrue(VectorEqualsVisitor.vectorEquals(from, to)); |
| 947 | + } |
| 948 | + } |
| 949 | + |
| 950 | + @Test |
| 951 | + public void testCopyStructVectorWithDurationValue() { |
| 952 | + try (StructVector from = StructVector.empty("v", allocator); |
| 953 | + StructVector to = StructVector.empty("v", allocator)) { |
| 954 | + |
| 955 | + from.allocateNew(); |
| 956 | + |
| 957 | + NullableStructWriter structWriter = from.getWriter(); |
| 958 | + DurationHolder durationHolder = new DurationHolder(); |
| 959 | + |
| 960 | + for (int i = 0; i < COUNT; i++) { |
| 961 | + structWriter.setPosition(i); |
| 962 | + structWriter.start(); |
| 963 | + durationHolder.value = 123456789L; |
| 964 | + structWriter.duration("durationField").write(durationHolder); // 使用 DurationHolder |
| 965 | + structWriter.end(); |
| 966 | + } |
| 967 | + |
| 968 | + from.setValueCount(COUNT); |
| 969 | + |
| 970 | + FieldReader in = from.getReader(); |
| 971 | + FieldWriter out = to.getWriter(); |
| 972 | + for (int i = 0; i < COUNT; i++) { |
| 973 | + in.setPosition(i); |
| 974 | + out.setPosition(i); |
| 975 | + ComplexCopier.copy(in, out); |
| 976 | + } |
| 977 | + to.setValueCount(COUNT); |
| 978 | + |
| 979 | + assertTrue(VectorEqualsVisitor.vectorEquals(from, to)); |
| 980 | + } |
| 981 | + } |
848 | 982 | } |
0 commit comments