Support for reading Array records to ListVector from JDBC

Change-Id: Id6788671f5b30ed2f33f4317d00b5e6dc92ee889

java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcFieldInfo.java

@@ -0,0 +1,114 @@
+/*
+ * 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.adapter.jdbc;
+
+import java.sql.ResultSetMetaData;
+import java.sql.SQLException;
+import java.sql.Types;
+
+import com.google.common.base.Preconditions;
+
+/**
+ * This class represents the information about a JDBC ResultSet Field that is
+ * needed to construct an {@link org.apache.arrow.vector.types.pojo.ArrowType}.
+ * Currently, this is:
+ * <ul>
+ *   <li>The JDBC {@link java.sql.Types} type.</li>
+ *   <li>The field's precision (used for {@link java.sql.Types#DECIMAL} and {@link java.sql.Types#NUMERIC} types)</li>
+ *   <li>The field's scale (used for {@link java.sql.Types#DECIMAL} and {@link java.sql.Types#NUMERIC} types)</li>
+ * </ul>
+ */
+public class JdbcFieldInfo {
+  private final int jdbcType;
+  private final int precision;
+  private final int scale;
+
+  /**
+   * Builds a <code>JdbcFieldInfo</code> using only the {@link java.sql.Types} type.  Do not use this constructor
+   * if the field type is {@link java.sql.Types#DECIMAL} or {@link java.sql.Types#NUMERIC}; the precision and
+   * scale will be set to <code>0</code>.
+   *
+   * @param jdbcType The {@link java.sql.Types} type.
+   * @throws IllegalArgumentException if jdbcType is {@link java.sql.Types#DECIMAL} or {@link java.sql.Types#NUMERIC}.
+   */
+  public JdbcFieldInfo(int jdbcType) {
+    Preconditions.checkArgument(
+        (jdbcType != Types.DECIMAL && jdbcType != Types.NUMERIC),
+        "DECIMAL and NUMERIC types require a precision and scale; please use another constructor.");
+
+    this.jdbcType = jdbcType;
+    this.precision = 0;
+    this.scale = 0;
+  }
+
+  /**
+   * Builds a <code>JdbcFieldInfo</code> from the {@link java.sql.Types} type, precision, and scale.
+   * Use this constructor for {@link java.sql.Types#DECIMAL} and {@link java.sql.Types#NUMERIC} types.
+   *
+   * @param jdbcType The {@link java.sql.Types} type.
+   * @param precision The field's numeric precision.
+   * @param scale The field's numeric scale.
+   */
+  public JdbcFieldInfo(int jdbcType, int precision, int scale) {
+    this.jdbcType = jdbcType;
+    this.precision = precision;
+    this.scale = scale;
+  }
+
+  /**
+   * Builds a <code>JdbcFieldInfo</code> from the corresponding {@link java.sql.ResultSetMetaData} column.
+   *
+   * @param rsmd The {@link java.sql.ResultSetMetaData} to get the field information from.
+   * @param column The column to get the field information for (on a 1-based index).
+   * @throws SQLException If the column information cannot be retrieved.
+   * @throws NullPointerException if <code>rsmd</code> is <code>null</code>.
+   * @throws IllegalArgumentException if <code>column</code> is out of bounds.
+   */
+  public JdbcFieldInfo(ResultSetMetaData rsmd, int column) throws SQLException {
+    Preconditions.checkNotNull(rsmd, "ResultSetMetaData cannot be null.");
+    Preconditions.checkArgument(column > 0, "ResultSetMetaData columns have indices starting at 1.");
+    Preconditions.checkArgument(
+        column <= rsmd.getColumnCount(),
+        "The index must be within the number of columns (1 to %s, inclusive)", rsmd.getColumnCount());
+
+    this.jdbcType = rsmd.getColumnType(column);
+    this.precision = rsmd.getPrecision(column);
+    this.scale = rsmd.getScale(column);
+  }
+
+  /**
+   * The {@link java.sql.Types} type.
+   */
+  public int getJdbcType() {
+    return jdbcType;
+  }
+
+  /**
+   * The numeric precision, for {@link java.sql.Types#NUMERIC}  and {@link java.sql.Types#DECIMAL} types.
+   */
+  public int getPrecision() {
+    return precision;
+  }
+
+  /**
+   * The numeric scale, for {@link java.sql.Types#NUMERIC}  and {@link java.sql.Types#DECIMAL} types.
+   */
+  public int getScale() {
+    return scale;
+  }
+}

