[7/7] Enable insert batching by default

NEXT_CHANGELOG.md

@@ -3,6 +3,7 @@
 ## [Unreleased]
 
 ### Added
+- Added native server-side batching (supported on DBR 18.2 and later) for parameterized `INSERT INTO ... VALUES (?, ...)` statements executed with `PreparedStatement`. The driver sends all parameter sets in one request instead of rewriting INSERT statements on the client. Each batch is limited to 10,000 parameters or 1 MB of parameter data, whichever limit is reached first. Native batching is enabled by default; set `EnableNativeBatching=0` to use legacy batching. `EnableBatchedInserts` now defaults to `1`, enabling client-side batching for compute running DBR versions earlier than 18.2. The `BINARY` data type is not supported in native batching; to batch binary values, set `EnableNativeBatching=0` and `supportManyParameters=1` to use legacy client-side batching.
 
 ### Updated
 - `DatabaseMetaData.getColumns(...)` with a `null` catalog now issues a single `SHOW COLUMNS IN ALL CATALOGS` statement (consistent with `getSchemas`/`getTables`) instead of enumerating every catalog and issuing a per-catalog `SHOW COLUMNS`. Older DBR versions that do not support the syntax transparently fall back to the previous enumerate-and-fan-out behavior.

docs/JDBC_METHOD_INVENTORY.md

