Iceberg Compaction: How Much Faster Are TPC-H Queries?

package com.olake.tpch;

import org.apache.iceberg.DeleteFile;
import org.apache.iceberg.Schema;
import org.apache.iceberg.Snapshot;
import org.apache.iceberg.Table;
import org.apache.iceberg.aws.glue.GlueCatalog;
import org.apache.iceberg.catalog.Namespace;
import org.apache.iceberg.catalog.TableIdentifier;
import org.apache.iceberg.data.GenericRecord;
import org.apache.iceberg.data.Record;
import org.apache.iceberg.deletes.EqualityDeleteWriter;
import org.apache.iceberg.io.OutputFile;
import org.apache.iceberg.data.parquet.GenericParquetWriter;
import org.apache.iceberg.parquet.Parquet;
import org.apache.iceberg.types.Types;

import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;

import java.io.IOException;
import java.time.Instant;
import java.time.LocalDateTime;
import java.time.ZoneId;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.stream.Collectors;

/**
 * Generates Iceberg V2 equality delete files for TPCH tables.
 *
 * v2 — Multi-batch optimisation:
 *   Groups N consecutive batches (default 50) into one "group".
 *   Each group performs:
 *     1. ONE full-table scan (reading from the original snapshot via time-travel)
 *        to collect primary-key rows for all batches in the group.
 *     2. Writes N separate equality delete files, one per batch, and commits each.
 *     3. ONE combined re-insertion from the original snapshot so logical data
 *        is unchanged.
 *
 * Supports --start-batch to resume from a specific batch (1-indexed, matches
 * the "batch N/1000" numbering in logs).
 *
 * End result: table data is logically identical, but the table now contains
 * ~{num-batches} equality delete files, which degrade read performance
 * until compaction is run.
 *
 * SAFETY: Before processing each table the script logs the current snapshot ID.
 *         If anything goes wrong you can roll back using:
 *           CALL catalog.system.rollback_to_snapshot('db.table', snapshot_id);
 */
public class EqualityDeleteGenerator {

    private static final DateTimeFormatter TS_FMT =
            DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss");

    // ── TPCH primary keys ──────────────────────────────────────────────
    private static final Map<String, List<String>> TABLE_KEYS = new LinkedHashMap<>();
    static {
        TABLE_KEYS.put("lineitem",  Arrays.asList("l_orderkey", "l_linenumber"));
        TABLE_KEYS.put("orders",    Collections.singletonList("o_orderkey"));
        TABLE_KEYS.put("customer",  Collections.singletonList("c_custkey"));
        TABLE_KEYS.put("part",      Collections.singletonList("p_partkey"));
        TABLE_KEYS.put("supplier",  Collections.singletonList("s_suppkey"));
        TABLE_KEYS.put("partsupp",  Arrays.asList("ps_partkey", "ps_suppkey"));
        TABLE_KEYS.put("nation",    Collections.singletonList("n_nationkey"));
        TABLE_KEYS.put("region",    Collections.singletonList("r_regionkey"));
    }

