Интеграция с Arrow и Polars

Колоночный обмен данными без копирования

NextStat нативно поддерживает Apache Arrow. Вы можете загружать гистограммные таблицы из PyArrow, Polars, DuckDB или Spark и возвращать результаты обратно в Arrow, не превращая данные в Python-объекты. Под капотом используется Arrow IPC, а реализация опирается на Rust-крейты arrow 57.3 и parquet 57.3.

PyArrow Table / Polars DataFrame / DuckDB Result
        │
        ▼  .to_arrow() or native
  Arrow RecordBatch
        │
        ▼  IPC serialize (~1 memcpy)
  nextstat.from_arrow(table) ──► Rust arrow crate
        │                            │
        ▼                            ▼
  HistFactoryModel            Arrow RecordBatch
        │                            │
        ▼                            ▼  IPC deserialize
  nextstat.to_arrow(model)    PyArrow Table

Быстрый старт

import pyarrow as pa
import nextstat

# Гистограммные данные в виде Arrow-таблицы
table = pa.table({
    "channel": ["SR", "SR", "CR"],
    "sample":  ["signal", "background", "background"],
    "yields":  [[5., 10., 15.], [100., 200., 150.], [500., 600.]],
    "stat_error": [[1., 2., 3.], [10., 14., 12.], None],
})

# Модель и фит
model = nextstat.from_arrow(table, poi="mu")
result = nextstat.fit(model)
print(result)

Схема таблицы

Входная Arrow-таблица должна соответствовать контракту ниже. Каждая строка описывает одну пару (channel, sample).

Polars

import polars as pl
import nextstat

# Прочитать гистограммы из Parquet через Polars
df = pl.read_parquet("histograms.parquet")
model = nextstat.from_arrow(df.to_arrow(), poi="mu")

# Или читать Parquet напрямую (нативный Rust, без Python-оверхеда)
model = nextstat.from_parquet("histograms.parquet", poi="mu")

DuckDB

import duckdb
import nextstat

con = duckdb.connect()
table = con.sql("""
    SELECT channel, sample, yields
    FROM 'histograms.parquet'
""").arrow()

model = nextstat.from_arrow(table)

Экспорт

Экспортируйте данные модели обратно в Arrow для последующего анализа, дашбордов или ML-пайплайнов.

model = nextstat.from_pyhf(workspace_json)

# Ожидаемые выходы по каналам
yields = nextstat.to_arrow(model, what="yields")
print(yields.to_pandas())
#   channel sample              yields
# 0      CR  total      [500.0, 600.0]
# 1      SR  total  [105.0, 210.0, 165.0]

# Метаданные параметров
params = nextstat.to_arrow(model, what="params")
print(params.to_pandas())
#              name  index  value  bound_lo  bound_hi  init
# 0              mu      0    1.0       0.0      10.0   1.0
# 1  staterror_SR[0]    1    1.0       1e-10    10.0   1.0

Наблюдаемые данные

# По умолчанию используется набор данных Асимова (Asimov dataset): сумма ожидаемых выходов.
# Передайте наблюдаемые значения явно:
model = nextstat.from_arrow(
    table,
    poi="mu",
    observations={
        "SR": [110., 215., 170.],
        "CR": [510., 590.],
    },
)

Низкоуровневый IPC API

Для максимального контроля можно работать напрямую с IPC-байтами. Именно так устроены from_arrow() и to_arrow() внутри.

import pyarrow as pa

# Сериализовать таблицу в IPC-байты
sink = pa.BufferOutputStream()
writer = pa.ipc.new_stream(sink, table.schema)
for batch in table.to_batches():
    writer.write_batch(batch)
writer.close()
ipc_bytes = sink.getvalue().to_pybytes()

# Принять IPC-байты напрямую
model = nextstat.from_arrow_ipc(ipc_bytes, poi="mu")

# Экспортировать как IPC-байты
yields_bytes = nextstat.to_arrow_yields_ipc(model)
params_bytes = nextstat.to_arrow_params_ipc(model)

# Десериализовать в Python
yields_table = pa.ipc.open_stream(yields_bytes).read_all()

Справка API

• nextstat.from_arrow(table, *, poi, observations) — таблица PyArrow/RecordBatch → HistFactoryModel.
• nextstat.to_arrow(model, *, params, what) — HistFactoryModel → таблица PyArrow. what="yields" или "params".
• nextstat.from_parquet(path, *, poi, observations) — файл Parquet → HistFactoryModel (нативный Rust-ридер).
• nextstat.from_arrow_ipc(bytes, poi, observations) — байты сырого IPC-потока → HistFactoryModel.
• nextstat.to_arrow_yields_ipc(model, params) — HistFactoryModel → IPC-байты (yields).
• nextstat.to_arrow_params_ipc(model, params) — HistFactoryModel → IPC-байты (параметры).