@@ -700,7 +700,7 @@ OUT/INOUT parameters, named parameters, and return-value syntax (`{? = call ...}
 | `updatesAreDetected(int)` | YES | - | No | RARE | Returns false |
 | `deletesAreDetected(int)` | YES | - | No | RARE | Returns false |
 | `insertsAreDetected(int)` | YES | - | No | RARE | Returns false |
-| `supportsBatchUpdates()` | YES | - | No | RARE | Returns false |
+| `supportsBatchUpdates()` | YES | - | No | RARE | Returns true when native or legacy batching is enabled |
 | `getUDTs(String, String, String, int[])` | YES | - | No | RARE | Returns empty ResultSet |
 | `getConnection()` | YES | - | No | OCCASIONAL | Returns parent connection |
 | `supportsSavepoints()` | YES | - | No | RARE | Returns false |

src/main/java/com/databricks/jdbc/api/impl/BatchParameterSet.java

@@ -3,7 +3,9 @@
 import java.sql.Date;
 import java.sql.Time;
 import java.sql.Timestamp;
+import java.util.Collections;
 import java.util.Comparator;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -12,16 +14,21 @@
 /**
  * Immutable, position-ordered snapshot of one prepared-statement parameter set.
  *
- * <p>This model normalizes JDBC's one-based parameter indexes to zero-based wire ordinals. It does
- * not validate parameter completeness, index continuity, or consistency with other parameter sets;
- * those validations remain the backend's responsibility.
+ * <p>This model preserves JDBC's one-based parameter indexes. Transport adapters are responsible
+ * for converting them to protocol-specific wire ordinals. It does not validate parameter
+ * completeness, index continuity, or consistency with other parameter sets; those validations
+ * remain the backend's responsibility.
  */
 public final class BatchParameterSet {
 
   private final List<ImmutableSqlParameter> parameters;
+  private final Map<Integer, ImmutableSqlParameter> parameterBindings;
 
   private BatchParameterSet(List<ImmutableSqlParameter> parameters) {
     this.parameters = List.copyOf(parameters);
+    Map<Integer, ImmutableSqlParameter> bindings = new LinkedHashMap<>();
+    this.parameters.forEach(parameter -> bindings.put(parameter.cardinal(), parameter));
+    this.parameterBindings = Collections.unmodifiableMap(bindings);
   }
 
   public static BatchParameterSet from(Map<Integer, ImmutableSqlParameter> parameterBindings) {
@@ -38,6 +45,10 @@ public List<ImmutableSqlParameter> getParameters() {
     return parameters;
   }
 
+  public Map<Integer, ImmutableSqlParameter> getParameterBindings() {
+    return parameterBindings;
+  }
+
   public int size() {
     return parameters.size();
   }
@@ -50,7 +61,7 @@ private static ImmutableSqlParameter snapshotParameter(
       Map.Entry<Integer, ImmutableSqlParameter> entry) {
     ImmutableSqlParameter parameter = entry.getValue();
     return ImmutableSqlParameter.builder()
-        .cardinal(entry.getKey() - 1)
+        .cardinal(entry.getKey())
         .type(parameter.type())
         .value(snapshotValue(parameter.value()))
         .build();

src/main/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaData.java

@@ -1279,9 +1279,10 @@ public boolean insertsAreDetected(int type) throws SQLException {
   public boolean supportsBatchUpdates() throws SQLException {
     LOGGER.debug("public boolean supportsBatchUpdates()");
     throwExceptionIfConnectionIsClosed();
-    // Advertise batch support only when the multi-row INSERT optimization is enabled, so
-    // batch-aware clients use executeBatch() instead of one executeUpdate() per row.
-    return session.getConnectionContext().isBatchedInsertsEnabled();
+    // Advertise batch support when either native parameter batching or the legacy multi-row
+    // optimization is enabled, so batch-aware clients use executeBatch().
+    return session.getConnectionContext().isNativeBatchingEnabled()
+        || session.getConnectionContext().isBatchedInsertsEnabled();
   }
 
   @Override

src/main/java/com/databricks/jdbc/api/impl/DatabricksPreparedStatement.java

@@ -33,7 +33,7 @@ public class DatabricksPreparedStatement extends DatabricksStatement implements
       JdbcLoggerFactory.getLogger(DatabricksPreparedStatement.class);
   private final String sql;
   private DatabricksParameterMetaData databricksParameterMetaData;
-  private List<DatabricksParameterMetaData> databricksBatchParameterMetaData;
+  private List<BatchParameterSet> batchParameterSets;
   private final boolean interpolateParameters;
   private final int CHUNK_SIZE = 8192;
 
@@ -43,7 +43,7 @@ public DatabricksPreparedStatement(DatabricksConnection connection, String sql)
     this.sql = sql;
     this.interpolateParameters = connection.getConnectionContext().supportManyParameters();
     this.databricksParameterMetaData = new DatabricksParameterMetaData(sql);
-    this.databricksBatchParameterMetaData = new ArrayList<>();
+    this.batchParameterSets = new ArrayList<>();
     // Cache whether this statement should return a ResultSet (based on SQL and config)
     this.shouldReturnResultSet = shouldReturnResultSetWithConfig(sql);
   }
@@ -58,7 +58,7 @@ public DatabricksPreparedStatement(DatabricksConnection connection, String sql)
     this.sql = sql;
     this.interpolateParameters = interpolateParameters;
     this.databricksParameterMetaData = databricksParameterMetaData;
-    this.databricksBatchParameterMetaData = new ArrayList<>();
+    this.batchParameterSets = new ArrayList<>();
     // Cache whether this statement should return a ResultSet (based on SQL and config)
     this.shouldReturnResultSet = shouldReturnResultSetWithConfig(sql);
   }
@@ -94,7 +94,7 @@ public int executeUpdate() throws SQLException {
   }
 
   @Override
-  public int[] executeBatch() throws DatabricksBatchUpdateException {
+  public int[] executeBatch() throws SQLException {
     LOGGER.debug("public int executeBatch()");
     long[] largeUpdateCount = executeLargeBatch();
     int[] updateCount = new int[largeUpdateCount.length];
@@ -107,10 +107,10 @@ public int[] executeBatch() throws DatabricksBatchUpdateException {
   }
 
   @Override
-  public long[] executeLargeBatch() throws DatabricksBatchUpdateException {
+  public long[] executeLargeBatch() throws SQLException {
     LOGGER.debug("public long executeLargeBatch()");
 
-    if (databricksBatchParameterMetaData.isEmpty()) {
+    if (batchParameterSets.isEmpty()) {
       return new long[0];
     }
 
@@ -121,18 +121,41 @@ public long[] executeLargeBatch() throws DatabricksBatchUpdateException {
             connection,
             interpolateParameters,
             (sqlToExecute, params, statementType, closeStatement) ->
-                executeInternal(sqlToExecute, params, statementType, closeStatement));
-
-    long[] updateCounts = batchExecutor.executeBatch(databricksBatchParameterMetaData);
+                executeInternal(sqlToExecute, params, statementType, closeStatement),
+            new PreparedStatementBatchExecutor.NativeBatchExecutor() {
+              @Override
+              public boolean isSupported() {
+                return supportsNativeParameterBatching();
+              }
+
+              @Override
+              public long[] execute(String sql, List<BatchParameterSet> parameterSets)
+                  throws SQLException {
+                return executeNativeBatchInternal(sql, parameterSets);
+              }
+            });
+
+    long[] updateCounts;
+    try {
+      updateCounts = batchExecutor.executeBatch(batchParameterSets);
+    } catch (NativeBatchResultException e) {
+      // The backend already completed the batch. Clear it before propagating the count-read error
+      // so a caller retry cannot insert the same rows again.
+      clearBatchAfterExecution();
+      throw e;
+    }
 
     // Clear the batch after successful execution per JDBC spec
+    clearBatchAfterExecution();
+    return updateCounts;
+  }
+
+  private void clearBatchAfterExecution() {
     try {
       clearBatch();
     } catch (SQLException e) {
-      LOGGER.error("Failed to clear batch after successful execution", e);
+      LOGGER.error("Failed to clear batch after execution", e);
     }
-
-    return updateCounts;
   }
 
   @Override
@@ -371,7 +394,8 @@ public boolean execute() throws SQLException {
   @Override
   public void addBatch() {
     LOGGER.debug("public void addBatch()");
-    this.databricksBatchParameterMetaData.add(databricksParameterMetaData);
+    this.batchParameterSets.add(
+        BatchParameterSet.from(databricksParameterMetaData.getParameterBindings()));
     this.databricksParameterMetaData = new DatabricksParameterMetaData(sql);
   }
 
@@ -380,7 +404,7 @@ public void clearBatch() throws DatabricksSQLException {
     LOGGER.debug("public void clearBatch()");
     checkIfClosed();
     this.databricksParameterMetaData = new DatabricksParameterMetaData(sql);
-    this.databricksBatchParameterMetaData = new ArrayList<>();
+    this.batchParameterSets = new ArrayList<>();
   }
 
   @Override
@@ -755,7 +779,7 @@ private void checkLength(long targetLength, long sourceLength) throws SQLExcepti
   }
 
   private void checkIfBatchOperation() throws DatabricksSQLException {
-    if (!this.databricksBatchParameterMetaData.isEmpty()) {
+    if (!this.batchParameterSets.isEmpty()) {
       String errorMessage =
           "Batch must either be executed with executeBatch() or cleared with clearBatch()";
       LOGGER.error(errorMessage);

src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java

@@ -62,6 +62,7 @@ enum ResultSetType {
 
   private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(DatabricksResultSet.class);
   protected static final String AFFECTED_ROWS_COUNT = "num_affected_rows";
+  private static final String REPEAT_COUNT = "repeat";
   private final ExecutionStatus executionStatus;
   private final StatementId statementId;
   private final IExecutionResult executionResult;
@@ -2310,6 +2311,44 @@ public long getUpdateCount() throws SQLException {
     return updateCount;
   }
 
+  long[] getBatchUpdateCounts(int expectedCount) throws SQLException {
+    checkIfClosed();
+    if (resultSetMetaData.getColumnNameIndex(AFFECTED_ROWS_COUNT) < 1) {
+      throw new DatabricksSQLException(
+          "Native batch result is missing column " + AFFECTED_ROWS_COUNT,
+          DatabricksDriverErrorCode.RESULT_SET_ERROR);
+    }
+
+    long[] counts = new long[expectedCount];
+    int index = 0;
+    boolean hasRepeatCount = resultSetMetaData.getColumnNameIndex(REPEAT_COUNT) > 0;
+    countingUpdateRows = true;
+    try {
+      while (next()) {
+        long repeatCount = hasRepeatCount ? getLong(REPEAT_COUNT) : 1;
+        if (repeatCount < 1 || repeatCount > expectedCount - index) {
+          throw new DatabricksSQLException(
+              "Native batch returned an invalid repeat count: " + repeatCount,
+              DatabricksDriverErrorCode.RESULT_SET_ERROR);
+        }
+        long affectedRows = getLong(AFFECTED_ROWS_COUNT);
+        for (long repeated = 0; repeated < repeatCount; repeated++) {
+          counts[index++] = affectedRows;
+        }
+      }
+    } finally {
+      countingUpdateRows = false;
+    }
+
+    if (index != expectedCount) {
+      throw new DatabricksSQLException(
+          String.format(
+              "Native batch returned %d update counts for %d parameter sets", index, expectedCount),
+          DatabricksDriverErrorCode.RESULT_SET_ERROR);
+    }
+    return counts;
+  }
+
   @Override
   public boolean hasUpdateCount() throws SQLException {
     checkIfClosed();

src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java

@@ -866,6 +866,15 @@ DatabricksResultSet executeInternal(
     LOGGER.debug(stackTraceMessage);
     CompletableFuture<DatabricksResultSet> futureResultSet =
         getFutureResult(sql, params, statementType);
+    return waitForExecutionResult(sql, stackTraceMessage, futureResultSet, closeStatement);
+  }
+
+  private DatabricksResultSet waitForExecutionResult(
+      String sql,
+      String stackTraceMessage,
+      CompletableFuture<DatabricksResultSet> futureResultSet,
+      boolean closeStatement)
+      throws SQLException {
     try {
       resultSet =
           timeoutInSeconds == 0
@@ -938,6 +947,38 @@ DatabricksResultSet executeInternal(
     return result;
   }
 
+  boolean supportsNativeParameterBatching() {
+    try {
+      IDatabricksClient client = connection.getSession().getDatabricksClient();
+      return client.supportsNativeParameterBatching(connection.getSession().getComputeResource());
+    } catch (DatabricksSQLException e) {
+      LOGGER.warn("Unable to determine native batch capability, using legacy execution", e);
+      return false;
+    }
+  }
+
+  long[] executeNativeBatchInternal(String sql, List<BatchParameterSet> parameterSets)
+      throws SQLException {
+    resetForNewExecution();
+    DatabricksThreadContextHolder.setStatementType(StatementType.UPDATE);
+    String stackTraceMessage =
+        format(
+            "DatabricksResultSet executeNativeBatchInternal(String sql = %s, parameterSetCount = %s)",
+            sql, parameterSets.size());
+    LOGGER.debug(stackTraceMessage);
+    DatabricksResultSet result =
+        waitForExecutionResult(
+            sql,
+            stackTraceMessage,
+            getFutureBatchResult(sql, parameterSets, StatementType.UPDATE),
+            true);
+    try {
+      return result.getBatchUpdateCounts(parameterSets.size());
+    } catch (SQLException e) {
+      throw new NativeBatchResultException(e);
+    }
+  }
+
   CompletableFuture<DatabricksResultSet> getFutureResult(
       String sql, Map<Integer, ImmutableSqlParameter> params, StatementType statementType) {
     return CompletableFuture.supplyAsync(
@@ -954,6 +995,21 @@ CompletableFuture<DatabricksResultSet> getFutureResult(
         executor);
   }
 
+  private CompletableFuture<DatabricksResultSet> getFutureBatchResult(
+      String sql, List<BatchParameterSet> parameterSets, StatementType statementType) {
+    return CompletableFuture.supplyAsync(
+        () -> {
+          try {
+            String sqlString = escapeProcessing ? StringUtil.convertJdbcEscapeSequences(sql) : sql;
+            sqlString = StringUtil.removeRedundantEscapeClause(sqlString);
+            return getBatchResultFromClient(sqlString, parameterSets, statementType);
+          } catch (SQLException e) {
+            throw new RuntimeException(e);
+          }
+        },
+        executor);
+  }
+
   DatabricksResultSet getResultFromClient(
       String sql, Map<Integer, ImmutableSqlParameter> params, StatementType statementType)
       throws SQLException {
@@ -968,6 +1024,19 @@ DatabricksResultSet getResultFromClient(
         null /* metadataOperationType */);
   }
 
+  private DatabricksResultSet getBatchResultFromClient(
+      String sql, List<BatchParameterSet> parameterSets, StatementType statementType)
+      throws SQLException {
+    IDatabricksClient client = connection.getSession().getDatabricksClient();
+    return client.executeStatementBatch(
+        sql,
+        connection.getSession().getComputeResource(),
+        parameterSets,
+        statementType,
+        connection.getSession(),
+        this);
+  }
+
   void checkIfClosed() throws DatabricksSQLException {
     if (isClosed) {
       throw new DatabricksSQLException(