COCO-Caption Experiment #0 — Lakehouse Demo

cd data-plane
docker compose up -d   # RustFS :29000, PostgreSQL :55432

import os, sys
# Ensure data-plane root is on the path when running from notebooks/
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath('.')), ''))
 
# Use internal Docker endpoint when running inside a container,
# host endpoint when running on the host.
S3_ENDPOINT = os.getenv('MINIO_ENDPOINT', 'http://localhost:29000')
S3_KEY      = os.getenv('MINIO_ACCESS_KEY', 'minio')
S3_SECRET   = os.getenv('MINIO_SECRET_KEY', 'minio123')
# DUCKLAKE_POSTGRES_DSN is set by docker-compose (uses postgresql:5432 service name);
# DUCKLAKE_DSN is the fallback for host-side execution.
PG_DSN = (
    os.getenv('DUCKLAKE_POSTGRES_DSN')
    or os.getenv('DUCKLAKE_DSN', 'postgresql://ducklake:123456@localhost:55432/ducklake')
)
 
BUCKET  = 'landing'
PREFIX  = 'coco-caption/notebook'
S3_URI  = f's3://{BUCKET}/{PREFIX}/**/*.parquet'
MAX_SAMPLES = 500
 
print(f'S3 endpoint: {S3_ENDPOINT}')
print(f'S3 URI: {S3_URI}')
print(f'PG DSN: {PG_DSN}')

from lakehouse.sync import sync_from_hf
import boto3
 
result = sync_from_hf(
    dataset_id='yerevann/coco-karpathy',
    split='restval',
    bucket=BUCKET,
    prefix=PREFIX,
    max_samples=MAX_SAMPLES,
    s3_endpoint_url=S3_ENDPOINT,
    s3_access_key=S3_KEY,
    s3_secret_key=S3_SECRET,
)
print(f'Files uploaded : {result.files_uploaded}')
print(f'Rows           : {result.rows_total:,}')
print(f'Bytes          : {result.bytes_total / 1024:.1f} KiB')
print(f'Destination    : s3://{result.bucket}/{result.prefix}/')
 
# Verify: list objects in bucket under the prefix
print()
print(f'Bucket contents  s3://{BUCKET}/{PREFIX}/')
s3 = boto3.client(
    's3',
    endpoint_url=S3_ENDPOINT,
    aws_access_key_id=S3_KEY,
    aws_secret_access_key=S3_SECRET,
)
resp = s3.list_objects_v2(Bucket=BUCKET, Prefix=PREFIX)
objects = resp.get('Contents', [])
if not objects:
    print('  (empty)')
else:
    for obj in objects:
        size_kib = obj['Size'] / 1024
        print(f"  {obj['Key']}  ({size_kib:.1f} KiB)")
print(f'Total objects: {len(objects)}')

from lakehouse.catalog import LakehouseCatalog
 
s3_host = S3_ENDPOINT.replace('http://', '').replace('https://', '')
 
try:
    with LakehouseCatalog(
        PG_DSN,
        data_path='s3://warehouse',
        s3_endpoint=s3_host,
        s3_access_key=S3_KEY,
        s3_secret_key=S3_SECRET,
    ) as cat:
        exp = cat.register_experiment('coco-caption-nb', 'Notebook demo — Experiment #0')
        run = cat.register_run(exp.experiment_id, sim_index=0,
                               config={'split': 'restval', 'max_samples': MAX_SAMPLES},
                               s3_prefix=PREFIX)
        run = cat.complete_run(run.run_id)
    print(f'Experiment : {exp.experiment_id}')
    print(f'Run        : {run.run_id}  status={run.status}')
except Exception as exc:
    print(f'⚠ DuckLake catalog skipped ({exc.__class__.__name__}): {exc}')

from lakehouse.query import LakehouseQuery
import pandas as pd
 
q = LakehouseQuery(s3_endpoint=s3_host, s3_access_key=S3_KEY, s3_secret_key=S3_SECRET)
 
