Speeding Up JDBC Bulk Inserts in YugabyteDB YSQL

				
					import java.sql.*;
import java.time.Duration;
import java.time.Instant;
import java.util.Properties;
import java.util.concurrent.ThreadLocalRandom;

/**
 * YugabyteDB JDBC insert benchmark using the YugabyteDB Smart Driver.
 *
 * Demonstrates:
 *  - JDBC batching
 *  - reWriteBatchedInserts=true
 *  - ysql_session_max_batch_size
 *  - ysql_max_in_flight_ops
 *
 * Compile:
 *   javac -cp yugabyte-jdbc-42.7.4-yb-1.12.jar YBInsertBenchmark.java
 *
 * Run:
 *   java -cp .:yugabyte-jdbc-42.7.4-yb-1.12.jar YBInsertBenchmark ...
 */
public class YBInsertBenchmark {

  /* ---------------- Argument helpers ---------------- */

  private static String getArg(String[] args, String name, String def) {
    for (int i = 0; i < args.length; i++) {
      if (args[i].equals(name) && i + 1 < args.length) {
        return args[i + 1];
      }
    }
    return def;
  }

  private static boolean hasFlag(String[] args, String name) {
    for (String a : args) {
      if (a.equals(name)) return true;
    }
    return false;
  }

  /* ---------------- Schema helpers ---------------- */

  private static void createSchema(Connection c) throws SQLException {
    try (Statement st = c.createStatement()) {
      st.execute("""
        CREATE TABLE IF NOT EXISTS insert_test_jdbc (
          id BIGSERIAL PRIMARY KEY,
          k1 INT NOT NULL,
          k2 INT NOT NULL,
          payload TEXT NOT NULL
        )
      """);

      st.execute("""
        CREATE INDEX IF NOT EXISTS insert_test_jdbc_k1_idx
        ON insert_test_jdbc(k1)
      """);

      st.execute("""
        CREATE INDEX IF NOT EXISTS insert_test_jdbc_k2_idx
        ON insert_test_jdbc(k2)
      """);
    }
  }

  private static void truncateTable(Connection c) throws SQLException {
    try (Statement st = c.createStatement()) {
      st.execute("TRUNCATE TABLE insert_test_jdbc");
    }
  }

  /* ---------------- YB session tuning ---------------- */

  private static void applySessionSettings(Connection c, int batchSize, int inFlight)
      throws SQLException {
    try (Statement st = c.createStatement()) {
      if (batchSize > 0) {
        st.execute("SET ysql_session_max_batch_size = " + batchSize);
      }
      if (inFlight > 0) {
        st.execute("SET ysql_max_in_flight_ops = " + inFlight);
      }
    }
  }

  /* ---------------- Worker ---------------- */

  private static final class Worker implements Runnable {
    private final String url;
    private final Properties props;
    private final long rows;
    private final int jdbcBatch;
    private final int payloadBytes;
    private final int ybBatch;
    private final int ybInFlight;

    volatile long inserted = 0;
    volatile Exception error = null;

    Worker(
        String url,
        Properties props,
        long rows,
        int jdbcBatch,
        int payloadBytes,
        int ybBatch,
        int ybInFlight) {

      this.url = url;
      this.props = props;
      this.rows = rows;
      this.jdbcBatch = jdbcBatch;
      this.payloadBytes = payloadBytes;
      this.ybBatch = ybBatch;
      this.ybInFlight = ybInFlight;
    }

