Migrating Hive tables to Iceberg: snapshot first, migrate second

"""Snapshot a Hive Parquet table into Iceberg and see what the copy shares.

Run inside the datalake-playground stack:
  docker cp snapshot_migration.py spark-master:/tmp/
  docker exec spark-master bash -c '$SPARK_HOME/bin/spark-submit --master "local[2]" \
    --driver-memory 1g --jars $(ls $ICEBERG_HOME/iceberg-spark-runtime*.jar) \
    /tmp/snapshot_migration.py'
"""
import random
from datetime import datetime, timedelta
from decimal import Decimal

from pyspark.sql import SparkSession

spark = (
    SparkSession.builder.appName("hive-snapshot-migration")
    .config("spark.sql.extensions",
            "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions")
    .config("spark.serializer", "org.apache.spark.serializer.KryoSerializer")
    # The session catalog must be Iceberg's SparkSessionCatalog for snapshot and
    # migrate: it serves Hive tables as before and Iceberg tables next to them.
    .config("spark.sql.catalog.spark_catalog", "org.apache.iceberg.spark.SparkSessionCatalog")
    .config("spark.sql.catalog.spark_catalog.type", "hive")
    .config("spark.sql.catalog.spark_catalog.uri", "thrift://hive-metastore:9083")
    .config("spark.hadoop.hive.exec.dynamic.partition.mode", "nonstrict")
    .config("spark.eventLog.enabled", "false")
    .config("spark.ui.enabled", "false")
    .enableHiveSupport()
    .getOrCreate()
)
spark.sparkContext.setLogLevel("ERROR")
sql = spark.sql
DB = "blog_hive_migration"


def one(q):
    return sql(q).collect()[0][0]


def prop(table, key):
    rows = sql(f"SHOW TBLPROPERTIES {table}").collect()
    return {r[0]: r[1] for r in rows}.get(key)


def describe(table, key):
    return (sql(f"DESCRIBE TABLE EXTENDED {table}")
            .where(f"col_name = '{key}'").collect()[0][1])


def storage(path):
    """Parquet data files and their bytes under a location, via the Hadoop FileSystem."""
    jvm = spark._jvm
    p = jvm.org.apache.hadoop.fs.Path(path)
    it = p.getFileSystem(spark._jsc.hadoopConfiguration()).listFiles(p, True)
    files, size = 0, 0
    while it.hasNext():
        f = it.next()
        if f.getPath().getName().endswith(".parquet") or f.getPath().getName().startswith("part-"):
            files += 1
            size += f.getLen()
    return files, size


def remove(path):
    """Delete a directory. The Hive 4 metastore stores tables created without
    LOCATION as EXTERNAL, so DROP TABLE leaves their files behind."""
    jvm = spark._jvm
    p = jvm.org.apache.hadoop.fs.Path(path)
    p.getFileSystem(spark._jsc.hadoopConfiguration()).delete(p, True)


def orders(n, day):
    rng = random.Random(day)
    start = datetime.strptime(day, "%Y-%m-%d")
    return [(i, rng.randint(1, 900), Decimal(str(round(rng.uniform(5, 500), 2))),
             rng.choice(["north", "south", "east", "west"]),
             start + timedelta(seconds=rng.randint(0, 86_399)), day)
            for i in range(n)]


SCHEMA = ("order_id BIGINT, customer_id BIGINT, amount DECIMAL(10,2), "
          "region STRING, event_ts TIMESTAMP, dt STRING")

sql(f"CREATE DATABASE IF NOT EXISTS {DB}")
DB_LOC = sql(f"DESCRIBE DATABASE {DB}").where("info_name = 'Location'").collect()[0][1]
for t in ("legacy_a_ice", "legacy_a"):          # make the script re-runnable
    sql(f"DROP TABLE IF EXISTS {DB}.{t}")
    remove(f"{DB_LOC}/{t}")
sql(f"""CREATE TABLE {DB}.legacy_a (
          order_id BIGINT, customer_id BIGINT, amount DECIMAL(10,2),
          region STRING, event_ts TIMESTAMP)
        PARTITIONED BY (dt STRING)
        STORED AS PARQUET""")