    // ── Entry point ────────────────────────────────────────────────────
    public static void main(String[] args) {
        String  catalogName     = getArg(args, "--catalog",           "olake_iceberg");
        String  database        = getArg(args, "--database",          "postgres_postgres_tpch");
        String  tablesArg       = getArg(args, "--tables",
                "lineitem,orders,customer,part,supplier,partsupp,nation,region");
        int     numBatches      = Integer.parseInt(
                getArg(args, "--num-batches",      "1000"));
        double  deletePct       = Double.parseDouble(
                getArg(args, "--delete-percent",    "2.0"));
        boolean continueErr     = hasFlag(args, "--continue-on-error");
        int     batchesPerGroup = Integer.parseInt(
                getArg(args, "--batches-per-group", "50"));
        String  startBatchArg   = getArg(args, "--start-batch",       "");

        // Parse --start-batch "lineitem=10,orders=5"
        // Values are 1-indexed (matches "batch N/1000" in logs)
        Map<String, Integer> startBatchMap = new HashMap<>();
        if (startBatchArg != null && !startBatchArg.isEmpty()) {
            for (String entry : startBatchArg.split(",")) {
                String[] kv = entry.split("=");
                if (kv.length == 2) {
                    startBatchMap.put(kv[0].trim(),
                            Integer.parseInt(kv[1].trim()));
                }
            }
        }

        if (database == null) {
            System.err.println("ERROR: --database is required");
            System.exit(1);
        }

        List<String> tableNames = Arrays.stream(tablesArg.split(","))
                .map(String::trim)
                .filter(s -> !s.isEmpty())
                .collect(Collectors.toList());

        SparkSession spark = SparkSession.builder()
                .appName("EqualityDeleteGenerator")
                .getOrCreate();

        GlueCatalog catalog = initGlueCatalog(spark, catalogName);

        log("════════════════════════════════════════════════════");
        log("  Iceberg Equality Delete File Generator  (v2 – multi-batch)");
        log("════════════════════════════════════════════════════");
        log("Catalog:           " + catalogName);
        log("Database:          " + database);
        log("Tables:            " + tableNames);
        log("Num batches:       " + numBatches);
        log("Delete percent:    " + deletePct + "%");
        log("Batches per group: " + batchesPerGroup);
        log("Start-batch map:   " + (startBatchMap.isEmpty() ? "(all from 1)" : startBatchMap));
        log("Continue on err:   " + continueErr);
        log("════════════════════════════════════════════════════");

        List<String> failures  = new ArrayList<>();
        List<String> summaries = new ArrayList<>();

        long globalStart = System.currentTimeMillis();

        for (int tIdx = 0; tIdx < tableNames.size(); tIdx++) {
            String tableName = tableNames.get(tIdx);
            // 1-indexed start batch (default 1 = from the beginning)
            int startBatch1Idx = startBatchMap.getOrDefault(tableName, 1);
            int startBatchIdx  = startBatch1Idx - 1;  // convert to 0-indexed

            log("");
            log("────────────────────────────────────────────────────");
            log(String.format("  TABLE %d/%d : %s  (start batch: %d)",
                    tIdx + 1, tableNames.size(), tableName, startBatch1Idx));
            log("────────────────────────────────────────────────────");
            try {
                String summary = processTable(spark, catalog, catalogName,
                        database, tableName, numBatches, deletePct,
                        startBatchIdx, batchesPerGroup);
                summaries.add(summary);
            } catch (Exception e) {
                String msg = "[" + tableName + "] FAILED: " + e.getMessage();
                logError(msg);
                e.printStackTrace();
                failures.add(msg);
                summaries.add(String.format("[%s] FAILED — %s",
                        tableName, e.getMessage()));
                if (!continueErr) break;
            }
        }

        try { catalog.close(); } catch (Exception ignored) { }

        long globalSecs = (System.currentTimeMillis() - globalStart) / 1000;

        // ── Final summary ──
        log("");
        log("════════════════════════════════════════════════════");
        log("  FINAL SUMMARY  (total time: " + formatDuration(globalSecs) + ")");
        log("════════════════════════════════════════════════════");
        for (String s : summaries) {
            log("  " + s);
        }
        log("════════════════════════════════════════════════════");

        if (!failures.isEmpty()) {
            logError("FAILURES (" + failures.size() + "):");
            failures.forEach(f -> logError("  ✗ " + f));
            System.exit(2);
        }
        log("ALL TABLES COMPLETED SUCCESSFULLY");
    }

    // ── Catalog bootstrap ──────────────────────────────────────────────

    private static GlueCatalog initGlueCatalog(SparkSession spark,
                                               String catalogName) {
        String prefix = "spark.sql.catalog." + catalogName + ".";
        Map<String, String> props = new HashMap<>();

        for (scala.Tuple2<String, String> kv :
                spark.sparkContext().getConf().getAll()) {
            if (kv._1().startsWith(prefix)) {
                String key = kv._1().substring(prefix.length());
                if (!key.isEmpty()) {
                    props.put(key, kv._2());
                }
            }
        }
        props.remove("catalog-impl");

        log("Initializing GlueCatalog with properties: " + props);

        GlueCatalog catalog = new GlueCatalog();
        catalog.setConf(spark.sparkContext().hadoopConfiguration());
        catalog.initialize(catalogName, props);

        log("GlueCatalog initialized successfully.");
        return catalog;
    }

    // ── Per-table processing (multi-batch groups) ─────────────────────

