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,18 @@
 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.sql.Statement;
 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;
@@ -33,8 +40,10 @@
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.util.AutoCloseables;
 import org.apache.arrow.util.Preconditions;
+import org.apache.arrow.util.VisibleForTesting;
 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 +54,11 @@ public final class ArrowFlightConnection extends AvaticaConnection {
   private final ArrowFlightConnectionConfigImpl config;
   private ExecutorService executorService;
   private int metadataResultSetCount;
+  // Track metadata result sets so connection close can release them.
   private Map<Integer, ArrowFlightJdbcFlightStreamResultSet> metadataResultSetMap = new HashMap<>();
+  // Track original statement owners when execution replaces their statement-map entries.
+  private final Set<AvaticaStatement> statementOwners =
+      Collections.newSetFromMap(new IdentityHashMap<>());
 
   /**
    * Creates a new {@link ArrowFlightConnection}.
@@ -142,7 +155,7 @@ 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);
     }
@@ -203,6 +216,33 @@ 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);
+  }
+
+  @VisibleForTesting
+  synchronized boolean isStatementOwner(final Statement statement) {
+    return statementOwners.contains(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 +261,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 +297,37 @@ BufferAllocator getBufferAllocator() {
   public ArrowFlightMetaImpl getMeta() {
     return (ArrowFlightMetaImpl) this.meta;
   }
+
+  @Override
+  public PreparedStatement prepareStatement(final String sql) throws SQLException {
+    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 {
+    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();
+    try {
+      return ArrowFlightPreparedStatement.builder(this)
+          .withQuery(sql)
+          .withGeneratedHandle()
+          .withResultSetType(resultSetType)
+          .withResultSetConcurrency(resultSetConcurrency)
+          .withResultSetHoldability(resultSetHoldability)
+          .build();
+    } catch (final RuntimeException e) {
+      throw HELPER.createException("while preparing SQL:" + sql, e);
+    }
+  }
 }

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;
@@ -82,17 +81,18 @@ public ArrowFlightPreparedStatement newPreparedStatement(
       final int resultSetHoldability)
       throws SQLException {
     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

@@ -21,17 +21,22 @@
 import java.sql.ResultSet;
 import java.sql.ResultSetMetaData;
 import java.sql.SQLException;
+import java.util.Collection;
+import java.util.Iterator;
 import java.util.Optional;
 import java.util.TimeZone;
 import java.util.concurrent.TimeUnit;
 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;
@@ -49,6 +54,7 @@ public final class ArrowFlightJdbcFlightStreamResultSet
   private final FlightInfo flightInfo;
   private CloseableEndpointStreamPair currentEndpointData;
   private FlightEndpointDataQueue flightEndpointDataQueue;
+  private volatile boolean cancelled;
 
   private VectorSchemaRootTransformer transformer;
   private VectorSchemaRoot currentVectorSchemaRoot;
@@ -106,7 +112,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);
@@ -143,7 +149,8 @@ protected AvaticaResultSet execute() throws SQLException {
 
   private void populateData() throws SQLException {
     loadNewQueue();
-    flightEndpointDataQueue.enqueue(connection.getClientHandler().getStreams(flightInfo));
+    enqueueEndpointData(
+        flightEndpointDataQueue, connection.getClientHandler().getStreams(flightInfo));
     loadNewFlightStream();
 
     // Ownership of the root will be passed onto the cursor.
@@ -192,23 +199,38 @@ 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;
+          }
+
+          enqueueEndpointData(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 ((cancelled || flightEndpointDataQueue.isClosed())
+            && e.status().code() == CallStatus.CANCELLED.code()) {
+          throw AvaticaConnection.HELPER.createException("Statement canceled");
+        }
+        throw e;
       }
 
-      currentEndpointData = getNextEndpointStream(false);
-
       if (currentEndpointData != null) {
         populateDataForCurrentFlightStream();
         continue;
       }
 
+      // Cancellation can wake the reader before queue cleanup completes.
+      if (cancelled || flightEndpointDataQueue.isClosed()) {
+        throw AvaticaConnection.HELPER.createException("Statement canceled");
+      }
+
       if (statement != null && statement.isCloseOnCompletion()) {
         statement.close();
       }
@@ -219,6 +241,7 @@ public boolean next() throws SQLException {
 
   @Override
   protected void cancel() {
+    cancelled = true;
     super.cancel();
     final CloseableEndpointStreamPair currentEndpoint = this.currentEndpointData;
     if (currentEndpoint != null) {
@@ -259,13 +282,61 @@ public synchronized void close() {
 
   private CloseableEndpointStreamPair getNextEndpointStream(final boolean canTimeout)
       throws SQLException {
-    if (canTimeout) {
-      final int statementTimeout = statement != null ? statement.getQueryTimeout() : 0;
-      return statementTimeout != 0
-          ? flightEndpointDataQueue.next(statementTimeout, TimeUnit.SECONDS)
-          : flightEndpointDataQueue.next();
-    } else {
-      return flightEndpointDataQueue.next();
+    try {
+      if (canTimeout) {
+        final int statementTimeout = statement != null ? statement.getQueryTimeout() : 0;
+        return statementTimeout != 0
+            ? flightEndpointDataQueue.next(statementTimeout, TimeUnit.SECONDS)
+            : flightEndpointDataQueue.next();
+      } else {
+        return flightEndpointDataQueue.next();
+      }
+    } catch (final IllegalStateException e) {
+      throw normalizeClosedQueueException(flightEndpointDataQueue, e);
+    }
+  }
+
+  static void enqueueEndpointData(
+      final FlightEndpointDataQueue queue,
+      final Collection<CloseableEndpointStreamPair> endpointRequests)
+      throws SQLException {
+    final Iterator<CloseableEndpointStreamPair> endpointIterator = endpointRequests.iterator();
+    while (endpointIterator.hasNext()) {
+      final CloseableEndpointStreamPair endpointRequest = endpointIterator.next();
+      try {
+        queue.enqueue(endpointRequest);
+      } catch (final IllegalStateException e) {
+        final SQLException exception = normalizeClosedQueueException(queue, e);
+        closeRejectedEndpoint(endpointRequest, exception);
+        endpointIterator.forEachRemaining(endpoint -> closeRejectedEndpoint(endpoint, exception));
+        throw exception;
+      }
+    }
+  }
+
+  private void enqueueEndpointData(final CloseableEndpointStreamPair endpointRequest)
+      throws SQLException {
+    try {
+      flightEndpointDataQueue.enqueue(endpointRequest);
+    } catch (final IllegalStateException e) {
+      throw normalizeClosedQueueException(flightEndpointDataQueue, e);
+    }
+  }
+
+  private static void closeRejectedEndpoint(
+      final CloseableEndpointStreamPair endpoint, final SQLException exception) {
+    try {
+      endpoint.close();
+    } catch (final Exception closeException) {
+      exception.addSuppressed(closeException);
+    }
+  }
+
+  private static SQLException normalizeClosedQueueException(
+      final FlightEndpointDataQueue queue, final IllegalStateException exception) {
+    if (queue.isClosed()) {
+      return AvaticaConnection.HELPER.createException("Statement canceled");
     }
+    throw exception;
   }
 }

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);