for day in ("2026-01-05", "2026-01-06", "2026-01-07"):
    spark.createDataFrame(orders(2_000, day), SCHEMA).createOrReplaceTempView("batch")
    sql(f"INSERT INTO {DB}.legacy_a PARTITION (dt) SELECT * FROM batch")

src_loc = describe(f"{DB}.legacy_a", "Location")
print("== the Hive table ==")
print("provider            :", describe(f"{DB}.legacy_a", "Provider"))
print("type in metastore   :", describe(f"{DB}.legacy_a", "Type"))
print("translated          :", prop(f"{DB}.legacy_a", "TRANSLATED_TO_EXTERNAL"),
      "| external.table.purge:", prop(f"{DB}.legacy_a", "external.table.purge"))
print("serde               :", describe(f"{DB}.legacy_a", "Serde Library"))
print("rows                :", one(f"SELECT count(*) FROM {DB}.legacy_a"))
print("partitions          :", sql(f"SHOW PARTITIONS {DB}.legacy_a").count())
print("files, bytes        :", storage(src_loc))

print("\n== CALL spark_catalog.system.snapshot ==")
sql(f"CALL spark_catalog.system.snapshot('{DB}.legacy_a', '{DB}.legacy_a_ice')").show()
copy_loc = describe(f"{DB}.legacy_a_ice", "Location")
print("copy provider       :", describe(f"{DB}.legacy_a_ice", "Provider"))
print("copy rows           :", one(f"SELECT count(*) FROM {DB}.legacy_a_ice"))
print("copy snapshots      :", one(f"SELECT count(*) FROM {DB}.legacy_a_ice.snapshots"))
print("copy gc.enabled     :", prop(f"{DB}.legacy_a_ice", "gc.enabled"))
print("copy 'snapshot'     :", prop(f"{DB}.legacy_a_ice", "snapshot"))
print("source location     :", src_loc)
print("copy location       :", copy_loc)
f = sql(f"SELECT count(*) AS n, sum(file_size_in_bytes) AS bytes, "
        f"sum(CASE WHEN file_path LIKE '{src_loc}/%' THEN 1 ELSE 0 END) AS in_source "
        f"FROM {DB}.legacy_a_ice.files").collect()[0]
print(f"copy data files     : {f.n} files, {f.bytes} bytes, {f.in_source} of them under the source location")
print("partition spec      :", sql(f"SELECT spec_id, partition FROM {DB}.legacy_a_ice.files LIMIT 1").collect()[0])
print("data files under the copy's location:", storage(copy_loc))

print("\n== verifying the copy against the source ==")
def per_partition(table):
    return [tuple(r) for r in sql(f"SELECT dt, count(*) AS n FROM {table} "
                                  f"GROUP BY dt ORDER BY dt").collect()]


def checksum(table):
    return tuple(sql(f"SELECT sum(amount), min(amount), max(amount) FROM {table}").collect()[0])


print("per partition, source:", per_partition(f"{DB}.legacy_a"))
print("per partition, copy  :", per_partition(f"{DB}.legacy_a_ice"))
print("checksums, source    :", checksum(f"{DB}.legacy_a"))
print("checksums, copy      :", checksum(f"{DB}.legacy_a_ice"))
print("file accounting, copy:", tuple(sql(f"SELECT count(*), sum(file_size_in_bytes) "
                                          f"FROM {DB}.legacy_a_ice.files").collect()[0]),
      "| source directory:", storage(src_loc))

print("\n== the two tables diverge ==")
spark.createDataFrame(orders(100, "2026-01-08"), SCHEMA).createOrReplaceTempView("late")
sql(f"INSERT INTO {DB}.legacy_a PARTITION (dt) SELECT * FROM late")
print("source rows after a Hive insert   :", one(f"SELECT count(*) FROM {DB}.legacy_a"))
print("copy rows after that insert       :", one(f"SELECT count(*) FROM {DB}.legacy_a_ice"))
sql(f"INSERT INTO {DB}.legacy_a_ice SELECT * FROM late WHERE order_id < 50")
print("copy rows after writing to the copy:", one(f"SELECT count(*) FROM {DB}.legacy_a_ice"))
print("source rows after that             :", one(f"SELECT count(*) FROM {DB}.legacy_a"))
new_file = sql(f"SELECT file_path FROM {DB}.legacy_a_ice.files "
               f"WHERE file_path NOT LIKE '{src_loc}/%'").collect()[0][0]
