diff --git a/java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/ArrowVectorIterator.java b/java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/ArrowVectorIterator.java index 6c3e9cf43e758..6e789009dd20a 100644 --- a/java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/ArrowVectorIterator.java +++ b/java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/ArrowVectorIterator.java @@ -139,7 +139,7 @@ private void initialize(VectorSchemaRoot root) throws SQLException { for (int i = 1; i <= consumers.length; i++) { final JdbcFieldInfo columnFieldInfo = JdbcToArrowUtils.getJdbcFieldInfoForColumn(rsmd, i, config); ArrowType arrowType = config.getJdbcToArrowTypeConverter().apply(columnFieldInfo); - consumers[i - 1] = JdbcToArrowUtils.getConsumer( + consumers[i - 1] = config.getJdbcConsumerGetter().apply( arrowType, i, isColumnNullable(resultSet.getMetaData(), i, columnFieldInfo), root.getVector(i - 1), config); } } diff --git a/java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowConfig.java b/java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowConfig.java index 012cd95c0b2b6..e23bad54afc14 100644 --- a/java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowConfig.java +++ b/java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowConfig.java @@ -22,8 +22,10 @@ import java.util.Map; import java.util.function.Function; +import org.apache.arrow.adapter.jdbc.consumer.JdbcConsumer; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Preconditions; +import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.types.pojo.ArrowType; /** @@ -76,6 +78,7 @@ public final class JdbcToArrowConfig { private final int targetBatchSize; private final Function jdbcToArrowTypeConverter; + private final JdbcConsumerFactory jdbcConsumerGetter; /** * Constructs a new configuration from the provided allocator and calendar. The allocator @@ -195,6 +198,38 @@ public final class JdbcToArrowConfig { Map schemaMetadata, Map> columnMetadataByColumnIndex, RoundingMode bigDecimalRoundingMode) { + this( + allocator, + calendar, + includeMetadata, + reuseVectorSchemaRoot, + arraySubTypesByColumnIndex, + arraySubTypesByColumnName, + targetBatchSize, + jdbcToArrowTypeConverter, + null, + explicitTypesByColumnIndex, + explicitTypesByColumnName, + schemaMetadata, + columnMetadataByColumnIndex, + bigDecimalRoundingMode); + } + + JdbcToArrowConfig( + BufferAllocator allocator, + Calendar calendar, + boolean includeMetadata, + boolean reuseVectorSchemaRoot, + Map arraySubTypesByColumnIndex, + Map arraySubTypesByColumnName, + int targetBatchSize, + Function jdbcToArrowTypeConverter, + JdbcConsumerFactory jdbcConsumerGetter, + Map explicitTypesByColumnIndex, + Map explicitTypesByColumnName, + Map schemaMetadata, + Map> columnMetadataByColumnIndex, + RoundingMode bigDecimalRoundingMode) { Preconditions.checkNotNull(allocator, "Memory allocator cannot be null"); this.allocator = allocator; this.calendar = calendar; @@ -212,6 +247,8 @@ public final class JdbcToArrowConfig { // set up type converter this.jdbcToArrowTypeConverter = jdbcToArrowTypeConverter != null ? jdbcToArrowTypeConverter : (jdbcFieldInfo) -> JdbcToArrowUtils.getArrowTypeFromJdbcType(jdbcFieldInfo, calendar); + + this.jdbcConsumerGetter = jdbcConsumerGetter != null ? jdbcConsumerGetter : JdbcToArrowUtils::getConsumer; } /** @@ -264,6 +301,13 @@ public Function getJdbcToArrowTypeConverter() { return jdbcToArrowTypeConverter; } + /** + * Gets the JDBC consumer getter. + */ + public JdbcConsumerFactory getJdbcConsumerGetter() { + return jdbcConsumerGetter; + } + /** * Returns the array sub-type {@link JdbcFieldInfo} defined for the provided column index. * @@ -338,4 +382,13 @@ public Map> getColumnMetadataByColumnIndex() { public RoundingMode getBigDecimalRoundingMode() { return bigDecimalRoundingMode; } + + /** + * Interface for a function that gets a JDBC consumer for the given values. + */ + @FunctionalInterface + public interface JdbcConsumerFactory { + JdbcConsumer apply(ArrowType arrowType, int columnIndex, boolean nullable, FieldVector vector, + JdbcToArrowConfig config); + } } diff --git a/java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowConfigBuilder.java b/java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowConfigBuilder.java index 2fe0492deb71f..7d88c23832067 100644 --- a/java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowConfigBuilder.java +++ b/java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowConfigBuilder.java @@ -26,6 +26,7 @@ import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Preconditions; +import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.types.pojo.ArrowType; /** @@ -44,6 +45,7 @@ public class JdbcToArrowConfigBuilder { private Map> columnMetadataByColumnIndex; private int targetBatchSize; private Function jdbcToArrowTypeConverter; + private JdbcToArrowConfig.JdbcConsumerFactory jdbcConsumerGetter; private RoundingMode bigDecimalRoundingMode; /** @@ -221,6 +223,18 @@ public JdbcToArrowConfigBuilder setJdbcToArrowTypeConverter( return this; } + /** + * Set the function used to get a JDBC consumer for a given type. + *

+ * Defaults to wrapping {@link + * JdbcToArrowUtils#getConsumer(ArrowType, Integer, Boolean, FieldVector, JdbcToArrowConfig)}. + */ + public JdbcToArrowConfigBuilder setJdbcConsumerGetter( + JdbcToArrowConfig.JdbcConsumerFactory jdbcConsumerGetter) { + this.jdbcConsumerGetter = jdbcConsumerGetter; + return this; + } + /** * Set whether to use the same {@link org.apache.arrow.vector.VectorSchemaRoot} instance on each iteration, * or to allocate a new one. @@ -274,6 +288,7 @@ public JdbcToArrowConfig build() { arraySubTypesByColumnName, targetBatchSize, jdbcToArrowTypeConverter, + jdbcConsumerGetter, explicitTypesByColumnIndex, explicitTypesByColumnName, schemaMetadata, diff --git a/java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowUtils.java b/java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowUtils.java index dc79f6efff324..f8a13b93b1ed8 100644 --- a/java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowUtils.java +++ b/java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowUtils.java @@ -208,7 +208,7 @@ public static ArrowType getArrowTypeFromJdbcType(final JdbcFieldInfo fieldInfo, return new ArrowType.Struct(); default: // no-op, shouldn't get here - return null; + throw new UnsupportedOperationException("Unmapped JDBC type: " + fieldInfo.getJdbcType()); } } @@ -489,7 +489,7 @@ static JdbcConsumer getConsumer(ArrowType arrowType, int columnIndex, boolean nu return new NullConsumer((NullVector) vector); default: // no-op, shouldn't get here - throw new UnsupportedOperationException(); + throw new UnsupportedOperationException("No consumer for Arrow type: " + arrowType); } } } diff --git a/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/example/FlightSqlExample.java b/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/example/FlightSqlExample.java index fe1e1445afc6e..3cc8f4a1c1bd5 100644 --- a/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/example/FlightSqlExample.java +++ b/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/example/FlightSqlExample.java @@ -299,9 +299,12 @@ private static boolean populateDerbyDatabase() { } private static ArrowType getArrowTypeFromJdbcType(final int jdbcDataType, final int precision, final int scale) { - final ArrowType type = - JdbcToArrowUtils.getArrowTypeFromJdbcType(new JdbcFieldInfo(jdbcDataType, precision, scale), DEFAULT_CALENDAR); - return isNull(type) ? ArrowType.Utf8.INSTANCE : type; + try { + return JdbcToArrowUtils.getArrowTypeFromJdbcType(new JdbcFieldInfo(jdbcDataType, precision, scale), + DEFAULT_CALENDAR); + } catch (UnsupportedOperationException ignored) { + return ArrowType.Utf8.INSTANCE; + } } private static void saveToVector(final Byte data, final UInt1Vector vector, final int index) {