nini-datacomp 0.1.0__py3-none-win_amd64.whl

nini_datacomp/__init__.py

@@ -0,0 +1,55 @@
+"""nini-datacomp: comparação de tabelas em escala, com kernels C++ sobre as colunas do Polars.
+
+Reescrita do datacompolars::
+
+    import nini_datacomp as dc
+
+    result = dc.compare("esquerda/", "direita.parquet", on=["id"], abs_tol={"valor": 0.005}, memory_limit="2GB")
+    result.is_match, result.rows.mismatched, result.columns["valor"].mismatches
+"""
+
+from __future__ import annotations
+
+from importlib import metadata
+
+# antes de qualquer import do Polars: o alocador dele lê a configuração uma vez só, quando carrega
+from .core.allocator import configure as _configure_allocator
+
+ALLOCATOR_STATUS = _configure_allocator()
+
+from .api import compare  # noqa: E402
+from .core.memory import pool_limit, pool_stats, release_memory  # noqa: E402
+from .errors import AllocatorWarning, NativeLibraryError, UnsupportedTypeError  # noqa: E402
+from .options import CompareOptions  # noqa: E402
+from .results import (  # noqa: E402
+    ColumnResult,
+    ColumnSummary,
+    ComparisonResult,
+    ExecutionInfo,
+    RowSummary,
+    TypeMismatch,
+)
+
+try:
+    __version__ = metadata.version("nini-datacomp")
+except metadata.PackageNotFoundError:  # pragma: no cover (rodando de src/ sem instalar)
+    __version__ = "0+desconhecida"
+
+__all__ = [
+    "ALLOCATOR_STATUS",
+    "AllocatorWarning",
+    "ColumnResult",
+    "ColumnSummary",
+    "CompareOptions",
+    "ComparisonResult",
+    "ExecutionInfo",
+    "NativeLibraryError",
+    "RowSummary",
+    "TypeMismatch",
+    "UnsupportedTypeError",
+    "__version__",
+    "compare",
+    "pool_limit",
+    "pool_stats",
+    "release_memory",
+]

nini_datacomp/_native.py