    private static String processTable(SparkSession spark,
                                       GlueCatalog catalog,
                                       String catalogName,
                                       String database,
                                       String tableName,
                                       int numBatches,
                                       double deletePct,
                                       int startBatch,
                                       int batchesPerGroup) throws IOException {

        List<String> keyColumns = TABLE_KEYS.get(tableName);
        if (keyColumns == null) {
            log("[" + tableName + "] No key columns defined — skipping.");
            return String.format("[%s] SKIPPED — no key columns defined", tableName);
        }

        String fullSparkName = catalogName + "." + database + "." + tableName;
        TableIdentifier tableId =
                TableIdentifier.of(Namespace.of(database), tableName);
        Table table = catalog.loadTable(tableId);

        // ── Capture original snapshot — used for ALL reads ──
        // Reading from this fixed snapshot avoids the overhead of applying
        // accumulated delete files during key-collection and reinsertion scans.
        table.refresh();
        Snapshot currentSnap = table.currentSnapshot();
        long originalSnapId = -1;
        if (currentSnap != null) {
            originalSnapId = currentSnap.snapshotId();
            String snapTime = Instant.ofEpochMilli(currentSnap.timestampMillis())
                    .atZone(ZoneId.systemDefault()).format(TS_FMT);
            log(String.format(
                    "[%s] *** ROLLBACK POINT ***  snapshot_id = %d  (taken at %s)",
                    tableName, originalSnapId, snapTime));
            log(String.format(
                    "[%s]     To undo all changes:  CALL %s.system.rollback_to_snapshot"
                    + "('%s.%s', %d);",
                    tableName, catalogName, database, tableName, originalSnapId));
        } else {
            log("[" + tableName + "] WARNING: table has no snapshots yet.");
        }

        // ── Row count (from current snapshot) ──
        long totalRows = spark.sql("SELECT count(*) FROM " + fullSparkName)
                              .first().getLong(0);
        log(String.format("[%s] Total rows: %,d", tableName, totalRows));
        if (totalRows == 0) {
            log("[" + tableName + "] Empty — skipping.");
            return String.format("[%s] SKIPPED — empty table", tableName);
        }

        // ── Modulo base (handle small tables) ──
        long standardModulo = Math.round(numBatches / (deletePct / 100.0));
        long moduloBase;
        if (totalRows < standardModulo) {
            moduloBase = numBatches;
            log(String.format(
                    "[%s] Small table (%,d rows) — modulo=%d (covers all rows)",
                    tableName, totalRows, numBatches));
        } else {
            moduloBase = standardModulo;
            long expectedDeletes = Math.round(totalRows * deletePct / 100.0);
            long perBatch = expectedDeletes / numBatches;
            log(String.format(
                    "[%s] Target: ~%,d rows (%.1f%%), ~%,d rows/batch, modulo=%,d",
                    tableName, expectedDeletes, deletePct, perBatch, moduloBase));
        }

        // ── Iceberg schema info ──
        Schema iceSchema    = table.schema();
        int[]  eqFieldIds   = keyColumns.stream()
                .mapToInt(c -> iceSchema.findField(c).fieldId())
                .toArray();
        Schema deleteSchema = iceSchema.select(keyColumns);

        // ── Hash expression for Spark SQL ──
        String keyColsJoined = String.join(", ", keyColumns);
        String hashExpr = "pmod(hash(" + keyColsJoined + "), " + moduloBase + ")";

        int totalGroups = (int) Math.ceil(
                (numBatches - startBatch) / (double) batchesPerGroup);

        log(String.format(
                "[%s] Config: startBatch=%d  batchesPerGroup=%d  totalGroups=%d",
                tableName, startBatch + 1, batchesPerGroup, totalGroups));
        log(String.format("[%s] Hash expression: %s", tableName, hashExpr));
        log(String.format("[%s] Original snapshot for reads: %d",
                tableName, originalSnapId));
        log(String.format("[%s] Table location: %s", tableName, table.location()));
        log(String.format("[%s] Starting group loop...", tableName));

        long totalDeletedRows   = 0;
        int  deleteFilesCreated = 0;
        int  emptyBatches       = 0;
        long tableStartMs       = System.currentTimeMillis();
        int  groupNum           = 0;

        for (int groupStart = startBatch;
             groupStart < numBatches;
             groupStart += batchesPerGroup) {

            int groupEnd  = Math.min(groupStart + batchesPerGroup, numBatches);
            int groupSize = groupEnd - groupStart;
            groupNum++;

            long groupStartMs = System.currentTimeMillis();

            log("");
            log(String.format(
                    "[%s] ── GROUP %d/%d  (batches %d–%d, %d batches) ──",
                    tableName, groupNum, totalGroups,
                    groupStart + 1, groupEnd, groupSize));

            // WHERE clause — consecutive batch range
            String groupWhere;
            if (groupSize == 1) {
                groupWhere = hashExpr + " = " + groupStart;
            } else {
                groupWhere = hashExpr + " >= " + groupStart
                        + " AND " + hashExpr + " < " + groupEnd;
            }

            // ════════════════════════════════════════════════════
            // STEP 1:  ONE scan of the original snapshot to collect all keys
            // ════════════════════════════════════════════════════
            String collectSql = String.format(
                    "SELECT %s, %s AS __batch_num FROM %s VERSION AS OF %d WHERE %s",
                    keyColsJoined, hashExpr, fullSparkName,
                    originalSnapId, groupWhere);

            log(String.format(
                    "[%s]   Step 1/3: Collecting keys for %d batches "
                    + "from snap %d ...",
                    tableName, groupSize, originalSnapId));

            long t1 = System.currentTimeMillis();
            List<Row> allKeyRows = spark.sql(collectSql).collectAsList();
            long collectSecs = (System.currentTimeMillis() - t1) / 1000;

            log(String.format("[%s]   Collected %,d key rows in %s",
                    tableName, allKeyRows.size(), formatDuration(collectSecs)));

            if (allKeyRows.isEmpty()) {
                emptyBatches += groupSize;
                log(String.format("[%s]   Group entirely empty — skipping",
                        tableName));
                continue;
            }

            // Partition rows by batch number
            Map<Integer, List<Row>> batchToKeys = new HashMap<>();
            for (Row row : allKeyRows) {
                Object batchVal = row.get(row.size() - 1); // __batch_num
                int batchNum = ((Number) batchVal).intValue();
                batchToKeys.computeIfAbsent(batchNum, k -> new ArrayList<>())
                           .add(row);
            }
            allKeyRows = null; // free memory

            // ════════════════════════════════════════════════════
            // STEP 2:  Write separate equality delete files + commit each
            // ════════════════════════════════════════════════════
            log(String.format(
                    "[%s]   Step 2/3: Writing %d equality delete files ...",
                    tableName, batchToKeys.size()));

            long t2 = System.currentTimeMillis();
            int groupDeletedRows = 0;
            int groupFiles       = 0;

            for (int batch = groupStart; batch < groupEnd; batch++) {
                List<Row> keyRows =
                        batchToKeys.getOrDefault(batch, Collections.emptyList());
                if (keyRows.isEmpty()) {
                    emptyBatches++;
                    continue;
                }

                String delPath = table.location() + "/data/eq-del-" + tableName
                        + "-b" + String.format("%04d", batch)
                        + "-" + UUID.randomUUID().toString() + ".parquet";

                OutputFile outFile = table.io().newOutputFile(delPath);

                EqualityDeleteWriter<Record> writer =
                        Parquet.writeDeletes(outFile)
                        .forTable(table)
                        .rowSchema(deleteSchema)
                        .createWriterFunc(msgType ->
                                GenericParquetWriter.create(deleteSchema, msgType))
                        .equalityFieldIds(eqFieldIds)
                        .buildEqualityWriter();

                for (Row row : keyRows) {
                    GenericRecord rec = GenericRecord.create(deleteSchema);
                    for (int col = 0; col < keyColumns.size(); col++) {
                        rec.set(col, coerceType(
                                row.get(col),
                                deleteSchema.columns().get(col).type()));
                    }
                    writer.write(rec);
                }
                writer.close();

                DeleteFile deleteFile = writer.toDeleteFile();
                table.newRowDelta()
                        .addDeletes(deleteFile)
                        .commit();
                table.refresh();

                totalDeletedRows += keyRows.size();
                deleteFilesCreated++;
                groupDeletedRows += keyRows.size();
                groupFiles++;

                // Short per-batch log
                String fileName =
                        delPath.substring(delPath.lastIndexOf('/') + 1);
                log(String.format(
                        "[%s]     batch %d/%d — %,d rows → %s",
                        tableName, batch + 1, numBatches,
                        keyRows.size(), fileName));
            }

            long writeSecs = (System.currentTimeMillis() - t2) / 1000;
            log(String.format(
                    "[%s]   Wrote %d delete files (%,d rows) in %s",
                    tableName, groupFiles, groupDeletedRows,
                    formatDuration(writeSecs)));

            batchToKeys = null; // free memory

            // ════════════════════════════════════════════════════
            // STEP 3:  ONE combined reinsertion from the original snapshot
            // ════════════════════════════════════════════════════
            log(String.format(
                    "[%s]   Step 3/3: Reinserting %,d rows from original "
                    + "snapshot %d ...",
                    tableName, groupDeletedRows, originalSnapId));

            long t3 = System.currentTimeMillis();
            spark.sql("REFRESH TABLE " + fullSparkName);
            spark.sql(String.format(
                    "INSERT INTO %s SELECT * FROM %s VERSION AS OF %d WHERE %s",
                    fullSparkName, fullSparkName, originalSnapId, groupWhere));

            long reinsertSecs = (System.currentTimeMillis() - t3) / 1000;
            log(String.format("[%s]   Reinsertion complete in %s",
                    tableName, formatDuration(reinsertSecs)));

            // ── Group summary ──
            long groupSecs   = (System.currentTimeMillis() - groupStartMs) / 1000;
            double pctDone   = (groupEnd * 100.0) / numBatches;
            double pctDeleted = (totalDeletedRows * 100.0) / totalRows;
            long elapsedSecs = (System.currentTimeMillis() - tableStartMs) / 1000;

            // ETA
            double batchesDone = groupEnd - startBatch;
            double batchesLeft = numBatches - groupEnd;
            long etaSecs = (batchesDone > 0)
                    ? Math.round(elapsedSecs * batchesLeft / batchesDone) : 0;

            log(String.format(
                "[%s]   GROUP %d/%d ✓  files=%d  rows=%,d  "
                + "collect=%s  write=%s  reinsert=%s  total=%s",
                tableName, groupNum, totalGroups,
                groupFiles, groupDeletedRows,
                formatDuration(collectSecs), formatDuration(writeSecs),
                formatDuration(reinsertSecs), formatDuration(groupSecs)));

            log(String.format(
                "[%s]   Cumulative: %,d files, %,d rows (%.2f%%)  "
                + "progress=%.1f%%  elapsed=%s  ETA=%s",
                tableName, deleteFilesCreated, totalDeletedRows, pctDeleted,
                pctDone, formatDuration(elapsedSecs),
                formatDuration(etaSecs)));

            // Milestone every 100 batches or at the end
            if (groupEnd % 100 == 0 || groupEnd == numBatches) {
                log(String.format(
                    "[%s]   ──── MILESTONE: %d/%d batches "
                    + "(%,d files, %,d rows, %.2f%%) ────",
                    tableName, groupEnd, numBatches,
                    deleteFilesCreated, totalDeletedRows, pctDeleted));
            }
        }

        long tableSecs = (System.currentTimeMillis() - tableStartMs) / 1000;
        double actualPct = (totalDeletedRows * 100.0) / totalRows;

        log("");
        log(String.format("[%s] ════════ TABLE COMPLETE ════════", tableName));
        log(String.format("[%s]   Total rows in table:     %,d",
                tableName, totalRows));
        log(String.format("[%s]   Equality-deleted rows:   %,d (%.2f%%)",
                tableName, totalDeletedRows, actualPct));
        log(String.format("[%s]   Delete files created:    %d",
                tableName, deleteFilesCreated));
        log(String.format("[%s]   Empty batches skipped:   %d",
                tableName, emptyBatches));
        log(String.format("[%s]   Total time:              %s",
                tableName, formatDuration(tableSecs)));
        if (originalSnapId >= 0) {
            log(String.format("[%s]   Rollback snapshot ID:    %d",
                    tableName, originalSnapId));
        }
        log(String.format("[%s] ════════════════════════════════", tableName));

        return String.format("[%s] %,d delete files, %,d rows (%.2f%%), %s",
                tableName, deleteFilesCreated, totalDeletedRows, actualPct,
                formatDuration(tableSecs));
    }

    // ── Helpers ────────────────────────────────────────────────────────

    private static Object coerceType(Object value,
                                     org.apache.iceberg.types.Type iceType) {
        if (value == null) return null;
        if (iceType instanceof Types.IntegerType && value instanceof Long) {
            return ((Long) value).intValue();
        }
        if (iceType instanceof Types.LongType && value instanceof Integer) {
            return ((Integer) value).longValue();
        }
        return value;
    }

    private static void log(String message) {
        String ts = LocalDateTime.now().format(TS_FMT);
        System.out.println("[" + ts + "] " + message);
    }

    private static void logError(String message) {
        String ts = LocalDateTime.now().format(TS_FMT);
        System.err.println("[" + ts + "] ERROR: " + message);
    }

    private static String formatDuration(long totalSeconds) {
        if (totalSeconds < 60) return totalSeconds + "s";
        long hours = totalSeconds / 3600;
        long mins  = (totalSeconds % 3600) / 60;
        long secs  = totalSeconds % 60;
        if (hours > 0) return String.format("%dh %dm %ds", hours, mins, secs);
        return String.format("%dm %ds", mins, secs);
    }

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

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