Hidden partitioning and partition evolution: changing layout without rewriting history

"""Hidden partitioning on Iceberg: derived partition values, pruning, and bucket.

Run inside the datalake-playground stack:
  docker cp hidden_partitioning.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/hidden_partitioning.py'
"""
import json
from datetime import datetime, timedelta

from pyspark.sql import SparkSession
from pyspark.sql.types import (DoubleType, LongType, StringType, StructField,
                               StructType, TimestampType)

spark = (
    SparkSession.builder.appName("iceberg-hidden-partitioning")
    .config("spark.sql.extensions",
            "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions")
    .config("spark.serializer", "org.apache.spark.serializer.KryoSerializer")
    .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.sql.catalog.ice.warehouse", "s3a://warehouse/")
    .config("spark.sql.session.timeZone", "UTC")
    .config("spark.eventLog.enabled", "false")
    .config("spark.ui.enabled", "false")
    .getOrCreate()
)
spark.sparkContext.setLogLevel("ERROR")
sql = spark.sql
NS = "ice.blog_partitioning"
EVENTS = f"{NS}.events"
ORDERS = f"{NS}.orders"
jvm = spark._jvm
conf = spark._jsc.hadoopConfiguration()


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


def partition_dirs(table):
    """Directories that hold at least one data file (S3A can keep empty markers)."""
    path = jvm.org.apache.hadoop.fs.Path(location(table) + "/data")
    files = path.getFileSystem(conf).listFiles(path, True)
    dirs = set()
    while files.hasNext():
        dirs.add(files.next().getPath().getParent().getName())
    return sorted(dirs)


def partition_specs(table):
    """The partition-specs array, read from the table's current metadata.json."""
    meta = sql(f"SELECT file FROM {table}.metadata_log_entries "
               f"ORDER BY timestamp DESC LIMIT 1").collect()[0][0]
    path = jvm.org.apache.hadoop.fs.Path(meta)
    stream = path.getFileSystem(conf).open(path)
    text = jvm.org.apache.commons.io.IOUtils.toString(stream, "UTF-8")
    stream.close()
    doc = json.loads(text)
    schema = next(x for x in doc["schemas"] if x["schema-id"] == doc["current-schema-id"])
    field_ids = {f["name"]: f["id"] for f in schema["fields"]}
    return doc["partition-specs"], doc["default-spec-id"], field_ids


def scan(df):
    """Collect df, then read the Iceberg scan node's metrics from the plan that ran."""
    qe = df._jdf.queryExecution()
    rows = df._jdf.collectAsList().size()
    plan = qe.executedPlan()
    if plan.getClass().getSimpleName() == "AdaptiveSparkPlanExec":
        plan = plan.executedPlan()
    metrics = plan.collectLeaves().apply(0).metrics()
    get = lambda k: metrics.apply(k).value()
    return rows, get("resultDataFiles"), get("skippedDataFiles")


sql(f"CREATE NAMESPACE IF NOT EXISTS {NS}")
sql(f"DROP TABLE IF EXISTS {EVENTS} PURGE")
sql(f"DROP TABLE IF EXISTS {ORDERS} PURGE")

