Skip to content

GH-62: Optimize Flight SQL JDBC Statement execution paths - #1090

Open
xborder wants to merge 25 commits into
apache:mainfrom
xborder:optimize-queries-3
Open

xborder wants to merge 25 commits into
apache:mainfrom
xborder:optimize-queries-3

Conversation

@xborder

@xborder xborder commented Mar 24, 2026 •

Copy link
Copy Markdown
Contributor

What Changed

  • This change simplifies the Flight SQL JDBC request flow for Statement.executeQuery/executeUpdate avoiding the PreparedStatement request flow
  • Preserved the current requirement that Statement.execute(String) still flows through PreparedStatement
    • Right now, Arrow Flight does not offer a good solution for this use case. There is no way for the client to disambiguate between a CommandStatementUpdate or CommandStatemenQuery without doing a server call
  • Moved prepared-statement-specific creation and execution concerns into ArrowFlightPreparedStatement
  • Similarly, created ArrowFlightStatement with specific logic to handle Statements
  • Removed statement/prepared-statement construction details and logic from ArrowFlightMetaImpl. ArrowFlightMetaImpl was left it focused on orchestration

Closes #62.

This PR was assisted by AI

@xborder
xborder marked this pull request as ready for review March 24, 2026 22:43
@github-actions

This comment has been minimized.

@jbonofre jbonofre added this to the 20.0.0 milestone Apr 13, 2026
.withQuery(query)
.withExistingStatement(this)
.build()
.prepareAndExecute(callback);

ghost Apr 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Perhaps a nitpick, but if we run multiple queries with the same statement using execute we will only close one of the handles in the server.

Code for repro:

    try (Connection connection = DriverManager.getConnection(jdbcUrl, properties);
         Statement statement = connection.createStatement()) {

      for (int i = 1; i <= n; i++) {
        String sql = "SELECT " + i;
        log("Executing: " + sql);

        boolean isResultSet = statement.execute(sql);
        if (isResultSet) {
          try (ResultSet rs = statement.getResultSet()) {
            // consume result set
          }
        } else {
          System.out.println("Updated rows: " + statement.getUpdateCount());
        }
      }
    }

A new prepared statement will be created on each query, but only one ClosePreparedStatement is sent to the server. I believe most Flight SQL servers are stateless and this won't be problematic in that scenario, but in case the server is holding some resources associated with the handle(s) of the prepared statement(s) that aren't closed there will be a leak.

Could we send the close after execution?

ghost Apr 27, 2026 •

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch. This issue existed before the refactor. I've created #1129 to be addressed separately

@jbonofre

ghost commented Sep 28, 2026

Copy link
Copy Markdown
Member

@xborder thanks for the update! I will do the review.

@xborder
xborder force-pushed the optimize-queries-3 branch from 146cd30 to b7bd7c5 Compare October 2, 2026 13:30
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 =

ghost Nov 4, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder, could we separate this change out into its own PR so it shows up more clearly in the changelog/history and is easier to review?

ghost Nov 4, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Albeit if it can't be separated then that's fine

ghost Nov 10, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is to ensure that we close the result sets for the original statements when the connection is closed or reset, even if their entries in Avatica’s statement map have been replaced.

final PrepareCallback callback)
throws SQLException {

return ArrowFlightPreparedStatement.builder(getConnection())

ghost Nov 10, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Calling Statement.execute(String) again on the same statement leaks the previous server-side prepared statement.

On the second call, statementMap holds the internal ArrowFlightPreparedStatement created by the first call, so ArrowFlightMetaImpl.prepareAndExecute routes here. We build a new prepared statement under the same handle, and its constructor replaces the statementMap entry, but this.preparedStatement is never closed. No ClosePreparedStatement is sent for it. A loop of N execute() calls leaves N-1 prepared statements open on the server, and stmt.close() only closes the last one.

The old statementHandlePreparedStatementMap.put(...) dropped the previous entry the same way, so this is not a regression. Since this method is new, could we close the previous one once the replacement is built?

final ArrowFlightPreparedStatement replacement =
    ArrowFlightPreparedStatement.builder(getConnection())
        .withQuery(query)
        .withExistingStatement(this)
        .build();
closePreparedResources();
return replacement.prepareAndExecuteInternal(callback);

A test that calls execute() a few times on one statement and checks how many prepared statements the mock producer closed would cover this.

ghost Nov 10, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I had already identified this issue. It already exists on main. I suggest addressing it in a follow-up PR after this one merges.

ghost Nov 10, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That works for me. Thanks!

} catch (final FlightRuntimeException e) {
throw new SQLException(e.getMessage(), e);
}
final Signature signature =

ghost Nov 10, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If anything throws between prepare() and the end of build(), the server side prepared statement is never cloud. One concrete case: a parameter schema with a ListView, LargeListView or RunEndEncoded field makes ConvertUtils.convertArrowFieldsToAvaticaParameters throw UnsupportedOperationException from buildSignature.

The old prepareForHandle had the same gap. Now that build() owns the whole lifecycle, could we close it on failure?

final ArrowFlightPreparedStatement statement;
try {
  final Signature signature =
      ArrowFlightMetaImpl.buildSignature(
          query,
          preparedStatement.getDataSetSchema(),
          preparedStatement.getParameterSchema(),
          preparedStatement.isUpdate());
  statement =
      new ArrowFlightPreparedStatement(
          connection,
          generateHandle ? null : handle,
          signature,
          preparedStatement,
          resultSetType,
          resultSetConcurrency,
          resultSetHoldability);
} catch (final Exception e) {
  try {
    preparedStatement.close();
  } catch (final Exception closeException) {
    e.addSuppressed(closeException);
  }
  throw e;
}

ghost Nov 10, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Remove redundant API visibility and interface identity tests, retain the checked prepareStatement implementation, and clarify connection ownership comments.

ghost left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the iterations @xborder!

Great work!

JB Onofré and others added 3 commits November 10, 2026 14:56
testShouldInterruptFlightStreamsIfQueryIsCancelledMidQuerying called
statement.cancel() right after starting the worker thread. It relied on
the prepare round trip in executeQuery to make the cancel land before a
result set existed, so only Avatica's cancel flag was set and next()
threw "Statement canceled".

Statement.executeQuery no longer prepares first, so the cancel can now
land while the result set is executing and surface as "Error while
executing SQL ...: exception while executing query: Statement
canceled", which failed the RC Binaries job.

Wait until executeQuery returns before cancelling, and replace the
300 ms sleep with a latch.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement PRs that add or improve features.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Java][FlightSQL] Arrow Flight SQL JDBC Driver makes ActionCreatePreparedStatementRequest even for Statement.executeQuery() and Statement.exeute().

4 participants