    @Override
    public void run() {
      final String sql =
          "INSERT INTO insert_test_jdbc(k1, k2, payload) VALUES (?, ?, ?)";
      final ThreadLocalRandom rnd = ThreadLocalRandom.current();
      final String payload = "x".repeat(Math.max(1, payloadBytes));

      try (Connection c = DriverManager.getConnection(url, props)) {
        c.setAutoCommit(false);
        applySessionSettings(c, ybBatch, ybInFlight);

        try (PreparedStatement ps = c.prepareStatement(sql)) {
          long local = 0;
          int batchCount = 0;

          while (local < rows) {
            ps.setInt(1, rnd.nextInt(0, 1000));
            ps.setInt(2, rnd.nextInt(0, 2000));
            ps.setString(3, payload);

            ps.addBatch();
            batchCount++;
            local++;
            inserted = local;

            if (batchCount >= jdbcBatch) {
              ps.executeBatch();
              c.commit();
              batchCount = 0;
            }
          }

          if (batchCount > 0) {
            ps.executeBatch();
            c.commit();
          }
        }
      } catch (Exception e) {
        error = e;
      }
    }
  }

  private static String rate(double rps) {
    if (rps > 1_000_000) return String.format("%.2f M rows/sec", rps / 1_000_000);
    if (rps > 1_000) return String.format("%.2f K rows/sec", rps / 1_000);
    return String.format("%.2f rows/sec", rps);
  }

  /* ---------------- Main ---------------- */

  public static void main(String[] args) throws Exception {

    // Force-load YugabyteDB Smart Driver
    Class.forName("com.yugabyte.Driver");

    String url = getArg(args, "--url", null);
    String user = getArg(args, "--user", "yugabyte");
    String password = getArg(args, "--password", "yugabyte");

    long rows = Long.parseLong(getArg(args, "--rows", "500000"));
    int jdbcBatch = Integer.parseInt(getArg(args, "--batchSize", "512"));
    int threads = Integer.parseInt(getArg(args, "--threads", "4"));
    int payloadBytes = Integer.parseInt(getArg(args, "--payloadBytes", "200"));

    int ybBatch = Integer.parseInt(getArg(args, "--ybBatch", "0"));
    int ybInFlight = Integer.parseInt(getArg(args, "--ybInFlight", "0"));

    boolean create = hasFlag(args, "--create");
    boolean truncate = hasFlag(args, "--truncate");

    if (url == null) {
      System.err.println("Missing --url");
      System.exit(2);
    }

    Properties props = new Properties();
    props.setProperty("user", user);
    props.setProperty("password", password);

    if (create || truncate) {
      try (Connection c = DriverManager.getConnection(url, props)) {
        if (create) createSchema(c);
        if (truncate) truncateTable(c);
      }
    }

    long perThread = rows / threads;
    long remainder = rows % threads;


    Worker[] workers = new Worker[threads];
    Thread[] threadsArr = new Thread[threads];

    for (int i = 0; i < threads; i++) {
      long r = perThread + (i == 0 ? remainder : 0);
      workers[i] =
          new Worker(url, props, r, jdbcBatch, payloadBytes, ybBatch, ybInFlight);
      threadsArr[i] = new Thread(workers[i], "worker-" + i);
    }

    System.out.println("==== YugabyteDB JDBC Insert Benchmark ====");
    System.out.println("URL: " + url);
    System.out.println("Rows: " + rows);
    System.out.println("Threads: " + threads);
    System.out.println("JDBC batch size: " + jdbcBatch);
    System.out.println(
        "Session overrides: ysql_session_max_batch_size="
            + ybBatch
            + ", ysql_max_in_flight_ops="
            + ybInFlight);
    System.out.println("=========================================");

    Instant start = Instant.now();
    for (Thread t : threadsArr) t.start();
    for (Thread t : threadsArr) t.join();
    Instant end = Instant.now();

    for (Worker w : workers) {
      if (w.error != null) throw w.error;
    }

    long total = 0;
    for (Worker w : workers) total += w.inserted;

    double seconds =
        Math.max(0.001, Duration.between(start, end).toMillis() / 1000.0);

    System.out.println("\n==== Results ====");
    System.out.printf("Inserted: %,d rows%n", total);
    System.out.printf("Elapsed:  %.3f seconds%n", seconds);
    System.out.printf("Rate:     %s%n", rate(total / seconds));
    System.out.println("================");
  }
}