Architecture of coola.equality¶
This page describes the design of the
coola.equality package for contributors: the
invariants, the contracts between the parts, and the trade-offs of the main design choices. The rules
that users rely on are in the user guide; the API is in the
reference.
Overview¶
The package is small on purpose: one engine, one extension point, one numeric rule and one table algorithm.
Constraints¶
- Performance is a requirement. Both the passing and the failing path are measured by the benchmarks; a simpler design that is slower is rejected.
- Early exit. Two objects are different as soon as one difference is found: nothing after it is visited, converted or formatted.
- First difference only, on demand.
objects_are_equalandobjects_are_allclosebuild no message and no path; onlycompareandassert_*explain, and only the first difference. - Equality drives hashing, not the reverse.
hash_objectis aligned with equality, never the other way round.
Architecture¶
objects_are_equal / objects_are_allclose / compare / assert_objects_equal / assert_objects_allclose
| one EqualityConfig (frozen) and one Engine per call
v
Engine.run(actual, expected) explicit stack of child iterators; one loop
| identity -> type -> tester lookup (cached per call) -> verdict or children
v
Tester.objects_are_equal(actual, expected, ctx) -> bool | Iterator[(actual, expected, path part)]
| leaves: numbers, arrays, scalars containers: sequence, mapping, set, object, table
v
numeric kernel (NumPy, torch, JAX) adapters: pandas, polars, pyarrow -> table
Four rules follow from it:
- One extension point: the tester. A tester is a class with one method.
- The engine owns the invariants: identity, strict type, depth, cycles, path and first failure. A tester cannot forget them because it never recurses.
- A container decomposes, it does not recurse: it returns the child pairs that must all be equal, lazily and in the order the first difference is wanted.
- coola owns the semantics, the libraries own the structure: every number goes through the
kernel; a library says whether two frames have the same columns, never whether
1is close to2.
The engine¶
The engine is one loop over a stack of child iterators (depth first, so the first difference in iteration order is the one reported). For each pair:
actual is expected: equal.type(actual) is not type(expected): different.- The tester is looked up in a per-call
dict[type, tester]filled from the registry the first time a type is met, so the registry lookup is paid once per type. A registration made during a comparison is seen by the next call. - The tester is called. A
boolis the verdict. An iterator is a container: if its pair is already on the stack the pair is equal (cycle); ifmax_depthis set and the stack is that deep,MaxDepthExceededErroris raised; otherwise it is pushed and the loop continues with its first child. - On
False, the path is built once from the path parts of the stack, the part of the failing pair and the parts given toctx.fail.
Consequences: the stack use is constant, RecursionError cannot happen, and explaining costs nothing
in the loop. ctx.fail(reason, *path) returns False; reason can be a callable that is called only
when the difference is explained. The full loop is in the appendix.
Reporting¶
There is one record, the first failure (path, reason): no list of reports and no message per parent
level. Values in a reason go through one bounded formatter, so a failure in a million-element array
formats a few values, not the array. A path part is a key or a position (['a'], [3]),
AttributeName("x") (.x) or Element(2) ({2}). An array reports its index. A table is a leaf
(the path stops at the table); its reason names the part of the structure that differs or the first
differing column, row and values.
The numeric kernel¶
NumPy, torch and JAX share one formula:
close = (a == b) | (finite & (|a - b| <= atol + rtol * max(|a|, |b|))) | (isnan(a) & isnan(b) if equal_nan)
| Dtype | Preparation | Rule |
|---|---|---|
bool |
none | exact |
| signed and unsigned integers | exact difference in integers | kernel, tolerance in float64 |
float16, bfloat16, float32 |
cast to float64 |
kernel |
float64, complex128 |
none (complex64 is cast up) |
kernel |
timedelta64 |
integer nanoseconds; NaT mask apart |
kernel; NaT by equal_nan |
datetime64, str, bytes, void |
none | exact; NaT by equal_nan |
object |
none | cells are children of the engine |
Python scalars use the same formula; Decimal, Fraction and timedelta keep a difference-based
rule.
Tables¶
pandas, polars and PyArrow share one algorithm in TableEqualityTester, which takes an adapter per
library: check the structure (shape, names, dtypes, index), then compare the columns lazily, one by
one, each by its kind, and stop at the first column that differs. A Series or an Array is a
one-column table (its failure says "series" or "arrays" through the kind_name of the adapter).
| Kind | Dtypes | Comparison |
|---|---|---|
numeric |
integers, floats, nullable and Arrow numerics | null masks equal, then the kernel |
duration |
timedelta64, polars Duration, Arrow duration |
kernel on integer nanoseconds |
exact |
booleans, strings, categories, instants, nested and object |
the library's equals; never sees the tolerance or equal_nan |
Nested and object columns are the one place where the library decides: None against NaN in an
object column is equal, and a NaN inside a nested cell equals a NaN without equal_nan. attrs,
flags and Arrow schema metadata are ignored. Adding a library means writing one adapter of about 40
lines.
Module layout¶
coola/equality/
interface.py objects_are_equal, objects_are_allclose
result.py ComparisonResult, compare, assert_objects_equal, assert_objects_allclose
config.py EqualityConfig
tolerance.py Tolerance
engine.py Engine, Element
tester/ base, registry, one module per family (default, scalar, collection, difference,
dataclass, object, equal, tolerant), per library (numpy, torch, jax, xarray,
pydantic, scipy), and table.py with the adapters pandas.py, polars.py, pyarrow.py
_helpers/ private helpers (formatting, numeric kernel, comparison keys)
Design choices and trade-offs¶
Each row is a choice the current design makes, with the alternative it was preferred to. The sections after the table give the longer arguments.
| Choice | Reason | Alternative and its cost |
|---|---|---|
| Iterative engine | constant stack, one place for depth, cycles and path | a recursive engine is faster, but needs a try/finally per level |
| One tester method returns a verdict or an iterator | one extension point | two methods (equal, expand) are explicit, but one can be forgotten |
No show_difference and no logger; compare and assert_* explain |
no cost on the fast path, thread safe | a flag on the fast function makes it format messages |
bool and every Enum are exact |
they are labels, not quantities | a "label" list for str and bytes is more to maintain |
| Masked arrays compare the data under the mask, the mask and fill value | no hidden difference | ignoring the masked data matches how the array prints |
Small floats are cast to float64 |
no overflow of the tolerance at the edge of the range | keeping the dtype is faster for large float32 arrays but wrong at edges |
| Integer difference is exact | no rounding of large integers | exact Fraction arithmetic is exact everywhere but much slower |
| Tables are one leaf, reported by column and row | simple and fast to compare; the row is found only when explaining | a second pass per library, and no row for a nested column |
| Time resolution is part of the dtype | as in NumPy | casting to the finer resolution is friendlier but hides a dtype difference |
| Library testers are public | users can compose, subclass or register them | private testers shrink the public surface |
max_depth (default None) and MaxDepthExceededError |
opt-in guard on pathological structures; a deep valid structure is compared | by default, no limit runs until memory is exhausted on an infinite structure |
compare(rtol, atol) instead of a Tolerance argument |
the five functions share the same keywords | a Tolerance argument is reusable, but unlike the array libraries |
No root repr in ComparisonResult |
the failure path does not depend on the size of roots | repr of the roots costs 20x on a failing million-element comparison |
Iterative or recursive engine¶
The iterative engine is about 1.5x slower than a recursive one on the happy path, and a tester
returns an iterator, which is less familiar than calling ctx.compare. In exchange there is no
RecursionError, explaining is free, and container testers are pure functions of two objects. A
recursive engine keeps the interpreter limit and max_depth, and needs a try/finally to keep the
depth, the cycle set and the path right.
Tester return protocol¶
One method that returns a bool or an iterator of children has one name and no flag; the engine
tests the type of the result. Two methods (equal for a leaf and expand for a container, plus a
container flag) are explicit and type-checkable, but are two entry points to document and to forget.
An always-iterator protocol has one shape, but every leaf pays for a generator.
Explanation lives in compare and assert_*¶
The fast functions have no show_difference flag and no logger. A flag would force every tester to
guard its formatting behind if config.show_difference: (a forgotten guard formats on every call),
require a reporter callback, and format the message before a logger can drop it. With the
explanation in compare and assert_*:
objects_are_equalis guaranteed to format nothing, and there is no global state, so it is thread safe;- the engine has one failure record, and the five functions have the same signature;
- the explanation is a structured
ComparisonResult(path,reason), not log text.
The cost is that a caller who wants "the answer and a log line in one call" writes two lines
(log.info(compare(a, b).format_message())). Passing a callable (explain=print) would route the
explanation, but it adds a parameter and a callback invoked inside a comparison (re-entrancy,
exceptions) for a use that compare already covers.
Result and tolerance shape¶
ComparisonResulthas no rootrepr. Thereprof the whole roots does not say where the difference is (the path and the two differing values do), and building it on every failure made the failing path depend on the size of the roots (231 ms forcompareof 1M floats against 10 ms without). The caller already holds the roots. A lazy property would keep references to the roots and give a differentreprif they are mutated.comparetakesrtolandatol, likeobjects_are_allcloseand the array libraries, so a call moves from one function to the other by renaming. The defaults differ on purpose:compareis exact unless a tolerance is given,objects_are_allcloseuses1e-5and1e-8.Toleranceis public.EqualityConfigholds one validated, hashable value instead of two loose floats, and a custom tester or a user-built config can readctx.config.tolerance. A caller of the five functions never needs it.
Data under a masked array¶
A masked array compares the data under the mask, the mask and the fill_value, so no difference is
hidden. Ignoring the data under the mask matches how two arrays print, and is what the hasher does;
the price of the current rule is that two arrays that print the same can be different. fill_value is
compared because it is visible through filled(). The hasher ignores the masked data, which keeps
"equal implies same hash" true; including it would make the hash as discriminating as equality at the
cost of changing every hash of a masked array.
Nested and object cells of frames¶
The library's equals decides, because it is fast and a cell-by-cell comparison through the engine
would be slow. The price is three rules (the library's, the engine's and the hasher's): the tolerance
and equal_nan do not reach nested cells, None equals NaN in an object column, and the hasher
has to treat them the same way. This exception is part of the Guarantees table of the user guide.
Comparing each cell through the engine gives one rule per value (a nested float gets the tolerance,
None is not NaN) at the cost of speed.
Time resolution¶
timedelta64[s] and timedelta64[ns] are different dtypes, as in NumPy. pandas 3 infers different
resolutions for values built in different ways, so two frames of the same instants can differ; the fix
is an astype, which the guide documents. Casting both sides to the finer resolution would be
friendlier, but it breaks the "same dtype" rule.
Precision¶
Small floats (float16, bfloat16, float32) are cast to float64 before the kernel. Keeping the
dtype is faster for large float32 arrays, but overflows the tolerance at the edge of the range (a
RuntimeWarning and a wrong True for 60000 against -60000 with rtol=1.5); the cast doubles
the temporaries. Integers use an exact integer difference and the tolerance in float64; exact
Fraction arithmetic for every integer array would be exact everywhere and much slower.
What is ignored¶
requires_grad, pandas attrs and flags, Arrow schema metadata and pydantic private attributes do
not take part in a comparison. They are documented in the user guide instead of being opt-in flags.
bool and every enum.Enum member are labels: the tolerance never applies to them.
Public surface¶
The library testers (NumpyArrayEqualityTester, TorchTensorEqualityTester, ...) are public so that
a user can compose, subclass or register them. Pandas, polars and PyArrow share the public
TableEqualityTester(adapter) rather than one subclass per library, which keeps the number of names
small. Engine is the public type of ctx, and Element and AttributeName are exported because
they appear in ComparisonResult.path. The helpers are private, except the kernel close.
Tables are reported by column and row¶
A table is a leaf, so the path stops at the table. The loop over the columns knows the name of the
first differing column at no cost, so the reason always names it. The row and the two values are found
by a second pass, only when the comparison explains a difference (ctx.explaining): the numeric and
duration columns use the NumPy kernel, the other columns use a function of the adapter
(Column.difference). A nested column has none, so its reason has no row. An adapter can say what
differs in the structure with structure_difference, which is called only to explain.
Lookup cache¶
The tester cache is a per-call dict: there is no lock on the hot path, and a registration during a
call is seen by the next call. A lock-free immutable snapshot inside the registry would also benefit
hashing and summary, but changes a component that every package shares.
Public API¶
objects_are_equal(actual, expected, *, equal_nan=False, max_depth=None, registry=None) -> bool
objects_are_allclose(actual, expected, *, rtol=1e-5, atol=1e-8, equal_nan=False, max_depth=None, registry=None) -> bool
compare(actual, expected, *, rtol=0.0, atol=0.0, equal_nan=False, max_depth=None, registry=None) -> ComparisonResult
assert_objects_equal(actual, expected, *, equal_nan=False, max_depth=None, registry=None, root="actual") -> None
assert_objects_allclose(actual, expected, *, rtol=1e-5, atol=1e-8, equal_nan=False, max_depth=None, registry=None, root="actual") -> None
The five functions have the same keywords with the same meaning. compare is exact unless rtol or
atol is given.
@dataclass(frozen=True, slots=True)
class Tolerance: # validated: finite, non-negative, not a bool
rtol: float = 0.0
atol: float = 0.0
@dataclass(frozen=True, slots=True)
class EqualityConfig: # shareable, hashable, thread-safe
registry: EqualityTesterRegistry
tolerance: Tolerance = Tolerance()
equal_nan: bool = False
max_depth: int | None = None # validated: None (no limit) or a positive int, not a bool
@dataclass(frozen=True, slots=True)
class ComparisonResult:
equal: bool
path: tuple[object, ...] = () # keys, positions, AttributeName, Element
reason: str | None = None
def __bool__(self) -> bool: ...
def format_message(self, root: str = "actual") -> str: ...
Exported by coola.equality: the five functions, ComparisonResult, EqualityConfig, Tolerance,
MaxDepthExceededError, AttributeName, Element, BaseEqualityTester, Engine (the type of
ctx), EqualityTesterRegistry, get_default_registry, register_equality_testers, and the named
testers (ObjectAttributesEqualityTester, EqualEqualityTester, OrderedMappingEqualityTester,
SequenceEqualityTester, MappingEqualityTester, ... and the library testers
NumpyArrayEqualityTester, TorchTensorEqualityTester, TableEqualityTester, ...).
The engine and the testers¶
The tester contract¶
class BaseEqualityTester(ABC, Generic[T]):
@abstractmethod
def objects_are_equal(
self, actual: T, expected: T, ctx: Engine
) -> bool | Iterator[tuple[object, object, object]]:
"""A verdict (a leaf), or the child pairs (actual, expected, path part) that must all be
equal (a container), produced lazily and in the order the first difference is wanted."""
# default __eq__: same type and same vars
ctx gives ctx.atol, ctx.rtol, ctx.equal_nan (read once from the config), ctx.fail(reason,
*path), ctx.explaining and ctx.warn_once(key, message). ctx.fail returns False, so a tester
writes return ctx.fail(...). reason may be a callable; it is called only when ctx.explaining
(compare, assert_*), so objects_are_equal formats nothing. path adds parts after the ones the
engine knows (the index in an array).
A tester has no state but its configuration; __eq__ comes from the base class. The engine runs the
loop described in The engine: the registry lookup costs about 300 ns, so it is paid
once per type and call. Stack use is constant (100 000 nested dicts compare under the default
interpreter limit of 1 000), so max_depth is a policy, not a safety net. MaxDepthExceededError is a
plain Exception, not a RecursionError.
Testers¶
class SequenceEqualityTester(BaseEqualityTester[Sequence]):
def objects_are_equal(self, actual, expected, ctx):
if len(actual) != len(expected):
return ctx.fail(lambda: f"lengths differ: {len(actual)} vs {len(expected)}")
return ((a, e, i) for i, (a, e) in enumerate(zip(actual, expected)))
class MappingEqualityTester(BaseEqualityTester[Mapping]):
def objects_are_equal(self, actual, expected, ctx):
keys_a, keys_e = (
key_index(actual, ctx),
key_index(expected, ctx),
) # comparison_key: types, NaN keys
if keys_a.keys() != keys_e.keys():
return ctx.fail(lambda: describe_key_difference(keys_a, keys_e))
return ((actual[keys_a[k]], expected[keys_e[k]], keys_a[k]) for k in keys_a)
class NumpyArrayEqualityTester(BaseEqualityTester[np.ndarray]):
def objects_are_equal(self, actual, expected, ctx):
if actual.dtype != expected.dtype:
return ctx.fail(lambda: f"dtypes differ: {actual.dtype} vs {expected.dtype}")
if actual.shape != expected.shape:
return ctx.fail(lambda: f"shapes differ: {actual.shape} vs {expected.shape}")
if actual.dtype.kind == "O": # the cells are objects: children of the engine
return (
(a, e, i)
for i, (a, e) in enumerate(zip(actual.ravel().tolist(), expected.ravel().tolist()))
)
close = close_numpy(actual, expected, ctx) # boolean array, the kernel of
if close.all():
return True
index = np.unravel_index(np.argmin(close), close.shape) # looked up only when it fails
return ctx.fail(lambda: describe_values(actual[index], expected[index]), *index)
class ScalarEqualityTester(BaseEqualityTester[float]): # int, float, complex, NumPy scalars
def objects_are_equal(self, actual, expected, ctx):
if isinstance(actual, (bool, enum.Enum)): # labels, not quantities
return actual == expected or ctx.fail(lambda: describe_values(actual, expected))
return number_close(actual, expected, ctx) or ctx.fail(
lambda: describe_values(actual, expected)
)
A user tester for their own class returns the pairs of its attributes:
class PointEqualityTester(BaseEqualityTester[Point]):
def objects_are_equal(self, actual, expected, ctx):
return ((actual.coords, expected.coords, AttributeName("coords")),)
Type rules¶
| Type | Tester | Rule |
|---|---|---|
object (unregistered) |
DefaultEqualityTester |
==; a raise or a non-bool result is a TypeError naming the type |
int, float, complex, NumPy scalars |
ScalarEqualityTester |
kernel formula; bool and enums exact |
Decimal, Fraction, timedelta |
DifferenceCloseEqualityTester |
difference-based tolerance; Decimal NaN by equal_nan |
list, tuple, deque, Sequence, array.array, dict values |
SequenceEqualityTester |
length, then items in order |
dict, Mapping, MappingProxyType |
MappingEqualityTester |
same keys (by type, NaN by equal_nan), then values |
OrderedDict |
OrderedMappingEqualityTester |
and the same key order |
set, frozenset |
SetEqualityTester |
elements matched by hash and type; reports Element |
| dataclasses, pydantic models | DataclassEqualityTester, PydanticModelEqualityTester |
fields as AttributeName children |
np.ndarray, np.ma.MaskedArray, NumPy datetime64/timedelta64 scalars |
Numpy*EqualityTester |
dtype, shape, kernel; for a masked array the data (masked values included), the mask and fill_value; NaT by equal_nan |
torch.Tensor, PackedSequence, jax.Array, xarray |
Torch*, Jax*, Xarray*EqualityTester |
device, dtype, shape, layout, kernel; xarray parts as children |
scipy.sparse matrices and arrays |
ScipySparseEqualityTester |
class, shape, dtype, non-zero positions, then the stored values with the kernel; duplicates summed, stored zeros ignored |
| pandas, polars, pyarrow frames, series, arrays | TableEqualityTester |
structure, then columns by kind |
The registry¶
EqualityTesterRegistry is a TypeRegistry (coola.registry):
find_equality_tester(data_type)returns the tester for the most specific registered type, walking the Method Resolution Order (MRO);objectis registered withDefaultEqualityTester, so a tester is always found;- more specific testers override general ones, and a subclass inherits the tester of its parent (a
tester for
Sequenceapplies tolistandtuple); - testers for optional libraries are registered as lazy groups, and loaded when a type of the library is looked up.
get_default_registry() returns the global registry and register_equality_testers(mapping, *,
exist_ok=False) adds testers to it. A test that registers a tester should not leak it: use
coola.registry.default.temporary_default_registries() (see the testing guide), or work on
get_default_registry().copy().
Error handling¶
When a specific tester is not available, the lookup falls back to a more general tester (MRO), then to
DefaultEqualityTester (registered for object), which uses ==. When the comparison explains the
difference (compare, assert_*), the reason says what differs, where in the structure, and the
values that differ.
Writing a tester¶
To add support for a custom type:
- Implement a tester:
from coola.equality import BaseEqualityTester
class MyTypeEqualityTester(BaseEqualityTester[MyType]):
def objects_are_equal(self, actual, expected, ctx):
# the engine checked the identity and the type: both objects are a MyType
return actual.compare_to(expected) or ctx.fail("my objects differ")
- Register it:
from coola.equality import register_equality_testers
register_equality_testers({MyType: MyTypeEqualityTester()})
To avoid modifying the global registry, create an EqualityTesterRegistry, register the tester on
it and pass it with registry=.
- Use it:
objects_are_equal(obj1, obj2).
See the extending guide for more details.
Checklist:
- Return a verdict (
ctx.fail(reason)for a difference) or the child pairs(actual, expected, path part); never recurse: the engine checks the identity, the type, the depth and the cycles. - Use
AttributeNamefor an attribute andctx.fail(reason, *path)for the location inside a leaf (the index in an array). - Build the reason lazily (a callable), so
objects_are_equalformats nothing. - Honour
ctx.equal_nanand the tolerance (ctx.rtol,ctx.atol:|a - e| <= atol + rtol * max(|a|, |e|)), or document that your tester ignores them. - Check the cheap properties first (metadata such as shape, dtype, device), then the values.
Strict types¶
coola enforces strict type checking:
1(int) â‰1.0(float) â‰True(bool)listâ‰tupledictâ‰OrderedDict- two
OrderedDicts with the same items in a different order are different (the order of the keys of adictis ignored)
This prevents subtle bugs from type coercion.
Not covered¶
- Per-library corner cases beyond the tested ones: sparse and quantized torch tensors, nested tensors,
dask-backed xarray, pandas
Interval,PeriodandSparsecolumns on theexactpath. - The GPU path, Windows, fork behaviour.
- A registration made while a comparison runs is seen by the next call only (the lookup cache is per call).
- Mutation testing of the package.
Appendix: the engine loop¶
def run(self, a, e): # Engine.run
stack, keys, parts, visiting = [], [], [], set() # child iterators, cycle keys, path parts
pair = (a, e, None)
while True:
a, e, part = pair
ok = True
if a is not e:
t = type(a)
if t is not type(e):
ok = self.fail("types differ")
else:
result = self.find(t).objects_are_equal(a, e, self) # a verdict, or children
if result is True or result is False:
ok = result
elif (key := (id(a), id(e))) not in visiting: # a pair on the stack is equal
if self.max_depth is not None and len(stack) >= self.max_depth:
raise MaxDepthExceededError
visiting.add(key)
stack.append(result)
keys.append(key)
parts.append(part)
if ok is False:
return self.record_failure(
parts, part
) # path = parts[1:] + (part,) + the parts of fail()
while True: # next pair: the top iterator, popping the exhausted ones
if not stack:
return True
if (nxt := next(stack[-1], None)) is not None:
pair = nxt
break
stack.pop()
visiting.discard(keys.pop())
parts.pop()
Relation to hashing¶
The property "equal objects have the same hash_object" holds with this equality. The hashers are
allowed to be coarser than equality (different values may share a hash), never finer. The cases that
need care:
| Case | Equality | Hasher |
|---|---|---|
pandas object column, None against NaN |
equal (the library's equals) |
must hash a None and a NaN of an object column the same |
NaN inside nested cells (lists, structs) of a frame |
equal without equal_nan (the library) |
the hash of nested NaN must not depend on the NaN bit pattern |
| Nested floats of a frame, tolerance | exact (no tolerance) | nothing to do: equal values have equal hashes |
| Masked array, data under the mask | compared | ignored by the hasher; allowed, because different values may collide |
Integers above 2^53 in arrays, bool and enums |
exact | the hash is exact (no float64 conversion) |
IntEnum member against int |
different types | the hashes may be equal: allowed |
MappingProxyType, array.array |
compared as a mapping and a sequence | hashed consistently with a dict and a list of the same content |
Time resolution (timedelta64[s] against [ns]) |
different dtypes | allowed to share a hash |
The hashing property tests (tests/properties/test_hashing*.py) take the equal pairs that the
equality tests generate and check that they hash equally. Where the hasher needs a rule that is not
obvious, the rule is taken from the equality guide, not the other way round. See
the hashing page.
Packages that depend on the equality API¶
| Package | Use |
|---|---|
coola.io, coola.nested |
import objects_are_equal from coola.equality, not from its private modules |
coola.registry (equality of two registries) |
testers are compared with their __eq__ |
coola.testing |
fixtures and helpers compare with compare and assert_* |