# Schema
schema_df = q.execute(f"DESCRIBE SELECT * FROM read_parquet('{S3_URI}') LIMIT 0").to_pandas()
print('Schema:')
print(schema_df[['column_name', 'column_type']].to_string(index=False))

# Row count
total = q.execute(f"SELECT COUNT(*) FROM read_parquet('{S3_URI}')").to_arrow().column(0)[0].as_py()
print(f'Total rows: {total:,}')

# Caption length distribution
hist_df = q.execute(f"""
    SELECT
        CASE
            WHEN len(sentences[1]) < 40 THEN '<40'
            WHEN len(sentences[1]) < 60 THEN '40-59'
            WHEN len(sentences[1]) < 80 THEN '60-79'
            ELSE '>=80'
        END AS caption_length_bucket,
        COUNT(*) AS images,
        ROUND(100.0 * COUNT(*) / SUM(COUNT(*)) OVER (), 1) AS pct
    FROM read_parquet('{S3_URI}')
    GROUP BY 1 ORDER BY 1
""").to_pandas()
print('Caption length distribution:')
hist_df

# Captions per image
cpimg_df = q.execute(f"""
    SELECT len(sentences) AS num_captions, COUNT(*) AS images
    FROM read_parquet('{S3_URI}')
    GROUP BY 1 ORDER BY 1
""").to_pandas()
print('Captions per image:')
cpimg_df

# Sample rows (predicate pushdown demo)
sample_df = q.execute(f"""
    SELECT cocoid, filename, sentences[1] AS caption
    FROM read_parquet('{S3_URI}')
    WHERE split = 'restval'
    LIMIT 5
""").to_pandas()
sample_df['caption'] = sample_df['caption'].str[:80]
sample_df

from lakehouse.stream import as_iterable_dataset
 
stream_sql = f"""
    SELECT cocoid, filename, sentences[1] AS caption
    FROM read_parquet('{S3_URI}') LIMIT 5
"""
 
ds = as_iterable_dataset(
    source=S3_URI, sql=stream_sql,
    s3_endpoint=s3_host, s3_access_key=S3_KEY, s3_secret_key=S3_SECRET,
)
print(f'Type: {type(ds).__name__}  (identical API to load_dataset(..., streaming=True))')
print()
for row in ds:
    print(f"{row['cocoid']}  {row['filename']}")
    print(f"  {row['caption'][:80]}")

from lakehouse.stream import stream_batches
 
shuffle_sql = f"SELECT cocoid FROM read_parquet('{S3_URI}') LIMIT 8"
 
normal   = [b.column('cocoid')[i].as_py()
            for b in stream_batches(source=S3_URI, sql=shuffle_sql,
                                    s3_endpoint=s3_host, s3_access_key=S3_KEY, s3_secret_key=S3_SECRET)
            for i in range(b.num_rows)]
 
shuffled = [b.column('cocoid')[i].as_py()
            for b in stream_batches(source=S3_URI, sql=shuffle_sql, shuffle=True,
                                    s3_endpoint=s3_host, s3_access_key=S3_KEY, s3_secret_key=S3_SECRET)
            for i in range(b.num_rows)]
 
print(f'Normal   order: {normal}')
print(f'Shuffled order: {shuffled}')
print(f'Same rows, different order: {sorted(normal) == sorted(shuffled)}')

try:
    import wandb
    cap_df = q.execute(f"""
        SELECT len(sentences[1]) AS caption_length
        FROM read_parquet('{S3_URI}') LIMIT 50
    """).to_pandas()
 
    run = wandb.init(project='auraison-coco-nb', name='coco-nb-exp0', reinit=True, mode='offline')
    for step_i, row in cap_df.iterrows():
        wandb.log({'caption_length': row['caption_length']}, step=int(step_i))
    run.finish()
    print(f'Logged {len(cap_df)} steps to W&B (offline). Sync with: wandb sync <run-dir>')
except Exception as exc:
    print(f'⚠ W&B skipped: {exc}')

