probe5.py

#!/usr/bin/env python3
"""probe5.py - sweep DataBlockCompression codes across type x cardinality x nulls.

probe2.py sampled a few cardinalities and produced a scheme table whose gaps
were filled by assuming each family is a contiguous triple of code widths. That
assumption turned out to be wrong: scheme 39 is not the 32-bit dictionary of
the {36,37,38} family, it is an *uncompressed* float32 array. So sweep widely
and record only what is actually observed.

The cardinality ladder walks a column from constant (single-value) through the
bit-packed and byte-packed dictionary widths and out the far end, where every
value is distinct and a dictionary costs more than the raw data - which is what
forces the truncation and uncompressed encodings into the open.

One file per type. Row count stays under the 2**17 block cap so each column is
a single block.
"""
import datetime
import decimal
import os
import shutil
import sys

from tableauhyperapi import (HyperProcess, Connection, Telemetry, CreateMode,
                             TableDefinition, SqlType, Inserter, NULLABLE,
                             TableName, Interval)

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, "probes5")
N = 70_000
CARDS = [1, 2, 3, 5, 20, 300, 5000, N]

EPOCH = datetime.datetime(2000, 1, 1)
TYPES = {
    "bool":  (SqlType.bool(), lambda k: bool(k % 2), 2),
    "sml":   (SqlType.small_int(), lambda k: -30000 + (k % 60000), 60000),
    "int":   (SqlType.int(), lambda k: 100000 + k, None),
    "big":   (SqlType.big_int(), lambda k: 10 ** 12 + k, None),
    "oid":   (SqlType.oid(), lambda k: 70000 + k, None),
    "date":  (SqlType.date(), lambda k: datetime.date(2000, 1, 1) +
              datetime.timedelta(days=k), None),
    "time":  (SqlType.time(), lambda k: (EPOCH + datetime.timedelta(
        microseconds=k * 1_000_003)).time(), None),
    "ts":    (SqlType.timestamp(), lambda k: EPOCH + datetime.timedelta(seconds=k), None),
    "tstz":  (SqlType.timestamp_tz(), lambda k: EPOCH.replace(
        tzinfo=datetime.timezone.utc) + datetime.timedelta(seconds=k), None),
    "dbl":   (SqlType.double(), lambda k: 1000.5 + k, None),
    "dblx":  (SqlType.double(), lambda k: 1000.1 + k * 1e-7, None),  # needs f64
    "flt":   (SqlType.float(), lambda k: 1000.5 + k, None),
    "num9":  (SqlType.numeric(9, 2), lambda k: decimal.Decimal(f"{1000 + k}.25"), None),
    "num18": (SqlType.numeric(18, 4), lambda k: decimal.Decimal(f"{10**12 + k}.0625"), None),
    "num38": (SqlType.numeric(38, 6), lambda k: decimal.Decimal(f"{10**24 + k}.015625"), None),
    "txt":   (SqlType.text(), lambda k: f"value{k:07d}", None),
    "txts":  (SqlType.text(), lambda k: f"{k % 1000:03d}", None),   # <= 3 B, inline
    "char":  (SqlType.char(10), lambda k: f"c{k:09d}", None),
    "bytea": (SqlType.bytes(), lambda k: k.to_bytes(6, "little"), None),
    "json":  (SqlType.json(), lambda k: f'{{"k":{k}}}', None),
    "intvl": (SqlType.interval(), lambda k: Interval(0, 0, k * 1000), None),
}


def build(hp, label, sqltype, make, cap):
    cards = sorted({min(c, cap) if cap else c for c in CARDS})
    cols = []
    for card in cards:
        for nulls in (False, True):
            vals = [make(i % card) for i in range(N)]
            if nulls:
                for i in range(0, N, 7):
                    vals[i] = None
            cols.append((f"c{card}_{'n' if nulls else 'x'}", vals))
    path = os.path.join(OUT, f"{label}.hyper")
    tdef = TableDefinition(TableName("Extract", "E"),
                           [TableDefinition.Column(n, sqltype, NULLABLE)
                            for n, _ in cols])
    with Connection(hp.endpoint, path, CreateMode.CREATE_AND_REPLACE) as c:
        c.catalog.create_schema("Extract")
        c.catalog.create_table(tdef)
        with Inserter(c, tdef) as ins:
            ins.add_rows([[v[i] for _, v in cols] for i in range(N)])
            ins.execute()
    return path


def main():
    shutil.rmtree(OUT, ignore_errors=True)
    os.makedirs(OUT)
    ok, bad = 0, []
    with HyperProcess(telemetry=Telemetry.DO_NOT_SEND_USAGE_DATA_TO_TABLEAU,
                      parameters={"default_database_version": "4"}) as hp:
        for label, (sqltype, make, cap) in TYPES.items():
            try:
                build(hp, label, sqltype, make, cap)
                ok += 1
            except Exception as e:
                bad.append(f"{label}: {str(e).splitlines()[0][:80]}")
    for b in bad:
        print(f"  FAIL {b}")
    print(f"{ok} built, {len(bad)} failed -> {OUT}")
    return 0


if __name__ == "__main__":
    sys.exit(main())