print("the copy's new file lands in       :", new_file.rsplit("/", 2)[0] + "/...")

print("\n== dropping the copy ==")
try:
    sql(f"DROP TABLE {DB}.legacy_a_ice PURGE")
    print("DROP TABLE ... PURGE: succeeded")
except Exception as e:
    msg = e.java_exception.getMessage() if hasattr(e, "java_exception") else str(e)
    print("DROP TABLE ... PURGE:", msg.splitlines()[0][:160])
sql(f"DROP TABLE {DB}.legacy_a_ice")
print("after DROP TABLE (no PURGE), source rows:", one(f"SELECT count(*) FROM {DB}.legacy_a"))
print("source files, bytes                     :", storage(src_loc))
spark.stop()

"""What migrate and snapshot say when Spark's session catalog is not Iceberg's.

Run inside the datalake-playground stack, after snapshot_migration.py:
  docker cp migrate_needs_session_catalog.py spark-master:/tmp/
  docker exec spark-master bash -c '$SPARK_HOME/bin/spark-submit --master "local[2]" \
    --driver-memory 1g --jars $(ls $ICEBERG_HOME/iceberg-spark-runtime*.jar) \
    /tmp/migrate_needs_session_catalog.py'
"""
from pyspark.sql import SparkSession

spark = (
    SparkSession.builder.appName("migrate-needs-session-catalog")
    .config("spark.sql.extensions",
            "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions")
    .config("spark.serializer", "org.apache.spark.serializer.KryoSerializer")
    # An Iceberg catalog on the same metastore, but spark_catalog left as Spark's own.
    .config("spark.sql.catalog.ice", "org.apache.iceberg.spark.SparkCatalog")
    .config("spark.sql.catalog.ice.type", "hive")
    .config("spark.sql.catalog.ice.uri", "thrift://hive-metastore:9083")
    .config("spark.eventLog.enabled", "false")
    .config("spark.ui.enabled", "false")
    .enableHiveSupport()
    .getOrCreate()
)
spark.sparkContext.setLogLevel("ERROR")
spark.sql("DROP TABLE IF EXISTS ice.blog_hive_migration.legacy_a_copy")   # re-runnable

for q in ("CALL ice.system.migrate('blog_hive_migration.legacy_a')",
          "CALL ice.system.snapshot('spark_catalog.blog_hive_migration.legacy_a', "
          "'ice.blog_hive_migration.legacy_a_copy')",
          "CALL spark_catalog.system.migrate('blog_hive_migration.legacy_a')"):
    try:
        spark.sql(q).collect()
        print(f"{q}\n  -> succeeded")
    except Exception as e:
        msg = e.java_exception.getMessage() if hasattr(e, "java_exception") else str(e)
        print(f"{q}\n  -> {type(e).__name__}: {msg.splitlines()[0][:200]}")
spark.stop()

"""Migrate a Hive Parquet table to Iceberg in place, roll it back, and see what PURGE does.

Run inside the datalake-playground stack:
  docker cp migrate_in_place.py spark-master:/tmp/
  docker exec spark-master bash -c '$SPARK_HOME/bin/spark-submit --master "local[2]" \
    --driver-memory 1g --jars $(ls $ICEBERG_HOME/iceberg-spark-runtime*.jar) \
    /tmp/migrate_in_place.py'
"""
import random
from datetime import datetime, timedelta
from decimal import Decimal

from pyspark.sql import SparkSession

spark = (
    SparkSession.builder.appName("hive-migrate-in-place")
    .config("spark.sql.extensions",
            "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions")
    .config("spark.serializer", "org.apache.spark.serializer.KryoSerializer")
    .config("spark.sql.catalog.spark_catalog", "org.apache.iceberg.spark.SparkSessionCatalog")
    .config("spark.sql.catalog.spark_catalog.type", "hive")
    .config("spark.sql.catalog.spark_catalog.uri", "thrift://hive-metastore:9083")
    .config("spark.hadoop.hive.exec.dynamic.partition.mode", "nonstrict")
    .config("spark.eventLog.enabled", "false")
    .config("spark.ui.enabled", "false")
    .enableHiveSupport()
    .getOrCreate()
)
spark.sparkContext.setLogLevel("ERROR")
sql = spark.sql
DB = "blog_hive_migration"


