Skip to content

Commit c1d4407

Browse files
vfragakylep-dremiojcralmeidarafael-tellesiurysalino
authored
Add fixes from PRs apache#233, apache#234, and apache#235 (apache#237)
* Set netty property when JDBC driver is loaded (apache#235) Co-authored-by: Vinicius Fraga <sxvinifp@gmail.com> * [Java] [JDBC] Change the statement implementation to reduce the job numbers. (apache#234) * Properly read incoming PreparedStatement columnmetadata * Make map Thread-safe, rename map and change getPreparedStatement return type Co-authored-by: Vinicius Fraga <sxvinifp@gmail.com> * Add column metadata when creating PreparedStatement (apache#233) Co-authored-by: Vinicius Fraga <sxvinifp@gmail.com> * Minor fixes Co-authored-by: Kyle Porter <kporter@dremio.com> Co-authored-by: Jose Almeida <almeidajcr90@gmail.com> Co-authored-by: Rafael Telles <rafael@telles.dev> Co-authored-by: Vinicius Fraga <sxvinifp@gmail.com> Co-authored-by: kylep-dremio <38920967+kylep-dremio@users.noreply.github.com> Co-authored-by: iurysalino <iurysalino@gmail.com> Co-authored-by: Gabriel Escobar <51451696+escobargabriel@users.noreply.github.com>
1 parent 928435a commit c1d4407

12 files changed

Lines changed: 343 additions & 88 deletions

‎java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -171,4 +171,8 @@ public void close() throws SQLException {
171171
BufferAllocator getBufferAllocator() {
172172
return allocator;
173173
}
174+
175+
public ArrowFlightMetaImpl getMeta() {
176+
return (ArrowFlightMetaImpl) this.meta;
177+
}
174178
}

‎java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriver.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,16 @@ public class ArrowFlightJdbcDriver extends UnregisteredDriver {
5050
private static DriverVersion version;
5151

5252
static {
53+
// Special code for supporting Java9 and higher.
54+
// Netty requires some extra properties to unlock some native memory management api
55+
// Setting this property if not already set externally
56+
// This has to be done before any netty class is being loaded
57+
final String key = "cfjd.io.netty.tryReflectionSetAccessible";
58+
final String tryReflectionSetAccessible = System.getProperty(key);
59+
if (tryReflectionSetAccessible == null) {
60+
System.setProperty(key, Boolean.TRUE.toString());
61+
}
62+
5363
new ArrowFlightJdbcDriver().register();
5464
}
5565

‎java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java‎

Lines changed: 20 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import org.apache.arrow.flight.FlightStream;
3333
import org.apache.arrow.util.AutoCloseables;
3434
import org.apache.arrow.vector.VectorSchemaRoot;
35+
import org.apache.arrow.vector.types.pojo.Schema;
3536
import org.apache.calcite.avatica.AvaticaResultSet;
3637
import org.apache.calcite.avatica.AvaticaResultSetMetaData;
3738
import org.apache.calcite.avatica.AvaticaStatement;
@@ -52,6 +53,8 @@ public final class ArrowFlightJdbcFlightStreamResultSet
5253
private VectorSchemaRootTransformer transformer;
5354
private VectorSchemaRoot currentVectorSchemaRoot;
5455

56+
private Schema schema;
57+
5558
ArrowFlightJdbcFlightStreamResultSet(final AvaticaStatement statement,
5659
final QueryState state,
5760
final Meta.Signature signature,
@@ -103,12 +106,8 @@ static ArrowFlightJdbcFlightStreamResultSet fromFlightInfo(
103106
return resultSet;
104107
}
105108

106-
FlightStreamQueue getFlightStreamQueue() {
107-
return flightStreamQueue;
108-
}
109-
110109
private void loadNewQueue() {
111-
Optional.ofNullable(getFlightStreamQueue()).ifPresent(AutoCloseables::closeNoChecked);
110+
Optional.ofNullable(flightStreamQueue).ifPresent(AutoCloseables::closeNoChecked);
112111
flightStreamQueue = createNewQueue(connection.getExecutorService());
113112
}
114113

@@ -124,21 +123,21 @@ protected AvaticaResultSet execute() throws SQLException {
124123
final FlightInfo flightInfo = ((ArrowFlightInfoStatement) statement).executeFlightInfoQuery();
125124

126125
if (flightInfo != null) {
126+
schema = flightInfo.getSchema();
127127
execute(flightInfo);
128128
}
129129
return this;
130130
}
131131

132-
private AvaticaResultSet execute(final FlightInfo flightInfo) throws SQLException {
132+
private void execute(final FlightInfo flightInfo) throws SQLException {
133133
loadNewQueue();
134-
getFlightStreamQueue().enqueue(connection.getClientHandler().getStreams(flightInfo));
134+
flightStreamQueue.enqueue(connection.getClientHandler().getStreams(flightInfo));
135135
loadNewFlightStream();
136136

137137
// Ownership of the root will be passed onto the cursor.
138138
if (currentFlightStream != null) {
139139
executeForCurrentFlightStream();
140140
}
141-
return this;
142141
}
143142

144143
private void executeForCurrentFlightStream() throws SQLException {
@@ -153,7 +152,12 @@ private void executeForCurrentFlightStream() throws SQLException {
153152
} else {
154153
currentVectorSchemaRoot = originalRoot;
155154
}
156-
execute(currentVectorSchemaRoot);
155+
156+
if (schema != null) {
157+
execute(currentVectorSchemaRoot, schema);
158+
} else {
159+
execute(currentVectorSchemaRoot);
160+
}
157161
}
158162

159163
@Override
@@ -208,7 +212,6 @@ protected void cancel() {
208212
currentFlightStream.cancel("Cancel", null);
209213
}
210214

211-
final FlightStreamQueue flightStreamQueue = getFlightStreamQueue();
212215
if (flightStreamQueue != null) {
213216
try {
214217
flightStreamQueue.close();
@@ -221,7 +224,13 @@ protected void cancel() {
221224
@Override
222225
public synchronized void close() {
223226
try {
224-
AutoCloseables.close(currentVectorSchemaRoot, currentFlightStream, getFlightStreamQueue());
227+
if (flightStreamQueue != null) {
228+
// flightStreamQueue should close currentFlightStream internally
229+
flightStreamQueue.close();
230+
} else if (currentFlightStream != null) {
231+
// close is only called for currentFlightStream if there's no queue
232+
currentFlightStream.close();
233+
}
225234
} catch (final Exception e) {
226235
throw new RuntimeException(e);
227236
} finally {

‎java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java‎

Lines changed: 12 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -24,18 +24,14 @@
2424
import java.sql.SQLException;
2525
import java.util.HashSet;
2626
import java.util.List;
27-
import java.util.Map;
2827
import java.util.Set;
2928
import java.util.TimeZone;
30-
import java.util.stream.Collectors;
31-
import java.util.stream.Stream;
3229

33-
import org.apache.arrow.driver.jdbc.utils.SqlTypes;
34-
import org.apache.arrow.flight.sql.FlightSqlColumnMetadata;
30+
import org.apache.arrow.driver.jdbc.utils.ConvertUtils;
3531
import org.apache.arrow.util.AutoCloseables;
3632
import org.apache.arrow.vector.VectorSchemaRoot;
37-
import org.apache.arrow.vector.types.pojo.ArrowType;
3833
import org.apache.arrow.vector.types.pojo.Field;
34+
import org.apache.arrow.vector.types.pojo.Schema;
3935
import org.apache.calcite.avatica.AvaticaResultSet;
4036
import org.apache.calcite.avatica.AvaticaResultSetMetaData;
4137
import org.apache.calcite.avatica.AvaticaStatement;
@@ -44,7 +40,6 @@
4440
import org.apache.calcite.avatica.Meta.Frame;
4541
import org.apache.calcite.avatica.Meta.Signature;
4642
import org.apache.calcite.avatica.QueryState;
47-
import org.apache.calcite.avatica.proto.Common;
4843
import org.slf4j.Logger;
4944
import org.slf4j.LoggerFactory;
5045

@@ -92,79 +87,23 @@ public static ArrowFlightJdbcVectorSchemaRootResultSet fromVectorSchemaRoot(
9287
return resultSet;
9388
}
9489

95-
private static List<ColumnMetaData> convertArrowFieldsToColumnMetaDataList(final List<Field> fields) {
96-
return Stream.iterate(0, Math::incrementExact).limit(fields.size())
97-
.map(index -> {
98-
final Field field = fields.get(index);
99-
final ArrowType.ArrowTypeID fieldTypeId = field.getType().getTypeID();
100-
101-
final Common.ColumnMetaData.Builder builder = Common.ColumnMetaData.newBuilder();
102-
builder.setOrdinal(index);
103-
builder.setColumnName(field.getName());
104-
builder.setLabel(field.getName());
105-
106-
setOnColumnMetaDataBuilder(builder, field.getMetadata());
107-
108-
builder.setType(Common.AvaticaType.newBuilder()
109-
.setId(SqlTypes.getSqlTypeIdFromArrowType(field.getType()))
110-
.setName(fieldTypeId.name())
111-
.build());
112-
113-
return ColumnMetaData.fromProto(builder.build());
114-
}).collect(Collectors.toList());
115-
}
116-
117-
private static void setOnColumnMetaDataBuilder(final Common.ColumnMetaData.Builder builder,
118-
final Map<String, String> metadataMap) {
119-
final FlightSqlColumnMetadata columnMetadata = new FlightSqlColumnMetadata(metadataMap);
120-
final String catalogName = columnMetadata.getCatalogName();
121-
if (catalogName != null) {
122-
builder.setCatalogName(catalogName);
123-
}
124-
final String schemaName = columnMetadata.getSchemaName();
125-
if (schemaName != null) {
126-
builder.setSchemaName(schemaName);
127-
}
128-
final String tableName = columnMetadata.getTableName();
129-
if (tableName != null) {
130-
builder.setTableName(tableName);
131-
}
132-
133-
final Integer precision = columnMetadata.getPrecision();
134-
if (precision != null) {
135-
builder.setPrecision(precision);
136-
}
137-
final Integer scale = columnMetadata.getScale();
138-
if (scale != null) {
139-
builder.setScale(scale);
140-
}
141-
142-
final Boolean isAutoIncrement = columnMetadata.isAutoIncrement();
143-
if (isAutoIncrement != null) {
144-
builder.setAutoIncrement(isAutoIncrement);
145-
}
146-
final Boolean caseSensitive = columnMetadata.isCaseSensitive();
147-
if (caseSensitive != null) {
148-
builder.setCaseSensitive(caseSensitive);
149-
}
150-
final Boolean readOnly = columnMetadata.isReadOnly();
151-
if (readOnly != null) {
152-
builder.setReadOnly(readOnly);
153-
}
154-
final Boolean searchable = columnMetadata.isSearchable();
155-
if (searchable != null) {
156-
builder.setSearchable(searchable);
157-
}
158-
}
159-
16090
@Override
16191
protected AvaticaResultSet execute() throws SQLException {
16292
throw new RuntimeException();
16393
}
16494

16595
void execute(final VectorSchemaRoot vectorSchemaRoot) {
16696
final List<Field> fields = vectorSchemaRoot.getSchema().getFields();
167-
final List<ColumnMetaData> columns = convertArrowFieldsToColumnMetaDataList(fields);
97+
final List<ColumnMetaData> columns = ConvertUtils.convertArrowFieldsToColumnMetaDataList(fields);
98+
signature.columns.clear();
99+
signature.columns.addAll(columns);
100+
101+
this.vectorSchemaRoot = vectorSchemaRoot;
102+
execute2(new ArrowFlightJdbcCursor(vectorSchemaRoot), this.signature.columns);
103+
}
104+
105+
void execute(final VectorSchemaRoot vectorSchemaRoot, final Schema schema) {
106+
final List<ColumnMetaData> columns = ConvertUtils.convertArrowFieldsToColumnMetaDataList(schema.getFields());
168107
signature.columns.clear();
169108
signature.columns.addAll(columns);
170109

‎java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@
2525
import java.util.ArrayList;
2626
import java.util.Collections;
2727
import java.util.List;
28+
import java.util.Map;
29+
import java.util.concurrent.ConcurrentHashMap;
2830

2931
import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler.PreparedStatement;
3032
import org.apache.arrow.util.Preconditions;
@@ -40,9 +42,15 @@
4042
* Metadata handler for Arrow Flight.
4143
*/
4244
public class ArrowFlightMetaImpl extends MetaImpl {
45+
private final Map<StatementHandle, PreparedStatement> statementHandlePreparedStatementMap;
4346

47+
/**
48+
* Constructs a {@link MetaImpl} object specific for Arrow Flight.
49+
* @param connection A {@link AvaticaConnection}.
50+
*/
4451
public ArrowFlightMetaImpl(final AvaticaConnection connection) {
4552
super(connection);
53+
this.statementHandlePreparedStatementMap = new ConcurrentHashMap<>();
4654
setDefaultConnectionProperties();
4755
}
4856

@@ -59,7 +67,11 @@ static Signature newSignature(final String sql) {
5967

6068
@Override
6169
public void closeStatement(final StatementHandle statementHandle) {
62-
// NO-OP.
70+
PreparedStatement preparedStatement = statementHandlePreparedStatementMap.remove(statementHandle);
71+
// Testing if the prepared statement was created because the statement can be not created until this moment
72+
if (preparedStatement != null) {
73+
preparedStatement.close();
74+
}
6375
}
6476

6577
@Override
@@ -131,6 +143,7 @@ public ExecuteResult prepareAndExecute(final StatementHandle handle,
131143
final PreparedStatement preparedStatement =
132144
((ArrowFlightConnection) connection).getClientHandler().prepare(query);
133145
final StatementType statementType = preparedStatement.getType();
146+
statementHandlePreparedStatementMap.put(handle, preparedStatement);
134147
final Signature signature = newSignature(query);
135148
final long updateCount =
136149
statementType.equals(StatementType.UPDATE) ? preparedStatement.executeUpdate() : -1;
@@ -180,4 +193,8 @@ void setDefaultConnectionProperties() {
180193
.setSchema(null)
181194
.setTransactionIsolation(Connection.TRANSACTION_NONE);
182195
}
196+
197+
PreparedStatement getPreparedStatement(StatementHandle statementHandle) {
198+
return statementHandlePreparedStatementMap.get(statementHandle);
199+
}
183200
}

‎java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,10 @@
2222
import java.sql.SQLException;
2323

2424
import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler;
25+
import org.apache.arrow.driver.jdbc.utils.ConvertUtils;
2526
import org.apache.arrow.flight.FlightInfo;
2627
import org.apache.arrow.util.Preconditions;
28+
import org.apache.arrow.vector.types.pojo.Schema;
2729
import org.apache.calcite.avatica.AvaticaPreparedStatement;
2830
import org.apache.calcite.avatica.Meta.Signature;
2931
import org.apache.calcite.avatica.Meta.StatementHandle;
@@ -67,8 +69,14 @@ static ArrowFlightPreparedStatement createNewPreparedStatement(
6769
final int resultSetType,
6870
final int resultSetConcurrency,
6971
final int resultSetHoldability) throws SQLException {
72+
73+
final ArrowFlightSqlClientHandler.PreparedStatement prepare = connection.getClientHandler().prepare(signature.sql);
74+
final Schema resultSetSchema = prepare.getDataSetSchema();
75+
76+
signature.columns.addAll(ConvertUtils.convertArrowFieldsToColumnMetaDataList(resultSetSchema.getFields()));
77+
7078
return new ArrowFlightPreparedStatement(
71-
connection, connection.getClientHandler().prepare(signature.sql), statementHandle,
79+
connection, prepare, statementHandle,
7280
signature, resultSetType, resultSetConcurrency, resultSetHoldability);
7381
}
7482

‎java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919

2020
import java.sql.SQLException;
2121

22+
import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler.PreparedStatement;
2223
import org.apache.arrow.flight.FlightInfo;
2324
import org.apache.calcite.avatica.AvaticaStatement;
2425
import org.apache.calcite.avatica.Meta;
@@ -42,11 +43,12 @@ public ArrowFlightConnection getConnection() throws SQLException {
4243

4344
@Override
4445
public FlightInfo executeFlightInfoQuery() throws SQLException {
46+
final PreparedStatement preparedStatement = getConnection().getMeta().getPreparedStatement(handle);
4547
final Meta.Signature signature = getSignature();
4648
if (signature == null) {
4749
return null;
4850
}
4951

50-
return getConnection().getClientHandler().getInfo(signature.sql);
52+
return preparedStatement.executeQuery();
5153
}
5254
}

‎java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,13 @@ public interface PreparedStatement extends AutoCloseable {
141141
*/
142142
StatementType getType();
143143

144+
/**
145+
* Gets the {@link Schema} of this {@link PreparedStatement}.
146+
*
147+
* @return {@link Schema}.
148+
*/
149+
Schema getDataSetSchema();
150+
144151
@Override
145152
void close();
146153
}
@@ -171,6 +178,11 @@ public StatementType getType() {
171178
return schema.getFields().isEmpty() ? StatementType.UPDATE : StatementType.SELECT;
172179
}
173180

181+
@Override
182+
public Schema getDataSetSchema() {
183+
return preparedStatement.getResultSetSchema();
184+
}
185+
174186
@Override
175187
public void close() {
176188
preparedStatement.close(getOptions());

0 commit comments

Comments
 (0)