java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrow.java

@@ -88,7 +88,7 @@ public static VectorSchemaRoot sqlToArrow(Connection connection, String query, B
     Preconditions.checkNotNull(allocator, "Memory allocator object can not be null");
 
     JdbcToArrowConfig config =
-            new JdbcToArrowConfig(allocator, JdbcToArrowUtils.getUtcCalendar(), false);
+            new JdbcToArrowConfig(allocator, JdbcToArrowUtils.getUtcCalendar());
     return sqlToArrow(connection, query, config);
   }
 
@@ -116,7 +116,7 @@ public static VectorSchemaRoot sqlToArrow(
     Preconditions.checkNotNull(allocator, "Memory allocator object can not be null");
     Preconditions.checkNotNull(calendar, "Calendar object can not be null");
 
-    return sqlToArrow(connection, query, new JdbcToArrowConfig(allocator, calendar, false));
+    return sqlToArrow(connection, query, new JdbcToArrowConfig(allocator, calendar));
   }
 
   /**
@@ -170,7 +170,7 @@ public static VectorSchemaRoot sqlToArrow(ResultSet resultSet, BaseAllocator all
     Preconditions.checkNotNull(allocator, "Memory Allocator object can not be null");
 
     JdbcToArrowConfig config = 
-            new JdbcToArrowConfig(allocator, JdbcToArrowUtils.getUtcCalendar(), false);
+            new JdbcToArrowConfig(allocator, JdbcToArrowUtils.getUtcCalendar());
     return sqlToArrow(resultSet, config);
   }
 
@@ -184,8 +184,7 @@ public static VectorSchemaRoot sqlToArrow(ResultSet resultSet, BaseAllocator all
    */
   public static VectorSchemaRoot sqlToArrow(ResultSet resultSet, Calendar calendar) throws SQLException, IOException {
     Preconditions.checkNotNull(resultSet, "JDBC ResultSet object can not be null");
-
-    return sqlToArrow(resultSet, new JdbcToArrowConfig(new RootAllocator(Integer.MAX_VALUE), calendar, false));
+    return sqlToArrow(resultSet, new JdbcToArrowConfig(new RootAllocator(Integer.MAX_VALUE), calendar));
   }
 
   /**
@@ -205,7 +204,7 @@ public static VectorSchemaRoot sqlToArrow(
     Preconditions.checkNotNull(resultSet, "JDBC ResultSet object can not be null");
     Preconditions.checkNotNull(allocator, "Memory Allocator object can not be null");
 
-    return sqlToArrow(resultSet, new JdbcToArrowConfig(allocator, calendar, false));
+    return sqlToArrow(resultSet, new JdbcToArrowConfig(allocator, calendar));
   }
 
   /**

java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowConfig.java

@@ -18,6 +18,7 @@
 package org.apache.arrow.adapter.jdbc;
 
 import java.util.Calendar;
+import java.util.Map;
 
 import org.apache.arrow.memory.BaseAllocator;
 
@@ -28,16 +29,29 @@
  * <p>
  * The allocator is used to construct the {@link org.apache.arrow.vector.VectorSchemaRoot},
  * and the calendar is used to define the time zone of any {@link org.apahe.arrow.vector.pojo.ArrowType.Timestamp}
- * fields that are created during the conversion.
+ * fields that are created during the conversion.  Neither field may be <code>null</code>.
  * </p>
  * <p>
- * Neither field may be <code>null</code>.
+ * If the <code>includeMetadata</code> flag is set, the Arrow field metadata will contain information
+ * from the corresponding {@link java.sql.ResultSetMetaData} that was used to create the
+ * {@link org.apache.arrow.vector.types.pojo.FieldType} of the corresponding
+ * {@link org.apache.arrow.vector.FieldVector}.
+ * </p>
+ * <p>
+ * If there are any {@link java.sql.Types#ARRAY} fields in the {@link java.sql.ResultSet}, the corresponding
+ * {@link JdbcFieldInfo} for the array's contents must be defined here.  Unfortunately, the sub-type
+ * information cannot be retrieved from all JDBC implementations (H2 for example, returns
+ * {@link java.sql.Types#NULL} for the array sub-type), so it must be configured here.  The column index
+ * or name can be used to map to a {@link JdbcFieldInfo}, and that will be used for the conversion.
  * </p>
  */
 public final class JdbcToArrowConfig {
+
   private Calendar calendar;
   private BaseAllocator allocator;
   private boolean includeMetadata;
+  private Map<Integer, JdbcFieldInfo> arraySubTypesByColumnIndex;
+  private Map<String, JdbcFieldInfo> arraySubTypesByColumnName;
 
   /**
    * Constructs a new configuration from the provided allocator and calendar.  The <code>allocator</code>
@@ -46,20 +60,47 @@ public final class JdbcToArrowConfig {
    *
    * @param allocator       The memory allocator to construct the Arrow vectors with.
    * @param calendar        The calendar to use when constructing Timestamp fields and reading time-based results.
-   * @param includeMetadata Whether to include JDBC field metadata in the Arrow Schema Field metadata.
    */
-  JdbcToArrowConfig(BaseAllocator allocator, Calendar calendar, boolean includeMetadata) {
+  JdbcToArrowConfig(BaseAllocator allocator, Calendar calendar) {
     Preconditions.checkNotNull(allocator, "Memory allocator cannot be null");
 
     this.allocator = allocator;
     this.calendar = calendar;
+    this.includeMetadata = false;
+    this.arraySubTypesByColumnIndex = null;
+    this.arraySubTypesByColumnName = null;
+  }
+
+  /**
+   * Constructs a new configuration from the provided allocator and calendar.  The <code>allocator</code>
+   * is used when constructing the Arrow vectors from the ResultSet, and the calendar is used to define
+   * Arrow Timestamp fields, and to read time-based fields from the JDBC <code>ResultSet</code>. 
+   *
+   * @param allocator       The memory allocator to construct the Arrow vectors with.
+   * @param calendar        The calendar to use when constructing Timestamp fields and reading time-based results.
+   * @param includeMetadata Whether to include JDBC field metadata in the Arrow Schema Field metadata.
+   * @param arraySubTypesByColumnIndex The type of the JDBC array at the column index (1-based).
+   * @param arraySubTypesByColumnName  The type of the JDBC array at the column name.
+   */
+  JdbcToArrowConfig(
+      BaseAllocator allocator,
+      Calendar calendar,
+      boolean includeMetadata,
+      Map<Integer, JdbcFieldInfo> arraySubTypesByColumnIndex,
+      Map<String, JdbcFieldInfo> arraySubTypesByColumnName) {
+
+    this(allocator, calendar);
+
     this.includeMetadata = includeMetadata;
+    this.arraySubTypesByColumnIndex = arraySubTypesByColumnIndex;
+    this.arraySubTypesByColumnName = arraySubTypesByColumnName;
   }
 
   /**
    * The calendar to use when defining Arrow Timestamp fields
    * and retrieving {@link Date}, {@link Time}, or {@link Timestamp}
    * data types from the {@link ResultSet}, or <code>null</code> if not converting.
+   *
    * @return the calendar.
    */
   public Calendar getCalendar() {
@@ -82,4 +123,32 @@ public BaseAllocator getAllocator() {
   public boolean shouldIncludeMetadata() {
     return includeMetadata;
   }
+
+  /**
+   * Returns the array sub-type {@link JdbcFieldInfo} defined for the provided column index.
+   *
+   * @param index The {@link java.sql.ResultSetMetaData} column index of an {@link java.sql.Types#ARRAY} type.
+   * @return The {@link JdbcFieldInfo} for that array's sub-type, or <code>null</code> if not defined.
+   */
+  public JdbcFieldInfo getArraySubTypeByColumnIndex(int index) {
+    if (arraySubTypesByColumnIndex == null) {
+      return null;
+    } else {
+      return arraySubTypesByColumnIndex.get(index);
+    }
+  }
+
+  /**
+   * Returns the array sub-type {@link JdbcFieldInfo} defined for the provided column name.
+   *
+   * @param index The {@link java.sql.ResultSetMetaData} column name of an {@link java.sql.Types#ARRAY} type.
+   * @return The {@link JdbcFieldInfo} for that array's sub-type, or <code>null</code> if not defined.
+   */
+  public JdbcFieldInfo getArraySubTypeByColumnName(String name) {
+    if (arraySubTypesByColumnName == null) {
+      return null;
+    } else {
+      return arraySubTypesByColumnName.get(name);
+    }
+  }
 }

java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowConfigBuilder.java

@@ -18,6 +18,7 @@
 package org.apache.arrow.adapter.jdbc;
 
 import java.util.Calendar;
+import java.util.Map;
 
 import org.apache.arrow.memory.BaseAllocator;
 
@@ -30,6 +31,8 @@ public class JdbcToArrowConfigBuilder {
   private Calendar calendar;
   private BaseAllocator allocator;
   private boolean includeMetadata;
+  private Map<Integer, JdbcFieldInfo> arraySubTypesByColumnIndex;
+  private Map<String, JdbcFieldInfo> arraySubTypesByColumnName;
 
   /**
    * Default constructor for the <code>JdbcToArrowConfigBuilder}</code>.
@@ -40,6 +43,8 @@ public JdbcToArrowConfigBuilder() {
     this.allocator = null;
     this.calendar = null;
     this.includeMetadata = false;
+    this.arraySubTypesByColumnIndex = null;
+    this.arraySubTypesByColumnName = null;
   }
 
   /**
@@ -126,6 +131,29 @@ public JdbcToArrowConfigBuilder setIncludeMetadata(boolean includeMetadata) {
     return this;
   }
 
+  /**
+   * Sets the mapping of column-index-to-{@link JdbcFieldInfo} used for columns of type {@link java.sql.Types#ARRAY}.
+   * The column index is 1-based, to match the JDBC column index.
+   *
+   * @param map The mapping.
+   * @return This instance of the <code>JdbcToArrowConfig</code>, for chaining.
+   */
+  public JdbcToArrowConfigBuilder setArraySubTypeByColumnIndexMap(Map<Integer, JdbcFieldInfo> map) {
+    this.arraySubTypesByColumnIndex = map;
+    return this;
+  }
+
+  /**
+   * Sets the mapping of column-name-to-{@link JdbcFieldInfo} used for columns of type {@link java.sql.Types#ARRAY}.
+   *
+   * @param map The mapping.
+   * @return This instance of the <code>JdbcToArrowConfig</code>, for chaining.
+   */
+  public JdbcToArrowConfigBuilder setArraySubTypeByColumnNameMap(Map<String, JdbcFieldInfo> map) {
+    this.arraySubTypesByColumnName = map;
+    return this;
+  }
+
   /**
    * This builds the {@link JdbcToArrowConfig} from the provided
    * {@link BaseAllocator} and {@link Calendar}.
@@ -134,6 +162,11 @@ public JdbcToArrowConfigBuilder setIncludeMetadata(boolean includeMetadata) {
    * @throws NullPointerException if either the allocator or calendar was not set.
    */
   public JdbcToArrowConfig build() {
-    return new JdbcToArrowConfig(allocator, calendar, includeMetadata);
+    return new JdbcToArrowConfig(
+        allocator,
+        calendar,
+        includeMetadata,
+        arraySubTypesByColumnIndex,
+        arraySubTypesByColumnName);
   }
 }