mykolapavlov-dremio · Oct 27, 2026
@@ -31,12 +31,16 @@ public final class SyncPutListener implements FlightClient.PutListener, AutoClos
 
   private final LinkedBlockingQueue<Object> queue;
   private final CompletableFuture<Void> 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

@@ -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

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

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

@@ -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);
     }
   }
 }
@@ -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 {

@@ -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);
+  }
+}
@@ -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<Long> 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.
      *