def one(q):
    return sql(q).collect()[0][0]


def props(table):
    return {r[0]: r[1] for r in sql(f"SHOW TBLPROPERTIES {table}").collect()}


def describe(table, key):
    return (sql(f"DESCRIBE TABLE EXTENDED {table}")
            .where(f"col_name = '{key}'").collect()[0][1])


def remove(path):
    jvm = spark._jvm
    p = jvm.org.apache.hadoop.fs.Path(path)
    p.getFileSystem(spark._jsc.hadoopConfiguration()).delete(p, True)


def attempt(label, q):
    try:
        result = sql(q).collect()
        print(f"{label}: {result[0][0] if result else 'ok'}")
    except Exception as e:
        msg = e.java_exception.getMessage() if hasattr(e, "java_exception") else str(e)
        # show the root cause when Spark wraps it in a job failure
        root = [l for l in msg.splitlines() if "FileNotFoundException" in l]
        line = root[0].split("): ", 1)[-1] if root else msg.splitlines()[0]
        print(f"{label}: {line[:230]}")


def orders(n, day):
    rng = random.Random(day)
    start = datetime.strptime(day, "%Y-%m-%d")
    return [(i, rng.randint(1, 900), Decimal(str(round(rng.uniform(5, 500), 2))),
             rng.choice(["north", "south", "east", "west"]),
             start + timedelta(seconds=rng.randint(0, 86_399)), day)
            for i in range(n)]


SCHEMA = ("order_id BIGINT, customer_id BIGINT, amount DECIMAL(10,2), "
          "region STRING, event_ts TIMESTAMP, dt STRING")


def hive_table(name, location=None):
    """A Hive Parquet table with 6,000 orders over three daily partitions."""
    where = f"LOCATION '{location}'" if location else ""
    kind = "EXTERNAL TABLE" if location else "TABLE"
    sql(f"""CREATE {kind} {DB}.{name} (
              order_id BIGINT, customer_id BIGINT, amount DECIMAL(10,2),
              region STRING, event_ts TIMESTAMP)
            PARTITIONED BY (dt STRING)
            STORED AS PARQUET {where}""")
    for day in ("2026-01-05", "2026-01-06", "2026-01-07"):
        spark.createDataFrame(orders(2_000, day), SCHEMA).createOrReplaceTempView("batch")
        sql(f"INSERT INTO {DB}.{name} PARTITION (dt) SELECT * FROM batch")


sql(f"CREATE DATABASE IF NOT EXISTS {DB}")
DB_LOC = sql(f"DESCRIBE DATABASE {DB}").where("info_name = 'Location'").collect()[0][1]
EXT_LOC = "s3a://warehouse/blog_hive_migration_ext/legacy_e"
for t in ("legacy_b", "legacy_b_BACKUP_", "legacy_c", "legacy_c_BACKUP_",
          "legacy_d", "legacy_e", "legacy_e_BACKUP_"):
    sql(f"DROP TABLE IF EXISTS {DB}.{t}")
    remove(f"{DB_LOC}/{t.lower()}")
remove(EXT_LOC)
hive_table("legacy_b")
loc = describe(f"{DB}.legacy_b", "Location")

print("== CALL spark_catalog.system.migrate ==")
sql(f"CALL spark_catalog.system.migrate('{DB}.legacy_b')").show()
print("tables now:", [r.tableName for r in sql(f"SHOW TABLES IN {DB} LIKE 'legacy_b*'").collect()])
p = props(f"{DB}.legacy_b")
print("legacy_b provider        :", describe(f"{DB}.legacy_b", "Provider"))
print("legacy_b 'migrated'      :", p.get("migrated"), "| gc.enabled:", p.get("gc.enabled", "(unset, default true)"))
print("legacy_b location        :", describe(f"{DB}.legacy_b", "Location"))
print("legacy_b_BACKUP_ provider:", describe(f"{DB}.legacy_b_BACKUP_", "Provider"))
print("legacy_b_BACKUP_ location:", describe(f"{DB}.legacy_b_BACKUP_", "Location"))
print("rows, Iceberg / backup   :", one(f"SELECT count(*) FROM {DB}.legacy_b"), "/",
      one(f"SELECT count(*) FROM {DB}.legacy_b_BACKUP_"))