@@ -0,0 +1,121 @@
+"""Carrega a biblioteca C++ (libdatacomp) via ctypes e aplica as assinaturas de `_signatures.py`.
+
+Ordem de busca:
+1. variável de ambiente ``NINI_DATACOMP_LIB`` (caminho do arquivo; útil para testar um build manual);
+2. a pasta do pacote (instalação normal por wheel);
+3. os arquivos registrados na instalação (instalação editável do ``uv sync``, em que o pacote aponta para ``src/`` e a
+   biblioteca compilada fica no site-packages).
+
+O ctypes solta o GIL durante cada chamada, então os blocos rodam em paralelo de verdade numa ThreadPool.
+Copiado do nini-datagen 0.3.0, sem a inicialização de tabelas (aqui não há tabelas a montar no carregamento).
+"""
+
+from __future__ import annotations
+
+import ctypes
+import os
+import sys
+import threading
+from importlib import metadata
+from pathlib import Path
+
+import numpy as np
+
+from ._signatures import ABI_VERSION, SIGNATURES
+from .errors import NativeLibraryError
+
+__all__ = ["ENV_VAR", "LIB_NAMES", "lib", "library_path", "ptr"]
+
+ENV_VAR = "NINI_DATACOMP_LIB"
+if sys.platform == "win32":
+    LIB_NAMES: tuple[str, ...] = ("datacomp.dll",)
+elif sys.platform == "darwin":
+    LIB_NAMES = ("libdatacomp.dylib", "libdatacomp.so")
+else:
+    LIB_NAMES = ("libdatacomp.so",)
+
+_lock = threading.Lock()
+_lib: ctypes.CDLL | None = None
+_path: Path | None = None
+
+
+def _candidates() -> list[Path]:
+    env = os.environ.get(ENV_VAR)
+    if env:
+        return [Path(env)]  # pedido explícito: não procura em outro lugar
+    here = Path(__file__).resolve().parent
+    found = [here / name for name in LIB_NAMES]
+    try:
+        files = metadata.files("nini-datacomp") or []
+    except metadata.PackageNotFoundError:
+        files = []
+    found += [Path(str(f.locate())) for f in files if f.name in LIB_NAMES]
+    return found
+
+
+def _load(path: Path) -> ctypes.CDLL:
+    try:
+        handle = ctypes.CDLL(str(path))
+    except OSError as e:
+        raise NativeLibraryError(f"não consegui carregar {path}: {e}") from e
+    try:
+        version = handle.dc_abi_version()
+    except AttributeError as e:
+        raise NativeLibraryError(f"{path} não é a biblioteca do nini-datacomp (falta dc_abi_version)") from e
+    if version != ABI_VERSION:
+        raise NativeLibraryError(
+            f"{path} tem ABI {version}, mas o pacote Python espera {ABI_VERSION}: recompile"
+            " (uv sync --reinstall-package nini-datacomp) ou reinstale o pacote"
+        )
+    missing = [name for name in SIGNATURES if not hasattr(handle, name)]
+    if missing:
+        raise NativeLibraryError(f"{path} não exporta {missing}")
+    for name, (restype, argtypes) in SIGNATURES.items():
+        fn = getattr(handle, name)
+        fn.restype = restype
+        fn.argtypes = argtypes
+    return handle
+
+
+def lib() -> ctypes.CDLL:
+    """A biblioteca carregada (carrega na primeira chamada; seguro entre threads)."""
+    global _lib, _path
+    if _lib is not None:
+        return _lib
+    with _lock:
+        if _lib is None:
+            tried = []
+            for path in _candidates():
+                if path.is_file():
+                    _lib, _path = _load(path), path
+                    break
+                tried.append(str(path))
+            else:
+                raise NativeLibraryError(
+                    "biblioteca C++ do nini-datacomp não encontrada (procurei em: "
+                    + ", ".join(tried)
+                    + "). Numa cópia do repositório, rode `uv sync`; para um build manual, aponte "
+                    + f"{ENV_VAR} para o arquivo."
+                )
+    return _lib
+
+
+def library_path() -> Path:
+    """Caminho do arquivo da biblioteca carregada."""
+    lib()
+    assert _path is not None
+    return _path
+
+
+def _reset() -> None:
+    """Só para testes: esquece a biblioteca carregada (a próxima chamada a `lib()` procura de novo)."""
+    global _lib, _path
+    with _lock:
+        _lib, _path = None, None
+
+
+def ptr(a: np.ndarray) -> int:
+    """Endereço do primeiro elemento de um array contíguo (o C++ escreve/lê nele como um vetor plano)."""
+    if not a.flags.c_contiguous:
+        raise ValueError("o array precisa ser contíguo (C order) para ir ao C++")
+    return a.ctypes.data

nini_datacomp/_signatures.py

