GH-62: Optimize Flight SQL JDBC Statement execution paths

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

@@ -19,11 +19,17 @@
 import static org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty.replaceSemiColons;
 
 import io.netty.util.concurrent.DefaultThreadFactory;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
+import java.util.IdentityHashMap;
 import java.util.Map;
 import java.util.Properties;
+import java.util.Set;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler;
@@ -35,6 +41,7 @@
 import org.apache.arrow.util.Preconditions;
 import org.apache.calcite.avatica.AvaticaConnection;
 import org.apache.calcite.avatica.AvaticaFactory;
+import org.apache.calcite.avatica.AvaticaStatement;
 import org.apache.calcite.avatica.DriverVersion;
 
 /** Connection to the Arrow Flight server. */
@@ -45,7 +52,10 @@ public final class ArrowFlightConnection extends AvaticaConnection {
   private final ArrowFlightConnectionConfigImpl config;
   private ExecutorService executorService;
   private int metadataResultSetCount;
+  // Track result sets and statements owned by this connection so close() can release them.
   private Map<Integer, ArrowFlightJdbcFlightStreamResultSet> metadataResultSetMap = new HashMap<>();
+  private final Set<AvaticaStatement> statementOwners =
+      Collections.newSetFromMap(new IdentityHashMap<>());
 
   /**
    * Creates a new {@link ArrowFlightConnection}.
@@ -142,12 +152,13 @@ private static ArrowFlightSqlClientHandler createNewClientHandler(
   void reset() throws SQLException {
     // Clean up any open Statements
     try {
-      AutoCloseables.close(statementMap.values());
+      AutoCloseables.close(getStatementsToClose());
     } catch (final Exception e) {
       throw AvaticaConnection.HELPER.createException(e.getMessage(), e);
     }
 
     statementMap.clear();
+    statementOwners.clear();
 
     // Reset Holdability
     this.setHoldability(this.metaData.getResultSetHoldability());
@@ -203,6 +214,28 @@ void onResultSetClose(Integer id) {
     metadataResultSetMap.remove(id);
   }
 
+  synchronized void registerStatementOwner(final AvaticaStatement statement) {
+    statementOwners.add(statement);
+  }
+
+  synchronized void unregisterStatementOwner(final AvaticaStatement statement) {
+    statementOwners.remove(statement);
+  }
+
+  private synchronized ArrayList<AutoCloseable> getStatementsToClose() {
+    final ArrayList<AutoCloseable> statements = new ArrayList<>(statementOwners);
+    final Set<Integer> ownedHandles = new HashSet<>();
+    for (AvaticaStatement statement : statementOwners) {
+      ownedHandles.add(statement.handle.id);
+    }
+    for (Map.Entry<Integer, AvaticaStatement> entry : statementMap.entrySet()) {
+      if (!ownedHandles.contains(entry.getKey())) {
+        statements.add(entry.getValue());
+      }
+    }
+    return statements;
+  }
+
   @Override
   public Properties getClientInfo() {
     final Properties copy = new Properties();
@@ -221,7 +254,7 @@ public void close() throws SQLException {
       topLevelException = e;
     }
     // copies of the collections are used to avoid concurrent modification problems
-    ArrayList<AutoCloseable> closeables = new ArrayList<>(statementMap.values());
+    ArrayList<AutoCloseable> closeables = getStatementsToClose();
     closeables.addAll(new ArrayList<>(metadataResultSetMap.values()));
     closeables.add(clientHandler);
     closeables.addAll(allocator.getChildAllocators());
@@ -257,4 +290,35 @@ BufferAllocator getBufferAllocator() {
   public ArrowFlightMetaImpl getMeta() {
     return (ArrowFlightMetaImpl) this.meta;
   }
+
+  @Override
+  public PreparedStatement prepareStatement(final String sql) throws SQLException {
+    checkOpen();
+    return prepareStatement(sql, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY);
+  }
+
+  @Override
+  public PreparedStatement prepareStatement(
+      final String sql, final int resultSetType, final int resultSetConcurrency)
+      throws SQLException {
+    checkOpen();
+    return prepareStatement(sql, resultSetType, resultSetConcurrency, getHoldability());
+  }
+
+  @Override
+  public PreparedStatement prepareStatement(
+      final String sql,
+      final int resultSetType,
+      final int resultSetConcurrency,
+      final int resultSetHoldability)
+      throws SQLException {
+    checkOpen();
+    return ArrowFlightPreparedStatement.builder(this)
+        .withQuery(sql)
+        .withGeneratedHandle()
+        .withResultSetType(resultSetType)
+        .withResultSetConcurrency(resultSetConcurrency)
+        .withResultSetHoldability(resultSetHoldability)
+        .build();
+  }
 }

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFactory.java

@@ -20,7 +20,6 @@
 import java.sql.SQLException;
 import java.util.Properties;
 import java.util.TimeZone;
-import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler;
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.calcite.avatica.AvaticaConnection;
 import org.apache.calcite.avatica.AvaticaFactory;
@@ -79,20 +78,20 @@ public ArrowFlightPreparedStatement newPreparedStatement(
       final Meta.Signature signature,
       final int resultType,
       final int resultSetConcurrency,
-      final int resultSetHoldability)
-      throws SQLException {
+      final int resultSetHoldability) {
     final ArrowFlightConnection flightConnection = (ArrowFlightConnection) connection;
-    ArrowFlightSqlClientHandler.PreparedStatement preparedStatement =
-        flightConnection.getMeta().getPreparedStatement(statementHandle);
+    final AvaticaStatement existingStatement =
+        flightConnection.statementMap.get(statementHandle.id);
+    if (existingStatement instanceof ArrowFlightPreparedStatement) {
+      return (ArrowFlightPreparedStatement) existingStatement;
+    }
+    if (existingStatement != null) {
+      throw new IllegalStateException(
+          "Unexpected statement type found for prepared statement handle: " + statementHandle);
+    }
 
-    return ArrowFlightPreparedStatement.newPreparedStatement(
-        flightConnection,
-        preparedStatement,
-        statementHandle,
-        signature,
-        resultType,
-        resultSetConcurrency,
-        resultSetHoldability);
+    throw new IllegalStateException(
+        "PreparedStatement was not pre-created for handle: " + statementHandle);
   }
 
   @Override

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

@@ -27,11 +27,14 @@
 import org.apache.arrow.driver.jdbc.client.CloseableEndpointStreamPair;
 import org.apache.arrow.driver.jdbc.utils.FlightEndpointDataQueue;
 import org.apache.arrow.driver.jdbc.utils.VectorSchemaRootTransformer;
+import org.apache.arrow.flight.CallStatus;
 import org.apache.arrow.flight.FlightInfo;
+import org.apache.arrow.flight.FlightRuntimeException;
 import org.apache.arrow.flight.FlightStream;
 import org.apache.arrow.util.AutoCloseables;
 import org.apache.arrow.vector.VectorSchemaRoot;
 import org.apache.arrow.vector.types.pojo.Schema;
+import org.apache.calcite.avatica.AvaticaConnection;
 import org.apache.calcite.avatica.AvaticaResultSet;
 import org.apache.calcite.avatica.AvaticaResultSetMetaData;
 import org.apache.calcite.avatica.AvaticaStatement;
@@ -67,7 +70,7 @@ public final class ArrowFlightJdbcFlightStreamResultSet
       throws SQLException {
     super(statement, state, signature, resultSetMetaData, timeZone, firstFrame);
     this.connection = (ArrowFlightConnection) statement.connection;
-    this.flightInfo = ((ArrowFlightInfoStatement) statement).executeFlightInfoQuery();
+    this.flightInfo = ((ArrowFlightMetaStatement) statement).executeFlightInfoQuery();
   }
 
   /** Private constructor for fromFlightInfo. */
@@ -106,7 +109,7 @@ static ArrowFlightJdbcFlightStreamResultSet fromFlightInfo(
     final TimeZone timeZone = TimeZone.getDefault();
     final QueryState state = new QueryState();
 
-    final Meta.Signature signature = ArrowFlightMetaImpl.newSignature(null, null, null, null);
+    final Meta.Signature signature = ArrowFlightMetaImpl.buildDefaultSignature();
 
     final AvaticaResultSetMetaData resultSetMetaData =
         new AvaticaResultSetMetaData(null, null, signature);
@@ -192,23 +195,40 @@ public boolean next() throws SQLException {
         return true;
       }
 
-      if (currentEndpointData != null) {
-        currentEndpointData.getStream().getRoot().clear();
-        if (currentEndpointData.getStream().next()) {
-          populateDataForCurrentFlightStream();
-          continue;
+      try {
+        if (currentEndpointData != null) {
+          currentEndpointData.getStream().getRoot().clear();
+          if (currentEndpointData.getStream().next()) {
+            populateDataForCurrentFlightStream();
+            continue;
+          }
+
+          flightEndpointDataQueue.enqueue(currentEndpointData);
         }
 
-        flightEndpointDataQueue.enqueue(currentEndpointData);
+        currentEndpointData = getNextEndpointStream(false);
+      } catch (final FlightRuntimeException e) {
+        // A concurrent statement.cancel() (or close) cancels in-flight FlightStreams,
+        // which surface here as CANCELLED. Normalize to Avatica's "Statement canceled".
+        if (flightEndpointDataQueue.isClosed()
+            && e.status().code() == CallStatus.CANCELLED.code()) {
+          throw AvaticaConnection.HELPER.createException("Statement canceled");
+        }
+        throw e;
       }
 
-      currentEndpointData = getNextEndpointStream(false);
-
       if (currentEndpointData != null) {
         populateDataForCurrentFlightStream();
         continue;
       }
 
+      // No more data. If the queue was closed concurrently (e.g. statement.cancel()
+      // racing with the reader past super.next()), surface as "Statement canceled"
+      // to match Avatica's cancellation semantics.
+      if (flightEndpointDataQueue.isClosed()) {
+        throw AvaticaConnection.HELPER.createException("Statement canceled");
+      }
+
       if (statement != null && statement.isCloseOnCompletion()) {
         statement.close();
       }

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

@@ -73,7 +73,7 @@ public static ArrowFlightJdbcVectorSchemaRootResultSet fromVectorSchemaRoot(
     final TimeZone timeZone = TimeZone.getDefault();
     final QueryState state = new QueryState();
 
-    final Meta.Signature signature = ArrowFlightMetaImpl.newSignature(null, null, null, null);
+    final Meta.Signature signature = ArrowFlightMetaImpl.buildDefaultSignature();
 
     final AvaticaResultSetMetaData resultSetMetaData =
         new AvaticaResultSetMetaData(null, null, signature);