# ------------------------------------------------ 1. a derived day partition
sql(f"""
    CREATE TABLE {EVENTS} (
        id BIGINT, event_ts TIMESTAMP, region STRING, amount DOUBLE)
    USING iceberg PARTITIONED BY (days(event_ts))
""")
schema = StructType([
    StructField("id", LongType(), False),
    StructField("event_ts", TimestampType(), False),
    StructField("region", StringType(), False),
    StructField("amount", DoubleType(), False),
])
regions = ["us-east", "us-west", "eu-central", "ap-south"]
start = datetime(2026, 1, 5)
rows = [(i, start + timedelta(days=i // 3000, seconds=(i % 3000) * 20),
         regions[i % 4], round(10 + (i % 500) * 0.37, 2)) for i in range(9000)]
spark.createDataFrame(rows, schema).writeTo(EVENTS).append()

specs, default_spec, field_ids = partition_specs(EVENTS)
print("== 1. hidden partitioning ==")
print("partition-specs:", json.dumps(specs))
print("field id of event_ts:", field_ids["event_ts"])
print("partition directories on disk:", partition_dirs(EVENTS))
print("data files in table:", sql(f"SELECT count(*) FROM {EVENTS}.files").collect()[0][0])

one_day = spark.table(EVENTS).where(
    "event_ts >= timestamp'2026-01-06 00:00:00' AND event_ts < timestamp'2026-01-07 00:00:00'")
rows_, read, skipped = scan(one_day)
print(f"rows matched: {rows_} | data files read: {read} | skipped by pruning: {skipped}")

# The partitions metadata table, the numbers the layout decision rests on
print("\n== 2. partitions per layout ==")
sql(f"""
    SELECT count(*) AS partitions,
           avg(file_count) AS avg_files,
           avg(record_count) AS avg_rows
    FROM {EVENTS}.partitions
""").show()

# ------------------------------------------------------------ 3. bucket(64)
sql(f"""
    CREATE TABLE {ORDERS} (order_id BIGINT, customer_id BIGINT, amount DOUBLE)
    USING iceberg
    PARTITIONED BY (bucket(64, customer_id))
""")
orders = [(100000 + i, 5000 + (i * 7919) % 20000, round(5 + (i % 300) * 1.1, 2))
          for i in range(20000)]
spark.createDataFrame(orders, "order_id BIGINT, customer_id BIGINT, amount DOUBLE") \
    .writeTo(ORDERS).append()
print("== 3. bucket(64, customer_id) ==")
print("distinct customers:", spark.table(ORDERS).select("customer_id").distinct().count())
sql(f"""
    SELECT count(*) AS partitions, min(record_count) AS min_rows,
           max(record_count) AS max_rows
    FROM {ORDERS}.partitions
""").show()
eq = spark.table(ORDERS).where("customer_id = 12345")
rng = spark.table(ORDERS).where("customer_id BETWEEN 12000 AND 12100")
print("data files in table:", sql(f"SELECT count(*) FROM {ORDERS}.files").collect()[0][0])
rows_, read, skipped = scan(eq)
print(f"equality filter: rows {rows_} | data files read: {read} | skipped: {skipped}")
rows_, read, skipped = scan(rng)
print(f"range filter   : rows {rows_} | data files read: {read} | skipped: {skipped}")

# ------------------------------------------ 4. two fields multiply the count
sql(f"DROP TABLE IF EXISTS {NS}.events_two_fields PURGE")
sql(f"""
    CREATE TABLE {NS}.events_two_fields (
        id BIGINT, event_ts TIMESTAMP, customer_id BIGINT)
    USING iceberg PARTITIONED BY (days(event_ts), bucket(16, customer_id))
""")
# One append per day. Writing the whole month at once opens a Parquet writer for
# every one of its 496 partitions in one task, which ran a 1 GB driver out of heap.
for day in range(31):
    rows_day = [(day * 480 + i, datetime(2026, 1, 1 + day) + timedelta(minutes=i * 3),
                 5000 + ((day * 480 + i) * 7919) % 20000) for i in range(480)]
    spark.createDataFrame(rows_day, "id BIGINT, event_ts TIMESTAMP, customer_id BIGINT") \
        .writeTo(f"{NS}.events_two_fields").append()
print("\n== 4. days(event_ts) x bucket(16, customer_id), one month ==")
sql(f"SELECT count(*) AS partitions, avg(record_count) AS avg_rows "
    f"FROM {NS}.events_two_fields.partitions").show()

# ------------------------------------------------ 5. sort order, not region
sql(f"ALTER TABLE {EVENTS} WRITE ORDERED BY region, event_ts")
print("== 5. write order ==")
print("sort-order after WRITE ORDERED BY:",
      spark._jvm.org.apache.iceberg.spark.Spark3Util.loadIcebergTable(
          spark._jsparkSession, EVENTS).sortOrder().toString().replace("\n", " "))

spark.stop()

"""Partition evolution on Iceberg: months to days without rewriting, then a rewrite.

Run inside the datalake-playground stack:
  docker cp partition_evolution.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/partition_evolution.py'
"""
import json
from datetime import datetime, timedelta

from pyspark.sql import SparkSession

spark = (
    SparkSession.builder.appName("iceberg-partition-evolution")
    .config("spark.sql.extensions",
            "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions")
    .config("spark.serializer", "org.apache.spark.serializer.KryoSerializer")
    .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.sql.catalog.ice.warehouse", "s3a://warehouse/")
    .config("spark.sql.session.timeZone", "UTC")
    .config("spark.eventLog.enabled", "false")
    .config("spark.ui.enabled", "false")
    .getOrCreate()
)
spark.sparkContext.setLogLevel("ERROR")
sql = spark.sql
NS = "ice.blog_partitioning"
TX = f"{NS}.transactions"
EV = f"{NS}.events_evolving"
DAILY = f"{NS}.events_daily"
jvm = spark._jvm
conf = spark._jsc.hadoopConfiguration()


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


def partition_dirs(table):
    """Directories that hold at least one data file (S3A can keep empty markers)."""
    path = jvm.org.apache.hadoop.fs.Path(location(table) + "/data")
    files = path.getFileSystem(conf).listFiles(path, True)
    dirs = set()
    while files.hasNext():
        dirs.add(files.next().getPath().getParent().getName())
    return sorted(dirs)


def metadata(table):
    meta = sql(f"SELECT file FROM {table}.metadata_log_entries "
               f"ORDER BY timestamp DESC LIMIT 1").collect()[0][0]
    path = jvm.org.apache.hadoop.fs.Path(meta)
    stream = path.getFileSystem(conf).open(path)
    text = jvm.org.apache.commons.io.IOUtils.toString(stream, "UTF-8")
    stream.close()
    return json.loads(text)


def scan(df):
    """Collect df, then read the Iceberg scan node's metrics from the plan that ran."""
    qe = df._jdf.queryExecution()
    rows = df._jdf.collectAsList().size()
    plan = qe.executedPlan()
    if plan.getClass().getSimpleName() == "AdaptiveSparkPlanExec":
        plan = plan.executedPlan()
    metrics = plan.collectLeaves().apply(0).metrics()
    return rows, metrics.apply("resultDataFiles").value(), metrics.apply("skippedDataFiles").value()


def attempt(label, statement):
    try:
        sql(statement)
        print(f"  OK    {label}")
    except Exception as e:
        jvm_e = getattr(e, "java_exception", None)
        cls = jvm_e.getClass().getName() if jvm_e is not None else type(e).__name__
        msg = " ".join((jvm_e.getMessage() if jvm_e is not None else str(e)).split())[:220]
        print(f"  ERROR {label}: {cls}: {msg}")


def load(table, first_day, days, per_day, start_id):
    rows = []
    for d in range(days):
        for i in range(per_day):
            rows.append((start_id + d * per_day + i,
                         first_day + timedelta(days=d, minutes=i * 4),
                         round(20 + (i % 90) * 1.5, 2)))
    spark.createDataFrame(rows, "txn_id BIGINT, ts TIMESTAMP, amount DOUBLE") \
        .writeTo(table).append()


sql(f"CREATE NAMESPACE IF NOT EXISTS {NS}")
sql(f"DROP TABLE IF EXISTS {TX} PURGE")
sql(f"DROP TABLE IF EXISTS {EV} PURGE")
sql(f"DROP TABLE IF EXISTS {DAILY} PURGE")

# ------------------------------------------- 1. months, then replace with days
sql(f"""CREATE TABLE {TX} (txn_id BIGINT, ts TIMESTAMP, amount DOUBLE)
        USING iceberg PARTITIONED BY (months(ts))""")
load(TX, datetime(2026, 1, 12), days=1, per_day=300, start_id=1)
load(TX, datetime(2026, 2, 16), days=1, per_day=300, start_id=301)
print("== 1. partition evolution ==")
print("dirs after months():", partition_dirs(TX))

sql(f"ALTER TABLE {TX} REPLACE PARTITION FIELD months(ts) WITH days(ts)")
load(TX, datetime(2026, 3, 9), days=1, per_day=300, start_id=601)
print("dirs after days()  :", partition_dirs(TX))

doc = metadata(TX)
for spec in doc["partition-specs"]:
    print("  spec", spec["spec-id"], json.dumps(spec["fields"]))
print("default-spec-id:", doc["default-spec-id"])
print("spec_id per data file:",
      [r[0] for r in sql(f"SELECT spec_id FROM {TX}.files ORDER BY spec_id").collect()])
print("rows total:", sql(f"SELECT count(*) FROM {TX}").collect()[0][0])

# One predicate on ts, pruned under whichever spec wrote each file
for label, where in [
        ("January 12 (month spec)", "ts >= timestamp'2026-01-12 00:00:00' AND ts < timestamp'2026-01-13 00:00:00'"),
        ("March 9 (day spec)     ", "ts >= timestamp'2026-03-09 00:00:00' AND ts < timestamp'2026-03-10 00:00:00'")]:
    rows, read, skipped = scan(spark.table(TX).where(where))
    print(f"  {label}: rows {rows} | data files read {read} | skipped {skipped}")

# ------------------------------------------------ 2. add, replace, drop fields
EVENTS_SCHEMA = "id BIGINT, event_ts TIMESTAMP, customer_id BIGINT"
print("\n== 2. evolving in practice ==")

# add a finer time field to a table already partitioned by days(event_ts)
sql(f"CREATE TABLE {DAILY} ({EVENTS_SCHEMA}) USING iceberg PARTITIONED BY (days(event_ts))")
attempt("ADD PARTITION FIELD hours(event_ts) on a days(event_ts) table",
        f"ALTER TABLE {DAILY} ADD PARTITION FIELD hours(event_ts)")

# replace months with days, and drop a bucket field, on a months + bucket table
sql(f"""CREATE TABLE {EV} ({EVENTS_SCHEMA})
        USING iceberg PARTITIONED BY (months(event_ts), bucket(16, customer_id))""")
spark.createDataFrame(
    [(i, datetime(2026, 1, 1) + timedelta(hours=i), 5000 + (i * 7919) % 20000) for i in range(1000)],
    EVENTS_SCHEMA).writeTo(EV).append()
attempt("REPLACE PARTITION FIELD months(event_ts) WITH days(event_ts)",
        f"ALTER TABLE {EV} REPLACE PARTITION FIELD months(event_ts) WITH days(event_ts)")
attempt("DROP PARTITION FIELD bucket(16, customer_id)",
        f"ALTER TABLE {EV} DROP PARTITION FIELD bucket(16, customer_id)")
spark.createDataFrame(
    [(1000 + i, datetime(2026, 3, 1) + timedelta(hours=i), 5000 + i) for i in range(200)],
    EVENTS_SCHEMA).writeTo(EV).append()
for spec in metadata(EV)["partition-specs"]:
    print("  spec", spec["spec-id"], [f"{f['transform']}({f['source-id']})" for f in spec["fields"]])
sql(f"""SELECT spec_id, count(*) AS files, sum(record_count) AS records
        FROM {EV}.files GROUP BY spec_id ORDER BY spec_id""").show()

# ---------------------------------------- 3. rewrite old data into the new spec
print("== 3. rewrite_data_files ==")
attempt("where => 'event_ts < date'2026-03-01'' (nested single quotes)",
        f"CALL ice.system.rewrite_data_files(table => 'blog_partitioning.events_evolving', "
        f"where => 'event_ts < date'2026-03-01'')")
sql("""CALL ice.system.rewrite_data_files(
         table => 'blog_partitioning.events_evolving',
         where => "event_ts < TIMESTAMP '2026-03-01 00:00:00'")""") \
    .select("rewritten_data_files_count", "added_data_files_count").show()
sql(f"""SELECT spec_id, count(*) AS files, sum(record_count) AS records
        FROM {EV}.files GROUP BY spec_id ORDER BY spec_id""").show()

spark.stop()