From 5b41198ba651d63131fb79b019309224d7cb935c Mon Sep 17 00:00:00 2001 From: Mykola Pavlov Date: Tue, 27 Oct 2026 17:44:03 +0100 Subject: [PATCH] DX-105241: Support Flight SQL batch updates Collect all prepared-statement update results and translate partial batch failures to BatchUpdateException. Close response buffers on success and failure to prevent allocator leaks. --- .../apache/arrow/flight/SyncPutListener.java | 43 +++++++++--- .../driver/jdbc/ArrowFlightMetaImpl.java | 12 +++- .../jdbc/ArrowFlightPreparedStatement.java | 18 +++++ .../client/ArrowFlightSqlClientHandler.java | 12 ++++ .../ArrowFlightPreparedStatementTest.java | 5 +- .../jdbc/utils/MockFlightSqlProducer.java | 20 +++--- .../sql/FlightSqlBatchUpdateException.java | 33 +++++++++ .../arrow/flight/sql/FlightSqlClient.java | 69 ++++++++++++++++--- 8 files changed, 178 insertions(+), 34 deletions(-) create mode 100644 flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlBatchUpdateException.java diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/SyncPutListener.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/SyncPutListener.java index 0dd460f1fd..2190489fd6 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/SyncPutListener.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/SyncPutListener.java @@ -31,12 +31,16 @@ public final class SyncPutListener implements FlightClient.PutListener, AutoClos private final LinkedBlockingQueue queue; private final CompletableFuture completed; + private final Object closeLock; + private boolean closed; private static final Object DONE = new Object(); private static final Object DONE_WITH_EXCEPTION = new Object(); public SyncPutListener() { queue = new LinkedBlockingQueue<>(); completed = new CompletableFuture<>(); + closeLock = new Object(); + closed = false; } private PutResult unwrap(Object queueItem) throws InterruptedException, ExecutionException { @@ -90,31 +94,48 @@ public void getResult() { @Override public void onNext(PutResult val) { - final ArrowBuf metadata = val.getApplicationMetadata(); - metadata.getReferenceManager().retain(); - queue.add(PutResult.metadata(metadata)); + synchronized (closeLock) { + if (closed) { + return; + } + + final ArrowBuf metadata = val.getApplicationMetadata(); + metadata.getReferenceManager().retain(); + queue.add(PutResult.metadata(metadata)); + } } @Override public void onError(Throwable t) { completed.completeExceptionally(StatusUtils.fromThrowable(t)); - queue.add(DONE_WITH_EXCEPTION); + synchronized (closeLock) { + if (!closed) { + queue.add(DONE_WITH_EXCEPTION); + } + } } @Override public void onCompleted() { completed.complete(null); - queue.add(DONE); + synchronized (closeLock) { + if (!closed) { + queue.add(DONE); + } + } } @Override public void close() { - queue.forEach( - o -> { - if (o instanceof PutResult) { - ((PutResult) o).close(); - } - }); + synchronized (closeLock) { + closed = true; + Object item; + while ((item = queue.poll()) != null) { + if (item instanceof PutResult) { + ((PutResult) item).close(); + } + } + } } @Override diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java index 64529b50c8..e697ecb124 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java @@ -156,9 +156,15 @@ public ExecuteBatchResult executeBatch( binder.bind(parameterValuesList.get(i), i); } - // Update query - long[] updatedCounts = {preparedStatement.executeUpdate()}; - return new ExecuteBatchResult(updatedCounts); + return new ExecuteBatchResult(preparedStatement.executeUpdateBatch()); + } + + static int[] toIntArray(long[] updateCounts) { + int[] result = new int[updateCounts.length]; + for (int i = 0; i < updateCounts.length; i++) { + result[i] = Math.toIntExact(updateCounts[i]); + } + return result; } @Override diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java index d7af6902f4..36057bf93f 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java @@ -16,12 +16,15 @@ */ package org.apache.arrow.driver.jdbc; +import java.sql.BatchUpdateException; import java.sql.PreparedStatement; import java.sql.SQLException; import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler; import org.apache.arrow.flight.FlightInfo; +import org.apache.arrow.flight.sql.FlightSqlBatchUpdateException; import org.apache.arrow.util.Preconditions; import org.apache.calcite.avatica.AvaticaPreparedStatement; +import org.apache.calcite.avatica.Meta.ExecuteBatchResult; import org.apache.calcite.avatica.Meta.Signature; import org.apache.calcite.avatica.Meta.StatementHandle; @@ -74,6 +77,21 @@ public synchronized void close() throws SQLException { super.close(); } + @Override + public long[] executeLargeBatch() throws SQLException { + checkOpen(); + try { + final ExecuteBatchResult result = + getConnection().getMeta().executeBatch(handle, getParameterValueBatch()); + return result.updateCounts; + } catch (FlightSqlBatchUpdateException e) { + throw new BatchUpdateException( + e.getMessage(), null, 0, ArrowFlightMetaImpl.toIntArray(e.getUpdateCounts()), e); + } finally { + parameterValueBatch.clear(); + } + } + @Override public FlightInfo executeFlightInfoQuery() throws SQLException { return preparedStatement.executeQuery(); diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java index f0ea284239..9e029dce66 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java @@ -367,6 +367,13 @@ public interface PreparedStatement extends AutoCloseable { */ long executeUpdate(); + /** + * Executes a {@link StatementType#UPDATE} query for all bound parameter rows. + * + * @return the update count for each bound parameter row. + */ + long[] executeUpdateBatch(); + /** * Gets the {@link StatementType} of this {@link PreparedStatement}. * @@ -454,6 +461,11 @@ public long executeUpdate() { return preparedStatement.executeUpdate(getOptions()); } + @Override + public long[] executeUpdateBatch() { + return preparedStatement.executeUpdateBatch(getOptions()); + } + @Override public StatementType getType() { final Schema schema = preparedStatement.getResultSetSchema(); diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java index 0369c3a162..cb74482cef 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java @@ -19,6 +19,7 @@ import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.jupiter.api.Assertions.assertAll; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -182,7 +183,7 @@ public void testReturnColumnCount() throws SQLException { @Test public void testUpdateQuery() throws SQLException { String query = "Fake update"; - PRODUCER.addUpdateQuery(query, /*updatedRows*/ 42); + PRODUCER.addUpdateQuery(query, /*updatedRows*/ 42, 42); try (final PreparedStatement stmt = connection.prepareStatement(query)) { int updated = stmt.executeUpdate(); assertEquals(42, updated); @@ -248,7 +249,7 @@ public void testUpdateQueryWithBatchedParameters() throws SQLException { stmt.setArray(2, connection.createArrayOf("INTEGER", new Integer[] {0, -1, 100000})); stmt.addBatch(); int[] updated = stmt.executeBatch(); - assertEquals(42, updated[0]); + assertArrayEquals(new int[] {42, 42}, updated); } } } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java index 45c2a96404..0866857e46 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java @@ -156,18 +156,22 @@ public void addSelectQuery( * Registers a new {@link StatementType#UPDATE} SQL query. * * @param sqlCommand the SQL command. - * @param updatedRows the number of rows affected. + * @param updatedRows the number of rows affected for each update result. */ - public void addUpdateQuery(final String sqlCommand, final long updatedRows) { + public void addUpdateQuery(final String sqlCommand, final long... updatedRows) { addUpdateQuery( sqlCommand, (flightStream, putResultStreamListener) -> { - final DoPutUpdateResult result = - DoPutUpdateResult.newBuilder().setRecordCount(updatedRows).build(); - try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); - final ArrowBuf buffer = allocator.buffer(result.getSerializedSize())) { - buffer.writeBytes(result.toByteArray()); - putResultStreamListener.onNext(PutResult.metadata(buffer)); + try { + for (long updatedRowCount : updatedRows) { + final DoPutUpdateResult result = + DoPutUpdateResult.newBuilder().setRecordCount(updatedRowCount).build(); + try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + final ArrowBuf buffer = allocator.buffer(result.getSerializedSize())) { + buffer.writeBytes(result.toByteArray()); + putResultStreamListener.onNext(PutResult.metadata(buffer)); + } + } } catch (final Throwable throwable) { putResultStreamListener.onError(throwable); } finally { diff --git a/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlBatchUpdateException.java b/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlBatchUpdateException.java new file mode 100644 index 0000000000..6989447d13 --- /dev/null +++ b/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlBatchUpdateException.java @@ -0,0 +1,33 @@ +/* + * 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.flight.sql; + +import java.util.Arrays; + +/** An error from a prepared-statement batch together with counts already received. */ +public final class FlightSqlBatchUpdateException extends RuntimeException { + private final long[] updateCounts; + + public FlightSqlBatchUpdateException(Throwable cause, long[] updateCounts) { + super(cause.getMessage(), cause); + this.updateCounts = Arrays.copyOf(updateCounts, updateCounts.length); + } + + public long[] getUpdateCounts() { + return Arrays.copyOf(updateCounts, updateCounts.length); + } +} diff --git a/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java b/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java index 623f9311e8..69dd5229d5 100644 --- a/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java +++ b/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java @@ -51,6 +51,7 @@ import java.io.ByteArrayInputStream; import java.io.IOException; import java.nio.channels.Channels; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.Iterator; @@ -1364,15 +1365,18 @@ public FlightInfo execute(final CallOption... options) { private SyncPutListener putParameters(FlightDescriptor descriptor, CallOption... options) { final SyncPutListener putListener = new SyncPutListener(); + try { + FlightClient.ClientStreamListener listener = + client.startPut(descriptor, parameterBindingRoot, putListener, options); - FlightClient.ClientStreamListener listener = - client.startPut(descriptor, parameterBindingRoot, putListener, options); - - listener.putNext(); - listener.completed(); - listener.getResult(); + listener.putNext(); + listener.completed(); - return putListener; + return putListener; + } catch (final RuntimeException | Error e) { + putListener.close(); + throw e; + } } /** @@ -1401,9 +1405,7 @@ public long executeUpdate(final CallOption... options) { .build()) .toByteArray()); setParameters(parameterBindingRoot == null ? VectorSchemaRoot.of() : parameterBindingRoot); - SyncPutListener putListener = putParameters(descriptor, options); - - try { + try (SyncPutListener putListener = putParameters(descriptor, options)) { final PutResult read = putListener.read(); try (final ArrowBuf metadata = read.getApplicationMetadata()) { final DoPutUpdateResult doPutUpdateResult = @@ -1419,6 +1421,53 @@ public long executeUpdate(final CallOption... options) { } } + /** + * Executes the prepared statement update on the server for all bound parameter rows. + * + * @param options RPC-layer hints for this call. + * @return the update count for each bound parameter row. + */ + public long[] executeUpdateBatch(final CallOption... options) { + checkOpen(); + final FlightDescriptor descriptor = + FlightDescriptor.command( + Any.pack( + CommandPreparedStatementUpdate.newBuilder() + .setPreparedStatementHandle( + preparedStatementResult.getPreparedStatementHandle()) + .build()) + .toByteArray()); + setParameters(parameterBindingRoot == null ? VectorSchemaRoot.of() : parameterBindingRoot); + + final List updateCounts = new ArrayList<>(); + try (SyncPutListener putListener = putParameters(descriptor, options)) { + while (true) { + final PutResult result = putListener.read(); + if (result == null) { + break; + } + try (result) { + final ArrowBuf metadata = result.getApplicationMetadata(); + final DoPutUpdateResult doPutUpdateResult = + DoPutUpdateResult.parseFrom(metadata.nioBuffer()); + updateCounts.add(doPutUpdateResult.getRecordCount()); + } + } + return updateCounts.stream().mapToLong(Long::longValue).toArray(); + } catch (final InterruptedException e) { + throw new FlightSqlBatchUpdateException( + CallStatus.CANCELLED.withCause(e).toRuntimeException(), + updateCounts.stream().mapToLong(Long::longValue).toArray()); + } catch (final ExecutionException e) { + throw new FlightSqlBatchUpdateException( + e.getCause(), updateCounts.stream().mapToLong(Long::longValue).toArray()); + } catch (final InvalidProtocolBufferException e) { + throw new FlightSqlBatchUpdateException( + CallStatus.INTERNAL.withCause(e).toRuntimeException(), + updateCounts.stream().mapToLong(Long::longValue).toArray()); + } + } + /** * Closes the client. *