from lakehouse.tools import LakehouseDeps, QualityReport
 
deps = LakehouseDeps(s3_endpoint=s3_host, s3_access_key=S3_KEY, s3_secret_key=S3_SECRET)
qe = deps.query_engine()
 
total_rows = qe.execute(f"SELECT COUNT(*) FROM read_parquet('{S3_URI}')").to_arrow().column(0)[0].as_py()
schema_tbl = qe.execute(f"DESCRIBE SELECT * FROM read_parquet('{S3_URI}') LIMIT 0").to_arrow()
columns    = schema_tbl.column('column_name').to_pylist()
null_counts = {
    col: qe.execute(f'SELECT COUNT(*) FROM read_parquet("{S3_URI}") WHERE "{col}" IS NULL'
                    ).to_arrow().column(0)[0].as_py()
    for col in columns
}
distinct_counts = {
    col: qe.execute(f'SELECT COUNT(DISTINCT "{col}") FROM read_parquet("{S3_URI}")'
                    ).to_arrow().column(0)[0].as_py()
    for col in columns
}
report = QualityReport(
    source=S3_URI, total_rows=total_rows, total_columns=len(columns),
    null_counts=null_counts, distinct_counts=distinct_counts, sample_rows=[]
)
 
summary_df = pd.DataFrame({
    'column': columns,
    'nulls': [null_counts[c] for c in columns],
    'distinct': [distinct_counts[c] for c in columns],
})
print(f'Total rows    : {report.total_rows:,}')
print(f'Total columns : {report.total_columns}')
print()
summary_df

# Ingest COCO parquet into DuckLake warehouse table, then evolve the schema.
# The underlying Parquet files in the lakehouse are never rewritten.
try:
    with LakehouseCatalog(
        PG_DSN, data_path='s3://warehouse',
        s3_endpoint=s3_host, s3_access_key=S3_KEY, s3_secret_key=S3_SECRET,
    ) as cat:
        con = cat._con
        con.execute('DROP TABLE IF EXISTS coco_caption')
        con.execute(f"CREATE TABLE coco_caption AS SELECT * FROM read_parquet('{S3_URI}')")
        v1_count = con.execute('SELECT COUNT(*) FROM coco_caption').fetchone()[0]
        print(f'Table created: {v1_count:,} rows  (snapshot v1)')
 
        # Add derived column — zero Parquet rewrite
        con.execute('ALTER TABLE coco_caption ADD COLUMN caption_length INTEGER')
        con.execute('UPDATE coco_caption SET caption_length = len(sentences[1])')
        print('ALTER TABLE ADD COLUMN caption_length  ✓  (snapshot v3)')
 
        top5 = con.execute(
            'SELECT cocoid, filename, caption_length '
            'FROM coco_caption ORDER BY caption_length DESC LIMIT 5'
        ).df()
        print('Top-5 longest captions after schema evolution:')
        print(top5.to_string(index=False))
except Exception as exc:
    print(f'⚠ DuckLake skipped ({exc.__class__.__name__}): {exc}')

# Every DDL/DML creates a new catalog snapshot.
# VERSION 1 = just after CREATE TABLE (no caption_length column yet).
try:
    with LakehouseCatalog(
        PG_DSN, data_path='s3://warehouse',
        s3_endpoint=s3_host, s3_access_key=S3_KEY, s3_secret_key=S3_SECRET,
    ) as cat:
        con = cat._con
        v1 = con.execute(
            'SELECT cocoid, filename, split FROM coco_caption AT (VERSION => 1) LIMIT 5'
        ).df()
        print('coco_caption AT VERSION 1 — before schema evolution:')
        print(v1.to_string(index=False))
        print(f'Columns at v1: {list(v1.columns)}')
 
        cur = con.execute(
            'SELECT cocoid, filename, caption_length FROM coco_caption LIMIT 5'
        ).df()
        print()
        print('Current snapshot — caption_length present:')
        print(cur.to_string(index=False))
