diff --git a/CMakeLists.txt b/CMakeLists.txt index 1895b275a..c6cc48ac3 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -15,8 +15,9 @@ endif() project(python-blosc2) # blosc2_ext.pyx calls blosc2_schunk_lock()/unlock(), added in c-blosc2 3.2.x -set(BLOSC2_MIN_VERSION 3.2.1) -set(BLOSC2_BUNDLED_VERSION v3.2.3) +set(BLOSC2_MIN_VERSION 3.3.0) +set(BLOSC2_BUNDLED_VERSION v3.3.0) +# set(BLOSC2_BUNDLED_VERSION bc074b228968d6121b3c8c1a38c0afc0bbf923f6) if(WIN32 AND NOT CMAKE_C_COMPILER_ID STREQUAL "Clang") message(FATAL_ERROR "Windows builds require clang-cl. Set CC/CXX to clang-cl or configure CMake with -T ClangCL.") @@ -109,7 +110,7 @@ endif() FetchContent_Declare(miniexpr GIT_REPOSITORY https://github.com/Blosc/miniexpr.git - GIT_TAG 58d2d0b4a3aee3d1ac84b213712cf982744196c8 + GIT_TAG bd2c602a652c50b306def625c5fe5491cbd13f76 # SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}/../miniexpr ) FetchContent_MakeAvailable(miniexpr) diff --git a/RELEASE_NOTES.md b/RELEASE_NOTES.md index dc98b3e53..cfd0dc8a7 100644 --- a/RELEASE_NOTES.md +++ b/RELEASE_NOTES.md @@ -6,6 +6,54 @@ XXX version-specific blurb XXX ### New features +- **String-valued expressions and DSL kernels** over fixed-width ` 2)")` used to raise `NotImplementedError`; + only the operator form `t[t.name == "x"]` was available. A variable-length + column cannot be an expression operand (its offsets and data have + independent chunk grids, so the prefilter contract does not apply), so + these are evaluated span by span, each span materialized to a fixed-width + array whose width is rounded up to a power of two and handed to miniexpr. + Nulls are materialized to `""` before any kernel sees them and re-masked + afterwards, so a null never satisfies a predicate — the same answer the + operator form gives. +- **Scalar comparisons on `utf8()` columns are 5-6x faster in expression + form.** `t.where("name == 'x'")` (and `!=`, `<`, `<=`, `>`, `>=`, either + operand order) is now answered by the same raw-byte scan the operator form + `t[t.name == "x"]` uses, instead of decoding the column to fixed-width + first: 156 -> 28 ms over 1M short values, 268 -> 56 ms over 1M ~31-byte + values. Mixed expressions get whatever they can -- in + `startswith(name, 'x') | (name == 'zz')` the comparison takes the fast + path and `startswith` still decodes. +- **New `blosc2.utf8_array(seq, spec=None)`** builds a `UTF8Array` from an + iterable of strings; `UTF8Array` is exported too. Previously the only + construction path was `UTF8Array(spec)` + `.extend()` + `.flush()`, which + was not exported at all. +- **`df.apply(f, axis=1, engine=blosc2.jit)` now runs `row["colname"]` + kernels that contain an `if`.** Neither dispatch route could before: + tracing evaluated the branch over a whole column (`truth value ... is + ambiguous`) and the DSL parser rejected the subscript. Such references are + now rewritten into named parameters, so the function is compiled and every + branch runs. This is not string-specific — numeric row kernels with a + branch were equally blocked. String columns reach this route too, which + makes the pandas-3 "format room info" kernel run unmodified. Nulls in a + string column are rejected rather than substituted, since a row-wise kernel + over a null raises in pandas as well. - New `blosc2.random` module: seedable, NumPy-quality random `NDArray` constructors. Each chunk gets its own independent `SeedSequence`-spawned stream and is generated concurrently in a thread pool, giving full `PCG64` @@ -33,6 +81,155 @@ XXX version-specific blurb XXX matching numpy. - Not implemented: `bytes` (returns raw `bytes`, not an `NDArray`). +- **`create_index()` now works on `utf8()` columns**, the last string flavour + without one. Both `utf8` and `dictionary` are indexed by the *alphabetical + rank* of each value: sorting by rank is sorting by the decoded string, so an + `int32` rank column drives the same machinery a numeric column uses. At 1M + rows / cardinality 20k: `sort_by` 424 ms -> 7.2 ms, `sorted_slice` 458 ms -> + 43 ms, and the index is the cheapest of the three flavours to build (277 ms + against 867 ms for ` 5.5 ms, `<` 34.6 ms -> 5.5 ms, and the dictionary operator + form `t[t.c == v]` 329.6 ms -> 8.4 ms. `startswith`/substring searches are + not accelerated (no index covers them), and ranks are frozen at build time, + so a value inserted ahead of existing ones sends the index stale until it is + rebuilt. + +- **`CTable.add_column()` accepts `values=`**, a sequence with one entry per + live row, as an alternative to backfilling from a declared default. This is + the supported way to land a result computed outside the table back into it, + which matters most for `utf8()` columns: string-returning expressions are + evaluated on fixed-width arrays, and the result previously had to be written + through the private `t._cols[name].set_all(...)`. A declared default is still + honoured for rows appended later, so the two can be combined. `values=` is + checked against the constraints declared on the spec, like the constructor + and `extend()` are: without that, coercion to a fixed-width dtype would + truncate an over-long string to `max_length` instead of complaining. +- **`blosc2.from_utf8()` / `blosc2.to_utf8()` and `UTF8Array.astype()`** make + the conversion between variable-length and fixed-width text an explicit, + documented pair. utf8 columns store and filter text compactly, but + string-*returning* expressions need miniexpr's compile-time output width, so + they run on fixed-width arrays; the rule is now written down (see "Computing + strings on a utf8 column" in the CTable reference) rather than left for + callers to discover. `from_utf8()` sizes the result to the longest value in + **codepoints**, counted from the raw bytes without decoding a row, so nothing + truncates and non-ASCII text does not over-allocate the 3-4x a byte-length + bound would. +- **The array constructors dispatch on NumPy's `StringDType`.** + `blosc2.asarray(np.array([...], dtype=StringDType()))` used to raise + `TypeError: data type 'StringDType()' not understood`, and + `blosc2.zeros(n, dtype=StringDType())` a `malformed node` `ValueError`; both + now return a `UTF8Array`, as do `empty`, `ones` and `full`, with the same + fill values NumPy uses (`''`, `''`, `'1'`, `str(fill_value)`). The dispatch + is on the *target* dtype, so `asarray(utf8_source, dtype=" 77.9 ms before, 6.6 ms after). +- **The utf8 compute refusals now route instead of merely refusing.** Every + path that cannot take a utf8 column — `add_computed_column`, + `add_generated_column`, `assign`, `apply`, `lazyudf`, with a string + expression or a DSL kernel — raises `NotImplementedError` naming the column + and printing the three-line conversion, echoing the user's own expression + where there is one. Two of those paths previously failed with a raw NumPy + `DTypePromotionError` and a `ValueError: malformed node or string ... + StringDType()`, neither of which named the column or the fix. + +### Bug fixes + +- **`@blosc2.jit` raised when a storage kwarg and an execution-tuning kwarg + were combined** and the decorated function returned a NumPy array — + `@blosc2.jit(jit=False, cparams=...)` ended in + `blosc2.asarray(retval, jit=False, ...)`, which rejects the tuning kwargs. + Only storage kwargs reach `asarray()` now; the function has already run, so + there is nothing left to tune. +- **A DSL kernel over a utf8 column registered as a computed column, then + broke the table.** `add_computed_column(name, kernel, inputs=["utf8_col"])` + was accepted, after which every read of that column *and* `str(table)` + raised `ValueError: malformed node or string`. The kernel is now refused at + registration, where the table is still untouched. +- **`min()`/`max()` read from a column index returned the wrong value.** Two + independent causes, both affecting every indexable dtype. The block summaries + cover the column's *physical* extent, so the capacity padding (zeros, empty + strings) was reduced along with the data and `min()` reported it — wrong on + any table whose row count is not exactly its slot capacity. And `delete()` + bumps a visibility epoch that nothing recorded, so deleted rows kept + contributing their values to the block they sat in. Whole blocks below the + live row count are still read from the sidecar; the block straddling the + boundary is now rescanned, and a deletion since the index was built makes the + shortcut stand down. +- **`create_index` on `utf8()` and `dictionary()` columns accepted any index + kind** and built one over the alphabetical ranks that no query would ever + consult — only `IndexKind.FULL` reaches a rank index. `kind` now defaults to + `FULL` for these two column kinds (`BUCKET` elsewhere, unchanged) and raises + `ValueError` when another kind is requested explicitly. Previously + `create_index("category")` on a dictionary column built an unused BUCKET + index by default. +- **Comparison operators on `UTF8Array`, dictionary and varlen scalar columns** + returned a plain `False`: none defined them, so `column == "value"` fell + through to object identity. Silently wrong rather than an error. All now + return boolean masks; `UTF8Array` and `DictionaryColumn` answer a scalar + without decoding any row. +- **`dictcol != value` raised `IndexError`** on any table with capacity + padding: the negation was applied after the live-row intersection, turning + every dead slot `True`. +- **Expressions over a bare `UTF8Array`** (`blosc2.lazyexpr("'x=' + a", {"a": + arr})`) produced correct values down the wrong path — widened to fixed-width + `typesize`, which the + chunk header contradicts once the typesize is capped. Across 153 slice + shapes at typesize 256, 150 raised `"Error while getting the slice"` and + the 3 single-element ones returned the wrong bytes with no error at all. + Reachable from ordinary data -- an ` 255` (31 works, 32 does not); `summary`, + `bucket`, `partial` and `full` indexes were affected, `opsi` was not. + A short span read now raises instead of leaving the destination partly + uninitialised. +- **Expressions over operands wider than 255 bytes returned wrong results.** + c-blosc2 caps a typesize above 255 to 1 in the chunk header so its split + machinery keeps working, and the miniexpr prefilter was asking + `blosc2_getitem_ctx()` for operand blocks in element units, which the chunk + then read as a byte range: every block past the first was uninitialised + memory. `arr == "hello"` over 1200 rows of ` +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +####################################################################### + +"""String workloads over the Chicago Taxi dataset: Blosc2 vs pandas/polars/DuckDB. + +The companion of `compare-query-methods.py`, for the *string* columns rather +than the numeric ones. Three tasks over `company` (` bool + transform 'co=' + company + '|pay=' + lower(payment_type) -> str + kernel the same, but branching on whether the company is a cab + company -- i.e. row-wise control flow, not one expression + +All three are timed; only `kernel` is plotted. It is the shape of the pandas-3 +blog kernel (datapythonista.me/blog/whats-new-in-pandas-3): every other engine +has to express it as a mask plus two fully-evaluated branches, whereas blosc2 +compiles it to a single masked pass with `@blosc2.dsl_kernel`. + +`blosc2 (raw)` is the same blosc2 path with `clevel=0` on operands *and* +result. Same container, same kernel, compression the only variable -- so the +gap between the two blosc2 bars is the price of compression, and the gap +between their footprints is what that price buys. + +Usage: + python string-ops.py # whole table, best of 3 + python string-ops.py --nrows 1000000 --apply + python string-ops.py --engines blosc2,numpy --nrows 1000000 +""" + +import argparse +import gc +import hashlib +import time + +import numpy as np + +import blosc2 + +PARQUET = "chicago-taxi-flat.parquet" +COLS = ["company", "payment.type"] + +# One chunk/block geometry for every blosc2 operand. Expressions combining two +# NDArrays only take the (miniexpr) fast path when the operands share a chunk +# grid, and asarray() picks the grid from the itemsize -- which differs between +# = nrows: + break + table = pa.Table.from_batches(batches).slice(0, nrows) + out = [] + for name in COLS: + col = table[name].combine_chunks().dictionary_decode() + out.append(col.to_numpy(zero_copy_only=False).astype(str)) + return out + + +# -------------------------------------------------------------------------- +# blosc2 +# -------------------------------------------------------------------------- + + +@blosc2.dsl_kernel +def taxi_label(company, ptype): + pay = ptype.lower() + c = company.lower() + if " cab" in c: + return "cab|" + c.removesuffix(" cab") + "|" + pay + return "other|" + c + "|" + pay + + +def blosc2_setup(co, pt, cparams=DEFAULT): + cp = {"cparams": cparams} + kw = {"chunks": CHUNKS, "blocks": BLOCKS, **cp} + return blosc2.asarray(co, **kw), blosc2.asarray(pt, **kw), cp + + +def blosc2_filter(a, b, cp): + # strict_miniexpr: a silent fallback to the NumPy path would still give the + # right answer, so without this the number below would not mean what it says. + e = blosc2.startswith(a, "Taxi") & (b != "Cash") + return e.compute(strict_miniexpr=True, **cp) + + +def blosc2_transform(a, b, cp): + e = "co=" + a + "|pay=" + blosc2.lower(b) + return e.compute(strict_miniexpr=True, **cp) + + +def blosc2_kernel(a, b, cp): + return blosc2.lazyudf(taxi_label, (a, b)).compute(**cp) + + +def blosc2_raw_setup(co, pt): + return blosc2_setup(co, pt, cparams=RAW) + + +blosc2_raw_filter = blosc2_filter +blosc2_raw_transform = blosc2_transform +blosc2_raw_kernel = blosc2_kernel + + +# -------------------------------------------------------------------------- +# NumPy +# -------------------------------------------------------------------------- + + +def numpy_setup(co, pt): + return co, pt + + +def numpy_filter(co, pt): + return np.strings.startswith(co, "Taxi") & (pt != "Cash") + + +def numpy_transform(co, pt): + return np.strings.add(np.strings.add("co=" + co, "|pay="), np.strings.lower(pt)) + + +def numpy_kernel(co, pt): + c = np.strings.lower(co) + tail = np.strings.add("|", np.strings.lower(pt)) + # np.strings has no removesuffix(); endswith + slice is the same thing. + trimmed = np.where( + np.strings.endswith(c, " cab"), np.strings.slice(c, 0, np.strings.str_len(c) - 4), c + ) + cab = np.strings.add("cab|" + trimmed, tail) + other = np.strings.add("other|" + c, tail) + return np.where(np.strings.find(c, " cab") >= 0, cab, other) + + +# -------------------------------------------------------------------------- +# pandas +# -------------------------------------------------------------------------- + + +def pandas_setup(co, pt): + import pandas as pd + + return pd.Series(co, dtype="str"), pd.Series(pt, dtype="str") + + +def pandas_filter(co, pt): + return co.str.startswith("Taxi") & (pt != "Cash") + + +def pandas_transform(co, pt): + return "co=" + co + "|pay=" + pt.str.lower() + + +def pandas_kernel(co, pt): + c = co.str.lower() + tail = "|" + pt.str.lower() + cab = "cab|" + c.str.removesuffix(" cab") + tail + return ("other|" + c + tail).where(~c.str.contains(" cab", regex=False), cab) + + +def pandas_kernel_apply(co, pt): + """The row-wise spelling of `kernel`, which is how it would first be written. + + Off the scale next to everything else, and reported separately for that + reason -- it is the baseline `@blosc2.dsl_kernel` exists to replace. + """ + import pandas as pd + + df = pd.DataFrame({"company": co, "ptype": pt}) + + def f(row): + pay = row["ptype"].lower() + c = row["company"].lower() + if " cab" in c: + return "cab|" + c.removesuffix(" cab") + "|" + pay + return "other|" + c + "|" + pay + + return df.apply(f, axis=1) + + +# -------------------------------------------------------------------------- +# polars +# -------------------------------------------------------------------------- + + +def polars_setup(co, pt): + import polars as pl + + return pl.DataFrame({"company": co, "ptype": pt}), None + + +def _pl(df, e): + return df.select(e.alias("r")).to_series() + + +def polars_filter(df, _): + import polars as pl + + return _pl(df, pl.col("company").str.starts_with("Taxi") & (pl.col("ptype") != "Cash")) + + +def polars_transform(df, _): + import polars as pl + + return _pl(df, pl.lit("co=") + pl.col("company") + "|pay=" + pl.col("ptype").str.to_lowercase()) + + +def polars_kernel(df, _): + import polars as pl + + c = pl.col("company").str.to_lowercase() + tail = pl.lit("|") + pl.col("ptype").str.to_lowercase() + return _pl( + df, + pl.when(c.str.contains(" cab", literal=True)) + .then(pl.lit("cab|") + c.str.strip_suffix(" cab") + tail) + .otherwise(pl.lit("other|") + c + tail), + ) + + +# -------------------------------------------------------------------------- +# DuckDB +# -------------------------------------------------------------------------- + +# No removesuffix() in SQL; ends_with + a slice is the literal equivalent and +# stays away from the regex engine, which would measure something else. +_DUCK_NOSUFFIX = "CASE WHEN ends_with(c, ' cab') THEN c[1:length(c) - 4] ELSE c END" + + +def _duck(con, q): + # .arrow() yields a RecordBatchReader from duckdb 1.5 on, a Table before it. + res = con.sql(q).arrow() + if hasattr(res, "read_all"): + res = res.read_all() + return res["r"] + + +def duckdb_setup(co, pt): + import duckdb + import pyarrow as pa + + con = duckdb.connect() + con.register("t", pa.table({"company": co, "ptype": pt})) + return con, None + + +def duckdb_filter(con, _): + return _duck(con, "SELECT starts_with(company, 'Taxi') AND ptype <> 'Cash' AS r FROM t") + + +def duckdb_transform(con, _): + return _duck(con, "SELECT 'co=' || company || '|pay=' || lower(ptype) AS r FROM t") + + +def duckdb_kernel(con, _): + return _duck( + con, + f""" + SELECT CASE WHEN contains(c, ' cab') + THEN 'cab|' || ({_DUCK_NOSUFFIX}) || tail + ELSE 'other|' || c || tail END AS r + FROM (SELECT lower(company) AS c, '|' || lower(ptype) AS tail FROM t) + """, + ) + + +# -------------------------------------------------------------------------- +# driver +# -------------------------------------------------------------------------- + +WINDOW = 1 << 19 # rows per verification window; bounds peak memory of the check + + +def _window(x, lo, hi): + """`x[lo:hi]` as a NumPy array, for any of the engines' native containers.""" + if type(x).__module__.startswith("pyarrow"): # NDArray has a .slice() too + return x.slice(lo, hi - lo).to_numpy(zero_copy_only=False) + part = x[lo:hi] + return np.asarray(part.to_numpy() if hasattr(part, "to_numpy") else part) + + +def digest(x, n): + """Memory-bounded fingerprint of a result, for cross-engine agreement. + + A 24 M-row ` 2 else ("ms", 1000) + panel(axes[0], [v * scale for v in t], "kernel: time", f"{unit}, lower is better", "{:.2f} " + unit) + panel(axes[1], s, "kernel: result footprint", "MB held in memory", "{:,.0f} MB") + + fig.suptitle( + f"Chicago Taxi row-wise string kernel, {nrows:,} rows (Nx = vs blosc2)\n" + "'blosc2 (raw)' is the identical path at clevel=0: compression is the only variable", + fontsize=11, + ) + fig.savefig(path, dpi=130) + print(f"wrote {path}") + + +if __name__ == "__main__": + main() diff --git a/bench/ndarray/jit-dsl-mandelbrot.py b/bench/ndarray/jit-dsl-mandelbrot.py index 6e9dfd5a4..f00317b9a 100644 --- a/bench/ndarray/jit-dsl-mandelbrot.py +++ b/bench/ndarray/jit-dsl-mandelbrot.py @@ -12,9 +12,11 @@ # doc/guides/optimization_tips.md ("Let @blosc2.jit compile control flow # instead of tracing it"). # -# Return paths are equalized (both calls end in a plain NumPy array): any -# non-None jit() kwarg flips the return from `retval[()]` to `.compute()`, -# which would otherwise skew the comparison. +# Return paths are equalized (both calls end in a plain NumPy array): a storage +# kwarg (cparams/chunks/urlpath/...) flips the return from `retval[()]` to +# `.compute()`, i.e. to an NDArray, which would otherwise skew the comparison. +# Execution-tuning kwargs (jit/jit_backend/fp_accuracy) do not, so `@blosc2.jit` +# is used bare here. from __future__ import annotations diff --git a/doc/guides/optim_tips/tip_11_dsl_random.png b/doc/guides/optim_tips/tip_11_dsl_random.png index 5fa677c26..43e864fd4 100644 Binary files a/doc/guides/optim_tips/tip_11_dsl_random.png and b/doc/guides/optim_tips/tip_11_dsl_random.png differ diff --git a/doc/guides/pandas_engine.md b/doc/guides/pandas_engine.md index 25d90c556..be671fc2c 100644 --- a/doc/guides/pandas_engine.md +++ b/doc/guides/pandas_engine.md @@ -152,6 +152,35 @@ single whole-column call rather than a per-row Python loop. Add a `for` or `while` to that idiom and it raises a `TypeError` pointing back here — that shape can only be compiled, not traced. +An `if`, on the other hand, is fine. `row["colname"]` references are rewritten +into named parameters, so the function is compiled as a DSL kernel and every +branch runs: + +```python +def format_room_info(row): + result = "property_type=" + row["property_type"] + desc = row["name"].lower() + if " with " not in desc: + return result + ", room_type=" + desc.removesuffix(" room") + before, after = desc.split(" with ", 1) + r2 = result + ", room_type=" + before.removesuffix(" room") + return r2 + ", amenity=" + after + + +df.apply(format_room_info, axis=1, engine=blosc2.jit) +``` + +String columns work here, and so do `.lower()`, `.removesuffix()` and +`str.split(sep, 1)` with tuple unpacking — they are lowered to the DSL's +function-call grammar for you. Two things to know: + +- A string local's width is fixed by its first assignment, so reassigning it + to a **longer** value is a compile error. Use a fresh name per step, as `r2` + does above. +- Nulls in a string column are **rejected**, because a row-wise kernel over a + null raises in pandas too (`"p=" + row["x"]` is a `TypeError`). Fill them + first if you want a value. + ## Gotchas **Your function normally runs only once.** The engine calls it a single time @@ -187,7 +216,9 @@ if it matters. ## Limitations -- Numeric dtypes only; anything else raises `ValueError`. +- Numeric, boolean and fixed-width string (`str`, `bytes`) columns; anything + else raises `ValueError`. Variable-width `utf8()` columns are not supported + yet. - `na_action="ignore"` is not supported for `map` (`NotImplementedError`): there is no per-element step at which to skip a value. - Only `DataFrame.apply` and `Series.map` reach the engine. pandas 3's diff --git a/doc/reference/classes.rst b/doc/reference/classes.rst index ff211f5ea..46c935617 100644 --- a/doc/reference/classes.rst +++ b/doc/reference/classes.rst @@ -18,6 +18,7 @@ variable-length, and object data. BatchArray ListArray ObjectArray + UTF8Array Proxies and External Data Sources diff --git a/doc/reference/ctable.rst b/doc/reference/ctable.rst index d37a0af2c..acc764086 100644 --- a/doc/reference/ctable.rst +++ b/doc/reference/ctable.rst @@ -458,6 +458,24 @@ For array-oriented grouped reductions without a :class:`CTable`, see Mutations --------- +:meth:`CTable.add_column` adds a physical column, filled either from a default +declared in the spec or from a ``values=`` sequence with one entry per live row:: + + t.add_column("weight", blosc2.field(blosc2.float64(), default=0.0)) + t.add_column("label", blosc2.utf8(), values=[f"row-{i}" for i in range(len(t))]) + +``values=`` is how a result computed outside the table lands back in it, which +matters most for :func:`utf8` columns: string-returning expressions are +evaluated on fixed-width arrays, so the result is written back explicitly rather +than by :meth:`CTable.add_computed_column`:: + + arr = t["name"][:].astype("``, ``>=``), but not the string-expression - form ``t.where("name == 'x'")`` yet. :meth:`CTable.create_index` on utf8 - columns is not supported yet either; both raise ``NotImplementedError`` - with a clear message. +.. [#rankindex] :func:`utf8` and :func:`dictionary` are indexed by the + *alphabetical rank* of each row's value: sorting by rank is sorting by the + decoded string, so an ``int32`` rank column drives the same index machinery + a numeric column uses. This accelerates :meth:`CTable.sort_by`, + :meth:`CTable.sorted_slice` and scalar comparisons (``==``/``!=`` on both, + plus ordering comparisons on utf8). It does **not** accelerate + ``startswith``/substring searches, which no index covers, and the ranks are + frozen at build time: a value inserted ahead of existing ones invalidates + all of them, so the index falls back to a full sort until rebuilt. + Only ``kind=IndexKind.FULL`` consults a rank index, so that is the default + for these two column kinds (elsewhere the default is ``BUCKET``) and any + other kind raises ``ValueError`` rather than building an unused index. + +.. [#utf8expr] utf8 columns support both the operator form + ``t[t.name == "x"]`` and the string-expression form + ``t.where("name == 'x'")``. Because a variable-length column cannot be an + expression operand directly, string expressions are evaluated span by span, + with each span materialized to a fixed-width array first; results are + identical to the operator form, including that a null never satisfies any + predicate. Nested (dotted) utf8 leaves are addressed by their dotted path + just like any other leaf. + +.. [#utf8compute] Expressions that *return* strings — concatenation, + ``upper``, ``replace``, DSL kernels — run on fixed-width arrays, so a utf8 + column is converted first and the result written back. See + :ref:`ComputingUtf8Strings` below. Note that a plain ``str`` annotation without an explicit :func:`field` spec -still maps to fixed-width ``string(max_length=32)`` for backward -compatibility; opt in to variable-length storage with -``blosc2.field(blosc2.utf8())``. +maps to fixed-width ``string(max_length=32)`` — the same width the decision +path above recommends, so the default is the fast path rather than a +compatibility accident. Opt in to variable-length storage with +``blosc2.field(blosc2.utf8())`` when the length is unbounded. + +.. _Utf8AndStringDType: + +utf8 and NumPy's ``StringDType`` +-------------------------------- + +:func:`utf8` is blosc2's answer to the same problem NumPy 2.0 solved with +``StringDType``, and the two interoperate on dtype: reads return +``StringDType`` arrays, and the array constructors dispatch on it:: + + blosc2.asarray(np.array(["a", "bb"], dtype=StringDType())) # -> UTF8Array + blosc2.zeros(3, dtype=StringDType()) # -> UTF8Array + blosc2.full(3, "x", dtype=StringDType()) # -> UTF8Array + +The fill values match NumPy's own (``''`` for ``zeros``/``empty``, ``'1'`` for +``ones``, ``str(fill_value)`` for ``full``), and the result satisfies the +:class:`blosc2.Array` protocol, so it can be used wherever a blosc2 array can. + +What blosc2 does **not** do is store ``StringDType`` in an +:class:`~blosc2.NDArray`, and it cannot: that dtype keeps each row's payload +outside the array buffer — a 100-character string still reports +``nbytes == 16`` — and supports no buffer protocol, so compressing the buffer +would persist pointers rather than text. A :class:`UTF8Array` holds the same +text as int64 offsets plus a UTF-8 blob, which is the layout Arrow uses for +``large_string`` and what makes :meth:`CTable.to_arrow` zero-copy. + +The dispatch is on the *target* dtype, so asking for a fixed width still gets +you a plain NDArray:: + + blosc2.asarray(utf8_source, dtype=" NDArray, fixed width + +Note the schema layer keeps its own vocabulary: ``blosc2.field()`` takes a +spec (:func:`utf8`, :func:`string`, :func:`int32`, ...) and not a raw NumPy +dtype, for any column type, because a spec also carries nullability, the null +sentinel, constraints and storage configuration. + +.. _ComputingUtf8Strings: + +Computing strings on a utf8 column +---------------------------------- + +The rule is: **utf8 stores and filters; fixed-width computes.** + +Everything on the *query* side works directly on a utf8 column — comparisons, +``where()`` including ``startswith``/``contains``, ``sum(where=)``, +``group_by``, ``sort_by`` and ``create_index``. Expressions that **return +strings** are different: miniexpr's string kernels need a compile-time output +width, which a variable-length column does not have. So those are converted, +computed, and written back:: + + fixed = blosc2.from_utf8(t["name"]) # -> ` and ranges, + because rank order *is* lexicographic order. +- ✗ `startswith` — nothing indexes prefixes on any flavour. +- ⚠ staleness: ranks shift when new values arrive. Dictionary handles this with a stored + `dict_hash` (`_dict_rank_hash`) and falls back to lexsort when stale; utf8 would need the same, + and a utf8 column is much more likely to gain new values than a category column is — which means + the fallback would fire often, and the fallback is a full lexsort. + +**Ordering shipped in `b1bbc54e`** (see ⁷); equality and ranges did not. + +Realistic remaining target: `=` | 34.16 ms | 8.07 ms | + +It hangs off `Column._utf8_scalar_mask`, which every scalar predicate already funnels through, and +returns `None` to fall back to the scan whenever the index cannot answer. Note the end-to-end +`where()` gain is smaller than these numbers — materializing the result rows out of the +offsets/blob dominates once the mask is cheap. + +**Dictionary followed in `1c22fdf5`**, minus the persistence — a dictionary already holds its +vocabulary in memory. The operator form `t[t.c == v]` goes **329.6 ms → 8.4 ms**. + +Two things surfaced while wiring it, both worth more than the feature: + +- **The staleness check cost more than the scan it saved.** `_dict_rank_index_stale` SHA1s the whole + dictionary — 24 ms for 20 k entries — on *every* query, including the ordering path that already + used the index. It now settles from the value epoch first, which also drops dictionary + `sorted_slice` from 51.3 ms to 37.7 ms. Wiring the index made queries *slower* until this was found. +- **`col != value` raised `IndexError`** on any table with capacity padding: `__ne__` negated the + result of `_dictionary_eq`, which had already been intersected with the live-row mask, so + `~(pred & valid)` turned every dead slot True. Pre-existing and unrelated to indexing; found by + fuzzing indexed against unindexed results. + +**The `where("c == 'x'")` string form is deliberately left alone** for both flavours. Rewriting to a +code comparison keeps it a single fused numeric expression; substituting a precomputed mask measured +*slower* (22.9 ms → 28.7 ms) even though the mask costs 4.8 ms. `plan_query` is still never consulted +for a utf8 or dictionary predicate — both routes bypass it rather than fix it, and accelerating the +string form needs the planner to consume index *positions* rather than a mask. + +**Also found here:** NumPy 2.4 does not match a lone `"\x00"` against a `StringDType` array +(`np.array(["\x00"], dtype=StringDType()) == "\x00"` is `False`), while `"\x00x"` and `"a\x00b"` +compare correctly. Every null mask in the utf8 paths is such a comparison, so that sentinel would +silently stop marking anything as null. `blosc2.utf8(null_value="\x00")` now rejects it. The +default sentinel is `'__BLOSC2_NULL__'`, so no shipped configuration was affected. + +--- + +## ⁸ The conversion pair — what shipped + +`f8af0714` and `5b31abe4`. The rule is now published rather than implied: **utf8 stores and filters; +fixed-width computes.** + +```python +fixed = blosc2.from_utf8(t["name"]) # -> "b"`) fails identically — the operand is widened to a `SimpleProxy` and the output + container is allocated from a `StringDType` the NDArray dtype round-trip cannot parse. So the + guard is on inputs, and the docs say so; an earlier draft of this document implied the output + type was the problem. +- **`lazyudf()` needed the guard twice.** The `DTypePromotionError` fires in the `lazyudf()` + function's dtype inference, before `LazyUDF.__init__` runs, so guarding the constructor alone + left `t.apply()` untouched. Both now check; `apply` also guards at the CTable level, because + `lazyudf` only ever sees the container and cannot say *which* column. + +Each printed recipe was run verbatim before the message shipped. + +Also fixed here: the `.. _Utf8Compute:` anchor added in `5b31abe4` collided with the +`[#utf8compute]` footnote label — docutils normalizes both to the same target name, which cost the +footnote its reference. Renamed to `ComputingUtf8Strings`. + +With this, items 1–4 of the priority list are done and only item 5 (drop G2/G3/G5, a decision +rather than work) remains. + +--- + +## ¹¹ NumPy `StringDType` convention — what was adopted, and what could not be + +`0ed38238`. The question was whether blosc2 should follow NumPy, which builds variable-length text +through a *dtype* (`np.array(v, dtype=StringDType())`) rather than through a separate constructor +(`blosc2.utf8_array(v)`). Answer: adopt the **dispatch**, not the dtype. + +**Why the dtype itself cannot be adopted.** `StringDType` is not a storage format: + +| | | +|---|---| +| `memoryview(arr)` | `ValueError: cannot include dtype 'StringDType' in a buffer` | +| `itemsize` | 16, whatever the content | +| `np.array(["x"*100], dtype=StringDType()).nbytes` | **16** — the payload is elsewhere | +| `.tobytes()` | a handle, not the text (≤15-byte strings are inlined; longer ones are pointers) | + +blosc2's NDArray compresses *buffers*, so `NDArray(dtype=StringDType())` would persist pointers — +garbage on reopen, in another process, or on another machine. Arrow reached the same conclusion, and +`UTF8Array`'s layout **is** Arrow's `large_string`, which is what makes `to_arrow` zero-copy. (The +`ast.literal_eval` failure in `NDArray.dtype` is a symptom, ~5 lines to fix, and fixing it buys +nothing.) + +**Why the schema layer was left alone.** `blosc2.field()` accepts a spec and never a raw dtype, for +*every* column type — `field(np.dtype("int32"))` is a `TypeError` too. Specs carry nullability, the +null sentinel, `ge`/`le`, storage config, `batch_rows`. Making utf8 the one dtype-addressable type +would have *broken* schema uniformity, not restored it. (Also: the runtime floor is `numpy>=1.26`, +where `StringDType` does not exist; `UTF8Spec.dtype = None` is deliberate.) + +**What shipped.** Constructors dispatch on the target dtype, matching NumPy's fill values exactly: + +```python +blosc2.asarray(np.array(["a", "bb"], dtype=StringDType())) # -> UTF8Array +blosc2.zeros(3, dtype=StringDType()) # -> UTF8Array, ['', '', ''] +blosc2.ones(3, dtype=StringDType()) # -> UTF8Array, ['1', '1', '1'] +blosc2.asarray(utf8_source, dtype=" NDArray, fixed width +``` + +Two container gaps closed along the way, both worth more than the dispatch: + +- **`UTF8Array` failed the `blosc2.Array` protocol**, and `.shape` was the *only* member it lacked — + for a container `CTable` uses throughout. It now has `.shape`/`.ndim`/`.size`. +- **`np.asarray(utf8_arr)` silently widened** to a fixed-width ` column` map, since they must still reach storage and the null sentinel by column name +while matching the alias in the expression. Both the raw-byte scalar-mask route and the span driver +are covered, including a leaf whose name is a prefix of another (`trip.who` under `trip.begin.who` — +longest-first ordering, same as the nested rewrite). + +**And the decision itself: G2, G3 and G4 are withdrawn**, recorded at the top of +`utf8-string-support.md` rather than here, so the plan cannot be picked up later without meeting the +verdict first. The short form: they would deliver an API indistinguishable from ``, `>=`) implementable on raw bytes without decoding. Python `str` comparison is code-point order, so byte-lex results match Python/StringDType semantics exactly. (The - same property already justifies `Utf8Factorizer`'s rank codes.) + same property already justifies `UTF8Factorizer`'s rank codes.) 4. **Null semantics are frozen.** A null (sentinel) value on either side never satisfies any comparison — SQL `WHERE` semantics, pinned by @@ -135,7 +135,7 @@ predicates into miniexpr only if a real workload later proves it pays. ### U1.a Equality (`==`, `!=`) against a `str` scalar -**Where:** a new method on `Utf8Array` (`src/blosc2/utf8_array.py`), plus +**Where:** a new method on `UTF8Array` (`src/blosc2/utf8_array.py`), plus wiring in `Column._utf8_compare` (`src/blosc2/ctable.py`, grep for `def _utf8_compare`). @@ -178,7 +178,7 @@ Key properties to preserve: (`idx = idx + 1` creates one new array per byte position; that is fine — the point is never materializing a `(k, L)` int64 index matrix). - **Pending rows:** call `self.flush()` at the start of the public entry - point (precedent: `Utf8Factorizer.__init__` and `factorize_span` flush; + point (precedent: `UTF8Factorizer.__init__` and `factorize_span` flush; it is a no-op unless there are buffered rows, and read-only tables cannot have any). @@ -209,7 +209,7 @@ currently duplicated. Behavior must not change (its tests pin it). **Algorithm:** per-byte vectorized lexicographic compare against the probe's bytes, grouped by row byte-length (the same grouping loop -`Utf8Array.factorize_span` uses — bincount on `np.diff(rel)`, then one +`UTF8Array.factorize_span` uses — bincount on `np.diff(rel)`, then one iteration per distinct length; distinct lengths are few in practice and each row is touched once regardless). @@ -371,7 +371,7 @@ extension. (own `add_custom_command`, `Python_add_library`, link/install rules). Measured at ~9 ns/row standalone (2e6-row synthetic column), matching the plan's 20-40 ns/row estimate. -- `Utf8Array._read_persisted_span` (`utf8_array.py`) tries the kernel via +- `UTF8Array._read_persisted_span` (`utf8_array.py`) tries the kernel via a new lazy `_pack_utf8_kernel()` helper (mirrors the `try: from blosc2 import groupby_ext / except ImportError: return None` pattern already used in `groupby.py`) and falls back to the old per-row @@ -385,10 +385,10 @@ extension. until both landed: 1. `Column._values_from_key`'s slice fast-path (`ctable.py`) excluded every `is_varlen_scalar` column, including utf8, from the - identity-position direct-slice shortcut, even though `Utf8Array` + identity-position direct-slice shortcut, even though `UTF8Array` slices itself efficiently. Changed the exclusion to `is_varlen_scalar and not is_utf8`. - 2. The real bottleneck: `Utf8Array._get_many` (used whenever + 2. The real bottleneck: `UTF8Array._get_many` (used whenever `_has_identity_positions()` is false — the common case, since a table's physical capacity is normally chunk-padded past its row count) always sorted the index array and did a fancy-indexed @@ -473,5 +473,5 @@ semantics for the C kernels — nothing built in U1 is throwaway. the root cause here and stop. - Never regress `string()`/`vlstring()` behavior or performance; the guard is `bench_string_kinds.py` plus the full test suite. -- No new public API: everything here is internal (`Utf8Array` methods, +- No new public API: everything here is internal (`UTF8Array` methods, `Column._utf8_compare` internals, an optional compiled helper). diff --git a/plans/utf8-string-support.md b/plans/utf8-string-support.md new file mode 100644 index 000000000..4b6b4a83c --- /dev/null +++ b/plans/utf8-string-support.md @@ -0,0 +1,239 @@ +# Compute parity for utf8 strings + +> ## Outcome: **G2, G3 and G4 withdrawn; G5 shipped.** Not superseded — decided against. +> +> The plan was written before the conversion pair existed. What shipped instead is the opposite +> rule: **utf8 stores and filters; fixed-width computes** (`f8af0714`, `5b31abe4`), with every +> compute-side refusal printing the two-line recipe that fixes it (`8e3868ba`). See +> `plans/string-flavours-assessment.md`, which measured the flavours end to end and is the +> document of record. +> +> | | verdict | why | +> |---|---|---| +> | G1 | moot | subsumed by G2, which is withdrawn | +> | **G2** computed columns | **withdrawn** | see below | +> | **G3** DSL kernels / `apply()` | **withdrawn** | same output-container problem, same recipe covers it | +> | **G4** bare `UTF8Array` | **withdrawn**, minus the two real bugs | `__eq__` fixed in `3692673f`; the wrong-path `lazyexpr` fixed in `0b486b07`. The remaining "lift the driver" work has no asked-for use case | +> | **G5** nested leaves | **shipped**, as a *query* fix | not a compute gap at all — see below | +> +> **Why G2 and G3 are withdrawn.** They would buy an API that looks identical to ` 3–5× slower (the span driver's decode + `astype` per span is unavoidable), paid for with a +> serialization hazard whose failure mode is an **unopenable table**: `_schema_dict_with_computed` +> saves `str(dtype)` and `np.dtype("StringDType()")` raises on load. The published rule is both +> cheaper and more honest — the `.astype()` the user writes *is* what the driver would have done +> silently. Reopen only on a concrete user request for utf8-typed computed columns; the ~1 week +> estimate below still stands, and the `"utf8"` dtype sentinel in §G2 is still the way to do it. +> +> **Why G5 was not withdrawn with them.** It is filed here under compute, but a nested utf8 leaf +> could not be *filtered* either — `t.where("trip.name == 'x'")` raised, while the same query on a +> ` nested columns, i.e. a hole in the rule the other four gaps were withdrawn in favour of. Fixed +> by aliasing dotted utf8 names in `_lazyexpr_over_cols`; the diagnosis in §G5 below was wrong +> about the mechanism (see the note there). + +Give `utf8()` columns (and `Utf8Array`) the same computing surface ` column` map so they can still reach storage and null sentinels. + +This is a **query** fix, not a compute one — hence shipping while G2/G3 are withdrawn. Scalar +comparisons, `startswith`/`upper`, mixed numeric predicates and `sum(where=)` all work on a dotted +utf8 leaf now, and the answers match the same data in a flat column. + +--- + +## Not gaps + +- **miniexpr.** Nothing. utf8 reaches it as fixed-width ` None: @@ -480,7 +480,7 @@ not investigated further; not gated by this plan. and stop. - Never regress `string()`/`vlstring()` ingest performance — guard with the full `bench_string_kinds.py` script, not just utf8's rows. -- No new public API — everything here is internal (`Utf8Array` methods, +- No new public API — everything here is internal (`UTF8Array` methods, one new lazy-import helper, one new compiled function in the existing `utf8_ext` module). @@ -488,7 +488,7 @@ not investigated further; not gated by this plan. ## Critical files -- `src/blosc2/utf8_array.py` — `Utf8Array.extend`, `_rewrite_from`, new +- `src/blosc2/utf8_array.py` — `UTF8Array.extend`, `_rewrite_from`, new `_encode_utf8_kernel()` helper (I1.a, I1.c, I2 caller-side wiring). - `src/blosc2/utf8_ext.pyx` — new `encode_utf8_span` function alongside the existing `pack_utf8_span` (I2 kernel). diff --git a/src/blosc2/__init__.py b/src/blosc2/__init__.py index cea161a22..764f8f78b 100644 --- a/src/blosc2/__init__.py +++ b/src/blosc2/__init__.py @@ -567,6 +567,7 @@ def _raise(exc): from .tree_store import TreeStore from .batch_array import Batch, BatchArray from .list_array import ListArray +from ._utf8_array import UTF8Array, from_utf8, to_utf8, utf8_array from .objectarray import ObjectArray, objectarray_from_cframe from .ref import Ref from .b2objects import open_b2object @@ -832,8 +833,12 @@ def _raise(exc): "uint32", "uint64", "utf8", + "utf8_array", "vlbytes", "vlstring", + # utf8 <-> fixed-width conversion + "from_utf8", + "to_utf8", # Grouped reductions "group_reduce", # Classes @@ -879,6 +884,7 @@ def _raise(exc): "Tuner", "URLPath", "ObjectArray", + "UTF8Array", # Version "__version__", # Utils diff --git a/src/blosc2/utf8_array.py b/src/blosc2/_utf8_array.py similarity index 57% rename from src/blosc2/utf8_array.py rename to src/blosc2/_utf8_array.py index 8c1d9a4a3..e7ad86182 100644 --- a/src/blosc2/utf8_array.py +++ b/src/blosc2/_utf8_array.py @@ -33,6 +33,7 @@ from __future__ import annotations import itertools +import operator from typing import TYPE_CHECKING, Any import numpy as np @@ -65,6 +66,241 @@ # groupby._factorize_fixed_width_str). _HASH_MIX = np.uint64(0x9E3779B97F4A7C15) +#: Nominal row span for evaluating an expression over utf8 operands. +UTF8_EXPR_SPAN = 65536 + +#: Byte ceiling for one span's fixed-width `` str: + """Message for every "utf8 does not compute here" refusal. + + The rule (utf8 stores and filters; fixed-width computes) is deliberate, so + the error's job is to route rather than merely refuse: *lead* states what + was rejected, and the shared tail spells out the three-line conversion, + parameterized on how the caller got here. + """ + write_back = ( + f" t.add_column('out', blosc2.utf8(), values=blosc2.to_utf8(res)) # or {source}.assign(res)" + if assignable + else " out = blosc2.to_utf8(res)" + ) + return ( + f"{lead}\n" + "utf8 stores and filters; fixed-width computes. Convert, compute, write back:\n" + f" fixed = blosc2.from_utf8({source})\n" + f" res = {compute}\n" + f"{write_back}\n" + "See 'Computing strings on a utf8 column' in the CTable reference docs." + ) + + +def utf8_span_dtype(span: np.ndarray) -> np.dtype: + """Fixed-width ``U`` dtype wide enough for every value in *span*. + + The width is span-local and data-dependent, whereas miniexpr bakes output + widths in at compile time, so round it up to a power of two: a column then + costs a handful of distinct compilations instead of one per span. + """ + longest = max((len(s) for s in span), default=0) + return np.dtype(f" 1: + raise ValueError( + f"utf8 operands carry different null sentinels ({distinct_sentinels}); " + "a string result can only have one, and picking one of them would " + "silently relabel the other's nulls. Give the operands a common " + "null_value, or compute on fixed-width arrays (blosc2.from_utf8)." + ) + if utf8_out is None: + utf8_out = UTF8Array(blosc2.utf8(null_value=null_value)) + # tolist() gives plain str, which is UTF8Array.extend's fast path. + utf8_out.extend(res.tolist()) + continue + if out is None: + out = np.zeros(n_phys, dtype=res.dtype) + out[start:stop] = res + if utf8_out is not None: + # Rows past the utf8 operands' logical length: the string zero value. + utf8_out.extend([""] * (n_phys - len(utf8_out))) + utf8_out.flush() + return utf8_out + if out is None: # no rows at all + out = np.zeros(n_phys, dtype=np.bool_) + return out + + +class UTF8LazyExpr: + """A deferred expression with at least one :class:`UTF8Array` operand. + + ``LazyExpr`` cannot hold one: a variable-width column is not an expression + operand, and wrapping it widens every row to a fixed `` None: + self.expression = expression + self.operands = dict(operands) + self._ne_args = ne_args + self._utf8 = {k: v for k, v in self.operands.items() if isinstance(v, UTF8Array)} + if not self._utf8: + raise ValueError("UTF8LazyExpr needs at least one UTF8Array operand") + + def __len__(self) -> int: + return min(len(v) for v in self._utf8.values()) + + @property + def shape(self) -> tuple[int, ...]: + return (len(self),) + + def compute(self, item=(), **kwargs): + """Evaluate the whole expression. + + Returns a :class:`UTF8Array` for a string result and a NumPy array for + a boolean or numeric one. ``strict_miniexpr=True`` asserts that + evaluation really did reach miniexpr rather than a NumPy fallback. + """ + if item not in ((), slice(None), Ellipsis): + raise NotImplementedError( + "expressions over a bare UTF8Array evaluate whole-array only; " + "call compute() and slice the result" + ) + strict = kwargs.pop("strict_miniexpr", False) + if kwargs: + raise TypeError(f"unexpected keyword arguments: {sorted(kwargs)}") + return utf8_span_eval( + self.expression, + {k: v for k, v in self.operands.items() if k not in self._utf8}, + self._utf8, + {k: v.spec.null_value for k, v in self._utf8.items()}, + len(self), + strict=strict, + # Read at call time, not bound as a default, so both are tunable. + span_rows=UTF8_EXPR_SPAN, + budget=UTF8_EXPR_BUDGET, + ) + + def __getitem__(self, item): + result = self.compute() + return result[item] + + def __str__(self) -> str: + return self.expression + + def __repr__(self) -> str: + return f"UTF8LazyExpr({self.expression!r}, shape={self.shape})" + + +# Fallback for comparisons against anything that is not a scalar str. +_COMPARE_OPS = { + "==": operator.eq, + "!=": operator.ne, + "<": operator.lt, + "<=": operator.le, + ">": operator.gt, + ">=": operator.ge, +} + def _factorize_byte_rows(mat: np.ndarray) -> tuple[np.ndarray, np.ndarray]: """Exact factorization of the rows of a ``(k, L)`` uint8 matrix. @@ -156,7 +392,7 @@ def _new_backend_arrays(cparams=None, dparams=None, *, offsets_urlpath=None, dat return offsets, data -class Utf8Array: +class UTF8Array: """Row-wise variable-length UTF-8 string array over offsets + bytes NDArrays. Provides the row-oriented interface expected by CTable columns: @@ -170,12 +406,12 @@ class Utf8Array: This class is internal; obtain instances via ``storage.create_varlen_scalar_column()`` or - ``storage.open_varlen_scalar_column()`` with a ``Utf8Spec``. + ``storage.open_varlen_scalar_column()`` with a ``UTF8Spec``. Parameters ---------- spec: - The :class:`~blosc2.schema.Utf8Spec` describing this column. + The :class:`~blosc2.schema.UTF8Spec` describing this column. offsets: ``int64`` NDArray of row offsets (length ``n + 1``). Created fresh (in memory) when ``None``. @@ -185,10 +421,10 @@ class Utf8Array: """ def __init__(self, spec, offsets=None, data=None) -> None: - from blosc2.schema import Utf8Spec + from blosc2.schema import UTF8Spec - if not isinstance(spec, Utf8Spec): - raise TypeError(f"Utf8Array requires a Utf8Spec, got {type(spec)!r}") + if not isinstance(spec, UTF8Spec): + raise TypeError(f"UTF8Array requires a UTF8Spec, got {type(spec)!r}") self._dtype = string_dtype() self._spec = spec if (offsets is None) != (data is None): @@ -251,6 +487,27 @@ def _read_persisted_span(self, a: int, b: int) -> np.ndarray: out[i] = blob[rel[i] : rel[i + 1]].decode("utf-8") return out + def _span_max_bytes(self, a: int, b: int) -> int: + """Longest UTF-8 byte length among rows ``[a, b)``. + + Read from the offsets alone -- no row is decoded. A byte length bounds + the codepoint count, so callers sizing a fixed-width ``U`` buffer can + use this directly. + """ + b = min(b, len(self)) + if b <= a: + return 0 + np_rows = self._persisted_rows + widest = 0 + if a < np_rows: + offs = np.asarray(self._offsets[a : min(b, np_rows) + 1], dtype=np.int64) + if offs.size > 1: + widest = int(np.diff(offs).max()) + if b > np_rows: + pending = self._pending[max(0, a - np_rows) : b - np_rows] + widest = max(widest, max((len(s.encode("utf-8")) for s in pending), default=0)) + return widest + def _gather_persisted(self, indices: np.ndarray) -> np.ndarray: """Gather persisted rows at *indices* (any order) as a StringDType array. @@ -291,7 +548,7 @@ def _get_many(self, indices: np.ndarray) -> np.ndarray: indices = np.where(indices < 0, indices + n, indices).astype(np.int64, copy=False) m = len(indices) if m and (indices.min() < 0 or indices.max() >= n): - raise IndexError("Utf8Array index out of range") + raise IndexError("UTF8Array index out of range") if m and indices[-1] - indices[0] == m - 1 and bool((np.diff(indices) == 1).all()): # A contiguous ascending run (e.g. a full-column read routed here # via an index array rather than a step-1 slice) is just a span @@ -390,7 +647,7 @@ def set_all(self, values: Iterable[Any]) -> None: Writes through the existing backing offsets/data NDArrays, so a store-backed column stays persistent (unlike building a fresh - in-memory ``Utf8Array``). Used by ``sort_by(inplace=True)`` and + in-memory ``UTF8Array``). Used by ``sort_by(inplace=True)`` and ``compact()`` to rewrite a column in a new row order. """ coerced = [self._coerce(v) for v in values] @@ -415,7 +672,7 @@ def __getitem__(self, index: int | slice | list | tuple | np.ndarray): if index < 0: index += n if not (0 <= index < n): - raise IndexError("Utf8Array index out of range") + raise IndexError("UTF8Array index out of range") if index >= self._persisted_rows: return self._pending[index - self._persisted_rows] return str(self._read_persisted_span(index, index + 1)[0]) @@ -434,7 +691,7 @@ def __getitem__(self, index: int | slice | list | tuple | np.ndarray): if isinstance(index, (list, tuple, np.ndarray)): return self._get_many(np.asarray(index, dtype=np.int64)) - raise TypeError(f"Utf8Array indices must be int, slice, or array; got {type(index)!r}") + raise TypeError(f"UTF8Array indices must be int, slice, or array; got {type(index)!r}") def __setitem__(self, index: int, value: Any) -> None: """Overwrite the value at *index*. @@ -444,14 +701,14 @@ def __setitem__(self, index: int, value: Any) -> None: an O(n - index) operation. """ if not isinstance(index, (int, np.integer)): - raise TypeError(f"Utf8Array assignment index must be int, got {type(index)!r}") + raise TypeError(f"UTF8Array assignment index must be int, got {type(index)!r}") value = self._coerce(value) n = len(self) index = int(index) if index < 0: index += n if not (0 <= index < n): - raise IndexError("Utf8Array index out of range") + raise IndexError("UTF8Array index out of range") if index >= self._persisted_rows: self._pending[index - self._persisted_rows] = value return @@ -481,6 +738,58 @@ def __setitem__(self, index: int, value: Any) -> None: ) self._bytes_used_cache = new_used + # ------------------------------------------------------------------ + # Comparisons + # ------------------------------------------------------------------ + + def _compare(self, other: Any, op: str) -> np.ndarray: + """Element-wise comparison, returning a boolean mask. + + A scalar ``str`` is answered by the raw-byte scanners, which never + decode a row. Anything else (a list, an ndarray, another + :class:`UTF8Array`) is materialized and handed to NumPy. + """ + if isinstance(other, str): + n = len(self) + if op in ("==", "!="): + mask = self.equal_mask_span(other, 0, n) + return ~mask if op == "!=" else mask + lt, gt = self.order_masks_span(other, 0, n) + return {"<": lt, ">": gt, "<=": ~gt, ">=": ~lt}[op] + right = other[:] if isinstance(other, UTF8Array) else other + return _COMPARE_OPS[op](np.asarray(self[:]), right) + + # Identity hashing is kept: these objects were hashable before __eq__ was + # defined, and an element-wise __eq__ never returns a bool for the hash + # contract to apply to. + __hash__ = object.__hash__ + + def __eq__(self, other: Any, /): + import blosc2 + + if blosc2._disable_overloaded_equal: + return self is other + return self._compare(other, "==") + + def __ne__(self, other: Any, /): + import blosc2 + + if blosc2._disable_overloaded_equal: + return self is not other + return self._compare(other, "!=") + + def __lt__(self, other: Any, /): + return self._compare(other, "<") + + def __le__(self, other: Any, /): + return self._compare(other, "<=") + + def __gt__(self, other: Any, /): + return self._compare(other, ">") + + def __ge__(self, other: Any, /): + return self._compare(other, ">=") + # ------------------------------------------------------------------ # Properties mirroring the interface expected by CTable # ------------------------------------------------------------------ @@ -494,6 +803,32 @@ def dtype(self): """The ``StringDType`` used for materialized reads.""" return self._dtype + @property + def shape(self) -> tuple[int, ...]: + """Row count as a 1-D shape. Completes the :class:`blosc2.Array` protocol.""" + return (len(self),) + + @property + def ndim(self) -> int: + """Always 1: a utf8 array is a flat sequence of strings.""" + return 1 + + @property + def size(self) -> int: + """Number of rows, matching NumPy's ``size`` for a 1-D array.""" + return len(self) + + def __array__(self, dtype=None, copy=None) -> np.ndarray: + """Materialize for NumPy, keeping ``StringDType`` unless asked otherwise. + + Without this, ``np.asarray`` falls back to iterating the rows and infers + a fixed-width `` float: return float("inf") return self.nbytes / cb - def factorizer(self) -> Utf8Factorizer: + def factorizer(self) -> UTF8Factorizer: """Return a fresh incremental factorizer over this column's rows.""" - return Utf8Factorizer(self) + return UTF8Factorizer(self) def factorize_span(self, a: int, b: int) -> tuple[np.ndarray, np.ndarray]: """Factorize rows ``[a, b)`` without decoding them. @@ -538,7 +873,7 @@ def factorize_span(self, a: int, b: int) -> tuple[np.ndarray, np.ndarray]: array of the distinct values sorted ascending and ``codes`` (int64, length ``b - a``) maps each row to its value — the same contract as ``np.unique(values, return_inverse=True)``, but computed from the raw - offsets/bytes buffers via :class:`Utf8Factorizer`: only the distinct + offsets/bytes buffers via :class:`UTF8Factorizer`: only the distinct values are ever decoded to ``str``. Pending rows are flushed first. """ fact = self.factorizer() @@ -648,18 +983,241 @@ def arrow_slice(self, pa, a: int, b: int, null_value: str | None = None): n, pa.py_buffer(rel), pa.py_buffer(data), validity, null_count ) - def copy(self, spec=None, **kwargs: Any) -> Utf8Array: + def _max_char_len(self, *, span_rows: int = UTF8_EXPR_SPAN) -> int: + """Longest value in codepoints, without decoding a row. + + A UTF-8 codepoint starts at every byte that is *not* a continuation + byte (``0b10xxxxxx``), so counting those over a row's byte range gives + its length in characters -- the width a fixed-width ``U`` array needs. + Byte lengths alone would only bound it, over-allocating up to 4x on + non-ASCII text. + + Byte lengths come from the offsets, so a span whose longest row cannot + beat the running best is skipped without reading any data at all, and + an all-ASCII span settles from the offsets alone. + """ + self.flush() + n = len(self) + widest = 0 + for start in range(0, n, span_rows): + stop = min(start + span_rows, n) + offs = np.asarray(self._offsets[start : stop + 1], dtype=np.int64) + byte_max = int(np.diff(offs).max(initial=0)) + if byte_max <= widest: + continue # bytes bound codepoints, so this span cannot win + raw = np.asarray(self._data[offs[0] : offs[-1]], dtype=np.uint8) + if not (raw & 0x80).any(): + widest = byte_max # pure ASCII: one byte per codepoint + continue + # Cumulative sums rather than reduceat: empty rows repeat an offset, + # which reduceat reads as "to the end" instead of as a zero count. + cumulative = np.concatenate(([0], np.cumsum((raw & 0xC0) != 0x80))) + rel = offs - offs[0] + widest = max(widest, int((cumulative[rel[1:]] - cumulative[rel[:-1]]).max(initial=0))) + return widest + + def astype(self, dtype=None, *, span_rows: int = UTF8_EXPR_SPAN) -> np.ndarray: + """Materialize as a fixed-width NumPy ``U`` array. + + This is the conversion half of the rule utf8 columns follow for + compute: they store and filter as variable-length text, and + string-returning expressions run on fixed-width arrays. See + :func:`blosc2.to_utf8` for the way back. + + Parameters + ---------- + dtype: + Target dtype. ``None`` or an unsized ``">> import blosc2 + >>> arr = blosc2.utf8_array(["hello", "café", "日本語"]) + >>> arr.astype().dtype + dtype('U", "=U")): + dtype = np.dtype(f" UTF8Array: """Return an in-memory copy.""" if spec is None: spec = self._spec - out = Utf8Array(spec) + out = UTF8Array(spec) out.extend(self) out.flush() return out -class Utf8Factorizer: - """Incremental factorizer over a :class:`Utf8Array`'s rows. +def utf8_array(seq, spec=None, **kwargs) -> UTF8Array: + """Build a :class:`UTF8Array` from an iterable of strings. + + Parameters + ---------- + seq: + Iterable of ``str`` (or ``None`` for a nullable *spec*). + spec: + The :class:`~blosc2.schema.UTF8Spec` describing the array. Defaults + to ``blosc2.utf8()`` (non-nullable). + kwargs: + Forwarded to :class:`UTF8Array` (``offsets``, ``data``). + + Returns + ------- + UTF8Array + + Examples + -------- + >>> import blosc2 + >>> arr = blosc2.utf8_array(["hello", "café", "日本語"]) + >>> str(arr[1]) + 'café' + """ + import blosc2 + + arr = UTF8Array(spec if spec is not None else blosc2.utf8(), **kwargs) + arr.extend(seq) + arr.flush() + return arr + + +def from_utf8(arr, dtype=None) -> np.ndarray: + """Convert variable-length UTF-8 text to a fixed-width NumPy ``U`` array. + + The outbound half of the utf8 compute rule: utf8 stores and filters text + compactly, while string-returning expressions and DSL kernels run on + fixed-width arrays. :func:`to_utf8` is the way back. + + Parameters + ---------- + arr: + A :class:`UTF8Array`, a utf8 :class:`~blosc2.CTable` column, a NumPy + ``StringDType`` array, or any iterable of ``str``. + dtype: + Target dtype. ``None`` (or an unsized ``">> import blosc2 + >>> arr = blosc2.utf8_array(["hello", "café"]) + >>> fixed = blosc2.from_utf8(arr) + >>> fixed.dtype + dtype('>> blosc2.to_utf8(fixed)[1] + 'café' + """ + raw = getattr(arr, "raw", arr) # a CTable Column exposes its container here + if isinstance(raw, UTF8Array): + return raw.astype(dtype) + values = np.asarray(raw if isinstance(raw, np.ndarray) else list(raw)) + if dtype is None or (isinstance(dtype, str) and dtype in ("U", "U", "=U")): + if values.dtype.kind == "U": + return values + dtype = np.dtype(f" UTF8Array: + """Build a :class:`UTF8Array` from fixed-width or otherwise decoded strings. + + The inbound half of the pair described in :func:`from_utf8`, and the way a + computed string result becomes storable again:: + + fixed = blosc2.from_utf8(t["name"]) + res = blosc2.lazyexpr("'x=' + a", {"a": fixed}).compute()[:] + t.add_column("prefixed", blosc2.utf8(), values=blosc2.to_utf8(res)) + + Parameters + ---------- + values: + NumPy ``U``/``StringDType`` array, or any iterable of ``str`` (or + ``None`` for a nullable *spec*). + spec: + The :class:`~blosc2.schema.UTF8Spec` describing the result. Defaults + to ``blosc2.utf8()`` (non-nullable). + + Returns + ------- + UTF8Array + """ + if isinstance(values, np.ndarray): + # tolist() yields plain str, which is UTF8Array.extend's fast path; + # iterating the array yields np.str_, which is not. + values = values.tolist() + return utf8_array(values, spec) + + +def is_string_dtype(dtype) -> bool: + """True for NumPy's variable-length ``StringDType`` (kind ``'T'``).""" + if dtype is None: + return False + try: + return np.dtype(dtype).kind == "T" + except TypeError: + # np.dtype() rejects StringDType passed as a class rather than instance. + return isinstance(dtype, type) and getattr(dtype, "kind", None) == "T" + + +def asarray_utf8(array, copy=None, **kwargs) -> UTF8Array: + """Back :func:`blosc2.asarray` when the *target* dtype is ``StringDType``. + + A ``StringDType`` array keeps its payload outside its own buffer (a 100 + character string still reports ``nbytes == 16``) and offers no buffer + protocol at all, so an :class:`~blosc2.NDArray` -- which compresses that + buffer -- cannot hold one: it would persist pointers. A + :class:`UTF8Array` holds the same text as offsets + UTF-8 bytes, the + layout Arrow uses for ``large_string``, so that is what this returns. + """ + if kwargs: + raise TypeError( + f"blosc2.asarray() does not accept {sorted(kwargs)!r} for variable-length " + "text; use blosc2.utf8_array(values, spec) to control its storage." + ) + ndim = getattr(array, "ndim", 1) + if ndim != 1: + raise ValueError( + f"Variable-length text is 1-D only, got a {ndim}-D array. Reshape it, or ask " + "for a fixed-width ' None: + def __init__(self, arr: UTF8Array) -> None: arr.flush() self._arr = arr self._values: list[str] = [] diff --git a/src/blosc2/blosc2_ext.pyx b/src/blosc2/blosc2_ext.pyx index ed325a565..3fc06cd8b 100644 --- a/src/blosc2/blosc2_ext.pyx +++ b/src/blosc2/blosc2_ext.pyx @@ -287,7 +287,7 @@ cdef extern from "blosc2.h": int blosc1_cbuffer_validate(const void* cbuffer, size_t cbytes, size_t* nbytes) - void blosc1_cbuffer_metainfo(const void* cbuffer, size_t* typesize, int* flags) + void blosc1_cbuffer_metainfo(const void* cbuffer, size_t* typesize, int* flags) nogil void blosc1_cbuffer_versions(const void* cbuffer, int* version, int* versionlz) @@ -396,6 +396,10 @@ cdef extern from "blosc2.h": int32_t srcsize, int start, int nitems, void* dest, int32_t destsize) nogil + int blosc2_getitem_bytes_ctx(blosc2_context* context, const void* src, + int32_t srcsize, int32_t start, int32_t nbytes, + void* dest, int32_t destsize) nogil + ctypedef struct blosc2_storage: @@ -674,6 +678,7 @@ cdef extern from "miniexpr.h": ME_COMPLEX64 ME_COMPLEX128 ME_STRING + ME_BYTES # typedef struct me_variable ctypedef struct me_variable: @@ -697,6 +702,9 @@ cdef extern from "miniexpr.h": int ncode void *parameters[1] + int me_compile(const char *expression, const me_variable *variables, + int var_count, me_dtype dtype, int *error, me_expr **out) + int me_compile_nd_jit(const char *expression, const me_variable *variables, int var_count, me_dtype dtype, int ndims, const int64_t *shape, const int32_t *chunkshape, @@ -740,6 +748,9 @@ cdef extern from "miniexpr.h": int me_nd_valid_nitems(const me_expr *expr, int64_t nchunk, int64_t nblock, int64_t *valid_nitems) nogil + me_dtype me_get_dtype(const me_expr *expr) nogil + size_t me_get_itemsize(const me_expr *expr) nogil + void me_print(const me_expr *n) nogil void me_free(me_expr *n) nogil @@ -901,9 +912,137 @@ cdef inline me_dtype _me_dtype_from_numpy_dtype(dtype_obj): f"miniexpr string operands require unicode dtype with UCS4 itemsize; got '{dtype}'" ) return ME_STRING + elif kind == "S": + if itemsize <= 0: + raise TypeError(f"miniexpr bytes operands require a non-empty itemsize; got '{dtype}'") + return ME_BYTES return -1 +cdef inline object _numpy_dtype_from_me_dtype(me_dtype dt): + if dt == ME_BOOL: + return np.dtype(np.bool_) + if dt == ME_INT8: + return np.dtype(np.int8) + if dt == ME_INT16: + return np.dtype(np.int16) + if dt == ME_INT32: + return np.dtype(np.int32) + if dt == ME_INT64: + return np.dtype(np.int64) + if dt == ME_UINT8: + return np.dtype(np.uint8) + if dt == ME_UINT16: + return np.dtype(np.uint16) + if dt == ME_UINT32: + return np.dtype(np.uint32) + if dt == ME_UINT64: + return np.dtype(np.uint64) + if dt == ME_FLOAT32: + return np.dtype(np.float32) + if dt == ME_FLOAT64: + return np.dtype(np.float64) + if dt == ME_COMPLEX64: + return np.dtype(np.complex64) + if dt == ME_COMPLEX128: + return np.dtype(np.complex128) + return None + + +def me_output_dtype(expression, operands): + """Ask miniexpr what dtype *expression* would produce over *operands*. + + ``operands`` maps operand name -> numpy dtype. Compiles with ME_AUTO, reads + the inferred result back, and throws the program away. python-blosc2 needs + this before evaluating, because the output container must be allocated with a + fixed itemsize and string widths are known only to miniexpr's own inference + (e.g. ` ` 0: + variables = malloc(sizeof(me_variable) * n) + if variables == NULL: + raise MemoryError() + + try: + for k, v in operands.items(): + var = &variables[built] + operand_dtype = np.dtype(v) + try: + var.dtype = _me_dtype_from_numpy_dtype(operand_dtype) + except TypeError: + return None + if var.dtype < 0: + return None + var_name = k.encode("utf-8") if isinstance(k, str) else k + var.name = malloc(strlen(var_name) + 1) + strcpy(var.name, var_name) + var.address = NULL + var.type = 0 + var.context = NULL + var.itemsize = operand_dtype.itemsize if operand_dtype.num in (18, 19) else 0 + built += 1 + + expression_bytes = ( + (expression).encode("utf-8") if isinstance(expression, str) else expression + ) + rc = me_compile(expression_bytes, variables, n, ME_AUTO, &error, &out_expr) + if rc != ME_COMPILE_SUCCESS or out_expr == NULL: + if out_expr != NULL: + me_free(out_expr) + return None + + out_dt = me_get_dtype(out_expr) + itemsize = me_get_itemsize(out_expr) + me_free(out_expr) + + if out_dt == ME_STRING: + if itemsize == 0 or itemsize % 4 != 0: + return None + return np.dtype(" blocknitems: raise ValueError("miniexpr: valid items exceed padded block size") - rc = blosc2_getitem_ctx(dctx, src, chunk_cbytes, start, blocknitems, - input_buffers[i], block_nbytes) + # Ask in bytes: blosc2_getitem_ctx() counts in the typesize the *chunk* + # records, which c-blosc2 caps to 1 above BLOSC_MAX_TYPESIZE (255), so its + # unit changes silently with the data -- every block past the first once + # came back as uninitialised memory here. Bytes are unambiguous, and a + # block offset is already one. + rc = blosc2_getitem_bytes_ctx(dctx, src, chunk_cbytes, nblock * block_nbytes, + block_nbytes, input_buffers[i], block_nbytes) blosc2_free_ctx(dctx) - if rc < 0: + if rc != block_nbytes: raise ValueError("miniexpr: error decompressing the chunk") # For reduction operations, we need to track which block we're processing # The linear_block_index should be based on the same grid the output shares @@ -2787,7 +2930,6 @@ cdef int aux_matmul(mm_udata *udata, int64_t nchunk, int32_t nblock, void *param cdef int32_t chunk_nbytes[2] cdef int32_t chunk_cbytes[2] cdef int32_t block_nbytes[2] - cdef int blocknitems[2] cdef int startA, startB, expected_blocknitems cdef blosc2_context* dctx cdef int i, j, block_i, block_j, chunk_i, chunk_j, ncols, block_ncols, Bblock_ncols, Bncols, Ablock_ncols, Ancols @@ -2871,21 +3013,24 @@ cdef int aux_matmul(mm_udata *udata, int64_t nchunk, int32_t nblock, void *param input_buffers[i] = malloc(block_nbytes[i]) if input_buffers[i] == NULL: raise MemoryError("miniexpr: cannot allocate input block buffer") - blocknitems[i] = block_nbytes[i] // ndarr.sc.typesize first_run = False nblockA = block_startA nblockB = block_startB while True: # block loop - startA = nblockA * blocknitems[0] - startB = nblockB * blocknitems[1] - rc = blosc2_getitem_ctx(dctx, src[0], chunk_cbytes[0], startA, blocknitems[0], - input_buffers[0], block_nbytes[0]) - if rc < 0: + startA = nblockA * block_nbytes[0] + startB = nblockB * block_nbytes[1] + # In bytes, for the reason given in aux_miniexpr(). matmul only ever + # sees numeric scalars (typesize <= 8), so the capped-typesize case is + # unreachable here today; asking in the unambiguous unit keeps that from + # becoming a latent trap if it ever stops being true. + rc = blosc2_getitem_bytes_ctx(dctx, src[0], chunk_cbytes[0], startA, + block_nbytes[0], input_buffers[0], block_nbytes[0]) + if rc != block_nbytes[0]: raise ValueError("matmul: error decompressing the A chunk") - rc = blosc2_getitem_ctx(dctx, src[1], chunk_cbytes[1], startB, blocknitems[1], - input_buffers[1], block_nbytes[1]) - if rc < 0: + rc = blosc2_getitem_bytes_ctx(dctx, src[1], chunk_cbytes[1], startB, + block_nbytes[1], input_buffers[1], block_nbytes[1]) + if rc != block_nbytes[1]: raise ValueError("matmul: error decompressing the B chunk") batch = 0 while batch < batches: @@ -3729,6 +3874,7 @@ cdef class NDArray: cdef int rc cdef int32_t lazychunk_cbytes cdef c_bool owns_dctx = False + cdef int32_t want_nbytes lazychunk_cbytes = blosc2_schunk_get_lazychunk(self.array.sc, nchunk, &chunk, &needs_free) if lazychunk_cbytes < 0: @@ -3759,10 +3905,16 @@ cdef class NDArray: if needs_free: free(chunk) raise RuntimeError("Could not create decompression context") + # In bytes, for the reason given in aux_miniexpr(): an index summary over a + # calloc(ninputs, sizeof(uint8_t*)) np_typesizes = calloc(ninputs, sizeof(int32_t)) if np_data == NULL or np_typesizes == NULL: - free(inputs_) - free(np_data) - free(np_typesizes) - free(udata) + _free_me_udata_tables(udata, inputs_, np_data, np_typesizes) raise MemoryError("Cannot allocate miniexpr raw-input tables") for i, operand in enumerate(operands): if isinstance(operand, np.ndarray): @@ -4034,8 +4189,7 @@ cdef class NDArray: if ninputs > 0: input_chunk_caches = calloc(ninputs, sizeof(me_input_cache_s)) if input_chunk_caches == NULL: - free(inputs_) - free(udata) + _free_me_udata_tables(udata, inputs_, np_data, np_typesizes) raise MemoryError("Cannot allocate miniexpr chunk caches") for i in range(ninputs): input_chunk_caches[i].nchunk = -1 @@ -4049,8 +4203,7 @@ cdef class NDArray: if input_chunk_caches[i].ready_lock != NULL: PyThread_free_lock(input_chunk_caches[i].ready_lock) free(input_chunk_caches) - free(inputs_) - free(udata) + _free_me_udata_tables(udata, inputs_, np_data, np_typesizes) raise MemoryError("Cannot allocate miniexpr chunk cache state lock") input_chunk_caches[i].ready_lock = PyThread_allocate_lock() if input_chunk_caches[i].ready_lock == NULL: @@ -4063,8 +4216,7 @@ cdef class NDArray: if input_chunk_caches[i].ready_lock != NULL: PyThread_free_lock(input_chunk_caches[i].ready_lock) free(input_chunk_caches) - free(inputs_) - free(udata) + _free_me_udata_tables(udata, inputs_, np_data, np_typesizes) raise MemoryError("Cannot allocate miniexpr chunk cache ready lock") udata.input_chunk_caches = input_chunk_caches eval_params = malloc(sizeof(me_eval_params)) @@ -4075,8 +4227,7 @@ cdef class NDArray: if input_chunk_caches[i].ready_lock != NULL: PyThread_free_lock(input_chunk_caches[i].ready_lock) free(input_chunk_caches) - free(inputs_) - free(udata) + _free_me_udata_tables(udata, inputs_, np_data, np_typesizes) raise MemoryError("Cannot allocate miniexpr eval params") eval_params.disable_simd = False eval_params.simd_ulp_mode = ME_SIMD_ULP_3_5 if fp_accuracy == blosc2.FPAccuracy.MEDIUM else ME_SIMD_ULP_1 @@ -4188,7 +4339,7 @@ cdef class NDArray: var.address = NULL # chunked compile: addresses provided later var.type = 0 # auto-set to ME_VARIABLE inside compiler var.context = NULL - var.itemsize = v.dtype.itemsize if v.dtype.num == 19 else 0 # only store item type if string + var.itemsize = v.dtype.itemsize if v.dtype.num in (18, 19) else 0 # only store item size for strings/bytes cdef int error = 0 cdef bytes expression_bytes @@ -4216,6 +4367,15 @@ cdef class NDArray: raise TypeError(f"miniexpr does not support operand or output dtype: {expression_display}; details: {me_error_msg}") if rc != ME_COMPILE_SUCCESS: raise NotImplementedError(f"Cannot compile expression: {expression_display}; details: {me_error_msg}") + # The output container was allocated before compiling, so a width miniexpr + # infers differently from the container's would overrun the block buffer. + cdef size_t inferred_itemsize = me_get_itemsize(out_expr) + if me_output_dtype in (ME_STRING, ME_BYTES) and inferred_itemsize != out_np_dtype.itemsize: + me_free(out_expr) + raise ValueError( + f"miniexpr infers a {inferred_itemsize}-byte string result for " + f"{expression_display}, but the output array is {out_np_dtype}" + ) udata.miniexpr_handle = out_expr # Free resources @@ -4279,7 +4439,7 @@ cdef class NDArray: var.address = NULL var.type = 0 var.context = NULL - var.itemsize = v.dtype.itemsize if v.dtype.num == 19 else 0 + var.itemsize = v.dtype.itemsize if v.dtype.num in (18, 19) else 0 cdef bytes expression_bytes = ( (expression).encode("utf-8") if isinstance(expression, str) else expression diff --git a/src/blosc2/ctable.py b/src/blosc2/ctable.py index 3794f9ce9..cab1513dd 100644 --- a/src/blosc2/ctable.py +++ b/src/blosc2/ctable.py @@ -15,6 +15,7 @@ import contextvars import copy import dataclasses +import itertools import json import operator import os @@ -26,7 +27,7 @@ from dataclasses import MISSING, dataclass from dataclasses import field as dataclass_field from textwrap import TextWrapper -from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar +from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, TypeVar import numpy as np @@ -54,7 +55,7 @@ ObjectSpec, SchemaSpec, StructSpec, - Utf8Spec, + UTF8Spec, VLBytesSpec, VLStringSpec, complex64, @@ -540,6 +541,33 @@ def __repr__(self) -> str: # --------------------------------------------------------------------------- +def _rank_index_row_lookup(values_path: str, positions_path: str, table, null_rank: int): + """Build ``rows_for_ranks(lo, hi)`` over a rank index's sorted sidecars. + + Shared by the utf8 and dictionary rank indexes, which differ only in how a + literal becomes a rank. Returns the physical rows whose rank lies in + ``[lo, hi)``; ``hi is None`` means "up to but excluding the nulls", which + carry the largest rank because a null satisfies no comparison. + """ + from blosc2.indexing import _open_sidecar_file + + vnd = _open_sidecar_file(values_path) + pnd = _open_sidecar_file(positions_path) + + def rows_for_ranks(rank_lo, rank_hi) -> np.ndarray: + start = table._sidecar_bisect(vnd, rank_lo, "left") + stop = ( + table._sidecar_bisect(vnd, null_rank, "left") + if rank_hi is None + else table._sidecar_bisect(vnd, rank_hi - 1, "right") + ) + if stop <= start: + return np.empty(0, dtype=np.int64) + return np.asarray(pnd[start:stop], dtype=np.int64) + + return rows_for_ranks + + def _find_physical_index(arr: blosc2.NDArray, logical_key: int) -> int: """Translate a logical (valid-row) index into a physical array index. @@ -1090,14 +1118,14 @@ def is_varlen_scalar(self) -> bool: """True if this column holds variable-length scalar strings or bytes.""" col = self._table._schema.columns_by_name.get(self._col_name) return col is not None and isinstance( - col.spec, (VLStringSpec, VLBytesSpec, StructSpec, ObjectSpec, Utf8Spec) + col.spec, (VLStringSpec, VLBytesSpec, StructSpec, ObjectSpec, UTF8Spec) ) @property def is_utf8(self) -> bool: """True if this column stores variable-length UTF-8 strings (offsets + bytes).""" col = self._table._schema.columns_by_name.get(self._col_name) - return col is not None and isinstance(col.spec, Utf8Spec) + return col is not None and isinstance(col.spec, UTF8Spec) @property def is_dictionary(self) -> bool: @@ -1227,7 +1255,7 @@ def _values_from_key(self, key, *, check_stale: bool = True): # noqa: C901 # letting NDArray's strided-gather fast path handle coarse steps. # Plain stored columns only; everything else falls through to the # position-gather path below. utf8 is a varlen-scalar kind but - # Utf8Array slices itself efficiently (offsets+bytes span read), + # UTF8Array slices itself efficiently (offsets+bytes span read), # so it takes the fast path too instead of the index-gather one. if ( not ( @@ -2014,10 +2042,7 @@ def __eq__(self, other): def __ne__(self, other): if self.is_dictionary: - result = self._dictionary_eq(other) - if isinstance(result, np.ndarray): - return ~result - return ~np.asarray(result, dtype=bool) + return self._dictionary_eq(other, negate=True) if self.is_utf8: return self._utf8_compare(np.not_equal, other) self._ensure_comparable() @@ -2029,7 +2054,7 @@ def _utf8_chunked_bool(self, fn, *, chunk_size: int = 65536) -> np.ndarray: """Apply ``fn(chunk, start, stop)`` over this utf8 column's logical rows. *fn* returns a boolean array for each ``StringDType`` chunk read from - the underlying :class:`~blosc2.utf8_array.Utf8Array`. Returns a + the underlying :class:`~blosc2._utf8_array.UTF8Array`. Returns a physical-length (``_valid_rows``-length) boolean NumPy array; rows beyond the column's logical length are left ``False``. """ @@ -2046,7 +2071,7 @@ def _utf8_chunked_bytes(self, fn, *, chunk_size: int = 65536) -> np.ndarray: """Apply ``fn(arr, start, stop)`` over this utf8 column's logical rows. Like :meth:`_utf8_chunked_bool`, but *fn* operates directly on the - underlying :class:`~blosc2.utf8_array.Utf8Array` (raw offsets/bytes) + underlying :class:`~blosc2._utf8_array.UTF8Array` (raw offsets/bytes) instead of a materialized ``StringDType`` chunk, so no per-row decode happens. Returns a physical-length boolean NumPy array; rows beyond the column's logical length are left ``False``. @@ -2104,14 +2129,21 @@ def fn(chunk, start, stop): raw = self._utf8_chunked_bool(fn) return blosc2.asarray(raw) & self._lazy_valid_rows() - def _utf8_compare_scalar(self, numpy_op, value: str): - """Scalar comparison, evaluated chunk by chunk directly on raw UTF-8 - bytes (no decode to ``StringDType``) via - :meth:`~blosc2.utf8_array.Utf8Array.equal_mask_span` / - :meth:`~blosc2.utf8_array.Utf8Array.order_masks_span`. + def _utf8_scalar_mask(self, numpy_op, value: str) -> np.ndarray: + """Raw physical-length boolean mask for ``column value``. + + Compares raw UTF-8 bytes with no decode to ``StringDType``, via + :meth:`~blosc2._utf8_array.UTF8Array.equal_mask_span` / + :meth:`~blosc2._utf8_array.UTF8Array.order_masks_span`. A null never + satisfies any comparison. Not intersected with the live-row mask -- + see :meth:`_utf8_compare_scalar` for that. """ nv = self.null_value + indexed = self._utf8_index_mask(numpy_op, value) + if indexed is not None: + return indexed + if numpy_op in (np.equal, np.not_equal): def fn(arr, start, stop): @@ -2137,10 +2169,110 @@ def fn(arr, start, stop): res = res & ~arr.equal_mask_span(nv, start, stop) return res - raw = self._utf8_chunked_bytes(fn) - return blosc2.asarray(raw) & self._lazy_valid_rows() + return self._utf8_chunked_bytes(fn) + + #: Rank predicate implied by each comparison, given the literal's insertion + #: point ``lo`` and whether the literal is itself in the vocabulary. + _UTF8_RANK_PREDICATE: ClassVar[dict] = { + "equal": lambda lo, hit: (lo, lo + 1) if hit else None, + "not_equal": lambda lo, hit: (lo, lo + 1) if hit else None, # inverted by caller + "less": lambda lo, hit: (0, lo), + "less_equal": lambda lo, hit: (0, lo + 1 if hit else lo), + "greater": lambda lo, hit: (lo + 1 if hit else lo, None), + "greater_equal": lambda lo, hit: (lo, None), + } + + def _utf8_index_mask(self, numpy_op, value: str) -> np.ndarray | None: + """Answer ``column value`` from the rank index, or ``None``. + + The index sorts rows by alphabetical rank, so a literal is located by + one ``searchsorted`` over the stored vocabulary and the matching rows + are a contiguous run of the sorted-positions sidecar — no scan of the + column at all. Returns ``None`` whenever the index cannot answer, and + the caller falls back to the raw-byte scan. + """ + table = self._table + descriptor = table._get_index_catalog().get(self._col_name) + if not descriptor or descriptor.get("kind") != "full" or descriptor.get("stale", False): + return None + full = descriptor.get("full") or {} + meta = full.get("utf8_rank") + if meta is None or table._utf8_rank_index_stale(self._col_name, meta): + return None + positions_path = full.get("positions_path") + values_path = full.get("values_path") + if positions_path is None or values_path is None: # in-memory sidecars + return None + + vocab = table._utf8_index_vocab(self._col_name, meta) + if vocab is None: + return None + lo = int(np.searchsorted(vocab, value, side="left")) + hit = lo < len(vocab) and vocab[lo] == value + bounds = self._UTF8_RANK_PREDICATE[numpy_op.__name__](lo, hit) + + null_rank = meta["null_rank"] + n_phys = len(table._valid_rows) + rows_for_ranks = _rank_index_row_lookup(values_path, positions_path, table, null_rank) + + mask = np.zeros(n_phys, dtype=bool) + if numpy_op is np.not_equal: + # Invert over the column's own rows only: the physical mask is + # capacity-padded, and padded slots must stay False, as they do on + # the scan path. Nulls are excluded rather than inverted into. + mask[: len(table._cols[self._col_name])] = True + if bounds is not None: + mask[rows_for_ranks(*bounds)] = False + mask[rows_for_ranks(null_rank, null_rank + 1)] = False + elif bounds is not None: + mask[rows_for_ranks(*bounds)] = True + return mask + + def _utf8_compare_scalar(self, numpy_op, value: str): + """Scalar comparison as a live-row-intersected boolean NDArray.""" + return blosc2.asarray(self._utf8_scalar_mask(numpy_op, value)) & self._lazy_valid_rows() + + def _dictionary_index_mask(self, value: str) -> np.ndarray | None: + """Answer ``column == value`` from the dict-rank index, or ``None``. + + The same lookup the utf8 rank index does, minus the persistence: a + dictionary already holds its own vocabulary in memory, so the literal's + rank is a ``searchsorted`` over the sorted dictionary. Returns ``None`` + whenever the index cannot answer, and the caller falls back to the + codes comparison. + """ + table = self._table + descriptor = table._get_index_catalog().get(self._col_name) + if not descriptor or descriptor.get("kind") != "full" or descriptor.get("stale", False): + return None + full = descriptor.get("full") or {} + meta = full.get("dict_rank") + if meta is None or table._dict_rank_index_stale(self._col_name, meta): + return None + values_path, positions_path = full.get("values_path"), full.get("positions_path") + if positions_path is None or values_path is None: # in-memory sidecars + return None - def _dictionary_eq(self, other): + # Ranks were assigned by argsort over the dictionary, so the rank of a + # literal is its position in the sorted dictionary. + dc = self._raw_col + cache = table.__dict__.setdefault("_dict_vocab_cache", {}) + key = (self._col_name, meta.get("dict_hash")) + sorted_vocab = cache.get(key) + if sorted_vocab is None: + cache.clear() + sorted_vocab = np.sort(np.asarray(list(dc.dictionary), dtype=np.str_)) + cache[key] = sorted_vocab + lo = int(np.searchsorted(sorted_vocab, value, side="left")) + if lo >= len(sorted_vocab) or sorted_vocab[lo] != value: + return np.zeros(len(table._valid_rows), dtype=bool) + + rows_for_ranks = _rank_index_row_lookup(values_path, positions_path, table, meta["null_rank"]) + mask = np.zeros(len(table._valid_rows), dtype=bool) + mask[rows_for_ranks(lo, lo + 1)] = True + return mask + + def _dictionary_eq(self, other, *, negate: bool = False): """Return a physical-slot boolean predicate for dictionary equality. Regular fixed-width columns build predicates against their raw physical @@ -2148,25 +2280,37 @@ def _dictionary_eq(self, other): need to use the same coordinate system so they can be combined with regular predicates before aggregate/view code intersects them with ``_valid_rows``. + + *negate* inverts the value test *before* the live-row intersection, so + ``!=`` stays a same-shaped predicate over live rows. Negating the + returned value instead would turn every dead slot True. """ + n_phys = len(self._table._valid_rows) dc = self._raw_col # DictionaryColumn spec = self._table._schema.columns_by_name[self._col_name].spec + valid = self._lazy_valid_rows() if other is None: target_code = spec.null_code elif isinstance(other, str): + indexed = self._dictionary_index_mask(other) + if indexed is not None: + return blosc2.asarray(~indexed if negate else indexed) & valid try: target_code = dc.value_to_code(other) except KeyError: - return blosc2.zeros(len(self._table._valid_rows), dtype=np.bool_) + # No row carries this value: nothing matches, everything differs. + if negate: + return blosc2.ones(n_phys, dtype=np.bool_) & valid + return blosc2.zeros(n_phys, dtype=np.bool_) else: raise TypeError( f"Dictionary column {self._col_name!r} can only be compared with str or None, " f"got {type(other).__name__!r}." ) - pred = dc.codes == np.int32(target_code) - valid = self._lazy_valid_rows() - if len(dc.codes) != len(self._table._valid_rows): - physical = blosc2.zeros(len(self._table._valid_rows), dtype=np.bool_) + code = np.int32(target_code) + pred = dc.codes != code if negate else dc.codes == code + if len(dc.codes) != n_phys: + physical = blosc2.zeros(n_phys, dtype=np.bool_) physical[: len(dc.codes)] = pred pred = physical return pred & valid @@ -2396,6 +2540,12 @@ def assign(self, data) -> None: root._mark_generated_columns_stale(self._col_name) root._mark_all_indexes_stale() return + if self.is_varlen_scalar: + self._assign_varlen_scalar(data) + root = self._table._root_table + root._mark_generated_columns_stale(self._col_name) + root._mark_all_indexes_stale() + return n_live = len(self) arr = np.asarray(data) if len(arr) != n_live: @@ -2410,6 +2560,30 @@ def assign(self, data) -> None: root._mark_generated_columns_stale(self._col_name) root._mark_all_indexes_stale() + def _assign_varlen_scalar(self, data) -> None: + """``assign()`` for utf8/vlstring/vlbytes/struct/object columns. + + These are rewritten whole rather than row by row: overwriting one row + of a utf8 column shifts every later offset, and one row of a batched + varlen column rewrites its whole batch, so the loop would be + quadratic. Dead slots keep their current contents. + """ + values = list(data) + n_live = len(self) + if len(values) != n_live: + raise ValueError(f"assign() requires {n_live} values (live rows), got {len(values)}.") + raw = self._raw_col + raw.flush() + n_phys = len(raw) + if n_live == n_phys: + raw.set_all(values) + return + current = list(raw[:]) + live_pos = np.flatnonzero(self._valid_rows[:n_phys]) + for pos, value in zip(live_pos, values, strict=True): + current[int(pos)] = value + raw.set_all(current) + # ------------------------------------------------------------------ # Null sentinel support # ------------------------------------------------------------------ @@ -2580,10 +2754,13 @@ def _normalize_sum_where(self, where): if where is None: return None if isinstance(where, str): - self._table._guard_varlen_scalar_expression(where) + self._table._guard_varlen_scalar_expression(where, allow_utf8=True) + utf8_names = self._table._utf8_names_in(where) operands = self._table._where_expression_operands(where) where, operands = self._table._rewrite_nested_expression(where, operands) - where = blosc2.lazyexpr(where, operands) + where = self._table._lazyexpr_over_cols(where, operands, utf8_names) + if isinstance(where, np.ndarray): + where = blosc2.asarray(where) if isinstance(where, np.ndarray) and where.dtype == np.bool_: where = blosc2.asarray(where) if isinstance(where, Column): @@ -2825,17 +3002,27 @@ def _lazy_aggregate_fastpath(self, op: str, *, where=None, dtype=None, ddof: int return NotImplemented def _summary_minmax_source(self): - """Return ``(sidecar_path, dtype, nullable)`` for a summary-readable - ``min``/``max``, or ``None`` when the index shortcut is not provably - correct. + """Return ``(sidecar_path, dtype, nullable, segment_len)`` for a + summary-readable ``min``/``max``, or ``None`` when the index shortcut is + not provably correct. Excluded: a view (its summary describes the base table); a column kind without numeric/string block extrema; a leaky null sentinel (only a non-nullable column, or a NaN-sentinel float — whose NaNs the summary drops — match the nulls-skipped contract of ``min()``); and a stale, - absent, or in-memory-only index. Deletions/appends are covered by the - stale flag (every mutation marks the index stale; a rebuild re-summarises - only the live rows, and capacity padding never enters the summaries). + absent, or in-memory-only index. + + Appends mark the index stale, so they are covered. Deletions are *not* + (``delete()`` tombstones in place and leaves the index usable for + queries), and the summaries are built over the column's *physical* + array, where a tombstoned row keeps contributing its value to its + block — so any hole at all disqualifies the shortcut, whether it was + punched before or after the build. With no holes the physical and + logical row numbers coincide, which is what lets the caller mix + summary blocks with a rescanned tail. Capacity padding *does* enter + the summaries; ``segment_len`` is returned so the caller can drop the + padded tail and rescan the one block that straddles the live/padded + boundary. """ table = self._table if table.base is not None: @@ -2858,9 +3045,16 @@ def _summary_minmax_source(self): is_nan_float = dtype.kind == "f" and isinstance(null_value, float) and np.isnan(null_value) if nullable and not is_nan_float: return None # non-NaN sentinel leaks into the block extrema - desc = table._root_table._get_index_catalog().get(self._col_name) + root = table._root_table + desc = root._get_index_catalog().get(self._col_name) if not desc or desc.get("stale", False): return None + # A tombstoned row still sits in its block and still contributes to that + # block's extrema, and the summaries index physical slots while min() + # reads logical rows. Both only line up while every slot below the + # watermark is live. + if root._n_rows is None or root._n_rows != root._resolve_last_pos(): + return None levels = desc.get("levels") or {} level = "block" if "block" in levels else next(iter(levels), None) if level is None: @@ -2868,7 +3062,10 @@ def _summary_minmax_source(self): path = levels[level].get("path") if path is None: return None # in-memory sidecar: the scan is already fast - return path, dtype, nullable + segment_len = levels[level].get("segment_len") + if not segment_len: + return None + return path, dtype, nullable, int(segment_len) def _index_summary_minmax(self, op: str): """Exact ``min``/``max`` from the column index's block summaries, or @@ -2878,11 +3075,21 @@ def _index_summary_minmax(self, op: str): Every index kind (SUMMARY/FULL/PARTIAL/BUCKET/OPSI) persists per-block ``(min, max, flags)``, so reducing those is decompression-free (~240x faster than scanning tens of millions of rows). + + The summaries cover the column's *physical* extent, which is padded out + to the slot capacity with zeros/empty strings — values that beat any real + datum on ``min``. Only whole blocks below ``n_rows`` are read from the + sidecar; the single block straddling the boundary is rescanned (one block + decompression, so the shortcut is preserved) and blocks past it dropped. """ source = self._summary_minmax_source() if source is None: return NotImplemented - path, dtype, nullable = source + path, dtype, nullable, segment_len = source + n_live = self._table._root_table._n_rows + if n_live is None or n_live == 0: + return NotImplemented + n_full = n_live // segment_len # blocks entirely within the live rows try: from blosc2.indexing import _INDEX_MMAP_MODE, FLAG_ALL_NAN, FLAG_HAS_NAN, _open_sidecar_file @@ -2897,14 +3104,37 @@ def _index_summary_minmax(self, op: str): return NotImplemented if vals.shape[0] == 0: return NotImplemented + # Drop the padded tail: keep only blocks lying wholly below n_rows. + flags = flags[:n_full] + vals = vals[:n_full] # A non-nullable float with NaN *data* makes numpy min/max return NaN, # but the summary dropped those NaNs — they would disagree, so bail. if dtype.kind == "f" and not nullable and bool((flags & (FLAG_HAS_NAN | FLAG_ALL_NAN)).any()): return NotImplemented valid = (flags & FLAG_ALL_NAN) == 0 - if not valid.any(): - return NotImplemented # whole column null → let the scan raise vals = vals[valid] + + # The straddling block is not summarisable (its tail is padding), so read + # just its live rows. This is also the whole answer when the column is + # shorter than one block, in which case no summary entry is usable. + tail = n_live - n_full * segment_len + if tail: + try: + seg = np.asarray(self[n_full * segment_len : n_live]) + except Exception: + return NotImplemented + if dtype.kind == "f": + seg = seg[~np.isnan(seg)] + if not nullable and seg.shape[0] != tail: + return NotImplemented # NaN data: see above + if seg.shape[0]: + seg_val = min(seg) if dtype.kind in "US" else seg.min() + if op == "max": + seg_val = max(seg) if dtype.kind in "US" else seg.max() + vals = np.concatenate([vals, np.asarray([seg_val], dtype=dtype)]) + + if vals.shape[0] == 0: + return NotImplemented # nothing usable → let the scan decide/raise if dtype.kind in "US": return min(vals) if op == "min" else max(vals) return vals.min() if op == "min" else vals.max() @@ -4050,11 +4280,11 @@ def _is_list_column(col: CompiledColumn) -> bool: @staticmethod def _is_varlen_scalar_column(col: CompiledColumn) -> bool: - return isinstance(col.spec, (VLStringSpec, VLBytesSpec, StructSpec, ObjectSpec, Utf8Spec)) + return isinstance(col.spec, (VLStringSpec, VLBytesSpec, StructSpec, ObjectSpec, UTF8Spec)) @staticmethod def _is_utf8_column(col: CompiledColumn) -> bool: - return isinstance(col.spec, Utf8Spec) + return isinstance(col.spec, UTF8Spec) @staticmethod def _is_dictionary_column(col: CompiledColumn) -> bool: @@ -4069,14 +4299,60 @@ def _dict_rank_index_stale(self, name: str, dict_rank_meta: dict) -> bool: """ from blosc2.ctable_indexing import _dict_rank_hash - col = self._root_table._cols.get(name) + root = self._root_table + col = root._cols.get(name) if col is None: return True + # Hashing the whole dictionary costs more than the scan this check is + # meant to let us skip (24 ms for 20k entries), so settle it from the + # value epoch first: unchanged epoch means nothing has been written + # since the index was built, so the ranks cannot have moved. + built_epoch = (root._get_index_catalog().get(name) or {}).get("built_value_epoch") + if built_epoch is not None and root._storage.get_epoch_counters()[0] == built_epoch: + return False dictionary = list(col.dictionary) if len(dictionary) != dict_rank_meta.get("dict_len"): return True return _dict_rank_hash(dictionary) != dict_rank_meta.get("dict_hash") + def _utf8_index_vocab(self, name: str, utf8_rank_meta: dict) -> np.ndarray | None: + """Rank-ordered vocabulary for a utf8 index, cached per table. + + Small relative to the column (one entry per distinct value) and read + once, so a literal→rank lookup costs a ``searchsorted`` rather than a + re-factorization of the column. + """ + cache = self.__dict__.setdefault("_utf8_vocab_cache", {}) + key = (name, utf8_rank_meta.get("n_rows"), utf8_rank_meta.get("nbytes")) + if key in cache: + return cache[key] + inline = utf8_rank_meta.get("vocab") + if inline is not None: + vocab = np.array(inline, dtype=np.str_) if inline else np.empty(0, dtype=np.str_) + else: + path = utf8_rank_meta.get("vocab_path") + if path is None or not os.path.exists(path): + return None + vocab = np.asarray(blosc2.open(path, mode="r")[:]) + cache.clear() # only the current build's vocabulary is ever of interest + cache[key] = vocab + return vocab + + def _utf8_rank_index_stale(self, name: str, utf8_rank_meta: dict) -> bool: + """True if a utf8-rank FULL index no longer matches the live column. + + The index encodes alphabetical ranks frozen at build time, so a value + appended ahead of existing ones invalidates every rank, not just the new + rows'. Checked with O(1) signals — re-deriving the vocabulary would mean + factorizing the column on every query. + """ + col = self._root_table._cols.get(name) + if col is None: + return True + return len(col) != utf8_rank_meta.get("n_rows") or int(col._bytes_used) != utf8_rank_meta.get( + "nbytes" + ) + @staticmethod def _is_ndarray_column(col: CompiledColumn) -> bool: return isinstance(col.spec, NDArraySpec) @@ -4160,7 +4436,7 @@ def _policy_null_value_for_spec(spec: SchemaSpec, policy: NullPolicy): return policy.float_value if isinstance(spec, b2_bool): return policy.bool_value - if isinstance(spec, (string, Utf8Spec)): + if isinstance(spec, (string, UTF8Spec)): return policy.string_value if isinstance(spec, b2_bytes): return policy.bytes_value @@ -4210,7 +4486,7 @@ def _validate_null_value_for_spec(name: str, spec: SchemaSpec, null_value) -> No if null_value != 255: raise ValueError(f"Null sentinel for nullable bool column {name!r} must be 255") return - if isinstance(spec, (string, Utf8Spec)): + if isinstance(spec, (string, UTF8Spec)): if not isinstance(null_value, str): raise TypeError(f"Null sentinel for string column {name!r} must be str") return @@ -6637,7 +6913,7 @@ def _resolve_arrow_columns(self, columns, include_computed: bool = True) -> list def _pa_type_from_spec(pa, spec): if isinstance(spec, DictionarySpec): return pa.dictionary(pa.int32(), pa.string(), ordered=spec.ordered) - if isinstance(spec, Utf8Spec): + if isinstance(spec, UTF8Spec): # Always large_string: 64-bit offsets match the int64 offsets array, # so multi-GB string columns export without int32-offset overflow. return pa.large_string() @@ -7056,7 +7332,7 @@ def _arrow_type_to_spec( # noqa: C901 if _is_arrow_string_type(pa, pa_type): if string_max_length is None: - from blosc2.utf8_array import have_string_dtype + from blosc2._utf8_array import have_string_dtype if not have_string_dtype(): # utf8 columns need numpy.dtypes.StringDType (NumPy >= 2.0). @@ -7127,7 +7403,7 @@ def _compiled_columns_from_arrow( # only binary columns keep the native-None varlen treatment. # On NumPy < 2.0 (no StringDType) utf8 columns are unavailable and # scalar strings keep the historical vlstring treatment instead. - from blosc2.utf8_array import have_string_dtype + from blosc2._utf8_array import have_string_dtype field_is_varlen_scalar = ( not field_is_list @@ -8770,8 +9046,9 @@ def add_column( # noqa: C901 self, name: str, spec: SchemaSpec | dataclasses.Field, + values=None, ) -> None: - """Add a new column filled from the default declared in *spec*. + """Add a new column filled from *values*, or from the default declared in *spec*. Parameters ---------- @@ -8780,16 +9057,30 @@ def add_column( # noqa: C901 spec: A schema descriptor such as ``b2.int64(ge=0)`` or a field descriptor such as ``b2.field(b2.int64(ge=0), default=0)``. - When the table already has live rows, use ``blosc2.field(...)`` - with a default declared so those rows can be backfilled. + When the table already has live rows and no *values* are given, + use ``blosc2.field(...)`` with a default declared so those rows + can be backfilled. + values: + Optional sequence with one entry per **live** row, in row order, + used to fill the new column. This is the supported way to land a + computed result back into the table:: + + res = blosc2.lazyexpr("'x=' + a", {"a": arr}).compute() + t.add_column("out", blosc2.utf8(), values=res[:]) + + A declared default is still honoured for rows appended later, so + *values* and ``blosc2.field(..., default=...)`` can be combined. Raises ------ ValueError If the table is read-only, is a view, the column already exists, - or a non-empty table is given a column with no default declared. + a non-empty table is given a column with neither *values* nor a + default declared, or ``len(values)`` does not match the number of + live rows. TypeError - If a declared default cannot be coerced to *spec*'s dtype. + If a declared default, or *values*, cannot be coerced to *spec*'s + dtype. """ if self._read_only: raise ValueError("Table is read-only (opened with mode='r').") @@ -8804,10 +9095,10 @@ def add_column( # noqa: C901 spec, default, column_config = self._column_spec_default_and_config(spec) n_live = self.nrows - if default is MISSING and n_live > 0: + if values is None and default is MISSING and n_live > 0: raise ValueError( "add_column() requires a default declared as blosc2.field(..., default=...) " - "when the table has live rows." + "or a values= sequence when the table has live rows." ) compiled_col = self._compiled_column_from_spec(name, spec) @@ -8817,9 +9108,22 @@ def add_column( # noqa: C901 validate_column_null_values=False, ) spec = compiled_col.spec + if self._is_list_column(compiled_col): + raise TypeError( + "add_column() does not support list columns; use the constructor with a full schema." + ) + if self._is_dictionary_column(compiled_col): + raise TypeError( + "add_column() does not support dictionary columns; use the constructor with a full schema." + ) + if values is not None: + values = self._add_column_values(name, compiled_col, values, n_live) if self._is_varlen_scalar_column(compiled_col): - # Varlen scalar columns don't use fixed-width NDArray storage. + # Varlen scalar columns don't use fixed-width NDArray storage, but the + # table still indexes them by *physical* position, so a new one has to + # span the physical extent rather than just the live rows -- otherwise + # any table with deleted rows reads past its end. col_storage = self._resolve_column_storage(compiled_col, None, None) new_col = self._storage.create_varlen_scalar_column( name, @@ -8827,13 +9131,19 @@ def add_column( # noqa: C901 cparams=col_storage.get("cparams"), dparams=col_storage.get("dparams"), ) - for _ in range(n_live): - new_col.append(default) + n_phys = self._resolve_last_pos() + filler = self._varlen_filler(spec, default) + if values is None: + new_col.extend([filler] * n_phys) + elif n_live == n_phys: + new_col.extend(values) + else: + padded = [filler] * n_phys + live_pos = np.flatnonzero(self._valid_rows[:n_phys]) + for pos, value in zip(live_pos, values, strict=True): + padded[int(pos)] = value + new_col.extend(padded) new_col.flush() - elif self._is_list_column(compiled_col): - raise TypeError( - "add_column() does not support list columns; use the constructor with a full schema." - ) else: if default is not MISSING: try: @@ -8862,7 +9172,14 @@ def add_column( # noqa: C901 dparams=col_storage.get("dparams"), ) if n_live > 0: - if self._is_ndarray_column(compiled_col): + if values is not None: + # No holes (the common case) means the live rows are the + # leading slots, so a contiguous write beats a scatter. + if n_live == self._resolve_last_pos(): + new_col[:n_live] = values + else: + new_col[np.where(self._valid_rows[:])[0]] = values + elif self._is_ndarray_column(compiled_col): new_col[self._valid_rows] = np.broadcast_to(default_val, (n_live, *spec.item_shape)) else: new_col[self._valid_rows] = default_val @@ -8881,6 +9198,65 @@ def add_column( # noqa: C901 if isinstance(self._storage, FileTableStorage): self._storage.save_schema(self._schema_dict_with_computed()) + @staticmethod + def _varlen_filler(spec, default): + """Value written into the dead slots of a freshly added varlen column. + + Never read back -- the table only ever indexes live positions -- so it + just has to be something the spec accepts. + """ + if default is not MISSING: + return default + null_value = getattr(spec, "null_value", None) + if null_value is not None: + return null_value + if isinstance(spec, VLBytesSpec): + return b"" + if isinstance(spec, (UTF8Spec, VLStringSpec)): + return "" + return None + + def _add_column_values(self, name: str, col: CompiledColumn, values, n_live: int): + """Validate and coerce the ``values=`` argument of :meth:`add_column`. + + Returns a list for varlen scalar columns (which are fed row by row) and + a dtype-coerced ndarray for the fixed-width ones. + + Constraints declared on the spec are checked here, *before* the + ``astype`` below: coercing to a fixed-width dtype truncates a too-long + string to ``max_length`` instead of complaining, so skipping the check + would silently drop characters. + """ + from blosc2.schema_vectorized import validate_column_values + + if self._is_varlen_scalar_column(col): + values = list(values) + if len(values) != n_live: + raise ValueError( + f"add_column() values= for {name!r} requires {n_live} entries " + f"(live rows), got {len(values)}." + ) + validate_column_values(col, values) + return values + + arr = values[:] if isinstance(values, blosc2.NDArray) else np.asarray(values) + if len(arr) != n_live: + raise ValueError( + f"add_column() values= for {name!r} requires {n_live} entries (live rows), got {len(arr)}." + ) + expected = (n_live, *col.spec.item_shape) if self._is_ndarray_column(col) else (n_live,) + if arr.shape != expected: + raise ValueError( + f"add_column() values= for {name!r} must have shape {expected}, got {arr.shape}." + ) + validate_column_values(col, arr) + try: + return arr.astype(col.spec.dtype) + except (ValueError, OverflowError) as exc: + raise TypeError( + f"Cannot coerce values= for {name!r} to dtype {col.spec.dtype!r}: {exc}" + ) from exc + def drop_column(self, name: str) -> None: """Remove a column from the table. @@ -9841,6 +10217,31 @@ def _resolve_dsl_kernel(self, kernel, inputs) -> tuple[Any, list[str]]: self._validate_transformer_dep(d) return kernel, col_deps + def _guard_utf8_kernel_deps(self, col_deps) -> None: + """Refuse a UDF/DSL kernel that reads a utf8 column, naming the column. + + :func:`blosc2.lazyudf` refuses the operand on its own, but only ever + sees the container, so it cannot say *which* column. Called from the + column-registration paths as well, where the kernel would otherwise be + accepted and then fail on every read -- and on ``str(table)`` -- when + the output container is allocated from a ``StringDType`` the NDArray + dtype round-trip cannot parse. It is the utf8 *operand* that does + this, whatever the kernel returns. + """ + from blosc2._utf8_array import utf8_compute_error + + for dep in col_deps: + col = self._schema.columns_by_name.get(dep) + if col is not None and self._is_utf8_column(col): + raise NotImplementedError( + utf8_compute_error( + f"Column {dep!r} is a variable-length utf8 column and cannot be a UDF " + "or DSL kernel operand.", + source=f"t[{dep!r}]", + compute="blosc2.lazyudf(kernel, (fixed,)).compute()[:]", + ) + ) + def _normalize_transformer(self, expr, inputs=None) -> dict: """Resolve *expr* into a transformer descriptor. @@ -9856,6 +10257,7 @@ def _normalize_transformer(self, expr, inputs=None) -> dict: """ if isinstance(expr, blosc2.DSLKernel): kernel, col_deps = self._resolve_dsl_kernel(expr, inputs) + self._guard_utf8_kernel_deps(col_deps) return {"kind": "dsl", "kernel": kernel, "col_deps": col_deps} # Resolve a callable once (a lambda may return a LazyExpr or a LazyUDF). obj = expr(self._cols) if (callable(expr) and not isinstance(expr, blosc2.LazyExpr)) else expr @@ -9867,10 +10269,12 @@ def _normalize_transformer(self, expr, inputs=None) -> dict: kernel = obj.func if kernel.dsl_error is not None: raise blosc2.DSLSyntaxError(f"Invalid DSL kernel: {kernel.dsl_error}") + col_deps = self._dsl_deps_from_lazyudf(obj) + self._guard_utf8_kernel_deps(col_deps) return { "kind": "dsl", "kernel": kernel, - "col_deps": self._dsl_deps_from_lazyudf(obj), + "col_deps": col_deps, "jit_backend": obj.kwargs.get("jit_backend"), } lazy, col_deps = self._normalize_expression_transformer(obj) @@ -10109,6 +10513,9 @@ def apply( # inputs add_computed_column()/add_generated_column() pass to # lazyudf() for DSL/UDF columns -- so the live-row mask is applied # once, here, to the result rather than to every operand. + # lazyudf() refuses a utf8 operand too, but only sees the container, so + # settle it here where the column name is still known. + self._guard_utf8_kernel_deps(names) operands = tuple(self._cols[self._logical_to_physical_name(name)] for name in names) result = blosc2.lazyudf(func, operands, dtype=dtype, jit=jit).compute() return result[self._valid_rows] @@ -11137,11 +11544,14 @@ def _sorted_positions_from_full_index(self, name: str, ascending: bool) -> np.nd descriptor = None else: dict_rank_meta = descriptor.get("full", {}).get("dict_rank") + utf8_rank_meta = descriptor.get("full", {}).get("utf8_rank") if dict_rank_meta is not None: if self._dict_rank_index_stale(name, dict_rank_meta): descriptor = None # ranks no longer match dictionary → lexsort else: is_dict_rank = True + elif utf8_rank_meta is not None and self._utf8_rank_index_stale(name, utf8_rank_meta): + descriptor = None # ranks no longer match the column → lexsort elif name in root._computed_cols: cc = root._computed_cols[name] for _lookup_key, candidate in catalog.items(): @@ -11249,13 +11659,22 @@ def _build_lex_keys( ascending: list[bool], live_pos: np.ndarray, n: int, + gathered: dict[str, np.ndarray] | None = None, ) -> list[np.ndarray]: """Build the key list for np.lexsort (innermost = last = primary key). For nullable columns a null-indicator key (0=non-null, 1=null) is inserted immediately after the value key, making it more significant. This ensures nulls sort last regardless of ascending/descending order. + + *gathered* lets a caller that has already read the columns at + *live_pos* hand them over instead of paying for a second gather; a + dictionary column is expected there as its **codes**, which is what it + is cheap to gather. """ + from blosc2.ctable_indexing import _dict_code_to_rank + + gathered = gathered or {} lex_keys = [] for name, asc in zip(reversed(cols), reversed(ascending), strict=True): cc = self._computed_cols.get(name) @@ -11267,14 +11686,24 @@ def _build_lex_keys( else: is_dict_col = col_info is not None and self._is_dictionary_column(col_info) if is_dict_col: - # Sort dictionary columns by decoded string values. - decoded = self._cols[name][live_pos] - raw = np.array(decoded, dtype=object) - # Replace None with placeholder so lexsort never compares None. - # Null indicator key (below) already places nulls last. - raw[raw == None] = "" # noqa: E711 + # Sort a dictionary column by the alphabetical rank of each + # row's code -- the same trick the FULL index plays. Ranks + # order exactly as the decoded values do, and sorting int32 + # skips both the decode and lexsort's string comparisons. + dict_col = self._cols[name] + dictionary = list(dict_col.dictionary) + code_to_rank = _dict_code_to_rank(dictionary) + raw_codes = gathered[name] if name in gathered else dict_col.codes[live_pos] + codes = np.asarray(raw_codes, dtype=np.int32) + # The null code is reserved (-1), not a dictionary entry, so + # it cannot be looked up; nulls take the largest rank. The + # null indicator key below is what actually places them. + is_null = codes == col_info.spec.null_code + raw = np.empty(len(codes), dtype=np.int32) + raw[~is_null] = code_to_rank[codes[~is_null]] + raw[is_null] = np.int32(len(dictionary)) else: - raw = self._cols[name][live_pos] + raw = gathered[name] if name in gathered else self._cols[name][live_pos] nv = getattr(col_info.spec, "null_value", None) if col_info else None # Value key @@ -11293,10 +11722,7 @@ def _build_lex_keys( # Null indicator key — more significant than the value key above, # so nulls always sort last (0 before 1 → non-null before null). if is_dict_col and col_info.spec.nullable: - null_code = col_info.spec.null_code - codes_at_pos = np.asarray(self._cols[name].codes[live_pos], dtype=np.int32) - null_ind = (codes_at_pos == null_code).astype(np.intp) - lex_keys.append(null_ind) + lex_keys.append(is_null.astype(np.intp)) elif nv is not None: if isinstance(nv, float) and np.isnan(nv): null_ind = np.isnan(raw).astype(np.intp) @@ -11454,12 +11880,18 @@ def _sorted_slice_positions(self, name: str, ascending: bool, key: slice) -> np. col_info = self._schema.columns_by_name.get(name) null_value = getattr(col_info.spec, "null_value", None) if col_info is not None else None - # Dict-rank index: use null_rank (int32) as sentinel for null-block location. + # Rank index: the sidecar holds int32 ranks, so the null block is located + # by null_rank, not by the column's own sentinel (a string, for utf8). dict_rank = full.get("dict_rank") if dict_rank is not None: if self._dict_rank_index_stale(name, dict_rank): return None # ranks no longer match dictionary → lexsort null_value = dict_rank["null_rank"] + utf8_rank = full.get("utf8_rank") + if utf8_rank is not None: + if self._utf8_rank_index_stale(name, utf8_rank): + return None # ranks no longer match the column → lexsort + null_value = utf8_rank["null_rank"] if null_value is not None else None # Numeric / NaN / string sentinels keep the null rows in one contiguous block # once sorted; other non-numeric sentinels (e.g. object) would need a # different locator. @@ -11564,43 +11996,9 @@ def _sorted_small_copy_from_live_positions( else: gathered[col.name] = arr[live_pos] - lex_keys = [] - for name, asc in zip(reversed(cols), reversed(ascending), strict=True): - col_info = self._schema.columns_by_name.get(name) - is_dict_col = col_info is not None and self._is_dictionary_column(col_info) - if is_dict_col: - raw = np.array(self._cols[name][live_pos], dtype=object) - # Replace None with placeholder so lexsort never compares None. - raw[raw == None] = "" # noqa: E711 - else: - raw = gathered[name] - - if not asc: - if raw.dtype.kind in "USO": - rank = np.argsort(np.argsort(raw, kind="stable"), kind="stable") - lex_keys.append((n - 1 - rank).astype(np.intp)) - elif np.issubdtype(raw.dtype, np.unsignedinteger): - lex_keys.append(-raw.astype(np.int64)) - else: - lex_keys.append(-raw) - else: - lex_keys.append(raw) - - if is_dict_col and col_info.spec.nullable: - null_code = col_info.spec.null_code - codes_at_pos = np.asarray(self._cols[name].codes[live_pos], dtype=np.int32) - null_ind = (codes_at_pos == null_code).astype(np.intp) - lex_keys.append(null_ind) - else: - nv = getattr(col_info.spec, "null_value", None) if col_info else None - if nv is not None: - if isinstance(nv, float) and np.isnan(nv): - null_ind = np.isnan(raw).astype(np.intp) - else: - null_ind = (raw == nv).astype(np.intp) - lex_keys.append(null_ind) - - order = np.lexsort(lex_keys) + # The gather above already read every column at live_pos, dictionary + # columns as codes -- exactly what the key builder wants. + order = np.lexsort(self._build_lex_keys(cols, ascending, live_pos, n, gathered)) result = self._empty_copy(capacity=n) for col in self._schema.columns: col_name = col.name @@ -12167,7 +12565,7 @@ def _dtype_info_label(dtype: np.dtype | None, spec: SchemaSpec | None = None) -> if isinstance(spec, DictionarySpec): ordered_tag = ", ordered" if spec.ordered else "" return f"dictionary[str{ordered_tag}]" - if isinstance(spec, Utf8Spec): + if isinstance(spec, UTF8Spec): return "utf8" if isinstance(spec, VLStringSpec): return "vlstring" @@ -12587,6 +12985,13 @@ def _rewrite_dictionary_predicates( def eq_repl(match: re.Match, _dc=dc, _name=name) -> str: value = ast.literal_eval(match.group(2)) + # Deliberately *not* served from the rank index, unlike the + # operator form: rewriting to a code comparison keeps this a + # single fused numeric expression, and substituting a + # precomputed mask instead measured consistently slower + # (22.9 ms -> 28.7 ms at 1M rows) even though the mask itself + # costs only 4.8 ms. Accelerating this form needs the planner + # to consume index positions, not a mask. try: code = int(_dc.value_to_code(value)) except KeyError: @@ -12626,6 +13031,25 @@ def in_repl(match: re.Match, _dc=dc, _name=name) -> str: rewritten = new_expr return rewritten, new_operands + @staticmethod + def _alias_dotted(expr: str, names: list[str], prefix: str) -> tuple[str, dict[str, str]]: + """Replace each dotted name in *expr* with ``{prefix}{i}``. + + Returns the rewritten expression and the ``alias -> original name`` + map, holding only the names that actually occurred in *expr*. + """ + rewritten = expr + aliases = {} + # Longest names first so trip.begin.lon is rewritten before trip.begin. + for i, name in enumerate(sorted((n for n in names if "." in n), key=len, reverse=True)): + alias = f"{prefix}{i}" + pattern = rf"(? tuple[str, dict[str, blosc2.NDArray | blosc2.LazyExpr]]: @@ -12635,28 +13059,27 @@ def _rewrite_nested_expression( columns are naturally addressed as dotted paths (e.g. ``trip.begin.lon``). This maps them to temporary aliases and returns rewritten expression and operand mapping. + + Only names present in *operands* are rewritten; the flavours excluded + from the operand namespace (utf8, dictionary, ...) are aliased by + whichever driver evaluates them -- see :meth:`_lazyexpr_over_cols`. """ - dotted = [name for name in operands if "." in name] - if not dotted: + rewritten, aliases = self._alias_dotted(expr, list(operands), "__nf") + if not aliases: return expr, operands - rewritten = expr new_operands = dict(operands) - # Longest names first so trip.begin.lon is rewritten before trip.begin. - for i, name in enumerate(sorted(dotted, key=len, reverse=True)): - alias = f"__nf{i}" - pattern = rf"(? bool: return re.search(rf"(? None: + def _guard_scalar_expression(self, expr: str, *, allow_utf8: bool = False) -> None: + from blosc2._utf8_array import utf8_compute_error + for name, meta in self._root_table._materialized_cols.items(): if meta.get("stale", False) and self._expression_references_name(expr, name): raise ValueError( @@ -12671,9 +13094,15 @@ def _guard_scalar_expression(self, expr: str) -> None: "support scalar columns. Use an element projection or a row-wise reduction first." ) if self._is_utf8_column(col) and self._expression_references_name(expr, col.name): + if allow_utf8: + continue raise NotImplementedError( - f"Column {col.name!r} is a variable-length utf8 column; " - "string expressions on utf8 columns are not supported yet." + utf8_compute_error( + f"Column {col.name!r} is a variable-length utf8 column; string expressions " + "that reference one are not supported here.", + source=f"t[{col.name!r}]", + compute=f"blosc2.lazyexpr({expr!r}, {{{col.name!r}: fixed}}).compute()[:]", + ) ) if self._is_varlen_scalar_column(col) and self._expression_references_name(expr, col.name): raise NotImplementedError( @@ -12681,8 +13110,160 @@ def _guard_scalar_expression(self, expr: str) -> None: "lazy expressions are not supported yet." ) - def _guard_varlen_scalar_expression(self, expr: str) -> None: - self._guard_scalar_expression(expr) + def _guard_varlen_scalar_expression(self, expr: str, *, allow_utf8: bool = False) -> None: + self._guard_scalar_expression(expr, allow_utf8=allow_utf8) + + # ------------------------------------------------------------------ + # utf8 string expressions: span-loop driver + # ------------------------------------------------------------------ + + #: Rows materialized per span by the utf8 expression driver. Matches the + #: chunk size of :meth:`Column._utf8_chunked_bool`. + _UTF8_EXPR_SPAN = 65536 + + #: Byte ceiling for one span's fixed-width `` list[str]: + """utf8 column names referenced by *expr*, in schema order. + + Call this on the *original* expression, before the dictionary/nested + rewrites: a nested utf8 leaf is aliased away by + :meth:`_rewrite_nested_expression` and would no longer be findable. + """ + return [ + col.name + for col in self._schema.columns + if self._is_utf8_column(col) and self._expression_references_name(expr, col.name) + ] + + #: Comparisons a utf8 column supports against a string literal, longest + #: spelling first so ``<=`` is not matched as ``<``. + _UTF8_CMP_OPS: ClassVar[dict] = { + "==": np.equal, + "!=": np.not_equal, + "<=": np.less_equal, + ">=": np.greater_equal, + "<": np.less, + ">": np.greater, + } + #: Mirrored operator for a reversed comparison (``'x' < name``). + _UTF8_CMP_MIRROR: ClassVar[dict] = {"==": "==", "!=": "!=", "<=": ">=", ">=": "<=", "<": ">", ">": "<"} + + def _rewrite_utf8_predicates( + self, expr: str, operands: dict, utf8_names: list[str], aliases: dict[str, str] | None = None + ) -> tuple[str, dict, list[str]]: + """Replace ``utf8col 'literal'`` terms with precomputed masks. + + A scalar comparison is answered by a raw-byte scan of the offsets/data + pair (:meth:`Column._utf8_scalar_mask`) with no decode at all, which is + several times cheaper than the span driver's decode -> `` + miniexpr round trip and is exactly what the operator form + ``t[t.name == "x"]`` already does. Substituting the mask as a boolean + operand keeps the rest of the expression (``&``/``|``, numeric terms) a + single native expression. + + Returns the rewritten expression, the extended operands, and the utf8 + names still referenced -- a name drops out only when *every* one of its + occurrences was rewritten, so anything else (``startswith(name, 'x')``, + ``upper(name)``) still routes to the span driver. + + Names are as they appear in *expr*; a nested leaf appears under an + alias and resolves to its column through *aliases*. + """ + rewritten = expr + new_operands = dict(operands) + remaining = [] + col_of = (aliases or {}).get + ops = "|".join(re.escape(o) for o in self._UTF8_CMP_OPS) + for i, name in enumerate(utf8_names): + column = self[col_of(name, name)] + counter = itertools.count() + + def repl(match: re.Match, _col=column, _i=i, _c=counter, reverse=False) -> str: + op = match.group(1) + literal = ast.literal_eval(match.group(2) if not reverse else match.group(1)) + if reverse: + op = self._UTF8_CMP_MIRROR[match.group(2)] + alias = f"__u8{_i}_{next(_c)}" + new_operands[alias] = blosc2.asarray(_col._utf8_scalar_mask(self._UTF8_CMP_OPS[op], literal)) + return alias + + escaped = r"(? bool: col = self[name] @@ -12801,11 +13382,12 @@ def where( # noqa: C901 if isinstance(expr_result, ColExpr): expr_result = expr_result._bind(self) if isinstance(expr_result, str): - self._guard_varlen_scalar_expression(expr_result) + self._guard_varlen_scalar_expression(expr_result, allow_utf8=True) + utf8_names = self._utf8_names_in(expr_result) operands = self._where_expression_operands(expr_result) expr_result, operands = self._rewrite_dictionary_predicates(expr_result, operands) expr_result, operands = self._rewrite_nested_expression(expr_result, operands) - expr_result = blosc2.lazyexpr(expr_result, operands) + expr_result = self._lazyexpr_over_cols(expr_result, operands, utf8_names) if isinstance(expr_result, np.ndarray) and expr_result.dtype == np.bool_: expr_result = blosc2.asarray(expr_result) if isinstance(expr_result, Column): diff --git a/src/blosc2/ctable_indexing.py b/src/blosc2/ctable_indexing.py index 7ec5f7a37..800f17dc2 100644 --- a/src/blosc2/ctable_indexing.py +++ b/src/blosc2/ctable_indexing.py @@ -13,6 +13,7 @@ import ast import contextlib import os +import pathlib from typing import TYPE_CHECKING, Any import numpy as np @@ -25,7 +26,7 @@ NDArraySpec, ObjectSpec, StructSpec, - Utf8Spec, + UTF8Spec, VLBytesSpec, VLStringSpec, ) @@ -57,6 +58,22 @@ def __init__(self): self.vlmeta = _FakeVlMeta() +def _dict_code_to_rank(dictionary) -> np.ndarray: + """``code -> alphabetical rank`` lookup for a dictionary column. + + Sorting by rank is sorting by decoded value, which is what lets an int32 + array stand in for the strings — in the FULL index (:class:`_DictRankWrapper`) + and in ``CTable._build_lex_keys``. The reserved null code is *not* a + dictionary entry, so callers assign nulls a rank of their own (``len``, + the largest, so nulls sort last). + """ + n_entries = len(dictionary) + order = np.argsort(dictionary, kind="stable") + code_to_rank = np.empty(n_entries, dtype=np.int32) + code_to_rank[order] = np.arange(n_entries, dtype=np.int32) + return code_to_rank + + def _dict_rank_hash(dictionary) -> str: """Stable hash of a dictionary's entries (code position + value). @@ -74,6 +91,79 @@ def _dict_rank_hash(dictionary) -> str: return h.hexdigest() +#: Rows factorized per pass when building a utf8 rank index. Bounds the +#: transient code buffer without changing the result. +_UTF8_RANK_SPAN = 1 << 20 + + +def _utf8_rank_arrays(col, n_phys: int, null_value: str | None): + """Alphabetical rank per row for a utf8 column, plus its staleness metadata. + + Sorting by rank is sorting by decoded string, so an ``int32`` rank column + can drive the whole numeric index machinery unchanged — the same trick + :class:`_DictRankWrapper` plays for dictionary columns. Unlike a dictionary + there is no stored code array, so the column is factorized here; the + factorizer hashes raw bytes and only ever decodes the distinct values. + + Null rows carry a sentinel *string*, so the sentinel is just another + vocabulary entry; it is given the largest rank so nulls sort last, matching + both the dictionary index and ``_build_lex_keys``. + """ + fact = col.factorizer() + codes = np.empty(n_phys, dtype=np.int64) + for start in range(0, n_phys, _UTF8_RANK_SPAN): + stop = min(start + _UTF8_RANK_SPAN, n_phys) + codes[start:stop] = fact.codes_for_span(start, stop) + uniques = fact.uniques() + n_entries = len(uniques) + + is_null = uniques == null_value if null_value is not None else np.zeros(n_entries, dtype=bool) + non_null = np.flatnonzero(~is_null) + order = non_null[np.argsort(uniques[non_null], kind="stable")] + code_to_rank = np.empty(max(n_entries, 1), dtype=np.int32) + code_to_rank[order] = np.arange(len(order), dtype=np.int32) + null_rank = np.int32(len(order)) + if is_null.any(): + code_to_rank[is_null] = null_rank + + # Rank order == alphabetical order, so this doubles as the lookup table that + # turns a query literal into a rank (np.searchsorted) without touching data. + sorted_vocab = uniques[order] + ranks = code_to_rank[codes] if n_entries else np.zeros(n_phys, dtype=np.int32) + # Staleness signals must be O(1) to check: re-deriving the vocabulary would + # mean factorizing the column again on every query. Any write already marks + # every index stale, so these only have to catch a rebuilt-but-changed + # column, for which row count plus blob size is enough. + meta = { + "null_rank": int(null_rank), + "vocab_len": int(n_entries), + "n_rows": int(n_phys), + "nbytes": int(col._bytes_used), + } + return ranks.astype(np.int32, copy=False), meta, sorted_vocab + + +def _persist_utf8_vocab(full: dict, meta: dict, sorted_vocab: np.ndarray) -> None: + """Store the rank-ordered vocabulary so a query literal can be turned into a rank. + + Written beside the index's own sidecars when the table is persistent, and + inlined into the descriptor otherwise — the in-memory index path is for + small tables by construction. Without it a literal→rank lookup would mean + factorizing the column again on every query. + """ + if len(sorted_vocab) == 0: + meta["vocab"] = [] + return + values_path = full.get("values_path") + if values_path is None: # in-memory index + meta["vocab"] = sorted_vocab.tolist() + return + width = max(len(v) for v in sorted_vocab) + vocab_path = str(pathlib.Path(values_path).with_suffix("")) + ".utf8_vocab.b2nd" + blosc2.asarray(sorted_vocab.astype(f" rank @@ -99,15 +189,17 @@ def __init__( self._null_code = null_code self._nullable = nullable self.dtype = np.dtype(np.int32) - # The codes array carries capacity padding beyond the live rows; expose only - # the live range so the index sidecars match n_rows (no padding → the - # zero-permutation window read engages instead of falling back). - self.shape = (n_live,) + # The codes array carries capacity padding beyond the written rows; expose + # only the physical extent so the index sidecars match n_rows (no padding → + # the zero-permutation window read engages instead of falling back). It has + # to be the physical extent and not the live count: tombstoned rows keep + # their positions, so live rows can sit past the live count. + self.shape = (n_phys,) self.ndim = 1 - chunk0 = codes.chunks[0] if codes.chunks else n_live - block0 = codes.blocks[0] if codes.blocks else n_live - self.chunks = (min(chunk0, n_live),) - self.blocks = (min(block0, n_live),) + chunk0 = codes.chunks[0] if codes.chunks else n_phys + block0 = codes.blocks[0] if codes.blocks else n_phys + self.chunks = (min(chunk0, n_phys),) + self.blocks = (min(block0, n_phys),) def __getitem__(self, key): codes_slice = np.asarray(self._codes[key], dtype=np.int32) @@ -588,7 +680,7 @@ def create_index( # noqa: C901 field: str | None = None, expression: str | None = None, operands: dict | None = None, - kind: blosc2.IndexKind = blosc2.IndexKind.BUCKET, + kind: blosc2.IndexKind | None = None, optlevel: int = 5, name: str | None = None, build: str = "auto", @@ -636,6 +728,12 @@ def create_index( # noqa: C901 lightest kind; it may still skip segments for broad range queries but cannot accelerate ``sort_by``. + When *kind* is omitted it defaults to ``BUCKET``, except on ``utf8()`` + and ``dictionary()`` columns, which are indexed by alphabetical rank and + only ever consulted through a ``FULL`` index — there the default is + ``FULL``, and asking for any other kind raises ``ValueError`` rather + than building an index that nothing would use. + .. rubric:: SUMMARY granularity For ``kind=SUMMARY``, ``granularity`` controls the segment size of the @@ -677,6 +775,22 @@ def create_index( # noqa: C901 if kwargs: raise TypeError(f"unexpected keyword argument(s): {', '.join(sorted(kwargs))}") + # Rank-indexed flavours are only ever consulted through a FULL index, so + # the BUCKET default would hand them an index nothing can use. Default + # per flavour, and remember whether the caller chose the kind: an + # explicit non-FULL request is an error rather than a silent promotion. + explicit_kind = kind is not None + if kind is None: + spec = None + if col_name is not None: + col_info = self._schema.columns_by_name.get(col_name) + spec = col_info.spec if col_info is not None else None + kind = ( + blosc2.IndexKind.FULL + if isinstance(spec, (UTF8Spec, DictionarySpec)) + else blosc2.IndexKind.BUCKET + ) + kind_str = _normalize_index_kind(kind) build_str = _normalize_build_mode(build) method_str = _normalize_full_build_method(method) if kind_str == "full" else None @@ -750,12 +864,6 @@ def create_index( # noqa: C901 ) if isinstance(self._schema.columns_by_name[col_name].spec, ListSpec): raise ValueError(f"Cannot create an index on list column {col_name!r} in V1.") - if isinstance(self._schema.columns_by_name[col_name].spec, Utf8Spec): - raise NotImplementedError( - f"Cannot create an index on variable-length utf8 column {col_name!r}: " - "indexing for utf8 columns is not supported yet. " - "Use a fixed-width string(max_length=N) column if you need an index." - ) if isinstance( self._schema.columns_by_name[col_name].spec, (VLStringSpec, VLBytesSpec, StructSpec, ObjectSpec) ): @@ -763,24 +871,54 @@ def create_index( # noqa: C901 f"Cannot create an index on variable-length scalar column {col_name!r}: " "indexing for vlstring/vlbytes/struct/object columns is not supported yet." ) + # Rank-indexed flavours (utf8, dictionary) are queried through their own + # literal->rank lookup rather than through plan_query, and both that path + # and the ordering path require kind="full". The other kinds build over + # the ranks without error and are then never consulted, so refuse them + # here rather than charge for an index nothing can use. + rank_spec = self._schema.columns_by_name[col_name].spec + if explicit_kind and kind_str != "full" and isinstance(rank_spec, (UTF8Spec, DictionarySpec)): + flavour = "utf8" if isinstance(rank_spec, UTF8Spec) else "dictionary" + raise ValueError( + f"Column {col_name!r} is a {flavour} column, which is indexed by alphabetical rank; " + f"only kind='full' consults that index, so kind={kind_str!r} would build but never " + "be used. Use kind='full'." + ) + # utf8 columns: index the alphabetical rank of each row's value. There is + # no stored code array to wrap lazily, so the ranks are materialized here + # (int32, 4 B/row) and handed to the builder as an ordinary array. + is_utf8 = isinstance(self._schema.columns_by_name[col_name].spec, UTF8Spec) + utf8_rank_meta = None + if is_utf8: + # Span the physical extent, not the live count: delete() tombstones in + # place, so live rows sit past _n_rows. A utf8 column carries no + # capacity padding (__len__ is persisted + pending), and this is the + # same length _utf8_rank_index_stale() compares the meta against. + n_phys = len(col_arr) + ranks_arr, utf8_rank_meta, utf8_vocab = _utf8_rank_arrays( + col_arr, n_phys, self[col_name].null_value + ) + col_arr = blosc2.asarray(ranks_arr) + # Dictionary columns: index by alphabetical rank instead of insertion-order codes. is_dictionary = isinstance(self._schema.columns_by_name[col_name].spec, DictionarySpec) dict_rank_meta = None if is_dictionary: dict_col = col_arr - n_live = self._n_rows if self._n_rows is not None else len(self._valid_rows) + # Physical extent again, but len(dict_col) is the slot *capacity*, so + # take the live-data watermark instead: it covers every live row while + # still excluding the trailing padding. + n_phys = self._resolve_last_pos() dictionary = list(dict_col.dictionary) n_entries = len(dictionary) - order = np.argsort(dictionary, kind="stable") - code_to_rank = np.empty(n_entries, dtype=np.int32) - code_to_rank[order] = np.arange(n_entries, dtype=np.int32) + code_to_rank = _dict_code_to_rank(dictionary) null_code = dict_col.spec.null_code null_rank = np.int32(n_entries) # Hash for staleness detection. dict_hash = _dict_rank_hash(dictionary) dict_rank_meta = {"null_rank": int(null_rank), "dict_hash": dict_hash, "dict_len": n_entries} col_arr = _DictRankWrapper( - dict_col.codes, code_to_rank, null_rank, null_code, dict_col.spec.nullable, n_live + dict_col.codes, code_to_rank, null_rank, null_code, dict_col.spec.nullable, n_phys ) is_persistent = self._storage.index_anchor_path(col_name) is not None @@ -802,7 +940,7 @@ def create_index( # noqa: C901 else: # In-memory path: materialise ranks as a proper NDArray (small tables only). if is_dictionary: - codes = np.asarray(dict_col.codes[:n_live], dtype=np.int32) + codes = np.asarray(dict_col.codes[:n_phys], dtype=np.int32) ranks_arr = code_to_rank[codes] if dict_col.spec.nullable: ranks_arr[codes == null_code] = null_rank @@ -823,11 +961,15 @@ def create_index( # noqa: C901 ) store = _IN_MEMORY_INDEXES[id(col_arr)] descriptor = _copy_descriptor(store["indexes"]["__self__"]) - if dict_rank_meta is not None: + if dict_rank_meta is not None or utf8_rank_meta is not None: full = descriptor.setdefault("full", {}) if full is None: full = descriptor["full"] = {} - full["dict_rank"] = dict_rank_meta + if dict_rank_meta is not None: + full["dict_rank"] = dict_rank_meta + else: + _persist_utf8_vocab(full, utf8_rank_meta, utf8_vocab) + full["utf8_rank"] = utf8_rank_meta value_epoch, _ = self._storage.get_epoch_counters() descriptor["built_value_epoch"] = value_epoch diff --git a/src/blosc2/ctable_storage.py b/src/blosc2/ctable_storage.py index 96359498c..0d01b4d44 100644 --- a/src/blosc2/ctable_storage.py +++ b/src/blosc2/ctable_storage.py @@ -28,6 +28,7 @@ import numpy as np import blosc2 +from blosc2._utf8_array import UTF8Array, _new_backend_arrays from blosc2.batch_array import BatchArray from blosc2.dictionary_column import DictionaryColumn from blosc2.list_array import ListArray @@ -37,9 +38,8 @@ _ScalarVarLenArray, _validate_role_metadata, ) -from blosc2.schema import Utf8Spec +from blosc2.schema import UTF8Spec from blosc2.schunk import process_opened_object -from blosc2.utf8_array import Utf8Array, _new_backend_arrays if TYPE_CHECKING: from blosc2.schema import ListSpec @@ -251,9 +251,9 @@ def open_list_column(self, name): raise RuntimeError("In-memory tables have no on-disk representation to open.") def create_varlen_scalar_column(self, name, *, spec, cparams=None, dparams=None): - if isinstance(spec, Utf8Spec): + if isinstance(spec, UTF8Spec): offsets, data = _new_backend_arrays(cparams, dparams) - return Utf8Array(spec, offsets, data) + return UTF8Array(spec, offsets, data) return _ScalarVarLenArray(spec) def open_varlen_scalar_column(self, name, spec): @@ -484,10 +484,10 @@ def open_list_column(self, name: str) -> ListArray: return self._estore[self._col_key(name)] def open_varlen_scalar_column(self, name: str, spec) -> _ScalarVarLenArray: - if isinstance(spec, Utf8Spec): + if isinstance(spec, UTF8Spec): offsets = self._estore[self._col_key(name)] data = self._estore[self._col_key(name) + _UTF8_DATA_SUFFIX] - return Utf8Array(spec, offsets, data) + return UTF8Array(spec, offsets, data) backend = self._estore[self._col_key(name)] return _ScalarVarLenArray(spec, backend) @@ -721,24 +721,24 @@ def open_list_column(self, name: str) -> ListArray: return blosc2.open(self._list_col_path(name), mode=self._mode) def create_varlen_scalar_column(self, name, *, spec, cparams=None, dparams=None) -> _ScalarVarLenArray: - if isinstance(spec, Utf8Spec): + if isinstance(spec, UTF8Spec): offsets, data = _new_backend_arrays(cparams, dparams) store = self._open_store() key = self._col_key(name) data_key = key + _UTF8_DATA_SUFFIX store[key] = offsets store[data_key] = data - return Utf8Array(spec, store[key], store[data_key]) + return UTF8Array(spec, store[key], store[data_key]) urlpath = self._list_col_path(name) backend = _make_persistent_backend(spec, urlpath, "w", cparams=cparams, dparams=dparams) return _ScalarVarLenArray(spec, backend) def open_varlen_scalar_column(self, name: str, spec) -> _ScalarVarLenArray: - if isinstance(spec, Utf8Spec): + if isinstance(spec, UTF8Spec): store = self._open_store() offsets = store[self._col_key(name)] data = store[self._col_key(name) + _UTF8_DATA_SUFFIX] - return Utf8Array(spec, offsets, data) + return UTF8Array(spec, offsets, data) store = self._open_store() path = self._list_col_path(name) if store.is_zip_store and self._mode == "r": @@ -1301,7 +1301,7 @@ def create_varlen_scalar_column( cparams=None, dparams=None, ) -> _ScalarVarLenArray: - if isinstance(spec, Utf8Spec): + if isinstance(spec, UTF8Spec): logical_key = self._col_logical_key(name) offsets_path = self._dest_path(logical_key, ".b2nd") data_path = self._dest_path(logical_key + _UTF8_DATA_SUFFIX, ".b2nd") @@ -1316,17 +1316,17 @@ def create_varlen_scalar_column( rel_path = os.path.relpath(dest_path, self._working_dir()).replace(os.sep, "/") self._store.map_tree[self._table_key(logical)] = rel_path self._store._modified = True - return Utf8Array(spec, offsets, data) + return UTF8Array(spec, offsets, data) urlpath = self._list_col_path(name) os.makedirs(os.path.dirname(urlpath), exist_ok=True) return _make_persistent_backend(spec, urlpath, "w", cparams=cparams, dparams=dparams) def open_varlen_scalar_column(self, name: str, spec) -> _ScalarVarLenArray: - if isinstance(spec, Utf8Spec): + if isinstance(spec, UTF8Spec): logical_key = self._col_logical_key(name) offsets = self._open_leaf(logical_key) data = self._open_leaf(logical_key + _UTF8_DATA_SUFFIX) - return Utf8Array(spec, offsets, data) + return UTF8Array(spec, offsets, data) if self._store.is_zip_store and self._mode == "r": rel = self._table_key(self._col_logical_key(name)).lstrip("/") + ".b2b" if rel not in self._store.offsets: diff --git a/src/blosc2/dictionary_column.py b/src/blosc2/dictionary_column.py index 537978337..dc57d19eb 100644 --- a/src/blosc2/dictionary_column.py +++ b/src/blosc2/dictionary_column.py @@ -57,24 +57,32 @@ def __init__(self, spec: DictionarySpec, codes, dict_store: _ScalarVarLenArray) self._dict_store = dict_store # _ScalarVarLenArray of vlstring (unique values) # Cache: str → int32 code. Built lazily from dict_store on first access. self._value_to_code: dict[str, int] | None = None + # Reverse cache: code → str, same lazy build. Indexing dict_store per + # code decompresses a whole msgpack batch every time, so decoding N rows + # that way costs O(N) batch decompressions instead of one dictionary read. + self._code_to_value: list[str] | None = None # ------------------------------------------------------------------ # Cache management # ------------------------------------------------------------------ def _ensure_cache(self) -> None: - """Build the value→code mapping from the persisted dict_store.""" + """Build the value→code and code→value mappings from the persisted dict_store.""" if self._value_to_code is not None: return self._dict_store.flush() cache: dict[str, int] = {} + values: list[str] = [] for code, value in enumerate(self._dict_store): + values.append(value) if value is not None: cache[value] = code self._value_to_code = cache + self._code_to_value = values def _invalidate_cache(self) -> None: self._value_to_code = None + self._code_to_value = None # ------------------------------------------------------------------ # Encoding / decoding @@ -101,6 +109,8 @@ def encode(self, value: str | None) -> int: ) self._dict_store.append(value) self._value_to_code[value] = new_code + assert self._code_to_value is not None + self._code_to_value.append(value) return new_code def decode(self, code: int) -> str | None: @@ -108,7 +118,8 @@ def decode(self, code: int) -> str | None: if code == self._spec.null_code: return None self._ensure_cache() - return self._dict_store[int(code)] + assert self._code_to_value is not None + return self._code_to_value[int(code)] def decode_batch(self, codes) -> list[str | None]: """Decode an array of int32 *codes* to a list of strings (``None`` for null codes). @@ -119,8 +130,8 @@ def decode_batch(self, codes) -> list[str | None]: is dramatically cheaper than looping over :meth:`decode`. """ codes = np.asarray(codes) - self._dict_store.flush() - all_strings = np.asarray(self._dict_store[:]) # D unique values, no nulls + self._ensure_cache() + all_strings = np.asarray(self._code_to_value) # D unique values, no nulls null_code = int(self._spec.null_code) result: list[str | None] = [None] * len(codes) non_null_idx = np.nonzero(codes != null_code)[0] @@ -236,6 +247,43 @@ def __getitem__(self, key) -> str | None | list: return [self.decode(int(codes_arr))] raise TypeError(f"DictionaryColumn indices must be int, slice, or array; got {type(key)!r}") + # Identity hashing is kept: these objects were hashable before __eq__ was + # defined, and an element-wise __eq__ never returns a bool for the hash + # contract to apply to. + __hash__ = object.__hash__ + + def __eq__(self, other): + """Element-wise equality mask, over the same slots ``self[:]`` exposes. + + Without this the comparison fell through to object identity and + ``column == "value"`` was a plain ``False`` — silently wrong. A scalar + string is answered by comparing *codes*, so no row is decoded; null + slots hold ``null_code`` and so never match. + """ + import blosc2 + + if blosc2._disable_overloaded_equal: + return self is other + return self._equality_mask(other, invert=False) + + def __ne__(self, other): + import blosc2 + + if blosc2._disable_overloaded_equal: + return self is not other + return self._equality_mask(other, invert=True) + + def _equality_mask(self, other, *, invert: bool): + codes = np.asarray(self._codes[:], dtype=np.int32) + if isinstance(other, str): + self._ensure_cache() + assert self._value_to_code is not None + code = self._value_to_code.get(other) + mask = np.zeros(len(codes), dtype=bool) if code is None else codes == code + else: + mask = np.asarray(self[:], dtype=object) == other + return ~mask if invert else mask + def __setitem__(self, key, value) -> None: """Encode *value* (str/None or list thereof) and write the code(s).""" if isinstance(key, (int, np.integer)): diff --git a/src/blosc2/dsl_kernel.py b/src/blosc2/dsl_kernel.py index 82c11b373..da7f6ce5a 100644 --- a/src/blosc2/dsl_kernel.py +++ b/src/blosc2/dsl_kernel.py @@ -10,6 +10,7 @@ import ast import contextlib import inspect +import keyword import os import textwrap import tokenize @@ -39,6 +40,176 @@ class DSLSyntaxError(ValueError): } +# Python string methods the DSL grammar exposes as plain functions. The DSL +# parser has no attribute syntax, so `desc.lower()` has to become `lower(desc)` +# before the source reaches miniexpr. +_STRING_METHOD_TO_DSL_FUNC = { + "lower": "lower", + "upper": "upper", + "strip": "strip", + "lstrip": "lstrip", + "rstrip": "rstrip", + "removeprefix": "removeprefix", + "removesuffix": "removesuffix", + "replace": "replace", + "startswith": "startswith", + "endswith": "endswith", +} + + +class _StringSyntaxRewriter(ast.NodeTransformer): + """Make ordinary Python string syntax parseable by the DSL grammar. + + Three rewrites, all shape-preserving: + + ``s.lower()`` -> ``lower(s)`` + ``x in s`` -> ``contains(s, x)`` (``not in`` negated) + ``a, b = s.split(sep, 1)`` + -> ``a = split_part(s, sep, 0)`` + ``b = split_part(s, sep, 1)`` + + The point is that a pandas UDF written in normal Python runs unmodified; + without this, `df.apply(f, axis=1, engine=blosc2.jit)` would require the + user to rewrite their function into function-call form first. + """ + + def __init__(self): + self.rewrote_any = False + + def visit_Call(self, node: ast.Call) -> ast.AST: + self.generic_visit(node) + func = node.func + if isinstance(func, ast.Attribute) and func.attr in _STRING_METHOD_TO_DSL_FUNC: + # `np.foo(...)` is handled by _NumpyAttrCallRewriter; leave it alone. + dsl_name = _STRING_METHOD_TO_DSL_FUNC[func.attr] + new_call = ast.Call( + func=ast.Name(id=dsl_name, ctx=ast.Load()), + args=[func.value, *node.args], + keywords=node.keywords, + ) + self.rewrote_any = True + return ast.copy_location(new_call, node) + return node + + def visit_Compare(self, node: ast.Compare) -> ast.AST: + self.generic_visit(node) + if len(node.ops) != 1 or not isinstance(node.ops[0], (ast.In, ast.NotIn)): + return node + needle, haystack = node.left, node.comparators[0] + call = ast.Call( + func=ast.Name(id="contains", ctx=ast.Load()), + args=[haystack, needle], + keywords=[], + ) + self.rewrote_any = True + if isinstance(node.ops[0], ast.NotIn): + call = ast.UnaryOp(op=ast.Not(), operand=call) + return ast.copy_location(call, node) + + def visit_Assign(self, node: ast.Assign): + self.generic_visit(node) + if len(node.targets) != 1 or not isinstance(node.targets[0], ast.Tuple): + return node + targets = node.targets[0].elts + parts = self._split_call_parts(node.value, len(targets)) + if parts is None: + # Leave it for the validator to reject with a proper message. + return node + self.rewrote_any = True + out = [] + for target, value in zip(targets, parts, strict=True): + assign = ast.Assign(targets=[target], value=value) + out.append(ast.copy_location(assign, node)) + return out + + @staticmethod + def _split_call_parts(value, count): + """Turn ``s.split(sep, 1)`` into ``count`` split_part() calls, or None. + + Only the maxsplit=1 form is supported, which is what tuple unpacking can + consume; a general N-way split has no fixed arity to unpack into. + """ + if count != 2 or not isinstance(value, ast.Call): + return None + func = value.func + name = func.attr if isinstance(func, ast.Attribute) else getattr(func, "id", None) + if name != "split": + return None + if isinstance(func, ast.Attribute): + subject, args = func.value, list(value.args) + else: + if not value.args: + return None + subject, args = value.args[0], list(value.args[1:]) + if len(args) != 2: + return None + sep, maxsplit = args + if not (isinstance(maxsplit, ast.Constant) and maxsplit.value == 1): + return None + return [ + ast.Call( + func=ast.Name(id="split_part", ctx=ast.Load()), + args=[subject, sep, ast.Constant(value=k)], + keywords=[], + ) + for k in range(count) + ] + + +class _RowSubscriptRewriter(ast.NodeTransformer): + """Turn the ``df.apply(f, axis=1)`` row idiom into named DSL parameters. + + ``def f(row): ... row["a"] ...`` becomes ``def f(a): ... a ...``, so a + kernel written the textbook way compiles as a DSL kernel. Without this, + any such kernel containing control flow has nowhere to go: the tracing + route evaluates the `if` on a whole column ("truth value ... is ambiguous") + and the DSL parser rejects the subscript. This is not string-specific -- + numeric row kernels with an `if` hit exactly the same wall. + + Applies only when the function takes one positional parameter and *every* + mention of it is ``param[]``; anything else (positional + indexing, iteration, attribute access) is left for the validator to reject. + + ``columns`` maps the generated parameter names back to the original column + labels, which need not be identifiers. + """ + + def __init__(self, param: str): + self.param = param + self.columns: dict[str, str] = {} + self.bailed = False + self._names: dict[str, str] = {} + + def _param_for(self, label: str) -> str: + if label in self._names: + return self._names[label] + candidate = label if label.isidentifier() and not keyword.iskeyword(label) else "" + if not candidate or candidate in self.columns or candidate == self.param: + candidate = f"_col{len(self._names)}" + self._names[label] = candidate + self.columns[candidate] = label + return candidate + + def visit_Subscript(self, node: ast.Subscript) -> ast.AST: + # Match before descending: visiting `node.value` would see the bare row + # name and trip the bail-out below. + if isinstance(node.value, ast.Name) and node.value.id == self.param: + key = node.slice + if not (isinstance(key, ast.Constant) and isinstance(key.value, str)): + self.bailed = True + return node + return ast.copy_location(ast.Name(id=self._param_for(key.value), ctx=ast.Load()), node) + self.generic_visit(node) + return node + + def visit_Name(self, node: ast.Name) -> ast.AST: + # Any surviving bare mention of the row parameter is a use we cannot + # turn into a column reference. + if node.id == self.param: + self.bailed = True + return node + + class _NumpyAttrCallRewriter(ast.NodeTransformer): """Rewrite `alias.foo(...)` calls to the bare `foo(...)` form the DSL grammar requires, for every *alias* bound to the real NumPy module. Also applies @@ -459,7 +630,7 @@ def _expr(self, node: ast.AST): # noqa: C901 return if isinstance(node, ast.Constant): val = node.value - if isinstance(val, bool | int | float | str): + if isinstance(val, bool | int | float | str | bytes): return self._err(node, "Unsupported constant in DSL expression") if isinstance(node, ast.UnaryOp): @@ -558,6 +729,9 @@ def __init__(self, func): self.dsl_source = None self.input_names = None self.dsl_error = None + # Set by _rewrite_row_subscripts when the kernel takes a row proxy. + self.row_param = None + self.row_columns = None try: dsl_source, input_names = self._extract_dsl(func) except DSLSyntaxError as e: @@ -599,10 +773,20 @@ def _extract_dsl(self, func, validate: bool = True): raise ValueError("No function definition found in sliced DSL source") input_names = self._input_names_from_signature(dsl_func) + dsl_source, dsl_tree, dsl_func, row_columns = self._rewrite_row_subscripts( + dsl_source, dsl_tree, dsl_func, input_names + ) + if row_columns is not None: + self.row_param = input_names[0] + self.row_columns = row_columns + input_names = self._input_names_from_signature(dsl_func) + dsl_source, dsl_tree, dsl_func = self._rewrite_numpy_attr_calls( func, dsl_source, dsl_tree, dsl_func, input_names ) + dsl_source, dsl_tree, dsl_func = self._rewrite_string_syntax(dsl_source, dsl_tree, dsl_func) + if validate: DSLValidator(dsl_source, input_names=input_names).validate(dsl_func) if _PRINT_DSL_KERNEL: @@ -611,6 +795,49 @@ def _extract_dsl(self, func, validate: bool = True): print(dsl_source) return dsl_source, input_names + @staticmethod + def _rewrite_row_subscripts(dsl_source, dsl_tree, dsl_func, input_names): + """Rewrite ``row["col"]`` into named parameters; see _RowSubscriptRewriter. + + Returns ``(source, tree, func, columns)`` with *columns* mapping the new + parameter names to the original column labels, or None as the fourth + element when the rewrite does not apply and nothing was changed. + """ + if len(input_names) != 1: + return dsl_source, dsl_tree, dsl_func, None + rewriter = _RowSubscriptRewriter(input_names[0]) + rewritten = rewriter.visit(ast.parse(dsl_source)) + if rewriter.bailed or not rewriter.columns: + return dsl_source, dsl_tree, dsl_func, None + + new_func = next((n for n in rewritten.body if isinstance(n, ast.FunctionDef)), None) + if new_func is None: + return dsl_source, dsl_tree, dsl_func, None + new_func.args.args = [ast.arg(arg=name) for name in rewriter.columns] + new_func.args.posonlyargs = [] + ast.fix_missing_locations(rewritten) + new_source = ast.unparse(rewritten) + new_tree = ast.parse(new_source) + new_func = next((n for n in new_tree.body if isinstance(n, ast.FunctionDef)), None) + return new_source, new_tree, new_func, dict(rewriter.columns) + + @staticmethod + def _rewrite_string_syntax(dsl_source, dsl_tree, dsl_func): + """Lower Python string syntax to the DSL's function-call grammar. + + No-op (returning the inputs unchanged) when there is nothing to rewrite. + """ + rewriter = _StringSyntaxRewriter() + rewritten = rewriter.visit(ast.parse(dsl_source)) + if not rewriter.rewrote_any: + return dsl_source, dsl_tree, dsl_func + + ast.fix_missing_locations(rewritten) + new_source = ast.unparse(rewritten) + new_tree = ast.parse(new_source) + new_func = next((node for node in new_tree.body if isinstance(node, ast.FunctionDef)), None) + return new_source, new_tree, new_func + @staticmethod def _rewrite_numpy_attr_calls(func, dsl_source, dsl_tree, dsl_func, input_names): """Rewrite `np.foo(...)` calls to bare `foo(...)`, for every name in *func*'s diff --git a/src/blosc2/groupby.py b/src/blosc2/groupby.py index 6aaf4747d..4eb1646dd 100644 --- a/src/blosc2/groupby.py +++ b/src/blosc2/groupby.py @@ -61,7 +61,7 @@ class _Utf8KeyChunk: ascending), so null detection, live-row masking, and per-chunk ``np.unique`` all run on int64 codes; only the (few) distinct strings are ever decoded. Produced by :meth:`CTableGroupBy._read_key_chunk` via - ``Utf8Array.factorizer``. + ``UTF8Array.factorizer``. """ codes: np.ndarray @@ -151,7 +151,7 @@ def __init__( self.dropna = bool(dropna) self.engine = engine self.chunk_size = chunk_size - # Per-key incremental Utf8Factorizer instances, shared across the + # Per-key incremental UTF8Factorizer instances, shared across the # chunk loop so the string vocabulary is built once (see # _read_key_chunk). self._utf8_factorizers: dict[str, Any] = {} @@ -1567,7 +1567,7 @@ def _read_key_chunk(self, name: str, start: int, stop: int) -> np.ndarray: # is decoded, only the distinct values (codes flow through the # rest of the pipeline). The factorizer is shared across chunks # so values seen before are hash-matched instead of re-sorted. - # Utf8Array is sized to the logical row count, not the physical + # UTF8Array is sized to the logical row count, not the physical # valid_rows capacity, so a chunk boundary can run past its end; # rows beyond it are never live (the row can't have been written # without this column), so the padding code is never read live. diff --git a/src/blosc2/indexing.py b/src/blosc2/indexing.py index e4e5f9711..6826cf086 100644 --- a/src/blosc2/indexing.py +++ b/src/blosc2/indexing.py @@ -6223,6 +6223,26 @@ def _bucket_match_from_span(span: np.ndarray, plan: IndexPlan) -> np.ndarray: return match +def _coalesce_spans(spans: list[tuple[int, int]], max_gap: int) -> list[tuple[int, int]]: + """Merge spans separated by fewer than *max_gap* elements. + + A read decompresses whole blocks, so two spans landing in the same block pay + for it twice unless they are merged. Widening a span only hands more rows to + the predicate, which rejects them; merged spans stay disjoint and ordered, so + positions remain unique and sorted. + """ + if max_gap <= 0 or len(spans) < 2: + return spans + merged = [spans[0]] + for start, stop in spans[1:]: + last_start, last_stop = merged[-1] + if start - last_stop < max_gap: + merged[-1] = (last_start, max(last_stop, stop)) + else: + merged.append((start, stop)) + return merged + + def _process_bucket_chunk_batch( chunk_ids: np.ndarray, where_x, @@ -6233,15 +6253,19 @@ def _process_bucket_chunk_batch( value_parts = [] position_parts = [] local_where_x = _bucket_worker_source(where_x) + blocks = getattr(local_where_x, "blocks", None) + block_len = int(blocks[0]) if blocks else 0 for chunk_id in chunk_ids: bucket_mask = plan.bucket_masks[int(chunk_id)] chunk_start = int(chunk_id) * plan.chunk_len chunk_stop = min(chunk_start + plan.chunk_len, total_len) + spans = [] for run_start, run_stop in _contiguous_true_runs(np.asarray(bucket_mask, dtype=bool)): start = chunk_start + run_start * plan.bucket_len stop = min(chunk_start + run_stop * plan.bucket_len, chunk_stop) - if start >= stop: - continue + if start < stop: + spans.append((start, stop)) + for start, stop in _coalesce_spans(spans, block_len): if _supports_block_reads(local_where_x): span = np.empty(stop - start, dtype=local_where_x.dtype) _read_ndarray_linear_span(local_where_x, start, span) @@ -6823,6 +6847,33 @@ def _plan_multi_exact_query(plans: list[ExactPredicatePlan]) -> IndexPlan | None return None +#: Decline a bucket plan that would touch more than this fraction of the column's +#: blocks. Buckets are far smaller than blocks, so a mask can select few buckets +#: and still force a read of nearly every block — at which point the scattered +#: reads cost more than the linear scan the index is meant to replace. +_BUCKET_MAX_BLOCK_FRACTION = 0.5 + + +def _bucket_block_fraction(bucket_masks: np.ndarray, bucket: dict) -> float: + """Fraction of the column's blocks that *bucket_masks* forces a read of. + + Selectivity in buckets overstates what the index saves: a read decompresses a + whole block, so the cost unit is the block, not the bucket. + """ + masks = np.asarray(bucket_masks, dtype=bool) + if masks.size == 0: + return 0.0 + # A bucket at least as wide as a block covers whole blocks, so the clamp to 1 + # is not a special case: the grouping below then reduces to the mask itself, + # and the fraction of blocks read equals the fraction of buckets selected. + per_block = max(1, int(bucket["nav_segment_len"]) // int(bucket["bucket_len"])) + n_blocks = math.ceil(masks.shape[-1] / per_block) + padded = np.zeros((*masks.shape[:-1], n_blocks * per_block), dtype=bool) + padded[..., : masks.shape[-1]] = masks + blocks_hit = padded.reshape(*masks.shape[:-1], n_blocks, per_block).any(axis=-1) + return float(blocks_hit.mean()) + + def _plan_single_exact_query(exact_plan: ExactPredicatePlan) -> IndexPlan: kind = exact_plan.descriptor["kind"] if kind in {"full", "opsi"}: @@ -6872,7 +6923,10 @@ def _plan_single_exact_query(exact_plan: ExactPredicatePlan) -> IndexPlan: bucket = exact_plan.descriptor["bucket"] total_units = bucket_masks.size selected_units = _bit_count_sum(bucket_masks) - if selected_units < total_units: + if ( + selected_units < total_units + and _bucket_block_fraction(bucket_masks, bucket) <= _BUCKET_MAX_BLOCK_FRACTION + ): return IndexPlan( True, "bucket approximate-order index selected", diff --git a/src/blosc2/lazyexpr.py b/src/blosc2/lazyexpr.py index 515c7cd4e..4ee2eadd6 100644 --- a/src/blosc2/lazyexpr.py +++ b/src/blosc2/lazyexpr.py @@ -231,6 +231,26 @@ def _get_result(expression, chunk_operands, ne_args, where=None, indices=None, _ _constructor_call_patterns = {name: re.compile(rf"\b{re.escape(name)}\s*\(") for name in constructors} +def _restore_code_unit_shuffle(cparams: blosc2.CParams, dtype) -> None: + """Give SHUFFLE the code-unit width that constructing a CParams erases. + + Left to itself, ``blosc2.uninit()`` shuffles a `` bool: return _constructor_call_patterns[constructor].search(expression) is not None @@ -1761,7 +1781,7 @@ def _miniexpr_eligible_operand(op): and op.ndim > 0 and op.shape == shape and op.dtype.isnative - and op.dtype.kind in "biufc" + and op.dtype.kind in "biufcUS" ) return False @@ -1806,12 +1826,19 @@ def _miniexpr_eligible_operand(op): if use_miniexpr: cparams = kwargs.pop("cparams", None) - if cparams is None: + if cparams is None and getitem: # getitem output is throwaway scratch (returned as a NumPy array and # discarded), so compressing it buys nothing but a round trip. - cparams = blosc2.CParams(clevel=0) if getitem else blosc2.CParams() + cparams = blosc2.CParams(clevel=0) + # Otherwise leave cparams unset rather than passing CParams(): its + # filters_meta defaults to all zeros, which *overrides* the dtype-aware + # shuffle width uninit() would pick. For None: + """Reject variable-length utf8 operands in a UDF, naming the conversion.""" + from blosc2._utf8_array import UTF8Array, utf8_compute_error + + for operand in inputs or (): + raw = getattr(operand, "raw", operand) # a CTable Column exposes its container here + if not isinstance(raw, UTF8Array): + continue + name = getattr(operand, "_col_name", None) + source = f"t[{name!r}]" if name else "arr" + raise NotImplementedError( + utf8_compute_error( + ( + f"Column {name!r} is a variable-length utf8 column and cannot be a UDF operand." + if name + else "A variable-length UTF8Array cannot be a UDF operand." + ), + source=source, + compute="blosc2.lazyudf(kernel, (fixed,)).compute()[:]", + assignable=name is not None, + ) + ) + + class LazyUDF(LazyArray): def __init__( self, func, inputs, dtype, shape=None, chunked_eval=True, jit=None, jit_backend=None, **kwargs ): + # A utf8 operand only duck-types as an array: convert_inputs() would wrap + # it in a SimpleProxy widened to a fixed blosc2.LazyExpr: return blosc2.LazyExpr(new_op=(self, "+", value)) def __radd__(self, value: int | float | blosc2.Array, /) -> blosc2.LazyExpr: - return self.__add__(value) + # Order matters: `+` on strings is concatenation, not commutative. + _check_allowed_dtypes(value) + return blosc2.LazyExpr(new_op=(value, "+", self)) def __iadd__(self, value: int | float | blosc2.Array, /) -> blosc2.LazyExpr: return self.__add__(value) @@ -5731,7 +5734,62 @@ def _check_dtype(dtype): return dtype -def empty(shape: int | tuple | list, dtype: np.dtype | str | None = np.float64, **kwargs: Any) -> NDArray: +def _is_string_dtype(dtype) -> bool: + """True for NumPy's variable-length ``StringDType``; see ``_utf8_array``.""" + from blosc2._utf8_array import is_string_dtype + + return is_string_dtype(dtype) + + +def _asarray_string_dispatch(array, copy, kwargs): + """Route variable-length text out of :func:`asarray`'s NDArray path. + + Returns ``(result, array)``. *result* is a :class:`UTF8Array` when the + **target** dtype is NumPy's ``StringDType``, and ``None`` otherwise -- in + which case *array* comes back ready for the fixed-width path, so + ``asarray(utf8_source, dtype=" NDArray | blosc2.UTF8Array: """Create an empty array. Parameters @@ -5785,6 +5843,8 @@ def empty(shape: int | tuple | list, dtype: np.dtype | str | None = np.float64, >>> array.dtype dtype('int32') """ + if _is_string_dtype(dtype): + return _utf8_filled(shape, "", **kwargs) dtype = _check_dtype(dtype) shape = _check_shape(shape) kwargs = _check_ndarray_kwargs(**kwargs) @@ -5862,7 +5922,9 @@ def nans(shape: int | tuple | list, dtype: np.dtype | str = np.float64, **kwargs return blosc2_ext.nans(shape, chunks, blocks, dtype, **kwargs) -def zeros(shape: int | tuple | list, dtype: np.dtype | str = np.float64, **kwargs: Any) -> NDArray: +def zeros( + shape: int | tuple | list, dtype: np.dtype | str = np.float64, **kwargs: Any +) -> NDArray | blosc2.UTF8Array: """Create an array with zero as the default value for uninitialized portions of the array. @@ -5893,6 +5955,8 @@ def zeros(shape: int | tuple | list, dtype: np.dtype | str = np.float64, **kwarg >>> array.dtype dtype('float64') """ + if _is_string_dtype(dtype): + return _utf8_filled(shape, "", **kwargs) dtype = _check_dtype(dtype) shape = _check_shape(shape) kwargs = _check_ndarray_kwargs(**kwargs) @@ -5907,7 +5971,7 @@ def full( fill_value: bytes | int | float | bool, dtype: np.dtype | str = None, **kwargs: Any, -) -> NDArray: +) -> NDArray | blosc2.UTF8Array: """Create an array, with :paramref:`fill_value` being used as the default value for uninitialized portions of the array. @@ -5947,6 +6011,8 @@ def full( >>> array.dtype dtype('bool') """ + if _is_string_dtype(dtype): + return _utf8_filled(shape, str(fill_value), **kwargs) if isinstance(fill_value, bytes): dtype = np.dtype(f"S{len(fill_value)}") if dtype is None: @@ -5962,7 +6028,9 @@ def full( return blosc2_ext.full(shape, chunks, blocks, fill_value, dtype, **kwargs) -def ones(shape: int | tuple | list, dtype: np.dtype | str = None, **kwargs: Any) -> NDArray: +def ones( + shape: int | tuple | list, dtype: np.dtype | str = None, **kwargs: Any +) -> NDArray | blosc2.UTF8Array: """Create an array with one as values. The parameters and keyword arguments are the same as for the @@ -6669,7 +6737,9 @@ def _ndarray_asarray_requires_copy( return builtins.any(key in user_kwargs for key in copy_keys) -def asarray(array: Sequence | blosc2.Array, copy: bool | None = None, **kwargs: Any) -> NDArray: +def asarray( + array: Sequence | blosc2.Array, copy: bool | None = None, **kwargs: Any +) -> NDArray | blosc2.UTF8Array: """Convert the `array` to an `NDArray`. Parameters @@ -6689,9 +6759,11 @@ def asarray(array: Sequence | blosc2.Array, copy: bool | None = None, **kwargs: Returns ------- - out: :ref:`NDArray` + out: :ref:`NDArray` or :class:`UTF8Array` A new :ref:`NDArray` made of :paramref:`array`, or the original - array when a copy is not required. + array when a copy is not required. When the target dtype is NumPy's + variable-length ``StringDType``, a :class:`UTF8Array` is returned + instead -- see the Notes. Notes ----- @@ -6700,6 +6772,16 @@ def asarray(array: Sequence | blosc2.Array, copy: bool | None = None, **kwargs: be used for ingesting e.g. disk or network based arrays very effectively and without consuming lots of memory. + ``StringDType`` cannot back an NDArray: it keeps each row's payload outside + the array buffer (a 100-character string still reports ``nbytes == 16``) + and offers no buffer protocol, so compressing that buffer would persist + pointers. Such input is therefore stored as a :class:`UTF8Array`, which + holds the same text as offsets + UTF-8 bytes -- the layout Arrow uses for + ``large_string``. The dispatch is on the *target* dtype, so + ``asarray(utf8_source, dtype=">> import blosc2 @@ -6716,6 +6798,9 @@ def asarray(array: Sequence | blosc2.Array, copy: bool | None = None, **kwargs: raise ValueError("Only unsafe casting is supported at the moment.") if not hasattr(array, "shape"): array = np.asarray(array) # defaults if dtype=None + utf8_out, array = _asarray_string_dispatch(array, copy, kwargs) + if utf8_out is not None: + return utf8_out dtype_ = blosc2.proxy.convert_dtype(array.dtype) dtype = blosc2.proxy.convert_dtype(kwargs.pop("dtype", dtype_)) # check if dtype provided kwargs = _check_ndarray_kwargs(**kwargs) diff --git a/src/blosc2/proxy.py b/src/blosc2/proxy.py index a827f00b1..425416359 100644 --- a/src/blosc2/proxy.py +++ b/src/blosc2/proxy.py @@ -673,7 +673,17 @@ class SimpleProxy(blosc2.Operand): """ def __init__(self, src, chunks: tuple | None = None, blocks: tuple | None = None): - if not hasattr(src, "shape") or not hasattr(src, "dtype"): + from blosc2._utf8_array import UTF8Array + + if isinstance(src, UTF8Array): + # The compute engine indexes chunk-wise into fixed-width elements, + # which a variable-length utf8 array has not got, so widen it here. + # (lazyexpr() routes utf8 operands to the span driver instead; this + # is the fallback for the entry points that do not.) Until this + # array grew a .shape, the branch below did the same thing by + # accident, via np.asarray. + src = src.astype() + elif not hasattr(src, "shape") or not hasattr(src, "dtype"): # If the source is not an array, convert it to NumPy src = np.asarray(src) if not hasattr(src, "__getitem__"): @@ -762,6 +772,44 @@ def as_simpleproxy(*arrs: Sequence[blosc2.Array]) -> tuple[SimpleProxy | blosc2. return out[0] if len(out) == 1 else out +def _is_pandas_string_series(col) -> bool: + """True for a pandas string column. + + pandas 3's `str` dtype reports `kind == "O"`, so the kind is useless here + and `pd.api.types.is_string_dtype` is the only reliable test. + """ + try: + import pandas as pd + except ImportError: + return False + return pd.api.types.is_string_dtype(getattr(col, "dtype", None)) + + +def _string_series_to_numpy(col, label=None): + """A pandas string column as a fixed-width ` TypeError, + `row['x'].lower()` -> AttributeError), so quietly turning one into `""` + would invent a value pandas never produces. + """ + if label is None: + label = col.name + if col.isna().any(): + raise ValueError( + f"blosc2.jit: string column {label!r} contains nulls, and a row-wise kernel " + "over a null raises in pandas too. Fill them first, e.g. " + f"df[{label!r}] = df[{label!r}].fillna('')." + ) + values = col.to_numpy(dtype=object) + return values.astype(str) + + class _PandasRowProxy(blosc2.Operand): """Row proxy for `PandasUdfEngine.apply`'s axis=1 route. @@ -777,6 +825,17 @@ def __init__(self, df): self._df = df self._cache = {} + def _raw_column(self, key): + """The column as a plain array, for the DSL route. + + The tracing route wants a SimpleProxy operand; a DSL kernel wants the + array itself, and accepts string columns the traced one does not. + """ + col = self._df[key] + if _is_pandas_string_series(col): + return _string_series_to_numpy(col, key) + return col.to_numpy() + def __getitem__(self, key): if not isinstance(key, str): raise TypeError( @@ -796,11 +855,11 @@ def __getitem__(self, key): f"row[{key!r}]: column label is duplicated ({n_matches} matches); " "axis=1 row proxies require unique column labels" ) - col = self._df[key].to_numpy() - if col.dtype.kind not in "biufc": + col = self._raw_column(key) + if col.dtype.kind not in "biufcUS": raise ValueError( - f"row[{key!r}]: column has dtype {col.dtype!r}, which is not numeric. " - "The Blosc2 engine only supports vectorized numeric computations." + f"row[{key!r}]: column has dtype {col.dtype!r}, which the Blosc2 engine " + "cannot vectorize. Numeric, boolean and string columns are supported." ) proxy = SimpleProxy(col) self._cache[key] = proxy @@ -904,6 +963,47 @@ def _signature_params(func) -> list: return [] +def _row_column(row, label): + """The raw column array behind *label*, from a row proxy or a DataFrame.""" + getter = getattr(row, "_raw_column", None) + if getter is not None: + return getter(label) + col = row[label] + if isinstance(col, np.ndarray | blosc2.NDArray): + return col + if _is_pandas_string_series(col): + return _string_series_to_numpy(col, label) + return np.asarray(col) + + +def _dsl_operand_values(kernel: DSLKernel, sig, args, func_kwargs) -> tuple: + """The kernel's operands, one per DSL input name, in declaration order.""" + if kernel.row_columns and len(args) == 1 and not func_kwargs: + # The `row["colname"]` kernel: its signature still says one row, but the + # compiled kernel takes one parameter per referenced column. + values = tuple(_row_column(args[0], label) for label in kernel.row_columns.values()) + else: + try: + bound = sig.bind(*args, **func_kwargs) + except TypeError as e: + # sig.bind's message names no function; prefix it, and point at the + # subsetting fix when a wide DataFrame was unpacked into the call. + hint = _wide_frame_hint(e, kernel.__name__, kernel.input_names or sig.parameters) + raise TypeError(f"{kernel.__name__}() {e}" + (f"\n{hint}" if hint else "")) from None + bound.apply_defaults() + values = tuple(bound.arguments[name] for name in kernel.input_names) + # Accept array-protocol operands (pandas Series, polars Series, ...) the same + # way the tracing route already does; zero-copy when the source is numpy-backed. + return tuple( + np.asarray(v) + if not isinstance(v, np.ndarray | blosc2.NDArray) + and hasattr(v, "__array__") + and getattr(v, "ndim", 0) > 0 + else v + for v in values + ) + + def _jit_dsl_wrapper(kernel: DSLKernel, out, decorator_kwargs: dict): """Build the call wrapper for the DSL (control-flow) dispatch route of `jit`. @@ -917,26 +1017,7 @@ def dsl_wrapper(*args, **func_kwargs): sig = kernel._sig if sig is None: raise TypeError(f"@blosc2.jit: cannot introspect the signature of {kernel.__name__!r}") - try: - bound = sig.bind(*args, **func_kwargs) - except TypeError as e: - # sig.bind's message names no function; prefix it, and point at the - # subsetting fix when a wide DataFrame was unpacked into the call. - hint = _wide_frame_hint(e, kernel.__name__, kernel.input_names or sig.parameters) - raise TypeError(f"{kernel.__name__}() {e}" + (f"\n{hint}" if hint else "")) from None - bound.apply_defaults() - values = tuple(bound.arguments[name] for name in kernel.input_names) - # Accept array-protocol operands (pandas Series, polars Series, ...) the - # same way the tracing route already does; zero-copy when the source is - # numpy-backed. - values = tuple( - np.asarray(v) - if not isinstance(v, np.ndarray | blosc2.NDArray) - and hasattr(v, "__array__") - and getattr(v, "ndim", 0) > 0 - else v - for v in values - ) + values = _dsl_operand_values(kernel, sig, args, func_kwargs) array_shapes = { v.shape @@ -992,6 +1073,9 @@ def dsl_wrapper(*args, **func_kwargs): return out if storage_kwargs and any(v is not None for v in storage_kwargs.values()): + # Execution-tuning kwargs go along too: compute() names all three, + # while lazyudf() above only names jit/jit_backend, so fp_accuracy + # would otherwise be dropped on this path. return lexpr.compute(**decorator_kwargs) return lexpr[()] @@ -1171,26 +1255,26 @@ def wrapper(*args, **func_kwargs): try: retval = func(*new_args, **func_kwargs) except Exception as e: - hints = [ - hint - for hint in ( - _wide_frame_hint( - e, getattr(func, "__name__", "the function"), _signature_params(func) - ), - _trace_hint, - ) - if hint is not None - ] - if hints: - raise type(e)("\n".join([str(e), *hints])) from e + # Notes rather than a re-raise: type(e)(msg) assumes a one-argument + # constructor, and any exception needing more (or rejecting a bare + # string) would surface as a TypeError instead of the real failure. + for hint in ( + _wide_frame_hint(e, getattr(func, "__name__", "the function"), _signature_params(func)), + _trace_hint, + ): + if hint is not None: + e.add_note(hint) raise # Treat return value # If it is a numpy array, return it as is if isinstance(retval, np.ndarray): if storage_kwargs and any(v is not None for v in storage_kwargs.values()): - # But if storage kwargs are provided, return a NDArray instead - return blosc2.asarray(retval, **kwargs) + # But if storage kwargs are provided, return a NDArray instead. + # Only storage kwargs: asarray() rejects the execution-tuning + # ones, and there is nothing left to tune -- the function has + # already run. + return blosc2.asarray(retval, **storage_kwargs) return retval # In some instances, the return value is not a LazyExpr @@ -1266,13 +1350,21 @@ def apply(cls, data, func, args, kwargs, decorator, axis): function once for each column or row. """ orig = data - values = cls._ensure_numpy_data(data) func_name = getattr(func, "__name__", "the function") uses_subscript, has_loop = ( _analyze_row_func(_undecorated(func)) if hasattr(orig, "columns") else (False, False) ) + # The row-proxy route reads columns one at a time and never needs the + # whole frame as one array, so a non-numeric column (a pandas string + # column, say) is fine there and only `nrows` is wanted. + if uses_subscript and axis in (1, "columns"): + values = None + nrows = len(orig) + else: + values = cls._ensure_numpy_data(data) + nrows = values.shape[0] func = _decorate_once(func, decorator) - if values.ndim == 1 or axis is None: + if values is not None and (values.ndim == 1 or axis is None): # pandas Series.apply or pipe result = func(values, *args, **kwargs) elif axis in (0, "index"): @@ -1302,14 +1394,10 @@ def apply(cls, data, func, args, kwargs, decorator, axis): # per-column dtypes survive. row_proxy = _PandasRowProxy(orig) result = func(row_proxy, *args, **kwargs) - if not ( - isinstance(result, np.ndarray) - and result.ndim == 1 - and result.shape[0] == values.shape[0] - ): + if not (isinstance(result, np.ndarray) and result.ndim == 1 and result.shape[0] == nrows): raise TypeError( '@blosc2.jit engine=... axis=1: functions using row["colname"] must ' - f"return one scalar per row (shape ({values.shape[0]},)); got " + f"return one scalar per row (shape ({nrows},)); got " f"{result!r}. Returning multiple values per row is not supported here." ) else: diff --git a/src/blosc2/scalar_array.py b/src/blosc2/scalar_array.py index 5865f5c01..c9d5c0bd5 100644 --- a/src/blosc2/scalar_array.py +++ b/src/blosc2/scalar_array.py @@ -24,6 +24,8 @@ from collections import defaultdict from typing import TYPE_CHECKING, Any +import numpy as np + if TYPE_CHECKING: from collections.abc import Iterable, Iterator @@ -256,6 +258,22 @@ def flush(self) -> None: self._pending.clear() self._invalidate_prefix_cache() + def set_all(self, values: Iterable[Any]) -> None: + """Replace the whole content, keeping the current row count. + + Writes each backing batch exactly once, where the equivalent loop over + :meth:`__setitem__` would rewrite a whole batch per row. Mirrors + ``UTF8Array.set_all`` so callers can treat both the same way. + """ + coerced = [self._coerce(v) for v in values] + if len(coerced) != len(self): + raise ValueError(f"set_all() expects {len(self)} values, got {len(coerced)}.") + prefix = self._persisted_prefix_sums() + for batch_index in range(len(prefix) - 1): + self._backend[batch_index] = coerced[prefix[batch_index] : prefix[batch_index + 1]] + # Batch lengths are unchanged, so the prefix cache stays valid. + self._pending = coerced[self._persisted_row_count :] + # ------------------------------------------------------------------ # Public read interface # ------------------------------------------------------------------ @@ -266,6 +284,32 @@ def __len__(self) -> int: def __iter__(self) -> Iterator[Any]: yield from self[:] + # Identity hashing is kept: these objects were hashable before __eq__ was + # defined, and an element-wise __eq__ never returns a bool for the hash + # contract to apply to. + __hash__ = object.__hash__ + + def __eq__(self, other): + """Element-wise equality mask over the stored rows. + + Without this the comparison fell through to object identity and + ``column == "value"`` was a plain ``False`` — silently wrong. Rows hold + arbitrary msgpack payloads, so the comparison is handed to NumPy over + the decoded values. + """ + import blosc2 + + if blosc2._disable_overloaded_equal: + return self is other + return np.asarray(self[:], dtype=object) == other + + def __ne__(self, other): + import blosc2 + + if blosc2._disable_overloaded_equal: + return self is not other + return np.asarray(self[:], dtype=object) != other + def __getitem__(self, index: int | slice | list | tuple) -> Any | list[Any]: if isinstance(index, int): n = len(self) diff --git a/src/blosc2/schema.py b/src/blosc2/schema.py index ad77e1a44..bdb131ae0 100644 --- a/src/blosc2/schema.py +++ b/src/blosc2/schema.py @@ -597,7 +597,7 @@ def to_metadata_dict(self) -> dict[str, Any]: return d -class Utf8Spec(SchemaSpec): +class UTF8Spec(SchemaSpec): """Variable-length UTF-8 string column stored Arrow-style as offsets + bytes. Unlike :class:`string`, this spec does not use a fixed-width NumPy dtype: @@ -622,6 +622,17 @@ class Utf8Spec(SchemaSpec): def __init__(self, *, nullable: _builtin_bool = False, null_value: str | None = None): if null_value is not None and not isinstance(null_value, str): raise TypeError(f"utf8 null_value must be str, got {type(null_value).__name__!r}") + if null_value == "\x00": + # NumPy 2.4 compares a lone NUL against StringDType as no-match: + # np.array(["\x00"], dtype=StringDType()) == "\x00" is False, while + # "\x00x" and "a\x00b" both compare correctly. Every null mask here + # is such a comparison, so this sentinel would silently stop marking + # anything as null. Reject it rather than mis-handle it. + raise ValueError( + "utf8 null_value cannot be a single NUL character: NumPy does not " + "match it against StringDType arrays, so nulls would go undetected. " + "Use a longer sentinel (the default is '__BLOSC2_NULL__')." + ) self.nullable = nullable or null_value is not None self.null_value = _normalize_scalar_value(null_value) @@ -640,6 +651,12 @@ def display_label(self) -> str: return "utf8" +#: Deprecated alias kept for the name this class shipped under in 4.9.1. +#: Persisted schemas are unaffected either way -- they record ``kind: "utf8"``, +#: never the class name. +Utf8Spec = UTF8Spec + + class ObjectSpec(SchemaSpec): """Schema-less Python object column backed by batched msgpack storage. @@ -833,7 +850,7 @@ def vlstring( ) -def utf8(*, nullable: bool = False, null_value: str | None = None) -> Utf8Spec: +def utf8(*, nullable: bool = False, null_value: str | None = None) -> UTF8Spec: """Build a variable-length UTF-8 string schema descriptor. Use this for high-cardinality or free-text string columns: values are @@ -847,14 +864,14 @@ def utf8(*, nullable: bool = False, null_value: str | None = None) -> Utf8Spec: inferred automatically from plain ``str`` annotations. utf8 columns support vectorized comparisons (``==``, ``!=``, ``<``, - ``<=``, ``>``, ``>=``), :meth:`CTable.group_by` keys, - :meth:`CTable.sort_by`, and Arrow/Parquet interop. Current limitations: - :meth:`CTable.create_index` is not supported yet (use a fixed-width - :class:`string` column if you need an index), and string-*expression* - filters such as ``t.where("name == 'x'")`` are not supported yet — use - the operator form ``t[t.name == 'x']`` instead. See - :ref:`ChoosingStringType` for a full comparison with :class:`string` - and :func:`vlstring`. + ``<=``, ``>``, ``>=``), string-expression filters such as + ``t.where("name == 'x'")`` and ``t.where("startswith(name, 'x')")``, + :meth:`CTable.group_by` keys, :meth:`CTable.sort_by`, Arrow/Parquet + interop, and :meth:`CTable.create_index`, which indexes the alphabetical + rank of each value and accelerates sorting and scalar comparisons (but + not ``startswith``/substring searches, which no index covers). See + :ref:`ChoosingStringType` for a full comparison with :class:`string` and + :func:`vlstring`. Parameters ---------- @@ -873,10 +890,10 @@ def utf8(*, nullable: bool = False, null_value: str | None = None) -> Utf8Spec: ... name: str = b2.field(b2.utf8()) ... note: str = b2.field(b2.utf8(nullable=True)) """ - from blosc2.utf8_array import string_dtype + from blosc2._utf8_array import string_dtype string_dtype() # fail early with a clear error on NumPy < 2.0 - return Utf8Spec(nullable=nullable, null_value=null_value) + return UTF8Spec(nullable=nullable, null_value=null_value) def object( diff --git a/src/blosc2/schema_compiler.py b/src/blosc2/schema_compiler.py index b295dd51e..2475474d7 100644 --- a/src/blosc2/schema_compiler.py +++ b/src/blosc2/schema_compiler.py @@ -28,7 +28,7 @@ ObjectSpec, SchemaSpec, StructSpec, - Utf8Spec, + UTF8Spec, VLBytesSpec, VLStringSpec, complex64, @@ -77,7 +77,7 @@ "bytes": b2_bytes, "vlstring": VLStringSpec, "vlbytes": VLBytesSpec, - "utf8": Utf8Spec, + "utf8": UTF8Spec, "object": ObjectSpec, "timestamp": timestamp, # dictionary @@ -112,7 +112,7 @@ def compute_display_width(spec: SchemaSpec) -> int: """Return a reasonable terminal display width for *spec*'s column.""" if isinstance(spec, DictionarySpec): return 32 - if isinstance(spec, (VLStringSpec, VLBytesSpec, ObjectSpec, Utf8Spec)): + if isinstance(spec, (VLStringSpec, VLBytesSpec, ObjectSpec, UTF8Spec)): return 40 if isinstance(spec, NDArraySpec): return max(20, len(spec.display_label()) + 4) diff --git a/tests/b2view/test_cli.py b/tests/b2view/test_cli.py index 7c877e66d..9b7dd7f5f 100644 --- a/tests/b2view/test_cli.py +++ b/tests/b2view/test_cli.py @@ -39,14 +39,14 @@ def test_download_skipped_when_file_already_in_cwd(): assert info_url is None -def test_download_urls_keep_relative_path_dest_is_basename(): +def test_download_url_dest_is_basename(): urlpath, url, info_url = resolve_source(None, "sub/dir/bundle.b2z", exists=lambda p: False) assert urlpath == "bundle.b2z" assert url == DOWNLOAD_BASE_URL + "sub/dir/bundle.b2z" assert info_url == INFO_BASE_URL + "sub/dir/bundle.b2z" -def test_download_and_positional_are_mutually_exclusive(): +def test_download_and_positional_exclusive(): with pytest.raises(ValueError, match="cannot be combined"): resolve_source("local.b2z", "foo.b2z") diff --git a/tests/b2view/test_group.py b/tests/b2view/test_group.py index cd8c30314..8affbd314 100644 --- a/tests/b2view/test_group.py +++ b/tests/b2view/test_group.py @@ -169,7 +169,7 @@ def test_group_sort_noop_when_not_grouped(group_store): assert browser.get_group_sort("/ctable") is None -def test_group_bars_categorical_is_bar_sorted_desc_and_capped(group_store): +def test_group_bars_categorical_sorted_and_capped(group_store): """A dictionary key yields capped bars ranked by the aggregate descending.""" path, _ = group_store with StoreBrowser(path) as browser: @@ -195,7 +195,7 @@ def test_group_bars_numeric_is_line_pareto_by_default(group_store): assert "rank" in bars["xlabel"] -def test_group_bars_numeric_sorted_by_key_uses_key_on_x(group_store): +def test_group_bars_numeric_sorted_by_key(group_store): """Sorting a numeric-key result by the key puts key values on X in that order.""" path, _ = group_store with StoreBrowser(path) as browser: diff --git a/tests/b2view/test_plot_model.py b/tests/b2view/test_plot_model.py index 38017a828..44d272db8 100644 --- a/tests/b2view/test_plot_model.py +++ b/tests/b2view/test_plot_model.py @@ -99,7 +99,7 @@ def test_stream_envelope_matches_full_read_ctable(plot_store, monkeypatch): _assert_exact(env, vals) -def test_stream_envelope_captures_spike_a_sample_would_miss(plot_store, monkeypatch): +def test_stream_envelope_captures_spike(plot_store, monkeypatch): path, vals = plot_store _force_stream(monkeypatch) with StoreBrowser(path) as browser: @@ -222,7 +222,7 @@ def test_read_series_clamps_range(plot_store): assert clamped["y"].shape == (N,) -def test_locked_row_window_confines_plot_and_read_series(plot_store): +def test_locked_row_window_confines_plot(plot_store): """A locked row window (the 'v' action) takes precedence over the full series in both plot_series and read_series, matching preview()/read_cell() (PR #663 review): a plot/hi-res of a windowed CTable shows only its rows.""" diff --git a/tests/b2view/test_sort.py b/tests/b2view/test_sort.py index f59448b04..ba721e321 100644 --- a/tests/b2view/test_sort.py +++ b/tests/b2view/test_sort.py @@ -185,7 +185,7 @@ async def _wait_for_table(pilot) -> None: @pytest.mark.asyncio @pytest.mark.tui -async def test_sort_key_opens_screen_applies_and_escape_clears(sort_store): +async def test_sort_key_applies_and_escape_clears(sort_store): path, _, _ = sort_store app = B2ViewApp(path, start_panel="data") async with app.run_test(size=TERM_SIZE) as pilot: diff --git a/tests/ctable/test_arrow_interop.py b/tests/ctable/test_arrow_interop.py index ed687c9a6..044686343 100644 --- a/tests/ctable/test_arrow_interop.py +++ b/tests/ctable/test_arrow_interop.py @@ -351,7 +351,7 @@ def test_from_arrow_string_fixed_width_with_max_length(): assert t["name"][:].tolist() == ["hi", "hello world", "!"] -def test_from_arrow_list_struct_nullable_values_roundtrip(): +def test_from_arrow_list_struct_nullable(): nutrient_type = pa.struct( [ pa.field("name", pa.string()), @@ -454,7 +454,7 @@ def test_from_arrow_dictionary_codes_use_aligned_grid(): assert list(t["c"][:5]) == c.to_pylist()[:5] -def test_to_arrow_dictionary_multi_batch_with_deletions(): +def test_to_arrow_dict_multi_batch_deletions(): """Dictionary-column export across several batches, with holes in the live-row mask from a deletion, still maps each batch to the correct physical positions. diff --git a/tests/ctable/test_column.py b/tests/ctable/test_column.py index 33d922cc0..8d880206e 100644 --- a/tests/ctable/test_column.py +++ b/tests/ctable/test_column.py @@ -439,7 +439,7 @@ def test_sum_empty_filtered_view_returns_zero(): assert t[t.id < 0]["id"].sum() == 0 -def test_sum_where_skips_valid_rows_mask_when_all_rows_visible(): +def test_sum_where_skips_mask_when_all_visible(): t = CTable(Row, new_data=DATA20, expected_size=len(DATA20)) mask = t["id"]._lazy_nonnull_mask(where=t["score"] < 100) assert mask.expression == "(o0 < 100)" @@ -904,7 +904,7 @@ def test_column_repr_shows_preview_values(): assert "..." in r -def test_info_omits_capacity_and_read_only_for_in_memory_table(): +def test_info_omits_capacity_for_in_memory(): t = CTable(Row, new_data=DATA20) info = repr(t.info) assert "capacity" not in info @@ -1077,7 +1077,7 @@ def test_ctable_setitem_view_raises(): view["score"] = np.zeros(len(view)) -def test_column_setitem_ndarray_fast_path_on_disk_table(tmp_path): +def test_setitem_ndarray_fast_path_on_disk(tmp_path): """Fast path fires for a disk-opened table (not just freshly-built in-memory tables).""" n = 60 urlpath = str(tmp_path / "tbl.b2") @@ -1113,7 +1113,7 @@ class R: np.testing.assert_allclose(t["val"][:], np.arange(n, dtype=np.float64) * 3.14) -def test_column_setitem_blosc2_ndarray_no_holes_uneven_chunks(): +def test_setitem_b2_ndarray_no_holes_uneven(): """Fast path works when nrows is not a multiple of chunk_size.""" n = 70 diff --git a/tests/ctable/test_column_ndarray_like.py b/tests/ctable/test_column_ndarray_like.py index d12de3e76..847c16fb4 100644 --- a/tests/ctable/test_column_ndarray_like.py +++ b/tests/ctable/test_column_ndarray_like.py @@ -26,7 +26,7 @@ def test_column_logical_metadata(): assert view.x.size == 3 -def test_column_boolean_operators_build_lazy_expressions(): +def test_boolean_operators_build_lazy_exprs(): t = blosc2.CTable(Row, new_data=DATA) view = t.where(t.flag & (t.x > 0)) diff --git a/tests/ctable/test_column_slice_fastpath.py b/tests/ctable/test_column_slice_fastpath.py index 48c52e243..d8465528f 100644 --- a/tests/ctable/test_column_slice_fastpath.py +++ b/tests/ctable/test_column_slice_fastpath.py @@ -148,7 +148,7 @@ def test_deletions_use_position_path_and_stay_correct(key): @pytest.mark.parametrize("key", [np.s_[::5], np.s_[::-2], np.s_[::-1]]) -def test_filtered_view_uses_position_path_and_stays_correct(key): +def test_filtered_view_uses_position_path(key): table, arr = _make() view = table.where("a >= 500") expected = arr["b"][arr["a"] >= 500] diff --git a/tests/ctable/test_csv_interop.py b/tests/ctable/test_csv_interop.py index 9a2063c95..8a904d807 100644 --- a/tests/ctable/test_csv_interop.py +++ b/tests/ctable/test_csv_interop.py @@ -356,7 +356,7 @@ def test_from_csv_ndarray_wrong_shape_raises(tmp_csv): CTable.from_csv(tmp_csv, NdarrayRow) -def test_from_csv_nonnullable_ndarray_empty_cell_raises(tmp_csv): +def test_csv_nonnullable_empty_cell_raises(tmp_csv): with open(tmp_csv, "w") as f: f.write("id,embedding\n") f.write("1,\n") diff --git a/tests/ctable/test_ctable_computed_cols.py b/tests/ctable/test_ctable_computed_cols.py index 9df61616b..9c90040ef 100644 --- a/tests/ctable/test_ctable_computed_cols.py +++ b/tests/ctable/test_ctable_computed_cols.py @@ -164,7 +164,7 @@ def test_computed_column_where_via_col(): assert len(view) == 3 # 9, 16, 25 -def test_getitem_boolean_lazyexpr_matches_where_for_computed_column(): +def test_getitem_bool_lazyexpr_matches_where(): t = _make_invoice_table(5) t.add_computed_column("total", lambda cols: cols["price"] * cols["qty"]) expr = t.total >= 9 @@ -499,7 +499,7 @@ def test_materialize_computed_column_extend_autofill(): np.testing.assert_allclose(t["total_stored"][:], [1.0, 4.0, 9.0, 16.0]) -def test_materialize_computed_column_explicit_append_value_wins(): +def test_materialize_explicit_append_wins(): t = _make_invoice_table(2) t.add_computed_column("total", lambda cols: cols["price"] * cols["qty"]) t.materialize_computed_column("total", new_name="total_stored") @@ -745,7 +745,7 @@ def test_materialize_computed_column_open_roundtrip(tmp_path): assert t2.index("total_stored").kind == "full" -def test_materialize_computed_column_open_append_autofill(tmp_path): +def test_materialize_open_append_autofill(tmp_path): path = str(tmp_path / "tbl") t = CTable(Invoice, [(1.0, 1, 0.1), (2.0, 2, 0.1)], urlpath=path, mode="w") t.add_computed_column("total", lambda cols: cols["price"] * cols["qty"]) @@ -1066,7 +1066,7 @@ def patched(self, expr): t.add_computed_column("total", lambda cols: cols["price"] * cols["qty"]) -def test_add_computed_column_malformed_expression_raises(monkeypatch): +def test_add_computed_malformed_expr_raises(monkeypatch): """add_computed_column raises ValueError when the expression string cannot be re-parsed.""" t = _make_invoice_table() diff --git a/tests/ctable/test_ctable_dataclass_schema.py b/tests/ctable/test_ctable_dataclass_schema.py index 8a2741eae..28e8a2116 100644 --- a/tests/ctable/test_ctable_dataclass_schema.py +++ b/tests/ctable/test_ctable_dataclass_schema.py @@ -165,7 +165,7 @@ class ArrayRow: assert np.array_equal(reopened.matrix[:], data) -def test_fixed_shape_ndarray_column_rejects_wrong_shape(): +def test_fixed_shape_ndarray_rejects_wrong_shape(): @dataclass class ArrayRow: matrix: np.ndarray = blosc2.field(blosc2.ndarray((2, 3), dtype=np.float64)) # noqa: RUF009 diff --git a/tests/ctable/test_ctable_indexing.py b/tests/ctable/test_ctable_indexing.py index ab2e01161..e2cc901ab 100644 --- a/tests/ctable/test_ctable_indexing.py +++ b/tests/ctable/test_ctable_indexing.py @@ -103,7 +103,7 @@ def test_where_with_index_matches_scan_in_memory(): @pytest.mark.heavy -def test_indexed_where_view_sort_by_reuses_cached_live_positions(monkeypatch): +def test_indexed_where_sort_by_reuses_live_pos(monkeypatch): t = _make_table(200) t.create_index("id", kind=blosc2.IndexKind.FULL) @@ -128,7 +128,7 @@ def test_create_expression_index_in_memory(): @pytest.mark.heavy -def test_where_with_expression_index_matches_scan_in_memory(): +def test_where_expr_index_matches_scan(): t = _make_table(200) t.create_index(expression="value * category", kind=blosc2.IndexKind.FULL, name="vc") result_idx = t.where((t._cols["value"] * t._cols["category"]) >= 150) @@ -196,7 +196,7 @@ def test_stale_on_column_assign_in_memory(): assert t.index("id").stale -def test_delete_bumps_visibility_epoch_not_stale_in_memory(): +def test_delete_bumps_epoch_not_stale(): t = _make_table(20) t.create_index("id") t.delete(0) @@ -226,7 +226,7 @@ def test_compact_index_in_memory(): @pytest.mark.heavy -def test_multi_column_conjunction_uses_multiple_indexes_in_memory(): +def test_conjunction_uses_multiple_indexes(): t = _make_table(200) t.create_index("id", kind=blosc2.IndexKind.FULL) t.create_index("category", kind=blosc2.IndexKind.FULL) @@ -240,7 +240,7 @@ def test_multi_column_conjunction_uses_multiple_indexes_in_memory(): assert ids_idx == ids_scan -def test_full_index_large_ctable_column_matches_scan_in_memory(): +def test_full_index_large_column_matches_scan(): @dataclasses.dataclass class SensorRow: sensor_id: int = blosc2.field(blosc2.int32()) @@ -288,7 +288,7 @@ def test_create_index_persistent(tmpdir): assert sidecars, "No sidecar .b2nd files found" -def test_create_index_persistent_does_not_cache_sidecar_handles(tmpdir): +def test_create_index_does_not_cache_sidecars(tmpdir): import blosc2.indexing as indexing path = str(tmpdir / "table.b2d") @@ -303,7 +303,7 @@ def test_create_index_persistent_does_not_cache_sidecar_handles(tmpdir): assert cached == [] -def test_persistent_ctable_releases_immediately_without_gc(tmpdir): +def test_persistent_releases_without_gc(tmpdir): path = str(tmpdir / "table.b2d") def build_table(): @@ -372,7 +372,7 @@ def test_where_with_index_matches_scan_persistent(tmpdir): @pytest.mark.heavy -def test_relative_b2d_ctable_index_sidecars_survive_reopen(tmp_path, monkeypatch): +def test_relative_b2d_sidecars_survive_reopen(tmp_path, monkeypatch): monkeypatch.chdir(tmp_path) t = _make_table(200, persistent_path="table.b2d") t.create_index("id", kind=blosc2.IndexKind.BUCKET) @@ -385,7 +385,7 @@ def test_relative_b2d_ctable_index_sidecars_survive_reopen(tmp_path, monkeypatch @pytest.mark.heavy -def test_persistent_index_drop_releases_sidecars_without_gc(tmpdir): +def test_index_drop_releases_sidecars_no_gc(tmpdir): import gc def run_query_and_drop(): @@ -434,7 +434,7 @@ def test_expression_index_persistent_roundtrip(tmpdir): assert len(result) > 0 -def test_sort_by_computed_column_with_expression_full_index(): +def test_sort_by_computed_col_full_index(): t = _make_table(40) t.add_computed_column("score", "value * category") t.create_index(expression="value * category", kind=blosc2.IndexKind.FULL, name="score_expr") @@ -462,7 +462,7 @@ def test_drop_index_persistent_catalog_cleared(tmpdir): assert len(t2.indexes) == 0 -def test_drop_indexed_column_removes_persistent_sidecars(tmpdir): +def test_drop_indexed_col_removes_sidecars(tmpdir): path = str(tmpdir / "table.b2d") t = _make_table(30, persistent_path=path) t.create_index("id") @@ -539,7 +539,7 @@ def test_query_after_reopen_persistent(tmpdir): assert ids == list(range(91, 100)) -def test_rename_indexed_column_rebuilds_catalog_persistent(tmpdir): +def test_rename_indexed_col_rebuilds_catalog(tmpdir): path = str(tmpdir / "table.b2d") t = _make_table(40, persistent_path=path) t.create_index("id") @@ -629,7 +629,7 @@ def test_indexes_multiple_columns(): assert col_names == {"id", "category"} -def test_indexed_ctable_b2z_double_open_append_no_corruption(tmp_path): +def test_b2z_double_open_append_no_corruption(tmp_path): """Opening an indexed CTable .b2z in append mode twice must not corrupt it. Regression test: GC of a CTable opened from .b2z was calling close() → @@ -688,7 +688,8 @@ def test_indexing_purges_stale_persistent_caches(): assert all(tmpdir not in path for path in indexing._GATHER_MMAP_HANDLES) -def test_indexing_purge_tolerates_reentrant_sidecar_handle_cache_mutation(monkeypatch): +def test_purge_tolerates_reentrant_cache_change(monkeypatch): + """Purging survives a sidecar handle cache mutated re-entrantly mid-purge.""" import blosc2.indexing as indexing stale_scope = ("persistent", "/tmp/stale-index.b2nd") @@ -713,7 +714,7 @@ def mutating_exists(path): indexing._SIDECAR_HANDLE_CACHE.pop(injected_key, None) -def test_summary_index_compact_store_no_cross_column_confusion(tmp_path): +def test_summary_compact_no_cross_column_mixup(tmp_path): """Regression: a SUMMARY index on one column of a compact (.b2z) store must not be applied to a *different* column's predicate. @@ -756,7 +757,7 @@ class Aligned: assert got == expected, f"index returned {got}, expected {expected} (scan)" -def test_sidecar_handle_cache_no_cross_column_collision(tmp_path): +def test_sidecar_cache_no_cross_col_collision(tmp_path): """Regression: in a compact (.b2z) multi-column store, reading the SUMMARY block sidecar handle for each column must return *that* column's data, not a sibling's. @@ -900,7 +901,7 @@ def test_incremental_summary_matches_ooc_build(tmp_path): assert np.allclose(a["max"], b["max"], equal_nan=True) -def test_incremental_summary_invalidated_by_inplace_update(tmp_path): +def test_incremental_summary_stale_on_inplace(tmp_path): """An in-place column write before close must invalidate the accumulator so the builder falls back to a correct full rescan.""" f, i = _build_incr_data(n=4000) @@ -934,7 +935,7 @@ def test_granularity_only_valid_for_summary(): @pytest.mark.heavy @pytest.mark.parametrize("threshold", [5.0, 50.0, 99.0, 99.99]) -def test_summary_cost_gate_correctness_across_selectivity(threshold): +def test_summary_cost_gate_across_selectivity(threshold): """The SUMMARY cost gate may use the index (selective query) or fall back to a scan (broad query); both branches must return scan-correct results.""" t, _ = _make_gran_table(n=6000) @@ -1103,7 +1104,7 @@ def test_cross_column_or_prunes_segments_compact_b2z(tmp_path, monkeypatch): assert pruned, "cross-column OR fell back to a full scan instead of pruning" -def test_cross_column_predicates_match_scan_compact_b2z(tmp_path): +def test_cross_column_preds_match_scan_b2z(tmp_path): """Cross-column AND/OR over two SUMMARY-indexed columns must match the boolean-mask (no-index) result across selective, non-selective, empty, and mixed-direction predicates.""" @@ -1141,7 +1142,7 @@ def _seg_plan(units, *, base_nrows=1000, segment_len=250, level="block"): ) -def test_merge_segment_plans_intersection_union_and_fallback(): +def test_merge_segment_plans_and_fallback(): """Unit-level guard for the cross-column merge semantics.""" from blosc2.indexing import _merge_segment_plans @@ -1160,3 +1161,318 @@ def test_merge_segment_plans_intersection_union_and_fallback(): assert _merge_segment_plans(coarse, fine, "and") is fine # fine prunes more assert _merge_segment_plans(fine, coarse, "and") is fine assert _merge_segment_plans(coarse, fine, "or") is None + + +# --------------------------------------------------------------------------- +# String index summaries wider than 255 bytes +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("max_length", [16, 31, 32, 64, 100]) +@pytest.mark.parametrize("kind", ["summary", "bucket", "partial", "full", "opsi"]) +def test_string_index_matches_unindexed_scan(max_length, kind): + """An index must never change a query's answer. + + A segment summary is a ``(min, max, flags)`` record, so a ``= 'c10'", + "(name > 'c05') & (name <= 'c08')", + ] + expected = {q: sorted(int(v) for v in t.where(q)["x"][:]) for q in queries} + assert expected["name == 'c07'"], "fixture should match some rows" + + t.create_index(col_name="name", kind=blosc2.IndexKind(kind)) + for q in queries: + assert sorted(int(v) for v in t.where(q)["x"][:]) == expected[q], q + + +def test_wide_sidecar_span_read_is_not_short(): + """get_1d_span_numpy() must fill the whole destination, not part of it. + + A short read used to leave the tail uninitialised rather than raise. + """ + dtype = np.dtype([("min", " 255, "the point of this test is a capped typesize" + values = np.zeros(500, dtype=dtype) + values["min"] = [f"lo-{i:04d}" for i in range(500)] + values["max"] = [f"hi-{i:04d}" for i in range(500)] + values["flags"] = np.arange(500) % 251 + + arr = blosc2.asarray(values, chunks=(128,)) + out = np.empty(100, dtype=dtype) + arr.get_1d_span_numpy(out, 1, 5, 100) + assert out.tolist() == values[128 + 5 : 128 + 105].tolist() + + +def test_coalesce_spans_merges_within_a_block(): + """Spans closer than one block must merge: reading them apart re-reads the block.""" + coalesce = blosc2.indexing._coalesce_spans + spans = [(0, 100), (200, 300), (50_000, 50_100)] + assert coalesce(spans, 1024) == [(0, 300), (50_000, 50_100)] + assert coalesce(spans, 0) == spans # unknown block size → leave alone + assert coalesce([(0, 10)], 1024) == [(0, 10)] + # merged spans stay disjoint and ordered, so gathered positions stay unique + merged = coalesce([(0, 100), (10, 400), (401, 402)], 1024) + assert merged == [(0, 402)] + + +def test_bucket_gate_counts_blocks_not_buckets(): + """Selectivity in buckets overstates the saving; the read unit is the block.""" + frac = blosc2.indexing._bucket_block_fraction + geom = {"nav_segment_len": 16384, "bucket_len": 256} # 64 buckets per block + + scattered = np.zeros((1, 640), dtype=bool) + scattered[0, ::64] = True # 1.6% of buckets, but one in every block + assert frac(scattered, geom) == 1.0 + + clustered = np.zeros((1, 640), dtype=bool) + clustered[0, 0:64] = True # 10% of buckets, all inside one block + assert frac(clustered, geom) == 0.1 + + assert frac(np.zeros((1, 640), dtype=bool), geom) == 0.0 + + # A bucket at least as wide as a block covers whole blocks, so the fraction of + # blocks read is just the fraction of buckets selected -- not, as it once was, + # the fraction of *chunks* touched (or a flat 1.0 for a 1-D mask), both of which + # overstate the cost and decline plans the index exists to serve. + for bucket_len in (16384, 32768): # one block per bucket, and two + wide = {"nav_segment_len": 16384, "bucket_len": bucket_len} + selective = np.zeros((1, 10), dtype=bool) + selective[0, 0] = True + assert frac(selective, wide) == 0.1 + assert frac(selective[0], wide) == 0.1 # 1-D mask, same answer + + +def test_bucket_plan_gate_matches_block_fraction(): + """The planner must take a bucket plan only when it prunes actual blocks.""" + rng = np.random.default_rng(0) + n, card = 200_000, 5_000 + pool = sorted(f"v-{i:05d}" for i in range(card)) + + @dataclasses.dataclass + class Row: + c: str = blosc2.field(blosc2.string(max_length=8)) + v: float = blosc2.field(blosc2.float64()) + + def table(values, kind): + t = blosc2.CTable(Row) + t.extend({"c": values, "v": rng.random(n)}, validate=False) + if kind: + t.create_index("c", kind=kind) + return t + + query = f"(c >= '{pool[100]}') & (c < '{pool[120]}')" + values = [pool[i] for i in rng.integers(0, card, n)] + + seen = [] + original = blosc2.indexing._plan_single_exact_query + + def capture(exact_plan): + plan = original(exact_plan) + if plan.bucket_masks is not None: + fraction = blosc2.indexing._bucket_block_fraction( + plan.bucket_masks, exact_plan.descriptor["bucket"] + ) + seen.append((plan.usable, fraction)) + return plan + + blosc2.indexing._plan_single_exact_query = capture + try: + indexed = sorted(table(values, "bucket").where(query)["c"][:].tolist()) + finally: + blosc2.indexing._plan_single_exact_query = original + + assert seen, "no bucket plan was considered" + for usable, fraction in seen: + gate = blosc2.indexing._BUCKET_MAX_BLOCK_FRACTION + assert usable == (fraction <= gate), f"took={usable} at block fraction {fraction}" + + # Whichever way the gate goes, the answer is the same as an unindexed scan. + assert indexed == sorted(table(values, None).where(query)["c"][:].tolist()) + + +# --------------------------------------------------------------------------- +# Summary min()/max() shortcut: live rows only +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class MinMaxRow: + c: str = blosc2.field(blosc2.string(max_length=11)) + n: int = blosc2.field(blosc2.int64()) + f: float = blosc2.field(blosc2.float64()) + + +def _minmax_table(path, n, kind="summary"): + t = blosc2.CTable(MinMaxRow, urlpath=str(path), mode="w") + strings = [f"taxi-{i % 997:05d}" for i in range(n)] + t.extend({"c": strings, "n": np.arange(n) + 5, "f": (np.arange(n) + 5) * 1.5}) + if kind is not None: + for col in ("c", "n", "f"): + t.create_index(col, kind=kind) + return t, strings + + +# 16384 is the block length these columns get, so these straddle, exactly fill, +# and fall short of a block boundary respectively. +@pytest.mark.parametrize("n", [5, 1000, 16384, 16385, 100_000]) +def test_summary_minmax_ignores_capacity_padding(tmpdir, n): + """Padded slots hold 0/'' and must not be reported as the column minimum.""" + t, strings = _minmax_table(tmpdir / f"pad{n}.b2t", n) + assert len(t._valid_rows) > t._n_rows or n == len(t._valid_rows) # padding present + assert t["c"].min() == min(strings) + assert t["c"].max() == max(strings) + assert t["n"].min() == 5 + assert t["n"].max() == n + 4 + assert t["f"].min() == 7.5 + + +@pytest.mark.parametrize("n", [5, 1000, 16385, 100_000]) +def test_summary_minmax_matches_unindexed_scan(tmpdir, n): + indexed, _ = _minmax_table(tmpdir / f"i{n}.b2t", n) + scan, _ = _minmax_table(tmpdir / f"s{n}.b2t", n, kind=None) + for col in ("c", "n", "f"): + assert indexed[col].min() == scan[col].min() + assert indexed[col].max() == scan[col].max() + + +def test_summary_minmax_declines_after_delete(tmpdir): + """delete() leaves the index usable for queries but the deleted row still + sits in its block, so the summary shortcut must stand down.""" + t, _ = _minmax_table(tmpdir / "del.b2t", 100_000) + assert t["n"].min() == 5 + t.delete(0) # drop the unique minimum + assert t["n"].min() == 6 + t.delete(t._n_rows - 1) # drop the unique maximum (values now run 6..100003) + assert t["n"].max() == 100_003 + assert t["c"].min() == min(t["c"][:].tolist()) + + +def test_summary_minmax_declines_when_built_over_holes(tmpdir): + """A row deleted *before* the build is still in its block when the summary is + written, so a fresh index over a holey column must not enable the shortcut.""" + t, _ = _minmax_table(tmpdir / "prebuilt.b2t", 100_000, kind=None) + t.delete(slice(0, 1000)) # drop the 1000 smallest + for col in ("c", "n", "f"): + t.create_index(col, kind="summary") + assert t["n"]._summary_minmax_source() is None + assert t["n"].min() == 1005 + assert t["n"].max() == 100_004 + assert t["c"].min() == min(t["c"][:].tolist()) + + +def test_summary_minmax_shortcut_still_taken(tmpdir): + """The padding fix must not disable the shortcut on the common padded table.""" + t, _ = _minmax_table(tmpdir / "fast.b2t", 100_000) + assert t["n"]._summary_minmax_source() is not None + assert t["n"]._index_summary_minmax("min") is not NotImplemented + t.delete(0) + assert t["n"]._summary_minmax_source() is None + + +def test_summary_minmax_nullable_nan_float(tmpdir): + """A NaN-sentinel float is the one nullable column the shortcut accepts; + padding is 0.0 there, which is not NaN and would pass as a real value.""" + + @dataclasses.dataclass + class NanRow: + f: float = blosc2.field(blosc2.float64(nullable=True, null_value=float("nan"))) + + t = blosc2.CTable(NanRow, urlpath=str(tmpdir / "nan.b2t"), mode="w") + vals = (np.arange(50_000) + 5) * 1.5 + vals[:10] = np.nan # leading nulls + t.extend({"f": vals}) + t.create_index("f", kind="summary") + assert t["f"].min() == np.nanmin(vals) + assert t["f"].max() == np.nanmax(vals) + + +# --------------------------------------------------------------------------- +# Rank-indexed flavours accept kind="full" only +# --------------------------------------------------------------------------- + + +_needs_string_dtype = pytest.mark.skipif( + not hasattr(np.dtypes, "StringDType"), + reason="utf8 columns require NumPy >= 2.0 (StringDType)", +) + +if hasattr(np.dtypes, "StringDType"): + + @dataclasses.dataclass + class UTF8Row: + c: str = blosc2.field(blosc2.utf8()) + +else: + # blosc2.utf8() raises on NumPy < 2.0, so the class cannot even be defined + # there. Every parametrization using it carries _needs_string_dtype, so + # this placeholder is never dereferenced. + UTF8Row = None + + +@dataclasses.dataclass +class DictRow: + c: str = blosc2.field(blosc2.dictionary()) + + +#: The two flavours whose indexes are rank-based, utf8 skipped on NumPy 1.x. +RANK_FLAVOURS = [ + pytest.param(UTF8Row, "utf8", marks=_needs_string_dtype), + pytest.param(DictRow, "dictionary"), +] + + +@pytest.mark.parametrize(("row_cls", "flavour"), RANK_FLAVOURS) +@pytest.mark.parametrize("kind", ["summary", "bucket", "partial", "opsi"]) +def test_rank_index_rejects_non_full_kind(tmpdir, row_cls, flavour, kind): + """These build over the int32 ranks without error and are then never + consulted, so they must be refused rather than silently useless.""" + t = blosc2.CTable(row_cls, urlpath=str(tmpdir / f"{flavour}_{kind}.b2t"), mode="w") + t.extend({"c": [f"v{i % 50:03d}" for i in range(2000)]}) + with pytest.raises(ValueError, match=f"{flavour} column.*kind='full'"): + t.create_index("c", kind=kind) + assert "c" not in t._get_index_catalog() + + +@pytest.mark.parametrize(("row_cls", "flavour"), RANK_FLAVOURS) +def test_rank_index_accepts_full_kind(tmpdir, row_cls, flavour): + t = blosc2.CTable(row_cls, urlpath=str(tmpdir / f"{flavour}_full.b2t"), mode="w") + values = [f"v{i % 50:03d}" for i in range(2000)] + t.extend({"c": values}) + t.create_index("c", kind="full") + assert t._get_index_catalog()["c"]["kind"] == "full" + # and it still answers correctly through the rank path + assert sorted(t[t["c"] == "v007"]["c"][:]) == [v for v in values if v == "v007"] + + +@pytest.mark.parametrize(("row_cls", "flavour"), RANK_FLAVOURS) +def test_rank_index_default_kind_is_full(tmpdir, row_cls, flavour): + """The BUCKET default would hand these flavours an unusable index.""" + t = blosc2.CTable(row_cls, urlpath=str(tmpdir / f"{flavour}_def.b2t"), mode="w") + t.extend({"c": [f"v{i % 50:03d}" for i in range(2000)]}) + t.create_index("c") # no kind + assert t._get_index_catalog()["c"]["kind"] == "full" + + +def test_default_kind_unchanged_for_other_columns(tmpdir): + t = _make_table(200, persistent_path=str(tmpdir / "def.b2t")) + t.create_index("id") + assert t._get_index_catalog()["id"]["kind"] == "bucket" diff --git a/tests/ctable/test_ctable_ndarray_columns.py b/tests/ctable/test_ctable_ndarray_columns.py index bb7da1d1c..f3281ccff 100644 --- a/tests/ctable/test_ctable_ndarray_columns.py +++ b/tests/ctable/test_ctable_ndarray_columns.py @@ -40,7 +40,7 @@ def test_ndarray_column_metadata_and_tuple_indexing(): np.testing.assert_array_equal(t.image[:, :, :, 0], np.stack([np.ones((2, 2)), np.full((2, 2), 2)])) -def test_ndarray_column_comparison_and_scalar_operation_guards(): +def test_ndarray_col_comparison_scalar_guards(): t = table() with pytest.raises(TypeError, match="Cannot compare ndarray column 'embedding' directly"): @@ -55,7 +55,7 @@ def test_ndarray_column_comparison_and_scalar_operation_guards(): t.create_index("embedding") -def test_ndarray_column_axis_reductions_and_where_projection(): +def test_ndarray_col_axis_reductions_and_where(): t = table() assert t.embedding.sum() == np.float32(21) @@ -71,7 +71,8 @@ def test_ndarray_column_axis_reductions_and_where_projection(): np.testing.assert_array_equal(filtered.id[:], np.array([2], dtype=np.int32)) -def test_generated_column_row_transformer_append_refresh_and_vector_output(): +def test_generated_col_transformer_lifecycle(): + """Append, refresh, and a vector-returning transformer, in one lifecycle.""" t = table() t.add_generated_column( @@ -105,7 +106,7 @@ def test_generated_column_row_transformer_append_refresh_and_vector_output(): np.testing.assert_allclose(t.image_mean_rgb[:], t.image[:].mean(axis=(1, 2))) -def test_stale_generated_column_raises_and_read_stale_escape_hatch(): +def test_stale_generated_col_read_stale_hatch(): t = table() t.add_generated_column( "embedding_sum", @@ -141,7 +142,7 @@ class NullableNDArrayRow: codes: object = blosc2.field(blosc2.ndarray((2,), dtype=blosc2.int16(), nullable=True)) -def test_nullable_ndarray_columns_append_extend_assign_and_reduce(): +def test_nullable_ndarray_cols_write_and_reduce(): t = blosc2.CTable(NullableNDArrayRow) t.append((1, np.array([1, 2, 3], dtype=np.float32), [4, 5])) @@ -231,7 +232,7 @@ def test_nullable_ndarray_arrow_roundtrip(): np.testing.assert_array_equal(rt.codes.is_null(), t.codes.is_null()) -def test_ndarray_column_setitem_blosc2_ndarray_no_holes(): +def test_ndarray_col_setitem_b2_no_holes(): """col[:] = blosc2.NDArray fast path works for fixed-shape ndarray columns.""" n = 50 diff --git a/tests/ctable/test_ctable_take.py b/tests/ctable/test_ctable_take.py index 3d36d81d3..0b9281b49 100644 --- a/tests/ctable/test_ctable_take.py +++ b/tests/ctable/test_ctable_take.py @@ -36,7 +36,7 @@ def make_table(n=8): return t -def test_ctable_take_preserves_order_duplicates_and_negative_indices(): +def test_take_keeps_order_dups_and_negatives(): t = make_table(8) t.delete(2) t.delete(5) @@ -83,7 +83,7 @@ def test_ctable_take_handles_varlen_and_list_columns(): assert list(result["tags"][:]) == [[2, 20], [0], [2, 20], None] -def test_column_take_preserves_order_duplicates_and_negative_indices(): +def test_col_take_keeps_order_dups_and_negs(): t = make_table(8) t.delete(2) t.delete(5) @@ -152,7 +152,7 @@ def test_column_take_rejects_bad_indices(): col.take([4]) -def test_top_level_take_rejects_axis_for_ctable_and_column(): +def test_top_level_take_rejects_axis(): t = make_table(4) with pytest.raises(ValueError, match="axis"): @@ -193,7 +193,7 @@ def test_slice_copy_false_is_a_zero_copy_view(): np.testing.assert_array_equal(view["id"][:], np.arange(2, 6, dtype=np.int32)) -def test_slice_copy_true_is_an_independent_compact_table(): +def test_slice_copy_true_is_independent(): t = make_table(8) sub = t.slice(2, 6) # copy=True by default assert sub._cols is not t._cols diff --git a/tests/ctable/test_dictionary_column.py b/tests/ctable/test_dictionary_column.py index 542c2f268..644828271 100644 --- a/tests/ctable/test_dictionary_column.py +++ b/tests/ctable/test_dictionary_column.py @@ -10,6 +10,7 @@ from dataclasses import dataclass +import numpy as np import pytest import blosc2 @@ -214,7 +215,7 @@ class Row: assert ct.where('"Acme" in company and amount > 8')["amount"][:].tolist() == [9.0] assert ct.where('"Acme" in company or "Beta" in company').nrows == 3 - def test_string_where_dictionary_literal_with_special_chars(self): + def test_where_dict_literal_special_chars(self): # Literals with commas/spaces/dashes (e.g. chicago-taxi company names). @dataclass class Row: @@ -228,7 +229,7 @@ class Row: assert ct.where(f'company == "{name}"')["n"][:].tolist() == [1, 3] assert ct.where(f"company == '{name}'")["n"][:].tolist() == [1, 3] # single quotes too - def test_dictionary_predicate_combines_with_regular_predicate_in_aggregate(self): + def test_dict_pred_combines_in_aggregate(self): ct = CTable(TripRow) ct.extend(DATA_TUPLES) assert ct["fare"].sum(where=(ct["fare"] > 6) & (ct["vendor"] == "Uber")) == pytest.approx(25.5) @@ -494,7 +495,7 @@ def test_cli_preserves_dict_by_default(tmp_path): def test_cli_decode_dictionaries_flag(tmp_path): from blosc2.cli.parquet_to_blosc2 import main - from blosc2.schema import Utf8Spec, VLStringSpec + from blosc2.schema import UTF8Spec, VLStringSpec path = tmp_path / "dict.parquet" out = tmp_path / "dict_decoded.b2d" @@ -506,10 +507,10 @@ def test_cli_decode_dictionaries_flag(tmp_path): assert main(["--decode-dictionaries", str(path), str(out)]) == 0 ct = CTable.open(str(out), mode="r") - from blosc2.utf8_array import have_string_dtype + from blosc2._utf8_array import have_string_dtype # Decoded strings become utf8 columns on NumPy >= 2.0, vlstring on older NumPy. - expected_spec = Utf8Spec if have_string_dtype() else VLStringSpec + expected_spec = UTF8Spec if have_string_dtype() else VLStringSpec assert isinstance(ct._schema.columns_by_name["vendor"].spec, expected_spec) assert list(ct["vendor"][:]) == ["Uber", "Lyft", "Uber"] ct.close() @@ -538,5 +539,234 @@ def test_cli_dict_export_roundtrip(tmp_path): assert rt.column("score").to_pylist() == [1, 2, 3, 4] +def test_decode_reads_dictionary_once_not_per_row(): + """Decoding N rows must not index the dict store N times. + + Each ``dict_store[code]`` decompresses a whole msgpack batch, so a per-code + decode makes reads and lexsort-based ``sort_by`` cost O(N) decompressions. + """ + + @dataclass + class Row: + c: str = blosc2.field(blosc2.dictionary()) + + t = CTable(Row) + t.extend({"c": [f"v{i % 50}" for i in range(2000)]}, validate=False) + t._flush_varlen_columns() + + col = t._cols["c"] + col._invalidate_cache() + store = col._dict_store + original = type(store).__getitem__ + calls = 0 + + def counting_getitem(self, key): + nonlocal calls + calls += 1 + return original(self, key) + + type(store).__getitem__ = counting_getitem + try: + values = col[0:2000] + finally: + type(store).__getitem__ = original + + assert values == [f"v{i % 50}" for i in range(2000)] + assert calls <= 1, f"decoded 2000 rows with {calls} dict-store reads" + + if __name__ == "__main__": pytest.main(["-v", __file__]) + + +def test_dictionary_column_comparisons_are_elementwise(): + """``column == value`` must not fall through to object identity. + + It used to return a plain ``False`` — silently wrong rather than an error. + """ + import numpy as np + + @dataclass + class Row: + c: str = blosc2.field(blosc2.dictionary()) + + t = CTable(Row) + t.extend({"c": ["hello", "world", "hello"]}, validate=False) + t._flush_varlen_columns() + col = t._cols["c"] + + np.testing.assert_array_equal((col == "hello")[:3], [True, False, True]) + np.testing.assert_array_equal((col != "hello")[:3], [False, True, False]) + # A value absent from the dictionary matches nothing rather than raising. + np.testing.assert_array_equal((col == "absent")[:3], [False, False, False]) + # Defining __eq__ must not have made the container unhashable. + assert isinstance(hash(col), int) + + +def test_dictionary_ne_predicate_matches_live_rows(): + """``col != value`` must negate the value test, not the live-row mask. + + Negating afterwards turned every dead capacity slot True, which then failed + with an IndexError when used to select rows. + """ + import numpy as np + + @dataclass + class Row: + c: str = blosc2.field(blosc2.dictionary()) + + values = ["a1", "b2", "c3"] * 13 # 39 live rows in a padded slot array + t = CTable(Row) + t.extend({"c": values}, validate=False) + t._flush_varlen_columns() + + assert sorted(t[t["c"] != "a1"]["c"][:]) == sorted(v for v in values if v != "a1") + assert len(t[t["c"] == "a1"]["c"][:]) == 13 + # A value no row carries: nothing matches, everything differs. + assert len(t[t["c"] == "absent"]["c"][:]) == 0 + assert len(t[t["c"] != "absent"]["c"][:]) == len(values) + assert np.asarray((t["c"] != "a1")[:]).sum() == 26 + + +def test_dictionary_index_answers_equality(tmp_path): + """With a rank index, ``col == value`` is a sidecar lookup, not a codes scan.""" + + @dataclass + class Row: + c: str = blosc2.field(blosc2.dictionary()) + + values = ["pear", "apple", "cherry", "apple", "banana"] + results = {} + for tag in ("scan", "index"): + t = CTable(Row, urlpath=str(tmp_path / f"{tag}.b2t"), mode="w") + t.extend({"c": values}, validate=False) + t._flush_varlen_columns() + if tag == "index": + t.create_index("c", kind="full") + assert t["c"]._dictionary_index_mask("apple") is not None + # A value absent from the dictionary still answers, matching nothing. + assert not t["c"]._dictionary_index_mask("absent").any() + results[tag] = { + probe: ( + sorted(t[t["c"] == probe]["c"][:]), + sorted(t[t["c"] != probe]["c"][:]), + ) + for probe in ("apple", "pear", "absent") + } + del t + + assert results["index"] == results["scan"] + assert results["scan"]["apple"][0] == ["apple", "apple"] + + +def test_dictionary_index_spans_deleted_rows(tmp_path): + """The rank index has to cover the physical extent, not the live row count. + + delete() tombstones in place and only decrements the live count, so live + rows can sit past it. An index sized by that count silently drops them + from equality results. + """ + + @dataclass + class Row: + c: str = blosc2.field(blosc2.dictionary()) + + values = ["pear", "apple", "cherry", "apple", "banana", "apple"] + t = CTable(Row, urlpath=str(tmp_path / "t.b2t"), mode="w") + t.extend({"c": values}, validate=False) + t._flush_varlen_columns() + t.delete(0) + t.create_index("c", kind="full") + + assert t["c"]._dictionary_index_mask("apple") is not None + # The trailing "apple" is live and must not be lost to the index. + assert sorted(t[t["c"] == "apple"]["c"][:]) == ["apple"] * 3 + assert sorted(t[t["c"] != "apple"]["c"][:]) == ["banana", "cherry"] + + +def test_dict_rank_staleness_uses_value_epoch(tmp_path): + """The staleness check must not re-hash the whole dictionary per query.""" + from blosc2.ctable_indexing import _dict_rank_hash + + @dataclass + class Row: + c: str = blosc2.field(blosc2.dictionary()) + + t = CTable(Row, urlpath=str(tmp_path / "t.b2t"), mode="w") + t.extend({"c": [f"v{i % 50}" for i in range(500)]}, validate=False) + t._flush_varlen_columns() + t.create_index("c", kind="full") + meta = t._get_index_catalog()["c"]["full"]["dict_rank"] + + calls = 0 + import blosc2.ctable_indexing as ci + + def counting_hash(dictionary): + nonlocal calls + calls += 1 + return _dict_rank_hash(dictionary) + + ci._dict_rank_hash = counting_hash + try: + assert not t._dict_rank_index_stale("c", meta) + finally: + ci._dict_rank_hash = _dict_rank_hash + assert calls == 0, "value epoch was unchanged, so no hash should have been needed" + + +def test_sort_by_keys_on_ranks_not_decoded_strings(): + """sort_by must not decode a dictionary column to sort it. + + Sorting by alphabetical rank is sorting by decoded value, so the key can + stay int32 -- which skips both the decode and lexsort's string + comparisons. Correctness alone would not notice the difference, so + assert on the key dtype as well as the order. + """ + + @dataclass + class Row: + c: str = blosc2.field(blosc2.dictionary()) + x: int = blosc2.field(blosc2.int64()) + + values = ["delta", "alpha", "Zeta", "beta", "alpha"] + t = CTable(Row, new_data={"c": values, "x": list(range(len(values)))}) + + live = np.arange(len(values)) + keys = t._build_lex_keys(["c"], [True], live, len(values)) + assert keys[0].dtype == np.int32 + + # Ranks must order exactly as Python orders the decoded strings. + assert list(t.sort_by("c")["c"][:]) == sorted(values) + assert list(t.sort_by("c", ascending=False)["c"][:]) == sorted(values, reverse=True) + + +def test_sort_by_dictionary_nulls_and_multiple_keys(): + """Nulls sort last in both directions, and rank keys compose with others.""" + + @dataclass + class Row: + c: str = blosc2.field(blosc2.dictionary(nullable=True)) + x: int = blosc2.field(blosc2.int64()) + + values = ["b", None, "a", None, "b"] + t = CTable(Row, new_data={"c": values, "x": [0, 1, 2, 3, 4]}) + + assert list(t.sort_by("c")["c"][:]) == ["a", "b", "b", None, None] + assert list(t.sort_by("c", ascending=False)["c"][:]) == ["b", "b", "a", None, None] + # Secondary key breaks the "b" tie; nulls still trail. + assert list(t.sort_by(["c", "x"], [True, False])["x"][:]) == [2, 4, 0, 3, 1] + + +def test_sort_by_dictionary_view_and_small_copy_agree(): + """The filtered small-copy path builds its keys the same way sort_by does.""" + + @dataclass + class Row: + c: str = blosc2.field(blosc2.dictionary(nullable=True)) + x: int = blosc2.field(blosc2.int64()) + + values = ["b", None, "a", "c", "b", None, "a"] + t = CTable(Row, new_data={"c": values, "x": list(range(len(values)))}) + filtered = t[t.x > 1] # small enough to take _sorted_small_copy_from_live_positions + assert list(filtered.sort_by("c")["c"][:]) == ["a", "a", "b", "c", None] + assert list(filtered.sort_by("c", ascending=False)["c"][:]) == ["c", "b", "a", "a", None] diff --git a/tests/ctable/test_getitem_access.py b/tests/ctable/test_getitem_access.py index 85b664dd2..6ba76965e 100644 --- a/tests/ctable/test_getitem_access.py +++ b/tests/ctable/test_getitem_access.py @@ -32,7 +32,7 @@ class AccessRow: ] -def test_display_rows_printoption_truncates_to_five_head_and_tail_rows(): +def test_display_rows_truncates_head_and_tail(): previous = blosc2.get_printoptions() try: t = CTable(AccessRow, new_data=[(i, float(i), True, str(i), [i]) for i in range(60)]) @@ -66,7 +66,7 @@ def test_display_rows_printoption_truncates_to_five_head_and_tail_rows(): ) -def test_rename_column_recomputes_display_width_for_shorter_name(): +def test_rename_col_recomputes_display_width(): @dataclass class WidthRow: very_long_temporary_name: float = blosc2.field(blosc2.float64()) @@ -80,7 +80,7 @@ class WidthRow: assert t._col_widths["x"] == max(len("x"), t._schema.columns_by_name["x"].display_width) -def test_display_precision_printoption_formats_float_values(): +def test_display_precision_formats_floats(): previous = blosc2.get_printoptions() try: t = CTable(AccessRow, new_data=[(1, 1.23456789, True, "x", [1])]) @@ -232,7 +232,7 @@ def test_getitem_slice_returns_view(): assert sub.base is t -def test_getitem_integer_list_and_bool_mask_return_views(): +def test_getitem_int_list_and_bool_mask_views(): t = CTable(AccessRow, new_data=DATA) gathered = t[[3, 0, 2]] assert isinstance(gathered, CTable) @@ -288,7 +288,7 @@ def test_getitem_non_boolean_expression_raises(): _ = t["id + 1"] -def test_ctable_array_materialization_uses_structured_dtype(): +def test_array_materialization_structured(): t = CTable(AccessRow, new_data=DATA) arr = np.asarray(t) assert arr.dtype.fields is not None diff --git a/tests/ctable/test_groupby.py b/tests/ctable/test_groupby.py index 3830d7efb..9d1a76620 100644 --- a/tests/ctable/test_groupby.py +++ b/tests/ctable/test_groupby.py @@ -76,7 +76,7 @@ def test_groupby_agg_numeric_reductions(): assert got[2] == ("Rome", 60.0, 30.0, 20.0, 40.0, 2) -def test_groupby_argmin_argmax_return_logical_positions(): +def test_groupby_argmin_argmax_logical_pos(): t = CTable(SalesRow, new_data=DATA) out = t.group_by("city", sort=True).agg({"sales": ["argmin", "argmax"]}) @@ -85,7 +85,7 @@ def test_groupby_argmin_argmax_return_logical_positions(): assert rows(out) == [("Berlin", -1, -1), ("Paris", 0, 3), ("Rome", 2, 4)] -def test_groupby_argmin_argmax_convenience_methods_and_view_positions(): +def test_groupby_argmin_methods_and_view_pos(): t = CTable(SalesRow, new_data=DATA) view = t.where("qty >= 3") @@ -161,7 +161,7 @@ class DictRow: sales: int = blosc2.field(blosc2.int32()) -def test_groupby_dictionary_key_groups_by_decoded_value(): +def test_groupby_dict_key_groups_by_value(): t = CTable(DictRow, new_data=[("Paris", 10), ("Rome", 20), ("Paris", 30)]) out = t.group_by("city", sort=True).agg({"sales": "sum"}) @@ -170,7 +170,7 @@ def test_groupby_dictionary_key_groups_by_decoded_value(): assert rows(out) == [("Paris", 40), ("Rome", 20)] -def test_groupby_dictionary_key_sorted_by_string_not_code_order(): +def test_groupby_dict_key_sorted_by_string(): """Dict groups come out alphabetical even when codes are assigned otherwise. Regression for the always-sorted contract: with "Rome" seen before "Paris" @@ -188,7 +188,7 @@ def test_groupby_dictionary_key_sorted_by_string_not_code_order(): assert rows(out) == [("Paris", 1, 3), ("Rome", 2, 0)] -def test_groupby_dictionary_key_sorted_matches_python_sorted(): +def test_groupby_dict_key_matches_python_sort(): """Vectorized dict-key ordering matches a Python sorted() reference.""" rng = np.random.default_rng(0) labels = [f"city_{i:03d}" for i in range(200)] @@ -206,7 +206,7 @@ def test_groupby_string_key_sorted_without_sort_flag(): assert [r[0] for r in rows(out)] == ["Berlin", "Paris", "Rome"] -def test_groupby_dictionary_key_argmin_argmax_positions(): +def test_groupby_dict_key_argmin_argmax_pos(): # Dictionary key drives the dense-position fast path; verify it returns the # logical row positions of the extremes (chicago-taxi "company" shape). t = CTable(DictRow, new_data=[("Paris", 10), ("Rome", 50), ("Paris", 30), ("Rome", 20)]) @@ -218,7 +218,7 @@ def test_groupby_dictionary_key_argmin_argmax_positions(): assert rows(out) == [("Paris", 0, 2), ("Rome", 3, 1)] -def test_groupby_dictionary_key_beyond_default_code_capacity(): +def test_groupby_dict_key_beyond_capacity(): data = [("Paris" if i % 2 == 0 else "Rome", 1) for i in range(5000)] t = CTable(DictRow, new_data=data) @@ -313,7 +313,7 @@ def test_groupby_fast_path_sum_variants(row_type, data, expected): assert rows(out) == expected -def test_groupby_float_integral_fast_path_falls_back_for_non_integral_keys(): +def test_groupby_float_path_falls_back_fractional(): t = CTable(Float64KeyRow, new_data=[(0.5, 1.0), (1.5, 2.0), (0.5, 3.0)]) # Float keys are not key-sorted by default (sort=None); request sort=True to @@ -323,7 +323,7 @@ def test_groupby_float_integral_fast_path_falls_back_for_non_integral_keys(): assert rows(out) == [(0.5, 4.0), (1.5, 2.0)] -def test_groupby_float_integral_fast_path_falls_back_for_nan_group_when_kept(): +def test_groupby_float_path_falls_back_nan_group(): t = CTable(Float64KeyRow, new_data=[(0.0, 1.0), (np.nan, 2.0), (0.0, 3.0)]) out = t.group_by("key", dropna=False).agg({"value": "sum"}) @@ -349,7 +349,7 @@ def test_groupby_integral_float_key_dense_min_max(row_type): assert out_max._cols["key"][:].dtype == t._cols["key"][:].dtype -def test_groupby_integral_float_key_falls_back_for_negative_keys(): +def test_groupby_float_key_falls_back_negative(): # Negative keys cannot use the dense (non-negative) mapping; the generic # path must still produce correct max results. t = CTable(Float64KeyRow, new_data=[(-1.0, 5.0), (-1.0, 8.0), (2.0, 3.0)]) @@ -368,7 +368,7 @@ def test_group_reduce_object_keys_sort_with_none(): assert sizes.tolist() == [1, 1, 2] -def test_group_reduce_object_numeric_keys_sort_with_none(): +def test_group_reduce_numeric_keys_with_none(): groups, sizes = blosc2.group_reduce(np.array([None, 2, 1, 2], dtype=object), sort=True, dropna=False) assert groups.tolist() == [None, 1, 2] @@ -474,7 +474,7 @@ def test_groupby_cython_integer_key_more_integer_aggs(): assert rows(out) == [(0, 2, 2, 3, 1.5, -2, 5), (1, 2, 2, 30, 15.0, 10, 20), (2, 1, 1, 7, 7.0, 7, 7)] -def test_groupby_cython_integer_key_nullable_float_aggs(): +def test_groupby_cython_int_key_null_aggs(): row_type = make_dataclass( "IntKeyNullableFloatAggsRow", [ @@ -514,7 +514,7 @@ def test_groupby_cython_arbitrary_float_key_aggs(): ] -def test_groupby_cython_arbitrary_float_key_nan_and_signed_zero(): +def test_groupby_cython_float_key_nan_and_zero(): t = CTable(Float64KeyRow, new_data=[(-0.0, 1.0), (0.0, 2.0), (np.nan, 3.0), (np.nan, 4.0)]) dropped = t.group_by("key").agg({"value": "sum"}) @@ -595,7 +595,7 @@ def test_groupby_persistent_output_urlpath(tmp_path): assert rows(reopened) == [("Berlin", 6), ("Paris", 7), ("Rome", 8)] -def test_groupby_persistent_output_urlpath_on_convenience_method(tmp_path): +def test_groupby_persistent_urlpath_shorthand(tmp_path): t = CTable(SalesRow, new_data=DATA) path = tmp_path / "grouped_mean.b2d" @@ -621,7 +621,7 @@ def _keys(out): return [out._cols[out.col_names[0]][i] for i in range(out.nrows)] -def test_groupby_int_key_always_ascending_regardless_of_sort(): +def test_groupby_int_key_always_ascending(): # Integer/dense keys come out ascending under every sort= value -- nonzero # ordering is free and unavoidable. t = CTable(Int32FloatRow, new_data=_INT_SORT_DATA) @@ -638,7 +638,7 @@ def test_groupby_dict_key_sorted_under_auto_and_true(): assert _keys(t.group_by("key", sort=False).sum("value")) == ["zeta", "alpha", "mike"] -def test_groupby_float_key_unsorted_under_auto_sorted_under_true(): +def test_groupby_float_key_auto_vs_sorted(): # Float keys only sort via a Python list.sort, so None (auto) leaves them # unsorted; True sorts. The unsorted order must be deterministic across runs. t = CTable(Float64KeyRow, new_data=_FLOAT_SORT_DATA) @@ -649,7 +649,7 @@ def test_groupby_float_key_unsorted_under_auto_sorted_under_true(): assert sorted(auto1) == [1.5, 2.5, 3.5] # same groups, order unspecified -def test_groupby_multikey_unsorted_under_auto_sorted_under_true(): +def test_groupby_multikey_auto_vs_sorted(): # Multi-key results only sort via a Python list.sort, so None (auto) leaves # them unsorted (deterministic but unspecified order); True sorts. data = [("z", 2, 1.0), ("a", 1, 2.0), ("z", 1, 3.0), ("a", 1, 4.0)] @@ -915,7 +915,7 @@ def inconsistent(values): g.agg(x=("sales", inconsistent)) -def test_agg_udf_unsupported_result_dtype_raises_clear_error(): +def test_agg_udf_bad_result_dtype_raises(): t = CTable(SalesRow, new_data=DATA) g = t.group_by("city") @@ -1065,7 +1065,7 @@ def test_factorize_fixed_width_str_matches_np_unique(): np.testing.assert_array_equal(got_inv, ref_inv) -def test_factorize_fixed_width_str_collision_falls_back(monkeypatch): +def test_factorize_str_collision_falls_back(monkeypatch): """With the mix constant forced to 0, the row hash degenerates to the last uint32 word, so strings differing only in earlier characters collide -- the verify pass must detect it and fall back to exact np.unique.""" diff --git a/tests/ctable/test_nested_access_storage.py b/tests/ctable/test_nested_access_storage.py index 485d305a6..5b35b2486 100644 --- a/tests/ctable/test_nested_access_storage.py +++ b/tests/ctable/test_nested_access_storage.py @@ -25,7 +25,7 @@ class PersistRow: a: int -def test_dotted_column_attribute_namespace_and_where_string(): +def test_dotted_col_attribute_and_where_string(): t = blosc2.CTable(AccessRow) t.append((1.0, 10.0)) t.append((2.0, 30.0)) @@ -44,7 +44,7 @@ def test_dotted_column_attribute_namespace_and_where_string(): assert view2.nrows == 2 -def test_dotted_column_persists_under_hierarchical_cols(tmp_path): +def test_dotted_col_persists_hierarchical(tmp_path): t = blosc2.CTable(PersistRow) t.append((1,)) t.rename_column("a", "trip.begin.lon") @@ -69,7 +69,7 @@ def test_select_struct_prefix_expands_descendants(): assert s.col_names == ["trip.begin.lon"] -def test_from_arrow_flattens_struct_columns_to_dotted_leaves(): +def test_from_arrow_flattens_struct_to_dotted(): trip_type = pa.struct([("begin", pa.struct([("lon", pa.float64()), ("lat", pa.float64())]))]) schema = pa.schema([pa.field("trip", trip_type)]) batch = pa.record_batch( @@ -107,7 +107,7 @@ def test_from_arrow_flattens_struct_columns_to_dotted_leaves(): row0["nope"] -def test_nested_field_name_escaping_for_literal_dot_and_slash(tmp_path): +def test_field_name_escaping_dot_and_slash(tmp_path): trip_type = pa.struct([pa.field("begin/point", pa.struct([pa.field("lon.deg", pa.float64())]))]) schema = pa.schema([pa.field("trip.info", trip_type)]) batch = pa.record_batch( diff --git a/tests/ctable/test_nested_metadata_root.py b/tests/ctable/test_nested_metadata_root.py index 5a6989df1..5b82794e2 100644 --- a/tests/ctable/test_nested_metadata_root.py +++ b/tests/ctable/test_nested_metadata_root.py @@ -18,7 +18,7 @@ def _table_with_empty_root_alias(): return blosc2.CTable.from_arrow(schema, [batch]) -def test_schema_version_2_with_nested_metadata_roundtrip(): +def test_schema_v2_nested_metadata_roundtrip(): schema = pa.schema([pa.field("x.y", pa.float64())]) batch = pa.record_batch([pa.array([1.0, 2.0])], schema=schema) t = blosc2.CTable.from_arrow(schema, [batch]) @@ -31,13 +31,13 @@ def test_schema_version_2_with_nested_metadata_roundtrip(): assert restored.metadata["nested"]["logical_to_physical"]["x.y"] == "x.y" -def test_empty_root_metadata_exports_back_to_empty_arrow_name(): +def test_empty_root_exports_empty_arrow_name(): t = _table_with_empty_root_alias() out = t.to_arrow() assert out.schema.names == [""] -def test_empty_root_logical_alias_getitem_select_and_index(): +def test_empty_root_alias_getitem_and_select(): t = _table_with_empty_root_alias() assert t[""][0] == 1.0 s = t.select([""]) diff --git a/tests/ctable/test_null_expressions.py b/tests/ctable/test_null_expressions.py index c17dc2400..a97ce22a7 100644 --- a/tests/ctable/test_null_expressions.py +++ b/tests/ctable/test_null_expressions.py @@ -64,7 +64,7 @@ def test_eq_sentinel_literal_does_not_match_null(): assert t[t.score == NULL_I64]["id"][:].tolist() == [] -def test_ne_sentinel_literal_does_not_match_null_either(): +def test_ne_sentinel_literal_no_null_match(): t = CTable(IntRow, new_data=[(1, 10, 0), (2, NULL_I64, 0)]) # A null never satisfies `!=` either — it isn't "not equal", it's unknown. assert t[t.score != NULL_I64]["id"][:].tolist() == [1] @@ -75,7 +75,7 @@ def test_is_null_still_finds_nulls(): assert list(t.score.is_null()) == [False, True] -def test_comparison_between_two_nullable_columns_excludes_either_null(): +def test_two_nullable_cols_exclude_either_null(): t = CTable( IntRow, new_data=[ @@ -94,7 +94,7 @@ def test_ge_le_also_exclude_nulls(): assert t[t.score <= -20]["id"][:].tolist() == [3] -def test_comparison_against_nan_scalar_does_not_crash_and_matches_nothing(): +def test_nan_scalar_comparison_matches_nothing(): """Regression: ``t.f == np.nan`` used to crash with NameError inside the lazyexpr evaluator (the scalar was embedded as the bare literal ``nan``). Now it evaluates -- and matches nothing, since a null satisfies no @@ -236,7 +236,7 @@ def test_reduction_on_derived_expression_skips_nulls(): assert (t.score + 1).std() == pytest.approx(15.0) -def test_reduction_on_chained_and_mixed_expressions_skips_nulls(): +def test_chained_expr_reduction_skips_nulls(): t = CTable(IntRow, new_data=[(1, 10, 5), (2, NULL_I64, 5), (3, -20, NULL_I64)]) assert ((t.score + 1) * 2).sum() == pytest.approx(2 * (11 - 19)) # nullable + nullable: null wherever either operand is null -> only row 1 live @@ -248,7 +248,7 @@ def test_reduction_on_chained_and_mixed_expressions_skips_nulls(): assert (t.score**0).sum() == pytest.approx(2.0) # nan**0 must not resurrect the null -def test_reduction_on_derived_expression_matches_pandas(): +def test_derived_expr_reduction_vs_pandas(): pd = pytest.importorskip("pandas") t = CTable(IntRow, new_data=[(1, 10, 0), (2, NULL_I64, 0), (3, -20, 0), (4, 7, 0)]) s = pd.Series([10, None, -20, 7], dtype="Int64") @@ -256,7 +256,7 @@ def test_reduction_on_derived_expression_matches_pandas(): assert (t.score + 1).mean() == pytest.approx(float((s + 1).mean())) -def test_derived_expression_reductions_respect_deleted_rows_and_views(): +def test_derived_expr_respects_deletes_and_views(): t = CTable(IntRow, new_data=[(1, 10, 0), (2, NULL_I64, 0), (3, -20, 0), (4, 7, 0)]) t.delete([0]) # drop score=10 assert (t.score + 1).sum() == pytest.approx(-19 + 8) @@ -264,7 +264,7 @@ def test_derived_expression_reductions_respect_deleted_rows_and_views(): assert (view.score + 1).sum() == pytest.approx(8.0) -def test_derived_expression_all_null_reduction_semantics(): +def test_derived_expr_all_null_reduction(): t = CTable(IntRow, new_data=[(1, NULL_I64, 0), (2, NULL_I64, 0)]) assert (t.score + 1).sum() == 0.0 # same convention as Column.sum() assert math.isnan((t.score + 1).mean()) @@ -274,7 +274,7 @@ def test_derived_expression_all_null_reduction_semantics(): (t.score + 1).max() -def test_derived_expression_ne_comparison_excludes_nulls(): +def test_derived_expr_ne_excludes_nulls(): t = CTable(IntRow, new_data=[(1, 10, 0), (2, NULL_I64, 0), (3, -20, 0)]) assert t[(t.score + 1) != 11]["id"][:].tolist() == [3] assert t[(t.score + 1) > 0]["id"][:].tolist() == [1] diff --git a/tests/ctable/test_nullable.py b/tests/ctable/test_nullable.py index ee5cb807d..779fe34c3 100644 --- a/tests/ctable/test_nullable.py +++ b/tests/ctable/test_nullable.py @@ -78,7 +78,7 @@ def test_null_value_property_set(): assert t["score"].null_value == -1 -def test_numpy_nan_null_value_skips_scalar_validation_constraints(): +def test_nan_null_value_skips_validation(): @dataclass class NumpyNaNFloatRow: value: float = blosc2.field(blosc2.float32(ge=0, null_value=np.float32(np.nan))) @@ -127,7 +127,7 @@ class Row: assert t["b"].dtype.itemsize >= len(b"__BLOSC2_NULL__") -def test_nullable_true_uses_null_policy_context_and_column_null_values(): +def test_nullable_uses_policy_and_col_nulls(): @dataclass class Row: i: int = blosc2.field(blosc2.int32(nullable=True)) @@ -163,7 +163,7 @@ def test_add_column_nullable_true_uses_null_policy(): assert t["extra"].null_value == np.iinfo(np.int32).max -def test_nullable_policy_rejects_out_of_range_integer_sentinel(): +def test_policy_rejects_out_of_range_sentinel(): @dataclass class Row: x: int = blosc2.field(blosc2.int8(nullable=True)) @@ -173,7 +173,7 @@ class Row: CTable(Row) -def test_nullable_policy_rejects_wrong_string_sentinel_type(): +def test_policy_rejects_wrong_sentinel_type(): @dataclass class Row: s: str = blosc2.field(blosc2.string(nullable=True)) @@ -636,7 +636,7 @@ def test_all_nulls_value_counts_empty(): assert len(vc) == 0 -def test_null_value_does_not_affect_non_nullable_column(): +def test_null_value_ignored_if_not_nullable(): t = CTable(IntRow, new_data=[(1, 10), (2, 20)]) # id column has no null_value — aggregates work normally assert t["id"].sum() == 3 diff --git a/tests/ctable/test_object_spec.py b/tests/ctable/test_object_spec.py index 9b6154dc6..0a072a41f 100644 --- a/tests/ctable/test_object_spec.py +++ b/tests/ctable/test_object_spec.py @@ -59,7 +59,7 @@ class StrictObjectRow: t.append([None]) -def test_object_column_rejects_non_msgpack_value_on_flush(): +def test_object_col_rejects_non_msgpack(): t = CTable(ObjectRow) t.append([1, {"not-msgpack": {1, 2, 3}}]) with pytest.raises(TypeError): diff --git a/tests/ctable/test_parquet_interop.py b/tests/ctable/test_parquet_interop.py index 16676c74a..908becedb 100644 --- a/tests/ctable/test_parquet_interop.py +++ b/tests/ctable/test_parquet_interop.py @@ -267,7 +267,7 @@ class StructRow: reopened = CTable.open(str(path), mode="r") assert reopened["props"][:] == [{"a": 1, "b": "x"}, None, {"a": 2, "b": "yy"}] - def test_from_arrow_object_fallback_for_unsupported_type(self): + def test_from_arrow_object_fallback(self): map_type = pa.map_(pa.string(), pa.int32()) batch = pa.record_batch( [pa.array([[("a", 1)], None, [("b", 2), ("c", 3)]], type=map_type)], names=["attrs"] @@ -408,7 +408,7 @@ def test_from_arrow_blosc2_batch_size_default(self): assert t["vals"][0] == [1] assert t["vals"][1] == [2, 3] - def test_from_arrow_blosc2_batch_size_override_and_none(self): + def test_from_arrow_batch_size_override(self): at = pa.table({"vals": pa.array([[1], [2], [3]], type=pa.list_(pa.int64()))}) t = CTable.from_arrow(at.schema, at.to_batches(max_chunksize=1), blosc2_batch_size=2) assert t._schema.columns_by_name["vals"].spec.batch_rows == 2 @@ -576,7 +576,7 @@ def test_null_policy_controls_default_sentinels(self): assert t._schema.columns_by_name["i"].spec.null_value == np.iinfo(np.int32).max assert t["i"].null_count() == 1 - def test_null_policy_string_value_applies_to_fixed_width_strings(self): + def test_null_policy_applies_to_fixed_width(self): """string_value in NullPolicy applies when string_max_length is given explicitly.""" at = pa.table( { @@ -613,14 +613,14 @@ def test_null_values_override_policy_and_auto_false(self, tmp_path): t2 = CTable.from_parquet(path, auto_null_sentinels=False) assert t2._schema.columns_by_name["i"].spec.null_value == -1 - def test_null_policy_rejects_vlbytes_column_null_values(self): + def test_null_policy_rejects_vlbytes_nulls(self): """Passing column_null_values for a vlbytes column raises TypeError.""" at = pa.table({"b": pa.array([b"a", None, b"c"], type=pa.large_binary())}) policy = blosc2.NullPolicy(column_null_values={"b": b"NA"}) with blosc2.null_policy(policy), pytest.raises(TypeError, match="vlbytes"): CTable.from_arrow(at.schema, at.to_batches()) - def test_null_policy_column_null_values_applies_to_utf8(self): + def test_null_policy_col_nulls_apply_to_utf8(self): """Passing column_null_values for a utf8 (scalar string) column sets its sentinel. On NumPy < 2.0 utf8 columns are unavailable, strings import as @@ -722,7 +722,7 @@ def test_max_rows_from_parquet_limits_rows(self, tmp_path): assert len(out) == 6 np.testing.assert_array_equal(out["id"][:], np.arange(6)) - def test_max_rows_zero_from_parquet_imports_empty_table(self, tmp_path): + def test_max_rows_zero_imports_empty_table(self, tmp_path): t = CTable(Row, new_data=DATA10) path = tmp_path / "x.parquet" t.to_parquet(path) @@ -784,7 +784,7 @@ def test_parquet_cli_nested_progress_skips_write_lines(tmp_path, capsys): assert " write" not in captured.out -def test_parquet_cli_separate_nested_flattens_top_level_structs(tmp_path, capsys): +def test_cli_separate_nested_flattens_structs(tmp_path, capsys): from blosc2.cli.parquet_to_blosc2 import main trip_type = pa.struct( @@ -819,7 +819,8 @@ def test_parquet_cli_separate_nested_flattens_top_level_structs(tmp_path, capsys ct.close() -def test_parquet_cli_no_separate_nested_preserves_top_level_struct_as_list(tmp_path): +def test_cli_no_separate_nested_keeps_struct(tmp_path): + """Without --separate-nested a top-level struct stays one list column.""" from blosc2.cli.parquet_to_blosc2 import main trip_type = pa.struct([pa.field("sec", pa.float32())]) diff --git a/tests/ctable/test_schema_mutations.py b/tests/ctable/test_schema_mutations.py index b8e7f99bf..b7a9e6319 100644 --- a/tests/ctable/test_schema_mutations.py +++ b/tests/ctable/test_schema_mutations.py @@ -114,7 +114,7 @@ def test_view_blocks_assign(): assert t["score"][5] == pytest.approx(50.0) -def test_take_from_view_yields_independent_writable_table(): +def test_take_from_view_is_independent(): t = CTable(Row, new_data=DATA10) view = t.where(t["id"] > 4) independent = view.take([0, 1]) @@ -190,7 +190,7 @@ def test_blosc2_open_raw_treestore_without_manifest(): assert np.array_equal(opened["/group/node"][:], np.arange(5)) -def test_blosc2_open_raw_treestore_for_unknown_manifest_kind(): +def test_open_raw_treestore_unknown_manifest(): path = table_path("unknown_manifest") with blosc2.TreeStore(path, mode="w", threshold=0) as tstore: meta = blosc2.SChunk() @@ -204,7 +204,7 @@ def test_blosc2_open_raw_treestore_for_unknown_manifest_kind(): assert np.array_equal(opened["/payload"][:], np.arange(3)) -def test_extensionless_ctable_path_uses_extensionless_store(): +def test_extensionless_path_uses_that_store(): path = os.path.join(TABLE_ROOT, "alias_ctable") t = CTable(Row, urlpath=path, mode="w", new_data=DATA10) t.close() @@ -271,14 +271,14 @@ def test_add_column_fills_default_for_existing_rows(): np.testing.assert_array_equal(t["weight"][:], np.full(10, 5.5)) -def test_add_column_without_default_allowed_for_empty_table(): +def test_add_col_no_default_ok_when_empty(): t = CTable(Row) t.add_column("weight", blosc2.float64()) t.append((1, 2.0, True, 3.0)) assert t["weight"][0] == pytest.approx(3.0) -def test_add_column_without_default_on_non_empty_table_raises(): +def test_add_col_no_default_raises_non_empty(): t = CTable(Row, new_data=DATA10) with pytest.raises(ValueError, match="requires a default"): t.add_column("weight", blosc2.float64()) @@ -342,6 +342,132 @@ def test_add_column_skips_deleted_rows(): assert all(v == 3.0 for v in vals) +# =========================================================================== +# add_column(values=) +# =========================================================================== + + +def test_add_column_values_fills_live_rows(): + t = CTable(Row, new_data=DATA10) + t.add_column("weight", blosc2.float64(), values=np.arange(10, dtype=np.float64)) + np.testing.assert_array_equal(t["weight"][:], np.arange(10, dtype=np.float64)) + + +def test_add_column_values_needs_no_default(): + """values= is the second way to satisfy a non-empty table.""" + t = CTable(Row, new_data=DATA10) + t.add_column("weight", blosc2.float64(), values=[1.0] * 10) + assert t["weight"][0] == pytest.approx(1.0) + + +def test_add_column_values_coerced_to_spec_dtype(): + t = CTable(Row, new_data=DATA10) + t.add_column("n", blosc2.int8(), values=list(range(10))) + assert t["n"][:].dtype == np.int8 + + +def test_add_column_values_wrong_length_raises(): + t = CTable(Row, new_data=DATA10) + with pytest.raises(ValueError, match="requires 10 entries"): + t.add_column("weight", blosc2.float64(), values=[1.0, 2.0]) + + +def test_add_column_values_uncoercible_raises(): + t = CTable(Row, new_data=DATA10) + with pytest.raises(TypeError, match="Cannot coerce values="): + t.add_column("n", blosc2.int8(), values=["nope"] * 10) + + +def test_add_column_values_enforces_declared_constraints(): + """values= must not slip past the constraints the spec declares. + + Coercing to a fixed-width dtype truncates an over-long string instead of + complaining, so an unchecked values= would silently drop characters -- + the same check runs for numeric bounds, hence both cases here. + """ + t = CTable(Row, new_data=DATA10) + with pytest.raises(ValueError, match="exceeds max_length=4"): + t.add_column("code", blosc2.string(max_length=4), values=["toolongvalue"] * 10) + with pytest.raises(ValueError, match="violates constraint le="): + t.add_column("bounded", blosc2.int64(le=100), values=[999] * 10) + # A value that does fit is still accepted, uncut. + t.add_column("code", blosc2.string(max_length=4), values=["abcd"] * 10) + assert list(t["code"][:]) == ["abcd"] * 10 + + +def test_add_column_values_skips_deleted_rows(): + """values= is positional over *live* rows, not physical slots.""" + t = CTable(Row, new_data=DATA10) + t.delete([0, 1]) # 8 live rows + t.add_column("weight", blosc2.float64(), values=np.arange(8, dtype=np.float64)) + np.testing.assert_array_equal(t["weight"][:], np.arange(8, dtype=np.float64)) + np.testing.assert_array_equal(t["id"][:], np.arange(2, 10)) + + +def test_add_col_values_keeps_default_later(): + t = CTable(Row, new_data=DATA10) + t.add_column("weight", blosc2.field(blosc2.float64(), default=9.0), values=[1.0] * 10) + t.append((10, 0.0, True, 0.0)) + np.testing.assert_array_equal(t["weight"][:], [*([1.0] * 10), 0.0]) + + +def test_add_column_values_ndarray_column(): + t = CTable(Row, new_data=DATA10) + vals = np.arange(20, dtype=np.float32).reshape(10, 2) + t.add_column("v", blosc2.ndarray((2,), np.float32), values=vals) + np.testing.assert_array_equal(t["v"][:], vals) + + +def test_add_column_values_ndarray_bad_shape_raises(): + t = CTable(Row, new_data=DATA10) + with pytest.raises(ValueError, match=r"must have shape \(10, 2\)"): + t.add_column("v", blosc2.ndarray((2,), np.float32), values=np.zeros(10, dtype=np.float32)) + + +def test_add_column_values_persists_on_disk(): + path = table_path("add_col_values") + t = CTable(Row, urlpath=path, mode="w", new_data=DATA10) + t.add_column("weight", blosc2.float64(), values=np.arange(10, dtype=np.float64)) + t.close() + t2 = CTable.open(path) + np.testing.assert_array_equal(t2["weight"][:], np.arange(10, dtype=np.float64)) + + +def test_add_column_values_vlstring(): + t = CTable(Row, new_data=DATA10) + vals = [f"s{i}" for i in range(10)] + t.add_column("s", blosc2.vlstring(), values=vals) + assert list(t["s"][:]) == vals + + +def test_add_column_values_vlstring_skips_deleted_rows(): + """Varlen columns are indexed physically, so the dead slots need filling too.""" + t = CTable(Row, new_data=DATA10) + t.delete([0, 1]) + vals = [f"s{i}" for i in range(8)] + t.add_column("s", blosc2.vlstring(), values=vals) + assert list(t["s"][:]) == vals + + +def test_add_col_vlstring_skips_deleted_rows(): + t = CTable(Row, new_data=DATA10) + t.delete([0, 1]) + t.add_column("s", blosc2.field(blosc2.vlstring(), default="z")) + assert list(t["s"][:]) == ["z"] * 8 + + +def test_add_column_values_list_column_raises(): + t = CTable(Row, new_data=DATA10) + with pytest.raises(TypeError, match="does not support list columns"): + t.add_column("l", blosc2.list(blosc2.int64()), values=[[1]] * 10) + + +def test_add_column_values_dictionary_column_raises(): + t = CTable(Row, new_data=DATA10) + with pytest.raises(TypeError, match="does not support dictionary columns"): + t.add_column("c", blosc2.dictionary(), values=["a"] * 10) + + # =========================================================================== # drop_column # =========================================================================== diff --git a/tests/ctable/test_schema_specs.py b/tests/ctable/test_schema_specs.py index d78227ef7..349999320 100644 --- a/tests/ctable/test_schema_specs.py +++ b/tests/ctable/test_schema_specs.py @@ -200,7 +200,7 @@ def test_complex128_metadata_dict(): assert complex128().to_metadata_dict() == {"kind": "complex128"} -def test_ndarray_metadata_dict_normalizes_numpy_scalar_null_value(): +def test_ndarray_metadata_normalizes_np_scalar(): spec = blosc2.ndarray((2,), dtype=np.int16, null_value=np.int16(123)) d = spec.to_metadata_dict() diff --git a/tests/ctable/test_schema_validation.py b/tests/ctable/test_schema_validation.py index b170caf6f..882b5bf57 100644 --- a/tests/ctable/test_schema_validation.py +++ b/tests/ctable/test_schema_validation.py @@ -66,7 +66,7 @@ def test_append_default_fill(): assert t[0].id == 5 -def test_append_omitted_no_default_column_raises_clear_error(): +def test_append_omitted_no_default_raises(): t = CTable(Row, expected_size=5) with pytest.raises(ValueError, match="no default declared"): t.append(()) @@ -109,14 +109,14 @@ def test_extend_le_violation(): t.extend(data) -def test_extend_omitted_columns_with_defaults_are_filled(): +def test_extend_omitted_defaults_are_filled(): t = CTable(Row, expected_size=10) t.extend({"id": [1, 2]}) assert list(t["score"][:]) == [0.0, 0.0] assert list(t["active"][:]) == [True, True] -def test_extend_omitted_no_default_column_raises_clear_error(): +def test_extend_omitted_no_default_raises(): t = CTable(Row, expected_size=10) with pytest.raises(ValueError, match="no default declared"): t.extend({"score": [1.0, 2.0]}) diff --git a/tests/ctable/test_sort_by.py b/tests/ctable/test_sort_by.py index dc860875a..a0c2d6f3d 100644 --- a/tests/ctable/test_sort_by.py +++ b/tests/ctable/test_sort_by.py @@ -82,7 +82,7 @@ def test_sort_accepts_nested_column_selector_from_view(): np.testing.assert_array_equal(s["trip.sec"][:], [1, 2, 3, 4]) -def test_sort_projected_view_with_dictionary_column_above_default_capacity(): +def test_sort_projected_view_dict_over_capacity(): n = 5000 data = [(i, n - i, f"label-{i % 7}") for i in range(n)] t = CTable(DictSortRow, new_data=data) @@ -97,7 +97,7 @@ def test_sort_projected_view_with_dictionary_column_above_default_capacity(): assert "label" in str(sorted_view) -def test_sort_accepts_column_selectors_in_multi_key_list(): +def test_sort_accepts_col_selectors_multi_key(): t = CTable(Row, new_data=DATA) s = t.sort_by([t.score, t.id], ascending=[True, False]) @@ -377,7 +377,7 @@ def _loaded_columns(table) -> set[str]: return {name for name in table.col_names if dict.__contains__(table._cols, name)} -def test_sort_unprojected_view_opens_only_needed_columns(tmp_path): +def test_sort_unprojected_opens_needed_cols(tmp_path): """``where(cond).sort_by(key)`` without ``columns=`` used to gather every column of the view (~30x slower than projecting first). It must open only the condition and sort-key columns, deferring the rest until read.""" diff --git a/tests/ctable/test_table_persistency.py b/tests/ctable/test_table_persistency.py index 00e4d5c4c..71abf81b5 100644 --- a/tests/ctable/test_table_persistency.py +++ b/tests/ctable/test_table_persistency.py @@ -944,7 +944,7 @@ class VLRow: assert t2["data"][:] == [b"bin2", b"bin3"] -def test_ctable_from_cframe_rejects_non_embedstore_cframe(): +def test_from_cframe_rejects_non_embedstore(): """Passing an NDArray cframe raises ValueError.""" nd = blosc2.arange(10) cframe = nd.to_cframe() diff --git a/tests/ctable/test_utf8.py b/tests/ctable/test_utf8.py index fa2089957..13bdd3a07 100644 --- a/tests/ctable/test_utf8.py +++ b/tests/ctable/test_utf8.py @@ -77,7 +77,7 @@ def test_utf8_spec_metadata_round_trip(): assert d["null_value"] == "" restored = spec_from_metadata_dict(d) - assert type(restored).__name__ == "Utf8Spec" + assert type(restored).__name__ == "UTF8Spec" assert restored.nullable is True assert restored.null_value == "" @@ -105,14 +105,14 @@ class Plain: # --------------------------------------------------------------------------- -# Utf8Array internal adapter +# UTF8Array internal adapter # --------------------------------------------------------------------------- def test_utf8_array_basic_roundtrip(): - from blosc2.utf8_array import Utf8Array + from blosc2._utf8_array import UTF8Array - arr = Utf8Array(blosc2.utf8()) + arr = UTF8Array(blosc2.utf8()) arr.extend(SAMPLE) assert len(arr) == len(SAMPLE) assert list(arr[:]) == SAMPLE @@ -124,9 +124,9 @@ def test_utf8_array_basic_roundtrip(): def test_utf8_array_reads_across_pending_boundary(): - from blosc2.utf8_array import Utf8Array + from blosc2._utf8_array import UTF8Array - arr = Utf8Array(blosc2.utf8()) + arr = UTF8Array(blosc2.utf8()) arr.extend(SAMPLE[:4]) arr.flush() arr.extend(SAMPLE[4:]) # stays pending @@ -141,9 +141,9 @@ def test_utf8_array_reads_across_pending_boundary(): def test_utf8_array_setitem_shifts_offsets(): - from blosc2.utf8_array import Utf8Array + from blosc2._utf8_array import UTF8Array - arr = Utf8Array(blosc2.utf8()) + arr = UTF8Array(blosc2.utf8()) arr.extend(["aa", "bb", "cc"]) arr.flush() arr[1] = "a longer replacement value" @@ -153,9 +153,9 @@ def test_utf8_array_setitem_shifts_offsets(): def test_utf8_array_rejects_non_str(): - from blosc2.utf8_array import Utf8Array + from blosc2._utf8_array import UTF8Array - arr = Utf8Array(blosc2.utf8()) + arr = UTF8Array(blosc2.utf8()) with pytest.raises(TypeError, match="Expected str"): arr.append(42) with pytest.raises(TypeError, match="not nullable"): @@ -168,9 +168,9 @@ def test_utf8_array_rejects_non_str(): def test_utf8_array_extend_empty_iterable_is_noop(): - from blosc2.utf8_array import Utf8Array + from blosc2._utf8_array import UTF8Array - arr = Utf8Array(blosc2.utf8()) + arr = UTF8Array(blosc2.utf8()) arr.extend([]) assert len(arr) == 0 arr.extend(iter([])) @@ -184,11 +184,11 @@ def test_utf8_array_extend_many_rows_no_dropped_rows(): `self._pending` to a fresh list rather than mutating it, so an `extend()` spanning several internal flushes must re-read `self._pending` after each one instead of caching a reference.""" - from blosc2.utf8_array import _FLUSH_ROWS, Utf8Array + from blosc2._utf8_array import _FLUSH_ROWS, UTF8Array n = _FLUSH_ROWS * 3 + 7 values = [f"row{i}" for i in range(n)] - arr = Utf8Array(blosc2.utf8()) + arr = UTF8Array(blosc2.utf8()) arr.extend(values) assert len(arr) == n arr.flush() @@ -196,22 +196,22 @@ def test_utf8_array_extend_many_rows_no_dropped_rows(): assert list(arr[:]) == values -def test_utf8_array_extend_none_straddles_chunk_boundary(): - from blosc2.utf8_array import _FLUSH_ROWS, Utf8Array +def test_utf8_array_extend_none_straddles_chunk(): + from blosc2._utf8_array import _FLUSH_ROWS, UTF8Array values = [f"v{i}" for i in range(_FLUSH_ROWS + 2)] values[_FLUSH_ROWS - 1] = None # last row of first chunk values[_FLUSH_ROWS + 1] = None # second row of second chunk - arr = Utf8Array(blosc2.utf8(null_value="")) + arr = UTF8Array(blosc2.utf8(null_value="")) arr.extend(values) expected = [v if v is not None else "" for v in values] assert list(arr[:]) == expected -def test_utf8_array_extend_append_interleaved_before_flush(): - from blosc2.utf8_array import Utf8Array +def test_utf8_array_extend_append_interleaved(): + from blosc2._utf8_array import UTF8Array - arr = Utf8Array(blosc2.utf8()) + arr = UTF8Array(blosc2.utf8()) arr.append("first") arr.extend(["second", "third"]) arr.append("fourth") @@ -220,25 +220,25 @@ def test_utf8_array_extend_append_interleaved_before_flush(): def test_utf8_array_extend_ascii_nul_byte_preserved(): - from blosc2.utf8_array import Utf8Array + from blosc2._utf8_array import UTF8Array values = ["nul\x00in", "plain", "\x00leading", "trailing\x00"] assert all(v.isascii() for v in values) - arr = Utf8Array(blosc2.utf8()) + arr = UTF8Array(blosc2.utf8()) arr.extend(values) arr.flush() assert list(arr[:]) == values -def test_utf8_array_extend_multi_mb_strings_bounded_flush(): +def test_utf8_array_extend_multi_mb_bounded(): """~20 multi-MB ASCII strings: char-count flush bound is checked once per _FLUSH_ROWS-sized chunk (not per row), so this overshoots _FLUSH_CHARS by at most one chunk before flushing -- confirm read-back is still correct despite the coarser check.""" - from blosc2.utf8_array import Utf8Array + from blosc2._utf8_array import UTF8Array values = [f"{i:06d}" + "x" * (2 * 1024 * 1024) for i in range(20)] - arr = Utf8Array(blosc2.utf8()) + arr = UTF8Array(blosc2.utf8()) arr.extend(values) arr.flush() assert list(arr[:]) == values @@ -255,7 +255,7 @@ def force_kernel_mode(request, monkeypatch): pure-Python per-row fallback, so the fallback stays covered even on a build where the compiled extension is available.""" if request.param == "fallback": - monkeypatch.setattr("blosc2.utf8_array._pack_utf8_kernel", lambda: None) + monkeypatch.setattr("blosc2._utf8_array._pack_utf8_kernel", lambda: None) return request.param @@ -282,9 +282,9 @@ def test_pack_utf8_span_rejects_malformed_rel(): def test_utf8_array_bulk_read_kernel_and_fallback(force_kernel_mode): - from blosc2.utf8_array import Utf8Array + from blosc2._utf8_array import UTF8Array - arr = Utf8Array(blosc2.utf8()) + arr = UTF8Array(blosc2.utf8()) arr.extend(SAMPLE) arr.flush() got = arr[:] @@ -292,21 +292,21 @@ def test_utf8_array_bulk_read_kernel_and_fallback(force_kernel_mode): assert list(got) == SAMPLE -def test_utf8_array_bulk_read_matches_python_ground_truth(force_kernel_mode): +def test_utf8_array_bulk_read_matches_python(force_kernel_mode): """A wider mix of byte lengths and edge cases than SAMPLE: many distinct ASCII/multi-byte/empty/NUL-bearing values, read back in one bulk span.""" - from blosc2.utf8_array import Utf8Array + from blosc2._utf8_array import UTF8Array rng = np.random.default_rng(5) pool = ["", "a", "café", "日本語", "x" * 5000, "nul\x00in", "nul\x00INSIDE", "emoji 🎉🚀"] values = [pool[i] for i in rng.integers(0, len(pool), 3000)] - arr = Utf8Array(blosc2.utf8()) + arr = UTF8Array(blosc2.utf8()) arr.extend(values) arr.flush() assert list(arr[:]) == values -def test_ctable_utf8_extend_and_read_kernel_and_fallback(force_kernel_mode): +def test_ctable_utf8_extend_read_two_routes(force_kernel_mode): t = make_table() values = t["name"][:] assert values.dtype == STRING_DTYPE @@ -314,7 +314,7 @@ def test_ctable_utf8_extend_and_read_kernel_and_fallback(force_kernel_mode): @pytest.mark.parametrize("ext", [".b2z", ".b2d"]) -def test_ctable_utf8_persistence_roundtrip_kernel_and_fallback(tmp_path, ext, force_kernel_mode): +def test_ctable_utf8_persist_two_routes(tmp_path, ext, force_kernel_mode): urlpath = str(tmp_path / f"utf8_kernel_mode{ext}") t = make_table(urlpath=urlpath, mode="w") t.close() @@ -336,14 +336,14 @@ def force_write_kernel_mode(request, monkeypatch): join+encode fallback, so the fallback stays covered even on a build where the compiled extension is available.""" if request.param == "fallback": - monkeypatch.setattr("blosc2.utf8_array._encode_utf8_kernel", lambda: None) + monkeypatch.setattr("blosc2._utf8_array._encode_utf8_kernel", lambda: None) return request.param def test_utf8_array_extend_kernel_and_fallback(force_write_kernel_mode): - from blosc2.utf8_array import Utf8Array + from blosc2._utf8_array import UTF8Array - arr = Utf8Array(blosc2.utf8()) + arr = UTF8Array(blosc2.utf8()) arr.extend(SAMPLE) arr.flush() assert list(arr[:]) == SAMPLE @@ -352,34 +352,34 @@ def test_utf8_array_extend_kernel_and_fallback(force_write_kernel_mode): def test_utf8_array_extend_matches_python_ground_truth(force_write_kernel_mode): """Same wider mix of byte lengths and edge cases as the read-side ground-truth test, exercised through the write path this time.""" - from blosc2.utf8_array import Utf8Array + from blosc2._utf8_array import UTF8Array rng = np.random.default_rng(7) pool = ["", "a", "café", "日本語", "x" * 5000, "nul\x00in", "nul\x00INSIDE", "emoji 🎉🚀"] values = [pool[i] for i in rng.integers(0, len(pool), 3000)] - arr = Utf8Array(blosc2.utf8()) + arr = UTF8Array(blosc2.utf8()) arr.extend(values) arr.flush() assert list(arr[:]) == values -def test_utf8_array_extend_ascii_nul_byte_kernel_and_fallback(force_write_kernel_mode): - from blosc2.utf8_array import Utf8Array +def test_utf8_array_extend_nul_two_routes(force_write_kernel_mode): + from blosc2._utf8_array import UTF8Array values = ["nul\x00in", "plain", "\x00leading", "trailing\x00"] - arr = Utf8Array(blosc2.utf8()) + arr = UTF8Array(blosc2.utf8()) arr.extend(values) arr.flush() assert list(arr[:]) == values -def test_utf8_array_extend_multi_mb_string_kernel_and_fallback(force_write_kernel_mode): +def test_utf8_array_extend_mb_two_routes(force_write_kernel_mode): """A single multi-MB value alongside short ones -- sanity-checks the total-length/offset accumulation in the compiled kernel's two passes.""" - from blosc2.utf8_array import Utf8Array + from blosc2._utf8_array import UTF8Array values = ["head", "x" * (8 * 1024 * 1024), "tail", "café" * 100_000] - arr = Utf8Array(blosc2.utf8()) + arr = UTF8Array(blosc2.utf8()) arr.extend(values) arr.flush() assert list(arr[:]) == values @@ -390,14 +390,14 @@ def test_ctable_utf8_extend_kernel_and_fallback(force_write_kernel_mode): assert list(t["name"][:]) == SAMPLE -def test_utf8_array_extend_lone_surrogate_raises_and_recovers(force_write_kernel_mode): +def test_utf8_array_extend_surrogate_recovers(force_write_kernel_mode): """A lone surrogate is invalid UTF-8: flush() must raise UnicodeEncodeError, matching str.encode('utf-8')'s own behavior, and the array must remain usable afterwards -- a regression test for the compiled kernel's temp-buffer cleanup on the error path.""" - from blosc2.utf8_array import Utf8Array + from blosc2._utf8_array import UTF8Array - arr = Utf8Array(blosc2.utf8()) + arr = UTF8Array(blosc2.utf8()) arr.extend(["first"]) arr.flush() arr.extend(["ok", "bad\udc80value"]) @@ -509,6 +509,39 @@ def test_ctable_utf8_add_and_drop_column(): assert "note" not in t.col_names +def test_ctable_utf8_add_column_values(): + t = make_table(["a", "b", "c"]) + t.add_column("note", blosc2.utf8(), values=["x", "yy", "zzz"]) + assert list(t["note"][:]) == ["x", "yy", "zzz"] + + +def test_ctable_utf8_add_col_values_from_expr(): + """The documented round trip: compute on = ""]["name"][:]) == ["", "a", "zzz"] -def test_ctable_utf8_ordering_multibyte_byte_length_boundaries(): +def test_ctable_utf8_ordering_multibyte_bounds(): """1-, 2-, and 3-byte UTF-8 encodings must byte-compare in code-point order (code points 0x7A < 0xE9 < 0x65E5).""" assert "z" < "é" < "日" @@ -831,7 +864,7 @@ def test_ctable_utf8_ordering_probe_equals_sentinel(): assert None not in got -def test_ctable_utf8_scalar_comparison_view_and_deleted_rows(): +def test_ctable_utf8_scalar_cmp_view_deletes(): """The predicate mask is physical-length; it must stay correct through a view and after rows have been deleted (live-row mask intersection).""" t = make_table(["paris", "london", "paris", "tokyo", "berlin", "paris"]) @@ -858,18 +891,18 @@ def test_ctable_utf8_startswith_endswith(): # --------------------------------------------------------------------------- -def test_utf8_factorize_span_matches_np_unique_contract(): +def test_utf8_factorize_span_matches_np_unique(): """The raw-bytes factorization keeps the np.unique contract: uniques sorted ascending, codes indexing them. Ground truth is Python's set — numpy's np.unique on StringDType merges strings differing only after an embedded NUL (numpy bug), which the byte-exact factorization does not. """ - from blosc2.utf8_array import Utf8Array + from blosc2._utf8_array import UTF8Array rng = np.random.default_rng(7) pool = ["", "a", "ab", "café", "日本語", "x" * 3000, "nul\x00in", "nul\x00IN", "Wien", "wien"] values = [pool[i] for i in rng.integers(0, len(pool), 5000)] - arr = Utf8Array(blosc2.utf8()) + arr = UTF8Array(blosc2.utf8()) arr.extend(values) codes, uniques = arr.factorize_span(0, len(values)) assert list(uniques) == sorted(set(values)) @@ -877,9 +910,9 @@ def test_utf8_factorize_span_matches_np_unique_contract(): def test_utf8_factorizer_cross_span_codes_are_global(): - from blosc2.utf8_array import Utf8Array + from blosc2._utf8_array import UTF8Array - arr = Utf8Array(blosc2.utf8()) + arr = UTF8Array(blosc2.utf8()) arr.extend(["b", "a", "b", "c", "a", "d"]) fact = arr.factorizer() c1 = fact.codes_for_span(0, 3) # b, a, b @@ -891,7 +924,7 @@ def test_utf8_factorizer_cross_span_codes_are_global(): assert c1[1] == c2[1] -def test_ctable_utf8_groupby_many_byte_lengths_and_non_ascii(): +def test_ctable_utf8_groupby_lengths_non_ascii(): rng = np.random.default_rng(3) pool = ["", "a", "bb", "café", "日本語のテキスト", "x" * 2000, "münchen"] names = [pool[i] for i in rng.integers(0, len(pool), 3000)] @@ -1068,7 +1101,7 @@ def test_ctable_utf8_sort_inplace(): assert list(t["name"][:]) == ["a", "b", "c"] -def test_ctable_utf8_sort_multi_key_with_bystander_utf8_column(): +def test_ctable_utf8_sort_multi_key_bystander(): """A non-key utf8 column in the same table must be reordered along with the sort, not just the sort key itself.""" @@ -1104,7 +1137,7 @@ class TwoCols: @pytest.mark.parametrize("ext", [".b2z", ".b2d"]) -def test_ctable_utf8_sort_inplace_persists_after_reopen(tmp_path, ext): +def test_ctable_utf8_sort_inplace_persists(tmp_path, ext): """Regression: sort_by(inplace=True) on a file-backed table must write the sorted utf8 rows through to the store, keeping them aligned with the other (on-disk-sorted) columns after close/reopen.""" @@ -1141,7 +1174,7 @@ def test_ctable_utf8_compact_persists_after_reopen(tmp_path, ext): t2.close() -def test_ctable_utf8_setitem_persisted_shifts_survive_reopen(tmp_path): +def test_ctable_utf8_setitem_shifts_reopen(tmp_path): """__setitem__ on persisted rows shifts the byte blob in place; longer, shorter, equal-length, and empty replacements must all round-trip.""" urlpath = str(tmp_path / "utf8_setitem.b2d") @@ -1172,16 +1205,24 @@ def test_ctable_utf8_sort_non_ascii(): # --------------------------------------------------------------------------- -def test_ctable_utf8_where_expression_raises_clearly(): - t = make_table() - with pytest.raises(NotImplementedError, match="utf8"): - t.where("name == 'hello'") +def test_ctable_utf8_create_index_builds_a_rank_index(): + """utf8 is indexed by the alphabetical rank of each row's value. + Sorting by rank is sorting by decoded string, so an int32 rank column drives + the existing numeric index machinery unchanged. + """ + values = ["pear", "apple", "café", "banana", "apple"] + t = make_table(values) + index = t.create_index(col_name="name", kind="full") -def test_ctable_utf8_create_index_raises_clearly(): - t = make_table() - with pytest.raises(NotImplementedError, match="utf8"): - t.create_index(col_name="name") + assert index.kind == "full" + meta = t._get_index_catalog()["name"]["full"]["utf8_rank"] + assert meta["vocab_len"] == len(set(values)) + assert meta["n_rows"] == len(values) + + # Ordering through the index must match a plain sort. + assert list(t.sort_by("name", view=True)["name"][:]) == sorted(values) + assert list(t.sorted_slice("name", slice(0, 2))["name"][:]) == sorted(values)[:2] def test_ctable_utf8_arrow_export_large_string(): @@ -1264,7 +1305,7 @@ def test_utf8_from_arrow_nulls_use_sentinel(): assert t["name"].null_count() == 1 -def test_utf8_from_arrow_fixed_width_when_max_length_given(): +def test_utf8_from_arrow_fixed_width_max_len(): pa = pytest.importorskip("pyarrow") at = pa.table({"name": pa.array(["hi", "there"], type=pa.string())}) t = CTable.from_arrow(at.schema, at.to_batches(), string_max_length=32) @@ -1281,3 +1322,1068 @@ def test_utf8_duckdb_query(): "SELECT name, count(*) AS n FROM arrow_tbl WHERE name = 'paris' GROUP BY name" ).fetchall() assert result == [("paris", 2)] + + +# --------------------------------------------------------------------------- +# utf8_array() constructor +# --------------------------------------------------------------------------- + + +def test_utf8_array_constructor(): + arr = blosc2.utf8_array(SAMPLE) + assert isinstance(arr, blosc2.UTF8Array) + assert len(arr) == len(SAMPLE) + assert list(arr[:]) == SAMPLE + + +def test_utf8_array_constructor_with_spec_and_nulls(): + arr = blosc2.utf8_array(["a", None, "c"], blosc2.utf8(nullable=True, null_value="")) + assert list(arr[:]) == ["a", "", "c"] + + +def test_utf8_string_expr_rejects_clashing_sentinels(): + """One result column means one sentinel, so operands must agree on it; the + old first-wins pick relabelled the other operand's nulls silently.""" + a = blosc2.utf8_array(["x", None, "z"], blosc2.utf8(null_value="")) + b = blosc2.utf8_array(["1", "2", None], blosc2.utf8(null_value="")) + with pytest.raises(ValueError, match="different null sentinels"): + blosc2.lazyexpr("a + b", {"a": a, "b": b}).compute() + # A boolean result never carries a sentinel, so it is unaffected. + assert list(blosc2.lazyexpr("a > b", {"a": a, "b": b}).compute()) == [True, False, False] + + +def test_utf8_string_expr_shared_sentinel_survives(): + spec = blosc2.utf8(null_value="") + a = blosc2.utf8_array(["x", None, "z"], spec) + b = blosc2.utf8_array(["1", "2", None], spec) + res = blosc2.lazyexpr("a + b", {"a": a, "b": b}).compute() + assert list(res[:]) == ["x1", "", ""] + assert res.spec.null_value == "" + + +def test_utf8_array_ctor_rejects_none_if_not_null(): + with pytest.raises(TypeError, match="not nullable"): + blosc2.utf8_array(["a", None]) + + +def test_utf8_array_span_max_bytes_reads_only_offsets(): + arr = blosc2.utf8_array(["a", "café", "日本語"]) # 1, 5 and 9 UTF-8 bytes + assert arr._span_max_bytes(0, 3) == 9 + assert arr._span_max_bytes(0, 2) == 5 + assert arr._span_max_bytes(0, 0) == 0 + # Pending (unflushed) rows are measured too. + arr.append("x" * 20) + assert arr._span_max_bytes(0, 4) == 20 + + +# --------------------------------------------------------------------------- +# String expressions over utf8 columns (span-loop driver) +# --------------------------------------------------------------------------- + + +def test_ctable_utf8_where_expression_equality(): + t = make_table(["hello", "help", "world", "café"]) + assert list(t.where("name == 'hello'")["name"][:]) == ["hello"] + assert list(t.where("name != 'hello'")["name"][:]) == ["help", "world", "café"] + + +def test_ctable_utf8_where_expr_vs_operator(): + t = make_table(["paris", "london", "tokyo", "paris"]) + for value in ("paris", "tokyo", "absent"): + expr = list(t.where(f"name == '{value}'")["x"][:]) + operator = list(t[t.name == value]["x"][:]) + assert expr == operator, value + + +def test_ctable_utf8_where_expression_predicates(): + t = make_table(["hello", "help", "world"]) + assert list(t.where("startswith(name, 'hel')")["name"][:]) == ["hello", "help"] + assert list(t.where("endswith(name, 'lo')")["name"][:]) == ["hello"] + assert list(t.where("contains(name, 'l')")["name"][:]) == ["hello", "help", "world"] + + +def test_ctable_utf8_where_expr_mixes_numeric(): + t = make_table(["a", "b", "c", "d"]) + assert list(t.where("(name == 'b') | (x > 2)")["name"][:]) == ["b", "d"] + assert list(t.where("(name != 'a') & (x < 2)")["name"][:]) == ["b"] + + +def test_ctable_utf8_where_expression_runs_on_miniexpr(): + """A silent NumPy fallback would produce the same values, so pin the engine. + + ``strict_miniexpr`` raises rather than falling back, which is the only + assertion that distinguishes the two. + """ + t = make_table(["hello", "help", "world"]) + got = t._utf8_span_eval("startswith(name, 'hel')", {}, ["name"], strict=True) + assert list(got[:3]) == [True, True, False] + + +def test_ctable_utf8_where_expr_many_widths(): + # Exercises the power-of-two width bucketing: values straddle several + # buckets and one of them is past the 255-byte typesize cap. + values = ["a", "bb", "x" * 40, "y" * 300, "café", ""] * 30 + t = make_table(values) + assert list(t.where("name == 'café'")["x"][:]) == [i for i, v in enumerate(values) if v == "café"] + assert list(t.where("startswith(name, 'y')")["x"][:]) == [ + i for i, v in enumerate(values) if v.startswith("y") + ] + + +def test_ctable_utf8_where_expr_splits_spans(): + # A single long row would size the whole span's 1] + assert list(view.where("name == 'paris'")["x"][:]) == [2] + + +# --------------------------------------------------------------------------- +# Null policy (3c): nulls are materialized to "" and re-masked afterwards +# --------------------------------------------------------------------------- + + +def _nullable_table(values): + return CTable( + NullableRow, + new_data={"name": list(values), "x": list(range(len(values)))}, + ) + + +def test_ctable_utf8_where_expr_nulls_no_match(): + t = _nullable_table(["hello", None, "help", None, "world"]) + assert list(t.where("name == 'hello'")["x"][:]) == [0] + assert list(t.where("startswith(name, 'hel')")["x"][:]) == [0, 2] + # Not even against the sentinel string itself: a null is not a value. + assert list(t.where("name == ''")["x"][:]) == [] + + +def test_ctable_utf8_where_expr_nulls_operator(): + t = _nullable_table(["hello", None, "help", None, "world"]) + for value in ("hello", "world", ""): + assert list(t.where(f"name == '{value}'")["x"][:]) == list(t[t.name == value]["x"][:]) + assert list(t.where(f"name != '{value}'")["x"][:]) == list(t[t.name != value]["x"][:]) + + +def test_ctable_utf8_where_expression_all_null_column(): + t = _nullable_table([None] * 5) + assert list(t.where("name == 'hello'")["x"][:]) == [] + assert list(t.where("name != 'hello'")["x"][:]) == [] + + +def test_ctable_utf8_where_expr_no_nulls_fast(): + # A nullable column with no actual nulls must not mask anything away. + t = _nullable_table(["hello", "help", "world"]) + assert list(t.where("startswith(name, 'hel')")["x"][:]) == [0, 1] + + +def test_ctable_utf8_sum_where_expression(): + t = make_table(["hello", "help", "world"]) + assert t["x"].sum(where="startswith(name, 'hel')") == 1 + + +# --------------------------------------------------------------------------- +# Scalar predicates take the raw-byte path instead of the span driver +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("expr", "predicate"), + [ + ("name == 'help'", lambda v: v == "help"), + ("name != 'help'", lambda v: v != "help"), + ("name < 'm'", lambda v: v < "m"), + ("name <= 'help'", lambda v: v <= "help"), + ("name > 'm'", lambda v: v > "m"), + ("name >= 'help'", lambda v: v >= "help"), + ("'help' == name", lambda v: v == "help"), + ("'help' != name", lambda v: v != "help"), + ("'m' > name", lambda v: v < "m"), + ("'m' <= name", lambda v: v >= "m"), + ], +) +def test_ctable_utf8_scalar_predicates_match_python(expr, predicate): + values = ["hello", "help", "world", "café", "日本語", "", "zz"] + t = make_table(values) + assert list(t.where(expr)["x"][:]) == [i for i, v in enumerate(values) if predicate(v)] + + +@pytest.mark.parametrize( + ("expr", "rewritten_away"), + [ + ("name == 'help'", True), + ("'help' == name", True), + ("name < 'm'", True), + ("(name == 'help') | (x > 4)", True), + ("(name == 'a') & (name != 'b')", True), + # Not a scalar comparison: these still need the span driver. + ("startswith(name, 'hel')", False), + ("contains(name, 'l')", False), + ("startswith(name, 'hel') | (name == 'zz')", False), + ], +) +def test_ctable_utf8_preds_skip_span_driver(expr, rewritten_away): + """The raw-byte scan is several times cheaper than decode -> miniexpr. + + Correctness alone would not notice the difference, so assert on which route + the expression takes: a utf8 name survives the rewrite only when something + other than a scalar comparison still references it. + """ + t = make_table(["hello", "help", "world", "zz", "a", "b"]) + operands = t._where_expression_operands(expr) + _, _, remaining = t._rewrite_utf8_predicates(expr, operands, t._utf8_names_in(expr)) + assert (remaining == []) is rewritten_away + + +def test_ctable_utf8_rewritten_pred_vs_driver(): + """Both routes must agree, including on nulls and on the sentinel spelling.""" + values = ["hello", None, "help", None, "world"] + t = _nullable_table(values) + for expr in ("name == 'hello'", "name != 'hello'", "name < 'm'", "name == ''"): + fast = list(t.where(expr)["x"][:]) + slow = list(np.flatnonzero(t._utf8_span_eval(expr, {}, ["name"])[: len(values)])) + assert fast == slow, expr + + +def test_ctable_utf8_pred_literal_with_ops(): + # The literal is parsed with ast.literal_eval, so quoted operators and + # spaces inside it must not be mistaken for expression syntax. + values = ["a == b", "x > y", "plain"] + t = make_table(values) + assert list(t.where("name == 'a == b'")["x"][:]) == [0] + assert list(t.where("name == 'x > y'")["x"][:]) == [1] + + +def test_ctable_utf8_pred_view_and_delete(): + t = make_table(["paris", "london", "paris", "tokyo"]) + t.delete([0]) + assert list(t.where("name == 'paris'")["x"][:]) == [2] + view = t[t.x > 1] + assert list(view.where("name == 'paris'")["x"][:]) == [2] + + +def test_ctable_utf8_two_preds_same_col(): + t = make_table(["a", "b", "c", "d"]) + assert list(t.where("(name > 'a') & (name < 'd')")["name"][:]) == ["b", "c"] + + +@pytest.mark.parametrize("op", ["==", "!=", "<", "<=", ">", ">="]) +def test_utf8_array_comparisons_match_numpy(op): + """Comparisons must be element-wise, not object identity. + + Without ``__eq__`` these fell through to identity, so ``arr == "hello"`` + was a plain ``False`` — silently wrong rather than an error. + """ + import operator + + values = ["hello", "world", "héllo", "abc", "", "hello"] + arr = blosc2.utf8_array(values) + ref = np.array(values, dtype=arr.dtype) + fn = getattr(operator, {"==": "eq", "!=": "ne", "<": "lt", "<=": "le", ">": "gt", ">=": "ge"}[op]) + + for probe in ("hello", "héllo", "", "zzz"): + got = fn(arr, probe) + assert isinstance(got, np.ndarray) + assert got.dtype == np.bool_ + np.testing.assert_array_equal(got, fn(ref, probe)) + + +def test_utf8_array_comparison_against_array_likes(): + values = ["a", "bb", "ccc"] + arr = blosc2.utf8_array(values) + ref = np.array(values, dtype=arr.dtype) + + np.testing.assert_array_equal(arr == values, np.ones(3, dtype=bool)) + np.testing.assert_array_equal(arr == ref, np.ones(3, dtype=bool)) + np.testing.assert_array_equal(arr == blosc2.utf8_array(values), np.ones(3, dtype=bool)) + np.testing.assert_array_equal(arr != blosc2.utf8_array(["a", "x", "ccc"]), [False, True, False]) + + +def test_utf8_array_comparison_edge_cases(): + # Unflushed pending rows take part in the comparison. + arr = blosc2.utf8_array(["a"]) + arr.append("b") + np.testing.assert_array_equal(arr == "b", [False, True]) + + # Empty array yields an empty mask rather than raising. + empty = blosc2.utf8_array([]) + assert (empty == "x").shape == (0,) + + # Defining __eq__ must not have made the container unhashable. + assert isinstance(hash(arr), int) + + +def test_bare_utf8_expr_uses_span_driver(): + """A bare UTF8Array must not evaluate through the NumPy slices_eval path. + + That path returns correct-looking values while never reaching miniexpr, + ignoring the span budget, and widening the result to a fixed `` 1)", {"a": arr, "n": nums}) + np.testing.assert_array_equal(mixed.compute(strict_miniexpr=True), [False, False, True]) + + +@pytest.mark.parametrize(("span_rows", "budget"), [(7, 64 << 20), (65536, 512)]) +def test_bare_utf8_array_expression_splits_spans(span_rows, budget, monkeypatch): + """Both the row-span and the byte-budget splits must hold over a bare array.""" + from blosc2 import _utf8_array + + values = [f"row-{i}" for i in range(50)] + arr = blosc2.utf8_array(values) + + spans = [] + original = _utf8_array.utf8_spans + monkeypatch.setattr( + _utf8_array, + "utf8_spans", + lambda a, n, s, b: [spans.append(x) or x for x in original(a, n, s, b)], + ) + monkeypatch.setattr(_utf8_array, "UTF8_EXPR_SPAN", span_rows) + monkeypatch.setattr(_utf8_array, "UTF8_EXPR_BUDGET", budget) + + result = blosc2.lazyexpr("'x=' + a", {"a": arr}).compute(strict_miniexpr=True) + assert len(spans) > 1, f"expected a split, got {spans}" + assert list(result[:]) == ["x=" + v for v in values] + + +def test_bare_utf8_expr_rejects_unsupported(): + arr = blosc2.utf8_array(["a", "b"]) + lazy = blosc2.lazyexpr("upper(a)", {"a": arr}) + + assert lazy.shape == (2,) + assert len(lazy) == 2 + assert list(lazy[0:1]) == ["A"] + + with pytest.raises(NotImplementedError, match="whole-array only"): + lazy.compute(item=slice(0, 1)) + with pytest.raises(NotImplementedError, match="not supported"): + blosc2.lazyexpr("upper(a)", {"a": arr}, where=(arr, arr)) + + +def test_ctable_utf8_index_reopen_nulls_last(tmp_path): + """A persisted utf8 rank index must reopen and keep nulls at the end.""" + from dataclasses import make_dataclass + + path = str(tmp_path / "utf8_index.b2t") + row_cls = make_dataclass("Row", [("name", str, blosc2.field(blosc2.utf8(nullable=True)))]) + values = ["pear", "apple", None, "banana"] + t = blosc2.CTable(row_cls, urlpath=path, mode="w") + t.extend({"name": values}, validate=False) + t._flush_varlen_columns() + t.create_index("name", kind="full") + del t + + reopened = blosc2.open(path) + assert reopened._get_index_catalog()["name"]["full"]["utf8_rank"]["null_rank"] == 3 + ordered = list(reopened.sort_by("name", view=True)["name"][:]) + assert ordered[:3] == ["apple", "banana", "pear"] + assert ordered[3] == reopened["name"].null_value # the null sentinel sorts last + + +def test_ctable_utf8_index_stale_on_change(): + """Appending a value ahead of existing ones invalidates every rank.""" + t = make_table(["pear", "banana"]) + t.create_index("name", kind="full") + meta = t._get_index_catalog()["name"]["full"]["utf8_rank"] + assert not t._utf8_rank_index_stale("name", meta) + + t.append({"name": "apple", "x": 2}) + t._flush_varlen_columns() + assert t._utf8_rank_index_stale("name", meta) + # The answer stays correct via the lexsort fallback. + assert list(t.sort_by("name", view=True)["name"][:]) == ["apple", "banana", "pear"] + + +def test_utf8_rejects_a_lone_nul_null_value(): + """NumPy will not match a lone NUL against StringDType, so nulls would vanish.""" + import numpy as np + + probe = np.array(["\x00"], dtype=np.dtypes.StringDType()) + assert not (probe == "\x00")[0], "numpy started matching lone NUL; the guard can go" + + with pytest.raises(ValueError, match="NUL"): + blosc2.utf8(null_value="\x00") + # A NUL that is not the whole string is fine — numpy matches those. + assert blosc2.utf8(null_value="\x00x").null_value == "\x00x" + + +@pytest.mark.parametrize("nullable", [False, True]) +def test_ctable_utf8_index_answers_scalar_predicates(nullable, tmp_path): + """With a rank index, a scalar comparison is a sidecar lookup, not a scan. + + The literal is located by one searchsorted over the stored vocabulary and + the matching rows are a contiguous run of the sorted-positions sidecar. + Results must be identical to the raw-byte scan, including that a null + satisfies no comparison. + """ + from dataclasses import make_dataclass + + import numpy as np + + values = ["pear", "apple", "café", "banana", "apple", "pear"] + if nullable: + values = [*values, None] + spec = blosc2.utf8(nullable=True) if nullable else blosc2.utf8() + row_cls = make_dataclass("Row", [("c", str, blosc2.field(spec))]) + + masks = {} + for tag in ("scan", "index"): + t = blosc2.CTable(row_cls, urlpath=str(tmp_path / f"{tag}.b2t"), mode="w") + t.extend({"c": values}, validate=False) + t._flush_varlen_columns() + if tag == "index": + t.create_index("c", kind="full") + col = t["c"] + got = {} + for name, op in ( + ("==", np.equal), + ("!=", np.not_equal), + ("<", np.less), + ("<=", np.less_equal), + (">", np.greater), + (">=", np.greater_equal), + ): + for probe in ("apple", "pear", "zzz-absent", ""): + got[(name, probe)] = col._utf8_scalar_mask(op, probe).copy() + masks[tag] = got + if tag == "index": + # The fast path must really have been taken, not silently skipped. + assert col._utf8_index_mask(np.equal, "apple") is not None + del t + + for key, scanned in masks["scan"].items(): + np.testing.assert_array_equal(masks["index"][key], scanned, err_msg=f"{key}") + + +def test_ctable_utf8_index_pred_falls_back(tmp_path): + """A stale rank index must not answer predicates from frozen ranks.""" + from dataclasses import make_dataclass + + import numpy as np + + row_cls = make_dataclass("Row", [("c", str, blosc2.field(blosc2.utf8()))]) + t = blosc2.CTable(row_cls, urlpath=str(tmp_path / "t.b2t"), mode="w") + t.extend({"c": ["pear", "banana"]}, validate=False) + t._flush_varlen_columns() + t.create_index("c", kind="full") + assert t["c"]._utf8_index_mask(np.equal, "pear") is not None + + t.append({"c": "apple"}) # a value ahead of the others invalidates every rank + t._flush_varlen_columns() + assert t["c"]._utf8_index_mask(np.equal, "pear") is None + np.testing.assert_array_equal(t["c"]._utf8_scalar_mask(np.equal, "apple")[:3], [False, False, True]) + + +def test_ctable_utf8_index_spans_deleted_rows(tmp_path): + """The index has to cover the physical extent, not the live row count. + + delete() tombstones in place, so live rows sit past the live count. Sizing + the index by that count leaves it permanently stale -- built, paid for, and + never consulted. + """ + from dataclasses import make_dataclass + + import numpy as np + + values = ["pear", "apple", "cherry", "apple", "banana", "apple"] + row_cls = make_dataclass("Row", [("c", str, blosc2.field(blosc2.utf8()))]) + t = blosc2.CTable(row_cls, urlpath=str(tmp_path / "t.b2t"), mode="w") + t.extend({"c": values}, validate=False) + t._flush_varlen_columns() + t.delete(0) + t.create_index("c", kind="full") + + meta = t._get_index_catalog()["c"]["full"]["utf8_rank"] + assert not t._utf8_rank_index_stale("c", meta) + assert t["c"]._utf8_index_mask(np.equal, "apple") is not None + # Every live "apple" comes back, including the one at the last position. + assert sorted(t[t["c"] == "apple"]["c"][:]) == ["apple"] * 3 + + +def test_ctable_utf8_index_ne_on_all_null_column(tmp_path): + """A null satisfies no comparison, so ``!= x`` on an all-null column is empty. + + With no non-null distinct values the null rank is 0, which used to send the + null-exclusion lookup down an always-empty branch. + """ + from dataclasses import make_dataclass + + row_cls = make_dataclass("Row", [("c", str, blosc2.field(blosc2.utf8(nullable=True)))]) + t = blosc2.CTable(row_cls, urlpath=str(tmp_path / "t.b2t"), mode="w") + t.extend({"c": [None] * 6}, validate=False) + t._flush_varlen_columns() + t.create_index("c", kind="full") + + assert t._get_index_catalog()["c"]["full"]["utf8_rank"]["null_rank"] == 0 + assert len(t[t["c"] != "x"]["c"][:]) == 0 + assert len(t[t["c"] == "x"]["c"][:]) == 0 + + +# --------------------------------------------------------------------------- +# Fixed-width conversion pair: astype / from_utf8 / to_utf8 +# --------------------------------------------------------------------------- + + +def test_utf8_astype_infers_exact_width(): + arr = blosc2.utf8_array(["hello", "café", "日本語", ""]) + out = arr.astype() + assert out.dtype == np.dtype(""] + + +def test_from_utf8_accepts_array_column_and_iterables(): + values = ["hello", "café"] + arr = blosc2.utf8_array(values) + t = make_table(values) + + for source in (arr, t["name"], np.array(values, dtype=STRING_DTYPE), values): + out = blosc2.from_utf8(source) + assert out.dtype == np.dtype("" + + +def test_utf8_conversion_round_trip(): + values = ["", "a", "日本語のテキスト", "x" * 40, "🎉"] + arr = blosc2.utf8_array(values) + assert list(blosc2.to_utf8(blosc2.from_utf8(arr))[:]) == values + + +def test_utf8_conversion_round_trip_via_expr(): + """The documented compute rule, end to end.""" + t = make_table(["a", "bb", "ccc"]) + fixed = blosc2.from_utf8(t["name"]) + res = blosc2.lazyexpr("'x=' + a", {"a": fixed}).compute()[:] + t.add_column("prefixed", blosc2.utf8(), values=blosc2.to_utf8(res)) + assert list(t["prefixed"][:]) == ["x=a", "x=bb", "x=ccc"] + + +# --------------------------------------------------------------------------- +# Column.assign on utf8 +# --------------------------------------------------------------------------- + + +def test_ctable_utf8_column_assign(): + t = make_table(["a", "bb", "ccc"]) + t["name"].assign(["X", "YY", "ZZZ"]) + assert list(t["name"][:]) == ["X", "YY", "ZZZ"] + + +def test_ctable_utf8_col_assign_from_computed(): + t = make_table(["a", "bb"]) + fixed = blosc2.from_utf8(t["name"]) + res = blosc2.lazyexpr("a + '!'", {"a": fixed}).compute()[:] + t["name"].assign(res) + assert list(t["name"][:]) == ["a!", "bb!"] + + +def test_ctable_utf8_column_assign_skips_deleted_rows(): + t = make_table(["a", "b", "c", "d"]) + t.delete([0, 2]) + t["name"].assign(["P", "Q"]) + assert list(t["name"][:]) == ["P", "Q"] + assert list(t["x"][:]) == [1, 3] + + +def test_ctable_utf8_column_assign_wrong_length_raises(): + t = make_table(["a", "bb"]) + with pytest.raises(ValueError, match="requires 2 values"): + t["name"].assign(["only-one"]) + + +def test_ctable_utf8_column_assign_persists(tmp_path): + path = str(tmp_path / "utf8_assign.b2d") + t = make_table(["a", "bb"], urlpath=path, mode="w") + t["name"].assign(["hello", "wörld"]) + t.close() + t2 = CTable.open(path) + assert list(t2["name"][:]) == ["hello", "wörld"] + + +# --------------------------------------------------------------------------- +# Compute-side refusals: they must name the column and route to the conversion +# --------------------------------------------------------------------------- + + +@blosc2.dsl_kernel +def _shout(name): + return name.upper() + + +def _assert_routes(message, source): + """Every utf8 compute refusal must hand back a usable recipe.""" + assert "blosc2.from_utf8(" in message + assert "blosc2.to_utf8(" in message + assert source in message + assert "Computing strings on a utf8 column" in message + + +def test_utf8_computed_col_names_workaround(): + t = make_table(["a", "bb"]) + with pytest.raises(NotImplementedError) as exc: + t.add_computed_column("up", "upper(name)") + _assert_routes(str(exc.value), "t['name']") + assert "upper(name)" in str(exc.value) # the user's own expression is echoed + + +def test_utf8_assign_expression_names_the_workaround(): + t = make_table(["a", "bb"]) + with pytest.raises(NotImplementedError) as exc: + t.assign(up="upper(name)") + _assert_routes(str(exc.value), "t['name']") + + +def test_utf8_generated_col_names_workaround(): + t = make_table(["a", "bb"]) + with pytest.raises(NotImplementedError) as exc: + t.add_generated_column("g", values="upper(name)") + _assert_routes(str(exc.value), "t['name']") + + +def test_utf8_kernel_refused_at_registration(): + """Regression: it used to register, then break every read *and* str(t).""" + t = make_table(["a", "bb"]) + with pytest.raises(NotImplementedError) as exc: + t.add_computed_column("up", _shout, inputs=["name"]) + _assert_routes(str(exc.value), "t['name']") + # The table must be left untouched and usable. + assert "up" not in t.col_names + assert list(t["name"][:]) == ["a", "bb"] + assert "name" in str(t) + + +def test_utf8_kernel_refused_whatever_returned(): + """It is the utf8 operand that cannot work, not the string output.""" + + @blosc2.dsl_kernel + def is_long(name): + return name > "b" + + t = make_table(["a", "bb"]) + with pytest.raises(NotImplementedError, match="cannot be a UDF"): + t.add_computed_column("flag", is_long, inputs=["name"]) + + +def test_utf8_apply_names_the_column(): + t = make_table(["a", "bb"]) + with pytest.raises(NotImplementedError) as exc: + t.apply(_shout, columns=["name"]) + _assert_routes(str(exc.value), "t['name']") + + +def test_utf8_lazyudf_over_a_column_names_the_column(): + t = make_table(["a", "bb"]) + with pytest.raises(NotImplementedError) as exc: + blosc2.lazyudf(_shout, (t["name"],)) + _assert_routes(str(exc.value), "t['name']") + + +def test_utf8_lazyudf_over_a_bare_array_routes_too(): + arr = blosc2.utf8_array(["a", "bb"]) + with pytest.raises(NotImplementedError) as exc: + blosc2.lazyudf(_shout, (arr,)) + msg = str(exc.value) + assert "blosc2.from_utf8(arr)" in msg + assert "blosc2.to_utf8(" in msg + # No table to assign into, so the message must not suggest one. + assert ".assign(" not in msg + + +def test_utf8_refusal_recipe_actually_works(): + """The recipe the error prints must run as printed.""" + t = make_table(["a", "bb"]) + fixed = blosc2.from_utf8(t["name"]) + res = blosc2.lazyexpr("upper(name)", {"name": fixed}).compute()[:] + t.add_column("out", blosc2.utf8(), values=blosc2.to_utf8(res)) + assert list(t["out"][:]) == ["A", "BB"] + + res = blosc2.lazyudf(_shout, (fixed,)).compute()[:] + assert list(blosc2.to_utf8(res)[:]) == ["A", "BB"] + + +def test_non_utf8_dsl_kernel_column_still_works(): + """The guard must not catch ordinary columns.""" + + @blosc2.dsl_kernel + def double(x): + return x * 2 + + t = make_table(["a", "bb"]) + t.add_computed_column("dbl", double, inputs=["x"]) + np.testing.assert_array_equal(t["dbl"][:], [0, 2]) + + +# --------------------------------------------------------------------------- +# NumPy StringDType interop: dtype-based dispatch to UTF8Array +# --------------------------------------------------------------------------- + + +def test_utf8_array_satisfies_array_protocol(): + arr = blosc2.utf8_array(["a", "bb", "ccc"]) + assert isinstance(arr, blosc2.Array) + assert arr.shape == (3,) + assert arr.ndim == 1 + assert arr.size == 3 + assert arr.dtype == STRING_DTYPE + + +def test_utf8_array_np_asarray_keeps_string_dtype(): + """np.asarray() used to iterate the rows and infer a fixed-width 1)', [2]), + # Both leaves at once, and the shorter name is a prefix of the longer. + ('(trip.begin.who == "bob") | (trip.who == "CAROL")', [1, 2]), + ], +) +def test_ctable_utf8_nested_leaf_filters(expr, expected): + # Dotted utf8 leaves are outside the operand namespace, so they reach the + # utf8 driver still spelled with dots -- which no expression engine parses. + t = _nested_table() + assert list(t.where(expr)["x"][:]) == expected + + +def test_ctable_utf8_nested_leaf_matches_flat(): + """A dotted name must not change the answer the same data gives flat.""" + values = ["hello", "help", "world", "zz"] + flat = make_table(values) + nested = make_table(values) + nested.rename_column("name", "trip.who") + for flat_expr, nested_expr in ( + ("name == 'hello'", 'trip.who == "hello"'), + ("startswith(name, 'hel')", 'startswith(trip.who, "hel")'), + ("name < 'w'", 'trip.who < "w"'), + ): + assert list(flat.where(flat_expr)["x"][:]) == list(nested.where(nested_expr)["x"][:]) + + +def test_ctable_utf8_nested_leaf_sum_persist(tmp_path): + urlpath = str(tmp_path / "utf8_nested.b2z") + t = _nested_table(urlpath=urlpath, mode="w") + assert t["x"].sum(where='startswith(trip.begin.who, "c")') == 2 + t.close() + + reopened = CTable.open(urlpath, mode="r") + try: + assert list(reopened.where('trip.begin.who == "dave"')["x"][:]) == [3] + finally: + reopened.close() diff --git a/tests/ctable/test_varlen_columns.py b/tests/ctable/test_varlen_columns.py index 4a527d6ea..4abb5a871 100644 --- a/tests/ctable/test_varlen_columns.py +++ b/tests/ctable/test_varlen_columns.py @@ -48,7 +48,7 @@ def test_list_column_display(): assert "['x', 'y']" not in col_text -def test_ctable_varlen_where_select_head_tail_and_compact(): +def test_varlen_where_select_head_tail_compact(): t = blosc2.CTable(Product, new_data=DATA) view = t.where(t.qty >= 2) assert view.tags[:] == [[], None, ["z"]] diff --git a/tests/ctable/test_vlstring_vlbytes.py b/tests/ctable/test_vlstring_vlbytes.py index 8f3bb3665..62947526e 100644 --- a/tests/ctable/test_vlstring_vlbytes.py +++ b/tests/ctable/test_vlstring_vlbytes.py @@ -206,7 +206,7 @@ def test_scalar_varlen_array_nullable(): assert sva[3] is None -def test_scalar_varlen_array_rejects_none_when_not_nullable(): +def test_varlen_array_rejects_none_not_nullable(): spec = blosc2.vlstring(nullable=False) sva = _make_sva(spec) with pytest.raises(TypeError, match="not nullable"): @@ -313,7 +313,7 @@ def test_ctable_vlstring_column_is_not_list(): assert ct.text.is_varlen_scalar -def test_ctable_vlstring_column_null_count_non_nullable(): +def test_vlstring_null_count_non_nullable(): ct = blosc2.CTable(VLRow, new_data=ROWS) # Non-nullable: no Nones → null_count = 0 assert ct.text.null_count() == 0 @@ -395,7 +395,7 @@ def test_ctable_vlstring_copy_with_deletions_compact(): assert list(copied.text) == expected -def test_ctable_vlstring_copy_noncompact_preserves_tombstones(): +def test_vlstring_copy_keeps_tombstones(): ct = blosc2.CTable(VLRow, new_data=ROWS) ct.delete([1, 3]) copied = ct.copy(compact=False) @@ -441,7 +441,7 @@ def test_ctable_vlstring_backend_role_metadata(tmp_path): } -def test_ctable_constructor_reopens_vlstring_persistent_table(tmp_path): +def test_ctor_reopens_vlstring_persistent(tmp_path): urlpath = str(tmp_path / "vl_ctor_reopen.b2d") ct = blosc2.CTable(VLRow, new_data=ROWS[:2], urlpath=urlpath, mode="w") ct.close() @@ -642,3 +642,58 @@ def test_ctable_vlstring_repr(): # repr is now the tabular view (same as str); a small table shows no footer. assert r == str(ct) assert "id" in r.splitlines()[0] # column header present + + +# --------------------------------------------------------------------------- +# Column.assign +# --------------------------------------------------------------------------- + + +@dataclass +class SmallBatchRow: + text: str = blosc2.field(blosc2.vlstring(batch_rows=4)) + + +def test_ctable_vlstring_column_assign(): + ct = blosc2.CTable(VLRow, new_data=ROWS) + ct["text"].assign([f"new-{i}" for i in range(len(ROWS))]) + assert list(ct["text"][:]) == [f"new-{i}" for i in range(len(ROWS))] + # the sibling column is untouched + assert ct["data"][0] == b"bin0" + + +def test_ctable_vlbytes_column_assign(): + ct = blosc2.CTable(VLRow, new_data=ROWS) + ct["data"].assign([bytes([i]) for i in range(len(ROWS))]) + assert list(ct["data"][:]) == [bytes([i]) for i in range(len(ROWS))] + + +def test_vlstring_assign_skips_deleted_rows(): + ct = blosc2.CTable(VLRow, new_data=ROWS) + ct.delete([1, 3]) + ct["text"].assign(["p", "q", "r"]) + assert list(ct["text"][:]) == ["p", "q", "r"] + assert list(ct["id"][:]) == [0, 2, 4] + + +def test_vlstring_assign_wrong_length_raises(): + ct = blosc2.CTable(VLRow, new_data=ROWS) + with pytest.raises(ValueError, match="requires 5 values"): + ct["text"].assign(["too", "few"]) + + +def test_ctable_vlstring_column_assign_spans_batches(): + """set_all() rewrites each backing batch once, so cross-batch rows must land.""" + n = 23 # several full batches of 4, plus a partial one + ct = blosc2.CTable(SmallBatchRow, new_data={"text": [f"v{i}" for i in range(n)]}) + ct["text"].assign([f"w{i}" for i in range(n)]) + assert list(ct["text"][:]) == [f"w{i}" for i in range(n)] + + +def test_ctable_vlstring_column_assign_persists(tmp_path): + path = str(tmp_path / "vl_assign.b2d") + ct = blosc2.CTable(VLRow, urlpath=path, mode="w", new_data=ROWS) + ct["text"].assign([f"new-{i}" for i in range(len(ROWS))]) + ct.close() + ct2 = blosc2.CTable.open(path) + assert list(ct2["text"][:]) == [f"new-{i}" for i in range(len(ROWS))] diff --git a/tests/ctable/test_where_expressions.py b/tests/ctable/test_where_expressions.py index 0850ef9a5..8f4c9d3a3 100644 --- a/tests/ctable/test_where_expressions.py +++ b/tests/ctable/test_where_expressions.py @@ -43,7 +43,7 @@ def test_where_column_arithmetic_can_be_composed(): np.testing.assert_array_equal(view.value[:], np.array([20, 30, 2], dtype=np.int32)) -def test_where_column_expression_accepts_transcendental_functions(): +def test_where_col_expr_accepts_transcendentals(): t = blosc2.CTable(Row, new_data=DATA) view = t.where(((t.value + 2) * blosc2.sin(t.category)) >= 10) @@ -51,7 +51,7 @@ def test_where_column_expression_accepts_transcendental_functions(): np.testing.assert_array_equal(view.value[:], np.array([10, 20], dtype=np.int32)) -def test_where_string_expression_accepts_transcendental_functions(): +def test_where_str_expr_accepts_transcendentals(): t = blosc2.CTable(Row, new_data=DATA) view = t.where("(value + 2) * sin(category) >= 10") @@ -59,7 +59,7 @@ def test_where_string_expression_accepts_transcendental_functions(): np.testing.assert_array_equal(view.value[:], np.array([10, 20], dtype=np.int32)) -def test_where_string_expression_can_reference_computed_columns(): +def test_where_str_expr_uses_computed_cols(): t = blosc2.CTable(Row, new_data=DATA) t.add_computed_column("score", "value * category") diff --git a/tests/ndarray/test_dsl_kernels.py b/tests/ndarray/test_dsl_kernels.py index a001a474e..0023147f5 100644 --- a/tests/ndarray/test_dsl_kernels.py +++ b/tests/ndarray/test_dsl_kernels.py @@ -235,7 +235,7 @@ def test_dsl_kernel_loop_kept_as_full_dsl_function(): np.testing.assert_allclose(res[...], expected, rtol=1e-5, atol=1e-6) -def test_dsl_kernel_integer_ops_kept_as_full_dsl_function(): +def test_kernel_integer_ops_kept_as_dsl(): assert kernel_integer_ops.dsl_source is not None assert "def kernel_integer_ops(x, y):" in kernel_integer_ops.dsl_source assert kernel_integer_ops.input_names == ["x", "y"] @@ -286,7 +286,7 @@ def wrapped_set_pref_expr(self, expression, inputs, fp_accuracy, aux_reduc=None, assert res.shape == shape -def test_dsl_kernel_with_no_inputs_works_with_explicit_shape(): +def test_kernel_no_inputs_with_explicit_shape(): assert kernel_index_ramp_no_inputs.dsl_source is not None assert "def kernel_index_ramp_no_inputs():" in kernel_index_ramp_no_inputs.dsl_source assert kernel_index_ramp_no_inputs.input_names == [] @@ -308,12 +308,12 @@ def test_dsl_kernel_with_no_inputs_sum_returns_scalar(): np.testing.assert_allclose(result, expected, rtol=0.0, atol=0.0) -def test_dsl_kernel_with_no_inputs_requires_shape_or_out(): +def test_kernel_no_inputs_needs_shape_or_out(): with pytest.raises(ValueError, match="shape"): _ = blosc2.lazyudf(kernel_index_ramp_no_inputs, (), dtype=np.float32) -def test_dsl_kernel_with_no_inputs_handles_windows_dtype_policy(monkeypatch): +def test_kernel_no_inputs_windows_dtype_policy(monkeypatch): import importlib lazyexpr_mod = importlib.import_module("blosc2.lazyexpr") @@ -326,7 +326,7 @@ def test_dsl_kernel_with_no_inputs_handles_windows_dtype_policy(monkeypatch): np.testing.assert_equal(res, expected) -def test_dsl_kernel_index_symbols_float_cast_matches_expected_ramp(): +def test_kernel_index_float_cast_matches_ramp(): shape = (32, 5) x2 = blosc2.zeros(shape, dtype=np.float32) expr = blosc2.lazyudf(kernel_index_ramp_float_cast, (x2,), dtype=np.float32) @@ -335,7 +335,7 @@ def test_dsl_kernel_index_symbols_float_cast_matches_expected_ramp(): np.testing.assert_allclose(res, expected, rtol=0.0, atol=0.0) -def test_dsl_kernel_index_symbols_float_cast_uses_miniexpr_fast_path(monkeypatch): +def test_kernel_index_float_cast_fast_path(monkeypatch): original_set_pref_expr = blosc2.NDArray._set_pref_expr captured = {"calls": 0, "expr": None} @@ -364,7 +364,7 @@ def wrapped_set_pref_expr(self, expression, inputs, fp_accuracy, aux_reduc=None, ) -def test_dsl_kernel_index_symbols_int_cast_matches_expected_ramp(): +def test_kernel_index_int_cast_matches_ramp(): shape = (32, 5) x2 = blosc2.zeros(shape, dtype=np.float32) expr = blosc2.lazyudf(kernel_index_ramp_int_cast, (x2,), dtype=np.int64) @@ -390,7 +390,7 @@ def test_dsl_kernel_bool_cast_numeric_matches_expected(): np.testing.assert_equal(res, expected) -def test_dsl_kernel_full_control_flow_kept_as_dsl_function(): +def test_kernel_control_flow_kept_as_dsl(): assert kernel_control_flow_full.dsl_source is not None assert "def kernel_control_flow_full(x, y):" in kernel_control_flow_full.dsl_source assert "for i in range(4):" in kernel_control_flow_full.dsl_source @@ -453,7 +453,7 @@ def test_dsl_kernel_accepts_scalar_param_per_call(): np.testing.assert_allclose(res[...], expected, rtol=1e-5, atol=1e-6) -def test_dsl_kernel_scalar_param_keeps_miniexpr_fast_path(monkeypatch): +def test_kernel_scalar_param_keeps_fast_path(monkeypatch): import importlib lazyexpr_mod = importlib.import_module("blosc2.lazyexpr") @@ -501,7 +501,7 @@ def wrapped_set_pref_expr(self, expression, inputs, fp_accuracy, aux_reduc=None, lazyexpr_mod.try_miniexpr = old_try_miniexpr -def test_dsl_kernel_scalar_float_cast_inlined_without_float_call(monkeypatch): +def test_kernel_scalar_float_cast_inlined(monkeypatch): import importlib lazyexpr_mod = importlib.import_module("blosc2.lazyexpr") @@ -536,7 +536,7 @@ def wrapped_set_pref_expr(self, expression, inputs, fp_accuracy, aux_reduc=None, lazyexpr_mod.try_miniexpr = old_try_miniexpr -def test_dsl_kernel_scalar_only_inputs_route_through_fast_eval(monkeypatch): +def test_kernel_scalar_only_via_fast_eval(monkeypatch): import importlib lazyexpr_mod = importlib.import_module("blosc2.lazyexpr") @@ -568,7 +568,8 @@ def fail_slices_eval(*args, **kwargs): np.testing.assert_equal(res[...], np.zeros(shape, dtype=np.float32)) -def test_dsl_kernel_scalar_only_inputs_specialization_injects_dummy_operand(monkeypatch): +def test_kernel_scalar_only_injects_dummy(monkeypatch): + """Scalar-only inputs specialize to a kernel with a dummy operand injected.""" import importlib lazyexpr_mod = importlib.import_module("blosc2.lazyexpr") @@ -600,7 +601,7 @@ def failing_set_pref_expr(self, expression, inputs, fp_accuracy, aux_reduc=None, lazyexpr_mod.try_miniexpr = old_try_miniexpr -def test_dsl_kernel_two_scalar_params_start_step_linear_ramp(): +def test_kernel_two_scalar_params_ramp(): shape = (9, 7) start = np.float32(2.5) step = np.float32(0.75) @@ -612,7 +613,7 @@ def test_dsl_kernel_two_scalar_params_start_step_linear_ramp(): np.testing.assert_allclose(res[...], expected, rtol=0.0, atol=0.0) -def test_dsl_kernel_three_scalar_params_start_stop_nitems_ramp(): +def test_kernel_three_scalar_params_ramp(): shape = (20, 25) start = np.float64(1.0) stop = np.float64(2.0) @@ -628,7 +629,7 @@ def test_dsl_kernel_three_scalar_params_start_stop_nitems_ramp(): np.testing.assert_allclose(res[...], expected, rtol=0.0, atol=0.0) -def test_dsl_kernel_float_cast_with_negative_scalar_param(): +def test_kernel_float_cast_negative_scalar(): shape = (10, 100) start = -10 stop = 10 @@ -643,7 +644,7 @@ def test_dsl_kernel_float_cast_with_negative_scalar_param(): np.testing.assert_allclose(res[...], expected, rtol=1e-6, atol=1e-6) -def test_dsl_kernel_float_cast_with_flat_idx_no_segfault_subprocess(): +def test_kernel_float_cast_flat_idx_no_crash(): if blosc2.IS_WASM: pytest.skip("subprocess is not supported on emscripten/wasm32") @@ -679,7 +680,7 @@ def kernel(start, stop, nitems): assert "ok" in result.stdout -def test_dsl_kernel_scalar_constant_subexpr_runtime_no_segfault(tmp_path): +def test_kernel_scalar_const_subexpr_no_crash(tmp_path): if blosc2.IS_WASM: pytest.skip("subprocess is not supported on emscripten/wasm32") @@ -709,7 +710,7 @@ def kernel_const_subexpr(x, start, step): assert "ok" in result.stdout -def test_dsl_kernel_miniexpr_failure_raises_even_with_strict_disabled(monkeypatch): +def test_kernel_failure_raises_strict_off(monkeypatch): import importlib lazyexpr_mod = importlib.import_module("blosc2.lazyexpr") @@ -736,7 +737,7 @@ def failing_set_pref_expr(self, expression, inputs, fp_accuracy, aux_reduc=None, lazyexpr_mod.try_miniexpr = old_try_miniexpr -def test_dsl_kernel_miniexpr_failure_includes_backend_error_details(monkeypatch): +def test_kernel_failure_includes_backend_error(monkeypatch): import importlib lazyexpr_mod = importlib.import_module("blosc2.lazyexpr") @@ -761,7 +762,7 @@ def failing_set_pref_expr(self, expression, inputs, fp_accuracy, aux_reduc=None, lazyexpr_mod.try_miniexpr = old_try_miniexpr -def test_dsl_kernel_miniexpr_failure_prefers_validate_dsl_error(monkeypatch): +def test_kernel_failure_prefers_validate_error(monkeypatch): import importlib lazyexpr_mod = importlib.import_module("blosc2.lazyexpr") @@ -843,7 +844,7 @@ def test_jit_backend_pragma_wrapping_dsl_source(): kernel_fallback_tuple_assign, ], ) -def test_dsl_kernel_flawed_syntax_detected_fallback_callable(kernel): +def test_kernel_flawed_syntax_falls_back(kernel): assert kernel.dsl_source is not None assert kernel.input_names == ["x", "y"] assert kernel.dsl_error is not None @@ -859,7 +860,7 @@ def test_dsl_kernel_flawed_syntax_detected_fallback_callable(kernel): ) -def test_dsl_kernel_ternary_rejected_with_actionable_error(): +def test_kernel_ternary_rejected_with_hint(): assert kernel_fallback_ternary.dsl_source is not None assert kernel_fallback_ternary.input_names == ["x"] assert kernel_fallback_ternary.dsl_error is not None @@ -1061,7 +1062,7 @@ def test_dsl_save_dictstore_operands(tmp_path): # G3 (variable name colliding with miniexpr codegen identifier) --- -def test_dsl_kernel_semicolon_joined_statements_rejected(): +def test_kernel_semicolon_statements_rejected(): # Source built from a string so the formatter cannot rewrite the ';'-join away. result = validate_dsl( kernel_from_source("def k(a, b):\n x = a * a; y = b * b\n return x + y\n", "k") @@ -1178,7 +1179,7 @@ def _numpy_operand_kernel(x, y): (13, 17, 19), # 3-D: odd shape ], ) -def test_dsl_kernel_numpy_operands_match_ndarray_reference(shape): +def test_kernel_numpy_operands_match_ndarray(shape): rng = np.random.default_rng(0) a = rng.random(shape).astype(np.float64) b = rng.random(shape).astype(np.float64) @@ -1186,7 +1187,7 @@ def test_dsl_kernel_numpy_operands_match_ndarray_reference(shape): np.testing.assert_array_equal(res, _dsl_reference(_numpy_operand_kernel, (a, b))) -def test_dsl_kernel_numpy_operands_mixed_dtype_promotes_output(): +def test_kernel_numpy_mixed_dtype_promotes(): rng = np.random.default_rng(1) a = (rng.random(10_007) * 10).astype(np.float32) b = (rng.random(10_007) * 10).astype(np.int64) @@ -1196,7 +1197,7 @@ def test_dsl_kernel_numpy_operands_mixed_dtype_promotes_output(): np.testing.assert_array_equal(res, ref) -def test_dsl_kernel_ndarray_operands_with_different_itemsize(): +def test_kernel_ndarray_different_itemsize(): # Blocks are sized in bytes, so a float32 and an int64 operand get different # chunks/blocks by default; the DSL path has no slow fallback, so it used to # raise "slicing is not supported" whenever the grids diverged (which depends @@ -1219,7 +1220,7 @@ def test_dsl_kernel_mixed_ndarray_and_numpy_operand(): np.testing.assert_array_equal(res, ref) -def test_dsl_kernel_numpy_operands_f_ordered_and_strided(): +def test_kernel_numpy_f_ordered_and_strided(): shape = (20, 10) b = np.arange(np.prod(shape), dtype=np.float64).reshape(shape) ref = _dsl_reference(_numpy_operand_kernel, (b, b)) @@ -1234,14 +1235,14 @@ def test_dsl_kernel_numpy_operands_f_ordered_and_strided(): np.testing.assert_array_equal(res_strided, ref_strided) -def test_dsl_kernel_numpy_operand_non_native_endian_requires_miniexpr(): +def test_kernel_non_native_endian_needs_miniexpr(): a = np.arange(100, dtype=">f8").reshape(10, 10) b = np.arange(100, dtype=np.float64).reshape(10, 10) with pytest.raises(RuntimeError, match="NDArray or NumPy inputs"): blosc2.lazyudf(_numpy_operand_kernel, (a, b), dtype=None)[()] -def test_dsl_kernel_zero_input_dummy_operand_injection_still_works(): +def test_kernel_zero_input_dummy_injection(): @blosc2.dsl_kernel def ramp(start, step): return start + step * _i0 # noqa: F821 # DSL index symbol resolved by miniexpr @@ -1251,7 +1252,8 @@ def ramp(start, step): np.testing.assert_allclose(res, expected) -def test_dsl_kernel_numpy_out_matches_compute_and_honors_explicit_cparams(): +def test_kernel_numpy_out_matches_compute(): + """A NumPy `out` matches compute(), and an explicit cparams is honoured.""" a = np.arange(1000, dtype=np.float64) b = np.arange(1000, dtype=np.float64) * 0.5 lexpr = blosc2.lazyudf(_numpy_operand_kernel, (a, b), dtype=None) @@ -1264,7 +1266,7 @@ def test_dsl_kernel_numpy_out_matches_compute_and_honors_explicit_cparams(): assert res_explicit.schunk.cparams.clevel == 5 -def test_dsl_kernel_numpy_attribute_calls_are_rewritten_to_bare_names(): +def test_kernel_numpy_attr_calls_rewritten(): @blosc2.dsl_kernel def k(x, y): if x >= 0: @@ -1294,7 +1296,7 @@ def k(x, y): ("np.absolute(x)", "abs"), ], ) -def test_dsl_kernel_numpy_func_aliases_map_to_dsl_names(numpy_call, expected_dsl_name): +def test_kernel_numpy_aliases_map_to_dsl(numpy_call, expected_dsl_name): src = f"def k(x):\n if x >= 0:\n return {numpy_call}\n else:\n return -x\n" k = kernel_from_source(src) @@ -1307,7 +1309,7 @@ def test_dsl_kernel_numpy_func_aliases_map_to_dsl_names(numpy_call, expected_dsl np.testing.assert_allclose(res, expected) -def test_dsl_kernel_numpy_alias_not_rewritten_when_shadowed_by_parameter(): +def test_kernel_numpy_alias_shadowed_by_param(): # A parameter literally named "np" shadows the module -- the rewrite must # not mistake a per-call NDArray/scalar input for the NumPy module. src = "def k(np, y):\n return np * y\n" @@ -1320,7 +1322,7 @@ def test_dsl_kernel_numpy_alias_not_rewritten_when_shadowed_by_parameter(): np.testing.assert_allclose(res, a * b) -def test_dsl_kernel_numpy_call_without_alias_left_untouched(): +def test_kernel_numpy_call_no_alias_untouched(): # No import of numpy bound in the kernel's defining scope -- nothing to # rewrite, and the plain bare-name form still works unaffected. @blosc2.dsl_kernel diff --git a/tests/ndarray/test_getitem.py b/tests/ndarray/test_getitem.py index 2d4581de0..92fdede28 100644 --- a/tests/ndarray/test_getitem.py +++ b/tests/ndarray/test_getitem.py @@ -157,7 +157,7 @@ def test_lazyexpr_where_full_slice_no_recursion(): np.testing.assert_allclose(a[a < 5][:], expected) -def test_lazyexpr_where_full_slice_persisted_reuses_shared_chunk_cache(tmp_path): +def test_where_full_slice_reuses_shared_cache(tmp_path): nitems = 60_000 expected = np.linspace(0, 1, nitems) a = blosc2.asarray( @@ -172,7 +172,7 @@ def test_lazyexpr_where_full_slice_persisted_reuses_shared_chunk_cache(tmp_path) blosc2.set_nthreads(old_nthreads) -def test_lazyexpr_where_full_slice_cached_repeat_avoids_full_mask_scan(monkeypatch): +def test_where_full_slice_repeat_avoids_scan(monkeypatch): nitems = 60_000 expected = np.arange(5, dtype=np.int64) a = blosc2.asarray(np.arange(nitems, dtype=np.int64), chunks=(20_000,)) @@ -188,7 +188,7 @@ def test_lazyexpr_where_full_slice_cached_repeat_avoids_full_mask_scan(monkeypat @pytest.mark.parametrize("mode", ["r", "a"]) -def test_lazyexpr_where_full_slice_persistent_uses_hot_cache_without_persisting(tmp_path, monkeypatch, mode): +def test_where_full_slice_hot_cache_no_persist(tmp_path, monkeypatch, mode): nitems = 60_000 expected = np.arange(5, dtype=np.int64) urlpath = tmp_path / "persisted_readonly.b2nd" @@ -303,7 +303,8 @@ def test_take_1d_sparse_path_negative_indices(): np.testing.assert_array_equal(a[idx], npa[idx]) -def test_take_1d_sparse_path_structured_non_behaved_partitions(): +def test_take_1d_sparse_structured_partitions(): + """The 1-D sparse path, on structured dtypes with non-behaved partitions.""" npa = np.empty((100,), dtype=[("a", np.int32), ("b", np.int32)]) npa["a"] = np.arange(1, 101) npa["b"] = np.arange(200, 100, -1) @@ -328,7 +329,7 @@ def test_ndarray_take_1d_matches_numpy(): np.testing.assert_array_equal(result[()], np.take(npa, idx)) -def test_ndarray_take_axis_with_nd_indices_matches_numpy(): +def test_take_axis_nd_indices_matches_numpy(): npa = np.arange(3 * 4 * 5, dtype=np.int32).reshape(3, 4, 5) a = blosc2.asarray(npa, chunks=(2, 2, 3)) idx = np.array([[3, 0], [1, -1]], dtype=np.int64) @@ -342,7 +343,7 @@ def test_ndarray_take_axis_with_nd_indices_matches_numpy(): np.testing.assert_array_equal(top_level_result[()], expected) -def test_ndarray_take_axis_none_nd_fallback_matches_numpy(): +def test_take_axis_none_nd_matches_numpy(): npa = np.arange(3 * 4 * 5, dtype=np.int32).reshape(3, 4, 5) a = blosc2.asarray(npa, chunks=(2, 2, 3)) idx = np.array([[0, -1], [17, 5]], dtype=np.int64) @@ -648,7 +649,7 @@ def test_getitem_integer_array_out_of_bounds(): _ = a[[-4]] -def test_getitem_integer_array_still_uses_fancy_for_boolean(): +def test_getitem_int_array_fancy_for_boolean(): """Boolean arrays should NOT be routed through the sparse path.""" a = blosc2.asarray(np.arange(12, dtype=np.int32).reshape(3, 4)) mask = np.array([True, False, True]) diff --git a/tests/ndarray/test_indexing.py b/tests/ndarray/test_indexing.py index 6aa54d91e..e8606c04a 100644 --- a/tests/ndarray/test_indexing.py +++ b/tests/ndarray/test_indexing.py @@ -47,7 +47,7 @@ def test_scalar_index_matches_scan(kind): np.testing.assert_array_equal(indexed, data[(data >= 120_000) & (data < 125_000)]) -def test_opsi_index_accepts_non_multiple_chunk_and_block_lengths(): +def test_opsi_accepts_non_multiple_chunk_block(): rng = np.random.default_rng(42) data = rng.random(5_000, dtype=np.float64) arr = blosc2.asarray(data, chunks=(781,), blocks=(160,)) @@ -105,7 +105,7 @@ def test_opsi_optlevel_controls_chunk_multiplier(optlevel, expected_multiplier): (9, 4), ], ) -def test_chunk_local_indexes_optlevel_controls_chunk_multiplier(kind, optlevel, expected_multiplier): +def test_chunk_local_optlevel_sets_multiplier(kind, optlevel, expected_multiplier): rng = np.random.default_rng(44) data = rng.integers(0, 100_000, size=20_000, dtype=np.int64) arr = blosc2.asarray(data, chunks=(1_000,), blocks=(200,)) @@ -142,7 +142,7 @@ def test_structured_field_index_matches_scan(kind): np.testing.assert_array_equal(indexed, data[(data["id"] >= 48_000) & (data["id"] < 51_000)]) -def test_module_level_will_use_index_matches_lazyexpr_method(): +def test_module_will_use_index_matches_method(): import blosc2.indexing as indexing indexed = blosc2.asarray(np.arange(100_000, dtype=np.int64), chunks=(10_000,), blocks=(2_000,)) @@ -218,7 +218,8 @@ def test_index_accessor_compact_updates_live_view(tmp_path): assert reopened.index("a")["full"]["runs"] == [] -def test_gather_positions_by_block_avoids_whole_chunk_fallback_for_multi_block_reads(monkeypatch): +def test_gather_by_block_avoids_chunk_fallback(monkeypatch): + """A read spanning several blocks gathers by block, not by whole chunk.""" import blosc2.indexing as indexing class FakeSource: @@ -265,7 +266,10 @@ def test_random_field_index_matches_scan(kind): arr.create_index(field="id", kind=_public_kind(kind)) expr = blosc2.lazyexpr("(id >= 70_000) & (id < 71_200)", arr.fields).where(arr) - assert expr.will_use_index() is True + # A shuffled column spreads the matches over every block, so a bucket mask + # prunes nothing worth reading and the planner declines it in favour of the + # scan. partial and full produce exact positions and are unaffected. + assert expr.will_use_index() is (kind != "bucket") indexed = expr.compute()[:] scanned = expr.compute(_use_index=False)[:] @@ -340,7 +344,7 @@ def test_bucket_numeric_dtype_query_matches_scan(dtype): np.testing.assert_array_equal(indexed, expected) -def test_numeric_unsupported_dtype_fallback_matches_scan(): +def test_unsupported_dtype_fallback_vs_scan(): values = (np.arange(2_000, dtype=np.float16) / np.float16(10)).astype(np.float16) arr = blosc2.asarray(values, chunks=(500,), blocks=(100,)) @@ -354,11 +358,11 @@ def test_numeric_unsupported_dtype_fallback_matches_scan(): def test_bucket_lossy_integer_values_match_scan(): - rng = np.random.default_rng(2) dtype = np.dtype([("id", np.int64), ("payload", np.float32)]) data = np.zeros(180_000, dtype=dtype) + # Ordered, so the matches sit in few blocks and the bucket evaluator — the + # thing under test — is actually reached rather than declined for a scan. data["id"] = np.arange(-90_000, 90_000, dtype=np.int64) - rng.shuffle(data["id"]) arr = blosc2.asarray(data, chunks=(18_000,), blocks=(3_000,)) descriptor = arr.create_index(field="id", kind=blosc2.IndexKind.BUCKET, optlevel=0) @@ -375,11 +379,10 @@ def test_bucket_lossy_integer_values_match_scan(): def test_bucket_lossy_float_values_match_scan(): - rng = np.random.default_rng(3) dtype = np.dtype([("x", np.float64), ("payload", np.float32)]) data = np.zeros(160_000, dtype=dtype) + # Ordered for the same reason as the integer case above. data["x"] = np.linspace(-5000.0, 5000.0, data.shape[0], dtype=np.float64) - rng.shuffle(data["x"]) arr = blosc2.asarray(data, chunks=(16_000,), blocks=(4_000,)) descriptor = arr.create_index(field="x", kind=blosc2.IndexKind.BUCKET, optlevel=0) @@ -497,7 +500,7 @@ def test_cross_column_exact_refinement_with_full_index(tmp_path, persistent): np.testing.assert_array_equal(indexed, expected) -def test_summary_threaded_downstream_order_matches_scan(monkeypatch): +def test_summary_threaded_order_matches_scan(monkeypatch): dtype = np.dtype([("id", np.int64), ("payload", np.int32)]) data = np.zeros(240_000, dtype=dtype) data["id"] = np.arange(data.shape[0], dtype=np.int64) @@ -537,7 +540,10 @@ def test_bucket_threaded_downstream_order_matches_scan(monkeypatch): monkeypatch.setattr(indexing, "INDEX_QUERY_MIN_CHUNKS_PER_THREAD", 1) monkeypatch.setattr(blosc2, "nthreads", 4) - expr = blosc2.lazyexpr("(id >= 60_000) & (id < 180_000)", arr.fields).where(arr) + # 4 of the 20 chunks: enough to fan out over the thread pool, few enough + # blocks that the plan is worth taking. A 50% span reads over half the + # column's blocks, at which point the planner rightly prefers the scan. + expr = blosc2.lazyexpr("(id >= 60_000) & (id < 108_000)", arr.fields).where(arr) explanation = expr.explain() assert explanation["will_use_index"] is True @@ -545,7 +551,7 @@ def test_bucket_threaded_downstream_order_matches_scan(monkeypatch): indexed = expr.compute()[:] scanned = expr.compute(_use_index=False)[:] - expected = data[(data["id"] >= 60_000) & (data["id"] < 180_000)] + expected = data[(data["id"] >= 60_000) & (data["id"] < 108_000)] np.testing.assert_array_equal(indexed, scanned) np.testing.assert_array_equal(indexed, expected) @@ -581,7 +587,7 @@ def test_persistent_index_survives_reopen(tmp_path, kind): @pytest.mark.parametrize("kind", ["bucket", "partial", "full"]) -def test_default_ooc_persistent_index_matches_scan_and_rebuilds(tmp_path, kind): +def test_ooc_persistent_matches_scan_rebuilds(tmp_path, kind): path = tmp_path / f"indexed_ooc_{kind}.b2nd" rng = np.random.default_rng(7) dtype = np.dtype([("id", np.int64), ("payload", np.float32)]) @@ -611,7 +617,7 @@ def test_default_ooc_persistent_index_matches_scan_and_rebuilds(tmp_path, kind): @pytest.mark.parametrize("kind", ["bucket", "partial"]) -def test_persistent_chunk_local_ooc_builds_do_not_use_temp_memmap(tmp_path, kind): +def test_persistent_chunk_ooc_no_temp_memmap(tmp_path, kind): path = tmp_path / f"persistent_no_memmap_{kind}.b2nd" data = np.arange(120_000, dtype=np.int64) indexing = __import__("blosc2.indexing", fromlist=["_segment_row_count"]) @@ -631,7 +637,7 @@ def test_persistent_chunk_local_ooc_builds_do_not_use_temp_memmap(tmp_path, kind @pytest.mark.parametrize("kind", ["bucket", "partial"]) -def test_in_memory_chunk_local_ooc_builds_do_not_use_temp_memmap(kind): +def test_in_memory_chunk_ooc_no_temp_memmap(kind): data = np.arange(120_000, dtype=np.int64) indexing = __import__("blosc2.indexing", fromlist=["_segment_row_count"]) assert not hasattr(indexing, "_open_temp_memmap") @@ -701,7 +707,7 @@ def test_in_mem_override_disables_ooc_builder(kind): @pytest.mark.parametrize("use_expression", [False, True]) -def test_ultralight_ooc_build_does_not_materialize_full_target(monkeypatch, tmp_path, use_expression): +def test_ultralight_ooc_no_full_materialize(monkeypatch, tmp_path, use_expression): path = tmp_path / ("indexed_expr_ultralight.b2nd" if use_expression else "indexed_ultralight.b2nd") if use_expression: data = np.zeros(120_000, dtype=[("x", np.int64)]) @@ -724,7 +730,8 @@ def fail_values_for_target(array, target): @pytest.mark.parametrize("kind", ["bucket", "partial"]) -def test_chunk_local_ooc_intra_chunk_build_uses_thread_pool_when_threads_forced(monkeypatch, kind): +def test_intra_chunk_ooc_uses_thread_pool(monkeypatch, kind): + """The intra-chunk OOC build uses the thread pool when threads are forced.""" if blosc2.IS_WASM: pytest.skip("wasm32 does not use Python thread pools for index building") data = np.arange(48_000, dtype=np.int64) @@ -756,7 +763,7 @@ def map(self, fn, iterable): @pytest.mark.parametrize("kind", ["bucket", "partial"]) -def test_in_memory_chunk_local_build_uses_cparams_nthreads(monkeypatch, kind): +def test_in_memory_chunk_uses_cparams_threads(monkeypatch, kind): if blosc2.IS_WASM: pytest.skip("wasm32 does not use Python thread pools for index building") data = np.arange(48_000, dtype=np.int64) @@ -807,7 +814,7 @@ def test_persistent_chunk_local_sidecars_use_cparams(tmp_path, kind): assert sidecar.cparams.clevel == 2 -def test_intra_chunk_sort_run_matches_numpy_stable_order(): +def test_intra_chunk_sort_matches_np_stable(): indexing_ext = __import__("blosc2.indexing_ext", fromlist=["intra_chunk_sort_run"]) values = np.array([4.0, np.nan, 2.0, 2.0, np.nan, 1.0, 4.0], dtype=np.float64) @@ -818,7 +825,7 @@ def test_intra_chunk_sort_run_matches_numpy_stable_order(): np.testing.assert_array_equal(positions, order.astype(np.uint16, copy=False)) -def test_intra_chunk_merge_sorted_slices_matches_lexsort_merge(): +def test_intra_chunk_merge_matches_lexsort(): indexing_ext = __import__("blosc2.indexing_ext", fromlist=["intra_chunk_merge_sorted_slices"]) left_values = np.array([1.0, 2.0, 2.0, np.nan], dtype=np.float64) left_positions = np.array([0, 2, 3, 6], dtype=np.uint16) @@ -836,7 +843,7 @@ def test_intra_chunk_merge_sorted_slices_matches_lexsort_merge(): np.testing.assert_array_equal(merged_positions, all_positions[order]) -def test_intra_chunk_merge_sorted_slices_validates_lengths(): +def test_intra_chunk_merge_validates_lengths(): indexing_ext = __import__("blosc2.indexing_ext", fromlist=["intra_chunk_merge_sorted_slices"]) values = np.array([1.0, 2.0], dtype=np.float64) positions = np.array([0, 1], dtype=np.uint16) @@ -847,7 +854,7 @@ def test_intra_chunk_merge_sorted_slices_validates_lengths(): ) -def test_index_search_boundary_bounds_validates_lengths(): +def test_search_boundary_validates_lengths(): indexing_ext = __import__("blosc2.indexing_ext", fromlist=["index_search_boundary_bounds"]) starts = np.array([1, 3], dtype=np.int64) ends = np.array([2], dtype=np.int64) @@ -856,7 +863,7 @@ def test_index_search_boundary_bounds_validates_lengths(): indexing_ext.index_search_boundary_bounds(starts, ends, None, True, None, True) -def test_mutation_marks_index_stale_and_rebuild_restores_it(): +def test_mutation_marks_stale_rebuild_restores(): data = np.arange(50_000, dtype=np.int64) arr = blosc2.asarray(data, chunks=(5_000,), blocks=(1_000,)) arr.create_index(kind=blosc2.IndexKind.FULL) @@ -873,7 +880,7 @@ def test_mutation_marks_index_stale_and_rebuild_restores_it(): assert expr.will_use_index() is True -def test_full_index_reuses_primary_order_for_indices_and_sort(): +def test_full_index_reuses_primary_order(): dtype = np.dtype([("a", np.int64), ("b", np.int64)]) data = np.array( [(2, 9), (1, 8), (2, 7), (1, 6), (2, 5), (1, 4), (2, 3), (1, 2), (2, 1), (1, 0)], @@ -909,7 +916,7 @@ def test_persistent_scalar_argsort_uses_full_index(tmp_path): np.testing.assert_array_equal(result[:], np.argsort(data, kind="stable")) -def test_filtered_ordered_queries_support_cross_field_exact_indexes(): +def test_filtered_ordered_cross_field_indexes(): dtype = np.dtype([("a", np.int64), ("b", np.int64), ("payload", np.int32)]) data = np.array( [ @@ -1065,7 +1072,8 @@ def test_persistent_full_index_runs_survive_reopen(tmp_path): np.testing.assert_array_equal(expr.compute()[:], expected[expected_mask]) -def test_persistent_compact_full_positional_query_avoids_whole_sidecar_load(monkeypatch, tmp_path): +def test_compact_positional_no_sidecar_load(monkeypatch, tmp_path): + """A positional query on a persistent compact index reads no whole sidecar.""" path = tmp_path / "full_selective_ooc.b2nd" rng = np.random.default_rng(12) data = np.arange(120_000, dtype=np.int64) @@ -1099,7 +1107,7 @@ def guarded_load(array, token, category, name, sidecar_path): ("full", {("full", "values"), ("full", "positions")}), ], ) -def test_in_memory_positional_queries_avoid_whole_loading_index_payloads(monkeypatch, kind, blocked): +def test_in_memory_positional_no_full_load(monkeypatch, kind, blocked): data = np.arange(120_000, dtype=np.int64) arr = blosc2.asarray(data, chunks=(12_000,), blocks=(2_000,)) arr.create_index(kind=_public_kind(kind)) @@ -1120,11 +1128,11 @@ def guarded_load(array, token, category, name, sidecar_path): @pytest.mark.parametrize("kind", ["bucket", "partial", "full"]) def test_expression_index_matches_scan(kind): - rng = np.random.default_rng(9) dtype = np.dtype([("x", np.int64), ("payload", np.int32)]) data = np.zeros(150_000, dtype=dtype) + # Ordered, so abs(x) puts the matches in two short runs rather than across + # every block, which is what lets the bucket plan be taken at all. data["x"] = np.arange(-75_000, 75_000, dtype=np.int64) - rng.shuffle(data["x"]) data["payload"] = np.arange(data.shape[0], dtype=np.int32) arr = blosc2.asarray(data, chunks=(15_000,), blocks=(3_000,)) @@ -1252,7 +1260,7 @@ def test_append_keeps_expression_index_current(kind): np.testing.assert_array_equal(arr.sort(order="abs(x)")[:], all_data[expected_positions]) -def test_repeated_appends_keep_full_expression_index_current(): +def test_repeated_appends_keep_expr_index(): dtype = np.dtype([("x", np.int64), ("payload", np.int32)]) data = np.array([(-10, 0), (7, 1), (-3, 2), (1, 3)], dtype=dtype) arr = blosc2.asarray(data, chunks=(2,), blocks=(2,)) @@ -1275,7 +1283,7 @@ def test_repeated_appends_keep_full_expression_index_current(): np.testing.assert_array_equal(expr.compute()[:], expected[expected_mask]) -def test_compact_full_index_clears_runs_and_preserves_results(tmp_path): +def test_compact_clears_runs_keeps_results(tmp_path): path = tmp_path / "compact_full_runs.b2nd" dtype = np.dtype([("a", np.int64), ("b", np.int64)]) data = np.array([(3, 9), (1, 8), (2, 7), (1, 6)], dtype=dtype) @@ -1316,7 +1324,7 @@ def test_compact_full_index_clears_runs_and_preserves_results(tmp_path): np.testing.assert_array_equal(expr.compute()[:], expected[expected_mask]) -def test_compact_full_expression_index_preserves_results(): +def test_compact_expr_index_keeps_results(): dtype = np.dtype([("x", np.int64), ("payload", np.int32)]) data = np.array([(-10, 0), (7, 1), (-3, 2), (1, 3)], dtype=dtype) arr = blosc2.asarray(data, chunks=(2,), blocks=(2,)) @@ -1338,7 +1346,7 @@ def test_compact_full_expression_index_preserves_results(): np.testing.assert_array_equal(expr.compute()[:], expected[expected_mask]) -def test_forced_ooc_full_index_merge_preserves_sorted_sidecars(monkeypatch, tmp_path): +def test_forced_ooc_merge_keeps_sidecars(monkeypatch, tmp_path): path = tmp_path / "forced_ooc_full_merge.b2nd" rng = np.random.default_rng(14) data = np.arange(4096, dtype=np.int64) @@ -1408,7 +1416,7 @@ def test_full_ooc_run_items_env_overrides_optlevel(monkeypatch, tmp_path, optlev assert full["ooc_run_item_budget_source"] == "env" -def test_create_index_full_ooc_defaults_tmpdir_to_array_directory(monkeypatch, tmp_path): +def test_full_ooc_tmpdir_defaults_to_array(monkeypatch, tmp_path): path = tmp_path / "default_tmpdir_full.b2nd" data = np.arange(4096, dtype=np.int64) arr = blosc2.asarray(data, urlpath=path, mode="w", chunks=(256,), blocks=(64,)) @@ -1428,7 +1436,7 @@ def tracking_temporary_directory(*args, **kwargs): assert recorded["dir"] == str(path.parent.resolve()) -def test_create_sorted_index_full_ooc_uses_explicit_tmpdir(monkeypatch, tmp_path): +def test_sorted_full_ooc_uses_given_tmpdir(monkeypatch, tmp_path): path = tmp_path / "explicit_tmpdir_full.b2nd" custom_tmpdir = tmp_path / "custom-index-tmp" custom_tmpdir.mkdir() @@ -1453,7 +1461,7 @@ def tracking_temporary_directory(*args, **kwargs): @pytest.mark.parametrize("persistent", [False, True]) -def test_compact_full_index_rebuilds_navigation_without_whole_loading(monkeypatch, tmp_path, persistent): +def test_compact_rebuilds_nav_without_full_load(monkeypatch, tmp_path, persistent): dtype = np.dtype([("a", np.int64), ("b", np.int64)]) data = np.array([(3, 9), (1, 8), (2, 7), (1, 6)], dtype=dtype) kwargs = {"chunks": (2,), "blocks": (2,)} @@ -1488,7 +1496,7 @@ def guarded_load(array, token, category, name, sidecar_path): np.testing.assert_array_equal(expr.compute()[:], expected) -def test_persistent_large_run_full_query_uses_bounded_fallback(monkeypatch, tmp_path): +def test_large_run_query_bounded_fallback(monkeypatch, tmp_path): path = tmp_path / "large_run_fallback.b2nd" dtype = np.dtype([("id", np.int64), ("payload", np.int32)]) data = np.array([(10, 0), (20, 1), (30, 2), (40, 3)], dtype=dtype) @@ -1520,7 +1528,7 @@ def guarded_load(array, token, category, name, sidecar_path): np.testing.assert_array_equal(expr.compute()[:], expected) -def test_large_run_full_expression_query_uses_bounded_fallback(monkeypatch): +def test_large_run_expr_query_bounded_fallback(monkeypatch): dtype = np.dtype([("x", np.int64), ("payload", np.int32)]) data = np.array([(-10, 0), (7, 1), (-3, 2), (1, 3)], dtype=dtype) arr = blosc2.asarray(data, chunks=(4,), blocks=(2,)) @@ -1605,7 +1613,7 @@ def test_canonical_digest_differs_on_order_change(): assert indexing._query_cache_digest(d1) != indexing._query_cache_digest(d2) -def test_canonical_digest_preserves_order_field_sequence(): +def test_canonical_digest_keeps_field_order(): d1 = indexing._normalize_query_descriptor("(id >= 3) & (id < 6)", ["__self__"], ["a", "b"]) d2 = indexing._normalize_query_descriptor("(id >= 3) & (id < 6)", ["__self__"], ["b", "a"]) assert indexing._query_cache_digest(d1) != indexing._query_cache_digest(d2) @@ -1736,7 +1744,7 @@ def test_in_memory_array_hot_cache_hit(): # --------------------------------------------------------------------------- -def test_persistent_arrays_do_not_create_query_cache_artifacts(tmp_path): +def test_persistent_arrays_no_cache_artifacts(tmp_path): arr, urlpath = _make_persistent_array(tmp_path) _clear_caches() @@ -1772,7 +1780,7 @@ def test_persistent_cache_helpers_are_disabled(tmp_path): assert not Path(indexing._query_cache_payload_path(arr)).exists() -def test_store_cached_coords_for_persistent_array_uses_hot_cache_only(tmp_path): +def test_cached_coords_use_hot_cache_only(tmp_path): arr, _ = _make_persistent_array(tmp_path, n=8_000) _clear_caches() @@ -1915,7 +1923,7 @@ def test_ordered_query_indices_cached(tmp_path, monkeypatch): np.testing.assert_array_equal(result1, result2) -def test_ordered_query_cache_distinguishes_order_sequences(tmp_path): +def test_query_cache_distinguishes_orders(tmp_path): path = tmp_path / "ordered_sequences.b2nd" dtype = np.dtype([("a", np.int64), ("b", np.int64)]) data = np.array([(1, 2), (1, 1), (2, 1), (2, 2)], dtype=dtype) @@ -2249,7 +2257,7 @@ def test_inmem_indices_cache_entries_are_dropped_on_gc(): assert indexing._HOT_CACHE == {} -def test_ondisk_indices_path_no_cross_array_hot_cache_contamination(tmp_path): +def test_ondisk_no_cross_array_cache_mixing(tmp_path): dtype = np.dtype([("id", np.int64), ("val", np.float32)]) data1 = np.empty(1_000, dtype=dtype) data2 = np.empty(1_000, dtype=dtype) diff --git a/tests/ndarray/test_jit.py b/tests/ndarray/test_jit.py index ba09af30b..70db6fedb 100644 --- a/tests/ndarray/test_jit.py +++ b/tests/ndarray/test_jit.py @@ -179,7 +179,7 @@ def reduc_std_jit_cparams(a, b, c): assert d_jit.schunk.cparams.filters == [blosc2.Filter.BITSHUFFLE] + [blosc2.Filter.NOFILTER] * 5 -def test_jit_execution_tuning_kwarg_alone_keeps_numpy_return(): +def test_tuning_kwarg_alone_keeps_numpy_return(): # jit/jit_backend/fp_accuracy tune *how* an expression runs, not what # container the result comes back in -- they must not by themselves flip # the return type from NumPy to NDArray (unlike storage kwargs). @@ -194,7 +194,7 @@ def f(a, b): np.testing.assert_allclose(res, a * 2.0 + b) -def test_jit_execution_tuning_kwarg_with_storage_kwarg_still_returns_ndarray(): +def test_tuning_plus_storage_kwarg_gives_ndarray(): @blosc2.jit(jit=False, cparams=blosc2.CParams(clevel=2)) def f(a, b): return a * 2.0 + b @@ -205,3 +205,19 @@ def f(a, b): assert isinstance(res, blosc2.NDArray) assert res.schunk.cparams.clevel == 2 np.testing.assert_allclose(res[:], a * 2.0 + b) + + +def test_numpy_return_with_both_kwarg_kinds(): + # A traced function whose return is already a NumPy array takes the + # asarray() branch, which accepts storage kwargs only -- forwarding the + # execution-tuning ones there raised instead of returning an NDArray. + @blosc2.jit(jit=False, cparams=blosc2.CParams(clevel=2)) + def f(a, b): + return np.sum(a * 2.0 + b, axis=0) + + a = np.arange(1000, dtype=np.float64).reshape(10, 100) + b = np.arange(1000, dtype=np.float64).reshape(10, 100) * 0.5 + res = f(a, b) + assert isinstance(res, blosc2.NDArray) + assert res.schunk.cparams.clevel == 2 + np.testing.assert_allclose(res[:], np.sum(a * 2.0 + b, axis=0)) diff --git a/tests/ndarray/test_jit_dsl_dispatch.py b/tests/ndarray/test_jit_dsl_dispatch.py index 6a7ba8751..e18b3a6ca 100644 --- a/tests/ndarray/test_jit_dsl_dispatch.py +++ b/tests/ndarray/test_jit_dsl_dispatch.py @@ -34,7 +34,7 @@ def _mandel_grid(): return cr, ci -def test_jit_control_flow_dispatches_to_dsl_and_matches_numpy(monkeypatch, capsys): +def test_control_flow_dispatches_to_dsl(monkeypatch, capsys): @blosc2.jit def mandel(cr, ci, max_iter): zr = 0.0 @@ -101,7 +101,7 @@ def elemwise(a, b): assert calls == [] # no control flow -> never routed through the DSL/lazyudf path -def test_jit_strict_true_on_elementwise_dsl_valid_function_uses_dsl(monkeypatch): +def test_strict_true_elementwise_uses_dsl(monkeypatch): calls = [] import blosc2.proxy as proxy_mod @@ -124,7 +124,7 @@ def elemwise(a, b): assert calls # dispatched through the DSL wrapper -def test_jit_strict_true_on_non_dsl_function_raises_at_decoration_time(): +def test_strict_true_non_dsl_raises_early(): with pytest.raises(Exception, match="axis"): @blosc2.jit(strict=True) @@ -145,7 +145,7 @@ def cf_func(a, b): np.testing.assert_allclose(res, a + b) -def test_jit_control_flow_on_python_scalar_flag_still_traces(): +def test_control_flow_scalar_flag_still_traces(): @blosc2.jit def scalar_flag(a, b, flag): if flag: @@ -158,6 +158,34 @@ def scalar_flag(a, b, flag): np.testing.assert_allclose(scalar_flag(a, b, False), a - b) +def test_trace_hint_keeps_the_original_error(): + """A hint must not replace the failure it annotates. + + The hint used to be re-raised as ``type(e)(msg)``, which assumes a + one-argument constructor; anything needing more surfaced as a TypeError + about that constructor and the real error was lost. + """ + + class TwoArgError(Exception): + def __init__(self, code, detail): + super().__init__(f"{code}: {detail}") + self.code = code + + @blosc2.jit + def cf_func(a, b, flag): + if flag: + raise TwoArgError(7, "boom") + return a - b + + a = np.arange(10, dtype=np.float64) + with pytest.raises(TwoArgError) as excinfo: + cf_func(a, a, True) + assert excinfo.value.code == 7 + assert "boom" in str(excinfo.value) + # The hint still reaches the user, as a note on the original exception. + assert any("control flow" in note for note in getattr(excinfo.value, "__notes__", [])) + + def test_jit_dsl_route_rejects_broadcasting(): @blosc2.jit def kernel(a, b, n): @@ -177,7 +205,7 @@ def _kernel_src(a, b, n): return acc -def test_jit_dsl_route_out_numpy_c_contiguous_filled_in_place(): +def test_out_numpy_contiguous_filled_in_place(): a = np.arange(1000, dtype=np.float64) b = np.arange(1000, dtype=np.float64) * 0.5 out = np.empty(1000, dtype=np.float64) @@ -187,7 +215,7 @@ def test_jit_dsl_route_out_numpy_c_contiguous_filled_in_place(): np.testing.assert_allclose(out, (a + b) * 3) -def test_jit_dsl_route_out_numpy_non_contiguous_uses_copyto_fallback(): +def test_out_numpy_non_contiguous_uses_copyto(): a = np.arange(1000, dtype=np.float64) b = np.arange(1000, dtype=np.float64) * 0.5 out = np.empty(2000, dtype=np.float64)[::2] @@ -198,7 +226,7 @@ def test_jit_dsl_route_out_numpy_non_contiguous_uses_copyto_fallback(): np.testing.assert_allclose(out, (a + b) * 3) -def test_jit_dsl_route_out_mismatched_shape_or_dtype_raises_typeerror(): +def test_out_mismatched_shape_or_dtype_raises(): a = np.arange(1000, dtype=np.float64) b = np.arange(1000, dtype=np.float64) * 0.5 @@ -209,7 +237,7 @@ def test_jit_dsl_route_out_mismatched_shape_or_dtype_raises_typeerror(): blosc2.jit(out=np.empty(1000, dtype=np.float32))(_kernel_src)(a, b, 3) -def test_jit_dsl_route_ndarray_out_raises_not_implemented_mentioning_urlpath(): +def test_ndarray_out_raises_and_names_urlpath(): a = np.arange(1000, dtype=np.float64) b = np.arange(1000, dtype=np.float64) * 0.5 nd_out = blosc2.zeros((1000,), dtype=np.float64) @@ -228,7 +256,7 @@ def test_jit_dsl_route_compute_urlpath_persists_result(tmp_path): np.testing.assert_allclose(reopened[:], (a + b) * 3) -def test_jit_dsl_route_ndarray_operands_match_numpy_operands(): +def test_ndarray_operands_match_numpy(): a = np.arange(1000, dtype=np.float64) b = np.arange(1000, dtype=np.float64) * 0.5 na = blosc2.asarray(a) @@ -239,7 +267,7 @@ def test_jit_dsl_route_ndarray_operands_match_numpy_operands(): np.testing.assert_array_equal(res_numpy, res_ndarray) -def test_jit_dsl_route_execution_tuning_kwarg_alone_keeps_numpy_return(): +def test_tuning_kwarg_alone_keeps_numpy_return(): # Same rule as the tracing route: jit/jit_backend/fp_accuracy tune execution, # not the return container, so they must not force an NDArray on their own. jit_f = blosc2.jit(jit=False)(_kernel_src) @@ -250,7 +278,7 @@ def test_jit_dsl_route_execution_tuning_kwarg_alone_keeps_numpy_return(): np.testing.assert_allclose(res, (a + b) * 3) -def test_jit_dsl_route_execution_tuning_kwarg_with_storage_kwarg_still_returns_ndarray(): +def test_tuning_plus_storage_kwarg_gives_ndarray(): jit_f = blosc2.jit(jit=False, cparams=blosc2.CParams(clevel=2))(_kernel_src) a = np.arange(1000, dtype=np.float64) b = np.arange(1000, dtype=np.float64) * 0.5 diff --git a/tests/ndarray/test_lazyexpr.py b/tests/ndarray/test_lazyexpr.py index 5055c5ab0..0ae650a1f 100644 --- a/tests/ndarray/test_lazyexpr.py +++ b/tests/ndarray/test_lazyexpr.py @@ -1578,7 +1578,7 @@ def test_numpy_funcs(array_fixture, func): pytest.skip("NumPy version has no cumulative_sum function.") -def test_lazyexpr_string_scalar_keeps_miniexpr_fast_path(monkeypatch): +def test_string_scalar_keeps_miniexpr_path(monkeypatch): import importlib lazyexpr_mod = importlib.import_module("blosc2.lazyexpr") @@ -1612,7 +1612,7 @@ def wrapped_set_pref_expr(self, expression, inputs, fp_accuracy, aux_reduc=None, lazyexpr_mod.try_miniexpr = old_try_miniexpr -def test_lazyexpr_unary_negative_literal_matches_subtraction(monkeypatch): +def test_unary_negative_literal_matches_sub(monkeypatch): import importlib lazyexpr_mod = importlib.import_module("blosc2.lazyexpr") @@ -1647,7 +1647,7 @@ def wrapped_set_pref_expr(self, expression, inputs, fp_accuracy, aux_reduc=None, lazyexpr_mod.try_miniexpr = old_try_miniexpr -def test_lazyexpr_miniexpr_failure_falls_back_by_default(monkeypatch): +def test_miniexpr_failure_falls_back_default(monkeypatch): import importlib lazyexpr_mod = importlib.import_module("blosc2.lazyexpr") diff --git a/tests/ndarray/test_linalg.py b/tests/ndarray/test_linalg.py index fec3077a0..d18a6901e 100644 --- a/tests/ndarray/test_linalg.py +++ b/tests/ndarray/test_linalg.py @@ -162,7 +162,7 @@ def test_matmul_uses_fast_path_for_supported_2d(monkeypatch, dtype): @pytest.mark.parametrize("dtype", [np.float32, np.float64]) -def test_matmul_uses_fast_path_with_multiple_inner_blocks(monkeypatch, dtype): +def test_matmul_fast_path_many_inner_blocks(monkeypatch, dtype): old_flag = utils_mod.try_miniexpr calls = _set_pref_matmul_call_recorder(monkeypatch) try: @@ -257,7 +257,7 @@ def test_matmul_falls_back_for_dtype_mismatch(monkeypatch): _toggle_miniexpr(old_flag) -def test_matmul_fast_path_limits_blas_threads_for_cblas(monkeypatch): +def test_matmul_limits_blas_threads_for_cblas(monkeypatch): old_flag = utils_mod.try_miniexpr calls = [] @@ -294,7 +294,7 @@ def __exit__(self, exc_type, exc, tb): _toggle_miniexpr(old_flag) -def test_matmul_fast_path_skips_blas_thread_limits_above_block_threshold(monkeypatch): +def test_matmul_keeps_blas_threads_over_limit(monkeypatch): old_flag = utils_mod.try_miniexpr def unexpected_threadpool_limits(*args, **kwargs): @@ -321,7 +321,7 @@ def unexpected_threadpool_limits(*args, **kwargs): _toggle_miniexpr(old_flag) -def test_matmul_fast_path_skips_blas_thread_limits_on_darwin(monkeypatch): +def test_matmul_keeps_blas_threads_on_darwin(monkeypatch): old_flag = utils_mod.try_miniexpr def unexpected_threadpool_limits(*args, **kwargs): @@ -347,7 +347,7 @@ def unexpected_threadpool_limits(*args, **kwargs): _toggle_miniexpr(old_flag) -def test_matmul_fast_path_skips_blas_thread_limits_for_non_cblas(monkeypatch): +def test_matmul_keeps_blas_threads_non_cblas(monkeypatch): old_flag = utils_mod.try_miniexpr def unexpected_threadpool_limits(*args, **kwargs): @@ -372,7 +372,8 @@ def unexpected_threadpool_limits(*args, **kwargs): _toggle_miniexpr(old_flag) -def test_matmul_fast_path_skips_blas_thread_limits_when_threadpoolctl_missing(monkeypatch): +def test_matmul_keeps_blas_threads_no_tpctl(monkeypatch): + """Without threadpoolctl installed, the BLAS thread limit is left alone.""" old_flag = utils_mod.try_miniexpr monkeypatch.setattr(blosc2_linalg, "threadpool_limits", None) monkeypatch.setattr(blosc2.blosc2_ext, "get_selected_matmul_block_backend", lambda: "cblas") diff --git a/tests/ndarray/test_ndarray.py b/tests/ndarray/test_ndarray.py index 672d4b545..270b9a7e0 100644 --- a/tests/ndarray/test_ndarray.py +++ b/tests/ndarray/test_ndarray.py @@ -103,7 +103,7 @@ def test_asarray(a): np.testing.assert_allclose(a, b[:]) -def test_asarray_ndarray_persists_copy_when_urlpath_requested(tmp_path): +def test_asarray_persists_copy_with_urlpath(tmp_path): array = blosc2.asarray(np.arange(10, dtype=np.int64), chunks=(5,), blocks=(2,)) path = tmp_path / "persisted_copy.b2nd" @@ -115,7 +115,8 @@ def test_asarray_ndarray_persists_copy_when_urlpath_requested(tmp_path): np.testing.assert_array_equal(persisted[:], array[:]) -def test_asarray_ndarray_copies_for_dtype_changes_and_rejects_copy_false(tmp_path): +def test_asarray_dtype_change_copies_or_raises(tmp_path): + """A dtype change copies; asking for copy=False with one is an error.""" array = blosc2.asarray(np.arange(10, dtype=np.int64), chunks=(5,), blocks=(2,)) cast = blosc2.asarray(array, dtype=np.float32) @@ -159,7 +160,7 @@ def test_array_copy_false_rejects_required_copy(): blosc2.array(a, dtype=np.float64, copy=False) -def test_array_copy_none_matches_asarray_for_compatible_ndarray(): +def test_array_copy_none_matches_asarray(): a = blosc2.asarray([1, 2, 3]) b = blosc2.array(a, copy=None) diff --git a/tests/ndarray/test_proxy.py b/tests/ndarray/test_proxy.py index d65e37a1f..17719b5dd 100644 --- a/tests/ndarray/test_proxy.py +++ b/tests/ndarray/test_proxy.py @@ -107,7 +107,7 @@ def test_open(urlpath, shape, chunks, blocks, slices, dtype): blosc2.remove_urlpath(proxy_urlpath) -def test_open_readonly_proxy_keeps_cache_and_source_readonly(tmp_path): +def test_readonly_proxy_keeps_both_readonly(tmp_path): source_path = tmp_path / "source.b2nd" proxy_path = tmp_path / "proxy.b2nd" data = np.arange(120, dtype=np.int32).reshape(12, 10) diff --git a/tests/ndarray/test_slice.py b/tests/ndarray/test_slice.py index 6a0e690db..3234df718 100644 --- a/tests/ndarray/test_slice.py +++ b/tests/ndarray/test_slice.py @@ -27,7 +27,7 @@ def test_detect_aligned_chunks_exact_multiple_shape(): assert detect_aligned_chunks((slice(5, 10), slice(0, 10)), (10, 20), (5, 10)) == [2] -def test_detect_aligned_chunks_non_exact_multiple_shape(): +def test_aligned_chunks_non_exact_multiple(): # The bug repro: dim 1 (100_003) isn't a multiple of its chunk (40_000), # so its true chunk count is 3, not 100_003 // 40_000 == 2. Before the # fix this returned [2] (row 0, col chunk 2) instead of the correct [3] @@ -35,14 +35,14 @@ def test_detect_aligned_chunks_non_exact_multiple_shape(): assert detect_aligned_chunks((slice(1, 2), slice(0, 40_000)), (2, 100_003), (1, 40_000)) == [3] -def test_detect_aligned_chunks_unaligned_slice_returns_empty(): +def test_aligned_chunks_unaligned_gives_empty(): # A slice boundary that isn't a chunk multiple must short-circuit to [], # regardless of the n_chunks bug (this check runs before n_chunks is # even computed). assert detect_aligned_chunks((slice(1, 2), slice(0, 40_001)), (2, 100_003), (1, 40_000)) == [] -def test_detect_aligned_chunks_middle_dim_non_exact_multiple(): +def test_aligned_chunks_middle_dim_non_exact(): # 3D, non-exact-multiple dim in the *middle* (not last) position, offset # in the first dim -- pins that the fix's multiplier chain is right in # general, not just for the 2D case (last dim == only non-first dim) @@ -51,7 +51,7 @@ def test_detect_aligned_chunks_middle_dim_non_exact_multiple(): assert detect_aligned_chunks(key, (4, 7, 5), (2, 3, 5)) == [3] -def test_detect_aligned_chunks_multiple_non_exact_multiple_dims(): +def test_aligned_chunks_several_non_exact_dims(): # Both non-first dims are non-exact-multiple at once. key = (slice(2, 4), slice(0, 3), slice(0, 4)) assert detect_aligned_chunks(key, (4, 7, 11), (2, 3, 4)) == [9] @@ -65,13 +65,13 @@ def test_detect_aligned_chunks_consecutive_true(): assert detect_aligned_chunks(key, (10, 20), (5, 10), consecutive=True) == [0, 1, 2, 3] -def test_detect_aligned_chunks_consecutive_true_not_consecutive(): +def test_aligned_chunks_consecutive_flag_false(): # Same grid, a region whose chunks are NOT consecutive in flat order. key = (slice(0, 5), slice(0, 10)) assert detect_aligned_chunks(key, (10, 30), (5, 10), consecutive=True) == [0] -def test_detect_aligned_chunks_consecutive_true_non_exact_multiple_shape(): +def test_aligned_chunks_consecutive_non_exact(): # The bug pattern (non-exact-multiple trailing dim) under # consecutive=True: before the fix, the corrupted flat indices could # come out consecutive when they shouldn't (or vice versa), since the diff --git a/tests/ndarray/test_string_output.py b/tests/ndarray/test_string_output.py new file mode 100644 index 000000000..c2a85278e --- /dev/null +++ b/tests/ndarray/test_string_output.py @@ -0,0 +1,266 @@ +####################################################################### +# Copyright (c) 2019-present, Blosc Development Team +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +####################################################################### + +"""String-valued expression results (miniexpr string output). + +The values are checked against NumPy, but the *width* matters just as much: the +output container is allocated before evaluation, so if it is sized from NumPy's +``result_type`` instead of miniexpr's inference the kernel silently truncates. +""" + +import numpy as np +import pytest + +import blosc2 +from blosc2 import blosc2_ext + +# NumPy 2.0 added the `np.strings` namespace and the `+` ufunc loop for U/S +# arrays; blosc2 still supports NumPy 1.26, where `np.char` is the equivalent +# and `arr + arr` raises "ufunc 'add' did not contain a loop". Only the +# *reference* values below need this -- what is under test runs on both. +np_strings = getattr(np, "strings", np.char) + + +def np_add(a, b): + """NumPy's own string concatenation, spelled to work on NumPy 1.26 too.""" + return np.char.add(a, b) + + +NAMES = [ + "Cozy Loft With City View", + "Small Single Room", + "Studio", + "Double Room", +] + + +@pytest.fixture +def names(): + return np.array(NAMES * 64, dtype="= expected.dtype.itemsize + + +@pytest.mark.parametrize("func", ["lower", "upper"]) +def test_case_matches_numpy(names, func): + arr = blosc2.asarray(names) + got = getattr(blosc2, func)(arr).compute(strict_miniexpr=True) + expected = getattr(np_strings, func)(names) + assert list(got[:]) == list(expected) + + +def test_case_expansion_matches_numpy(): + # NumPy uses full case mapping; a 1:1 table would give "STRAßE" here. + src = np.array(["straße", "fix"] * 64, dtype="= 0) + + +def test_bytes_dsl_kernel(raws): + @blosc2.dsl_kernel + def tag(x): + if b"o" in x: + return x + return b"long-" + x + + arr = blosc2.asarray(raws) + got = blosc2.lazyudf(tag, (arr,)).compute(strict_miniexpr=True) + expected = [v if b"o" in v else b"long-" + v for v in raws] + assert list(got[:]) == expected + + +def test_bytes_and_str_do_not_mix(raws): + # NumPy raises on `S` + `U` too; miniexpr must not silently pick one. + assert blosc2_ext.me_output_dtype("o0 + o1", {"o0": "S8", "o1": " 4096, "need a block wider than one miniexpr eval block" + + got = ("x=" + arr).compute(strict_miniexpr=True) + assert list(got[:]) == list(np_add("x=", values)) + + +def test_rebinding_wider_operands_rewidens(): + # lazyexpr(expr, operands) rebinds the operands in place and leaves the + # expression text alone. The inferred width follows the operand dtypes, so a + # width cached under the expression alone would size the output from the first, + # narrower binding and the concat would truncate. + narrow = np.array(["ab"] * 16, dtype="= expected.dtype.itemsize + got = expr.compute(strict_miniexpr=True) + assert list(got[:]) == list(expected) diff --git a/tests/test_b2view_model.py b/tests/test_b2view_model.py index 03611b693..582f26340 100644 --- a/tests/test_b2view_model.py +++ b/tests/test_b2view_model.py @@ -116,7 +116,7 @@ def test_preview_array_2d_returns_grid_preview(): np.testing.assert_array_equal(preview["data"]["4"], np.array([10, 16, 22])) -def test_store_browser_uses_grid_preview_for_2d_ndarray(tmp_path): +def test_browser_grid_preview_for_2d_ndarray(tmp_path): path = tmp_path / "bundle.b2z" with blosc2.TreeStore(str(path), mode="w") as store: store["/arr"] = np.arange(30).reshape(5, 6) @@ -154,7 +154,7 @@ def test_ctable_preview_buffer_reuses_loaded_rows(tmp_path): np.testing.assert_array_equal(page1["data"]["x"], np.arange(5, 10)) -def test_preview_ctable_skips_expensive_nested_columns_by_default(): +def test_preview_skips_expensive_nested_cols(): class Table: def __init__(self): self.col_names = ["path"] @@ -289,7 +289,7 @@ def __getitem__(self, name): assert preview["data"]["path"][1] == [{"x": 2}, {"x": 3}] -def test_ctable_preview_header_uses_column_names_without_dtype_labels(): +def test_preview_header_omits_dtype_labels(): preview = { "start": 0, "stop": 1, diff --git a/tests/test_batch_array.py b/tests/test_batch_array.py index 059d1599e..49981739d 100644 --- a/tests/test_batch_array.py +++ b/tests/test_batch_array.py @@ -143,7 +143,7 @@ def test_batcharray_arrow_ipc_roundtrip(): blosc2.remove_urlpath(urlpath) -def test_batcharray_inferred_layout_preserves_user_vlmeta(): +def test_batcharray_layout_keeps_user_vlmeta(): barray = blosc2.BatchArray() barray.vlmeta["user"] = {"x": 1} @@ -152,7 +152,7 @@ def test_batcharray_inferred_layout_preserves_user_vlmeta(): assert barray.vlmeta["user"] == {"x": 1} -def test_batcharray_arrow_layout_persistence_preserves_user_vlmeta(): +def test_batcharray_arrow_layout_keeps_vlmeta(): pa = pytest.importorskip("pyarrow") barray = blosc2.BatchArray(serializer="arrow") @@ -232,7 +232,7 @@ def fail_decode(*args, **kwargs): assert "items per batch: mean=" in items["nbatches"] -def test_batcharray_info_reports_exact_block_stats_from_lazy_chunks(): +def test_batcharray_info_block_stats_from_lazy(): barray = blosc2.BatchArray(items_per_block=2) barray.extend([[1, 2, 3, 4, 5], [6, 7], [8]]) @@ -240,7 +240,7 @@ def test_batcharray_info_reports_exact_block_stats_from_lazy_chunks(): assert items["nblocks"] == "5 (items per block: mean=1.60, max=2, min=1)" -def test_batcharray_pop_keeps_batch_lengths_metadata_in_sync(): +def test_batcharray_pop_keeps_lengths_in_sync(): barray = blosc2.BatchArray(items_per_block=2) barray.extend([[1, 2, 3], [4, 5], [6]]) @@ -253,7 +253,7 @@ def test_batcharray_pop_keeps_batch_lengths_metadata_in_sync(): assert items["nbatches"].startswith("2 (items per batch: mean=2.00") -def test_batcharray_clear_keeps_empty_store_vlmeta_readable(): +def test_batcharray_clear_keeps_vlmeta_ok(): urlpath = "test_batcharray_clear_empty_vlmeta.b2b" blosc2.remove_urlpath(urlpath) @@ -269,7 +269,7 @@ def test_batcharray_clear_keeps_empty_store_vlmeta_readable(): blosc2.remove_urlpath(urlpath) -def test_batcharray_delete_last_keeps_empty_store_vlmeta_readable(): +def test_batcharray_delete_last_keeps_vlmeta(): urlpath = "test_batcharray_delete_last_empty_vlmeta.b2b" blosc2.remove_urlpath(urlpath) @@ -360,7 +360,7 @@ def test_batcharray_iter_items(): assert list(barray.iter_items()) == [1, 2, 3, 4, 5, 6] -def test_batcharray_respects_explicit_use_dict_and_non_zstd(): +def test_batcharray_use_dict_and_non_zstd(): barray = blosc2.BatchArray(cparams={"codec": blosc2.Codec.LZ4, "clevel": 5}) assert barray.cparams.codec == blosc2.Codec.LZ4 assert barray.cparams.use_dict is False @@ -380,36 +380,26 @@ def test_batcharray_respects_explicit_use_dict_and_non_zstd(): assert barray.cparams.use_dict is False -def test_batcharray_guess_items_per_block_uses_1mib_budget_for_low_clevel(monkeypatch): - # Budgets are fixed; detected cache sizes must not influence the layout. +@pytest.mark.parametrize( + ("clevel", "payloads", "expected"), + [ + # 1 MiB budget: three 300 KiB payloads fit, a fourth would exceed it + (3, [300 * 1024] * 4, 3), + # 8 MiB budget: a single 5 MiB payload fits, a second would exceed it + (5, [5 * 2**20] * 4, 1), + # 16 MiB budget: two 6 MiB payloads fit, a third would exceed it + (7, [6 * 2**20] * 4, 2), + # clevel 9 takes the whole batch, whatever the payload sizes + (9, [100] * 4, 4), + ], + ids=["1mib-low", "8mib-default", "16mib-high", "full-batch"], +) +def test_batcharray_blocksize_budget(monkeypatch, clevel, payloads, expected): + """The per-clevel budget picks the block size, not the detected caches.""" monkeypatch.setitem(blosc2.cpu_info, "l1_data_cache_size", 100) monkeypatch.setitem(blosc2.cpu_info, "l2_cache_size", 1000) - barray = blosc2.BatchArray(cparams={"clevel": 3}) - # 1 MiB budget: three 300 KiB payloads fit, a fourth would exceed it - assert barray._guess_blocksize([300 * 1024] * 4) == 3 - - -def test_batcharray_guess_items_per_block_uses_8mib_budget_for_default_clevel(monkeypatch): - monkeypatch.setitem(blosc2.cpu_info, "l1_data_cache_size", 100) - monkeypatch.setitem(blosc2.cpu_info, "l2_cache_size", 150) - barray = blosc2.BatchArray(cparams={"clevel": 5}) - # 8 MiB budget: a single 5 MiB payload fits, a second would exceed it - assert barray._guess_blocksize([5 * 2**20] * 4) == 1 - - -def test_batcharray_guess_items_per_block_uses_16mib_budget_for_high_clevel(monkeypatch): - monkeypatch.setitem(blosc2.cpu_info, "l1_data_cache_size", 100) - monkeypatch.setitem(blosc2.cpu_info, "l2_cache_size", 150) - barray = blosc2.BatchArray(cparams={"clevel": 7}) - # 16 MiB budget: two 6 MiB payloads fit, a third would exceed it - assert barray._guess_blocksize([6 * 2**20] * 4) == 2 - - -def test_batcharray_guess_items_per_block_uses_full_batch_for_clevel_9(monkeypatch): - monkeypatch.setitem(blosc2.cpu_info, "l1_data_cache_size", 1) - monkeypatch.setitem(blosc2.cpu_info, "l2_cache_size", 1) - barray = blosc2.BatchArray(cparams={"clevel": 9}) - assert barray._guess_blocksize([100, 100, 100, 100]) == 4 + barray = blosc2.BatchArray(cparams={"clevel": clevel}) + assert barray._guess_blocksize(payloads) == expected def test_vlcompress_small_blocks_roundtrip(): @@ -636,7 +626,7 @@ def test_batcharray_copy(): blosc2.remove_urlpath(copy_path) -def test_batcharray_copy_with_storage_preserves_user_metadata(): +def test_batcharray_copy_keeps_user_metadata(): urlpath = "test_batcharray_copy_storage.b2b" copy_path = "test_batcharray_copy_storage_out.b2b" blosc2.remove_urlpath(urlpath) diff --git a/tests/test_dict_store.py b/tests/test_dict_store.py index 01aeb018f..92bd8bb15 100644 --- a/tests/test_dict_store.py +++ b/tests/test_dict_store.py @@ -117,7 +117,7 @@ def test_to_b2z_and_reopen(populated_dict_store): assert np.all(dstore_read["/nodeB"][:] == np.arange(6)) -def test_extensionless_dict_store_defaults_to_directory(tmp_path): +def test_extensionless_store_is_a_directory(tmp_path): path = tmp_path / "test_dstore_extless" with DictStore(str(path), mode="w") as dstore: @@ -421,7 +421,7 @@ def test_external_objectarray_file_and_reopen(tmp_path): @pytest.mark.parametrize("storage_type", ["b2d", "b2z"]) -def test_metadata_discovery_reopens_renamed_external_ndarray(storage_type, tmp_path): +def test_discovery_reopens_renamed_ndarray(storage_type, tmp_path): path = tmp_path / f"test_renamed_ndarray.{storage_type}" ext_path = tmp_path / "renamed_array_source.b2nd" @@ -445,7 +445,7 @@ def test_metadata_discovery_reopens_renamed_external_ndarray(storage_type, tmp_p @pytest.mark.parametrize("storage_type", ["b2d", "b2z"]) -def test_metadata_discovery_reopens_renamed_external_objectarray(storage_type, tmp_path): +def test_discovery_reopens_renamed_objectarray(storage_type, tmp_path): path = tmp_path / f"test_renamed_objectarray.{storage_type}" ext_path = tmp_path / "renamed_objectarray_source.b2frame" values = ["alpha", {"nested": True}, None, (1, 2, 3)] diff --git a/tests/test_group_reduce.py b/tests/test_group_reduce.py index 5e0478aea..1f803bfa5 100644 --- a/tests/test_group_reduce.py +++ b/tests/test_group_reduce.py @@ -18,7 +18,7 @@ def test_group_reduce_size_and_sum_integer_keys(): np.testing.assert_array_equal(sums, np.array([4, 90])) -def test_group_reduce_integer_keys_float_aggs_with_nan_values(): +def test_int_keys_float_aggs_with_nan_values(): keys = np.array([0, 1, 0, 1, 2], dtype=np.uint16) values = np.array([1.0, np.nan, 3.0, np.nan, 10.0]) @@ -40,7 +40,7 @@ def test_group_reduce_integer_keys_float_aggs_with_nan_values(): assert maxs[2] == 10.0 -def test_group_reduce_arbitrary_float_keys_and_nan_key_group(): +def test_float_keys_and_nan_key_group(): keys = np.array([0.5, np.nan, 0.5, -0.0, 0.0, np.nan]) values = np.array([1.0, 2.0, 3.0, 10.0, 20.0, 5.0]) @@ -54,7 +54,7 @@ def test_group_reduce_arbitrary_float_keys_and_nan_key_group(): assert sums[2] == 7.0 -def test_group_reduce_object_keys_sort_none_first_nan_last(): +def test_object_keys_sort_none_first_nan_last(): keys = np.array([np.nan, None, "b", "a", np.nan, None], dtype=object) groups, sizes = blosc2.group_reduce(keys, op="size", sort=True, dropna=False) diff --git a/tests/test_list_array.py b/tests/test_list_array.py index 2aba378b3..8342ff3d6 100644 --- a/tests/test_list_array.py +++ b/tests/test_list_array.py @@ -50,7 +50,7 @@ def test_listarray_append_extend_and_replace(storage, tmp_path): assert restored[:] == reopened[:] -def test_listarray_batch_pending_rows_visible_before_flush(): +def test_listarray_pending_rows_visible(): arr = blosc2.ListArray(item_spec=blosc2.int32(), storage="batch", batch_rows=4) arr.append([1, 2]) arr.append([]) @@ -86,7 +86,7 @@ def test_listarray_arrow_roundtrip(): assert arr.to_arrow().to_pylist() == [["a"], None, ["b", "c"]] -def test_listarray_extend_validate_false_preserves_none(): +def test_listarray_extend_no_validate_keeps_none(): arr = blosc2.ListArray(item_spec=blosc2.int32(), nullable=True, storage="batch", batch_rows=2) arr.extend([[1], None, [2, 3]], validate=False) assert arr[:] == [[1], None, [2, 3]] @@ -139,7 +139,7 @@ def test_listarray_copy_fast_path_empty(): assert dst[:] == [] -def test_listarray_copy_cparams_override_uses_slow_path(): +def test_listarray_copy_cparams_slow_path(): # Supplying cparams must bypass chunk_copy and still produce correct data. src = _make_batch_array() dst = src.copy(cparams={"codec": blosc2.Codec.LZ4, "clevel": 1}) diff --git a/tests/test_locking.py b/tests/test_locking.py index 0ad4de0a4..f8848dd09 100644 --- a/tests/test_locking.py +++ b/tests/test_locking.py @@ -685,7 +685,7 @@ def test_cross_process_multiwriter_ndarray_append(tmp_path): blosc2.remove_urlpath(urlpath) -def test_cross_process_multiwriter_ndarray_append_sparse_nonaligned(tmp_path): +def test_multiwriter_append_sparse_nonaligned(tmp_path): # Same bug class as test_cross_process_multiwriter_ndarray_append, but # on the two physical layouts that test didn't touch: sparse storage # (contiguous=False, each chunk its own file, a different rewrite path diff --git a/tests/test_objectarray.py b/tests/test_objectarray.py index a3ae8c0ce..f86a2c31b 100644 --- a/tests/test_objectarray.py +++ b/tests/test_objectarray.py @@ -261,7 +261,7 @@ def test_objectarray_msgpack_supports_lazyexpr(tmp_path): np.testing.assert_array_equal(restored[:], expected) -def test_objectarray_msgpack_supports_lazyudf_dslkernel(tmp_path): +def test_msgpack_supports_lazyudf_dslkernel(tmp_path): udf, expected = _make_persistent_lazyudf(tmp_path) oarr = blosc2.ObjectArray() @@ -272,7 +272,7 @@ def test_objectarray_msgpack_supports_lazyudf_dslkernel(tmp_path): np.testing.assert_allclose(restored[:], expected) -def test_objectarray_msgpack_rejects_lazyexpr_with_in_memory_operands(): +def test_msgpack_rejects_in_memory_lazyexpr(): expr = _make_in_memory_lazyexpr() oarr = blosc2.ObjectArray() @@ -280,7 +280,7 @@ def test_objectarray_msgpack_rejects_lazyexpr_with_in_memory_operands(): oarr.append(expr) -def test_objectarray_msgpack_rejects_plain_python_lazyudf(tmp_path): +def test_msgpack_rejects_plain_python_lazyudf(tmp_path): udf = _make_persistent_python_lazyudf(tmp_path) oarr = blosc2.ObjectArray() @@ -337,7 +337,7 @@ def test_objectarray_zstd_uses_dict_by_default(): assert oarr.cparams.use_dict is True -def test_objectarray_respects_explicit_use_dict_and_non_zstd(): +def test_objectarray_use_dict_and_non_zstd(): oarr = blosc2.ObjectArray(cparams={"codec": blosc2.Codec.LZ4, "clevel": 5}) assert oarr.cparams.codec == blosc2.Codec.LZ4 assert oarr.cparams.use_dict is False @@ -535,3 +535,18 @@ def test_objectarray_delete_negative_step_slice(): oarr2.extend(range(5)) del oarr2[::-1] assert len(oarr2) == 0 + + +def test_varlen_scalar_cmp_is_elementwise(): + """``column == value`` must not fall through to object identity.""" + from dataclasses import make_dataclass + + row_cls = make_dataclass("Row", [("c", str, blosc2.field(blosc2.vlstring()))]) + t = blosc2.CTable(row_cls) + t.extend({"c": ["hello", "world", "hello"]}, validate=False) + t._flush_varlen_columns() + col = t._cols["c"] + + np.testing.assert_array_equal(col == "hello", [True, False, True]) + np.testing.assert_array_equal(col != "hello", [False, True, False]) + assert isinstance(hash(col), int) diff --git a/tests/test_pandas_udf_engine.py b/tests/test_pandas_udf_engine.py index f9df6bc16..802668d6e 100644 --- a/tests/test_pandas_udf_engine.py +++ b/tests/test_pandas_udf_engine.py @@ -181,7 +181,7 @@ def test_apply_object_dtype_raises_clear_error(self): with pytest.raises(ValueError, match="numeric dtype"): df.apply(lambda x: x + 1, engine=blosc2.jit) - def test_apply_axis1_row_subscript_idiom_matches_default_engine(self): + def test_axis1_subscript_matches_default_engine(self): def add_people(row): return row["max_people"] + row["max_children"] @@ -190,7 +190,7 @@ def add_people(row): result = df.apply(add_people, engine=blosc2.jit, axis=1) pd.testing.assert_series_equal(result, expected) - def test_apply_axis1_row_subscript_args_kwargs_forwarded(self): + def test_axis1_subscript_args_kwargs_forwarded(self): def combine(row, num1, num2=0): return row["a"] + row["b"] + num1 + num2 @@ -199,7 +199,7 @@ def combine(row, num1, num2=0): result = df.apply(combine, engine=blosc2.jit, axis=1, args=(10,), num2=100) pd.testing.assert_series_equal(result, expected) - def test_apply_axis1_row_subscript_preserves_column_dtype(self): + def test_axis1_subscript_keeps_column_dtype(self): # a mixed-dtype frame would be upcast by DataFrame.values; the row # proxy must extract columns from the original frame instead. def add(row): @@ -209,7 +209,7 @@ def add(row): result = df.apply(add, engine=blosc2.jit, axis=1) np.testing.assert_allclose(result.to_numpy(), [1.5, 2.5, 3.5]) - def test_apply_axis1_row_subscript_with_loop_raises_clear_error(self): + def test_axis1_subscript_with_loop_raises(self): def kepler_row(row): m, ecc = row["m"], row["ecc"] e = m + ecc * np.sin(m) @@ -222,7 +222,7 @@ def kepler_row(row): with pytest.raises(TypeError, match="for/while loop"): df.apply(kepler_row, engine=blosc2.jit, axis=1) - def test_apply_axis1_row_subscript_duplicate_column_raises(self): + def test_axis1_subscript_duplicate_col_raises(self): def add(row): return row["a"] + 1 @@ -230,7 +230,7 @@ def add(row): with pytest.raises(KeyError, match="duplicated"): df.apply(add, engine=blosc2.jit, axis=1) - def test_apply_axis1_row_subscript_attribute_access_raises(self): + def test_axis1_subscript_attr_access_raises(self): def bad(row): return row["a"] + row.b @@ -238,19 +238,18 @@ def bad(row): with pytest.raises(AttributeError, match="row\\['b'\\]"): df.apply(bad, engine=blosc2.jit, axis=1) - def test_apply_axis1_row_subscript_non_numeric_column_raises(self): - # Whole-frame numeric-dtype validation (`_ensure_numpy_data`) already - # gates this ahead of row-proxy dispatch; `_PandasRowProxy` carries - # its own per-column check too, for callers that construct it - # directly. + def test_axis1_subscript_unvectorizable_raises(self): + # String columns are supported now, so the per-column check in + # `_PandasRowProxy` is what still rejects a dtype the engine cannot + # vectorize at all. def bad(row): - return row["a"] + len(row["b"]) + return row["a"] + row["b"] - df = pd.DataFrame({"a": [1.0, 2.0], "b": ["x", "y"]}) - with pytest.raises(ValueError, match="numeric dtype"): + df = pd.DataFrame({"a": [1.0, 2.0], "b": pd.to_datetime(["2020-01-01", "2020-01-02"])}) + with pytest.raises(ValueError, match="cannot vectorize"): df.apply(bad, engine=blosc2.jit, axis=1) - def test_apply_axis1_positional_idiom_still_uses_per_row_loop(self): + def test_axis1_positional_uses_per_row_loop(self): # No `row["..."]` subscript: falls back to the historical per-row # loop, unaffected by the row-proxy dispatch added for the subscript # idiom above. @@ -259,7 +258,7 @@ def test_apply_axis1_positional_idiom_still_uses_per_row_loop(self): result = df.apply(lambda row: row * 2, engine=blosc2.jit, axis=1) pd.testing.assert_frame_equal(result, expected) - def test_apply_already_jitted_function_is_not_decorated_twice(self): + def test_apply_jitted_func_not_decorated_twice(self): # Decorating and passing engine= both request the same thing. Applying # the decorator a second time used to wrap the array in a SimpleProxy # before the inner DSL kernel saw it, which then failed asking for @@ -278,7 +277,7 @@ def branch(col): result = df.apply(func, engine=blosc2.jit) pd.testing.assert_frame_equal(result, expected) - def test_map_already_jitted_function_is_not_decorated_twice(self): + def test_map_jitted_func_not_decorated_twice(self): def branch(col): if col >= 0: out = col * 2.0 @@ -369,7 +368,9 @@ def dsl(a, b): for name, func in (("traced", traced), ("dsl", dsl)): with pytest.raises(TypeError) as excinfo: func(**df) - message = str(excinfo.value) + # The guidance rides along as a note (it prints with the traceback); + # rebuilding the exception to append it would assume its constructor. + message = "\n".join([str(excinfo.value), *getattr(excinfo.value, "__notes__", [])]) assert name in message assert "'note'" in message assert "**df[['a', 'b']]" in message @@ -377,3 +378,108 @@ def dsl(a, b): # A missing operand is a different mistake and keeps its own message with pytest.raises(TypeError, match="missing a required argument"): dsl(a=df["a"]) + + +@pytest.mark.skipif(pd is None, reason="pandas not installed") +@pytest.mark.skipif(_pandas_too_old, reason="engine= integration targets pandas 3.x") +class TestRowKernelsWithControlFlow: + """`row["colname"]` combined with an `if`. + + Neither dispatch route could run these before: tracing evaluates the `if` + over a whole column ("truth value ... is ambiguous") and the DSL parser + rejected the subscript. They are now rewritten into named parameters. + """ + + def test_numeric_row_kernel_with_branch(self): + def pick(row): + if row["a"] > 2: + return row["a"] + row["b"] + return row["a"] - row["b"] + + df = pd.DataFrame({"a": [1.0, 2.0, 3.0, 4.0], "b": [10.0, 20.0, 30.0, 40.0]}) + expected = df.apply(pick, axis=1) + result = df.apply(pick, axis=1, engine=blosc2.jit) + pd.testing.assert_series_equal(result, expected) + + def test_blog_kernel_matches_default_engine(self): + """End-to-end acceptance: the pandas-3 blog kernel, run unmodified.""" + + def format_room_info(row): + result = "property_type=" + row["property_type"] + desc = row["name"].lower() + if " with " not in desc: + return result + ", room_type=" + desc.removesuffix(" room") + before, after = desc.split(" with ", 1) + r2 = result + ", room_type=" + before.removesuffix(" room") + return r2 + ", amenity=" + after + + df = pd.DataFrame( + { + "property_type": ["Entire home", "Private room", "Shared room", "Loft"] * 8, + "name": [ + "Cozy Loft With City View", + "Small Single Room", + "Studio with balcony", + "Double Room", + ] + * 8, + } + ) + expected = df.apply(format_room_info, axis=1) + result = df.apply(format_room_info, axis=1, engine=blosc2.jit) + pd.testing.assert_series_equal(result, expected) + + def test_column_named_like_a_dsl_function(self): + """A column name that shadows a DSL builtin must not change meaning. + + The rewrite turns row["sqrt"] into a parameter literally called + `sqrt`, which then coexists with a real sqrt() call in the same + expression; operands and calls are distinguished by position, so both + resolve correctly. Index symbols (`_i0`) are checked for the same + reason. + """ + + def collide(row): + return row["sqrt"] + np.sqrt(row["b"]) + + def index_symbol(row): + return row["_i0"] + row["b"] + + df = pd.DataFrame({"sqrt": [1.0, 2.0], "b": [4.0, 9.0], "_i0": [5.0, 6.0]}) + for fn in (collide, index_symbol): + pd.testing.assert_series_equal(df.apply(fn, axis=1, engine=blosc2.jit), df.apply(fn, axis=1)) + + def test_non_identifier_column_label(self): + def tag(row): + if row["room type"] == "loft": + return "L" + return "-" + + df = pd.DataFrame({"room type": ["loft", "studio", "loft"]}) + expected = df.apply(tag, axis=1) + result = df.apply(tag, axis=1, engine=blosc2.jit) + pd.testing.assert_series_equal(result, expected) + + def test_null_string_column_is_rejected(self): + # pandas raises on a row kernel over a null too; substituting "" would + # invent a value it never produces. + def concat(row): + return "p=" + row["x"] + + df = pd.DataFrame({"x": ["a", None, "c"]}) + with pytest.raises(TypeError): + df.apply(concat, axis=1) + with pytest.raises(ValueError, match="contains nulls"): + df.apply(concat, axis=1, engine=blosc2.jit) + + def test_positional_row_access_still_rejected(self): + # Only `row["literal"]` is rewritten; anything else must keep failing + # loudly rather than silently taking a different route. + def positional(row): + if row[0] > 1: + return row[0] + return row[1] + + df = pd.DataFrame({"a": [1.0, 2.0], "b": [3.0, 4.0]}) + with pytest.raises((TypeError, ValueError, RuntimeError)): + df.apply(positional, axis=1, engine=blosc2.jit) diff --git a/tests/test_proxy_schunk.py b/tests/test_proxy_schunk.py index dcd793ec5..3164e1d6e 100644 --- a/tests/test_proxy_schunk.py +++ b/tests/test_proxy_schunk.py @@ -77,7 +77,7 @@ def test_open(urlpath, chunksize, nchunks): blosc2.remove_urlpath(proxy_urlpath) -def test_open_readonly_proxy_keeps_schunk_cache_and_source_readonly(tmp_path): +def test_readonly_proxy_keeps_both_readonly(tmp_path): source_path = tmp_path / "source.b2frame" proxy_path = tmp_path / "proxy.b2frame" data = np.arange(200, dtype="int32") diff --git a/tests/test_python_blosc.py b/tests/test_python_blosc.py index 8d0b3d149..abcc631b2 100644 --- a/tests/test_python_blosc.py +++ b/tests/test_python_blosc.py @@ -184,7 +184,7 @@ def test_unpack_array_with_from_py27_exceptions(self): with pytest.raises(UnicodeDecodeError): blosc2.unpack_array(self.PY_27_INPUT) - def test_unpack_array_with_unicode_characters_from_py27(self): + def test_unpack_array_unicode_from_py27(self): import numpy as np out_array = np.array(["å", "ç", "ø", "π", "˚"]) diff --git a/tests/test_random.py b/tests/test_random.py index b5fdf9c0d..fc8902230 100644 --- a/tests/test_random.py +++ b/tests/test_random.py @@ -216,7 +216,7 @@ def test_choice_2d_a_not_implemented(): [(m, a, k) for m, (a, k) in _VECTOR_DIST_CASES.items()], ids=_VECTOR_DIST_CASES.keys(), ) -def test_vector_distribution_shape_reproducible_and_finite(method, args, k): +def test_vector_dist_reproducible_and_finite(method, args, k): def draw(): rng = getattr(blosc2.random.default_rng(0), method) return rng(*args, shape=(40,)) @@ -249,7 +249,7 @@ def test_vector_dist_numpy_integer_shape(): assert a.shape == (5, 3) -def test_permutation_int_is_a_permutation_and_reproducible(): +def test_permutation_int_is_valid_and_stable(): a = blosc2.random.default_rng(0).permutation(10) b = blosc2.random.default_rng(0).permutation(10) np.testing.assert_array_equal(a[:], b[:]) diff --git a/tests/test_schunk_get_slice.py b/tests/test_schunk_get_slice.py index 032105ebd..ae6dc23fc 100644 --- a/tests/test_schunk_get_slice.py +++ b/tests/test_schunk_get_slice.py @@ -116,3 +116,49 @@ def test_schunk_get_slice_raises(): assert schunk[start:stop] == b"" blosc2.remove_urlpath(kwargs["urlpath"]) + + +# --------------------------------------------------------------------------- +# Typesizes above BLOSC_MAX_TYPESIZE (255) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("typesize", [252, 256, 512]) +def test_get_slice_wide_typesize_matches_source(typesize): + """c-blosc2 records a typesize above 255 as 1 in the chunk header, but + blosc2_schunk_get_slice_buffer() still divides byte offsets by + schunk->typesize, so a partially covered chunk addressed the wrong range: + most slices failed outright and single-element ones returned the wrong + bytes with no error. An