|
| 1 | +####################################################################### |
| 2 | +# Copyright (c) 2019-present, Blosc Development Team <blosc@blosc.org> |
| 3 | +# All rights reserved. |
| 4 | +# |
| 5 | +# SPDX-License-Identifier: BSD-3-Clause |
| 6 | +####################################################################### |
| 7 | + |
| 8 | +"""Compare the three plain-string column representations head to head: |
| 9 | +``utf8()`` (offsets + bytes, StringDType reads), ``string(max_length=L)`` |
| 10 | +(fixed-width UTF-32), and ``vlstring()`` (msgpack cells). |
| 11 | +
|
| 12 | +Two workloads: |
| 13 | +- "taxi company": the real ``company`` column from the Chicago taxi dataset |
| 14 | + (medium-length, low-cardinality strings), when the parquet file is present. |
| 15 | +- "synthetic free text": high-cardinality random words of 0-60 chars with |
| 16 | + some multi-byte values — the workload utf8() is designed for. |
| 17 | +
|
| 18 | +For each representation: ingest, storage footprint, full read, equality |
| 19 | +filter, groupby-key aggregation, sort, and Arrow export. Operations a |
| 20 | +representation does not support are reported as such rather than skipped |
| 21 | +silently. |
| 22 | +""" |
| 23 | + |
| 24 | +import argparse |
| 25 | +import pathlib |
| 26 | +import time |
| 27 | +from dataclasses import make_dataclass |
| 28 | + |
| 29 | +import numpy as np |
| 30 | + |
| 31 | +import blosc2 |
| 32 | +from blosc2 import CTable |
| 33 | + |
| 34 | +N_TAXI = 10_000_000 |
| 35 | +N_SYNTH = 2_000_000 # fixed-width U~130 at 1e7 rows would need a ~5 GB ingest buffer |
| 36 | +REPS = 3 |
| 37 | +TAXI_PARQUET = pathlib.Path(__file__).parent.parent / "chicago-taxi" / "chicago-taxi-flat.parquet" |
| 38 | + |
| 39 | +parser = argparse.ArgumentParser(description=__doc__) |
| 40 | +parser.add_argument( |
| 41 | + "--ingest-only", |
| 42 | + action="store_true", |
| 43 | + help="Only measure ingest and storage footprint; skip full read, filter, " |
| 44 | + "groupby, sort, and to_arrow (the last of which alone can take over a " |
| 45 | + "minute per column kind on the full taxi workload). Meant for fast " |
| 46 | + "iteration when only ingest performance is under investigation.", |
| 47 | +) |
| 48 | +args = parser.parse_args() |
| 49 | + |
| 50 | +rng = np.random.default_rng(42) |
| 51 | + |
| 52 | + |
| 53 | +def bench(label, fn, reps=REPS): |
| 54 | + times = [] |
| 55 | + result = None |
| 56 | + for _ in range(reps): |
| 57 | + t0 = time.perf_counter() |
| 58 | + result = fn() |
| 59 | + times.append(time.perf_counter() - t0) |
| 60 | + print(f" {label:34s} {min(times) * 1e3:9.1f} ms") |
| 61 | + return result |
| 62 | + |
| 63 | + |
| 64 | +def load_taxi_company(n): |
| 65 | + import pyarrow.parquet as pq |
| 66 | + |
| 67 | + tbl = pq.read_table(TAXI_PARQUET, columns=["company"]) |
| 68 | + col = tbl.column("company").combine_chunks() |
| 69 | + values = col.to_pylist()[:n] |
| 70 | + return [v if v is not None else "" for v in values] |
| 71 | + |
| 72 | + |
| 73 | +def synth_free_text(n): |
| 74 | + words = np.array( |
| 75 | + ["taxi", "río", "航空", "boulevard", "x" * 40, "café", "", "downtown", "zürich", "o'hare"] |
| 76 | + ) |
| 77 | + # 2-6 words per row, high cardinality via a row counter suffix on ~half. |
| 78 | + parts = words[rng.integers(0, len(words), (n, 3))] |
| 79 | + joined = [" ".join(p) for p in parts] |
| 80 | + salt = rng.integers(0, 100_000, n) # ~100k distinct values: high cardinality, sane group count |
| 81 | + return [f"{s} #{salt[i]}" if i % 2 else s for i, s in enumerate(joined)] |
| 82 | + |
| 83 | + |
| 84 | +def run_workload(title, values, filter_value): |
| 85 | + n = len(values) |
| 86 | + max_len = max(len(v) for v in values) |
| 87 | + float_vals = rng.random(n) |
| 88 | + print(f"\n=== {title} ({n:.0e} rows, max length {max_len} chars) ===") |
| 89 | + |
| 90 | + specs = { |
| 91 | + "utf8": blosc2.utf8(), |
| 92 | + "string": blosc2.string(max_length=max_len), |
| 93 | + "vlstring": blosc2.vlstring(), |
| 94 | + } |
| 95 | + for kind, spec in specs.items(): |
| 96 | + row_cls = make_dataclass( |
| 97 | + "Row", [("s", str, blosc2.field(spec)), ("val", float, blosc2.field(blosc2.float64()))] |
| 98 | + ) |
| 99 | + print(f"[{kind}]") |
| 100 | + t0 = time.perf_counter() |
| 101 | + t = CTable(row_cls) |
| 102 | + t.extend({"s": values, "val": float_vals}, validate=False) |
| 103 | + t._flush_varlen_columns() |
| 104 | + print(f" {'ingest':34s} {(time.perf_counter() - t0) * 1e3:9.1f} ms") |
| 105 | + |
| 106 | + col = t._cols["s"] |
| 107 | + nbytes = getattr(col, "nbytes", None) |
| 108 | + cbytes = getattr(col, "cbytes", None) |
| 109 | + if nbytes is None: # NDArray-backed fixed-width column |
| 110 | + nbytes, cbytes = col.schunk.nbytes, col.schunk.cbytes |
| 111 | + print( |
| 112 | + f" {'storage nbytes -> cbytes':34s} {nbytes / 2**20:7.1f} MB -> {cbytes / 2**20:7.1f} MB (cratio {nbytes / cbytes:.1f}x)" |
| 113 | + ) |
| 114 | + |
| 115 | + if args.ingest_only: |
| 116 | + del t |
| 117 | + continue |
| 118 | + |
| 119 | + bench("full column read", lambda t=t: t["s"][:]) |
| 120 | + try: |
| 121 | + bench("filter: count(s == value)", lambda t=t: int((t.s == filter_value)[:].sum())) |
| 122 | + except (NotImplementedError, TypeError) as exc: |
| 123 | + print(f" {'filter: count(s == value)':34s} unsupported: {str(exc)[:60]}") |
| 124 | + try: |
| 125 | + bench("groupby key: sum(val)", lambda t=t: t.group_by("s").sum("val"), reps=1) |
| 126 | + except (NotImplementedError, TypeError) as exc: |
| 127 | + print(f" {'groupby key: sum(val)':34s} unsupported: {str(exc)[:60]}") |
| 128 | + try: |
| 129 | + bench("sort_by(s) (copy)", lambda t=t: t.sort_by("s"), reps=1) |
| 130 | + except (NotImplementedError, TypeError) as exc: |
| 131 | + print(f" {'sort_by(s) (copy)':34s} unsupported: {str(exc)[:60]}") |
| 132 | + try: |
| 133 | + bench("to_arrow()", lambda t=t: t.to_arrow(), reps=1) |
| 134 | + except Exception as exc: |
| 135 | + print(f" {'to_arrow()':34s} failed: {str(exc)[:60]}") |
| 136 | + del t |
| 137 | + |
| 138 | + |
| 139 | +if TAXI_PARQUET.exists(): |
| 140 | + print("loading taxi company column...", flush=True) |
| 141 | + taxi = load_taxi_company(N_TAXI) |
| 142 | + # the most frequent company value as the filter probe |
| 143 | + vals, counts = np.unique(np.array(taxi, dtype=np.dtypes.StringDType()), return_counts=True) |
| 144 | + run_workload("chicago-taxi company", taxi, str(vals[np.argmax(counts)])) |
| 145 | + del taxi |
| 146 | +else: |
| 147 | + print(f"({TAXI_PARQUET} not found; skipping the real-data workload)") |
| 148 | + |
| 149 | +print("building synthetic free text...", flush=True) |
| 150 | +synth = synth_free_text(N_SYNTH) |
| 151 | +run_workload("synthetic free text", synth, synth[123]) |
0 commit comments