except Exception as exc:
    print(f'⚠ DuckLake skipped ({exc.__class__.__name__}): {exc}')

# Join simulation_runs (DuckLake catalog) with coco_caption (lakehouse table)
# — no ETL, no separate MLflow lookup.
try:
    with LakehouseCatalog(
        PG_DSN, data_path='s3://warehouse',
        s3_endpoint=s3_host, s3_access_key=S3_KEY, s3_secret_key=S3_SECRET,
    ) as cat:
        con = cat._con
        joined = con.execute("""
            SELECT
                r.config->>'split'        AS split,
                r.config->>'max_samples'  AS max_samples,
                r.status,
                COUNT(c.cocoid)           AS images_in_lakehouse,
                ROUND(AVG(c.caption_length), 1) AS avg_caption_len
            FROM simulation_runs r
            JOIN coco_caption c ON r.s3_prefix = 'coco-caption/notebook'
            GROUP BY 1, 2, 3
        """).df()
        print('simulation_runs ⋈ coco_caption:')
        print(joined.to_string(index=False))
except Exception as exc:
    print(f'⚠ DuckLake skipped ({exc.__class__.__name__}): {exc}')

import time
 
# EXPLAIN reveals that DuckDB pushes the WHERE filter into the Parquet reader
# (PARQUET_SCAN → row-group pruning), avoiding a full scan.
plan = q.execute(f"""
    EXPLAIN SELECT cocoid, filename, sentences[1] AS caption
    FROM read_parquet('{S3_URI}')
    WHERE split = 'restval'
""").to_pandas()
print(plan['explain_value'].iloc[0])
 
# Timed comparison
t0 = time.perf_counter()
n_filtered = q.execute(
    f"SELECT COUNT(*) FROM read_parquet('{S3_URI}') WHERE split = 'restval'"
).to_arrow().column(0)[0].as_py()
t1 = time.perf_counter()
n_all = q.execute(
    f"SELECT COUNT(*) FROM read_parquet('{S3_URI}')"
).to_arrow().column(0)[0].as_py()
t2 = time.perf_counter()
print(f'Filtered scan ({n_filtered:,} rows): {(t1-t0)*1000:.1f} ms')
print(f'Full scan     ({n_all:,} rows): {(t2-t1)*1000:.1f} ms')
print('(On a 100 GB COCO dataset the gap widens from ms → seconds)')

# Anti-join INSERT: only rows whose cocoid is not already in the lakehouse.
# Running twice must produce the same final row count — MVCC guarantees
# concurrent readers see a consistent snapshot during the insert.
try:
    with LakehouseCatalog(
        PG_DSN, data_path='s3://warehouse',
        s3_endpoint=s3_host, s3_access_key=S3_KEY, s3_secret_key=S3_SECRET,
    ) as cat:
        con = cat._con
        before = con.execute('SELECT COUNT(*) FROM coco_caption').fetchone()[0]
 
        idempotent_insert = f"""
            INSERT INTO coco_caption
            SELECT src.*, len(src.sentences[1]) AS caption_length
            FROM read_parquet('{S3_URI}') src
            LEFT JOIN coco_caption dst ON src.cocoid = dst.cocoid
            WHERE dst.cocoid IS NULL
        """
        con.execute(idempotent_insert)
        after1 = con.execute('SELECT COUNT(*) FROM coco_caption').fetchone()[0]
        con.execute(idempotent_insert)  # second run
        after2 = con.execute('SELECT COUNT(*) FROM coco_caption').fetchone()[0]
 
        print(f'Before ingest : {before:,} rows')
        print(f'After 1st run : {after1:,} rows  (+{after1-before})')
        print(f'After 2nd run : {after2:,} rows  (+{after2-after1}) ← idempotent')
except Exception as exc:
    print(f'⚠ DuckLake skipped ({exc.__class__.__name__}): {exc}')

# No separate analytics database needed — DuckDB runs columnar OLAP
# over the same Parquet files used for training.
 
