[4/7] Implement native Thrift parameter batching

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

@@ -0,0 +1,88 @@
+package com.databricks.jdbc.api.impl;
+
+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;
+import java.util.stream.Collectors;
+
+/**
+ * Immutable, position-ordered snapshot of one prepared-statement parameter set.
+ *
+ * <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) {
+    Objects.requireNonNull(parameterBindings, "parameterBindings");
+    List<ImmutableSqlParameter> orderedParameters =
+        parameterBindings.entrySet().stream()
+            .sorted(Comparator.comparingInt(Map.Entry::getKey))
+            .map(BatchParameterSet::snapshotParameter)
+            .collect(Collectors.toList());
+    return new BatchParameterSet(orderedParameters);
+  }
+
+  public List<ImmutableSqlParameter> getParameters() {
+    return parameters;
+  }
+
+  public Map<Integer, ImmutableSqlParameter> getParameterBindings() {
+    return parameterBindings;
+  }
+
+  public int size() {
+    return parameters.size();
+  }
+
+  public boolean isEmpty() {
+    return parameters.isEmpty();
+  }
+
+  private static ImmutableSqlParameter snapshotParameter(
+      Map.Entry<Integer, ImmutableSqlParameter> entry) {
+    ImmutableSqlParameter parameter = entry.getValue();
+    return ImmutableSqlParameter.builder()
+        .cardinal(entry.getKey())
+        .type(parameter.type())
+        .value(snapshotValue(parameter.value()))
+        .build();
+  }
+
+  private static Object snapshotValue(Object value) {
+    if (value instanceof Timestamp) {
+      Timestamp timestamp = (Timestamp) value;
+      Timestamp copy = new Timestamp(timestamp.getTime());
+      copy.setNanos(timestamp.getNanos());
+      return copy;
+    }
+    if (value instanceof Date) {
+      return new Date(((Date) value).getTime());
+    }
+    if (value instanceof Time) {
+      return new Time(((Time) value).getTime());
+    }
+    if (value instanceof byte[]) {
+      return ((byte[]) value).clone();
+    }
+    return value;
+  }
+}

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

@@ -1480,6 +1480,11 @@ public boolean isBatchedInsertsEnabled() {
     return getParameter(DatabricksJdbcUrlParams.ENABLE_BATCHED_INSERTS).equals("1");
   }
 
+  @Override
+  public boolean isNativeBatchingEnabled() {
+    return getParameter(DatabricksJdbcUrlParams.ENABLE_NATIVE_BATCHING).equals("1");
+  }
+
   @Override
   public List<String> getNonRowcountQueryPrefixes() {
     String prefixesStr = getParameter(DatabricksJdbcUrlParams.NON_ROWCOUNT_QUERY_PREFIXES);

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(