apacheGH-40893: [Java][FlightRPC] Support IntervalMonthDayNanoVector …

…in FlightSQL JDBC Driver (apache#40894) ### Rationale for this change Fixes apache#40893. ### What changes are included in this PR? - Support IntervalMonthDayNanoVector in FlightSQL JDBC Driver - Return PeriodDuration as JDBC Object type, because there is no good java.time type for this interval - Return an ISO-8601 interval as the stringified version of PeriodDuration - Make PeriodDuration implement TemporalAccessor for standardization ### Are these changes tested? Unit tests have been added that match those for other interval types. I'm unaware of any other types of tests worth adding to, but I'd be happy to if pointed there. ### Are there any user-facing changes? The only change users should noticed is that the FlightSQL JDBC Driver can now handle more query responses. * GitHub Issue: apache#40893 Authored-by: paul <[email protected]> Signed-off-by: David Li <[email protected]>
R-JunmingChen · Mar 31, 2024 · 17a5368 · 17a5368
1 parent 9f0101e
commit 17a5368
Show file tree

Hide file tree

Showing 6 changed files with 217 additions and 4 deletions.
diff --git a/...e/src/main/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactory.java b/...e/src/main/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactory.java
@@ -51,6 +51,7 @@
 import org.apache.arrow.vector.Float8Vector;
 import org.apache.arrow.vector.IntVector;
 import org.apache.arrow.vector.IntervalDayVector;
+import org.apache.arrow.vector.IntervalMonthDayNanoVector;
 import org.apache.arrow.vector.IntervalYearVector;
 import org.apache.arrow.vector.LargeVarBinaryVector;
 import org.apache.arrow.vector.LargeVarCharVector;
@@ -176,6 +177,9 @@ public static ArrowFlightJdbcAccessor createAccessor(ValueVector vector,
     } else if (vector instanceof IntervalYearVector) {
       return new ArrowFlightJdbcIntervalVectorAccessor(((IntervalYearVector) vector), getCurrentRow,
           setCursorWasNull);
+    } else if (vector instanceof IntervalMonthDayNanoVector) {
+      return new ArrowFlightJdbcIntervalVectorAccessor(((IntervalMonthDayNanoVector) vector), getCurrentRow,
+          setCursorWasNull);
     } else if (vector instanceof StructVector) {
       return new ArrowFlightJdbcStructVectorAccessor((StructVector) vector, getCurrentRow,
           setCursorWasNull);

diff --git a/...pache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcIntervalVectorAccessor.java b/...pache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcIntervalVectorAccessor.java
@@ -30,8 +30,11 @@
 import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory;
 import org.apache.arrow.vector.BaseFixedWidthVector;
 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.holders.NullableIntervalDayHolder;
+import org.apache.arrow.vector.holders.NullableIntervalMonthDayNanoHolder;
 import org.apache.arrow.vector.holders.NullableIntervalYearHolder;
 
 /**
@@ -96,6 +99,35 @@ public ArrowFlightJdbcIntervalVectorAccessor(IntervalYearVector vector,
     objectClass = java.time.Period.class;
   }
 
+  /**
+   * Instantiate an accessor for a {@link IntervalMonthDayNanoVector}.
+   *
+   * @param vector             an instance of a IntervalMonthDayNanoVector.
+   * @param currentRowSupplier the supplier to track the rows.
+   * @param setCursorWasNull   the consumer to set if value was null.
+   */
+  public ArrowFlightJdbcIntervalVectorAccessor(IntervalMonthDayNanoVector vector,
+                                               IntSupplier currentRowSupplier,
+                                               ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) {
+    super(currentRowSupplier, setCursorWasNull);
+    this.vector = vector;
+    stringGetter = (index) -> {
+      final NullableIntervalMonthDayNanoHolder holder = new NullableIntervalMonthDayNanoHolder();
+      vector.get(index, holder);
+      if (holder.isSet == 0) {
+        return null;
+      } else {
+        final int months = holder.months;
+        final int days = holder.days;
+        final long nanos = holder.nanoseconds;
+        final Period period = Period.ofMonths(months).plusDays(days);
+        final Duration duration = Duration.ofNanos(nanos);
+        return new PeriodDuration(period, duration).toISO8601IntervalString();
+      }
+    };
+    objectClass = PeriodDuration.class;
+  }
+
   @Override
   public Class<?> getObjectClass() {
     return objectClass;

diff --git a/...c/test/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactoryTest.java b/...c/test/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactoryTest.java
@@ -41,6 +41,7 @@
 import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestRule;
 import org.apache.arrow.vector.DurationVector;
 import org.apache.arrow.vector.IntervalDayVector;
+import org.apache.arrow.vector.IntervalMonthDayNanoVector;
 import org.apache.arrow.vector.IntervalYearVector;
 import org.apache.arrow.vector.LargeVarCharVector;
 import org.apache.arrow.vector.ValueVector;
@@ -405,6 +406,19 @@ public void createAccessorForIntervalYearVector() {
     }
   }
 
+  @Test
+  public void createAccessorForIntervalMonthDayNanoVector() {
+    try (ValueVector valueVector = new IntervalMonthDayNanoVector("",
+        rootAllocatorTestRule.getRootAllocator())) {
+      ArrowFlightJdbcAccessor accessor =
+          ArrowFlightJdbcAccessorFactory.createAccessor(valueVector, GET_CURRENT_ROW,
+              (boolean wasNull) -> {
+              });
+
+      Assert.assertTrue(accessor instanceof ArrowFlightJdbcIntervalVectorAccessor);
+    }
+  }
+
   @Test
   public void createAccessorForUnionVector() {
     try (ValueVector valueVector = new UnionVector("", rootAllocatorTestRule.getRootAllocator(),

diff --git a/...e/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcIntervalVectorAccessorTest.java b/...e/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcIntervalVectorAccessorTest.java
@@ -24,6 +24,7 @@
 
 import java.time.Duration;
 import java.time.Period;
+import java.time.format.DateTimeParseException;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.function.Supplier;
@@ -32,7 +33,9 @@
 import org.apache.arrow.driver.jdbc.utils.AccessorTestUtils;
 import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestRule;
 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.After;
 import org.junit.Assert;
@@ -66,6 +69,9 @@ public class ArrowFlightJdbcIntervalVectorAccessorTest {
         } 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;
       };
@@ -98,6 +104,17 @@ public static Collection<Object[]> data() {
           }
           return vector;
         }, "IntervalYearVector"},
+        {(Supplier<ValueVector>) () -> {
+          IntervalMonthDayNanoVector vector =
+                  new IntervalMonthDayNanoVector("", rootAllocatorTestRule.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"},
     });
   }
 
@@ -137,13 +154,31 @@ public void testShouldGetObjectReturnNull() throws Exception {
   }
 
   private String getStringOnVector(ValueVector vector, int index) {
-    String object = getExpectedObject(vector, index).toString();
+    Object object = getExpectedObject(vector, index);
     if (object == null) {
       return null;
     } else if (vector instanceof IntervalDayVector) {
-      return formatIntervalDay(Duration.parse(object));
+      return formatIntervalDay(Duration.parse(object.toString()));
     } else if (vector instanceof IntervalYearVector) {
-      return formatIntervalYear(Period.parse(object));
+      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;
   }
@@ -225,6 +260,8 @@ private Class<?> getExpectedObjectClassForVector(ValueVector vector) {
       return Duration.class;
     } else if (vector instanceof IntervalYearVector) {
       return Period.class;
+    } else if (vector instanceof IntervalMonthDayNanoVector) {
+      return PeriodDuration.class;
     }
     return null;
   }
@@ -239,6 +276,10 @@ private void setAllNullOnVector(ValueVector vector) {
       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);
+      }
     }
   }
 