# Caption word-count distribution
wc_df = q.execute(f"""
    SELECT
        len(string_split(sentences[1], ' '))  AS word_count,
        COUNT(*)                               AS images
    FROM read_parquet('{S3_URI}')
    GROUP BY 1
    ORDER BY word_count
""").to_pandas()
print('Caption word-count distribution:')
print(wc_df.to_string(index=False))
 
# Multi-caption images (annotator coverage)
multi_df = q.execute(f"""
    SELECT
        len(sentences)        AS num_captions,
        COUNT(*)              AS images,
        ROUND(AVG(len(sentences[1])), 1)   AS avg_caption_chars
    FROM read_parquet('{S3_URI}')
    GROUP BY 1
    ORDER BY num_captions DESC
""").to_pandas()
print()
print('Images by caption count (annotator coverage):')
print(multi_df.to_string(index=False))

from lakehouse.stream import as_iterable_dataset
 
SPLITS = {
    'low':    "len(sentences[1]) < 40",
    'medium': "len(sentences[1]) BETWEEN 40 AND 59",
    'large':  "len(sentences[1]) >= 60",
}
 
# Catalog registration + streaming dataset creation for each tier
split_datasets = {}  # held open for downstream training loops
 
try:
    with LakehouseCatalog(
        PG_DSN, data_path='s3://warehouse',
        s3_endpoint=s3_host, s3_access_key=S3_KEY, s3_secret_key=S3_SECRET,
    ) as cat:
        for tier, predicate in SPLITS.items():
            # Count via predicate-pushdown — no full scan
            row_count = q.execute(
                f"SELECT COUNT(*) FROM read_parquet('{S3_URI}') WHERE {predicate}"
            ).to_arrow().column(0)[0].as_py()
 
            # Register experiment + completed run in DuckLake catalog
            exp = cat.register_experiment(
                f'coco-caption-{tier}',
                f'Caption-length split: {tier} ({predicate})'
            )
            run = cat.register_run(
                exp.experiment_id, sim_index=0,
                config={'tier': tier, 'predicate': predicate, 'rows': row_count},
                s3_prefix=PREFIX,
            )
            run = cat.complete_run(run.run_id)
 
            # Build IterableDataset — predicate pushed into Parquet reader
            sql = f"""
                SELECT cocoid, filename, sentences[1] AS caption
                FROM read_parquet('{S3_URI}')
                WHERE {predicate}
            """
            ds = as_iterable_dataset(
                source=S3_URI, sql=sql,
                s3_endpoint=s3_host, s3_access_key=S3_KEY, s3_secret_key=S3_SECRET,
            )
            split_datasets[tier] = ds
 
            # Sample one row to verify the stream is live
            sample = next(iter(ds))
            print(f'[{tier:6}]  rows={row_count:3d}'
                  f'  exp={exp.experiment_id[:8]}  run={run.run_id[:8]}')
            print(f'           "{sample["caption"][:72]}"')
 
except Exception as exc:
    print(f'⚠ DuckLake skipped ({exc.__class__.__name__}): {exc}')
    # Fallback: build datasets without catalog tracking
    for tier, predicate in SPLITS.items():
        sql = f"""
            SELECT cocoid, filename, sentences[1] AS caption
            FROM read_parquet('{S3_URI}')
            WHERE {predicate}
        """
        split_datasets[tier] = as_iterable_dataset(
            source=S3_URI, sql=sql,
            s3_endpoint=s3_host, s3_access_key=S3_KEY, s3_secret_key=S3_SECRET,
        )
        row_count = q.execute(
            f"SELECT COUNT(*) FROM read_parquet('{S3_URI}') WHERE {predicate}"
        ).to_arrow().column(0)[0].as_py()
        print(f'[{tier:6}]  rows={row_count:3d}  (no catalog — PostgreSQL not running)')
 
print(f'\nDatasets ready: {list(split_datasets)}')
print('Each is an IterableDataset — identical API to load_dataset(..., streaming=True)')

deps.close()
q.close()
print('Done — lakehouse connections closed.')