f = sql(f"SELECT count(*) AS n, sum(CASE WHEN file_path LIKE '{loc}/dt=%' THEN 1 ELSE 0 END) AS hive_dirs "
        f"FROM {DB}.legacy_b.files").collect()[0]
print(f"Iceberg data files       : {f.n}, {f.hive_dirs} of them still in the Hive dt=... directories")
print("snapshots                :", one(f"SELECT count(*) FROM {DB}.legacy_b.snapshots"))
print("partition spec           :", sql(f"SELECT partition FROM {DB}.legacy_b.files LIMIT 1").collect()[0][0])

print("\n== rolling back ==")
sql(f"DROP TABLE {DB}.legacy_b")          # Iceberg drop without PURGE keeps the data files
print("backup rows after DROP TABLE:", one(f"SELECT count(*) FROM {DB}.legacy_b_BACKUP_"))
attempt("rename the backup back", f"ALTER TABLE {DB}.legacy_b_BACKUP_ RENAME TO {DB}.legacy_b")
# The dropped Iceberg table left its metadata/ folder at the old location. The
# data files live under the backup's directory, so that folder is safe to remove.
remove(loc)
sql(f"ALTER TABLE {DB}.legacy_b_BACKUP_ RENAME TO {DB}.legacy_b")
print("after the rename, legacy_b provider:", describe(f"{DB}.legacy_b", "Provider"),
      "| rows:", one(f"SELECT count(*) FROM {DB}.legacy_b"))
print("after the rename, legacy_b location:", describe(f"{DB}.legacy_b", "Location"))
print("tables now:", [r.tableName for r in sql(f"SHOW TABLES IN {DB} LIKE 'legacy_b*'").collect()])

print("\n== migrate again, then the expensive mistake ==")
sql(f"CALL spark_catalog.system.migrate('{DB}.legacy_b')").collect()
sql(f"DROP TABLE {DB}.legacy_b PURGE")
attempt("backup rows after DROP TABLE ... PURGE", f"SELECT count(*) FROM {DB}.legacy_b_BACKUP_")
attempt("backup files still listed", f"SELECT count(*) FROM (SELECT DISTINCT input_file_name() FROM {DB}.legacy_b_BACKUP_)")

print("\n== a table created with an explicit LOCATION ==")
hive_table("legacy_e", EXT_LOC)
print("legacy_e type            :", describe(f"{DB}.legacy_e", "Type"))
sql(f"CALL spark_catalog.system.migrate('{DB}.legacy_e')").collect()
print("legacy_e location        :", describe(f"{DB}.legacy_e", "Location"))
print("legacy_e_BACKUP_ location:", describe(f"{DB}.legacy_e_BACKUP_", "Location"))
print("Iceberg files still under the original directory:",
      one(f"SELECT count(*) FROM {DB}.legacy_e.files WHERE file_path LIKE '{EXT_LOC}/dt=%'"), "of",
      one(f"SELECT count(*) FROM {DB}.legacy_e.files"))
sql(f"DROP TABLE {DB}.legacy_e")
attempt("rename legacy_e's backup back", f"ALTER TABLE {DB}.legacy_e_BACKUP_ RENAME TO {DB}.legacy_e")
print("legacy_e after rollback:", describe(f"{DB}.legacy_e", "Provider"),
      "| rows:", one(f"SELECT count(*) FROM {DB}.legacy_e"))

print("\n== migrate without engine.hive.enabled, for the Hive read check ==")
hive_table("legacy_d")
sql(f"CALL spark_catalog.system.migrate(table => '{DB}.legacy_d', drop_backup => true)").collect()
print("legacy_d provider:", describe(f"{DB}.legacy_d", "Provider"),
      "| engine.hive.enabled:", props(f"{DB}.legacy_d").get("engine.hive.enabled", "(unset, default false)"))

