ArrowFlightJdbcIntervalVectorAccessorTest.java
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.arrow.driver.jdbc.accessor.impl.calendar;
import static org.apache.arrow.driver.jdbc.utils.IntervalStringUtils.formatIntervalDay;
import static org.apache.arrow.driver.jdbc.utils.IntervalStringUtils.formatIntervalYear;
import static org.hamcrest.CoreMatchers.equalTo;
import static org.hamcrest.CoreMatchers.is;
import static org.junit.jupiter.api.Assertions.assertEquals;
import java.time.Duration;
import java.time.Period;
import java.time.format.DateTimeParseException;
import java.util.function.Supplier;
import java.util.stream.Stream;
import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory;
import org.apache.arrow.driver.jdbc.utils.AccessorTestUtils;
import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension;
import org.apache.arrow.vector.IntervalDayVector;
import org.apache.arrow.vector.IntervalMonthDayNanoVector;
import org.apache.arrow.vector.IntervalYearVector;
import org.apache.arrow.vector.PeriodDuration;
import org.apache.arrow.vector.ValueVector;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.extension.RegisterExtension;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;
public class ArrowFlightJdbcIntervalVectorAccessorTest {
@RegisterExtension
public static RootAllocatorTestExtension rootAllocatorTestExtension =
new RootAllocatorTestExtension();
private ValueVector vector;
private final AccessorTestUtils.AccessorSupplier<ArrowFlightJdbcIntervalVectorAccessor>
accessorSupplier =
(vector, getCurrentRow) -> {
ArrowFlightJdbcAccessorFactory.WasNullConsumer noOpWasNullConsumer =
(boolean wasNull) -> {};
if (vector instanceof IntervalDayVector) {
return new ArrowFlightJdbcIntervalVectorAccessor(
(IntervalDayVector) vector, getCurrentRow, noOpWasNullConsumer);
} else if (vector instanceof IntervalYearVector) {
return new ArrowFlightJdbcIntervalVectorAccessor(
(IntervalYearVector) vector, getCurrentRow, noOpWasNullConsumer);
} else if (vector instanceof IntervalMonthDayNanoVector) {
return new ArrowFlightJdbcIntervalVectorAccessor(
(IntervalMonthDayNanoVector) vector, getCurrentRow, noOpWasNullConsumer);
}
return null;
};
final AccessorTestUtils.AccessorIterator<ArrowFlightJdbcIntervalVectorAccessor> accessorIterator =
new AccessorTestUtils.AccessorIterator<>(accessorSupplier);
public static Stream<Arguments> data() {
return Stream.of(
Arguments.of(
(Supplier<ValueVector>)
() -> {
IntervalDayVector vector =
new IntervalDayVector("", rootAllocatorTestExtension.getRootAllocator());
int valueCount = 10;
vector.setValueCount(valueCount);
for (int i = 0; i < valueCount; i++) {
vector.set(i, i + 1, (i + 1) * 1000);
}
return vector;
},
"IntervalDayVector"),
Arguments.of(
(Supplier<ValueVector>)
() -> {
IntervalYearVector vector =
new IntervalYearVector("", rootAllocatorTestExtension.getRootAllocator());
int valueCount = 10;
vector.setValueCount(valueCount);
for (int i = 0; i < valueCount; i++) {
vector.set(i, i + 1);
}
return vector;
},
"IntervalYearVector"),
Arguments.of(
(Supplier<ValueVector>)
() -> {
IntervalMonthDayNanoVector vector =
new IntervalMonthDayNanoVector(
"", rootAllocatorTestExtension.getRootAllocator());
int valueCount = 10;
vector.setValueCount(valueCount);
for (int i = 0; i < valueCount; i++) {
vector.set(i, i + 1, (i + 1) * 10, (i + 1) * 100);
}
return vector;
},
"IntervalMonthDayNanoVector"));
}
public void setup(Supplier<ValueVector> vectorSupplier) {
this.vector = vectorSupplier.get();
}
@AfterEach
public void tearDown() {
if (this.vector != null) {
this.vector.close();
}
}
@ParameterizedTest
@MethodSource("data")
public void testShouldGetObjectReturnValidObject(Supplier<ValueVector> vectorSupplier)
throws Exception {
setup(vectorSupplier);
accessorIterator.assertAccessorGetter(
vector,
ArrowFlightJdbcIntervalVectorAccessor::getObject,
(accessor, currentRow) -> is(getExpectedObject(vector, currentRow)));
}
@ParameterizedTest
@MethodSource("data")
public void testShouldGetObjectPassingObjectClassAsParameterReturnValidObject(
Supplier<ValueVector> vectorSupplier) throws Exception {
setup(vectorSupplier);
Class<?> objectClass = getExpectedObjectClassForVector(vector);
accessorIterator.assertAccessorGetter(
vector,
accessor -> accessor.getObject(objectClass),
(accessor, currentRow) -> is(getExpectedObject(vector, currentRow)));
}
@ParameterizedTest
@MethodSource("data")
public void testShouldGetObjectReturnNull(Supplier<ValueVector> vectorSupplier) throws Exception {
setup(vectorSupplier);
setAllNullOnVector(vector);
accessorIterator.assertAccessorGetter(
vector,
ArrowFlightJdbcIntervalVectorAccessor::getObject,
(accessor, currentRow) -> equalTo(null));
}
private String getStringOnVector(ValueVector vector, int index) {
Object object = getExpectedObject(vector, index);
if (object == null) {
return null;
} else if (vector instanceof IntervalDayVector) {
return formatIntervalDay(Duration.parse(object.toString()));
} else if (vector instanceof IntervalYearVector) {
return formatIntervalYear(Period.parse(object.toString()));
} else if (vector instanceof IntervalMonthDayNanoVector) {
String iso8601IntervalString = ((PeriodDuration) object).toISO8601IntervalString();
String[] periodAndDuration = iso8601IntervalString.split("T");
if (periodAndDuration.length == 1) {
// If there is no 'T', then either Period or Duration is zero, and the other one will
// successfully parse it
String periodOrDuration = periodAndDuration[0];
try {
return new PeriodDuration(Period.parse(periodOrDuration), Duration.ZERO)
.toISO8601IntervalString();
} catch (DateTimeParseException e) {
return new PeriodDuration(Period.ZERO, Duration.parse(periodOrDuration))
.toISO8601IntervalString();
}
} else {
// If there is a 'T', both Period and Duration are non-zero, and we just need to prepend the
// 'PT' to the
// duration for both to parse successfully
Period parse = Period.parse(periodAndDuration[0]);
Duration duration = Duration.parse("PT" + periodAndDuration[1]);
return new PeriodDuration(parse, duration).toISO8601IntervalString();
}
}
return null;
}
@ParameterizedTest
@MethodSource("data")
public void testShouldGetIntervalYear(Supplier<ValueVector> vectorSupplier) {
setup(vectorSupplier);
assertEquals("-002-00", formatIntervalYear(Period.parse("P-2Y")));
assertEquals("-001-01", formatIntervalYear(Period.parse("P-1Y-1M")));
assertEquals("-001-02", formatIntervalYear(Period.parse("P-1Y-2M")));
assertEquals("-002-03", formatIntervalYear(Period.parse("P-2Y-3M")));
assertEquals("-002-04", formatIntervalYear(Period.parse("P-2Y-4M")));
assertEquals("-011-01", formatIntervalYear(Period.parse("P-11Y-1M")));
assertEquals("+002-00", formatIntervalYear(Period.parse("P+2Y")));
assertEquals("+001-01", formatIntervalYear(Period.parse("P+1Y1M")));
assertEquals("+001-02", formatIntervalYear(Period.parse("P+1Y2M")));
assertEquals("+002-03", formatIntervalYear(Period.parse("P+2Y3M")));
assertEquals("+002-04", formatIntervalYear(Period.parse("P+2Y4M")));
assertEquals("+011-01", formatIntervalYear(Period.parse("P+11Y1M")));
}
@ParameterizedTest
@MethodSource("data")
public void testShouldGetIntervalDay(Supplier<ValueVector> vectorSupplier) {
setup(vectorSupplier);
assertEquals("-001 00:00:00.000", formatIntervalDay(Duration.parse("PT-24H")));
assertEquals("+001 00:00:00.000", formatIntervalDay(Duration.parse("PT+24H")));
assertEquals("-000 01:00:00.000", formatIntervalDay(Duration.parse("PT-1H")));
// "JDK-8054978: java.time.Duration.parse() fails for negative duration with 0 seconds and
// nanos" not fixed on JDK8
// assertEquals("-000 01:00:00.001", formatIntervalDay(Duration.parse("PT-1H-0M-00.001S")));
assertEquals("-000 01:00:00.001", formatIntervalDay(Duration.ofHours(-1).minusMillis(1)));
assertEquals("-000 01:01:01.000", formatIntervalDay(Duration.parse("PT-1H-1M-1S")));
assertEquals("-000 02:02:02.002", formatIntervalDay(Duration.parse("PT-2H-2M-02.002S")));
assertEquals("-000 23:59:59.999", formatIntervalDay(Duration.parse("PT-23H-59M-59.999S")));
// "JDK-8054978: java.time.Duration.parse() fails for negative duration with 0 seconds and
// nanos" not fixed on JDK8
// assertEquals("-000 11:59:00.100", formatIntervalDay(Duration.parse("PT-11H-59M-00.100S")));
assertEquals(
"-000 11:59:00.100",
formatIntervalDay(Duration.ofHours(-11).minusMinutes(59).minusMillis(100)));
assertEquals("-000 05:02:03.000", formatIntervalDay(Duration.parse("PT-5H-2M-3S")));
assertEquals("-000 22:22:22.222", formatIntervalDay(Duration.parse("PT-22H-22M-22.222S")));
assertEquals("+000 01:00:00.000", formatIntervalDay(Duration.parse("PT+1H")));
assertEquals("+000 01:00:00.001", formatIntervalDay(Duration.parse("PT+1H0M00.001S")));
assertEquals("+000 01:01:01.000", formatIntervalDay(Duration.parse("PT+1H1M1S")));
assertEquals("+000 02:02:02.002", formatIntervalDay(Duration.parse("PT+2H2M02.002S")));
assertEquals("+000 23:59:59.999", formatIntervalDay(Duration.parse("PT+23H59M59.999S")));
assertEquals("+000 11:59:00.100", formatIntervalDay(Duration.parse("PT+11H59M00.100S")));
assertEquals("+000 05:02:03.000", formatIntervalDay(Duration.parse("PT+5H2M3S")));
assertEquals("+000 22:22:22.222", formatIntervalDay(Duration.parse("PT+22H22M22.222S")));
}
@ParameterizedTest
@MethodSource("data")
public void testIntervalDayWithJodaPeriodObject(Supplier<ValueVector> vectorSupplier) {
setup(vectorSupplier);
assertEquals("+1567 00:00:00.000", formatIntervalDay(Duration.ofDays(1567)));
assertEquals("-1567 00:00:00.000", formatIntervalDay(Duration.ofDays(-1567)));
}
@ParameterizedTest
@MethodSource("data")
public void testShouldGetStringReturnCorrectString(Supplier<ValueVector> vectorSupplier)
throws Exception {
setup(vectorSupplier);
accessorIterator.assertAccessorGetter(
vector,
ArrowFlightJdbcIntervalVectorAccessor::getString,
(accessor, currentRow) -> is(getStringOnVector(vector, currentRow)));
}
@ParameterizedTest
@MethodSource("data")
public void testShouldGetStringReturnNull(Supplier<ValueVector> vectorSupplier) throws Exception {
setup(vectorSupplier);
setAllNullOnVector(vector);
accessorIterator.assertAccessorGetter(
vector,
ArrowFlightJdbcIntervalVectorAccessor::getString,
(accessor, currentRow) -> equalTo(null));
}
@ParameterizedTest
@MethodSource("data")
public void testShouldGetObjectClassReturnCorrectClass(Supplier<ValueVector> vectorSupplier)
throws Exception {
setup(vectorSupplier);
Class<?> expectedObjectClass = getExpectedObjectClassForVector(vector);
accessorIterator.assertAccessorGetter(
vector,
ArrowFlightJdbcIntervalVectorAccessor::getObjectClass,
(accessor, currentRow) -> equalTo(expectedObjectClass));
}
private Class<?> getExpectedObjectClassForVector(ValueVector vector) {
if (vector instanceof IntervalDayVector) {
return Duration.class;
} else if (vector instanceof IntervalYearVector) {
return Period.class;
} else if (vector instanceof IntervalMonthDayNanoVector) {
return PeriodDuration.class;
}
return null;
}
private void setAllNullOnVector(ValueVector vector) {
int valueCount = vector.getValueCount();
if (vector instanceof IntervalDayVector) {
for (int i = 0; i < valueCount; i++) {
((IntervalDayVector) vector).setNull(i);
}
} else if (vector instanceof IntervalYearVector) {
for (int i = 0; i < valueCount; i++) {
((IntervalYearVector) vector).setNull(i);
}
} else if (vector instanceof IntervalMonthDayNanoVector) {
for (int i = 0; i < valueCount; i++) {
((IntervalMonthDayNanoVector) vector).setNull(i);
}
}
}
private Object getExpectedObject(ValueVector vector, int currentRow) {
if (vector instanceof IntervalDayVector) {
return Duration.ofDays(currentRow + 1).plusMillis((currentRow + 1) * 1000L);
} else if (vector instanceof IntervalYearVector) {
return Period.ofMonths(currentRow + 1);
} else if (vector instanceof IntervalMonthDayNanoVector) {
Period period = Period.ofMonths(currentRow + 1).plusDays((currentRow + 1) * 10L);
Duration duration = Duration.ofNanos((currentRow + 1) * 100L);
return new PeriodDuration(period, duration);
}
return null;
}
}