@@ -0,0 +1,27 @@
+"""Tipos dos argumentos de cada função da biblioteca C++ (GERADO por tools/gen_bindings.py; não edite)."""
+from __future__ import annotations
+
+from ctypes import c_double, c_float, c_int, c_int64, c_uint8, c_uint16, c_uint32, c_uint64, c_void_p
+
+ABI_VERSION = 8
+
+# nome: (tipo de retorno, [tipos dos argumentos])
+SIGNATURES = {
+    "dc_abi_version": (c_int, []),
+    "dc_count_unset": (c_int64, [c_void_p, c_int64, c_int64]),
+    "dc_describe_batch": (c_int, [c_void_p, c_void_p, c_int, c_void_p]),
+    "dc_compare_column": (c_int64, [c_void_p, c_void_p, c_void_p, c_void_p, c_void_p, c_int64, c_int64, c_int, c_double, c_double, c_double, c_void_p]),
+    "dc_compare_rows": (c_int, [c_int, c_void_p, c_void_p, c_int64, c_int64, c_int64, c_void_p]),
+    "dc_compare_pairs": (c_int, [c_int, c_int, c_void_p, c_void_p, c_int, c_void_p, c_void_p, c_void_p, c_void_p, c_void_p, c_int64, c_int64, c_int64, c_void_p]),
+    "dc_align_keys": (c_int, [c_int, c_void_p, c_int64, c_void_p, c_int64, c_int, c_void_p, c_void_p, c_void_p]),
+    "dc_pool_alloc": (c_int64, [c_int64]),
+    "dc_pool_free": (None, [c_int64]),
+    "dc_pool_limit": (c_int64, [c_int64]),
+    "dc_pool_trim": (None, []),
+    "dc_pool_stats": (None, [c_void_p]),
+    "dc_parquet_row_groups": (c_int64, [c_void_p, c_int64, c_void_p, c_int64]),
+    "dc_row_hash": (c_int, [c_int, c_void_p, c_int64, c_int64, c_int, c_void_p]),
+    "dc_diff_pairs": (c_int64, [c_int, c_void_p, c_void_p, c_void_p, c_void_p, c_void_p, c_void_p, c_int64, c_int64, c_int, c_void_p, c_void_p]),
+}
+
+__all__ = ['ABI_VERSION', 'SIGNATURES', 'c_double', 'c_float', 'c_int', 'c_int64', 'c_uint8', 'c_uint16', 'c_uint32', 'c_uint64', 'c_void_p']

nini_datacomp/api.py