@@ -247,6 +288,10 @@ private Object getExpectedObject(ValueVector vector, int currentRow) {
       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;
   }

diff --git a/java/vector/src/main/java/org/apache/arrow/vector/PeriodDuration.java b/java/vector/src/main/java/org/apache/arrow/vector/PeriodDuration.java
@@ -17,16 +17,33 @@
 
 package org.apache.arrow.vector;
 
+import static java.time.temporal.ChronoUnit.DAYS;
+import static java.time.temporal.ChronoUnit.MONTHS;
+import static java.time.temporal.ChronoUnit.NANOS;
+import static java.time.temporal.ChronoUnit.SECONDS;
+import static java.time.temporal.ChronoUnit.YEARS;
+
 import java.time.Duration;
 import java.time.Period;
+import java.time.temporal.ChronoUnit;
+import java.time.temporal.Temporal;
+import java.time.temporal.TemporalAmount;
+import java.time.temporal.TemporalUnit;
+import java.time.temporal.UnsupportedTemporalTypeException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
 
 import org.apache.arrow.util.Preconditions;
 
 /**
  * Combination of Period and Duration for representing this interval type
  * as a POJO.
  */
-public class PeriodDuration {
+public class PeriodDuration implements TemporalAmount {
+
+  private static final List<TemporalUnit> SUPPORTED_UNITS =
+          Collections.unmodifiableList(Arrays.<TemporalUnit>asList(YEARS, MONTHS, DAYS, SECONDS, NANOS));
   private final Period period;
   private final Duration duration;
 
@@ -43,6 +60,60 @@ public Duration getDuration() {
     return duration;
   }
 
+  @Override
+  public long get(TemporalUnit unit) {
+    if (unit instanceof ChronoUnit) {
+      switch ((ChronoUnit) unit) {
+        case YEARS:
+          return period.getYears();
+        case MONTHS:
+          return period.getMonths();
+        case DAYS:
+          return period.getDays();
+        case SECONDS:
+          return duration.getSeconds();
+        case NANOS:
+          return duration.getNano();
+        default:
+          break;
+      }
+    }
+    throw new UnsupportedTemporalTypeException("Unsupported TemporalUnit: " + unit);
+  }
+
+  @Override
+  public List<TemporalUnit> getUnits() {
+    return SUPPORTED_UNITS;
+  }
+
+  @Override
+  public Temporal addTo(Temporal temporal) {
+    return temporal.plus(period).plus(duration);
+  }
+
+  @Override
+  public Temporal subtractFrom(Temporal temporal) {
+    return temporal.minus(period).minus(duration);
+  }
+
+  /**
+   * Format this PeriodDuration as an ISO-8601 interval.
+   *
+   * @return An ISO-8601 formatted string representing the interval.
+   */
+  public String toISO8601IntervalString() {
+    if (duration.isZero()) {
+      return period.toString();
+    }
+    String durationString = duration.toString();
+    if (period.isZero()) {
+      return durationString;
+    }
+
+    // Remove 'P' from duration string and concatenate to produce an ISO-8601 representation
+    return period + durationString.substring(1);
+  }
+
   @Override
   public String toString() {
     return period.toString() + " " + duration.toString();

diff --git a/java/vector/src/test/java/org/apache/arrow/vector/TestPeriodDuration.java b/java/vector/src/test/java/org/apache/arrow/vector/TestPeriodDuration.java
@@ -21,7 +21,10 @@
 import static org.junit.Assert.assertNotEquals;
 
 import java.time.Duration;
+import java.time.LocalDate;
+import java.time.LocalDateTime;
 import java.time.Period;
+import java.time.temporal.ChronoUnit;
 
 import org.junit.Test;
 
@@ -43,4 +46,48 @@ public void testBasics() {
     assertNotEquals(pd1.hashCode(), pd3.hashCode());
   }
 
+  @Test
+  public void testToISO8601IntervalString() {
+    assertEquals("P0D",
+            new PeriodDuration(Period.ZERO, Duration.ZERO).toISO8601IntervalString());
+    assertEquals("P1Y2M3D",
+            new PeriodDuration(Period.of(1, 2, 3), Duration.ZERO).toISO8601IntervalString());
+    assertEquals("PT0.000000123S",
+            new PeriodDuration(Period.ZERO, Duration.ofNanos(123)).toISO8601IntervalString());
+    assertEquals("PT1.000000123S",
+            new PeriodDuration(Period.ZERO, Duration.ofSeconds(1).withNanos(123)).toISO8601IntervalString());
+    assertEquals("PT1H1.000000123S",
+            new PeriodDuration(Period.ZERO, Duration.ofSeconds(3601).withNanos(123)).toISO8601IntervalString());
+    assertEquals("PT24H1M1.000000123S",
+            new PeriodDuration(Period.ZERO, Duration.ofSeconds(86461).withNanos(123)).toISO8601IntervalString());
+    assertEquals("P1Y2M3DT24H1M1.000000123S",
+            new PeriodDuration(Period.of(1, 2, 3), Duration.ofSeconds(86461).withNanos(123)).toISO8601IntervalString());
+
+    assertEquals("P-1Y-2M-3D",
+            new PeriodDuration(Period.of(-1, -2, -3), Duration.ZERO).toISO8601IntervalString());
+    assertEquals("PT-0.000000123S",
+            new PeriodDuration(Period.ZERO, Duration.ofNanos(-123)).toISO8601IntervalString());
+    assertEquals("PT-24H-1M-0.999999877S",
+            new PeriodDuration(Period.ZERO, Duration.ofSeconds(-86461).withNanos(123)).toISO8601IntervalString());
+    assertEquals("P-1Y-2M-3DT-0.999999877S",
+            new PeriodDuration(Period.of(-1, -2, -3), Duration.ofSeconds(-1).withNanos(123)).toISO8601IntervalString());
+  }
+
+  @Test
+  public void testTemporalAccessor() {
+    LocalDate date = LocalDate.of(2024, 1, 2);
+    PeriodDuration pd1 = new PeriodDuration(Period.ofYears(1), Duration.ZERO);
+    assertEquals(LocalDate.of(2025, 1, 2), pd1.addTo(date));
+
+    LocalDateTime dateTime = LocalDateTime.of(2024, 1, 2, 3, 4);
+    PeriodDuration pd2 = new PeriodDuration(Period.ZERO, Duration.ofMinutes(1));
+    assertEquals(LocalDateTime.of(2024, 1, 2, 3, 3), pd2.subtractFrom(dateTime));
+
+    PeriodDuration pd3 = new PeriodDuration(Period.of(1, 2, 3), Duration.ofSeconds(86461).withNanos(123));
+    assertEquals(pd3.get(ChronoUnit.YEARS), 1);
+    assertEquals(pd3.get(ChronoUnit.MONTHS), 2);
+    assertEquals(pd3.get(ChronoUnit.DAYS), 3);
+    assertEquals(pd3.get(ChronoUnit.SECONDS), 86461);
+    assertEquals(pd3.get(ChronoUnit.NANOS), 123);
+  }
 }