diff --git a/rayforce/capi/raypy_kdb.c b/rayforce/capi/raypy_kdb.c index 8266c92..b6c9102 100644 --- a/rayforce/capi/raypy_kdb.c +++ b/rayforce/capi/raypy_kdb.c @@ -1,3 +1,4 @@ +#define _GNU_SOURCE /* * raypy_kdb.c — KDB+ IPC client. * diff --git a/rayforce/types/base.py b/rayforce/types/base.py index 0b21c29..159229d 100644 --- a/rayforce/types/base.py +++ b/rayforce/types/base.py @@ -140,6 +140,12 @@ def __rfloordiv__(self, other) -> t.Any: def __mod__(self, other) -> t.Any: return _eval_operation("MODULO", self, other) + def __pow__(self, other) -> t.Any: + return _eval_operation("POW", self, other) + + def __rpow__(self, other) -> t.Any: + return _eval_operation("POW", other, self) + class AriphmeticScalarMixin(Scalar, _AriphmeticMixin): ... @@ -246,6 +252,14 @@ def median(self) -> t.Any: def deviation(self) -> t.Any: return _eval_operation("DEVIATION", self) + def std(self) -> t.Any: + # Sample std (ddof=1), matching polars/pandas. Engine verb + # `stddev`. For population std (ddof=0) use `.deviation()`. + return _eval_operation("STDDEV", self) + + def pearson_corr(self, other: t.Any) -> t.Any: + return _eval_operation("PEARSON_CORR", self, other) + def min(self) -> t.Any: return _eval_operation("MIN", self) @@ -266,6 +280,14 @@ def take(self, i: int) -> t.Any: def at(self, index: t.Any) -> t.Any: return _eval_operation("AT", self, index) + def top(self, n: int) -> t.Any: + # n largest values, descending. Per-group inside .by(...) gives + # canonical H2O q8 (largest-N per id6) without a full sort. + return _eval_operation("TOP", self, n) + + def bot(self, n: int) -> t.Any: + return _eval_operation("BOT", self, n) + class SetOperationContainerMixin(Container): def except_(self, other: t.Any) -> t.Any: diff --git a/rayforce/types/operators.py b/rayforce/types/operators.py index 81d4a72..08e2d6b 100644 --- a/rayforce/types/operators.py +++ b/rayforce/types/operators.py @@ -20,6 +20,7 @@ class Operation(enum.StrEnum): MODULO = "%" DIV_INT = "div" NEGATE = "neg" + POW = "pow" # Comparison EQUALS = "==" @@ -45,7 +46,9 @@ class Operation(enum.StrEnum): FIRST = "first" LAST = "last" MEDIAN = "med" - DEVIATION = "dev" + DEVIATION = "dev" # population std, ddof=0 + STDDEV = "stddev" # sample std, ddof=1 (matches polars/pandas) + PEARSON_CORR = "pearson_corr" ROW = "row" # Statistical @@ -63,6 +66,8 @@ class Operation(enum.StrEnum): REVERSE = "reverse" GROUP = "group" TAKE = "take" + TOP = "top" + BOT = "bot" REMOVE = "remove" FILTER = "filter" FIND = "find" diff --git a/rayforce/types/table.py b/rayforce/types/table.py index 509a1e1..714e5a1 100644 --- a/rayforce/types/table.py +++ b/rayforce/types/table.py @@ -201,6 +201,28 @@ def min(self) -> Expression: def median(self) -> Expression: return Expression(Operation.MEDIAN, self) + def deviation(self) -> Expression: + # Population std (ddof=0). Engine verb `dev`. For sample std + # matching polars/pandas, use `.std()` instead. + return Expression(Operation.DEVIATION, self) + + def std(self) -> Expression: + # Sample std (ddof=1), matching polars `pl.std` and pandas + # `.std()`. Engine verb `stddev`. + return Expression(Operation.STDDEV, self) + + def pearson_corr(self, other: t.Any) -> Expression: + # Pearson correlation between two columns; per-group inside + # .by(...) closes canonical H2O q9 (with `**2` for r²). + return Expression(Operation.PEARSON_CORR, self, other) + + def top(self, n: int) -> Expression: + # n largest values per group (descending). Closes q8. + return Expression(Operation.TOP, self, n) + + def bot(self, n: int) -> Expression: + return Expression(Operation.BOT, self, n) + def distinct(self) -> Expression: return Expression(Operation.DISTINCT, self) @@ -253,6 +275,12 @@ def __floordiv__(self, other) -> Expression: def __mod__(self, other) -> Expression: return Expression(Operation.MODULO, self, other) + def __pow__(self, other) -> Expression: + return Expression(Operation.POW, self, other) + + def __rpow__(self, other) -> Expression: + return Expression(Operation.POW, other, self) + def __eq__(self, other) -> Expression: # type: ignore[override] return Expression(Operation.EQUALS, self, other)