From e05d0ce522b0da63b898977bdea89093b19079ed Mon Sep 17 00:00:00 2001
From: David Li
Date: Fri, 8 Jul 2022 15:11:50 -0400
Subject: [PATCH] ARROW-17005: [Java] Allow overriding column nullability in
arrow-jdbc
---
.../adapter/jdbc/ArrowVectorIterator.java | 6 +-
.../arrow/adapter/jdbc/JdbcFieldInfo.java | 32 +-
.../arrow/adapter/jdbc/JdbcToArrowUtils.java | 22 +-
.../arrow/adapter/jdbc/ResultSetUtility.java | 40 ++-
.../adapter/jdbc/ResultSetUtilityTest.java | 107 +++++++
.../adapter/jdbc/UnreliableMetaDataTest.java | 273 ++++++++++++++++++
.../adapter/jdbc/h2/JdbcToArrowTest.java | 247 ----------------
7 files changed, 463 insertions(+), 264 deletions(-)
create mode 100644 java/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/ResultSetUtilityTest.java
create mode 100644 java/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/UnreliableMetaDataTest.java
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 a8a65ff346b92..3ce921d859435 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
@@ -137,10 +137,10 @@ private VectorSchemaRoot createVectorSchemaRoot() throws SQLException {
private void initialize(VectorSchemaRoot root) throws SQLException {
for (int i = 1; i <= consumers.length; i++) {
- ArrowType arrowType = config.getJdbcToArrowTypeConverter()
- .apply(new JdbcFieldInfo(resultSet.getMetaData(), i));
+ final JdbcFieldInfo columnFieldInfo = JdbcToArrowUtils.getJdbcFieldInfoForColumn(rsmd, i, config);
+ ArrowType arrowType = config.getJdbcToArrowTypeConverter().apply(columnFieldInfo);
consumers[i - 1] = JdbcToArrowUtils.getConsumer(
- arrowType, i, isColumnNullable(resultSet, i), root.getVector(i - 1), config);
+ 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/JdbcFieldInfo.java b/java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcFieldInfo.java
index e3747bb04b8cd..3443a1e44c12d 100644
--- a/java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcFieldInfo.java
+++ b/java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcFieldInfo.java
@@ -29,12 +29,14 @@
* Currently, this is:
*
* - The JDBC {@link java.sql.Types} type.
- * - The field's precision (used for {@link java.sql.Types#DECIMAL} and {@link java.sql.Types#NUMERIC} types)
- * - The field's scale (used for {@link java.sql.Types#DECIMAL} and {@link java.sql.Types#NUMERIC} types)
+ * - The nullability.
+ * - The field's precision (used for {@link java.sql.Types#DECIMAL} and {@link java.sql.Types#NUMERIC} types).
+ * - The field's scale (used for {@link java.sql.Types#DECIMAL} and {@link java.sql.Types#NUMERIC} types).
*
*/
public class JdbcFieldInfo {
private final int jdbcType;
+ private final int nullability;
private final int precision;
private final int scale;
@@ -52,6 +54,7 @@ public JdbcFieldInfo(int jdbcType) {
"DECIMAL and NUMERIC types require a precision and scale; please use another constructor.");
this.jdbcType = jdbcType;
+ this.nullability = ResultSetMetaData.columnNullableUnknown;
this.precision = 0;
this.scale = 0;
}
@@ -66,6 +69,23 @@ public JdbcFieldInfo(int jdbcType) {
*/
public JdbcFieldInfo(int jdbcType, int precision, int scale) {
this.jdbcType = jdbcType;
+ this.nullability = ResultSetMetaData.columnNullableUnknown;
+ this.precision = precision;
+ this.scale = scale;
+ }
+
+ /**
+ * Builds a JdbcFieldInfo
from the {@link java.sql.Types} type, nullability, precision, and scale.
+ *
+ * @param jdbcType The {@link java.sql.Types} type.
+ * @param nullability The nullability. Must be one of {@link ResultSetMetaData#columnNoNulls},
+ * {@link ResultSetMetaData#columnNullable}, or {@link ResultSetMetaData#columnNullableUnknown}.
+ * @param precision The field's numeric precision.
+ * @param scale The field's numeric scale.
+ */
+ public JdbcFieldInfo(int jdbcType, int nullability, int precision, int scale) {
+ this.jdbcType = jdbcType;
+ this.nullability = nullability;
this.precision = precision;
this.scale = scale;
}
@@ -87,6 +107,7 @@ public JdbcFieldInfo(ResultSetMetaData rsmd, int column) throws SQLException {
"The index must be within the number of columns (1 to %s, inclusive)", rsmd.getColumnCount());
this.jdbcType = rsmd.getColumnType(column);
+ this.nullability = rsmd.isNullable(column);
this.precision = rsmd.getPrecision(column);
this.scale = rsmd.getScale(column);
}
@@ -98,6 +119,13 @@ public int getJdbcType() {
return jdbcType;
}
+ /**
+ * The nullability.
+ */
+ public int isNullable() {
+ return nullability;
+ }
+
/**
* The numeric precision, for {@link java.sql.Types#NUMERIC} and {@link java.sql.Types#DECIMAL} types.
*/
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 4bb6558e06fe1..43fed849a2298 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
@@ -260,7 +260,7 @@ public static Schema jdbcToArrowSchema(ResultSetMetaData rsmd, JdbcToArrowConfig
final ArrowType arrowType = config.getJdbcToArrowTypeConverter().apply(columnFieldInfo);
if (arrowType != null) {
final FieldType fieldType = new FieldType(
- isColumnNullable(rsmd, i), arrowType, /* dictionary encoding */ null, metadata);
+ isColumnNullable(rsmd, i, columnFieldInfo), arrowType, /* dictionary encoding */ null, metadata);
List children = null;
if (arrowType.getTypeID() == ArrowType.List.TYPE_TYPE) {
@@ -280,7 +280,7 @@ public static Schema jdbcToArrowSchema(ResultSetMetaData rsmd, JdbcToArrowConfig
return new Schema(fields, null);
}
- private static JdbcFieldInfo getJdbcFieldInfoForColumn(
+ static JdbcFieldInfo getJdbcFieldInfoForColumn(
ResultSetMetaData rsmd,
int arrayColumn,
JdbcToArrowConfig config)
@@ -347,12 +347,14 @@ public static void jdbcToArrowVectors(ResultSet rs, VectorSchemaRoot root, Calen
jdbcToArrowVectors(rs, root, new JdbcToArrowConfig(new RootAllocator(0), calendar));
}
- static boolean isColumnNullable(ResultSet resultSet, int index) throws SQLException {
- return isColumnNullable(resultSet.getMetaData(), index);
- }
-
- static boolean isColumnNullable(ResultSetMetaData resultSetMetadata, int index) throws SQLException {
- int nullableValue = resultSetMetadata.isNullable(index);
+ static boolean isColumnNullable(ResultSetMetaData resultSetMetadata, int index, JdbcFieldInfo info)
+ throws SQLException {
+ int nullableValue;
+ if (info != null && info.isNullable() != ResultSetMetaData.columnNullableUnknown) {
+ nullableValue = info.isNullable();
+ } else {
+ nullableValue = resultSetMetadata.isNullable(index);
+ }
return nullableValue == ResultSetMetaData.columnNullable ||
nullableValue == ResultSetMetaData.columnNullableUnknown;
}
@@ -375,7 +377,9 @@ public static void jdbcToArrowVectors(ResultSet rs, VectorSchemaRoot root, JdbcT
JdbcConsumer[] consumers = new JdbcConsumer[columnCount];
for (int i = 1; i <= columnCount; i++) {
FieldVector vector = root.getVector(rsmd.getColumnLabel(i));
- consumers[i - 1] = getConsumer(vector.getField().getType(), i, isColumnNullable(rs, i), vector, config);
+ final JdbcFieldInfo columnFieldInfo = getJdbcFieldInfoForColumn(rsmd, i, config);
+ consumers[i - 1] = getConsumer(
+ vector.getField().getType(), i, isColumnNullable(rsmd, i, columnFieldInfo), vector, config);
}
CompositeJdbcConsumer compositeConsumer = null;
diff --git a/java/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/ResultSetUtility.java b/java/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/ResultSetUtility.java
index 2ffba2e484501..339e120beea76 100644
--- a/java/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/ResultSetUtility.java
+++ b/java/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/ResultSetUtility.java
@@ -40,6 +40,7 @@
import java.sql.Types;
import java.util.ArrayList;
import java.util.Calendar;
+import java.util.List;
import java.util.Map;
public class ResultSetUtility {
@@ -266,9 +267,9 @@ public MockResultSet build() throws SQLException {
}
public static class MockResultSetMetaData extends ThrowingResultSetMetaData {
- private ArrayList columns;
+ private final List columns;
- public MockResultSetMetaData(ArrayList columns) {
+ public MockResultSetMetaData(List columns) {
this.columns = columns;
}
@@ -491,6 +492,9 @@ private int getSqlType() throws SQLException {
}
public BigDecimal getBigDecimal() throws SQLException {
+ if (value == null) {
+ return null;
+ }
try {
return new BigDecimal(getValueAsString());
} catch (Exception ex) {
@@ -499,10 +503,16 @@ public BigDecimal getBigDecimal() throws SQLException {
}
public String getString() throws SQLException {
+ if (value == null) {
+ return null;
+ }
return getValueAsString();
}
public boolean getBoolean() throws SQLException {
+ if (value == null) {
+ return false;
+ }
try {
return (boolean) value;
} catch (Exception ex) {
@@ -511,6 +521,9 @@ public boolean getBoolean() throws SQLException {
}
public int getInt() throws SQLException {
+ if (value == null) {
+ return 0;
+ }
try {
return Integer.parseInt(getValueAsString());
} catch (Exception ex) {
@@ -519,6 +532,9 @@ public int getInt() throws SQLException {
}
public long getLong() throws SQLException {
+ if (value == null) {
+ return 0L;
+ }
try {
return Long.parseLong(getValueAsString());
} catch (Exception ex) {
@@ -527,6 +543,9 @@ public long getLong() throws SQLException {
}
public double getDouble() throws SQLException {
+ if (value == null) {
+ return 0.0;
+ }
try {
return Double.parseDouble(getValueAsString());
} catch (Exception ex) {
@@ -535,6 +554,9 @@ public double getDouble() throws SQLException {
}
public Date getDate() throws SQLException {
+ if (value == null) {
+ return null;
+ }
try {
return Date.valueOf(getValueAsString());
} catch (Exception ex) {
@@ -543,6 +565,9 @@ public Date getDate() throws SQLException {
}
public Time getTime() throws SQLException {
+ if (value == null) {
+ return null;
+ }
try {
return Time.valueOf(getValueAsString());
} catch (Exception ex) {
@@ -551,6 +576,9 @@ public Time getTime() throws SQLException {
}
public Timestamp getTimestamp() throws SQLException {
+ if (value == null) {
+ return null;
+ }
try {
return Timestamp.valueOf(getValueAsString());
} catch (Exception ex) {
@@ -558,7 +586,10 @@ public Timestamp getTimestamp() throws SQLException {
}
}
- public Float getFloat() throws SQLException {
+ public float getFloat() throws SQLException {
+ if (value == null) {
+ return 0.0f;
+ }
try {
return Float.parseFloat(getValueAsString());
} catch (Exception ex) {
@@ -567,6 +598,9 @@ public Float getFloat() throws SQLException {
}
public short getShort() throws SQLException {
+ if (value == null) {
+ return 0;
+ }
try {
return Short.parseShort(getValueAsString());
} catch (Exception ex) {
diff --git a/java/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/ResultSetUtilityTest.java b/java/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/ResultSetUtilityTest.java
new file mode 100644
index 0000000000000..2424ed625248d
--- /dev/null
+++ b/java/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/ResultSetUtilityTest.java
@@ -0,0 +1,107 @@
+/*
+ * 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 static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Types;
+
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.junit.Test;
+
+/** Tests of the ResultSetUtility. */
+public class ResultSetUtilityTest {
+ @Test
+ public void testZeroRowResultSet() throws Exception {
+ for (boolean reuseVectorSchemaRoot : new boolean[]{false, true}) {
+ try (BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE)) {
+ ResultSet rs = ResultSetUtility.generateEmptyResultSet();
+ JdbcToArrowConfig config = new JdbcToArrowConfigBuilder(
+ allocator, JdbcToArrowUtils.getUtcCalendar(), /* include metadata */ false)
+ .setReuseVectorSchemaRoot(reuseVectorSchemaRoot)
+ .build();
+
+ ArrowVectorIterator iter = JdbcToArrow.sqlToArrowVectorIterator(rs, config);
+ assertTrue("Iterator on zero row ResultSet should haveNext() before use", iter.hasNext());
+ VectorSchemaRoot root = iter.next();
+ assertNotNull("VectorSchemaRoot from first next() result should never be null", root);
+ assertEquals("VectorSchemaRoot from empty ResultSet should have zero rows", 0, root.getRowCount());
+ assertFalse("hasNext() should return false on empty ResultSets after initial next() call", iter.hasNext());
+ }
+ }
+ }
+
+ @Test
+ public void testBasicResultSet() throws Exception {
+ try (ResultSetUtility.MockResultSet resultSet = ResultSetUtility.generateBasicResultSet(3)) {
+ // Before row 1:
+ assertTrue(resultSet.isBeforeFirst());
+ assertFalse(resultSet.isFirst());
+ assertFalse(resultSet.isLast());
+ assertFalse(resultSet.isAfterLast());
+ assertThrows(SQLException.class, () -> resultSet.getString(1));
+
+ // Row 1:
+ assertTrue(resultSet.next());
+ assertFalse(resultSet.isBeforeFirst());
+ assertTrue(resultSet.isFirst());
+ assertFalse(resultSet.isLast());
+ assertFalse(resultSet.isAfterLast());
+ assertEquals("row number: 1", resultSet.getString(1));
+
+ // Row 2:
+ assertTrue(resultSet.next());
+ assertFalse(resultSet.isBeforeFirst());
+ assertFalse(resultSet.isFirst());
+ assertFalse(resultSet.isLast());
+ assertFalse(resultSet.isAfterLast());
+ assertEquals("row number: 2", resultSet.getString(1));
+
+ // Row 3:
+ assertTrue(resultSet.next());
+ assertFalse(resultSet.isBeforeFirst());
+ assertFalse(resultSet.isFirst());
+ assertTrue(resultSet.isLast());
+ assertFalse(resultSet.isAfterLast());
+ assertEquals("row number: 3", resultSet.getString(1));
+
+ // After row 3:
+ assertFalse(resultSet.next());
+ assertFalse(resultSet.isBeforeFirst());
+ assertFalse(resultSet.isFirst());
+ assertFalse(resultSet.isLast());
+ assertTrue(resultSet.isAfterLast());
+ }
+ }
+
+ @Test
+ public void testMockDataTypes() throws SQLException {
+ ResultSetUtility.MockDataElement element = new ResultSetUtility.MockDataElement(1L, Types.NUMERIC);
+ assertEquals(1L, element.getLong());
+ assertEquals(1, element.getInt());
+ assertEquals("1", element.getString());
+ }
+}
diff --git a/java/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/UnreliableMetaDataTest.java b/java/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/UnreliableMetaDataTest.java
new file mode 100644
index 0000000000000..e4aa50c36bb15
--- /dev/null
+++ b/java/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/UnreliableMetaDataTest.java
@@ -0,0 +1,273 @@
+/*
+ * 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 static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import java.math.BigDecimal;
+import java.math.RoundingMode;
+import java.sql.ResultSet;
+import java.sql.ResultSetMetaData;
+import java.sql.SQLException;
+import java.sql.Types;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.IntVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.types.pojo.Schema;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+/**
+ * Test options for dealing with unreliable ResultSetMetaData from JDBC drivers.
+ */
+@RunWith(Parameterized.class)
+public class UnreliableMetaDataTest {
+ private final boolean reuseVectorSchemaRoot;
+ private BufferAllocator allocator;
+
+ public UnreliableMetaDataTest(boolean reuseVectorSchemaRoot) {
+ this.reuseVectorSchemaRoot = reuseVectorSchemaRoot;
+ }
+
+ @Before
+ public void beforeEach() {
+ allocator = new RootAllocator();
+ }
+
+ @After
+ public void afterEach() {
+ allocator.close();
+ }
+
+ @Parameterized.Parameters(name = "reuseVectorSchemaRoot = {0}")
+ public static Collection