print("\n== migrate with drop_backup, then improve the layout ==")
hive_table("legacy_c")
# engine.hive.enabled registers Hive's Iceberg storage handler, so Hive can still read it
sql(f"""CALL spark_catalog.system.migrate(
          table => '{DB}.legacy_c',
          properties => map('engine.hive.enabled', 'true'),
          drop_backup => true)""").show()
print("tables now:", [r.tableName for r in sql(f"SHOW TABLES IN {DB} LIKE 'legacy_c*'").collect()])
sql(f"""ALTER TABLE {DB}.legacy_c SET TBLPROPERTIES (
          'write.target-file-size-bytes' = '268435456',
          'write.object-storage.enabled' = 'true',
          'write.parquet.compression-codec' = 'zstd',
          'write.metadata.delete-after-commit.enabled' = 'true')""")
sql(f"ALTER TABLE {DB}.legacy_c WRITE ORDERED BY region, dt")
print("rewrite with the post's options:")
sql(f"""CALL spark_catalog.system.rewrite_data_files(
          table => '{DB}.legacy_c',
          options => map('min-input-files','5','target-file-size-bytes','268435456'))""") \
    .select("rewritten_data_files_count", "added_data_files_count").show()
print("rewrite with rewrite-all:")
sql(f"""CALL spark_catalog.system.rewrite_data_files(
          table => '{DB}.legacy_c',
          options => map('rewrite-all','true','target-file-size-bytes','268435456'))""") \
    .select("rewritten_data_files_count", "added_data_files_count").show()
sql(f"ALTER TABLE {DB}.legacy_c REPLACE PARTITION FIELD dt WITH days(event_ts)")
spark.createDataFrame(orders(500, "2026-01-08"), SCHEMA).createOrReplaceTempView("new_day")
sql(f"INSERT INTO {DB}.legacy_c SELECT * FROM new_day")
sql(f"SELECT spec_id, count(*) AS files, sum(record_count) AS rows "
    f"FROM {DB}.legacy_c.files GROUP BY spec_id ORDER BY spec_id").show()
print("rows:", one(f"SELECT count(*) FROM {DB}.legacy_c"))
spark.stop()

"""Move a Hive table into Iceberg one partition at a time with add_files.

Run inside the datalake-playground stack:
  docker cp add_files_in_pieces.py spark-master:/tmp/
  docker exec spark-master bash -c '$SPARK_HOME/bin/spark-submit --master "local[2]" \
    --driver-memory 1g --jars $(ls $ICEBERG_HOME/iceberg-spark-runtime*.jar) \
    /tmp/add_files_in_pieces.py'
"""
import random
from datetime import datetime, timedelta
from decimal import Decimal

from pyspark.sql import SparkSession

spark = (
    SparkSession.builder.appName("hive-add-files")
    .config("spark.sql.extensions",
            "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions")
    .config("spark.serializer", "org.apache.spark.serializer.KryoSerializer")
    .config("spark.sql.catalog.spark_catalog", "org.apache.iceberg.spark.SparkSessionCatalog")
    .config("spark.sql.catalog.spark_catalog.type", "hive")
    .config("spark.sql.catalog.spark_catalog.uri", "thrift://hive-metastore:9083")
    .config("spark.hadoop.hive.exec.dynamic.partition.mode", "nonstrict")
    .config("spark.eventLog.enabled", "false")
    .config("spark.ui.enabled", "false")
    .enableHiveSupport()
    .getOrCreate()
)
spark.sparkContext.setLogLevel("ERROR")
sql = spark.sql
DB = "blog_hive_migration"


def one(q):
    return sql(q).collect()[0][0]


def remove(path):
    jvm = spark._jvm
    p = jvm.org.apache.hadoop.fs.Path(path)
    p.getFileSystem(spark._jsc.hadoopConfiguration()).delete(p, True)


