This is an automated email from the ASF dual-hosted git repository.
lidavidm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/main by this push:
new 0e4fa8af7b GH-37021 [Java][arrow-jdbc] Pluggable getConsumer (#37085)
0e4fa8af7b is described below
commit 0e4fa8af7b72493ffe05388a7f3fed0d4181e7b2
Author: Diego Fernández Giraldo <[email protected]>
AuthorDate: Thu Aug 10 17:43:45 2023 -0400
GH-37021 [Java][arrow-jdbc] Pluggable getConsumer (#37085)
### Rationale for this change
This was discussed in [this
thread](https://github.com/apache/arrow/issues/37021#issuecomment-1666169351).
The `getConsumer` implementation depends on the implementation for
`getArrowTypeFromJdbcType`, especially for timestamp objects. Since we can
provide a different implementation for `JdbcToArrowTypeConverter` it follows we
should be able to provide a different implementation for `JdbcConsumerGetter`
### What changes are included in this PR?
Adds a way to configure an alternate `JdbcToArrowUtils.getConsumer`
function through `JdbcToArrowConfigBuilder.setJdbcConsumerGetter`.
It also throws a more helpful exception from the default implementations
when a provided type is not mapped.
### Are these changes tested?
No, the default behavior remains unchanged.
* Related: #37021
* Closes: #37021
Authored-by: Diego Fernandez <[email protected]>
Signed-off-by: David Li <[email protected]>
---
.../arrow/adapter/jdbc/ArrowVectorIterator.java | 2 +-
.../arrow/adapter/jdbc/JdbcToArrowConfig.java | 53 ++++++++++++++++++++++
.../adapter/jdbc/JdbcToArrowConfigBuilder.java | 15 ++++++
.../arrow/adapter/jdbc/JdbcToArrowUtils.java | 4 +-
.../arrow/flight/sql/example/FlightSqlExample.java | 9 ++--
5 files changed, 77 insertions(+), 6 deletions(-)
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 6c3e9cf43e..6e789009dd 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 @@ public class ArrowVectorIterator implements
Iterator<VectorSchemaRoot>, AutoClos
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 012cd95c0b..e23bad54af 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.Calendar;
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<JdbcFieldInfo, ArrowType> jdbcToArrowTypeConverter;
+ private final JdbcConsumerFactory jdbcConsumerGetter;
/**
* Constructs a new configuration from the provided allocator and calendar.
The <code>allocator</code>
@@ -195,6 +198,38 @@ public final class JdbcToArrowConfig {
Map<String, String> schemaMetadata,
Map<Integer, Map<String, String>> 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<Integer, JdbcFieldInfo> arraySubTypesByColumnIndex,
+ Map<String, JdbcFieldInfo> arraySubTypesByColumnName,
+ int targetBatchSize,
+ Function<JdbcFieldInfo, ArrowType> jdbcToArrowTypeConverter,
+ JdbcConsumerFactory jdbcConsumerGetter,
+ Map<Integer, JdbcFieldInfo> explicitTypesByColumnIndex,
+ Map<String, JdbcFieldInfo> explicitTypesByColumnName,
+ Map<String, String> schemaMetadata,
+ Map<Integer, Map<String, String>> 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 final class JdbcToArrowConfig {
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 final class JdbcToArrowConfig {
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 2fe0492deb..7d88c23832 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 java.util.function.Function;
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<Integer, Map<String, String>> columnMetadataByColumnIndex;
private int targetBatchSize;
private Function<JdbcFieldInfo, ArrowType> jdbcToArrowTypeConverter;
+ private JdbcToArrowConfig.JdbcConsumerFactory jdbcConsumerGetter;
private RoundingMode bigDecimalRoundingMode;
/**
@@ -221,6 +223,18 @@ public class JdbcToArrowConfigBuilder {
return this;
}
+ /**
+ * Set the function used to get a JDBC consumer for a given type.
+ * <p>
+ * 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 class JdbcToArrowConfigBuilder {
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 dc79f6efff..f8a13b93b1 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 class JdbcToArrowUtils {
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 @@ public class JdbcToArrowUtils {
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 fe1e1445af..3cc8f4a1c1 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 @@ public class FlightSqlExample implements
FlightSqlProducer, AutoCloseable {
}
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) {
