Skip to content

Commit

Permalink
ARROW-17005: [Java] Allow overriding column nullability in arrow-jdbc
Browse files Browse the repository at this point in the history
  • Loading branch information
lidavidm committed Jul 8, 2022
1 parent 8042f00 commit e05d0ce
Show file tree
Hide file tree
Showing 7 changed files with 463 additions and 264 deletions.
Expand Up @@ -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);
}
}

Expand Down
Expand Up @@ -29,12 +29,14 @@
* 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>
* <li>The nullability.</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 nullability;
private final int precision;
private final int scale;

Expand All @@ -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;
}
Expand All @@ -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 <code>JdbcFieldInfo</code> 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;
}
Expand All @@ -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);
}
Expand All @@ -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.
*/
Expand Down
Expand Up @@ -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<Field> children = null;
if (arrowType.getTypeID() == ArrowType.List.TYPE_TYPE) {
Expand All @@ -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)
Expand Down Expand Up @@ -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;
}
Expand All @@ -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;
Expand Down
Expand Up @@ -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 {
Expand Down Expand Up @@ -266,9 +267,9 @@ public MockResultSet build() throws SQLException {
}

public static class MockResultSetMetaData extends ThrowingResultSetMetaData {
private ArrayList<MockColumnMetaData> columns;
private final List<MockColumnMetaData> columns;

public MockResultSetMetaData(ArrayList<MockColumnMetaData> columns) {
public MockResultSetMetaData(List<MockColumnMetaData> columns) {
this.columns = columns;
}

Expand Down Expand Up @@ -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) {
Expand All @@ -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) {
Expand All @@ -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) {
Expand All @@ -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) {
Expand All @@ -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) {
Expand All @@ -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) {
Expand All @@ -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) {
Expand All @@ -551,14 +576,20 @@ public Time getTime() throws SQLException {
}

public Timestamp getTimestamp() throws SQLException {
if (value == null) {
return null;
}
try {
return Timestamp.valueOf(getValueAsString());
} catch (Exception ex) {
throw new SQLException(ex);
}
}

public Float getFloat() throws SQLException {
public float getFloat() throws SQLException {
if (value == null) {
return 0.0f;
}
try {
return Float.parseFloat(getValueAsString());
} catch (Exception ex) {
Expand All @@ -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) {
Expand Down
@@ -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());
}
}

0 comments on commit e05d0ce

Please sign in to comment.