def attempt(label, q):
    try:
        result = sql(q).collect()
        print(f"{label}: {result[0][0] if result else 'ok'}")
    except Exception as e:
        msg = e.java_exception.getMessage() if hasattr(e, "java_exception") else str(e)
        # show the root cause when Spark wraps it in a job failure
        root = [l for l in msg.splitlines() if "FileNotFoundException" in l]
        line = root[0].split("): ", 1)[-1] if root else msg.splitlines()[0]
        print(f"{label}: {line[:230]}")


def orders(n, day):
    rng = random.Random(day)
    start = datetime.strptime(day, "%Y-%m-%d")
    return [(i, rng.randint(1, 900), Decimal(str(round(rng.uniform(5, 500), 2))),
             rng.choice(["north", "south", "east", "west"]),
             start + timedelta(seconds=rng.randint(0, 86_399)), day)
            for i in range(n)]


SCHEMA = ("order_id BIGINT, customer_id BIGINT, amount DECIMAL(10,2), "
          "region STRING, event_ts TIMESTAMP, dt STRING")

sql(f"CREATE DATABASE IF NOT EXISTS {DB}")
DB_LOC = sql(f"DESCRIBE DATABASE {DB}").where("info_name = 'Location'").collect()[0][1]
for t in ("sales", "legacy_sales"):
    sql(f"DROP TABLE IF EXISTS {DB}.{t}")
    remove(f"{DB_LOC}/{t}")
sql(f"""CREATE TABLE {DB}.legacy_sales (
          order_id BIGINT, customer_id BIGINT, amount DECIMAL(10,2),
          region STRING, event_ts TIMESTAMP)
        PARTITIONED BY (dt STRING) STORED AS PARQUET""")
for day in ("2026-01-05", "2026-01-06", "2026-01-07"):
    spark.createDataFrame(orders(2_000, day), SCHEMA).createOrReplaceTempView("batch")
    sql(f"INSERT INTO {DB}.legacy_sales PARTITION (dt) SELECT * FROM batch")

sql(f"""CREATE TABLE {DB}.sales (
          order_id BIGINT, customer_id BIGINT, amount DECIMAL(10,2),
          region STRING, event_ts TIMESTAMP, dt STRING)
        USING iceberg PARTITIONED BY (dt)""")

print("== add one partition ==")
sql(f"""CALL spark_catalog.system.add_files(
          table => '{DB}.sales',
          source_table => '{DB}.legacy_sales',
          partition_filter => map('dt', '2026-01-06'))""").show()
print("Iceberg :", sql(f"SELECT count(*), sum(amount) FROM {DB}.sales WHERE dt = '2026-01-06'").collect()[0])
print("Hive    :", sql(f"SELECT count(*), sum(amount) FROM {DB}.legacy_sales WHERE dt = '2026-01-06'").collect()[0])
print("files shared:", one(f"SELECT count(*) FROM {DB}.sales.files WHERE file_path LIKE '{DB_LOC}/legacy_sales/%'"),
      "of", one(f"SELECT count(*) FROM {DB}.sales.files"))

print("\n== the same partition again ==")
attempt("add_files twice", f"""CALL spark_catalog.system.add_files(
          table => '{DB}.sales', source_table => '{DB}.legacy_sales',
          partition_filter => map('dt', '2026-01-06'))""")

print("\n== the rest, one partition at a time ==")
for day in ("2026-01-05", "2026-01-07"):
    r = sql(f"""CALL spark_catalog.system.add_files(
              table => '{DB}.sales', source_table => '{DB}.legacy_sales',
              partition_filter => map('dt', '{day}'))""").collect()[0]
    print(f"{day}: added {r.added_files_count} files")
print("rows, Iceberg / Hive:", one(f"SELECT count(*) FROM {DB}.sales"), "/",
      one(f"SELECT count(*) FROM {DB}.legacy_sales"))
print("gc.enabled on the Iceberg table:",
      {r[0]: r[1] for r in sql(f"SHOW TBLPROPERTIES {DB}.sales").collect()}.get("gc.enabled", "(unset, default true)"))

print("\n== dropping the Iceberg table with PURGE ==")
sql(f"DROP TABLE {DB}.sales PURGE")
attempt("Hive rows afterwards", f"SELECT count(*) FROM {DB}.legacy_sales")
spark.stop()