probe8.py

#!/usr/bin/env python3
"""probe8.py - break the multi-attribute (Relation_Sample) block layout.

`probes4/big.hyper` is the corpus's only sample block, and it is a benign one:
three attributes, and the *last* of them has an empty string heap. So the rule
"the last attribute's stringData runs to the end of the block" is never
actually loaded, and neither is the one-attribute case where the last
attribute's region ends come from its own group.

Each database here is built to make a different one of those rules matter:

  tail      last column is Text with long values, so the last attribute's
            string heap is the thing that runs to the block end
  one       a single column, so begin and end come from the same group
  many      twelve columns of mixed widths, so a wrong stride desynchronises
  nulls     every column nullable with NULLs, for the nullable schemes
  wide      16-byte types (BigNumeric, Interval) either side of a Varchar

Two checks per file, neither of which the layout can pass by luck:

  moments   the footer's sum, sum of squares and sum of cubes, recomputed from
            the values we decoded. Any wrong code width, region extent or
            dictionary index moves at least one of the three.
  tuples    every decoded sample row must be an actual row of the table, read
            back through the Hyper API. Columns that each hold plausible
            values but are misaligned into rows fail this and pass everything
            else.

  /usr/bin/python3 probe8.py [--rows N]
"""
import os
import shutil
import sys
from collections import Counter

import hyperparse as hp
from checkrows import normalise

try:
    from tableauhyperapi import (HyperProcess, Connection, Telemetry, CreateMode,
                                 TableDefinition, SqlType, NULLABLE, NOT_NULLABLE,
                                 TableName, Inserter)
except ImportError:
    sys.exit("needs tableauhyperapi; try /usr/bin/python3")

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, "probes8")
ROWS = 200_000            # big.hyper's size, known to make the engine sample


def build(hp_proc, name, cols, row):
    """Create OUT/name.hyper with one table of `cols`, rows from `row(i)`."""
    path = os.path.join(OUT, f"{name}.hyper")
    tbl = TableName("Extract", "T")
    with Connection(hp_proc.endpoint, path, CreateMode.CREATE_AND_REPLACE) as c:
        c.catalog.create_schema("Extract")
        c.catalog.create_table(TableDefinition(tbl, cols))
        with Inserter(c, TableDefinition(tbl, cols)) as ins:
            ins.add_rows([row(i) for i in range(ROWS)])
            ins.execute()
    return path


def col(name, t, null=True):
    return TableDefinition.Column(name, t, NULLABLE if null else NOT_NULLABLE)


def databases(proc):
    import datetime
    import decimal

    long_text = [f"label-{i:05d}-" + "x" * (i % 40) for i in range(1000)]

    yield build(proc, "tail",
                [col("n", SqlType.big_int()), col("t", SqlType.text())],
                lambda i: [i, long_text[i % 1000]])

    yield build(proc, "one", [col("t", SqlType.text())],
                lambda i: [long_text[i % 1000]])

    yield build(proc, "many", [
        col("b", SqlType.bool()), col("s", SqlType.small_int()),
        col("i", SqlType.int()), col("g", SqlType.big_int()),
        col("d", SqlType.date()), col("ts", SqlType.timestamp()),
        col("f", SqlType.double()), col("n", SqlType.numeric(18, 4)),
        col("c", SqlType.char(8)), col("v", SqlType.varchar(64)),
        col("j", SqlType.json()), col("y", SqlType.bytes()),
    ], lambda i: [
        bool(i % 2), (i % 30000) - 15000, i - 100000, i * 7,
        datetime.date(1900 + i % 200, 1 + i % 12, 1 + i % 28),
        datetime.datetime(2000, 1, 1) + datetime.timedelta(seconds=i),
        i * 1.5, decimal.Decimal(i) / 10000,
        f"c{i % 1000:06d}", long_text[i % 1000],
        f'{{"k":{i % 997}}}', bytes([i % 251, (i * 7) % 253, i % 17]),
    ])

    # Every fourth row NULL in every column, so a nullable scheme is forced and
    # code 0 stays reserved.
    yield build(proc, "nulls", [
        col("g", SqlType.big_int()), col("v", SqlType.varchar(64)),
        col("f", SqlType.double()), col("d", SqlType.date()),
    ], lambda i: ([None, None, None, None] if i % 4 == 0 else
                  [i, long_text[i % 1000], i * 0.25,
                   datetime.date(1990 + i % 30, 1 + i % 12, 1 + i % 28)]))

    yield build(proc, "wide", [
        col("q", SqlType.numeric(38, 6)), col("v", SqlType.varchar(64)),
        col("iv", SqlType.interval()), col("g", SqlType.big_int()),
    ], lambda i: [decimal.Decimal(i) / 1000000, long_text[i % 1000],
                  None, i])