@@ -0,0 +1,250 @@
+"""`compare`: a função pública.
+
+Abre as fontes, ajusta os nomes das colunas, confere a chave (os erros fatais vão para ``result.errors``, como no
+datacompolars) e roda o motor (`pipeline._run`): janelas que cabem no ``memory_limit``, hash de linha nas colunas
+exatas, coluna a coluna nas com tolerância, contagens por coluna e amostras.
+"""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Sequence
+from typing import Any
+
+import polars as pl
+
+from .io import Source, SourceLike, open_source
+from .options import CompareOptions
+from .pipeline import RowCounts, _run, _tol_of, align_columns, sample_key_names
+from .results import ColumnResult, ColumnSummary, ComparisonResult, ExecutionInfo, RowSummary, TypeMismatch
+
+__all__ = ["compare"]
+
+
+def compare(
+    left: SourceLike,
+    right: SourceLike,
+    on: str | Sequence[str],
+    *,
+    options: CompareOptions | None = None,
+    **kwargs: Any,
+) -> ComparisonResult:
+    """Compara ``left`` e ``right`` (caminho de Parquet, pasta, glob, ``pl.DataFrame`` ou ``pl.LazyFrame``) pela chave
+    ``on`` (uma coluna ou várias).
+
+    As opções vêm em ``options`` (`CompareOptions`) ou como argumentos nomeados (``abs_tol=0.01``); com os dois, os
+    argumentos nomeados ganham. Erros fatais (chave que falta, de tipos diferentes, repetida) não levantam exceção:
+    ficam em ``result.errors``, com ``result.aborted`` verdadeiro. Tipos que o pacote não compara (List, Struct)
+    levantam `UnsupportedTypeError`.
+    """
+    t0 = time.perf_counter()
+    if kwargs:
+        base = options.model_dump() if options is not None else {}
+        options = CompareOptions(**{**base, **kwargs})
+    opts = options if options is not None else CompareOptions()
+    lower = not opts.case_sensitive_columns
+    keys = _keys(on, lower=lower)
+    abs_tol, rel_tol = _names_of(opts.abs_tol, lower), _names_of(opts.rel_tol, lower)
+
+    ls, rs = open_source(left), open_source(right)
+    errors: list[str] = []
+    if lower:
+        ls, rs = _lowercase(ls, "left", errors), _lowercase(rs, "right", errors)
+    if errors:
+        return _failed(opts, keys, errors, left_only=(), right_only=(), t0=t0)
+    lschema, rschema = ls.lazy.collect_schema(), rs.lazy.collect_schema()
+    left_only = tuple(c for c in lschema.names() if c not in rschema)
+    right_only = tuple(c for c in rschema.names() if c not in lschema)
+    errors += _key_errors(lschema, rschema, keys, opts.cast)
+    if errors:
+        return _failed(opts, keys, errors, left_only=left_only, right_only=right_only, t0=t0)
+
+    cols = align_columns(ls, rs, keys, opts.cast)  # schemas só; UnsupportedTypeError para aninhadas
+    warnings: list[str] = []
+    if not cols.compared:
+        warnings.append("No common columns besides the join key: only row membership was compared.")
+    for label, tol in (("abs_tol", abs_tol), ("rel_tol", rel_tol)):
+        if isinstance(tol, dict) and (unknown := sorted(set(tol) - set(cols.compared))):
+            warnings.append(f"{label} names columns that are not compared: {unknown}")
+
+    res = _run(
+        left, right, keys, memory_limit=opts.memory_limit, window_rows=opts.window_rows, columns=None, prefetch=True,
+        engine="auto", hashing=True, casting=opts.cast, verify_hash=opts.verify_hash, threads=opts.threads,
+        abs_tol=abs_tol, rel_tol=rel_tol, ignore_case=opts.ignore_case, ignore_spaces=opts.ignore_whitespace,
+        details=opts.details, samples=opts.samples, sources=(ls, rs),
+    )  # fmt: skip
+    if res.dup_left:
+        errors.append(f"Duplicate join keys in the left table: {res.dup_left:,} surplus rows")
+    if res.dup_right:
+        errors.append(f"Duplicate join keys in the right table: {res.dup_right:,} surplus rows")
+
+    dtypes = {c: cols.left_casts.get(c, lschema[c]) for c in lschema.names()}
+    key_dtypes = [_sample_dtype(dtypes[k]) for k in keys]
+    details = opts.details and not errors
+    results = {}
+    for c in res.compared_columns:
+        n = res.column_mismatches.get(c, 0) if details else None
+        tol = (_tol_of(abs_tol, c), _tol_of(rel_tol, c)) if c in res.tolerance_columns else None
+        results[c] = ColumnResult(
+            name=c,
+            dtype=str(dtypes[c]),
+            mismatches=n,
+            rate=None if n is None else (n / res.common if res.common else 0.0),
+            samples=_samples(
+                res, c, keys=keys, key_dtypes=key_dtypes, dtype=_sample_dtype(dtypes[c]), names=opts.names
+            ),
+            tolerance=tol,
+        )
+    columns = ColumnSummary(
+        join_columns=tuple(keys),
+        results=results,
+        left_only=left_only,
+        right_only=right_only,
+        type_mismatches=tuple(TypeMismatch(c, a, b) for c, a, b in res.type_mismatches),
+    )
+    return ComparisonResult(
+        names=opts.names,
+        join_columns=tuple(keys),
+        rows=_rows(res, aborted=bool(errors)),
+        columns=columns,
+        errors=tuple(errors),
+        warnings=tuple(warnings),
+        execution=_execution(res, opts, details=details, t0=t0),
+    )
+
+
+# ----------------------------------------------------------------------------- nomes e chave
+def _keys(on: str | Sequence[str], *, lower: bool) -> list[str]:
+    keys = [on] if isinstance(on, str) else list(on)
+    if not keys:
+        raise ValueError("on precisa de ao menos uma coluna")
+    if any(not isinstance(k, str) or not k.strip() for k in keys):
+        raise ValueError("on não pode ter nomes vazios")
+    if lower:
+        keys = [k.lower() for k in keys]
+    if len(set(keys)) != len(keys):
+        raise ValueError(f"on tem colunas repetidas: {keys}")
+    return keys
+
+
+def _names_of(tol: float | dict[str, float], lower: bool) -> float | dict[str, float]:
+    if isinstance(tol, dict) and lower:
+        return {k.lower(): v for k, v in tol.items()}
+    return tol
+
+
+def _lowercase(src: Source, side: str, errors: list[str]) -> Source:
+    """Os nomes das colunas em minúsculas; dois que viram o mesmo são erro fatal (como no datacompolars), e a fonte
+    volta como estava (renomear com a colisão quebraria no Polars)."""
+    names = src.lazy.collect_schema().names()
+    seen: set[str] = set()
+    clash = False
+    for n in names:
+        low = n.lower()
+        if low in seen:
+            errors.append(f"Columns collide in the {side} table after lowercasing names: '{low}'")
+            clash = True
+        seen.add(low)
+    return src if clash else src.rename({n: n.lower() for n in names if n != n.lower()})
+
+
+def _key_errors(lschema: pl.Schema, rschema: pl.Schema, keys: Sequence[str], cast: str) -> list[str]:
+    errors = []
+    for side, schema in (("left", lschema), ("right", rschema)):
+        missing = [k for k in keys if k not in schema]
+        if missing:
+            errors.append(f"Join keys missing from the {side} table: {missing}")
+    if errors or cast != "none":
+        return errors
+    return [
+        f"Join key '{k}' has different types: left {lschema[k]} vs right {rschema[k]}"
+        for k in keys
+        if lschema[k] != rschema[k]
+    ]
+
+
+# ----------------------------------------------------------------------------- resultado
+def _sample_dtype(dtype: pl.DataType) -> pl.DataType:
+    """O tipo da coluna nas amostras: Categorical e Enum são comparados (e mostrados) como texto."""
+    return pl.String() if isinstance(dtype, (pl.Categorical, pl.Enum)) else dtype
+
+
+def _samples(
+    res: RowCounts,
+    column: str,
+    *,
+    keys: Sequence[str],
+    key_dtypes: Sequence[pl.DataType],
+    dtype: pl.DataType,
+    names: tuple[str, str],
+) -> pl.DataFrame:
+    headers = [*keys, f"{column} ({names[0]})", f"{column} ({names[1]})"]
+    frame = res.samples.get(column)
+    if frame is None:
+        return pl.DataFrame(schema=dict(zip(headers, [*key_dtypes, dtype, dtype], strict=True)))
+    raw = [*sample_key_names(len(keys)), "__left", "__right"]
+    return frame.rename(dict(zip(raw, headers, strict=True)))
+
+
+def _rows(res: RowCounts, *, aborted: bool) -> RowSummary:
+    if aborted:  # como no datacompolars: nada contado
+        return RowSummary(0, 0, 0, 0, 0, 0)
+    return RowSummary(
+        left_total=res.left_total,
+        right_total=res.right_total,
+        common=res.common,
+        left_only=res.left_only,
+        right_only=res.right_only,
+        mismatched=res.mismatched or 0,
+    )
+
+
+def _versions() -> tuple[str, str]:
+    from . import __version__  # noqa: PLC0415 - o pacote já carregou quando compare roda
+
+    return __version__, pl.__version__
+
+
+def _execution(res: RowCounts, opts: CompareOptions, *, details: bool, t0: float) -> ExecutionInfo:
+    version, polars_version = _versions()
+    return ExecutionInfo(
+        windows=res.windows,
+        rows_per_window=res.rows_per_window,
+        column_groups=res.column_groups,
+        memory_limit=res.memory_limit,
+        window_bytes=res.window_bytes,
+        tolerance_columns=res.tolerance_columns,
+        verify_hash=opts.verify_hash,
+        hash_collisions=res.hash_collisions,
+        details_computed=details,
+        version=version,
+        polars_version=polars_version,
+        total_seconds=round(time.perf_counter() - t0, 6),
+        timings=dict(res.timings),
+    )
+
+
+def _failed(
+    opts: CompareOptions,
+    keys: Sequence[str],
+    errors: Sequence[str],
+    *,
+    left_only: tuple[str, ...],
+    right_only: tuple[str, ...],
+    t0: float,
+) -> ComparisonResult:
+    """O resultado de uma comparação que não começou (erro fatal antes de ler os dados)."""
+    version, polars_version = _versions()
+    return ComparisonResult(
+        names=opts.names,
+        join_columns=tuple(keys),
+        rows=RowSummary(0, 0, 0, 0, 0, 0),
+        columns=ColumnSummary(join_columns=tuple(keys), results={}, left_only=left_only, right_only=right_only),
+        errors=tuple(errors),
+        execution=ExecutionInfo(
+            verify_hash=opts.verify_hash,
+            version=version,
+            polars_version=polars_version,
+            total_seconds=round(time.perf_counter() - t0, 6),
+        ),
+    )

