# /// 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()