def check(path, proc, show):
    d = open(path, "rb").read()
    f = hp.parse(path)
    cat = f.live_catalog
    rels = hp.relations(cat)
    ents = hp.relation_samples(f)
    print(f"\n{os.path.basename(path)}  {len(ents)} sample object(s)")
    if not ents:
        print("  !! no Relation_Sample - nothing probed")
        return 0, 1
    ok = bad = 0
    with Connection(proc.endpoint, path) as conn:
        for e in ents:
            blob = hp.read_object(d, e)
            rel = rels[e.relation]
            attrs = rel.get("attributes", [])
            tbl = f'"{rel["_schema"]}"."{rel["name"]}"'
            s = hp.parse_relation_sample(blob)
            if s is None:
                print(f"  !! {e}: does not parse"); bad += 1; continue
            v = s.block.violations()
            if v:
                print(f"  !! {e}: {v}"); bad += 1; continue
            cols, note = hp.decode_sample(
                blob, [hp.ColType.parse(a.get("type")) for a in attrs])
            print(f"  {tbl}: {s.block.tuple_count:,} of {s.relation_rows:,} rows, "
                  f"{len(s.block.groups)} attrs, schemes="
                  f"{[g[3] for g in s.block.groups]}")
            heaps = [len(s.block.regions(k)[3]) for k in range(len(s.block.groups))]
            print(f"    string heaps: {heaps}   (last must be non-zero to "
                  f"exercise end.stringData=size)" if heaps[-1] == 0 else
                  f"    string heaps: {heaps}")
            if note and not cols:
                print(f"    DECODE-FAIL {note}"); bad += 1; continue

            m = hp.check_sample_moments(blob, [hp.ColType.parse(a.get("type")) for a in attrs])
            if m:
                print(f"    MOMENT-FAIL"); [print(f"      {x}") for x in m]; bad += 1
            else:
                print("    moments OK"); ok += 1

            sel = ", ".join(f'"{a["name"]}"' for a in attrs)
            truth = Counter(tuple(map(normalise, r)) for r in
                            conn.execute_list_query(f"SELECT {sel} FROM {tbl}"))
            mine = Counter(tuple(normalise(c[r]) for c in cols)
                           for r in range(s.block.tuple_count))
            extra = mine - truth
            if extra:
                print(f"    TUPLE-FAIL {sum(extra.values())} sampled rows are not "
                      f"rows of the table")
                for t in list(extra)[:3]:
                    print(f"      {t}")
                bad += 1
            else:
                print(f"    every sampled row is a real row "
                      f"({len(mine)} distinct)")
                ok += 1
            for r in range(min(show, s.block.tuple_count)):
                print("      " + " | ".join(f"{c[r]!s:>14}"[:14] for c in cols))
    return ok, bad


def main():
    show = 0
    if "--rows" in sys.argv:
        show = int(sys.argv[sys.argv.index("--rows") + 1])
    shutil.rmtree(OUT, ignore_errors=True)
    os.makedirs(OUT, exist_ok=True)
    tot_ok = tot_bad = 0
    with HyperProcess(telemetry=Telemetry.DO_NOT_SEND_USAGE_DATA_TO_TABLEAU,
                      parameters={"default_database_version": "4"}) as proc:
        paths = list(databases(proc))
        for p in paths:
            o, b = check(p, proc, show)
            tot_ok += o
            tot_bad += b
    print(f"\nTOTAL: {tot_ok} checks pass, {tot_bad} fail")
    return 1 if tot_bad else 0


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