Skip to content

Commit

Permalink
Implement Java extension type
Browse files Browse the repository at this point in the history
  • Loading branch information
lidavidm committed Jun 21, 2024
1 parent 24f2636 commit 95ff4c9
Show file tree
Hide file tree
Showing 7 changed files with 674 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -211,6 +211,8 @@ public JdbcToArrowConfigBuilder setTargetBatchSize(int targetBatchSize) {
*
* <p>Defaults to wrapping {@link JdbcToArrowUtils#getArrowTypeFromJdbcType(JdbcFieldInfo,
* Calendar)}.
*
* @see JdbcToArrowUtils#reportUnsupportedTypesAsUnknown(Function)
*/
public JdbcToArrowConfigBuilder setJdbcToArrowTypeConverter(
Function<JdbcFieldInfo, ArrowType> jdbcToArrowTypeConverter) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@

import static org.apache.arrow.vector.types.FloatingPointPrecision.DOUBLE;
import static org.apache.arrow.vector.types.FloatingPointPrecision.SINGLE;
import static org.apache.arrow.vector.types.Types.MinorType;

import java.io.IOException;
import java.math.RoundingMode;
Expand All @@ -37,6 +38,7 @@
import java.util.Locale;
import java.util.Map;
import java.util.TimeZone;
import java.util.function.Function;
import org.apache.arrow.adapter.jdbc.consumer.ArrayConsumer;
import org.apache.arrow.adapter.jdbc.consumer.BigIntConsumer;
import org.apache.arrow.adapter.jdbc.consumer.BinaryConsumer;
Expand Down Expand Up @@ -80,6 +82,7 @@
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.complex.ListVector;
import org.apache.arrow.vector.complex.MapVector;
import org.apache.arrow.vector.extension.UnknownType;
import org.apache.arrow.vector.types.DateUnit;
import org.apache.arrow.vector.types.TimeUnit;
import org.apache.arrow.vector.types.pojo.ArrowType;
Expand Down Expand Up @@ -216,11 +219,28 @@ public static ArrowType getArrowTypeFromJdbcType(
case Types.STRUCT:
return new ArrowType.Struct();
default:
// no-op, shouldn't get here
throw new UnsupportedOperationException("Unmapped JDBC type: " + fieldInfo.getJdbcType());
}
}

/**
* Wrap a JDBC to Arrow type converter such that {@link UnsupportedOperationException} becomes
* {@link org.apache.arrow.vector.extension.UnknownType}.
*
* @param typeConverter The type converter to wrap.
* @param vendorName The database name to report as the Unknown type's vendor name.
*/
public static Function<JdbcFieldInfo, ArrowType> reportUnsupportedTypesAsUnknown(
Function<JdbcFieldInfo, ArrowType> typeConverter, String vendorName) {
return (final JdbcFieldInfo fieldInfo) -> {
try {
return typeConverter.apply(fieldInfo);
} catch (UnsupportedOperationException e) {
return new UnknownType(MinorType.NULL.getType(), fieldInfo.getTypeName(), vendorName);
}
};
}

/**
* Create Arrow {@link Schema} object for the given JDBC {@link java.sql.ResultSetMetaData}.
*
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -32,19 +32,28 @@
import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertTinyIntVectorValues;
import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertVarBinaryVectorValues;
import static org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper.assertVarcharVectorValues;
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertEquals;

import java.io.IOException;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.Arrays;
import java.util.Calendar;
import java.util.List;
import java.util.function.Function;
import java.util.stream.Stream;
import org.apache.arrow.adapter.jdbc.AbstractJdbcToArrowTest;
import org.apache.arrow.adapter.jdbc.JdbcFieldInfo;
import org.apache.arrow.adapter.jdbc.JdbcToArrowConfig;
import org.apache.arrow.adapter.jdbc.JdbcToArrowConfigBuilder;
import org.apache.arrow.adapter.jdbc.JdbcToArrowTestHelper;
import org.apache.arrow.adapter.jdbc.JdbcToArrowUtils;
import org.apache.arrow.adapter.jdbc.Table;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.BigIntVector;
import org.apache.arrow.vector.BitVector;
Expand All @@ -62,7 +71,12 @@
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.complex.ListVector;
import org.apache.arrow.vector.extension.UnknownType;
import org.apache.arrow.vector.types.Types;
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.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;
Expand Down Expand Up @@ -189,6 +203,44 @@ public void testJdbcSchemaMetadata(Table table) throws SQLException, ClassNotFou
JdbcToArrowTestHelper.assertFieldMetadataMatchesResultSetMetadata(rsmd, schema);
}

@Test
void testUnknownType() throws SQLException, ClassNotFoundException {
try (BufferAllocator allocator = new RootAllocator()) {
String url = "jdbc:h2:mem:JdbcToArrowTest";
String driver = "org.h2.Driver";
Class.forName(driver);
conn = DriverManager.getConnection(url);
try (Statement stmt = conn.createStatement()) {
stmt.executeUpdate("CREATE TABLE unknowntype (a GEOMETRY, b INT)");
}

String query = "SELECT * FROM unknowntype";
Calendar calendar = Calendar.getInstance();
Function<JdbcFieldInfo, ArrowType> typeConverter =
(field) -> JdbcToArrowUtils.getArrowTypeFromJdbcType(field, calendar);
JdbcToArrowConfig config =
new JdbcToArrowConfigBuilder()
.setAllocator(allocator)
.setJdbcToArrowTypeConverter(
JdbcToArrowUtils.reportUnsupportedTypesAsUnknown(typeConverter, "H2"))
.build();
Schema schema;
try (Statement stmt = conn.createStatement();
ResultSet rs = stmt.executeQuery(query)) {
schema =
assertDoesNotThrow(() -> JdbcToArrowUtils.jdbcToArrowSchema(rs.getMetaData(), config));
}

Schema expected =
new Schema(
List.of(
Field.nullable(
"A", new UnknownType(Types.MinorType.NULL.getType(), "GEOMETRY", "H2")),
Field.nullable("B", Types.MinorType.INT.getType())));
assertEquals(expected, schema);
}
}

/**
* This method calls the assert methods for various DataSets.
*
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
/*
* 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.vector.extension;

/** */
public class InvalidExtensionMetadataException extends RuntimeException {
public InvalidExtensionMetadataException(String message) {
super(message);
}

public InvalidExtensionMetadataException(String message, Throwable cause) {
super(message, cause);
}
}
Loading

0 comments on commit 95ff4c9

Please sign in to comment.