Measure the cost of encryption on columnar scans

# /// script
# requires-python = ">=3.10"
# dependencies = ["pyarrow", "duckdb", "numpy", "cryptography"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Measure the cost of encryption on columnar scans.

See how encryption removes the statistics that make queries fast. Three Parquet files of the
same rows: ts in plaintext (sorted), ts as a random permutation (a stand-in with the same effect
on statistics), and ts encrypted with AES-CTR and decrypted in our own code at query time.
Each query is timed as the best of several runs, and the row groups each must read are counted."""

import argparse
import time
import os
import tempfile
import json
import pyarrow as pa
import pyarrow.parquet as pq
import duckdb
import numpy as np
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes

def best(fn, repeats):
    fn()                                               # warm-up: file open, page cache
    times = []
    for _ in range(repeats):
        t = time.perf_counter(); out = fn(); times.append(time.perf_counter() - t)
    return min(times), out

def groups_read(path, col, a, b):
    """Row groups whose min/max statistics overlap [a, b]; None when the column has no usable order."""
    md = pq.ParquetFile(path).metadata
    i = md.schema.to_arrow_schema().get_field_index(col)
    st = [md.row_group(g).column(i).statistics for g in range(md.num_row_groups)]
    if any(s is None or not s.has_min_max or isinstance(s.min, bytes) for s in st):
        return md.num_row_groups, md.num_row_groups    # ciphertext: nothing can be skipped
    return sum(1 for s in st if not (s.max < a or s.min > b)), md.num_row_groups

def column_mb(path, col):
    md = pq.ParquetFile(path).metadata
    i = md.schema.to_arrow_schema().get_field_index(col)
    return sum(md.row_group(g).column(i).total_compressed_size for g in range(md.num_row_groups)) / 2**20

def main():
    parser = argparse.ArgumentParser(description=__doc__.split('\n')[0])
    parser.add_argument("--rows", type=int, default=20_000_000, help="Number of rows to generate.")
    parser.add_argument("--row-group", type=int, default=500_000, help="Rows per row group.")
    parser.add_argument("--repeats", type=int, default=5, help="Timed runs per query (the best is kept).")
    parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility.")
    parser.add_argument("--json", type=str, help="Also write the results to this JSON file.")
    args = parser.parse_args()

    rng = np.random.default_rng(args.seed)
    print(f"Generating {args.rows} rows...")
    ts = np.arange(args.rows, dtype=np.int64)
    value = rng.random(args.rows)
    ts_perm = rng.permutation(ts)                      # stand-in for ciphertext

    key, nonce = os.urandom(32), os.urandom(16)
    def ctr():
        return Cipher(algorithms.AES(key), modes.CTR(nonce))
    enc = ctr().encryptor().update(ts.tobytes())       # AES-CTR over the whole column: every 8-byte value looks random
    ts_cipher = pa.FixedSizeBinaryArray.from_buffers(pa.binary(8), args.rows, [None, pa.py_buffer(enc)])

    a, b = int(args.rows * 0.5), int(args.rows * 0.51)   # about 1% of the rows
    queries, sizes = [], []
    with tempfile.TemporaryDirectory() as d:
        paths = {}
        for label, col, arr in (("plaintext", "ts", ts), ("permuted (stand-in)", "ts", ts_perm), ("AES-CTR encrypted", "ts_enc", ts_cipher)):
            for comp in ("none", "snappy", "zstd"):
                p = os.path.join(d, f"{label[:4]}-{comp}.parquet")
                pq.write_table(pa.table({col: arr, "value": value}), p, compression=comp, row_group_size=args.row_group)
                sizes.append({"file": label, "compression": comp, "ts_mb": column_mb(p, col), "file_mb": os.path.getsize(p) / 2**20})
                if comp == "snappy":
                    paths[label] = (p, col)

        con = duckdb.connect()
        for label in ("plaintext", "permuted (stand-in)"):
            p, col = paths[label]
            q = f"SELECT sum(value) FROM '{p}' WHERE {col} BETWEEN {a} AND {b}"
            t, res = best(lambda: con.execute(q).fetchone()[0], args.repeats)
            g, total = groups_read(p, col, a, b)
            queries.append({"file": label, "seconds": t, "row_groups_read": g, "row_groups": total, "result": float(res)})

        p, col = paths["AES-CTR encrypted"]
        def decrypt_and_sum():                         # read everything, decrypt in our own code, then filter
            t = pq.read_table(p, columns=[col, "value"])
            raw = t.column(col).combine_chunks().buffers()[1].to_pybytes()
            plain = np.frombuffer(ctr().decryptor().update(raw), dtype=np.int64)
            v = t.column("value").to_numpy()
            return float(v[(plain >= a) & (plain <= b)].sum())
        t, res = best(decrypt_and_sum, args.repeats)
        g, total = groups_read(p, col, a, b)
        queries.append({"file": "AES-CTR encrypted", "seconds": t, "row_groups_read": g, "row_groups": total, "result": res})

    base = queries[0]["seconds"]
    for q in queries:
        q["slowdown"] = q["seconds"] / base
        print(f"{q['file']:<22} {q['seconds']*1000:8.1f} ms  {q['slowdown']:5.1f}x  row groups read {q['row_groups_read']}/{q['row_groups']}  sum={q['result']:.2f}")
    print()
    for s in sizes:
        print(f"{s['file']:<22} {s['compression']:<7} ts column {s['ts_mb']:7.1f} MB   file {s['file_mb']:7.1f} MB")

    size = {(s["file"], s["compression"]): s["ts_mb"] for s in sizes}
    results = {
        "summary": {
            "rows": args.rows,
            "plain_s": queries[0]["seconds"], "perm_s": queries[1]["seconds"], "enc_s": queries[2]["seconds"],
            "perm_slowdown": queries[1]["slowdown"], "enc_slowdown": queries[2]["slowdown"],
            "plain_groups": queries[0]["row_groups_read"], "row_groups": queries[0]["row_groups"],
            "ts_raw_mb": size[("plaintext", "none")],
            "plain_zstd_mb": size[("plaintext", "zstd")], "enc_zstd_mb": size[("AES-CTR encrypted", "zstd")],
            "enc_zstd_ratio": size[("AES-CTR encrypted", "zstd")] / size[("AES-CTR encrypted", "none")],
        },
        "tables": {"queries": queries, "sizes": sizes},
    }
    if args.json:
        with open(args.json, "w") as f:
            json.dump(results, f, indent=2)

if __name__ == "__main__":
    main()