Skip to content

Commit baa67ea

Browse files
henriquejszahenriquejsza
committed
GH-551: [Java] Add TimeStamp*TZ support to ComplexCopier
ComplexCopier skipped every minor type that has type parameters except Decimal, so copying a vector containing TimeStampSecTZ, TimeStampMilliTZ, TimeStampMicroTZ, or TimeStampNanoTZ threw UnsupportedOperationException. Generate the copy and writer-lookup cases for the timezone-aware timestamp types the same way as the other scalar types, and pass the timezone through to the struct writer so it is preserved. Add list, struct and map copy tests for TimeStampNanoTZ. Closes #551. Signed-off-by: henriquejsza <henriquejsza@gmail.com>
1 parent 2f35627 commit baa67ea

2 files changed

Lines changed: 142 additions & 3 deletions

File tree

‎vector/src/main/codegen/templates/ComplexCopier.java‎

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,10 @@
3030

3131
<#include "/@includes/vv_imports.ftl" />
3232

33+
<#function is_timestamp_tz type>
34+
<#return type?starts_with("TimeStamp") && type?ends_with("TZ")>
35+
</#function>
36+
3337
/*
3438
* This class is generated using freemarker and the ${.template_name} template.
3539
*/
@@ -121,7 +125,7 @@ public static void copy(FieldReader reader, FieldWriter writer) {
121125
<#assign fields = minor.fields!type.fields />
122126
<#assign uncappedName = name?uncap_first/>
123127

124-
<#if !minor.typeParams?? || minor.class?starts_with("Decimal") >
128+
<#if !minor.typeParams?? || minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) >
125129

126130
case ${name?upper_case}:
127131
if (reader.isSet()) {
@@ -158,6 +162,15 @@ private static FieldWriter getStructWriterForReader(FieldReader reader, StructWr
158162
return (FieldWriter) writer.${uncappedName}(name);
159163
}
160164
</#if>
165+
<#if is_timestamp_tz(minor.class)>
166+
case ${name?upper_case}:
167+
if (reader.getField().getType() instanceof ArrowType.Timestamp) {
168+
ArrowType.Timestamp type = (ArrowType.Timestamp) reader.getField().getType();
169+
return (FieldWriter) writer.${uncappedName}(name, type.getTimezone());
170+
} else {
171+
return (FieldWriter) writer.${uncappedName}(name);
172+
}
173+
</#if>
161174
162175
</#list></#list>
163176
case STRUCT:
@@ -182,7 +195,7 @@ private static FieldWriter getListWriterForReader(FieldReader reader, ListWriter
182195
<#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first />
183196
<#assign fields = minor.fields!type.fields />
184197
<#assign uncappedName = name?uncap_first/>
185-
<#if !minor.typeParams?? || minor.class?starts_with("Decimal") >
198+
<#if !minor.typeParams?? || minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) >
186199
case ${name?upper_case}:
187200
return (FieldWriter) writer.<#if name == "Int">integer<#else>${uncappedName}</#if>();
188201
</#if>
@@ -209,7 +222,7 @@ private static FieldWriter getMapWriterForReader(FieldReader reader, MapWriter w
209222
<#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first />
210223
<#assign fields = minor.fields!type.fields />
211224
<#assign uncappedName = name?uncap_first/>
212-
<#if !minor.typeParams?? || minor.class?starts_with("Decimal") >
225+
<#if !minor.typeParams?? || minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) >
213226
case ${name?upper_case}:
214227
return (FieldWriter) writer.<#if name == "Int">integer<#else>${uncappedName}</#if>();
215228
</#if>

‎vector/src/test/java/org/apache/arrow/vector/complex/impl/TestComplexCopier.java‎

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,13 @@
1616
*/
1717
package org.apache.arrow.vector.complex.impl;
1818

19+
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
20+
import static org.junit.jupiter.api.Assertions.assertEquals;
1921
import static org.junit.jupiter.api.Assertions.assertThrows;
2022
import static org.junit.jupiter.api.Assertions.assertTrue;
2123

2224
import java.math.BigDecimal;
25+
import java.util.Arrays;
2326
import java.util.UUID;
2427
import org.apache.arrow.memory.BufferAllocator;
2528
import org.apache.arrow.memory.RootAllocator;
@@ -36,8 +39,10 @@
3639
import org.apache.arrow.vector.complex.writer.FieldWriter;
3740
import org.apache.arrow.vector.extension.UuidType;
3841
import org.apache.arrow.vector.holders.DecimalHolder;
42+
import org.apache.arrow.vector.types.TimeUnit;
3943
import org.apache.arrow.vector.types.Types;
4044
import org.apache.arrow.vector.types.pojo.ArrowType;
45+
import org.apache.arrow.vector.types.pojo.Field;
4146
import org.apache.arrow.vector.types.pojo.FieldType;
4247
import org.apache.arrow.vector.util.DecimalUtility;
4348
import org.junit.jupiter.api.AfterEach;
@@ -954,4 +959,125 @@ public void testCopyStructVectorWithExtensionType() {
954959
assertTrue(VectorEqualsVisitor.vectorEquals(from, to));
955960
}
956961
}
962+
963+
@Test
964+
public void testCopyListOfTimeStampNanoTZVector() {
965+
try (ListVector from = ListVector.empty("v", allocator);
966+
ListVector to = ListVector.empty("v", allocator)) {
967+
from.addOrGetVector(FieldType.nullable(new ArrowType.Timestamp(TimeUnit.NANOSECOND, "UTC")));
968+
to.addOrGetVector(FieldType.nullable(new ArrowType.Timestamp(TimeUnit.NANOSECOND, "UTC")));
969+
970+
UnionListWriter listWriter = from.getWriter();
971+
listWriter.allocate();
972+
for (int i = 0; i < COUNT; i++) {
973+
listWriter.setPosition(i);
974+
listWriter.startList();
975+
listWriter.timeStampNanoTZ().writeTimeStampNanoTZ(i * 1_000_000L);
976+
listWriter.timeStampNanoTZ().writeTimeStampNanoTZ(i * 2_000_000L);
977+
listWriter.endList();
978+
}
979+
from.setValueCount(COUNT);
980+
981+
// copy values
982+
FieldReader in = from.getReader();
983+
FieldWriter out = to.getWriter();
984+
for (int i = 0; i < COUNT; i++) {
985+
in.setPosition(i);
986+
out.setPosition(i);
987+
ComplexCopier.copy(in, out);
988+
}
989+
to.setValueCount(COUNT);
990+
991+
// validate equals
992+
assertTrue(VectorEqualsVisitor.vectorEquals(from, to));
993+
}
994+
}
995+
996+
@Test
997+
public void testCopyStructOfTimeStampNanoTZVector() {
998+
try (final StructVector from = StructVector.empty("v", allocator);
999+
final StructVector to = StructVector.empty("v", allocator)) {
1000+
from.allocateNewSafe();
1001+
1002+
NullableStructWriter structWriter = from.getWriter();
1003+
for (int i = 0; i < COUNT; i++) {
1004+
structWriter.setPosition(i);
1005+
structWriter.start();
1006+
structWriter.timeStampNanoTZ("ts", "UTC").writeTimeStampNanoTZ(i * 1_000_000L);
1007+
structWriter.end();
1008+
}
1009+
from.setValueCount(COUNT);
1010+
1011+
// copy values
1012+
FieldReader in = from.getReader();
1013+
FieldWriter out = to.getWriter();
1014+
for (int i = 0; i < COUNT; i++) {
1015+
in.setPosition(i);
1016+
out.setPosition(i);
1017+
ComplexCopier.copy(in, out);
1018+
}
1019+
to.setValueCount(COUNT);
1020+
1021+
// validate equals
1022+
assertTrue(VectorEqualsVisitor.vectorEquals(from, to));
1023+
}
1024+
}
1025+
1026+
@Test
1027+
public void testCopyMapOfTimeStampNanoTZVector() {
1028+
// Before the fix ComplexCopier.getMapWriterForReader threw
1029+
// UnsupportedOperationException for a TimeStamp*TZ map value. The map value
1030+
// child is pre-typed so the reader reports TIMESTAMPNANOTZ and the map case
1031+
// is actually exercised. vectorEquals is not asserted because the MapWriter
1032+
// scalar API has no timezone overload (unlike StructWriter), so the copied
1033+
// value child does not round-trip to the same timezone; this test locks in
1034+
// that copying the map no longer throws.
1035+
Field key = new Field("key", FieldType.notNullable(new ArrowType.Int(32, true)), null);
1036+
Field value =
1037+
new Field(
1038+
"value",
1039+
FieldType.nullable(new ArrowType.Timestamp(TimeUnit.NANOSECOND, "UTC")),
1040+
null);
1041+
Field entries =
1042+
new Field(
1043+
MapVector.DATA_VECTOR_NAME,
1044+
FieldType.notNullable(ArrowType.Struct.INSTANCE),
1045+
Arrays.asList(key, value));
1046+
Field mapField =
1047+
new Field("v", FieldType.nullable(new ArrowType.Map(false)), Arrays.asList(entries));
1048+
1049+
try (final MapVector from = (MapVector) mapField.createVector(allocator);
1050+
final MapVector to = MapVector.empty("v", allocator, false)) {
1051+
1052+
from.allocateNew();
1053+
1054+
UnionMapWriter mapWriter = from.getWriter();
1055+
for (int i = 0; i < COUNT; i++) {
1056+
mapWriter.setPosition(i);
1057+
mapWriter.startMap();
1058+
mapWriter.startEntry();
1059+
mapWriter.key().integer().writeInt(i);
1060+
mapWriter.value().timeStampNanoTZ().writeTimeStampNanoTZ(i * 1_000_000L);
1061+
mapWriter.endEntry();
1062+
mapWriter.endMap();
1063+
}
1064+
1065+
from.setValueCount(COUNT);
1066+
1067+
// copy values
1068+
FieldReader in = from.getReader();
1069+
FieldWriter out = to.getWriter();
1070+
assertDoesNotThrow(
1071+
() -> {
1072+
for (int i = 0; i < COUNT; i++) {
1073+
in.setPosition(i);
1074+
out.setPosition(i);
1075+
ComplexCopier.copy(in, out);
1076+
}
1077+
});
1078+
to.setValueCount(COUNT);
1079+
1080+
assertEquals(from.getValueCount(), to.getValueCount());
1081+
}
1082+
}
9571083
}

0 commit comments

Comments
 (0)