nini_datacomp/core/__init__.py

@@ -0,0 +1 @@
+"""Infraestrutura interna: threads (`blocks`), memória (`memory`) e a interface C do Arrow (`arrow_c`, `arrow`)."""

nini_datacomp/core/allocator.py

@@ -0,0 +1,60 @@
+"""O alocador do Polars: no Windows ele é o mimalloc, que por padrão guarda a memória liberada em vez de devolvê-la.
+
+Medido na fase 1c (``tools/diag_memory.py``, 10M linhas, ``memory_limit = 1GB``): com as janelas de ~200 MB, o
+processo ia acumulando memória janela após janela e terminava com ~1,1 GB, quase toda já liberada pelo Polars. Com
+``MIMALLOC_PURGE_DELAY=0`` o mimalloc devolve a memória na hora: pico de ~0,6–0,7 GB (−40 a −48%), sem custo de tempo
+nas tabelas largas (10M x 50: 5,6 s → 4,2 s) e +0,3 s nas estreitas (10M x 10).
+
+O mimalloc lê a variável uma vez, quando a biblioteca do Polars carrega; depois não há como mudar. Por isso o
+``nini_datacomp`` a define no próprio import, antes de importar o Polars, e só se o Polars ainda não estiver carregado
+(e sem sobrescrever um valor que já exista). Se o Polars veio antes, avisa com `AllocatorWarning`.
+
+``NINI_DATACOMP_ALLOCATOR=0`` desliga tudo isso (nem define a variável, nem avisa).
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+import warnings
+from collections.abc import Container, MutableMapping
+from typing import Literal
+
+from ..errors import AllocatorWarning
+
+__all__ = ["ALLOCATOR_ENV", "SETTINGS", "Status", "configure"]
+
+ALLOCATOR_ENV = "NINI_DATACOMP_ALLOCATOR"
+SETTINGS = {"MIMALLOC_PURGE_DELAY": "0", "MIMALLOC_RESET_DELAY": "0"}  # RESET_DELAY: o nome antigo, mimalloc < 2.1
+
+# "set": definida agora; "user": já estava definida; "late": o Polars já tinha carregado; "disabled": desligado
+Status = Literal["set", "user", "late", "disabled"]
+
+_LATE = (
+    "o Polars foi importado antes do nini_datacomp, então o alocador dele (mimalloc) vai guardar a memória que libera "
+    "e o pico de uma comparação grande pode ficar bem maior (medido: até ~1,9x). Para evitar, importe o nini_datacomp "
+    "antes do Polars ou defina MIMALLOC_PURGE_DELAY=0 antes de iniciar o Python. NINI_DATACOMP_ALLOCATOR=0 silencia "
+    "este aviso."
+)
+
+
+def configure(
+    environ: MutableMapping[str, str] | None = None,
+    modules: Container[str] | None = None,
+    platform: str | None = None,
+) -> Status:
+    """Define as variáveis do mimalloc se ainda der tempo; avisa (só no Windows, onde foi medido) se não der."""
+    environ = os.environ if environ is None else environ
+    modules = sys.modules if modules is None else modules
+    platform = sys.platform if platform is None else platform
+    if environ.get(ALLOCATOR_ENV, "1") == "0":
+        return "disabled"
+    if "MIMALLOC_PURGE_DELAY" in environ:
+        return "user"
+    if "polars" in modules:
+        if platform == "win32":
+            warnings.warn(_LATE, AllocatorWarning, stacklevel=2)
+        return "late"
+    for name, value in SETTINGS.items():
+        environ.setdefault(name, value)
+    return "set"