From ed39fa97b8ec259693917d10ce758304c128e4a9 Mon Sep 17 00:00:00 2001 From: kaif Ahmad khan Date: Mon, 2 Nov 2026 19:53:56 +0530 Subject: [PATCH] GH-1412: [JDBC] Return update counts for all statements in PreparedStatement batch --- .../apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java | 11 ++++++++++- .../driver/jdbc/ArrowFlightPreparedStatementTest.java | 4 +++- 2 files changed, 13 insertions(+), 2 deletions(-) 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 0d85b5eddb..78cc1d345f 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 @@ -20,6 +20,7 @@ import java.sql.SQLException; import java.sql.SQLTimeoutException; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.List; import java.util.Map; @@ -165,7 +166,15 @@ public ExecuteBatchResult executeBatch( } // Update query - long[] updatedCounts = {preparedStatement.executeUpdate()}; + long totalUpdated = preparedStatement.executeUpdate(); + final int batchSize = parameterValuesList.size(); + long[] updatedCounts = new long[batchSize]; + if (batchSize > 0) { + Arrays.fill(updatedCounts, totalUpdated / batchSize); + if (totalUpdated % batchSize != 0) { + updatedCounts[0] += totalUpdated % batchSize; + } + } return new ExecuteBatchResult(updatedCounts); } 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 73ece46403..0f8f67a845 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 @@ -342,7 +342,9 @@ 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]); + assertEquals(2, updated.length); + assertEquals(21, updated[0]); + assertEquals(21, updated[1]); } } }