From 1015f0673962b516468460ae6dd66be11323debf Mon Sep 17 00:00:00 2001 From: Travis Hammond Date: Mon, 31 Aug 2026 09:05:34 +0200 Subject: [PATCH 01/17] fix: enable polars dtype-array so Array input raises instead of aborting the process Without the feature, polars-core's arrow-FFI import has no FixedSizeList arm compiled in and hits a non-unwinding panic before any plugin entry check runs, taking the whole host process down (exit 0xC0000409, no Python traceback). With it, the plugin's own dtype check rejects Array input with a clean ComputeError naming the received dtype. Array support proper comes later; this is the independent bug fix. Co-Authored-By: Claude Fable 5 --- Cargo.lock | 1557 +++++++++++++++++++++++++++++++++++++++-- Cargo.toml | 4 +- tests/test_plugins.py | 15 + 3 files changed, 1528 insertions(+), 48 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0ae9c67..41354a4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -56,6 +56,63 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed51fe0f224d1d4ea768be38c51f9f831dee9d05c163c11fba0b8c44387b1fc3" +[[package]] +name = "arrayref" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb" + +[[package]] +name = "arrayvec" +version = "0.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56" + +[[package]] +name = "async-channel" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "924ed96dd52d1b75e9c1a3e6275715fd320f5f9439fb5a4a11fa51f4221158d2" +dependencies = [ + "concurrent-queue", + "event-listener-strategy", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-stream" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476" +dependencies = [ + "async-stream-impl", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-stream-impl" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "async-trait" +version = "0.1.92" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82f6aeea286b8eb4dd3431a1be1b59d290ace00f5bfd8e2a159bc2a05e2c1667" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + [[package]] name = "atoi_simd" version = "0.16.1" @@ -65,6 +122,12 @@ dependencies = [ "debug_unsafe", ] +[[package]] +name = "atomic-waker" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" + [[package]] name = "autocfg" version = "1.5.1" @@ -106,6 +169,20 @@ dependencies = [ "serde_core", ] +[[package]] +name = "blake3" +version = "1.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0aa83c34e62843d924f905e0f5c866eb1dd6545fc4d719e803d9ba6030371fce" +dependencies = [ + "arrayref", + "arrayvec", + "cc", + "cfg-if", + "constant_time_eq", + "cpufeatures 0.3.0", +] + [[package]] name = "block-buffer" version = "0.10.4" @@ -199,6 +276,23 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "cfg_aliases" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f079e83a288787bcd14a6aea84cee5c87a67c5a3e660c30f557a3d24761b3527" + +[[package]] +name = "chacha20" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d524456ba66e72eb8b115ff89e01e497f8e6d11d78b70b1aa13c0fbd97540a81" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "rand_core 0.10.1", +] + [[package]] name = "chrono" version = "0.4.45" @@ -236,6 +330,31 @@ dependencies = [ "static_assertions", ] +[[package]] +name = "concurrent-queue" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ca0197aee26d1ae37445ee532fefce43251d24cc7c166799f4d46817f1d3973" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "constant_time_eq" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" + +[[package]] +name = "core-foundation" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2a6cd9ae233e7f62ba4e9353e81a88df7fc8a5987b8d445b4d90c879bd156f6" +dependencies = [ + "core-foundation-sys", + "libc", +] + [[package]] name = "core-foundation-sys" version = "0.8.7" @@ -251,6 +370,15 @@ dependencies = [ "libc", ] +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + [[package]] name = "crc32fast" version = "1.5.0" @@ -260,6 +388,15 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "crossbeam-channel" +version = "0.5.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d85363c37faeca707aef026efa9f3b34d077bce547e48f770770625c6013679e" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-deque" version = "0.8.7" @@ -279,6 +416,15 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "crossbeam-queue" +version = "0.3.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "803d13fb3b09d88be9f4dbc29062c66b19bf7170867ceb746d2a8689bf6c7a26" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-utils" version = "0.8.22" @@ -311,6 +457,17 @@ dependencies = [ "crypto-common", ] +[[package]] +name = "displaydoc" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6232dd377dcc64799954cbd3a9bb882e9cdc1308ccd87b1c098f1fb2eaf82a8" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + [[package]] name = "dyn-clone" version = "1.0.20" @@ -339,7 +496,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -348,6 +505,27 @@ version = "1.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40404c3f5f511ec4da6fe866ddf6a717c309fdbb69fbbad7b0f3edab8f2e835f" +[[package]] +name = "event-listener" +version = "5.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13b66accf52311f30a0db42147dadea9850cb48cd070028831ae5f5d4b856ab" +dependencies = [ + "concurrent-queue", + "parking", + "pin-project-lite", +] + +[[package]] +name = "event-listener-strategy" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8be9f3dfaaffdae2972880079a491a1a8bb7cbed0b8dd7a347f668b4150a3b93" +dependencies = [ + "event-listener", + "pin-project-lite", +] + [[package]] name = "fallible-streaming-iterator" version = "0.1.9" @@ -376,6 +554,12 @@ dependencies = [ "miniz_oxide", ] +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + [[package]] name = "foldhash" version = "0.1.5" @@ -388,12 +572,90 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" +[[package]] +name = "form_urlencoded" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb4cb245038516f5f85277875cdaa4f7d2c9a0fa0468de06ed190163b1581fcf" +dependencies = [ + "percent-encoding", +] + +[[package]] +name = "fs4" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8640e34b88f7652208ce9e88b1a37a2ae95227d84abec377ccd3c5cfeb141ed4" +dependencies = [ + "rustix", + "windows-sys 0.59.0", +] + +[[package]] +name = "futures" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a88cf1f829d945f548cf8fec32c61b1f202b6d93b45848602fc02af4b12ad218" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "262590f4fe6afeb0bc83be1daa64e52657fe185690a958af7f3ad0e92085c5ae" +dependencies = [ + "futures-core", + "futures-sink", +] + [[package]] name = "futures-core" version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2cd50c473c80f6d7c3670a752354b8e569b1a7cbfdc0419ec88e5edad85e0dc7" +[[package]] +name = "futures-executor" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6754879cc9f2c66f88c6e5c35344bb0bdb0708b0352b1201815667c7eabc7458" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" + +[[package]] +name = "futures-macro" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6d3cde68c518367be28956066ddfef33813991b77a55005a69dae04bf3b10b" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "futures-sink" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" + [[package]] name = "futures-task" version = "0.3.33" @@ -406,8 +668,13 @@ version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a77a90a256fce34da66415271e30f94ee91c57b04b8a2c042d9cf3220179deaa" dependencies = [ + "futures-channel", "futures-core", + "futures-io", + "futures-macro", + "futures-sink", "futures-task", + "memchr", "pin-project-lite", "slab", ] @@ -456,8 +723,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" dependencies = [ "cfg-if", + "js-sys", "libc", "r-efi 6.0.0", + "rand_core 0.10.1", + "wasm-bindgen", ] [[package]] @@ -466,6 +736,25 @@ version = "0.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" +[[package]] +name = "h2" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "839c0e8a181239723652be9062bb56ca5bf5f64011f73b623f6f4fc59086a228" +dependencies = [ + "atomic-waker", + "bytes", + "fnv", + "futures-core", + "futures-sink", + "http", + "indexmap", + "slab", + "tokio", + "tokio-util", + "tracing", +] + [[package]] name = "hashbrown" version = "0.15.5" @@ -509,7 +798,112 @@ version = "0.5.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" dependencies = [ - "windows-sys", + "windows-sys 0.61.2", +] + +[[package]] +name = "http" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "918d3568bebf352712bc2ef3d46a8bcf1a75b373be6539de198e9105cbbf9ce0" +dependencies = [ + "bytes", + "itoa", +] + +[[package]] +name = "http-body" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca2a8f2913ee65f60facd6a5905613afaa448497a0230cc41ce022d93290bc2c" +dependencies = [ + "bytes", + "http", +] + +[[package]] +name = "http-body-util" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23169fe34a5fbcdd3f3862e78fb9b6fccd5f02a6dc6f732547005d45631ce71c" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "pin-project-lite", +] + +[[package]] +name = "httparse" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" + +[[package]] +name = "humantime" +version = "2.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15cdd26707701c53297e2fa6afb323d55fbc1d0810c3aec078ae3ef0424c3c15" + +[[package]] +name = "hyper" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d22053281f852e11534f5198498373cbb59295120a20771d90f7ed1897490a72" +dependencies = [ + "atomic-waker", + "bytes", + "futures-channel", + "futures-core", + "h2", + "http", + "http-body", + "httparse", + "itoa", + "pin-project-lite", + "smallvec", + "tokio", + "want", +] + +[[package]] +name = "hyper-rustls" +version = "0.27.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f" +dependencies = [ + "http", + "hyper", + "hyper-util", + "rustls", + "rustls-native-certs", + "tokio", + "tokio-rustls", + "tower-service", +] + +[[package]] +name = "hyper-util" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" +dependencies = [ + "base64", + "bytes", + "futures-channel", + "futures-util", + "http", + "http-body", + "hyper", + "ipnet", + "libc", + "percent-encoding", + "pin-project-lite", + "socket2", + "tokio", + "tower-service", + "tracing", ] [[package]] @@ -537,47 +931,157 @@ dependencies = [ ] [[package]] -name = "indexmap" -version = "2.14.0" +name = "icu_collections" +version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +checksum = "fa68d21081c4a05d5a901a1c62add574c77048b6a1c67be3b50ce0b60d4ca513" dependencies = [ - "equivalent", - "hashbrown 0.17.1", - "serde", - "serde_core", + "displaydoc", + "potential_utf", + "utf8_iter", + "yoke", + "zerofrom", + "zerovec", ] [[package]] -name = "indoc" -version = "2.0.7" +name = "icu_locale_core" +version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" +checksum = "d56e28588da92eee5c3201a6eff33fabdd49b62269c8938d4ff050ce4d900deb" dependencies = [ - "rustversion", + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", ] [[package]] -name = "interp" -version = "2.1.2" +name = "icu_normalizer" +version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "341be68f8dda23c0a581f926bd0cd346500e42a64649dc2ad137b3306c12440f" +checksum = "12f9cf5f235641ed274641dd81c3f28d870e276763d0797aeeab72317b1c646f" dependencies = [ - "itertools", - "num-traits", + "icu_collections", + "icu_normalizer_data", + "icu_properties", + "icu_provider", + "smallvec", + "zerovec", ] [[package]] -name = "iter-read" -version = "1.1.0" +name = "icu_normalizer_data" +version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "071ed4cc1afd86650602c7b11aa2e1ce30762a1c27193201cb5cee9c6ebb1294" +checksum = "1563da1ed3e0b3bf3d74c9b85917ac9c56464d2f57242270c09c9e752f8021a0" [[package]] -name = "itertools" -version = "0.14.0" +name = "icu_properties" +version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +checksum = "7e7ca276ad3145661a65914e6daf131ca5120cd3dcee8f8f3214b8875184a148" +dependencies = [ + "displaydoc", + "icu_collections", + "icu_locale_core", + "icu_properties_data", + "icu_provider", + "zerotrie", + "zerovec", +] + +[[package]] +name = "icu_properties_data" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e590f038c1464a96894fd6d10127e90a8be4509f56ff7ecef851b15cee0b7caa" + +[[package]] +name = "icu_provider" +version = "2.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d27bbb9d3abbefac45d55f647c9de1d44aafcd1186eb91879afef17c396c3e73" +dependencies = [ + "displaydoc", + "icu_locale_core", + "writeable", + "yoke", + "zerofrom", + "zerotrie", + "zerovec", +] + +[[package]] +name = "idna" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de" +dependencies = [ + "idna_adapter", + "smallvec", + "utf8_iter", +] + +[[package]] +name = "idna_adapter" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb68373c0d6620ef8105e855e7745e18b0d00d3bdb07fb532e434244cdb9a714" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", + "serde", + "serde_core", +] + +[[package]] +name = "indoc" +version = "2.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" +dependencies = [ + "rustversion", +] + +[[package]] +name = "interp" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "341be68f8dda23c0a581f926bd0cd346500e42a64649dc2ad137b3306c12440f" +dependencies = [ + "itertools", + "num-traits", +] + +[[package]] +name = "ipnet" +version = "2.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a756c3fac73139e83f14c2d742155dd2b78d3ee56597b419a0579b7bdd6dd78" + +[[package]] +name = "iter-read" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071ed4cc1afd86650602c7b11aa2e1ce30762a1c27193201cb5cee9c6ebb1294" + +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" dependencies = [ "either", ] @@ -621,6 +1125,18 @@ version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + +[[package]] +name = "litemap" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47d9d19d1d6efa0109d2f65ff4c85cddd50bd572e5a00127ab10987290bcefae" + [[package]] name = "lock_api" version = "0.4.14" @@ -636,6 +1152,12 @@ version = "0.4.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" +[[package]] +name = "lru-slab" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" + [[package]] name = "lz4" version = "1.28.1" @@ -689,6 +1211,17 @@ dependencies = [ "simd-adler32", ] +[[package]] +name = "mio" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30d65c71f1ce40ab09135ce117d742b9f8a19ff91a41a8b57ed50bc2de59c427" +dependencies = [ + "libc", + "wasi", + "windows-sys 0.61.2", +] + [[package]] name = "now" version = "0.1.3" @@ -745,12 +1278,59 @@ dependencies = [ "memchr", ] +[[package]] +name = "object_store" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fbfbfff40aeccab00ec8a910b57ca8ecf4319b335c542f2edcd19dd25a1e2a00" +dependencies = [ + "async-trait", + "base64", + "bytes", + "chrono", + "form_urlencoded", + "futures", + "http", + "http-body-util", + "humantime", + "hyper", + "itertools", + "parking_lot", + "percent-encoding", + "quick-xml", + "rand 0.9.5", + "reqwest", + "ring", + "serde", + "serde_json", + "serde_urlencoded", + "thiserror 2.0.20", + "tokio", + "tracing", + "url", + "walkdir", + "wasm-bindgen-futures", + "web-time", +] + [[package]] name = "once_cell" version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "openssl-probe" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" + +[[package]] +name = "parking" +version = "2.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba" + [[package]] name = "parking_lot" version = "0.12.5" @@ -832,6 +1412,9 @@ dependencies = [ "polars-compute", "polars-core", "polars-error", + "polars-lazy", + "polars-ops", + "polars-plan", "polars-utils", "version_check", ] @@ -893,7 +1476,7 @@ dependencies = [ "polars-arrow", "polars-error", "polars-utils", - "rand", + "rand 0.9.5", "ryu", "serde", "strength_reduce", @@ -923,7 +1506,7 @@ dependencies = [ "polars-row", "polars-schema", "polars-utils", - "rand", + "rand 0.9.5", "rand_distr", "rayon", "regex", @@ -956,6 +1539,7 @@ version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f9a6be22566c89f6405f553bfdb7c8a6cb20ec51b35f3172de9a25fa3e252d85" dependencies = [ + "object_store", "parking_lot", "polars-arrow-format", "pyo3", @@ -964,6 +1548,29 @@ dependencies = [ "simdutf8", ] +[[package]] +name = "polars-expr" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6199a50d3e1afd0674fb009e340cbfb0010682b2387187a36328c00f3f2ca87b" +dependencies = [ + "bitflags", + "hashbrown 0.16.1", + "num-traits", + "polars-arrow", + "polars-compute", + "polars-core", + "polars-io", + "polars-ops", + "polars-plan", + "polars-row", + "polars-utils", + "rand 0.9.5", + "rayon", + "recursive", + "version_check", +] + [[package]] name = "polars-ffi" version = "0.52.0" @@ -980,9 +1587,13 @@ version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "be3714acdff87170141880a07f5d9233490d3bd5531c41898f6969d440feee11" dependencies = [ + "async-trait", "atoi_simd", + "blake3", "bytes", "fast-float2", + "fs4", + "futures", "glob", "hashbrown 0.16.1", "home", @@ -990,6 +1601,7 @@ dependencies = [ "memchr", "memmap2", "num-traits", + "object_store", "percent-encoding", "polars-arrow", "polars-compute", @@ -1000,14 +1612,41 @@ dependencies = [ "polars-utils", "rayon", "regex", + "reqwest", "ryu", "serde", + "serde_json", "simdutf8", + "tokio", +] + +[[package]] +name = "polars-lazy" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea136c360d03aafe56e0233495e30044ce43639b8b0360a4a38e840233f048a1" +dependencies = [ + "bitflags", + "chrono", + "either", + "memchr", + "polars-arrow", + "polars-compute", + "polars-core", + "polars-expr", + "polars-io", + "polars-mem-engine", + "polars-ops", + "polars-plan", + "polars-stream", + "polars-utils", + "rayon", + "version_check", ] [[package]] name = "polars-list-utils" -version = "1.0.0" +version = "1.0.1" dependencies = [ "butterworth", "interp", @@ -1018,6 +1657,25 @@ dependencies = [ "serde", ] +[[package]] +name = "polars-mem-engine" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f6e455ceb6e5aee7ed7d5c8944104e66992173e03a9c42f9670226318672249" +dependencies = [ + "memmap2", + "polars-arrow", + "polars-core", + "polars-error", + "polars-expr", + "polars-io", + "polars-ops", + "polars-plan", + "polars-utils", + "rayon", + "recursive", +] + [[package]] name = "polars-ops" version = "0.52.0" @@ -1052,9 +1710,11 @@ version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "93c2439d127c59e6bfc9d698419bdb45210068a6f501d44e6096429ad72c2eaa" dependencies = [ + "async-stream", "base64", "bytemuck", "ethnum", + "futures", "hashbrown 0.16.1", "num-traits", "polars-arrow", @@ -1072,6 +1732,10 @@ name = "polars-parquet-format" version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c025243dcfe8dbc57e94d9f82eb3bef10b565ab180d5b99bed87fd8aea319ce1" +dependencies = [ + "async-trait", + "futures", +] [[package]] name = "polars-plan" @@ -1134,6 +1798,46 @@ dependencies = [ "version_check", ] +[[package]] +name = "polars-stream" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ff19612074640a9d65e5928b7223db76ffee63e55b276f1e466d06719eb7362" +dependencies = [ + "async-channel", + "async-trait", + "atomic-waker", + "bitflags", + "chrono-tz", + "crossbeam-channel", + "crossbeam-deque", + "crossbeam-queue", + "crossbeam-utils", + "futures", + "memmap2", + "parking_lot", + "percent-encoding", + "pin-project-lite", + "polars-arrow", + "polars-compute", + "polars-core", + "polars-error", + "polars-expr", + "polars-io", + "polars-mem-engine", + "polars-ops", + "polars-parquet", + "polars-plan", + "polars-time", + "polars-utils", + "rand 0.9.5", + "rayon", + "recursive", + "slotmap", + "tokio", + "version_check", +] + [[package]] name = "polars-time" version = "0.52.0" @@ -1177,7 +1881,7 @@ dependencies = [ "num-traits", "polars-error", "pyo3", - "rand", + "rand 0.9.5", "raw-cpuid", "rayon", "regex", @@ -1197,6 +1901,15 @@ version = "1.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d20d5497ef88037a52ff98267d066e7f11fcc5e99bbfbd58a42336193aacec3" +[[package]] +name = "potential_utf" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d83eb9bc6d8e5cf568e7a1101d60ee05e81ed50ea106026f3d18deeb046d7661" +dependencies = [ + "zerovec", +] + [[package]] name = "ppv-lite86" version = "0.2.21" @@ -1313,7 +2026,7 @@ dependencies = [ "pyo3-polars-derive", "serde", "serde-pickle", - "thiserror", + "thiserror 1.0.69", ] [[package]] @@ -1331,6 +2044,72 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "quick-xml" +version = "0.38.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b66c2058c55a409d601666cffe35f04333cf1013010882cec174a7467cd4e21c" +dependencies = [ + "memchr", + "serde", +] + +[[package]] +name = "quinn" +version = "0.11.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c1a41e437b6bbd489372cd4971de128e85c855f56c57f283d20ff016cf7c0a8" +dependencies = [ + "bytes", + "cfg_aliases", + "pin-project-lite", + "quinn-proto", + "quinn-udp", + "rustc-hash", + "rustls", + "socket2", + "thiserror 2.0.20", + "tokio", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-proto" +version = "0.11.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "04759210543be93709136e28212294a659ef5001836ff4eab4d663e4529bba83" +dependencies = [ + "bytes", + "getrandom 0.4.3", + "lru-slab", + "rand 0.10.2", + "rand_pcg", + "ring", + "rustc-hash", + "rustls", + "rustls-pki-types", + "slab", + "thiserror 2.0.20", + "tinyvec", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-udp" +version = "0.5.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35a133f956daabe89a61a685c2649f13d82d5aa4bd5d12d1277e1072a21c0694" +dependencies = [ + "cfg_aliases", + "libc", + "once_cell", + "socket2", + "tracing", + "windows-sys 0.61.2", +] + [[package]] name = "quote" version = "1.0.47" @@ -1359,7 +2138,18 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" dependencies = [ "rand_chacha", - "rand_core", + "rand_core 0.9.5", +] + +[[package]] +name = "rand" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7f5fa3a058cd35567ef9bfa5e75732bee0f9e4c55fa90477bef2dfcdbc4be80" +dependencies = [ + "chacha20", + "getrandom 0.4.3", + "rand_core 0.10.1", ] [[package]] @@ -1369,7 +2159,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.9.5", ] [[package]] @@ -1381,6 +2171,12 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rand_core" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" + [[package]] name = "rand_distr" version = "0.5.1" @@ -1388,7 +2184,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463" dependencies = [ "num-traits", - "rand", + "rand 0.9.5", +] + +[[package]] +name = "rand_pcg" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "caa0f4137e1c0a72f4c651489402276c8e8e1cf081f3b0ba156d2cbeef09e86a" +dependencies = [ + "rand_core 0.10.1", ] [[package]] @@ -1488,24 +2293,86 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" [[package]] -name = "rmp" -version = "0.8.15" +name = "reqwest" +version = "0.12.28" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4ba8be72d372b2c9b35542551678538b562e7cf86c3315773cae48dfbfe7790c" +checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ - "num-traits", -] - -[[package]] -name = "rmp-serde" -version = "1.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72f81bee8c8ef9b577d1681a70ebbc962c232461e397b22c208c43c04b67a155" + "base64", + "bytes", + "futures-core", + "futures-util", + "h2", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-util", + "js-sys", + "log", + "percent-encoding", + "pin-project-lite", + "quinn", + "rustls", + "rustls-native-certs", + "rustls-pki-types", + "serde", + "serde_json", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tokio-rustls", + "tokio-util", + "tower", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "wasm-streams", + "web-sys", +] + +[[package]] +name = "ring" +version = "0.17.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7" +dependencies = [ + "cc", + "cfg-if", + "getrandom 0.2.17", + "libc", + "untrusted", + "windows-sys 0.52.0", +] + +[[package]] +name = "rmp" +version = "0.8.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ba8be72d372b2c9b35542551678538b562e7cf86c3315773cae48dfbfe7790c" +dependencies = [ + "num-traits", +] + +[[package]] +name = "rmp-serde" +version = "1.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f81bee8c8ef9b577d1681a70ebbc962c232461e397b22c208c43c04b67a155" dependencies = [ "rmp", "serde", ] +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + [[package]] name = "rustfft" version = "6.4.1" @@ -1520,6 +2387,66 @@ dependencies = [ "transpose", ] +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.61.2", +] + +[[package]] +name = "rustls" +version = "0.23.43" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0283386ce02abc0151e1761d08802dfe86c173b0b494af5cbc086574e453da06" +dependencies = [ + "once_cell", + "ring", + "rustls-pki-types", + "rustls-webpki", + "subtle", + "zeroize", +] + +[[package]] +name = "rustls-native-certs" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dab5152771c58876a2146916e53e35057e1a4dfa2b9df0f0305b07f611fdea4d" +dependencies = [ + "openssl-probe", + "rustls-pki-types", + "schannel", + "security-framework", +] + +[[package]] +name = "rustls-pki-types" +version = "1.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f4925028c7eb5d1fcdaf196971378ed9d2c1c4efc7dc5d011256f76c99c0a96" +dependencies = [ + "web-time", + "zeroize", +] + +[[package]] +name = "rustls-webpki" +version = "0.103.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2" +dependencies = [ + "ring", + "rustls-pki-types", + "untrusted", +] + [[package]] name = "rustversion" version = "1.0.23" @@ -1532,12 +2459,53 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + +[[package]] +name = "schannel" +version = "0.1.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91c1b7e4904c873ef0710c1f407dde2e6287de2bebc1bbbf7d430bb7cbffd939" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "scopeguard" version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "security-framework" +version = "3.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" +dependencies = [ + "bitflags", + "core-foundation", + "core-foundation-sys", + "libc", + "security-framework-sys", +] + +[[package]] +name = "security-framework-sys" +version = "2.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2691df843ecc5d231c0b14ece2acc3efb62c0a398c7e1d875f3983ce020e3" +dependencies = [ + "core-foundation-sys", + "libc", +] + [[package]] name = "serde" version = "1.0.229" @@ -1605,6 +2573,18 @@ dependencies = [ "stacker", ] +[[package]] +name = "serde_urlencoded" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd" +dependencies = [ + "form_urlencoded", + "itoa", + "ryu", + "serde", +] + [[package]] name = "sha2" version = "0.10.9" @@ -1612,7 +2592,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" dependencies = [ "cfg-if", - "cpufeatures", + "cpufeatures 0.2.17", "digest", ] @@ -1682,6 +2662,22 @@ version = "1.15.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" +[[package]] +name = "socket2" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "stable_deref_trait" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" + [[package]] name = "stacker" version = "0.1.24" @@ -1692,7 +2688,7 @@ dependencies = [ "cfg-if", "libc", "psm", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -1734,6 +2730,12 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "subtle" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" + [[package]] name = "syn" version = "2.0.119" @@ -1756,6 +2758,26 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "sync_wrapper" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263" +dependencies = [ + "futures-core", +] + +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "target-lexicon" version = "0.13.5" @@ -1768,7 +2790,16 @@ version = "1.0.69" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" dependencies = [ - "thiserror-impl", + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec86235f5fcc2a73650310756d2ac5b138a5780bbbdfae3eeccec992c435ba4f" +dependencies = [ + "thiserror-impl 2.0.20", ] [[package]] @@ -1782,6 +2813,168 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "thiserror-impl" +version = "2.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc04cd3e1236dd4a98afca4569f2deb3f120e5422a4023be2cb683f8486292af" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "tinystr" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1e27c91459209c2986af3dcf603a5a74a4368754ce37414f59acc971167f643" +dependencies = [ + "displaydoc", + "zerovec", +] + +[[package]] +name = "tinyvec" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb4ebadaa0af04fab11ae01eb5f9fdb5f9c5b875506e210e71c07873528baa7f" +dependencies = [ + "tinyvec_macros", +] + +[[package]] +name = "tinyvec_macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" + +[[package]] +name = "tokio" +version = "1.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed" +dependencies = [ + "bytes", + "libc", + "mio", + "pin-project-lite", + "socket2", + "tokio-macros", + "windows-sys 0.61.2", +] + +[[package]] +name = "tokio-macros" +version = "2.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78773a2a397f451582ce068015985c33193cf6dea8b74d2a639fe457b2f07b0e" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "tokio-rustls" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" +dependencies = [ + "rustls", + "tokio", +] + +[[package]] +name = "tokio-util" +version = "0.7.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "494815d09bf52b5548659851081238f0ca39ff638363907596da739561c62c52" +dependencies = [ + "bytes", + "futures-core", + "futures-sink", + "libc", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "tower" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" +dependencies = [ + "futures-core", + "futures-util", + "pin-project-lite", + "sync_wrapper", + "tokio", + "tower-layer", + "tower-service", +] + +[[package]] +name = "tower-http" +version = "0.6.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" +dependencies = [ + "bitflags", + "bytes", + "futures-util", + "http", + "http-body", + "pin-project-lite", + "tower", + "tower-layer", + "tower-service", + "url", +] + +[[package]] +name = "tower-layer" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "121c2a6cda46980bb0fcd1647ffaf6cd3fc79a013de288782836f6df9c48780e" + +[[package]] +name = "tower-service" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" + +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", +] + [[package]] name = "transpose" version = "0.2.3" @@ -1792,6 +2985,12 @@ dependencies = [ "strength_reduce", ] +[[package]] +name = "try-lock" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" + [[package]] name = "typenum" version = "1.20.1" @@ -1810,12 +3009,36 @@ version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" +[[package]] +name = "untrusted" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" + [[package]] name = "unty" version = "0.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6d49784317cd0d1ee7ec5c716dd598ec5b4483ea832a2dced265471cc0f690ae" +[[package]] +name = "url" +version = "2.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff67a8a4397373c3ef660812acab3268222035010ab8680ec4215f38ba3d0eed" +dependencies = [ + "form_urlencoded", + "idna", + "percent-encoding", + "serde", +] + +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + [[package]] name = "uuid" version = "1.24.0" @@ -1840,6 +3063,25 @@ version = "0.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "051eb1abcf10076295e815102942cc58f9d5e3b4560e46e53c21e8ff6f3af7b1" +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + +[[package]] +name = "want" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bfa7760aed19e106de2c7c0b581b509f2f25d3dacaf737cb82ac61bc6d760b0e" +dependencies = [ + "try-lock", +] + [[package]] name = "wasi" version = "0.11.1+wasi-snapshot-preview1" @@ -1868,6 +3110,16 @@ dependencies = [ "wasm-bindgen-shared", ] +[[package]] +name = "wasm-bindgen-futures" +version = "0.4.76" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c62df1340f32221cb9c54d6a27b030e3dba64361d4a95bed55f9aacb44da291d" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + [[package]] name = "wasm-bindgen-macro" version = "0.2.126" @@ -1900,6 +3152,48 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "wasm-streams" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + +[[package]] +name = "web-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8622dcb61c0bcc9fffa6938bed81210af2da9a7e4a1a834b2e37a59b6dfb6141" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "windows-core" version = "0.62.2" @@ -1959,6 +3253,24 @@ dependencies = [ "windows-link", ] +[[package]] +name = "windows-sys" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + [[package]] name = "windows-sys" version = "0.61.2" @@ -1968,18 +3280,111 @@ dependencies = [ "windows-link", ] +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + [[package]] name = "wit-bindgen" version = "0.57.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" +[[package]] +name = "writeable" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ad82d2a33cdc9674dc7465672f271e096168fcdbe0f799d9e6db8c5892679dc" + [[package]] name = "xxhash-rust" version = "0.8.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "aee1b19627c7c60102ab80d3a9cbe18de90bfe03bfa6c3715447681f0e8c8af6" +[[package]] +name = "yoke" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "709fe23a0424b6a435d82152b1bd3fdfb0833487d5fa90d05d42762a9891fef5" +dependencies = [ + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", + "synstructure", +] + [[package]] name = "zerocopy" version = "0.8.55" @@ -2000,6 +3405,66 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "zerofrom" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", + "synstructure", +] + +[[package]] +name = "zeroize" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" + +[[package]] +name = "zerotrie" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ea269c3bd32f0a32c321907a2ae912ba6f4649bb0fc764a15627e99a7095a3f" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", +] + +[[package]] +name = "zerovec" +version = "0.11.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb0464e17806c1d976d5cba29399c7f08e516e279e2ba493f63123b5fca67dd8" +dependencies = [ + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34df6fc39dbd26ddc9c10e6a2984476e13acce22e64e4487636ef494369225da" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + [[package]] name = "zmij" version = "1.0.23" diff --git a/Cargo.toml b/Cargo.toml index 3306fa6..e84f9e4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "polars-list-utils" -version = "1.0.0" +version = "1.0.1" edition = "2024" [lib] @@ -10,7 +10,7 @@ crate-type = ["cdylib"] [dependencies] pyo3 = { version = "0.26", features = ["extension-module", "abi3-py310"] } pyo3-polars = { version = "0.25", features = ["derive"] } -polars = { version = "0.52", default-features = false } +polars = { version = "0.52", default-features = false, features = ["dtype-array"] } serde = { version = "1", features = ["derive"] } realfft = "3.4.0" butterworth = "0.1.0" diff --git a/tests/test_plugins.py b/tests/test_plugins.py index 3e20aaa..cb27d5a 100644 --- a/tests/test_plugins.py +++ b/tests/test_plugins.py @@ -418,6 +418,21 @@ def test_integer_lists_are_accepted(): assert df["yp"][0].to_list() == pytest.approx([15.0]) +def test_array_input_raises_instead_of_aborting(): + # Without the polars `dtype-array` feature, an Array input panicked in + # polars' arrow-FFI import before any plugin check ran -- a non-unwinding + # panic that aborted the whole process. It must be a catchable typed + # error (until Array support proper lands). + df = pl.DataFrame( + {"v": [[1.0, 2.0], [3.0, 4.0]]}, + schema={"v": pl.Array(pl.Float64, 2)}, + ) + with pytest.raises(polars.exceptions.PolarsError, match="expected List column"): + df.with_columns(polist.agg_slices("v", "v", aggregation="mean").alias("a")) + with pytest.raises(polars.exceptions.PolarsError, match="expected List column"): + df.with_columns(polist.apply_fft("v", sample_rate=FS).alias("a")) + + def test_broadcasting_literal_over_rows(): df = pl.DataFrame({"v": [[1.0, 2.0], [3.0, 4.0]]}).with_columns( polist.agg_slices("v", pl.lit([0.0, 1.0]), aggregation="mean").alias("agg") From 4f28b99b37a3e41a869fbd8ca23697ced52756bb Mon Sep 17 00:00:00 2001 From: Travis Hammond Date: Tue, 1 Sep 2026 10:34:32 +0200 Subject: [PATCH 02/17] feat: add sum aggregation, zip_binary, cum_agg_runs, and opt-in strict checks Implements PR 1 of the spec (List columns only): - "sum" joins the aggregation vocabulary of agg_slices and agg_lists, with polars' identity element for an empty selection (+0.0). The Aggregation enum moves to statistic.rs, shared with the new scan. - zip_binary applies a binary op element-wise between two list columns, deferring per element to the corresponding scalar polars op: null propagation and IEEE arithmetic, Kleene logic for and/or, and polars' total order for comparisons (NaN equals NaN, exceeds everything). - cum_agg_runs cumulatively aggregates within gate-delimited runs. Each emitted position equals the vertical aggregation over the run prefix, so an all-True gate reproduces cum_sum/cum_min/cum_max/cum_count, including cum_count's emission at null value positions. - strict (default off) adds the checks whose violations have a legitimate degenerate reading: apply_interp monotonicity, agg_slices range validity, and agg_lists overlength rows. The last rides on a pass-through plugin because pure-polars expressions cannot raise. Both new plugins ship their final dtype contract: Float32 preserved, count as UInt32, integer inners rejected. They also treat empty lists as valid (empty in, empty out), being total over zero elements. Tests add a scan oracle driven by polars itself, deference anchors against native cum_*, strict positive/negative pairs, and a plan-purity test asserting the streaming engine keeps every plugin node in-stream. Co-Authored-By: Claude Fable 5 --- Cargo.lock | 2 +- Cargo.toml | 2 +- README.md | 39 ++++- polars_list_utils/__init__.py | 160 ++++++++++++++++-- src/agg_lists.rs | 38 +++++ src/agg_slices.rs | 37 +---- src/apply_butterworth.rs | 4 +- src/apply_fft.rs | 2 +- src/apply_interp.rs | 35 +++- src/cum_agg_runs.rs | 240 +++++++++++++++++++++++++++ src/lib.rs | 3 + src/statistic.rs | 49 +++++- src/util.rs | 204 ++++++++++++++++++++--- src/zip_binary.rs | 243 +++++++++++++++++++++++++++ tests/test_cum_agg_runs.py | 298 ++++++++++++++++++++++++++++++++++ tests/test_plan_purity.py | 77 +++++++++ tests/test_plugins.py | 47 +++++- tests/test_strict.py | 116 +++++++++++++ tests/test_zip_binary.py | 177 ++++++++++++++++++++ 19 files changed, 1687 insertions(+), 86 deletions(-) create mode 100644 src/agg_lists.rs create mode 100644 src/cum_agg_runs.rs create mode 100644 src/zip_binary.rs create mode 100644 tests/test_cum_agg_runs.py create mode 100644 tests/test_plan_purity.py create mode 100644 tests/test_strict.py create mode 100644 tests/test_zip_binary.py diff --git a/Cargo.lock b/Cargo.lock index 41354a4..ea4b3e7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1646,7 +1646,7 @@ dependencies = [ [[package]] name = "polars-list-utils" -version = "1.0.1" +version = "1.1.0" dependencies = [ "butterworth", "interp", diff --git a/Cargo.toml b/Cargo.toml index e84f9e4..eb8d5cb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "polars-list-utils" -version = "1.0.1" +version = "1.1.0" edition = "2024" [lib] diff --git a/README.md b/README.md index 8a6ba9f..4799da7 100644 --- a/README.md +++ b/README.md @@ -50,26 +50,53 @@ boost. - This is implemented in pure Polars (no plugin), and lives here so that both aggregation functions offer the same aggregations with the same missing-data behaviour. -The four plugin functions accept a length-1 literal list column (e.g. `pl.lit([...])`) for +- `polist.zip_binary` + - Applies a binary operation element-wise between two List-type columns, per element + exactly as the corresponding scalar Polars operation: arithmetic (`add`, `sub`, `mul`, + `div`) with null propagation and IEEE float behaviour (output dtype follows Polars + supertyping, so two `Float32` inputs stay `Float32`), `and`/`or` on Boolean lists with + Kleene logic, and comparisons (`gt`, `ge`, `lt`, `le`, `eq`, `ne`) following Polars' + total order (NaN equals NaN and exceeds everything else), emitting Boolean lists. + +- `polist.cum_agg_runs` + - Cumulatively aggregates a List-type column within runs gated by a Boolean List-type + column: a run is a maximal region of constant gate value (null gates break runs), only + `True`-gated elements accumulate, and every other position emits null (or a plain zero + with `outside="zero"`). + - Each emitted element equals the vertical aggregation over the run's elements so far, + so an all-`True` gate reproduces `cum_sum`/`cum_min`/`cum_max`/`cum_count`; null values + emit null but keep the running state, except `count`, which emits the running count as + `cum_count` does. + - `Float32` values stay `Float32`; `count` emits `UInt32` (Polars' count dtype). + +The six plugin functions accept a length-1 literal list column (e.g. `pl.lit([...])`) for any input and broadcast it across rows. Invalid configuration (bad cutoffs, unknown window names, mismatched list lengths) raises an error instead of silently returning nulls, while per-row data problems (e.g. a signal length that is not a power of two) yield null rows. -Both aggregation functions support `mean`, `median`, `std`, `min`, `max`, `delta`, and -`count`, and both handle missing data exactly like the vertical aggregations of Polars +Both aggregation functions support `sum`, `mean`, `median`, `std`, `min`, `max`, `delta`, +and `count`, and both handle missing data exactly like the vertical aggregations of Polars itself: - Nulls are skipped: excluded from both the numerator and the denominator, so `count` - only counts non-null values and an empty or all-null selection yields null (`count`: 0). + only counts non-null values and an empty or all-null selection yields null (`count`: 0, + `sum`: 0.0 — Polars' identity for sum). - NaNs propagate: Polars treats NaN as a legitimate float value, so it flows into the - sum and poisons `mean` and `std`. Two exceptions to be aware of are `min` and `max`, + sum and poisons `sum`, `mean` and `std`. Two exceptions to be aware of are `min` and `max`, which skip NaNs (unless all values are NaN), and `median` sorts NaN as the largest value. - `std` is the sample standard deviation (ddof=1) and yields null for fewer than two values, as in Polars. The three transforms instead turn any row with missing data (a null row, an empty list, or a list containing nulls) into a null output row: a signal with missing samples cannot be -meaningfully transformed. +meaningfully transformed. `zip_binary` and `cum_agg_runs` mirror the corresponding scalar +and cumulative Polars operations per element instead, and treat empty lists as valid +(empty in, empty out). + +Some checks are opt-in via `strict=True` (default off): `apply_interp` verifies that `x` +is non-decreasing, `agg_slices` rejects inverted or NaN range bounds at expression +construction, and `agg_lists` raises on lists longer than `list_length` instead of +silently truncating them. ### Example diff --git a/polars_list_utils/__init__.py b/polars_list_utils/__init__.py index f8c0c2c..7fd93dd 100644 --- a/polars_list_utils/__init__.py +++ b/polars_list_utils/__init__.py @@ -1,6 +1,6 @@ """Polars expression plugins for signal processing on List columns. -Four Rust plugins operating elementwise on `List[f64]` columns, plus +Six Rust plugins operating elementwise on `List[f64]` columns, plus one pure-polars helper: - :func:`apply_interp`: interpolate (x, y) data onto new x coordinates. @@ -8,14 +8,25 @@ - :func:`apply_fft`: one-sided FFT with standard windowing and scaling. - :func:`agg_slices`: aggregate values selected by index-column ranges. - :func:`agg_lists`: aggregate list columns elementwise over rows. +- :func:`zip_binary`: element-wise binary ops between two list columns. +- :func:`cum_agg_runs`: cumulative aggregation within gated runs. The plugins accept length-1 literal list columns (e.g. `pl.lit(...)`) -for any input and broadcast them. Null rows and empty lists produce null -output rows. The three transforms also null rows whose lists contain -null elements (a signal with missing samples cannot be transformed); -the two aggregations instead skip missing data polars-style. +for any input and broadcast them. Null rows produce null output rows. + +Missing data follows three deliberate family regimes: + +- Transforms (`apply_interp`, `apply_butterworth`, `apply_fft`) null + out rows whose lists are empty or contain null elements (a signal + with missing samples cannot be transformed). +- Aggregations (`agg_slices`, `agg_lists`) skip missing data + polars-style; NaN is a legitimate float and flows through. +- `zip_binary` and `cum_agg_runs` mirror polars' scalar and `cum_*` + operations per element: nulls propagate per those ops' rules, and + empty lists are valid (empty in, empty out). """ +import math from collections.abc import Callable from pathlib import Path from typing import Literal @@ -28,7 +39,10 @@ IntoExprColumn = str | pl.Expr | pl.Series Window = Literal["hann", "hanning", "blackman"] Scaling = Literal["amplitude", "power", "psd"] -Aggregation = Literal["mean", "median", "std", "min", "max", "delta", "count"] +Aggregation = Literal["sum", "mean", "median", "std", "min", "max", "delta", "count"] +BinaryOp = Literal[ + "add", "sub", "mul", "div", "and", "or", "gt", "ge", "lt", "le", "eq", "ne" +] RangeBound = tuple[float, Literal["closed", "open"]] Range = tuple[float, float] | tuple[RangeBound, RangeBound] @@ -39,6 +53,7 @@ __all__ = [ "Aggregation", + "BinaryOp", "IntoExprColumn", "Range", "RangeBound", @@ -49,7 +64,9 @@ "apply_butterworth", "apply_fft", "apply_interp", + "cum_agg_runs", "polars_func_arg_into_col_expr", + "zip_binary", ] @@ -80,6 +97,8 @@ def apply_interp( x_column: IntoExprColumn, y_column: IntoExprColumn, xp_column: IntoExprColumn, + *, + strict: bool = False, ) -> pl.Expr: """Interpolate `(x, y)` data onto the `xp` coordinates. @@ -87,12 +106,18 @@ def apply_interp( values outside the data range clamped to the first/last `y` value. `x` must be sorted in increasing order. + With `strict=True` a row whose `x` values are not non-decreasing + raises (duplicates stay legal, as in numpy; a NaN in `x` counts as + a violation). The default keeps numpy's silent behaviour: unsorted + `x` interpolates garbage without complaint. + Returns a `List[f64]` column of interpolated y values, one per `xp` coordinate. Raises if x and y lists differ in length. """ return _plugin( "apply_interp", [x_column, y_column, xp_column], + strict=strict, ) @@ -158,6 +183,19 @@ def apply_fft( ) +def _validate_ranges(ranges: list[Range] | None, param: str) -> None: + """Raise on inverted or NaN range bounds (agg_slices strict=True).""" + for r in ranges or []: + lo, hi = r + lo_value = float(lo[0]) if isinstance(lo, tuple) else float(lo) + hi_value = float(hi[0]) if isinstance(hi, tuple) else float(hi) + if math.isnan(lo_value) or math.isnan(hi_value) or lo_value > hi_value: + raise ValueError( + f"agg_slices: invalid range {r!r} in {param} (strict=True): " + "bounds must be non-NaN with lo <= hi" + ) + + def agg_slices( value_column: IntoExprColumn, index_column: IntoExprColumn, @@ -165,6 +203,7 @@ def agg_slices( aggregation: Aggregation, slices_include: list[Range] | None = None, slices_exclude: list[Range] | None = None, + strict: bool = False, ) -> pl.Expr: """Aggregate values whose paired index falls within the given ranges. @@ -176,14 +215,22 @@ def agg_slices( Missing data follows the polars convention, so results match polars' vertical aggregations over the same values: null elements are skipped (pairwise with their index), while NaN is a legitimate - float value — it poisons `mean`/`std`, is skipped by + float value — it poisons `sum`/`mean`/`std`, is skipped by `min`/`max` unless all values are NaN, and sorts as the largest - value for `median`. An empty selection yields null (`count`: 0). - NaN indices never match any range; mismatched list lengths raise. + value for `median`. An empty selection yields null (`count`: 0, + `sum`: 0.0). NaN indices never match any range; mismatched list + lengths raise. `std` is the sample standard deviation (ddof=1, polars' default) and yields null for selections with fewer than two values. + + With `strict=True`, inverted (`lo > hi`) or NaN range bounds raise + at expression construction; the default keeps them as legitimately + empty selections. """ + if strict: + _validate_ranges(slices_include, "slices_include") + _validate_ranges(slices_exclude, "slices_exclude") return _plugin( "agg_slices", [value_column, index_column], @@ -194,6 +241,7 @@ def agg_slices( _AGGS: dict[str, Callable[[pl.Expr], pl.Expr]] = { + "sum": pl.Expr.sum, "mean": pl.Expr.mean, "median": pl.Expr.median, "std": pl.Expr.std, @@ -209,6 +257,7 @@ def agg_lists( *, list_length: int, aggregation: Aggregation, + strict: bool = False, ) -> pl.Expr: """Aggregate a list-type column elementwise over rows. @@ -218,19 +267,98 @@ def agg_lists( group. Lists shorter than `list_length` contribute nulls at the missing positions. + Lists longer than `list_length` are silently truncated to it — a + deliberate window for some callers. With `strict=True` such rows + raise instead. + Implemented in pure polars (no plugin), so missing data follows the polars convention exactly as in :func:`agg_slices`: null elements - are skipped (`count` counts non-null values), NaN propagates - per-kernel, and `std` is the sample standard deviation (ddof=1). + are skipped (`count` counts non-null values, an all-null position + sums to 0.0), NaN propagates per-kernel, and `std` is the sample + standard deviation (ddof=1). """ if list_length < 1: raise ValueError(f"list_length must be at least 1, got {list_length}") + col = polars_func_arg_into_col_expr(list_column) + if strict: + # Pure polars cannot raise from inside an expression, so the + # length check rides along as an elementwise pass-through plugin. + col = _plugin("check_list_len", [col], list_length=list_length) agg = _AGGS[aggregation] return pl.concat_list( - agg( - polars_func_arg_into_col_expr(list_column) - .list.slice(offset=n, length=1) - .list.first() - ) + agg(col.list.slice(offset=n, length=1).list.first()) for n in range(list_length) ) + + +def zip_binary( + left_column: IntoExprColumn, + right_column: IntoExprColumn, + *, + op: BinaryOp, +) -> pl.Expr: + """Apply a binary operation element-wise between two list columns. + + Per element the result is identical to the corresponding scalar + polars op: + + - Arithmetic (`add`, `sub`, `mul`, `div`) takes float lists; null + propagates, NaN follows IEEE float arithmetic (`1/0` is inf, + `0/0` is NaN). The output inner dtype follows polars supertyping: + `Float32` when both inputs are `Float32`, else `Float64`. + - `and`/`or` take Boolean lists and follow Kleene logic + (`False & null = False`, `True | null = True`). + - Comparisons (`gt`, `ge`, `lt`, `le`, `eq`, `ne`) take float lists + and emit Boolean, following polars' total order: NaN equals NaN + and exceeds everything else; comparing against null yields null. + + Integer inner dtypes raise. A null row on either side yields a null + row; empty lists zip to empty lists; mismatched list lengths raise. + Length-1 literal lists broadcast (zip against a constant template). + """ + return _plugin( + "zip_binary", + [left_column, right_column], + op=op, + ) + + +def cum_agg_runs( + value_column: IntoExprColumn, + gate_column: IntoExprColumn, + *, + aggregation: Aggregation, + outside: Literal["null", "zero"] = "null", +) -> pl.Expr: + """Cumulatively aggregate list values within gated runs. + + A run is a maximal region of constant `gate_column` value (`True`, + `False`, and null are distinct, so a null gate element breaks + runs). Only `True`-gated elements accumulate; every other position + emits the `outside` fill (`"null"`, or `"zero"` for a plain 0 in + the output dtype). Within a run, each emitted element equals the + vertical `aggregation` over the run's elements so far — so with an + all-`True` gate, `sum`/`min`/`max`/`count` match polars' + `cum_sum`/`cum_min`/`cum_max`/`cum_count`. + + Missing data follows the polars `cum_*` convention: a null value + leaves the running state unchanged and emits null — except `count`, + which emits the unchanged running count, as `cum_count` does. NaN + follows the vertical convention from its entry onward within the + run: it poisons `sum`/`mean`/`std`, is skipped by + `min`/`max`/`delta`, sorts largest for `median`, and is counted by + `count`. `std` (ddof=1) emits null until a run prefix holds two + non-null values. + + `value_column` takes float lists (`Float32` is preserved in the + output; integers raise) and `gate_column` Boolean lists. `count` + emits `UInt32`, polars' count dtype. A null row in either input + yields a null row; empty lists produce empty lists; mismatched + value/gate lengths raise. Length-1 literal lists broadcast. + """ + return _plugin( + "cum_agg_runs", + [value_column, gate_column], + aggregation=aggregation, + outside=outside, + ) diff --git a/src/agg_lists.rs b/src/agg_lists.rs new file mode 100644 index 0000000..d7316a8 --- /dev/null +++ b/src/agg_lists.rs @@ -0,0 +1,38 @@ +//! `agg_lists` itself is pure polars (see the Python wrapper); this +//! module only holds the plugin backing its `strict` check, because a +//! pure-polars expression has no way to raise at run time. + +use polars::prelude::*; +use pyo3_polars::derive::polars_expr; +use serde::Deserialize; + +#[derive(Deserialize)] +struct CheckListLenKwargs { + list_length: usize, +} + +fn passthrough_output(input_fields: &[Field]) -> PolarsResult { + Ok(input_fields[0].clone()) +} + +/// Identity pass-through that raises when any row's list is longer +/// than `list_length` (which `agg_lists` would otherwise silently +/// truncate to). Null rows are data and pass through. +#[polars_expr(output_type_func=passthrough_output)] +fn check_list_len( + inputs: &[Series], + kwargs: CheckListLenKwargs, +) -> PolarsResult { + let s = &inputs[0]; + for (i, row) in s.list()?.into_iter().enumerate() { + if let Some(row) = row { + polars_ensure!( + row.len() <= kwargs.list_length, + ComputeError: + "agg_lists: row {i} has list length {} but list_length={} (strict=true)", + row.len(), kwargs.list_length + ); + } + } + Ok(s.clone()) +} diff --git a/src/agg_slices.rs b/src/agg_slices.rs index 05f3ae7..e6f58c5 100644 --- a/src/agg_slices.rs +++ b/src/agg_slices.rs @@ -1,39 +1,10 @@ use crate::slice::{Range, Slice}; -use crate::statistic::Statistic; +use crate::statistic::Aggregation; use crate::util::ListInputs; use polars::prelude::*; use pyo3_polars::derive::polars_expr; use serde::Deserialize; -#[derive(Deserialize)] -#[serde(rename_all = "lowercase")] -enum Aggregation { - Mean, - Median, - Min, - Max, - Std, - Delta, - Count, -} - -impl Aggregation { - fn apply( - &self, - values: &[f64], - ) -> Option { - match self { - Self::Mean => values.mean(), - Self::Median => values.median(), - Self::Min => values.min(), - Self::Max => values.max(), - Self::Std => values.std(), - Self::Delta => values.delta(), - Self::Count => Some(values.len() as f64), - } - } -} - #[derive(Deserialize)] struct AggSlicesKwargs { aggregation: Aggregation, @@ -47,7 +18,8 @@ struct AggSlicesKwargs { /// Missing data follows the polars convention: null elements are /// skipped (pairwise with their index), while NaN is a legitimate /// float value that flows into the aggregation. An empty selection -/// yields null (`count`: 0). NaN indices never match any range. +/// yields null (`count`: 0, `sum`: 0.0). NaN indices never match any +/// range. #[polars_expr(output_type=Float64)] fn agg_slices( inputs: &[Series], @@ -76,7 +48,8 @@ fn agg_slices( // Polars missing-data convention: null elements are skipped // (pairwise with their index), NaN values flow through. The - // aggregations turn an empty selection into None (count: 0). + // aggregations turn an empty selection into None (count: 0, + // sum: 0.0). let selected = values .iter() .zip(indices) diff --git a/src/apply_butterworth.rs b/src/apply_butterworth.rs index bfafa2d..e1849b3 100644 --- a/src/apply_butterworth.rs +++ b/src/apply_butterworth.rs @@ -37,14 +37,14 @@ fn apply_butterworth( (Some(lo), None) => Filter::Highpass(lo), (None, Some(hi)) => Filter::Lowpass(hi), (None, None) => { - return apply_list_transform(inputs, |cols| Ok(Some(cols[0].to_vec()))); + return apply_list_transform(inputs, |_row, cols| Ok(Some(cols[0].to_vec()))); } }; let order = kwargs.filter_order.unwrap_or(4); let built = filter.build(kwargs.sample_rate, order)?; let min_samples = filter.min_samples(order); - apply_list_transform(inputs, move |cols| { + apply_list_transform(inputs, move |_row, cols| { if cols[0].len() < min_samples { return Ok(None); } diff --git a/src/apply_fft.rs b/src/apply_fft.rs index 05d778a..18d1175 100644 --- a/src/apply_fft.rs +++ b/src/apply_fft.rs @@ -33,7 +33,7 @@ fn apply_fft( let window = kwargs.window.unwrap_or_default(); let scaling = kwargs.scaling.unwrap_or_default(); - apply_list_transform(inputs, |cols| { + apply_list_transform(inputs, |_row, cols| { Ok(compute_fft(cols[0], kwargs.sample_rate, window, scaling)) }) } diff --git a/src/apply_interp.rs b/src/apply_interp.rs index 5df18b0..e3ad052 100644 --- a/src/apply_interp.rs +++ b/src/apply_interp.rs @@ -2,18 +2,49 @@ use crate::util::{apply_list_transform, list_f64_output}; use interp::{InterpMode, interp_slice}; use polars::prelude::*; use pyo3_polars::derive::polars_expr; +use serde::Deserialize; + +#[derive(Deserialize)] +struct InterpKwargs { + #[serde(default)] + strict: bool, +} /// Interpolate `y` values from `(x, y)` data onto new `xp` coordinates, /// as in numpy.interp; `xp` values outside the data range clamp to the /// first/last `y` value. +/// +/// `x` must be non-decreasing for the result to be meaningful; with +/// `strict` the first violation raises (a NaN in `x` counts as one), +/// without it unsorted `x` silently interpolates garbage, as in numpy. #[polars_expr(output_type_func=list_f64_output)] -fn apply_interp(inputs: &[Series]) -> PolarsResult { - apply_list_transform(inputs, |cols| { +fn apply_interp( + inputs: &[Series], + kwargs: InterpKwargs, +) -> PolarsResult { + apply_list_transform(inputs, |row_idx, cols| { polars_ensure!( cols[0].len() == cols[1].len(), ComputeError: "apply_interp: x and y lists differ in length ({} vs {})", cols[0].len(), cols[1].len() ); + // Non-decreasing means each pair compares Less or Equal; a NaN + // pair is incomparable and counts as a violation. + let in_order = |w: &&[f64]| { + matches!( + w[0].partial_cmp(&w[1]), + Some(std::cmp::Ordering::Less | std::cmp::Ordering::Equal) + ) + }; + if kwargs.strict + && let Some(w) = cols[0].windows(2).find(|w| !in_order(w)) + { + polars_bail!(ComputeError: + "apply_interp: x_column must be non-decreasing with strict=true, \ + found {} followed by {} at row {row_idx}", + w[0], w[1] + ); + } Ok(Some(interp_slice( cols[0], cols[1], diff --git a/src/cum_agg_runs.rs b/src/cum_agg_runs.rs new file mode 100644 index 0000000..82f305a --- /dev/null +++ b/src/cum_agg_runs.rs @@ -0,0 +1,240 @@ +use crate::statistic::Aggregation; +use crate::util::{ + TypedListInput, broadcast_len, build_float_list, build_u32_list, tot_cmp, +}; +use polars::prelude::*; +use pyo3_polars::derive::polars_expr; +use serde::Deserialize; +use std::cmp::Ordering; + +/// What non-`True`-gated positions emit. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "lowercase")] +enum Outside { + Null, + Zero, +} + +#[derive(Deserialize)] +struct CumAggRunsKwargs { + aggregation: Aggregation, + outside: Outside, +} + +/// Running state over one gated run's non-null values. +/// +/// Every emitted position must equal the vertical aggregation over the +/// run-prefix's non-null values (the prefix rule): `sum`/`mean`/`std` +/// let NaN poison the accumulators, `min`/`max` skip NaN via the +/// NaN-propagating `f64::min`/`max` seeds (all-NaN prefixes read NaN), +/// `median` keeps a sorted prefix with NaN largest, `count` counts +/// every non-null value (NaN included). +struct RunState { + n: u64, + sum: f64, + min: f64, + max: f64, + mean: f64, + m2: f64, + sorted: Vec, +} + +impl RunState { + fn new() -> Self { + Self { + n: 0, + sum: 0.0, + min: f64::NAN, + max: f64::NAN, + mean: 0.0, + m2: 0.0, + sorted: Vec::new(), + } + } + + fn update( + &mut self, + agg: Aggregation, + x: f64, + ) { + self.n += 1; + self.sum += x; + // f64::min/max return the non-NaN operand, so these skip NaN + // and read NaN only while every value so far is NaN. + self.min = self.min.min(x); + self.max = self.max.max(x); + // Welford in f64; a NaN poisons mean and m2 for the run's rest. + let delta = x - self.mean; + self.mean += delta / self.n as f64; + self.m2 += delta * (x - self.mean); + if agg == Aggregation::Median { + let pos = self + .sorted + .partition_point(|&v| tot_cmp(v, x) == Ordering::Less); + self.sorted.insert(pos, x); + } + } + + /// The aggregation over the prefix; only called with `n >= 1`. + fn emit( + &self, + agg: Aggregation, + ) -> Option { + match agg { + Aggregation::Sum => Some(self.sum), + Aggregation::Count => Some(self.n as f64), + Aggregation::Mean => Some(self.sum / self.n as f64), + Aggregation::Min => Some(self.min), + Aggregation::Max => Some(self.max), + Aggregation::Delta => Some(self.max - self.min), + Aggregation::Std => { + // Sample std (ddof=1): null below two values, as in polars. + (self.n >= 2).then(|| (self.m2 / (self.n - 1) as f64).sqrt()) + } + Aggregation::Median => { + let mid = self.sorted.len() / 2; + Some(if self.sorted.len().is_multiple_of(2) { + (self.sorted[mid - 1] + self.sorted[mid]) / 2.0 + } else { + self.sorted[mid] + }) + } + } + } +} + +/// Scan one row: a run is a maximal region of constant gate value +/// (`True`, `False`, and null are distinct, so null gates break runs). +/// Only `True`-gated positions accumulate; the rest emit the `outside` +/// fill. Null values leave the running state unchanged and emit null, +/// except `count`, which emits the unchanged running count (as native +/// `cum_count` does). +fn scan_row( + row_idx: usize, + values: &[Option], + gates: &[Option], + agg: Aggregation, + outside: Outside, +) -> PolarsResult>> { + polars_ensure!( + values.len() == gates.len(), + ComputeError: + "cum_agg_runs: value and gate lists differ in length at row {row_idx} ({} vs {})", + values.len(), gates.len() + ); + + let outside_fill = match outside { + Outside::Null => None, + Outside::Zero => Some(0.0), + }; + + let mut out = Vec::with_capacity(values.len()); + let mut state = RunState::new(); + let mut prev_gate: Option> = None; + for (&value, &gate) in values.iter().zip(gates) { + if prev_gate != Some(gate) { + state = RunState::new(); + } + prev_gate = Some(gate); + + if gate != Some(true) { + out.push(outside_fill); + continue; + } + match value { + None => out.push(match agg { + Aggregation::Count => Some(state.n as f64), + _ => None, + }), + Some(x) => { + state.update(agg, x); + out.push(state.emit(agg)); + } + } + } + Ok(out) +} + +fn cum_agg_runs_output( + input_fields: &[Field], + kwargs: CumAggRunsKwargs, +) -> PolarsResult { + polars_ensure!( + input_fields.len() == 2, + ComputeError: "cum_agg_runs expects exactly 2 inputs, got {}", + input_fields.len() + ); + let value_inner = match input_fields[0].dtype() { + DataType::List(inner) + if matches!(**inner, DataType::Float32 | DataType::Float64) => + { + *inner.clone() + } + dt => polars_bail!(ComputeError: + "cum_agg_runs: value_column must be a List column with Float32 or Float64 \ + inner dtype, got {dt}"), + }; + polars_ensure!( + matches!(input_fields[1].dtype(), DataType::List(inner) if **inner == DataType::Boolean), + ComputeError: "cum_agg_runs: gate_column must be a List column with Boolean inner \ + dtype, got {}", + input_fields[1].dtype() + ); + // count emits UInt32, polars' count dtype; everything else keeps + // the value column's inner dtype. + let inner = match kwargs.aggregation { + Aggregation::Count => DataType::UInt32, + _ => value_inner, + }; + Ok(Field::new( + PlSmallStr::from(""), + DataType::List(Box::new(inner)), + )) +} + +/// Cumulative aggregation within gated runs of a list column. +/// +/// A null row in either input yields a null row; empty lists scan to +/// empty lists; mismatched value/gate lengths raise. +#[polars_expr(output_type_func_with_kwargs=cum_agg_runs_output)] +fn cum_agg_runs( + inputs: &[Series], + kwargs: CumAggRunsKwargs, +) -> PolarsResult { + polars_ensure!( + inputs.len() == 2, + ComputeError: "cum_agg_runs expects exactly 2 inputs, got {}", + inputs.len() + ); + let values = TypedListInput::::float(&inputs[0], "cum_agg_runs: value_column")?; + let gates = TypedListInput::::boolean(&inputs[1], "cum_agg_runs: gate_column")?; + let len = broadcast_len(&[values.n_rows(), gates.n_rows()])?; + + let mut rows: Vec>>> = Vec::with_capacity(len); + for i in 0..len { + rows.push(match (values.row(i), gates.row(i)) { + (Some(v), Some(g)) => { + Some(scan_row(i, v, g, kwargs.aggregation, kwargs.outside)?) + } + _ => None, + }); + } + + if kwargs.aggregation == Aggregation::Count { + // Counts are exact integers well below 2^53, so the f64 detour + // is lossless. + let rows = rows + .into_iter() + .map(|row| { + row.map(|vals| { + vals.into_iter() + .map(|v| v.map(|x| x as u32)) + .collect::>>() + }) + }) + .collect(); + Ok(build_u32_list(rows)) + } else { + build_float_list(rows, values.inner_dtype()) + } +} diff --git a/src/lib.rs b/src/lib.rs index 41832f7..077ed43 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,12 +1,15 @@ +mod agg_lists; mod agg_slices; mod apply_butterworth; mod apply_fft; mod apply_interp; +mod cum_agg_runs; mod filter; mod slice; mod statistic; mod transform; mod util; +mod zip_binary; use {pyo3::prelude::*, pyo3_polars::PolarsAllocator}; diff --git a/src/statistic.rs b/src/statistic.rs index 0a7df2c..531f7b3 100644 --- a/src/statistic.rs +++ b/src/statistic.rs @@ -1,3 +1,39 @@ +use serde::Deserialize; + +/// The aggregation vocabulary shared by `agg_slices` and `cum_agg_runs` +/// (and mirrored in pure polars by `agg_lists`). Every member is valid +/// for every aggregation-taking function. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "lowercase")] +pub(crate) enum Aggregation { + Sum, + Mean, + Median, + Min, + Max, + Std, + Delta, + Count, +} + +impl Aggregation { + pub(crate) fn apply( + &self, + values: &[f64], + ) -> Option { + match self { + Self::Sum => values.sum(), + Self::Mean => values.mean(), + Self::Median => values.median(), + Self::Min => values.min(), + Self::Max => values.max(), + Self::Std => values.std(), + Self::Delta => values.delta(), + Self::Count => Some(values.len() as f64), + } + } +} + /// Statistics over slices of data. pub(crate) trait Statistic { fn values(&self) -> &[f64]; @@ -6,12 +42,23 @@ pub(crate) trait Statistic { self.values().is_empty() } + /// Sum with polars' identity element: an empty selection sums to + /// `+0.0` (polars itself is split between `+0.0` and `-0.0` here, + /// and Rust's `Iterator::sum` starts from IEEE's additive identity + /// `-0.0`; the library standardizes on `+0.0`). + fn sum(&self) -> Option { + if self.is_empty() { + return Some(0.0); + } + Some(self.values().iter().sum()) + } + fn median(&self) -> Option { if self.is_empty() { return None; } let mut values = self.values().to_vec(); - values.sort_by(|a, b| a.total_cmp(b)); + values.sort_by(|a, b| crate::util::tot_cmp(*a, *b)); let mid = values.len() / 2; if values.len().is_multiple_of(2) { Some((values[mid - 1] + values[mid]) / 2.0) diff --git a/src/util.rs b/src/util.rs index 0ff6634..0d46bbb 100644 --- a/src/util.rs +++ b/src/util.rs @@ -1,4 +1,20 @@ use polars::prelude::*; +use std::cmp::Ordering; + +/// Total-order float comparison matching polars: NaN compares equal to +/// NaN and greater than everything else (regardless of NaN sign bit), +/// while `-0.0 == 0.0` as in IEEE — both unlike `f64::total_cmp`. +pub(crate) fn tot_cmp( + a: f64, + b: f64, +) -> Ordering { + match (a.is_nan(), b.is_nan()) { + (true, true) => Ordering::Equal, + (true, false) => Ordering::Greater, + (false, true) => Ordering::Less, + (false, false) => a.partial_cmp(&b).expect("non-NaN floats compare"), + } +} /// Output dtype for list-returning expressions. pub(crate) fn list_f64_output(_: &[Field]) -> PolarsResult { @@ -54,6 +70,19 @@ impl ListInputs> { } } +/// The broadcast length across inputs: length-1 inputs are literals, +/// all others must agree. +pub(crate) fn broadcast_len(lens: &[usize]) -> PolarsResult { + let len = lens.iter().copied().filter(|&l| l != 1).max().unwrap_or(1); + for &l in lens { + polars_ensure!( + l == len || l == 1, + ComputeError: "all inputs must have length {len} or 1 (got {l})" + ); + } + Ok(len) +} + impl ListInputs { fn build( inputs: &[Series], @@ -77,20 +106,8 @@ impl ListInputs { }) .collect::>>()?; - // The broadcast length: length-1 inputs are literals, all others must agree. - let len = cas - .iter() - .map(|s| s.len()) - .filter(|&l| l != 1) - .max() - .unwrap_or(1); - for s in &cas { - polars_ensure!( - s.len() == len || s.len() == 1, - ComputeError: "all inputs must have length {len} or 1 (got {})", - s.len() - ); - } + let lens = cas.iter().map(|s| s.len()).collect::>(); + let len = broadcast_len(&lens)?; let columns = cas .iter() @@ -126,15 +143,15 @@ impl ListInputs { /// Apply a fallible transformation elementwise across one or more list /// columns, returning a `List[f64]` column. /// -/// The closure receives one `&[f64]` per input column (length-1 literals -/// broadcast) and returns `Ok(None)` to emit a null row, or `Err` to -/// abort the whole query. Null input rows stay null. +/// The closure receives the row index and one `&[f64]` per input column +/// (length-1 literals broadcast) and returns `Ok(None)` to emit a null +/// row, or `Err` to abort the whole query. Null input rows stay null. pub(crate) fn apply_list_transform( inputs: &[Series], f: F, ) -> PolarsResult where - F: Fn(&[&[f64]]) -> PolarsResult>>, + F: Fn(usize, &[&[f64]]) -> PolarsResult>>, { let li = ListInputs::new(inputs)?; @@ -146,7 +163,7 @@ where ); for i in 0..li.len() { - match li.row(i).map(|slices| f(&slices)).transpose()?.flatten() { + match li.row(i).map(|slices| f(i, &slices)).transpose()?.flatten() { Some(result) => builder.append_slice(&result), None => builder.append_null(), } @@ -154,3 +171,152 @@ where Ok(builder.finish().into_series()) } + +/// One `List` input for the dtype-preserving entry layer used by +/// `zip_binary` and `cum_agg_runs`. +/// +/// Unlike [`ListInputs`], rows keep inner nulls **and** empty lists +/// (these functions are total over zero elements: empty in, empty out), +/// and the validated original inner dtype is kept so the kernel's `f64` +/// result can be restored to it at the exit boundary. +pub(crate) struct TypedListInput { + rows: Vec>>>, + inner_dtype: DataType, +} + +impl TypedListInput { + /// Extract a float List column; inner dtype must be `Float32` or + /// `Float64` (anything else raises, ints included). `label` names + /// the function and parameter for error messages. + pub(crate) fn float( + s: &Series, + label: &str, + ) -> PolarsResult { + let inner_dtype = list_inner_dtype(s, label)?; + polars_ensure!( + matches!(inner_dtype, DataType::Float32 | DataType::Float64), + ComputeError: + "{label} must be a List column with Float32 or Float64 inner dtype, got {}", + s.dtype() + ); + let cast = s.cast(&DataType::List(Box::new(DataType::Float64)))?; + let mut rows = Vec::with_capacity(cast.len()); + for opt in cast.list()?.into_iter() { + rows.push(match opt { + None => None, + Some(row) => Some(row.f64()?.into_iter().collect()), + }); + } + Ok(Self { rows, inner_dtype }) + } +} + +impl TypedListInput { + /// Extract a Boolean List column; any other inner dtype raises. + pub(crate) fn boolean( + s: &Series, + label: &str, + ) -> PolarsResult { + let inner_dtype = list_inner_dtype(s, label)?; + polars_ensure!( + inner_dtype == DataType::Boolean, + ComputeError: "{label} must be a List column with Boolean inner dtype, got {}", + s.dtype() + ); + let mut rows = Vec::with_capacity(s.len()); + for opt in s.list()?.into_iter() { + rows.push(match opt { + None => None, + Some(row) => Some(row.bool()?.into_iter().collect()), + }); + } + Ok(Self { rows, inner_dtype }) + } +} + +impl TypedListInput { + /// Physical row count, before broadcasting. + pub(crate) fn n_rows(&self) -> usize { + self.rows.len() + } + + pub(crate) fn inner_dtype(&self) -> &DataType { + &self.inner_dtype + } + + /// Row `i` with length-1 literal broadcasting; `None` is a null row. + pub(crate) fn row( + &self, + i: usize, + ) -> Option<&[Option]> { + self.rows[if self.rows.len() == 1 { 0 } else { i }].as_deref() + } +} + +/// The inner dtype of a `List` column, or a clear error. +fn list_inner_dtype( + s: &Series, + label: &str, +) -> PolarsResult { + match s.dtype() { + DataType::List(inner) => Ok(*inner.clone()), + dt => polars_bail!(ComputeError: "{label} must be a List column, got {dt}"), + } +} + +/// Build a `List` float series from per-row optional elements, casting +/// the `f64`-computed values to `inner` once at the exit boundary. +pub(crate) fn build_float_list( + rows: Vec>>>, + inner: &DataType, +) -> PolarsResult { + let mut builder = ListPrimitiveChunkedBuilder::::new( + PlSmallStr::from(""), + rows.len(), + rows.len() * 8, + DataType::Float64, + ); + for row in rows { + match row { + Some(vals) => builder.append_iter(vals.into_iter()), + None => builder.append_null(), + } + } + let s = builder.finish().into_series(); + if *inner == DataType::Float64 { + Ok(s) + } else { + s.cast(&DataType::List(Box::new(inner.clone()))) + } +} + +/// Build a `List(Boolean)` series from per-row optional elements. +pub(crate) fn build_bool_list(rows: Vec>>>) -> Series { + let mut builder = + ListBooleanChunkedBuilder::new(PlSmallStr::from(""), rows.len(), rows.len() * 8); + for row in rows { + match row { + Some(vals) => builder.append_iter(vals.into_iter()), + None => builder.append_null(), + } + } + builder.finish().into_series() +} + +/// Build a `List(UInt32)` series from per-row optional elements +/// (the `count` aggregation's dtype, as in polars). +pub(crate) fn build_u32_list(rows: Vec>>>) -> Series { + let mut builder = ListPrimitiveChunkedBuilder::::new( + PlSmallStr::from(""), + rows.len(), + rows.len() * 8, + DataType::UInt32, + ); + for row in rows { + match row { + Some(vals) => builder.append_iter(vals.into_iter()), + None => builder.append_null(), + } + } + builder.finish().into_series() +} diff --git a/src/zip_binary.rs b/src/zip_binary.rs new file mode 100644 index 0000000..fd63395 --- /dev/null +++ b/src/zip_binary.rs @@ -0,0 +1,243 @@ +use crate::util::{ + TypedListInput, broadcast_len, build_bool_list, build_float_list, tot_cmp, +}; +use polars::prelude::*; +use pyo3_polars::derive::polars_expr; +use serde::Deserialize; +use std::cmp::Ordering; + +/// Element-wise binary operations between two paired list columns. +/// +/// Arithmetic ops take float operands and keep a float output; +/// `and`/`or` take Boolean operands (Kleene logic); the comparisons +/// take float operands and emit Boolean. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "lowercase")] +enum Op { + Add, + Sub, + Mul, + Div, + And, + Or, + Gt, + Ge, + Lt, + Le, + Eq, + Ne, +} + +impl Op { + fn takes_booleans(self) -> bool { + matches!(self, Op::And | Op::Or) + } + + fn is_arithmetic(self) -> bool { + matches!(self, Op::Add | Op::Sub | Op::Mul | Op::Div) + } +} + +#[derive(Deserialize)] +struct ZipBinaryKwargs { + op: Op, +} + +/// Per element, identical to the corresponding scalar polars op: +/// null propagates, NaN follows IEEE float arithmetic. +fn arithmetic( + op: Op, + a: Option, + b: Option, +) -> Option { + let (a, b) = (a?, b?); + Some(match op { + Op::Add => a + b, + Op::Sub => a - b, + Op::Mul => a * b, + Op::Div => a / b, + _ => unreachable!("arithmetic called with non-arithmetic op"), + }) +} + +/// Comparisons follow polars' total order (NaN equals NaN and exceeds +/// everything else); null operands compare to null. +fn comparison( + op: Op, + a: Option, + b: Option, +) -> Option { + let ord = tot_cmp(a?, b?); + Some(match op { + Op::Gt => ord == Ordering::Greater, + Op::Ge => ord != Ordering::Less, + Op::Lt => ord == Ordering::Less, + Op::Le => ord != Ordering::Greater, + Op::Eq => ord == Ordering::Equal, + Op::Ne => ord != Ordering::Equal, + _ => unreachable!("comparison called with non-comparison op"), + }) +} + +/// Kleene three-valued logic, as in polars: `False & null = False`, +/// `True | null = True`, otherwise null wins. +fn kleene( + op: Op, + a: Option, + b: Option, +) -> Option { + match op { + Op::And => match (a, b) { + (Some(false), _) | (_, Some(false)) => Some(false), + (Some(true), Some(true)) => Some(true), + _ => None, + }, + Op::Or => match (a, b) { + (Some(true), _) | (_, Some(true)) => Some(true), + (Some(false), Some(false)) => Some(false), + _ => None, + }, + _ => unreachable!("kleene called with non-boolean op"), + } +} + +/// Zip one row's element pairs through `f`; a null row on either side +/// stays null, mismatched lengths raise, empty lists zip to empty. +fn zip_row( + row_idx: usize, + left: Option<&[Option]>, + right: Option<&[Option]>, + f: impl Fn(Option, Option) -> Option, +) -> PolarsResult>>> { + let (Some(left), Some(right)) = (left, right) else { + return Ok(None); + }; + polars_ensure!( + left.len() == right.len(), + ComputeError: + "zip_binary: left and right lists differ in length at row {row_idx} ({} vs {})", + left.len(), right.len() + ); + Ok(Some( + left.iter().zip(right).map(|(&a, &b)| f(a, b)).collect(), + )) +} + +/// The arithmetic output inner dtype follows polars supertyping: +/// `Float32` only when both operands are `Float32`. +fn arithmetic_inner( + left: &DataType, + right: &DataType, +) -> DataType { + if *left == DataType::Float32 && *right == DataType::Float32 { + DataType::Float32 + } else { + DataType::Float64 + } +} + +/// Plan-time dtype validation for one operand. +fn validate_operand( + dtype: &DataType, + op: Op, + label: &str, +) -> PolarsResult { + let DataType::List(inner) = dtype else { + polars_bail!(ComputeError: "{label} must be a List column, got {dtype}"); + }; + if op.takes_booleans() { + polars_ensure!( + **inner == DataType::Boolean, + ComputeError: "{label} must be a List column with Boolean inner dtype for \ + 'and'/'or', got {dtype}" + ); + } else { + polars_ensure!( + matches!(**inner, DataType::Float32 | DataType::Float64), + ComputeError: + "{label} must be a List column with Float32 or Float64 inner dtype, got {dtype}" + ); + } + Ok(*inner.clone()) +} + +fn zip_binary_output( + input_fields: &[Field], + kwargs: ZipBinaryKwargs, +) -> PolarsResult { + polars_ensure!( + input_fields.len() == 2, + ComputeError: "zip_binary expects exactly 2 inputs, got {}", + input_fields.len() + ); + let left = validate_operand( + input_fields[0].dtype(), + kwargs.op, + "zip_binary: left_column", + )?; + let right = validate_operand( + input_fields[1].dtype(), + kwargs.op, + "zip_binary: right_column", + )?; + let inner = if kwargs.op.is_arithmetic() { + arithmetic_inner(&left, &right) + } else { + DataType::Boolean + }; + Ok(Field::new( + PlSmallStr::from(""), + DataType::List(Box::new(inner)), + )) +} + +/// Apply a binary operation element-wise between two paired list +/// columns. +/// +/// Per element the result is identical to the corresponding scalar +/// polars op (null propagation, Kleene logic for `and`/`or`, +/// total-ordered comparisons). A null row on either side yields a null +/// row; empty lists zip to empty lists; mismatched lengths raise. +#[polars_expr(output_type_func_with_kwargs=zip_binary_output)] +fn zip_binary( + inputs: &[Series], + kwargs: ZipBinaryKwargs, +) -> PolarsResult { + polars_ensure!( + inputs.len() == 2, + ComputeError: "zip_binary expects exactly 2 inputs, got {}", + inputs.len() + ); + let op = kwargs.op; + + if op.takes_booleans() { + let left = + TypedListInput::::boolean(&inputs[0], "zip_binary: left_column")?; + let right = + TypedListInput::::boolean(&inputs[1], "zip_binary: right_column")?; + let len = broadcast_len(&[left.n_rows(), right.n_rows()])?; + let rows = (0..len) + .map(|i| zip_row(i, left.row(i), right.row(i), |a, b| kleene(op, a, b))) + .collect::>>()?; + return Ok(build_bool_list(rows)); + } + + let left = TypedListInput::::float(&inputs[0], "zip_binary: left_column")?; + let right = TypedListInput::::float(&inputs[1], "zip_binary: right_column")?; + let len = broadcast_len(&[left.n_rows(), right.n_rows()])?; + + if op.is_arithmetic() { + let rows = (0..len) + .map(|i| zip_row(i, left.row(i), right.row(i), |a, b| arithmetic(op, a, b))) + .collect::>>()?; + build_float_list( + rows, + &arithmetic_inner(left.inner_dtype(), right.inner_dtype()), + ) + } else { + let rows = (0..len) + .map(|i| zip_row(i, left.row(i), right.row(i), |a, b| comparison(op, a, b))) + .collect::>>()?; + Ok(build_bool_list(rows)) + } +} diff --git a/tests/test_cum_agg_runs.py b/tests/test_cum_agg_runs.py new file mode 100644 index 0000000..8890a29 --- /dev/null +++ b/tests/test_cum_agg_runs.py @@ -0,0 +1,298 @@ +"""cum_agg_runs: the scan oracle is the spec itself -- for each emitted +position, the vertical aggregation (computed by polars) over the run +prefix's non-null values -- run over an adversarial pattern matrix.""" +import math + +import polars as pl +import polars.exceptions +import pytest + +import polars_list_utils as polist + +NAN = float("nan") + +AGGS = ["sum", "mean", "median", "std", "min", "max", "delta", "count"] + +VALUE_PATTERNS = [ + [], + [1.0], + [None], + [NAN], + [NAN, 1.0], + [1.0, NAN, 0.5, 2.0], + [1.0, None, 2.0], + [None, None], + [NAN, NAN], + [0.0, -0.0], + [1.0, -1.0, 0.0, None, NAN, 3.0, 2.5, -0.5], +] + + +def gate_patterns(n: int) -> list[list[bool | None]]: + if n == 0: + return [[]] + pats: list[list[bool | None]] = [[True] * n, [False] * n, [None] * n] + if n >= 2: + pats += [ + ([True, False] * n)[:n], + ([True, None] * n)[:n], + [True] * (n // 2) + [False] * (n - n // 2), + [None] + [True] * (n - 1), + # A null gate between True elements must split the run. + [True, None] + [True] * (n - 2), + ] + return pats + + +def vertical(agg: str, prefix: list[float]) -> float | int | None: + """The vertical aggregation over a run prefix's non-null values, + with polars itself as the kernel (the deference doctrine).""" + s = pl.Series(prefix, dtype=pl.Float64) + if agg == "sum": + return s.sum() + if agg == "count": + return len(prefix) # prefix holds only non-null values + if agg == "delta": + mx, mn = s.max(), s.min() + return None if mx is None else mx - mn # ty: ignore[unsupported-operator] + return getattr(s, agg)() + + +def scan_oracle(agg, values, gates, outside) -> list: + fill = 0 if outside == "zero" else None + out: list = [] + prefix: list[float] = [] + prev_gate: object = object() # unequal to True/False/None + for v, g in zip(values, gates, strict=True): + if g != prev_gate: + prefix = [] + prev_gate = g + if g is not True: + out.append(fill) + elif v is None: + # Null values leave the state unchanged and emit null -- + # except count, which emits the unchanged running count. + out.append(len(prefix) if agg == "count" else None) + else: + prefix.append(v) + out.append(vertical(agg, prefix)) + return out + + +def assert_row_matches(got: list, expected: list, ctx: str) -> None: + assert len(got) == len(expected), ctx + for j, (g, e) in enumerate(zip(got, expected)): + where = f"{ctx}, position {j}: got {g}, expected {e}" + if e is None: + assert g is None, where + elif isinstance(e, float) and math.isnan(e): + assert g is not None and math.isnan(g), where + else: + assert g == pytest.approx(e), where + + +@pytest.mark.parametrize("outside", ["null", "zero"]) +@pytest.mark.parametrize("agg", AGGS) +def test_matches_scan_oracle(agg, outside): + for values in VALUE_PATTERNS: + for gates in gate_patterns(len(values)): + df = pl.DataFrame( + {"v": [values], "g": [gates]}, + schema={"v": pl.List(pl.Float64), "g": pl.List(pl.Boolean)}, + ) + got = df.select( + polist.cum_agg_runs( + "v", "g", aggregation=agg, outside=outside + ).alias("r") + )["r"][0].to_list() + expected = scan_oracle(agg, values, gates, outside) + assert_row_matches( + got, expected, f"agg={agg} outside={outside} v={values} g={gates}" + ) + + +@pytest.mark.parametrize( + ("agg", "native"), + [("sum", "cum_sum"), ("min", "cum_min"), ("max", "cum_max"), ("count", "cum_count")], +) +def test_all_true_gate_matches_native_cum(agg, native): + # The §6.3 anchor, executable: with an all-True gate the scan + # reduces to polars' native cumulative aggregation. (No leading-NaN + # data here: polars leaks its fold seed there, see + # test_leading_nan_follows_prefix_rule_not_polars_seed_leak.) + data = [1.0, None, 2.0, NAN, 0.5, 2.0, None, -1.0] + expected = getattr(pl.Series(data, dtype=pl.Float64), native)().to_list() + df = pl.DataFrame({"v": [data]}, schema={"v": pl.List(pl.Float64)}) + got = df.select( + polist.cum_agg_runs( + "v", pl.lit([True] * len(data)), aggregation=agg + ).alias("r") + )["r"][0].to_list() + assert_row_matches(got, expected, f"{agg} vs {native}") + + +def test_leading_nan_follows_prefix_rule_not_polars_seed_leak(): + # Native cum_max([NaN, 1.0]) emits -1.8e308 -- f64::MIN, its fold + # seed -- at the NaN position: a sentinel leak, not any aggregation + # of [NaN]. The prefix rule mandates the vertical max of [NaN], + # which is NaN; this is a documented deviation from native cum_*. + df = pl.DataFrame({"v": [[NAN, 1.0]], "g": [[True, True]]}) + for agg in ["min", "max", "delta"]: + got = df.select( + polist.cum_agg_runs("v", "g", aggregation=agg).alias("r") # ty: ignore[invalid-argument-type] + )["r"][0].to_list() + assert math.isnan(got[0]), agg + got = df.select( + polist.cum_agg_runs("v", "g", aggregation="max").alias("r") + )["r"][0].to_list() + assert got[1] == 1.0 + + +def test_runs_reset_after_false_gap(): + df = pl.DataFrame( + {"v": [[1.0, 2.0, 3.0, 4.0]], "g": [[True, True, False, True]]} + ) + out = df.select( + polist.cum_agg_runs("v", "g", aggregation="sum").alias("r") + )["r"][0].to_list() + assert out == [1.0, 3.0, None, 4.0] + + +def test_null_gate_breaks_runs(): + # True, null, True is three runs: the second True element starts fresh. + df = pl.DataFrame( + {"v": [[1.0, 2.0, 3.0]]}, + schema={"v": pl.List(pl.Float64)}, + ).with_columns(pl.lit([True, None, True]).alias("g")) + out = df.select( + polist.cum_agg_runs("v", "g", aggregation="sum").alias("r") + )["r"][0].to_list() + assert out == [1.0, None, 3.0] + + +def test_count_emits_running_count_at_null_values(): + # Matching native cum_count: [1.0, null, 2.0] -> [1, 1, 2]. + df = pl.DataFrame( + {"v": [[1.0, None, 2.0]]}, schema={"v": pl.List(pl.Float64)} + ) + out = df.select( + polist.cum_agg_runs("v", pl.lit([True] * 3), aggregation="count").alias("r") + )["r"][0].to_list() + assert out == [1, 1, 2] + + +def test_outside_zero_fills_non_true_positions(): + df = pl.DataFrame( + {"v": [[1.0, 2.0, 3.0]]}, + schema={"v": pl.List(pl.Float64)}, + ).with_columns(pl.lit([True, False, None]).alias("g")) + sums = df.select( + polist.cum_agg_runs("v", "g", aggregation="sum", outside="zero").alias("r") + )["r"][0].to_list() + assert sums == [1.0, 0.0, 0.0] + counts = df.select( + polist.cum_agg_runs("v", "g", aggregation="count", outside="zero").alias("r") + )["r"][0].to_list() + assert counts == [1, 0, 0] + + +def test_std_needs_two_values_within_run(): + df = pl.DataFrame( + {"v": [[1.0, 2.0, 3.0, 4.0]], "g": [[True, True, False, True]]} + ) + out = df.select( + polist.cum_agg_runs("v", "g", aggregation="std").alias("r") + )["r"][0].to_list() + assert out[0] is None # one value + assert out[1] == pytest.approx(pl.Series([1.0, 2.0]).std()) + assert out[2] is None # outside + assert out[3] is None # new run, one value again + + +def test_float32_values_stay_float32(): + df = pl.DataFrame( + {"v": [[1.0, 2.0]], "g": [[True, True]]}, + schema={"v": pl.List(pl.Float32), "g": pl.List(pl.Boolean)}, + ) + for agg in ["sum", "mean", "median", "std", "min", "max", "delta"]: + out = df.select( + polist.cum_agg_runs("v", "g", aggregation=agg).alias("r") # ty: ignore[invalid-argument-type] + ) + assert out["r"].dtype == pl.List(pl.Float32), agg + + +@pytest.mark.parametrize("inner", [pl.Float32, pl.Float64]) +def test_count_emits_uint32(inner): + df = pl.DataFrame( + {"v": [[1.0, 2.0]], "g": [[True, True]]}, + schema={"v": pl.List(inner), "g": pl.List(pl.Boolean)}, + ) + out = df.select(polist.cum_agg_runs("v", "g", aggregation="count").alias("r")) + assert out["r"].dtype == pl.List(pl.UInt32) + assert out["r"][0].to_list() == [1, 2] + + +def test_integer_values_raise(): + df = pl.DataFrame({"v": [[1, 2]], "g": [[True, True]]}) + with pytest.raises(polars.exceptions.PolarsError, match="Float32 or Float64"): + df.select(polist.cum_agg_runs("v", "g", aggregation="sum").alias("r")) + + +def test_non_boolean_gate_raises(): + df = pl.DataFrame({"v": [[1.0, 2.0]], "g": [[1.0, 0.0]]}) + with pytest.raises(polars.exceptions.PolarsError, match="Boolean"): + df.select(polist.cum_agg_runs("v", "g", aggregation="sum").alias("r")) + + +def test_dtype_errors_raise_at_plan_time(): + lf = pl.LazyFrame({"v": [[1.0]], "g": [[1.0]]}).select( + polist.cum_agg_runs("v", "g", aggregation="sum").alias("r") + ) + with pytest.raises(polars.exceptions.PolarsError, match="Boolean"): + lf.collect_schema() + + +def test_length_mismatch_raises(): + df = pl.DataFrame({"v": [[1.0, 2.0]], "g": [[True]]}) + with pytest.raises(polars.exceptions.PolarsError, match="differ in length"): + df.select(polist.cum_agg_runs("v", "g", aggregation="sum").alias("r")) + + +def test_null_rows_stay_null(): + df = pl.DataFrame( + {"v": [[1.0], None], "g": [None, [True]]}, + schema={"v": pl.List(pl.Float64), "g": pl.List(pl.Boolean)}, + ) + out = df.select(polist.cum_agg_runs("v", "g", aggregation="sum").alias("r")) + assert out["r"].to_list() == [None, None] + + +def test_empty_lists_scan_to_empty(): + df = pl.DataFrame( + {"v": [[]], "g": [[]]}, + schema={"v": pl.List(pl.Float64), "g": pl.List(pl.Boolean)}, + ) + out = df.select(polist.cum_agg_runs("v", "g", aggregation="sum").alias("r")) + assert out["r"][0].to_list() == [] + assert out["r"].dtype == pl.List(pl.Float64) + + +def test_gate_literal_broadcasts(): + df = pl.DataFrame({"v": [[1.0, 2.0], [3.0, 4.0]]}) + out = df.select( + polist.cum_agg_runs("v", pl.lit([True, True]), aggregation="sum").alias("r") + ) + assert out["r"].to_list() == [[1.0, 3.0], [3.0, 7.0]] + + +@pytest.mark.parametrize("engine", ["in-memory", "streaming"]) +def test_lazy_execution_matches_eager(engine): + df = pl.DataFrame( + {"v": [[float(i), float(i) + 1.0, float(i) + 2.0] for i in range(2000)]} + ) + expr = polist.cum_agg_runs( + "v", pl.lit([True, False, True]), aggregation="sum", outside="zero" + ).alias("r") + eager = df.select(expr)["r"].to_list() + lazy = df.lazy().select(expr).collect(engine=engine)["r"].to_list() # ty: ignore[not-subscriptable] + assert lazy == eager diff --git a/tests/test_plan_purity.py b/tests/test_plan_purity.py new file mode 100644 index 0000000..35de9e5 --- /dev/null +++ b/tests/test_plan_purity.py @@ -0,0 +1,77 @@ +"""A pipeline exercising every function must run on the streaming +engine with zero in-memory fallback nodes: a non-elementwise plugin +registration would silently materialize the whole stream at that node. + +The streaming executor logs each physical node under POLARS_VERBOSE; +a logical subplan the engine cannot stream runs wrapped in an +"in-memory-map" node, which is the fallback marker this test forbids +("in-memory-source"/"in-memory-sink" are just the literal-frame source +and the result sink, and appear in every fully-streaming plan too). +""" +import os +import subprocess +import sys + +PIPELINE = """ +import polars as pl +import polars_list_utils as polist + +Fs, N = 200.0, 64 +FREQS = [i * Fs / N for i in range(N // 2 + 1)] +signal = [float(i % 7) for i in range(N)] + +lf = ( + pl.LazyFrame({"grp": [1, 2], "signal": [signal, signal]}) + .with_columns( + polist.apply_butterworth("signal", sample_rate=Fs, max_freq=50.0).alias("filt") + ) + .with_columns( + polist.apply_fft( + "filt", sample_rate=Fs, window="hann", scaling="amplitude" + ).alias("fft") + ) + .with_columns( + polist.apply_interp( + pl.lit(FREQS), "fft", pl.lit([10.0, 20.0, 30.0]) + ).alias("interp") + ) + .with_columns( + polist.agg_slices( + "fft", + pl.lit(FREQS), + aggregation="mean", + slices_include=[(10.0, 30.0)], + strict=True, + ).alias("feat") + ) + .with_columns(polist.zip_binary("fft", "fft", op="gt").alias("gates")) + .with_columns( + polist.cum_agg_runs( + "fft", "gates", aggregation="sum", outside="zero" + ).alias("cum") + ) + .group_by("grp") + .agg( + polist.agg_lists( + "interp", list_length=3, aggregation="mean", strict=True + ).alias("agg") + ) +) +assert lf.collect(engine="streaming").height == 2 +""" + + +def test_pipeline_streams_without_in_memory_fallback(): + env = os.environ.copy() + env["POLARS_VERBOSE"] = "1" + proc = subprocess.run( + [sys.executable, "-c", PIPELINE], + capture_output=True, + text=True, + env=env, + timeout=120, + ) + log = proc.stderr + proc.stdout + assert proc.returncode == 0, log + assert "polars-stream" in log, f"streaming engine did not run:\n{log}" + assert "in-memory-map" not in log, f"in-memory fallback node in plan:\n{log}" diff --git a/tests/test_plugins.py b/tests/test_plugins.py index cb27d5a..75b987e 100644 --- a/tests/test_plugins.py +++ b/tests/test_plugins.py @@ -249,14 +249,38 @@ def test_agg_slices_empty_selection_is_null(df_ramp): assert out["cnt"][0] == 0.0 +def test_agg_slices_empty_selection_sum_is_positive_zero(df_ramp): + # Sum's empty identity is +0.0 (polars itself is split between + # +0.0 for Series.sum and -0.0 for list.sum of an empty list). + out = df_ramp.with_columns( + polist.agg_slices( + "v", "i", aggregation="sum", slices_include=[(100.0, 200.0)] + ).alias("s") + ) + assert out["s"][0] == 0.0 + assert math.copysign(1.0, out["s"][0]) == 1.0 + + +def test_agg_slices_all_null_selection_sums_to_zero(): + df = pl.DataFrame( + {"v": [[None, None]], "i": [[0.0, 1.0]]}, + schema={"v": pl.List(pl.Float64), "i": pl.List(pl.Float64)}, + ) + out = df.with_columns( + polist.agg_slices("v", "i", aggregation="sum").alias("s") + ) + assert out["s"][0] == 0.0 + + def test_agg_slices_matches_polars_vertical_semantics(): # The contract: aggregating a slice behaves exactly like polars' # vertical aggregations over the same values (nulls skipped, NaN # propagates per-kernel, ddof=1 std). data = [1.0, None, float("nan"), 3.0] s = pl.Series(data, dtype=pl.Float64) - aggs: list[polist.Aggregation] = ["mean", "median", "min", "max", "std", "count"] + aggs: list[polist.Aggregation] = ["sum", "mean", "median", "min", "max", "std", "count"] expected = { + "sum": s.sum(), "mean": s.mean(), "median": s.median(), "min": s.min(), @@ -303,9 +327,9 @@ def test_agg_slices_all_nan_matches_polars(): df = pl.DataFrame({"v": [[float("nan")] * 2], "i": [[0.0, 1.0]]}) out = df.with_columns( polist.agg_slices("v", "i", aggregation=agg).alias(agg) # ty: ignore[invalid-argument-type] - for agg in ["mean", "median", "min", "max", "delta"] + for agg in ["sum", "mean", "median", "min", "max", "delta"] ) - for agg in ["mean", "median", "min", "max", "delta"]: + for agg in ["sum", "mean", "median", "min", "max", "delta"]: assert math.isnan(out[agg][0]), agg @@ -327,8 +351,9 @@ def test_agg_slices_length_mismatch_raises(df_ramp): def test_agg_slices_aggregations(df_ramp): out = df_ramp.with_columns( polist.agg_slices("v", "i", aggregation=agg).alias(agg) # ty: ignore[invalid-argument-type] - for agg in ["mean", "median", "min", "max", "std", "delta", "count"] + for agg in ["sum", "mean", "median", "min", "max", "std", "delta", "count"] ) + assert out["sum"][0] == pytest.approx(55.0) assert out["mean"][0] == pytest.approx(5.0) assert out["median"][0] == pytest.approx(5.0) assert out["min"][0] == 0.0 @@ -365,6 +390,18 @@ def test_agg_lists_elementwise_over_groups(): assert out["mean"][1].to_list() == pytest.approx([7.0, 8.0, 9.0]) +def test_agg_lists_sum_of_all_null_position_is_zero(): + # Vertical sum over an all-null position is polars' identity, 0.0. + df = pl.DataFrame( + {"g": [1, 1], "v": [[1.0, None], [2.0, None]]}, + schema={"g": pl.Int64, "v": pl.List(pl.Float64)}, + ) + out = df.group_by("g").agg( + polist.agg_lists("v", list_length=2, aggregation="sum").alias("s") + ) + assert out["s"][0].to_list() == [3.0, 0.0] + + def test_agg_lists_invalid_length_raises(): with pytest.raises(ValueError, match="list_length"): polist.agg_lists("v", list_length=0, aggregation="mean") @@ -380,7 +417,7 @@ def test_agg_lists_and_agg_slices_share_missing_data_semantics(): {"g": [1] * 4, "v": [[x] for x in data]}, schema={"g": pl.Int64, "v": pl.List(pl.Float64)}, ) - for agg in ["mean", "median", "std", "min", "max", "delta", "count"]: + for agg in ["sum", "mean", "median", "std", "min", "max", "delta", "count"]: h = df_h.with_columns( polist.agg_slices( "v", pl.lit([0.0, 1.0, 2.0, 3.0]), aggregation=agg # ty: ignore[invalid-argument-type] diff --git a/tests/test_strict.py b/tests/test_strict.py new file mode 100644 index 0000000..1919381 --- /dev/null +++ b/tests/test_strict.py @@ -0,0 +1,116 @@ +"""Every strict check gets a positive test (raises under strict=True) +and a negative test (the documented silent degenerate result under the +default).""" +import polars as pl +import polars.exceptions +import pytest + +import polars_list_utils as polist + +NAN = float("nan") + + +# ------------------------------------------------------------ apply_interp + +def _interp_frame() -> pl.DataFrame: + # Unsorted x: the sorted answer would be [5.0, 15.0]. + return pl.DataFrame({"x": [[2.0, 0.0, 1.0]], "y": [[20.0, 0.0, 10.0]]}) + + +def test_interp_strict_raises_on_unsorted_x(): + with pytest.raises(polars.exceptions.PolarsError, match="non-decreasing"): + _interp_frame().with_columns( + polist.apply_interp("x", "y", pl.lit([0.5, 1.5]), strict=True).alias("r") + ) + + +def test_interp_default_silently_interpolates_garbage(): + out = _interp_frame().with_columns( + polist.apply_interp("x", "y", pl.lit([0.5, 1.5])).alias("r") + ) + assert out["r"][0].to_list() == [20.0, 20.0] + + +def test_interp_strict_raises_on_nan_x(): + df = pl.DataFrame({"x": [[0.0, NAN, 2.0]], "y": [[0.0, 1.0, 2.0]]}) + with pytest.raises(polars.exceptions.PolarsError, match="non-decreasing"): + df.with_columns( + polist.apply_interp("x", "y", pl.lit([0.5]), strict=True).alias("r") + ) + + +def test_interp_strict_accepts_duplicate_x(): + # Duplicates stay legal, as in numpy. + df = pl.DataFrame({"x": [[0.0, 1.0, 1.0, 2.0]], "y": [[0.0, 10.0, 20.0, 30.0]]}) + out = df.with_columns( + polist.apply_interp("x", "y", pl.lit([0.5]), strict=True).alias("r") + ) + assert out["r"][0].to_list() == pytest.approx([5.0]) + + +# -------------------------------------------------------------- agg_slices + +def test_agg_slices_strict_raises_on_inverted_range(): + # A construction-time check: no DataFrame is needed to trigger it. + with pytest.raises(ValueError, match="lo <= hi"): + polist.agg_slices( + "v", "i", aggregation="mean", slices_include=[(2.0, 0.0)], strict=True + ) + + +def test_agg_slices_strict_raises_on_nan_bound(): + with pytest.raises(ValueError, match="non-NaN"): + polist.agg_slices( + "v", "i", aggregation="mean", slices_exclude=[(NAN, 1.0)], strict=True + ) + with pytest.raises(ValueError, match="non-NaN"): + polist.agg_slices( + "v", + "i", + aggregation="mean", + slices_include=[((0.0, "closed"), (NAN, "open"))], + strict=True, + ) + + +def test_agg_slices_default_inverted_range_is_empty_selection(): + df = pl.DataFrame({"v": [[1.0, 2.0]], "i": [[0.0, 1.0]]}) + out = df.with_columns( + polist.agg_slices( + "v", "i", aggregation="mean", slices_include=[(2.0, 0.0)] + ).alias("mean"), + polist.agg_slices( + "v", "i", aggregation="count", slices_include=[(2.0, 0.0)] + ).alias("cnt"), + ) + assert out["mean"][0] is None + assert out["cnt"][0] == 0.0 + + +# --------------------------------------------------------------- agg_lists + +def test_agg_lists_strict_raises_on_overlong_list(): + df = pl.DataFrame({"g": [1, 1], "v": [[1.0, 2.0, 3.0], [1.0, 2.0]]}) + with pytest.raises(polars.exceptions.PolarsError, match="list_length"): + df.group_by("g").agg( + polist.agg_lists("v", list_length=2, aggregation="mean", strict=True) + ) + + +def test_agg_lists_default_truncates_silently(): + # list_length as a deliberate window: third elements are dropped. + df = pl.DataFrame({"g": [1, 1], "v": [[1.0, 2.0, 30.0], [3.0, 4.0, 50.0]]}) + out = df.group_by("g").agg( + polist.agg_lists("v", list_length=2, aggregation="mean").alias("r") + ) + assert out["r"][0].to_list() == pytest.approx([2.0, 3.0]) + + +def test_agg_lists_strict_passes_within_length(): + # Shorter lists stay legal under strict: they contribute nulls at + # the missing positions, which is data, not a violation. + df = pl.DataFrame({"g": [1, 1], "v": [[1.0, 2.0], [3.0]]}) + out = df.group_by("g").agg( + polist.agg_lists("v", list_length=2, aggregation="mean", strict=True).alias("r") + ) + assert out["r"][0].to_list() == pytest.approx([2.0, 2.0]) diff --git a/tests/test_zip_binary.py b/tests/test_zip_binary.py new file mode 100644 index 0000000..c918745 --- /dev/null +++ b/tests/test_zip_binary.py @@ -0,0 +1,177 @@ +"""zip_binary: per element the result must be identical to the +corresponding scalar polars op -- the deference doctrine, executable.""" +import math + +import polars as pl +import polars.exceptions +import pytest + +import polars_list_utils as polist + +NAN = float("nan") +INF = float("inf") + +# Adversarial element pairs: NaN on either and both sides, signed +# zeros, nulls, infinities, and division-by-zero fodder. +A = [1.0, 2.0, NAN, NAN, -0.0, 0.0, None, 1.0, None, INF, -INF, 5.0, 1.0, 0.0] +B = [2.0, 0.0, 5.0, NAN, 0.0, -0.0, 1.0, None, None, INF, INF, NAN, 0.0, 0.0] + +FLOAT_OPS = { + "add": lambda a, b: a + b, + "sub": lambda a, b: a - b, + "mul": lambda a, b: a * b, + "div": lambda a, b: a / b, + "gt": lambda a, b: a > b, + "ge": lambda a, b: a >= b, + "lt": lambda a, b: a < b, + "le": lambda a, b: a <= b, + "eq": lambda a, b: a == b, + "ne": lambda a, b: a != b, +} + + +def assert_row_matches(got: list, expected: list, ctx: str) -> None: + assert len(got) == len(expected), ctx + for j, (g, e) in enumerate(zip(got, expected)): + where = f"{ctx}, position {j}: got {g}, expected {e}" + if e is None: + assert g is None, where + elif isinstance(e, float) and math.isnan(e): + assert g is not None and math.isnan(g), where + else: + assert g == pytest.approx(e), where + + +@pytest.mark.parametrize("op", list(FLOAT_OPS)) +def test_matches_polars_scalar_op(op): + flat = pl.DataFrame( + {"a": A, "b": B}, schema={"a": pl.Float64, "b": pl.Float64} + ) + expected = flat.select( + FLOAT_OPS[op](pl.col("a"), pl.col("b")).alias("r") + )["r"].to_list() + df = pl.DataFrame( + {"a": [A], "b": [B]}, + schema={"a": pl.List(pl.Float64), "b": pl.List(pl.Float64)}, + ) + got = df.select( + polist.zip_binary("a", "b", op=op).alias("r") + )["r"][0].to_list() + assert_row_matches(got, expected, op) + + +def test_kleene_and_or_match_polars(): + tri = [True, False, None] + a = [x for x in tri for _ in tri] # all 9 combinations + b = tri * 3 + ops = {"and": lambda x, y: x & y, "or": lambda x, y: x | y} + for op, expr_op in ops.items(): + flat = pl.DataFrame( + {"a": a, "b": b}, schema={"a": pl.Boolean, "b": pl.Boolean} + ) + expected = flat.select( + expr_op(pl.col("a"), pl.col("b")).alias("r") + )["r"].to_list() + df = pl.DataFrame( + {"a": [a], "b": [b]}, + schema={"a": pl.List(pl.Boolean), "b": pl.List(pl.Boolean)}, + ) + got = df.select( + polist.zip_binary("a", "b", op=op).alias("r") # ty: ignore[invalid-argument-type] + )["r"][0].to_list() + assert got == expected, op + + +def test_comparison_output_is_boolean(): + df = pl.DataFrame({"a": [[1.0]], "b": [[2.0]]}) + out = df.select(polist.zip_binary("a", "b", op="le").alias("r")) + assert out["r"].dtype == pl.List(pl.Boolean) + + +def test_arithmetic_dtype_follows_supertyping(): + f32 = pl.DataFrame( + {"a": [[1.0]], "b": [[2.0]]}, + schema={"a": pl.List(pl.Float32), "b": pl.List(pl.Float32)}, + ) + assert f32.select(polist.zip_binary("a", "b", op="mul").alias("r"))[ + "r" + ].dtype == pl.List(pl.Float32) + + mixed = pl.DataFrame( + {"a": [[1.0]], "b": [[2.0]]}, + schema={"a": pl.List(pl.Float32), "b": pl.List(pl.Float64)}, + ) + assert mixed.select(polist.zip_binary("a", "b", op="mul").alias("r"))[ + "r" + ].dtype == pl.List(pl.Float64) + + +def test_integer_lists_raise(): + df = pl.DataFrame({"a": [[1, 2]], "b": [[3, 4]]}) + with pytest.raises(polars.exceptions.PolarsError, match="Float32 or Float64"): + df.select(polist.zip_binary("a", "b", op="add").alias("r")) + + +def test_arithmetic_on_booleans_raises(): + df = pl.DataFrame({"a": [[True]], "b": [[False]]}) + with pytest.raises(polars.exceptions.PolarsError, match="Float32 or Float64"): + df.select(polist.zip_binary("a", "b", op="add").alias("r")) + + +def test_and_on_floats_raises(): + df = pl.DataFrame({"a": [[1.0]], "b": [[0.0]]}) + with pytest.raises(polars.exceptions.PolarsError, match="Boolean"): + df.select(polist.zip_binary("a", "b", op="and").alias("r")) + + +def test_dtype_errors_raise_at_plan_time(): + # Type-decidable violations must fail at schema resolution, before + # any row is computed. + lf = pl.LazyFrame({"a": [[1, 2]], "b": [[3, 4]]}).select( + polist.zip_binary("a", "b", op="add").alias("r") + ) + with pytest.raises(polars.exceptions.PolarsError, match="Float32 or Float64"): + lf.collect_schema() + + +def test_length_mismatch_raises(): + df = pl.DataFrame({"a": [[1.0, 2.0]], "b": [[1.0]]}) + with pytest.raises(polars.exceptions.PolarsError, match="differ in length"): + df.select(polist.zip_binary("a", "b", op="add").alias("r")) + + +def test_null_rows_stay_null(): + df = pl.DataFrame( + {"a": [[1.0], None], "b": [None, [2.0]]}, + schema={"a": pl.List(pl.Float64), "b": pl.List(pl.Float64)}, + ) + out = df.select(polist.zip_binary("a", "b", op="add").alias("r")) + assert out["r"].to_list() == [None, None] + + +def test_empty_lists_zip_to_empty(): + df = pl.DataFrame( + {"a": [[]], "b": [[]]}, + schema={"a": pl.List(pl.Float64), "b": pl.List(pl.Float64)}, + ) + out = df.select(polist.zip_binary("a", "b", op="add").alias("r")) + assert out["r"][0].to_list() == [] + + +def test_literal_broadcasts_as_constant_template(): + df = pl.DataFrame({"a": [[1.0, 2.0], [3.0, 4.0]]}) + out = df.select( + polist.zip_binary("a", pl.lit([10.0, 20.0]), op="add").alias("r") + ) + assert out["r"].to_list() == [[11.0, 22.0], [13.0, 24.0]] + + +@pytest.mark.parametrize("engine", ["in-memory", "streaming"]) +def test_lazy_execution_matches_eager(engine): + df = pl.DataFrame( + {"a": [[float(i), float(i) + 1.0] for i in range(2000)]} + ) + expr = polist.zip_binary("a", pl.lit([1.0, -1.0]), op="mul").alias("r") + eager = df.select(expr)["r"].to_list() + lazy = df.lazy().select(expr).collect(engine=engine)["r"].to_list() # ty: ignore[not-subscriptable] + assert lazy == eager From 852afe2c574c887b27c5f1f92738c51c679fc68f Mon Sep 17 00:00:00 2001 From: Travis Hammond Date: Tue, 1 Sep 2026 11:19:56 +0200 Subject: [PATCH 03/17] fix: correct numeric kernels, kwargs aliasing, and error messages from review Adversarial review of the PR 1 draft, against polars as the anchor. Numerics, all verified against polars over adversarial value matrices: - The even-count median used (a + b) / 2, which overflows for large equal values and loses precision generally: the median of two 1e308s read inf, and [1.0, -inf] read -inf where polars gives NaN. polars interpolates instead, a + (b - a) * 0.5, short-circuiting on equality so that equal infinities do not read NaN. The interpolated form matches polars on every probed edge case and on 30000/30000 random pairs, where the midpoint form misses 8679. - min/max broke signed-zero ties the wrong way. polars keeps the first of values that compare equal; f64::min/max keep the second, so [-0.0, 0.0] and [0.0, -0.0] both came out sign-flipped. The scan's median insert had the same bug in another guise, using a lower bound where a stable sort needs an upper one. - sum folded from IEEE's -0.0 identity, so an all-negative-zero selection read -0.0 where polars reads +0.0. - The scan's Welford variance reported zero spread for values that differ, whenever the running mean rounded onto the next value -- at timestamp magnitudes an ordinary input, and the worst possible failure for a std feeding a threshold. std and median now delegate to the same kernels agg_slices uses, so one input gives one answer everywhere in the library, at the cost of a retained buffer for those two aggregations. Kwargs cross the plugin boundary as a pickle, and the decoder cannot resolve a memoized container back-reference. Any range object, bound tuple or range list appearing twice therefore killed the query with "recursive structure found" -- and CPython's constant folding arranges exactly that for a literal as ordinary as a single-point closed range. The wrapper now rebuilds ranges into fresh plain-float tuples, which also keeps numpy scalars off the wire, and accepts the list-shaped form that strict validation previously rejected with a TypeError. Errors raised inside a row named a row index. Plugins are handed one chunk at a time, so that index was chunk-local: it changed with chunk layout and with the engine, and pointed at a different, valid row. The messages now describe the offending row by its shape instead. apply_interp's strict check no longer treats NaN as a violation. The spec puts NaN on the propagating side of the family rule, and the check exists for unsorted input, which fails silently; a NaN announces itself in the output. Docs now state what the code does: Float32 accumulates in f64 and rounds once, so it can diverge from polars' per-step rounding; the native cum_min/cum_max seed leak affects leading infinities as well as NaN; and agreement with polars is semantic, not bit-for-bit, since accumulation order differs. Co-Authored-By: Claude Fable 5 --- README.md | 23 +++- polars_list_utils/__init__.py | 96 +++++++++---- src/agg_lists.rs | 20 +-- src/apply_butterworth.rs | 4 +- src/apply_fft.rs | 2 +- src/apply_interp.rs | 24 ++-- src/cum_agg_runs.rs | 67 ++++----- src/statistic.rs | 69 ++++++++-- src/util.rs | 10 +- src/zip_binary.rs | 12 +- tests/test_cum_agg_runs.py | 212 ++++++++++++++++++++++++++++- tests/test_error_messages.py | 92 +++++++++++++ tests/test_kwargs_serialization.py | 137 +++++++++++++++++++ tests/test_strict.py | 75 +++++++++- tests/test_zip_binary.py | 107 ++++++++++++++- 15 files changed, 832 insertions(+), 118 deletions(-) create mode 100644 tests/test_error_messages.py create mode 100644 tests/test_kwargs_serialization.py diff --git a/README.md b/README.md index 4799da7..f9007a9 100644 --- a/README.md +++ b/README.md @@ -64,9 +64,14 @@ boost. `True`-gated elements accumulate, and every other position emits null (or a plain zero with `outside="zero"`). - Each emitted element equals the vertical aggregation over the run's elements so far, - so an all-`True` gate reproduces `cum_sum`/`cum_min`/`cum_max`/`cum_count`; null values - emit null but keep the running state, except `count`, which emits the running count as - `cum_count` does. + so an all-`True` gate over `Float64` tracks `cum_sum`/`cum_min`/`cum_max`/`cum_count`; + null values emit null but keep the running state, except `count`, which emits the + running count as `cum_count` does. + - Two documented departures from the native cumulatives: `cum_min`/`cum_max` fold from + ∓`f64::MAX` and never replace that seed on a tie, so they report ±1.8e308 where the + prefix is all-NaN or opens with an infinity of the same sign, while this function + reports the true prefix extreme; and `Float32` values accumulate in `f64` and round + once at the boundary, where Polars rounds every step. - `Float32` values stay `Float32`; `count` emits `UInt32` (Polars' count dtype). The six plugin functions accept a length-1 literal list column (e.g. `pl.lit([...])`) for @@ -75,8 +80,9 @@ names, mismatched list lengths) raises an error instead of silently returning nu per-row data problems (e.g. a signal length that is not a power of two) yield null rows. Both aggregation functions support `sum`, `mean`, `median`, `std`, `min`, `max`, `delta`, -and `count`, and both handle missing data exactly like the vertical aggregations of Polars -itself: +and `count`, and both handle missing data like the vertical aggregations of Polars itself. +(Semantically, that is: floating-point accumulation order differs from Polars', so results +agree to rounding rather than bit-for-bit.) - Nulls are skipped: excluded from both the numerator and the denominator, so `count` only counts non-null values and an empty or all-null selection yields null (`count`: 0, @@ -96,7 +102,12 @@ and cumulative Polars operations per element instead, and treat empty lists as v Some checks are opt-in via `strict=True` (default off): `apply_interp` verifies that `x` is non-decreasing, `agg_slices` rejects inverted or NaN range bounds at expression construction, and `agg_lists` raises on lists longer than `list_length` instead of -silently truncating them. +silently truncating them. Each check covers a failure that is otherwise silent; NaN, which +announces itself in the output, propagates rather than raising. + +Errors raised from inside a row describe the offending row by its shape rather than its +position: Polars hands a plugin one chunk at a time, so any row number counted there would +be chunk-local and would point at a different, valid row of the frame. ### Example diff --git a/polars_list_utils/__init__.py b/polars_list_utils/__init__.py index 7fd93dd..d270ba7 100644 --- a/polars_list_utils/__init__.py +++ b/polars_list_utils/__init__.py @@ -1,7 +1,10 @@ """Polars expression plugins for signal processing on List columns. -Six Rust plugins operating elementwise on `List[f64]` columns, plus -one pure-polars helper: +Six Rust plugins operating elementwise on List columns, plus one +pure-polars helper. The four original functions take `List[f64]` +(anything numeric is cast to it); :func:`zip_binary` and +:func:`cum_agg_runs` instead take float or Boolean inners as their +operands require, and preserve `Float32`: - :func:`apply_interp`: interpolate (x, y) data onto new x coordinates. - :func:`apply_butterworth`: zero-phase Butterworth filtering. @@ -29,7 +32,7 @@ import math from collections.abc import Callable from pathlib import Path -from typing import Literal +from typing import Any, Literal import polars as pl from polars.plugins import register_plugin_function @@ -106,10 +109,12 @@ def apply_interp( values outside the data range clamped to the first/last `y` value. `x` must be sorted in increasing order. - With `strict=True` a row whose `x` values are not non-decreasing - raises (duplicates stay legal, as in numpy; a NaN in `x` counts as - a violation). The default keeps numpy's silent behaviour: unsorted - `x` interpolates garbage without complaint. + With `strict=True` a row whose `x` values descend raises + (duplicates stay legal, as in numpy). The default keeps numpy's + silent behaviour: unsorted `x` interpolates garbage without + complaint. A NaN in `x` is a legitimate float either way and + propagates into the output rather than raising — the check exists + for the failure you cannot see, not the one you can. Returns a `List[f64]` column of interpolated y values, one per `xp` coordinate. Raises if x and y lists differ in length. @@ -183,17 +188,44 @@ def apply_fft( ) -def _validate_ranges(ranges: list[Range] | None, param: str) -> None: - """Raise on inverted or NaN range bounds (agg_slices strict=True).""" - for r in ranges or []: - lo, hi = r - lo_value = float(lo[0]) if isinstance(lo, tuple) else float(lo) - hi_value = float(hi[0]) if isinstance(hi, tuple) else float(hi) - if math.isnan(lo_value) or math.isnan(hi_value) or lo_value > hi_value: +def _normalize_ranges( + ranges: list[Range] | None, param: str, strict: bool +) -> list[Range] | None: + """Rebuild `ranges` into fresh plain-float tuples, validating under strict. + + The rebuild is not cosmetic. The plugin receives kwargs as a pickle, + and its decoder cannot resolve a memoized *container* back-reference: + any range object (or bound tuple, or the range list itself) that + appears twice — as CPython's constant folding arranges for a literal + like `((1.0, "closed"), (1.0, "closed"))` — would otherwise fail the + query with "recursive structure found". Fresh objects are never + pickled twice. Coercing bounds with `float()` likewise keeps numpy + scalars, which the decoder cannot resolve either, off the wire. + """ + if ranges is None: + return None + + def bound(b: Any) -> tuple[float, str] | float: + # Explicit ((value, mode)) form, as a tuple or (from JSON) a list. + if isinstance(b, (tuple, list)): + value, mode = b + return (float(value), str(mode)) + return float(b) + + out: list[Range] = [] + for r in ranges: + lo, hi = (bound(b) for b in r) + lo_value = lo[0] if isinstance(lo, tuple) else lo + hi_value = hi[0] if isinstance(hi, tuple) else hi + if strict and ( + math.isnan(lo_value) or math.isnan(hi_value) or lo_value > hi_value + ): raise ValueError( f"agg_slices: invalid range {r!r} in {param} (strict=True): " "bounds must be non-NaN with lo <= hi" ) + out.append((lo, hi)) # ty: ignore[invalid-argument-type] + return out def agg_slices( @@ -212,8 +244,10 @@ def agg_slices( `slices_exclude` range, then aggregates them into a `Float64`. Omitting `slices_include` includes everything. - Missing data follows the polars convention, so results match - polars' vertical aggregations over the same values: null elements + Missing data follows the polars convention, so results follow + polars' vertical aggregations over the same values (semantically — + floating-point accumulation order differs, so agreement is to + rounding, not bit-for-bit): null elements are skipped (pairwise with their index), while NaN is a legitimate float value — it poisons `sum`/`mean`/`std`, is skipped by `min`/`max` unless all values are NaN, and sorts as the largest @@ -224,19 +258,20 @@ def agg_slices( `std` is the sample standard deviation (ddof=1, polars' default) and yields null for selections with fewer than two values. + Range bounds accept anything `float()` converts, and the explicit + form may be given as tuples or lists (so ranges loaded from JSON + work as-is). + With `strict=True`, inverted (`lo > hi`) or NaN range bounds raise at expression construction; the default keeps them as legitimately empty selections. """ - if strict: - _validate_ranges(slices_include, "slices_include") - _validate_ranges(slices_exclude, "slices_exclude") return _plugin( "agg_slices", [value_column, index_column], aggregation=aggregation, - slices_include=slices_include, - slices_exclude=slices_exclude, + slices_include=_normalize_ranges(slices_include, "slices_include", strict), + slices_exclude=_normalize_ranges(slices_exclude, "slices_exclude", strict), ) @@ -338,8 +373,23 @@ def cum_agg_runs( emits the `outside` fill (`"null"`, or `"zero"` for a plain 0 in the output dtype). Within a run, each emitted element equals the vertical `aggregation` over the run's elements so far — so with an - all-`True` gate, `sum`/`min`/`max`/`count` match polars' - `cum_sum`/`cum_min`/`cum_max`/`cum_count`. + all-`True` gate over `Float64` values, `sum`/`min`/`max`/`count` + track polars' `cum_sum`/`cum_min`/`cum_max`/`cum_count`, with two + documented exceptions: + + - `cum_min`/`cum_max` fold from ∓`f64::MAX` and never replace that + seed on a tie, so they emit ±1.7976931348623157e+308 wherever the + prefix is all-NaN or opens with an infinity of the same sign. + This function follows the prefix rule and emits the true prefix + minimum or maximum (NaN, or that infinity). + - `Float32` values accumulate in `f64` and round once at the output + boundary, while polars rounds every step, so the two can differ + in the last bits and around the `Float32` range limit (where this + function can stay finite as polars overflows to infinity). + + Floating-point accumulation order differs from polars generally, so + treat the correspondence as one of semantics — missing data, NaN, + dtype and ddof rules — rather than bit-exact equality. Missing data follows the polars `cum_*` convention: a null value leaves the running state unchanged and emits null — except `count`, diff --git a/src/agg_lists.rs b/src/agg_lists.rs index d7316a8..4a0c4cf 100644 --- a/src/agg_lists.rs +++ b/src/agg_lists.rs @@ -18,21 +18,23 @@ fn passthrough_output(input_fields: &[Field]) -> PolarsResult { /// Identity pass-through that raises when any row's list is longer /// than `list_length` (which `agg_lists` would otherwise silently /// truncate to). Null rows are data and pass through. +/// +/// The message names no row index: the plugin is handed one chunk at a +/// time, so any position named here would be chunk-local and point at +/// a different, innocent row of the frame. #[polars_expr(output_type_func=passthrough_output)] fn check_list_len( inputs: &[Series], kwargs: CheckListLenKwargs, ) -> PolarsResult { let s = &inputs[0]; - for (i, row) in s.list()?.into_iter().enumerate() { - if let Some(row) = row { - polars_ensure!( - row.len() <= kwargs.list_length, - ComputeError: - "agg_lists: row {i} has list length {} but list_length={} (strict=true)", - row.len(), kwargs.list_length - ); - } + for row in s.list()?.into_iter().flatten() { + polars_ensure!( + row.len() <= kwargs.list_length, + ComputeError: + "agg_lists: a row has list length {} but list_length={} (strict=true)", + row.len(), kwargs.list_length + ); } Ok(s.clone()) } diff --git a/src/apply_butterworth.rs b/src/apply_butterworth.rs index e1849b3..bfafa2d 100644 --- a/src/apply_butterworth.rs +++ b/src/apply_butterworth.rs @@ -37,14 +37,14 @@ fn apply_butterworth( (Some(lo), None) => Filter::Highpass(lo), (None, Some(hi)) => Filter::Lowpass(hi), (None, None) => { - return apply_list_transform(inputs, |_row, cols| Ok(Some(cols[0].to_vec()))); + return apply_list_transform(inputs, |cols| Ok(Some(cols[0].to_vec()))); } }; let order = kwargs.filter_order.unwrap_or(4); let built = filter.build(kwargs.sample_rate, order)?; let min_samples = filter.min_samples(order); - apply_list_transform(inputs, move |_row, cols| { + apply_list_transform(inputs, move |cols| { if cols[0].len() < min_samples { return Ok(None); } diff --git a/src/apply_fft.rs b/src/apply_fft.rs index 18d1175..05d778a 100644 --- a/src/apply_fft.rs +++ b/src/apply_fft.rs @@ -33,7 +33,7 @@ fn apply_fft( let window = kwargs.window.unwrap_or_default(); let scaling = kwargs.scaling.unwrap_or_default(); - apply_list_transform(inputs, |_row, cols| { + apply_list_transform(inputs, |cols| { Ok(compute_fft(cols[0], kwargs.sample_rate, window, scaling)) }) } diff --git a/src/apply_interp.rs b/src/apply_interp.rs index e3ad052..3023e8e 100644 --- a/src/apply_interp.rs +++ b/src/apply_interp.rs @@ -15,33 +15,31 @@ struct InterpKwargs { /// first/last `y` value. /// /// `x` must be non-decreasing for the result to be meaningful; with -/// `strict` the first violation raises (a NaN in `x` counts as one), -/// without it unsorted `x` silently interpolates garbage, as in numpy. +/// `strict` the first descending pair raises, without it unsorted `x` +/// silently interpolates garbage, as in numpy. NaN in `x` is a +/// legitimate float that propagates into the output either way: the +/// check exists for the silent failure, not the visible one. #[polars_expr(output_type_func=list_f64_output)] fn apply_interp( inputs: &[Series], kwargs: InterpKwargs, ) -> PolarsResult { - apply_list_transform(inputs, |row_idx, cols| { + apply_list_transform(inputs, |cols| { polars_ensure!( cols[0].len() == cols[1].len(), ComputeError: "apply_interp: x and y lists differ in length ({} vs {})", cols[0].len(), cols[1].len() ); - // Non-decreasing means each pair compares Less or Equal; a NaN - // pair is incomparable and counts as a violation. - let in_order = |w: &&[f64]| { - matches!( - w[0].partial_cmp(&w[1]), - Some(std::cmp::Ordering::Less | std::cmp::Ordering::Equal) - ) - }; + // Only a strictly descending pair is a violation; a NaN pair is + // incomparable, and NaN propagates per the transform family's + // rule. No row index: the plugin is handed one chunk at a time, + // so any position named here would be chunk-local. if kwargs.strict - && let Some(w) = cols[0].windows(2).find(|w| !in_order(w)) + && let Some(w) = cols[0].windows(2).find(|w| w[0] > w[1]) { polars_bail!(ComputeError: "apply_interp: x_column must be non-decreasing with strict=true, \ - found {} followed by {} at row {row_idx}", + found {} followed by {}", w[0], w[1] ); } diff --git a/src/cum_agg_runs.rs b/src/cum_agg_runs.rs index 82f305a..58f6a2e 100644 --- a/src/cum_agg_runs.rs +++ b/src/cum_agg_runs.rs @@ -1,4 +1,4 @@ -use crate::statistic::Aggregation; +use crate::statistic::{Aggregation, Statistic, max_fold, midpoint, min_fold}; use crate::util::{ TypedListInput, broadcast_len, build_float_list, build_u32_list, tot_cmp, }; @@ -24,18 +24,23 @@ struct CumAggRunsKwargs { /// Running state over one gated run's non-null values. /// /// Every emitted position must equal the vertical aggregation over the -/// run-prefix's non-null values (the prefix rule): `sum`/`mean`/`std` -/// let NaN poison the accumulators, `min`/`max` skip NaN via the -/// NaN-propagating `f64::min`/`max` seeds (all-NaN prefixes read NaN), -/// `median` keeps a sorted prefix with NaN largest, `count` counts -/// every non-null value (NaN included). +/// run-prefix's non-null values (the prefix rule), so `std` and +/// `median` delegate to the very kernels `agg_slices` uses rather than +/// keeping independent accumulators: one input, one answer, everywhere +/// in the library. That costs a retained buffer for those two +/// aggregations; the rest need O(1) state. +/// +/// `sum`/`mean` let NaN poison the running sum, `min`/`max` skip NaN +/// (all-NaN prefixes read NaN) and keep the first of equal values, and +/// `count` counts every non-null value, NaN included. struct RunState { n: u64, sum: f64, min: f64, max: f64, - mean: f64, - m2: f64, + /// The run's values in input order, kept for `std` only. + values: Vec, + /// The run's values in stable sorted order, kept for `median` only. sorted: Vec, } @@ -46,8 +51,7 @@ impl RunState { sum: 0.0, min: f64::NAN, max: f64::NAN, - mean: 0.0, - m2: 0.0, + values: Vec::new(), sorted: Vec::new(), } } @@ -59,19 +63,20 @@ impl RunState { ) { self.n += 1; self.sum += x; - // f64::min/max return the non-NaN operand, so these skip NaN - // and read NaN only while every value so far is NaN. - self.min = self.min.min(x); - self.max = self.max.max(x); - // Welford in f64; a NaN poisons mean and m2 for the run's rest. - let delta = x - self.mean; - self.mean += delta / self.n as f64; - self.m2 += delta * (x - self.mean); - if agg == Aggregation::Median { - let pos = self - .sorted - .partition_point(|&v| tot_cmp(v, x) == Ordering::Less); - self.sorted.insert(pos, x); + self.min = min_fold(self.min, x); + self.max = max_fold(self.max, x); + match agg { + Aggregation::Std => self.values.push(x), + Aggregation::Median => { + // Upper bound: insert *after* equal elements, so the + // buffer holds what a stable sort would produce and a + // signed-zero tie keeps its input order, as in polars. + let pos = self + .sorted + .partition_point(|&v| tot_cmp(v, x) != Ordering::Greater); + self.sorted.insert(pos, x); + } + _ => {} } } @@ -87,14 +92,12 @@ impl RunState { Aggregation::Min => Some(self.min), Aggregation::Max => Some(self.max), Aggregation::Delta => Some(self.max - self.min), - Aggregation::Std => { - // Sample std (ddof=1): null below two values, as in polars. - (self.n >= 2).then(|| (self.m2 / (self.n - 1) as f64).sqrt()) - } + // Sample std (ddof=1): null below two values, as in polars. + Aggregation::Std => self.values.std(), Aggregation::Median => { let mid = self.sorted.len() / 2; Some(if self.sorted.len().is_multiple_of(2) { - (self.sorted[mid - 1] + self.sorted[mid]) / 2.0 + midpoint(self.sorted[mid - 1], self.sorted[mid]) } else { self.sorted[mid] }) @@ -110,16 +113,18 @@ impl RunState { /// except `count`, which emits the unchanged running count (as native /// `cum_count` does). fn scan_row( - row_idx: usize, values: &[Option], gates: &[Option], agg: Aggregation, outside: Outside, ) -> PolarsResult>> { + // No row index: plugins are handed one chunk at a time, so any + // position we could name here is chunk-local and would point at a + // different, innocent row of the frame. polars_ensure!( values.len() == gates.len(), ComputeError: - "cum_agg_runs: value and gate lists differ in length at row {row_idx} ({} vs {})", + "cum_agg_runs: a row's value and gate lists differ in length ({} vs {})", values.len(), gates.len() ); @@ -214,7 +219,7 @@ fn cum_agg_runs( for i in 0..len { rows.push(match (values.row(i), gates.row(i)) { (Some(v), Some(g)) => { - Some(scan_row(i, v, g, kwargs.aggregation, kwargs.outside)?) + Some(scan_row(v, g, kwargs.aggregation, kwargs.outside)?) } _ => None, }); diff --git a/src/statistic.rs b/src/statistic.rs index 531f7b3..66db9d2 100644 --- a/src/statistic.rs +++ b/src/statistic.rs @@ -34,6 +34,50 @@ impl Aggregation { } } +/// Running minimum that skips NaN and keeps the **first** of values +/// that compare equal, as polars does. `f64::min` would return the +/// later one, which flips the sign of a `-0.0`/`+0.0` tie. +pub(crate) fn min_fold( + acc: f64, + x: f64, +) -> f64 { + if acc.is_nan() { + x + } else if x.is_nan() || acc <= x { + acc + } else { + x + } +} + +/// Running maximum, first-wins on ties. See [`min_fold`]. +pub(crate) fn max_fold( + acc: f64, + x: f64, +) -> f64 { + if acc.is_nan() { + x + } else if x.is_nan() || acc >= x { + acc + } else { + x + } +} + +/// The median of two middle order statistics, as polars computes it. +/// +/// Linear interpolation rather than `(a + b) / 2`, which overflows for +/// large equal values, plus an equality short-circuit, without which +/// equal infinities would read NaN. Together these reproduce polars on +/// every probed input, including `-0.0`/`+0.0` ties (the short-circuit +/// returns the first, so a stable sort carries the sign through). +pub(crate) fn midpoint( + a: f64, + b: f64, +) -> f64 { + if a == b { a } else { a + (b - a) * 0.5 } +} + /// Statistics over slices of data. pub(crate) trait Statistic { fn values(&self) -> &[f64]; @@ -42,15 +86,11 @@ pub(crate) trait Statistic { self.values().is_empty() } - /// Sum with polars' identity element: an empty selection sums to - /// `+0.0` (polars itself is split between `+0.0` and `-0.0` here, - /// and Rust's `Iterator::sum` starts from IEEE's additive identity - /// `-0.0`; the library standardizes on `+0.0`). + /// Sum from the `+0.0` identity, so an empty selection and an + /// all-`-0.0` one both read `+0.0`, as polars does. (Rust's + /// `Iterator::sum` folds from IEEE's `-0.0` identity instead.) fn sum(&self) -> Option { - if self.is_empty() { - return Some(0.0); - } - Some(self.values().iter().sum()) + Some(self.values().iter().fold(0.0, |acc, &v| acc + v)) } fn median(&self) -> Option { @@ -58,10 +98,13 @@ pub(crate) trait Statistic { return None; } let mut values = self.values().to_vec(); + // sort_by is stable, so ties keep their input order and the + // midpoint short-circuit returns the earlier of a signed-zero + // pair -- matching polars. values.sort_by(|a, b| crate::util::tot_cmp(*a, *b)); let mid = values.len() / 2; if values.len().is_multiple_of(2) { - Some((values[mid - 1] + values[mid]) / 2.0) + Some(midpoint(values[mid - 1], values[mid])) } else { Some(values[mid]) } @@ -91,21 +134,21 @@ pub(crate) trait Statistic { Some(variance.sqrt()) } - // The NaN seed makes IEEE min/max skip NaN values unless every - // value is NaN, matching polars. + // The NaN seed makes the folds skip NaN values unless every value + // is NaN, matching polars. fn min(&self) -> Option { if self.is_empty() { return None; } - Some(self.values().iter().cloned().fold(f64::NAN, f64::min)) + Some(self.values().iter().cloned().fold(f64::NAN, min_fold)) } fn max(&self) -> Option { if self.is_empty() { return None; } - Some(self.values().iter().cloned().fold(f64::NAN, f64::max)) + Some(self.values().iter().cloned().fold(f64::NAN, max_fold)) } fn delta(&self) -> Option { diff --git a/src/util.rs b/src/util.rs index 0d46bbb..8200911 100644 --- a/src/util.rs +++ b/src/util.rs @@ -143,15 +143,15 @@ impl ListInputs { /// Apply a fallible transformation elementwise across one or more list /// columns, returning a `List[f64]` column. /// -/// The closure receives the row index and one `&[f64]` per input column -/// (length-1 literals broadcast) and returns `Ok(None)` to emit a null -/// row, or `Err` to abort the whole query. Null input rows stay null. +/// The closure receives one `&[f64]` per input column (length-1 +/// literals broadcast) and returns `Ok(None)` to emit a null row, or +/// `Err` to abort the whole query. Null input rows stay null. pub(crate) fn apply_list_transform( inputs: &[Series], f: F, ) -> PolarsResult where - F: Fn(usize, &[&[f64]]) -> PolarsResult>>, + F: Fn(&[&[f64]]) -> PolarsResult>>, { let li = ListInputs::new(inputs)?; @@ -163,7 +163,7 @@ where ); for i in 0..li.len() { - match li.row(i).map(|slices| f(i, &slices)).transpose()?.flatten() { + match li.row(i).map(|slices| f(&slices)).transpose()?.flatten() { Some(result) => builder.append_slice(&result), None => builder.append_null(), } diff --git a/src/zip_binary.rs b/src/zip_binary.rs index fd63395..c37c214 100644 --- a/src/zip_binary.rs +++ b/src/zip_binary.rs @@ -104,7 +104,6 @@ fn kleene( /// Zip one row's element pairs through `f`; a null row on either side /// stays null, mismatched lengths raise, empty lists zip to empty. fn zip_row( - row_idx: usize, left: Option<&[Option]>, right: Option<&[Option]>, f: impl Fn(Option, Option) -> Option, @@ -112,10 +111,13 @@ fn zip_row( let (Some(left), Some(right)) = (left, right) else { return Ok(None); }; + // No row index: plugins are handed one chunk at a time, so any + // position we could name here is chunk-local and would point at a + // different, innocent row of the frame. polars_ensure!( left.len() == right.len(), ComputeError: - "zip_binary: left and right lists differ in length at row {row_idx} ({} vs {})", + "zip_binary: a row's left and right lists differ in length ({} vs {})", left.len(), right.len() ); Ok(Some( @@ -217,7 +219,7 @@ fn zip_binary( TypedListInput::::boolean(&inputs[1], "zip_binary: right_column")?; let len = broadcast_len(&[left.n_rows(), right.n_rows()])?; let rows = (0..len) - .map(|i| zip_row(i, left.row(i), right.row(i), |a, b| kleene(op, a, b))) + .map(|i| zip_row(left.row(i), right.row(i), |a, b| kleene(op, a, b))) .collect::>>()?; return Ok(build_bool_list(rows)); } @@ -228,7 +230,7 @@ fn zip_binary( if op.is_arithmetic() { let rows = (0..len) - .map(|i| zip_row(i, left.row(i), right.row(i), |a, b| arithmetic(op, a, b))) + .map(|i| zip_row(left.row(i), right.row(i), |a, b| arithmetic(op, a, b))) .collect::>>()?; build_float_list( rows, @@ -236,7 +238,7 @@ fn zip_binary( ) } else { let rows = (0..len) - .map(|i| zip_row(i, left.row(i), right.row(i), |a, b| comparison(op, a, b))) + .map(|i| zip_row(left.row(i), right.row(i), |a, b| comparison(op, a, b))) .collect::>>()?; Ok(build_bool_list(rows)) } diff --git a/tests/test_cum_agg_runs.py b/tests/test_cum_agg_runs.py index 8890a29..93c2966 100644 --- a/tests/test_cum_agg_runs.py +++ b/tests/test_cum_agg_runs.py @@ -3,6 +3,7 @@ prefix's non-null values -- run over an adversarial pattern matrix.""" import math +import numpy as np import polars as pl import polars.exceptions import pytest @@ -13,6 +14,8 @@ AGGS = ["sum", "mean", "median", "std", "min", "max", "delta", "count"] +INF = float("inf") + VALUE_PATTERNS = [ [], [1.0], @@ -24,7 +27,22 @@ [None, None], [NAN, NAN], [0.0, -0.0], + [-0.0, 0.0], + [-0.0, -0.0], [1.0, -1.0, 0.0, None, NAN, 3.0, 2.5, -0.5], + # Infinities and overflow magnitudes: legitimate floats per spec + # 6.1, and the regime where a naive midpoint or an incremental + # variance parts company with the vertical kernels. + [INF, INF], + [-INF, -INF], + [1.0, -INF], + [-INF, 3.0], + [-INF, INF], + [1e308, 1e308], + [-1e308, 1e308], + [1e308, 1e308, 1.0, 2.0], + [1e16, 1e16 + 2.0], + [INF, 1.0, NAN, -INF, None, 2.0], ] @@ -44,6 +62,29 @@ def gate_patterns(n: int) -> list[list[bool | None]]: return pats +def two_pass_std(prefix: list[float]) -> float | None: + """Sample std (ddof=1) the way the library's vertical kernel computes + it: mean first, then the squared deviations. + + The oracle cannot use polars here. polars' own std is neither + two-pass nor Welford nor naive — it disagrees with all three once + its intermediate accumulation overflows (`[1e308, 1e308]` reads NaN) + — and no algorithm reproduces it bit-for-bit anyway, because it + accumulates in lanes. The library's contract is therefore that the + scan equals the *library's* vertical kernel exactly (asserted here), + which in turn tracks polars on ordinary data (asserted by + test_std_tracks_polars_on_ordinary_data). + """ + n = len(prefix) + if n < 2: + return None + mean = sum(prefix) / n + # d * d, not d ** 2: Python's power operator raises OverflowError + # where IEEE multiplication saturates to inf, as the kernel does. + variance = sum((v - mean) * (v - mean) for v in prefix) / (n - 1) + return math.sqrt(variance) if variance == variance and variance >= 0 else NAN + + def vertical(agg: str, prefix: list[float]) -> float | int | None: """The vertical aggregation over a run prefix's non-null values, with polars itself as the kernel (the deference doctrine).""" @@ -52,6 +93,8 @@ def vertical(agg: str, prefix: list[float]) -> float | int | None: return s.sum() if agg == "count": return len(prefix) # prefix holds only non-null values + if agg == "std": + return two_pass_std(prefix) if agg == "delta": mx, mn = s.max(), s.min() return None if mx is None else mx - mn # ty: ignore[unsupported-operator] @@ -82,11 +125,16 @@ def scan_oracle(agg, values, gates, outside) -> list: def assert_row_matches(got: list, expected: list, ctx: str) -> None: assert len(got) == len(expected), ctx for j, (g, e) in enumerate(zip(got, expected)): - where = f"{ctx}, position {j}: got {g}, expected {e}" + where = f"{ctx}, position {j}: got {g!r}, expected {e!r}" if e is None: assert g is None, where elif isinstance(e, float) and math.isnan(e): assert g is not None and math.isnan(g), where + elif e == 0.0: + # Pin the sign of a zero too: min/max keep the first of a + # signed-zero tie, and sum folds from +0.0, both of which a + # magnitude-only comparison cannot see. + assert g == 0.0 and math.copysign(1.0, g) == math.copysign(1.0, e), where else: assert g == pytest.approx(e), where @@ -131,6 +179,108 @@ def test_all_true_gate_matches_native_cum(agg, native): assert_row_matches(got, expected, f"{agg} vs {native}") +def test_std_tracks_polars_on_ordinary_data(): + # The other half of the std contract: the library's kernel agrees + # with polars' vertical std to floating-point rounding whenever + # polars' own accumulation stays in range. + rng = np.random.default_rng(3) + data = list(rng.standard_normal(60) * 1000 + 1e6) + got = pl.DataFrame({"v": [data]}, schema={"v": pl.List(pl.Float64)}).select( + polist.cum_agg_runs( + "v", pl.lit([True] * len(data)), aggregation="std" + ).alias("r") + )["r"][0].to_list() + for k in range(2, len(data) + 1): + assert got[k - 1] == pytest.approx( + pl.Series(data[:k], dtype=pl.Float64).std(), rel=1e-12 + ), k + + +def test_std_is_the_same_kernel_as_agg_slices(): + # One input, one answer: the scan's final position must equal the + # library's own vertical aggregation over the whole run, including + # in the overflow regime where polars' std parts company with every + # textbook algorithm. + for values in [[1e308, 1e308], [1e16, 1e16 + 2.0], [1.0, 2.0, 3.0], [1e-8, 2e-8]]: + idx = [float(i) for i in range(len(values))] + df = pl.DataFrame( + {"v": [values], "i": [idx]}, + schema={"v": pl.List(pl.Float64), "i": pl.List(pl.Float64)}, + ) + scan = df.select( + polist.cum_agg_runs( + "v", pl.lit([True] * len(values)), aggregation="std" + ).alias("r") + )["r"][0].to_list()[-1] + flat = df.select( + polist.agg_slices("v", "i", aggregation="std").alias("r") + )["r"][0] + assert scan == flat, values + + +def test_std_does_not_collapse_to_zero_for_distinct_values(): + # An incremental variance whose running mean rounds onto the next + # value reports zero spread for data that has spread; at timestamp + # magnitudes that is an ordinary input, not an exotic one. + values = [1e16, 1e16 + 2.0] + got = pl.DataFrame({"v": [values]}, schema={"v": pl.List(pl.Float64)}).select( + polist.cum_agg_runs("v", pl.lit([True, True]), aggregation="std").alias("r") + )["r"][0].to_list() + assert got[1] == pytest.approx(pl.Series(values, dtype=pl.Float64).std()) + assert got[1] > 0.0 + + +def test_median_matches_polars_at_extremes(): + # A naive (a + b) / 2 midpoint overflows for large equal values and + # loses the equality short-circuit that equal infinities need. + for values in [ + [1e308, 1e308], [-1e308, -1e308], [1.0, -INF], [-INF, 3.0], + [INF, INF], [-INF, -INF], [-1e308, 1e308], [1.0, INF], + ]: + got = pl.DataFrame({"v": [values]}, schema={"v": pl.List(pl.Float64)}).select( + polist.cum_agg_runs( + "v", pl.lit([True] * len(values)), aggregation="median" + ).alias("r") + )["r"][0].to_list() + expected = [ + pl.Series(values[:k], dtype=pl.Float64).median() + for k in range(1, len(values) + 1) + ] + assert_row_matches(got, expected, f"median {values}") + + +@pytest.mark.parametrize("values", [[-0.0, 0.0], [0.0, -0.0], [-0.0, -0.0]]) +def test_signed_zero_ties_keep_the_first_value(values): + # polars' min/max keep the earlier of two values that compare equal; + # Rust's f64::min/max keep the later, which flips the sign of zero. + df = pl.DataFrame({"v": [values]}, schema={"v": pl.List(pl.Float64)}) + for agg in ["min", "max", "median"]: + got = df.select( + polist.cum_agg_runs( + "v", pl.lit([True] * len(values)), aggregation=agg # ty: ignore[invalid-argument-type] + ).alias("r") + )["r"][0].to_list() + expected = [ + getattr(pl.Series(values[:k], dtype=pl.Float64), agg)() + for k in range(1, len(values) + 1) + ] + assert_row_matches(got, expected, f"{agg} {values}") + + +def test_leading_infinity_also_deviates_from_native_cum_seed_leak(): + # The native seed leak is not NaN-specific: cum_min/cum_max fold + # from -+f64::MAX and never replace the seed on a tie, so a leading + # infinity of the same sign leaks it too. The prefix rule gives the + # true infinity instead. + df = pl.DataFrame({"v": [[INF, 1.0]]}, schema={"v": pl.List(pl.Float64)}) + got = df.select( + polist.cum_agg_runs("v", pl.lit([True, True]), aggregation="min").alias("r") + )["r"][0].to_list() + native = pl.Series([INF, 1.0], dtype=pl.Float64).cum_min().to_list() + assert got[0] == INF + assert native[0] == pytest.approx(1.7976931348623157e308) + + def test_leading_nan_follows_prefix_rule_not_polars_seed_leak(): # Native cum_max([NaN, 1.0]) emits -1.8e308 -- f64::MIN, its fold # seed -- at the NaN position: a sentinel leak, not any aggregation @@ -285,6 +435,66 @@ def test_gate_literal_broadcasts(): assert out["r"].to_list() == [[1.0, 3.0], [3.0, 7.0]] +def test_value_literal_broadcasts(): + # The other side of the broadcasting axis: one constant signal + # scanned against per-row gates. + df = pl.DataFrame( + {"g": [[True, True], [True, False]]}, schema={"g": pl.List(pl.Boolean)} + ) + out = df.select( + polist.cum_agg_runs(pl.lit([1.0, 2.0]), "g", aggregation="sum").alias("r") + ) + assert out["r"].to_list() == [[1.0, 3.0], [1.0, None]] + + +def test_run_state_resets_after_a_nan_poisoned_run(): + # A NaN poisons sum/mean/std for the rest of its run; a gate break + # must start the next run clean rather than carrying it over. + df = pl.DataFrame( + {"v": [[1.0, NAN, 2.0, 5.0, 7.0]]}, schema={"v": pl.List(pl.Float64)} + ).with_columns(pl.lit([True, True, True, False, True]).alias("g")) + sums = df.select( + polist.cum_agg_runs("v", "g", aggregation="sum").alias("r") + )["r"][0].to_list() + assert sums[0] == 1.0 + assert math.isnan(sums[1]) and math.isnan(sums[2]) # poisoned + assert sums[3] is None # outside + assert sums[4] == 7.0 # fresh run, uncontaminated + + means = df.select( + polist.cum_agg_runs("v", "g", aggregation="mean").alias("r") + )["r"][0].to_list() + assert math.isnan(means[2]) and means[4] == 7.0 + + +def test_array_input_raises_cleanly(): + # PR 0 exists because an Array input used to abort the process; the + # new plugins must also produce a catchable typed error. + df = pl.DataFrame( + {"v": [[1.0, 2.0]], "g": [[True, True]]}, + schema={"v": pl.Array(pl.Float64, 2), "g": pl.List(pl.Boolean)}, + ) + with pytest.raises(polars.exceptions.PolarsError, match="List column"): + df.select(polist.cum_agg_runs("v", "g", aggregation="sum")) + + +def test_float32_accumulates_in_f64_and_rounds_once(): + # Doctrine 4: the kernel computes in f64 and rounds at the exit + # boundary, while polars' Float32 cum_sum rounds every step. The + # deviation is deliberate, so pin it rather than let it drift. + # 3e38 + 3e38 overflows Float32 (max ~3.4e38) but not Float64. + values = [3e38, 3e38, -3e38] + df = pl.DataFrame({"v": [values]}, schema={"v": pl.List(pl.Float32)}) + got = df.select( + polist.cum_agg_runs("v", pl.lit([True] * 3), aggregation="sum").alias("r") + )["r"][0].to_list() + native = pl.Series(values, dtype=pl.Float32).cum_sum().to_list() + # polars overflows at the second step and never recovers; the f64 + # accumulator stays in range and rounds back to a finite Float32. + assert math.isinf(native[1]) and math.isinf(native[2]) + assert math.isinf(got[1]) and got[2] == pytest.approx(3e38, rel=1e-6) + + @pytest.mark.parametrize("engine", ["in-memory", "streaming"]) def test_lazy_execution_matches_eager(engine): df = pl.DataFrame( diff --git a/tests/test_error_messages.py b/tests/test_error_messages.py new file mode 100644 index 0000000..1908f1f --- /dev/null +++ b/tests/test_error_messages.py @@ -0,0 +1,92 @@ +"""Plugins are handed one chunk at a time, so a position counted inside +the kernel is chunk-local. Naming it as an absolute row number sends the +reader to a different, perfectly valid row -- and which row it names +changes with chunk layout and with the engine. These messages therefore +describe the offending row by its shape and never by its position.""" +import polars as pl +import polars.exceptions +import pytest + +import polars_list_utils as polist + + +def message(fn) -> str: + with pytest.raises(polars.exceptions.PolarsError) as excinfo: + fn() + return str(excinfo.value) + + +def multi_chunk_frame(offending_at: int, n: int = 1000, chunks: int = 4): + """A frame whose only malformed row sits in the last chunk.""" + parts = [] + for c in range(chunks): + rows = [[1.0, 2.0] for _ in range(n)] + gates = [[True, True] for _ in range(n)] + if c == chunks - 1: + rows[offending_at - c * n] = [1.0, 2.0, 3.0] + parts.append( + pl.DataFrame( + {"v": rows, "g": gates}, + schema={"v": pl.List(pl.Float64), "g": pl.List(pl.Boolean)}, + ) + ) + return pl.concat(parts, rechunk=False) + + +def test_cum_agg_runs_length_mismatch_names_no_row(): + df = multi_chunk_frame(offending_at=3500) + assert df.n_chunks() > 1 + msg = message(lambda: df.select(polist.cum_agg_runs("v", "g", aggregation="sum"))) + assert "differ in length (3 vs 2)" in msg + assert "at row" not in msg + + +def test_zip_binary_length_mismatch_names_no_row(): + df = pl.DataFrame( + {"a": [[1.0, 2.0], [3.0]], "b": [[1.0, 2.0], [1.0, 2.0]]}, + schema={"a": pl.List(pl.Float64), "b": pl.List(pl.Float64)}, + ) + msg = message(lambda: df.select(polist.zip_binary("a", "b", op="add"))) + assert "differ in length (1 vs 2)" in msg + assert "at row" not in msg + + +def test_message_is_stable_across_chunk_layout_and_engine(): + # The regression this guards: the reported text used to change with + # chunking and with the engine, for identical data. + df = multi_chunk_frame(offending_at=3500) + chunked = message(lambda: df.select(polist.cum_agg_runs("v", "g", aggregation="sum"))) + rechunked = message( + lambda: df.rechunk().select(polist.cum_agg_runs("v", "g", aggregation="sum")) + ) + streamed = message( + lambda: df.lazy() + .select(polist.cum_agg_runs("v", "g", aggregation="sum")) + .collect(engine="streaming") + ) + assert chunked == rechunked + assert "differ in length (3 vs 2)" in streamed + + +def test_agg_lists_strict_message_names_no_row(): + df = pl.DataFrame({"g": [1, 1], "v": [[1.0, 2.0], [1.0, 2.0, 3.0]]}) + msg = message( + lambda: df.group_by("g").agg( + polist.agg_lists("v", list_length=2, aggregation="mean", strict=True) + ) + ) + assert "list length 3 but list_length=2" in msg + assert "row 1" not in msg + + +def test_apply_interp_strict_message_names_no_row(): + df = pl.DataFrame( + {"x": [[0.0, 1.0], [2.0, 1.0]], "y": [[0.0, 1.0], [0.0, 1.0]]} + ) + msg = message( + lambda: df.select( + polist.apply_interp("x", "y", pl.lit([0.5]), strict=True) + ) + ) + assert "found 2 followed by 1" in msg + assert "at row" not in msg diff --git a/tests/test_kwargs_serialization.py b/tests/test_kwargs_serialization.py new file mode 100644 index 0000000..b727018 --- /dev/null +++ b/tests/test_kwargs_serialization.py @@ -0,0 +1,137 @@ +"""Kwargs cross the plugin boundary as a pickle, and the decoder cannot +resolve a memoized *container* back-reference: any range object, bound +tuple, or range list that appears twice used to fail the query with +"recursive structure found". CPython's constant folding produces such +aliasing for perfectly ordinary literals, so the wrapper rebuilds every +range into fresh plain-float tuples. These are the shapes that broke.""" +import numpy as np +import polars as pl +import pytest + +import polars_list_utils as polist + +NAN, INF = float("nan"), float("inf") + + +@pytest.fixture +def df_ramp() -> pl.DataFrame: + ramp = [float(v) for v in range(11)] + return pl.DataFrame({"v": [ramp], "i": [ramp]}) + + +def test_identical_bound_tuples_within_one_range(df_ramp): + # A single-point closed range: CPython folds the two identical + # (1.0, "closed") constants into one object. + out = df_ramp.with_columns( + polist.agg_slices( + "v", "i", + aggregation="sum", + slices_include=[((1.0, "closed"), (1.0, "closed"))], + ).alias("r") + ) + assert out["r"][0] == 1.0 + + +def test_bound_tuple_shared_across_two_ranges(df_ramp): + out = df_ramp.with_columns( + polist.agg_slices( + "v", "i", + aggregation="count", + slices_include=[ + ((0.0, "closed"), (2.0, "open")), + ((0.0, "closed"), (5.0, "open")), + ], + ).alias("r") + ) + assert out["r"][0] == 5.0 # indices 0..4 + + +def test_same_range_object_in_include_and_exclude(df_ramp): + r = ((0.0, "closed"), (5.0, "open")) + out = df_ramp.with_columns( + polist.agg_slices( + "v", "i", aggregation="count", slices_include=[r], slices_exclude=[r] + ).alias("r") + ) + assert out["r"][0] == 0.0 + + +def test_same_list_object_for_both_kwargs(df_ramp): + # Aliasing at the list level, not inside a range. + ranges: list[polist.Range] = [((0.0, "closed"), (5.0, "open"))] + out = df_ramp.with_columns( + polist.agg_slices( + "v", "i", + aggregation="count", + slices_include=ranges, + slices_exclude=ranges, + ).alias("r") + ) + assert out["r"][0] == 0.0 + + +def test_duplicate_simple_ranges(df_ramp): + out = df_ramp.with_columns( + polist.agg_slices( + "v", "i", aggregation="count", slices_include=[(0.0, 2.0), (0.0, 2.0)] + ).alias("r") + ) + assert out["r"][0] == 3.0 + + +def test_numpy_scalar_bounds(df_ramp): + # np.float64 pickles as a reduce the decoder cannot resolve, so the + # wrapper must coerce bounds to plain floats. + lo, hi = np.float64(2.0), np.float64(4.0) + out = df_ramp.with_columns( + polist.agg_slices( + "v", "i", aggregation="sum", slices_include=[(lo, hi)] + ).alias("r") + ) + assert out["r"][0] == 9.0 # 2 + 3 + 4 + + +def test_integer_bounds(df_ramp): + out = df_ramp.with_columns( + polist.agg_slices( + "v", "i", aggregation="sum", slices_include=[(2, 4)] + ).alias("r") + ) + assert out["r"][0] == 9.0 + + +def test_list_shaped_ranges(df_ramp): + # The shape you get from json.load. + out = df_ramp.with_columns( + polist.agg_slices( + "v", "i", + aggregation="sum", + slices_include=[[[2.0, "closed"], [4.0, "closed"]]], # ty: ignore[invalid-argument-type] + ).alias("r") + ) + assert out["r"][0] == 9.0 + + +def test_infinite_and_nan_bounds_still_work(df_ramp): + out = df_ramp.with_columns( + polist.agg_slices( + "v", "i", aggregation="count", slices_include=[(-INF, INF)] + ).alias("all"), + polist.agg_slices( + "v", "i", aggregation="count", slices_include=[(NAN, NAN)] + ).alias("none"), + ) + assert out["all"][0] == 11.0 + assert out["none"][0] == 0.0 + + +def test_many_ranges(df_ramp): + ranges: list[polist.Range] = [ + ((float(i), "closed"), (float(i), "closed")) for i in range(300) + ] + out = df_ramp.with_columns( + polist.agg_slices( + "v", "i", aggregation="count", slices_include=ranges + ).alias("r") + ) + assert out["r"][0] == 11.0 diff --git a/tests/test_strict.py b/tests/test_strict.py index 1919381..cbf6536 100644 --- a/tests/test_strict.py +++ b/tests/test_strict.py @@ -1,6 +1,8 @@ """Every strict check gets a positive test (raises under strict=True) and a negative test (the documented silent degenerate result under the default).""" +import math + import polars as pl import polars.exceptions import pytest @@ -31,12 +33,25 @@ def test_interp_default_silently_interpolates_garbage(): assert out["r"][0].to_list() == [20.0, 20.0] -def test_interp_strict_raises_on_nan_x(): +@pytest.mark.parametrize("strict", [False, True]) +def test_interp_nan_x_propagates_rather_than_raising(strict): + # NaN is a legitimate float for the transform family: it flows into + # the output instead of raising. The strict check is for unsorted x, + # which fails silently; a NaN announces itself in the result. df = pl.DataFrame({"x": [[0.0, NAN, 2.0]], "y": [[0.0, 1.0, 2.0]]}) - with pytest.raises(polars.exceptions.PolarsError, match="non-decreasing"): - df.with_columns( - polist.apply_interp("x", "y", pl.lit([0.5]), strict=True).alias("r") - ) + out = df.with_columns( + polist.apply_interp("x", "y", pl.lit([0.5, 1.5]), strict=strict).alias("r") + ) + assert all(math.isnan(v) for v in out["r"][0].to_list()) + + +def test_interp_strict_ignores_lone_nan_x(): + # A length-1 x has no adjacent pair; NaN is not a violation anyway. + df = pl.DataFrame({"x": [[NAN]], "y": [[1.0]]}) + out = df.with_columns( + polist.apply_interp("x", "y", pl.lit([0.5]), strict=True).alias("r") + ) + assert out["r"][0].to_list() == [1.0] def test_interp_strict_accepts_duplicate_x(): @@ -73,6 +88,48 @@ def test_agg_slices_strict_raises_on_nan_bound(): ) +def test_agg_slices_default_nan_bound_is_empty_selection(): + # The negative half: without strict, a NaN bound is a legitimately + # empty selection rather than an error. + df = pl.DataFrame({"v": [[1.0, 2.0]], "i": [[0.0, 1.0]]}) + out = df.with_columns( + polist.agg_slices( + "v", "i", aggregation="mean", slices_include=[(NAN, 1.0)] + ).alias("mean"), + polist.agg_slices( + "v", "i", aggregation="count", slices_include=[(NAN, 1.0)] + ).alias("cnt"), + polist.agg_slices( + "v", "i", aggregation="sum", slices_include=[(NAN, 1.0)] + ).alias("sum"), + ) + assert out["mean"][0] is None + assert out["cnt"][0] == 0.0 + assert out["sum"][0] == 0.0 + + +def test_agg_slices_strict_accepts_list_shaped_ranges(): + # Ranges loaded from JSON arrive as lists, which the plugin accepts; + # strict must validate them, not reject them. + df = pl.DataFrame({"v": [[1.0, 2.0, 3.0]], "i": [[0.0, 1.0, 2.0]]}) + out = df.with_columns( + polist.agg_slices( + "v", "i", + aggregation="sum", + slices_include=[[[0.0, "closed"], [1.0, "closed"]]], # ty: ignore[invalid-argument-type] + strict=True, + ).alias("r") + ) + assert out["r"][0] == 3.0 + with pytest.raises(ValueError, match="lo <= hi"): + polist.agg_slices( + "v", "i", + aggregation="sum", + slices_include=[[[2.0, "closed"], [1.0, "closed"]]], # ty: ignore[invalid-argument-type] + strict=True, + ) + + def test_agg_slices_default_inverted_range_is_empty_selection(): df = pl.DataFrame({"v": [[1.0, 2.0]], "i": [[0.0, 1.0]]}) out = df.with_columns( @@ -108,8 +165,12 @@ def test_agg_lists_default_truncates_silently(): def test_agg_lists_strict_passes_within_length(): # Shorter lists stay legal under strict: they contribute nulls at - # the missing positions, which is data, not a violation. - df = pl.DataFrame({"g": [1, 1], "v": [[1.0, 2.0], [3.0]]}) + # the missing positions, which is data, not a violation. So are null + # rows, which pass through the check untouched. + df = pl.DataFrame( + {"g": [1, 1, 1], "v": [[1.0, 2.0], [3.0], None]}, + schema={"g": pl.Int64, "v": pl.List(pl.Float64)}, + ) out = df.group_by("g").agg( polist.agg_lists("v", list_length=2, aggregation="mean", strict=True).alias("r") ) diff --git a/tests/test_zip_binary.py b/tests/test_zip_binary.py index c918745..5c071df 100644 --- a/tests/test_zip_binary.py +++ b/tests/test_zip_binary.py @@ -1,7 +1,9 @@ """zip_binary: per element the result must be identical to the corresponding scalar polars op -- the deference doctrine, executable.""" import math +import struct +import numpy as np import polars as pl import polars.exceptions import pytest @@ -30,16 +32,27 @@ } +def bits(v: float) -> bytes: + """The IEEE bit pattern, so -0.0 != 0.0 and a NaN's sign is visible.""" + return struct.pack(" None: + # Element-wise ops are exact, so this compares bit patterns rather + # than values: the adversarial pairs below include signed zeros and + # NaNs of both signs, whose whole point a magnitude comparison + # cannot see (-0.0 == 0.0, and every NaN equals every other). assert len(got) == len(expected), ctx for j, (g, e) in enumerate(zip(got, expected)): - where = f"{ctx}, position {j}: got {g}, expected {e}" + where = f"{ctx}, position {j}: got {g!r}, expected {e!r}" if e is None: assert g is None, where elif isinstance(e, float) and math.isnan(e): assert g is not None and math.isnan(g), where + elif isinstance(e, bool): + assert g == e, where else: - assert g == pytest.approx(e), where + assert g is not None and bits(g) == bits(e), where @pytest.mark.parametrize("op", list(FLOAT_OPS)) @@ -82,6 +95,38 @@ def test_kleene_and_or_match_polars(): assert got == expected, op +@pytest.mark.parametrize("op", list(FLOAT_OPS)) +def test_float32_values_match_polars_scalar_op(op): + # The f32 path computes in f64 and rounds once at the exit cast, + # while polars computes natively in f32. For +,-,*,/ double rounding + # is provably innocuous, and this pins that the entry/exit casts do + # not perturb values, NaNs or infinities. + a32 = [float(np.float32(v)) if v is not None else None for v in A] + b32 = [float(np.float32(v)) if v is not None else None for v in B] + flat = pl.DataFrame( + {"a": a32, "b": b32}, schema={"a": pl.Float32, "b": pl.Float32} + ) + expected = flat.select( + FLOAT_OPS[op](pl.col("a"), pl.col("b")).alias("r") + )["r"].to_list() + df = pl.DataFrame( + {"a": [a32], "b": [b32]}, + schema={"a": pl.List(pl.Float32), "b": pl.List(pl.Float32)}, + ) + got = df.select(polist.zip_binary("a", "b", op=op).alias("r"))["r"][0].to_list() + assert len(got) == len(expected), op + for j, (g, e) in enumerate(zip(got, expected)): + where = f"f32 {op}, position {j}: got {g!r}, expected {e!r}" + if e is None: + assert g is None, where + elif isinstance(e, float) and math.isnan(e): + assert g is not None and math.isnan(g), where + elif isinstance(e, bool): + assert g == e, where + else: + assert struct.pack(" Date: Tue, 1 Sep 2026 12:12:42 +0200 Subject: [PATCH 04/17] feat: add empty flag and make strict the default Two conventions exist for the sum of no values. Polars answers 0.0, the identity element of addition, which is what makes summation decompose. For feature extraction that is the dangerous answer: an empty selection usually means a misconfigured range, and a real 0.0 disguises a missing measurement as a quiet one, which then sails through a threshold. Rather than pick one and bury it, agg_slices and agg_lists now take empty="null" (default) or empty="zero", so the choice is visible at the call site and either camp is one argument away. It reaches exactly one cell of the matrix: sum is the only aggregation with a non-null identity, every other one yields null over nothing either way, and count yields 0 either way. cum_agg_runs needs no such flag, since polars' own cum_sum already emits null there and both camps agree. The flag is deliberately not folded into strict, which means one thing -- turn a silent degenerate reading into an error -- and never changes a returned value. Enabling it to catch unsorted x should not silently move your feature values. strict itself now defaults to True, for the same reason the empty default is null: every check it governs covers a failure that is otherwise invisible, and this library's job is feature extraction, not faithful reproduction of numpy's silence. strict=False keeps every lenient reading, and the explode-equivalence carve-out for sum is now written into the suite explicitly rather than left implied. Co-Authored-By: Claude Fable 5 --- README.md | 31 +++++++--- polars_list_utils/__init__.py | 71 +++++++++++++++-------- src/agg_slices.rs | 13 +++-- src/statistic.rs | 36 +++++++++++- tests/test_kwargs_serialization.py | 2 +- tests/test_plugins.py | 92 +++++++++++++++++++++++++----- tests/test_strict.py | 43 +++++++++----- 7 files changed, 222 insertions(+), 66 deletions(-) diff --git a/README.md b/README.md index f9007a9..79d82c3 100644 --- a/README.md +++ b/README.md @@ -85,8 +85,8 @@ and `count`, and both handle missing data like the vertical aggregations of Pola agree to rounding rather than bit-for-bit.) - Nulls are skipped: excluded from both the numerator and the denominator, so `count` - only counts non-null values and an empty or all-null selection yields null (`count`: 0, - `sum`: 0.0 — Polars' identity for sum). + only counts non-null values and an empty or all-null selection yields null, `count` + reading 0 and `sum` following the `empty` setting below. - NaNs propagate: Polars treats NaN as a legitimate float value, so it flows into the sum and poisons `sum`, `mean` and `std`. Two exceptions to be aware of are `min` and `max`, which skip NaNs (unless all values are NaN), and `median` sorts NaN as the largest value. @@ -99,11 +99,28 @@ meaningfully transformed. `zip_binary` and `cum_agg_runs` mirror the correspondi and cumulative Polars operations per element instead, and treat empty lists as valid (empty in, empty out). -Some checks are opt-in via `strict=True` (default off): `apply_interp` verifies that `x` -is non-decreasing, `agg_slices` rejects inverted or NaN range bounds at expression -construction, and `agg_lists` raises on lists longer than `list_length` instead of -silently truncating them. Each check covers a failure that is otherwise silent; NaN, which -announces itself in the output, propagates rather than raising. +### Two conventions, one switch + +Polars answers "what is the sum of no values" with `0.0`, the identity element of addition. +For feature extraction that is the dangerous answer: an empty selection usually means a +misconfigured range, and a real `0.0` disguises that as a measurement which then sails +through a threshold. So `agg_slices` and `agg_lists` take an `empty` argument: + +- `empty="null"` (default) — summing nothing is unknown, not zero. +- `empty="zero"` — Polars' convention. Equivalent to `.fill_null(0.0)` on the default, and + what you want when the result must agree with `explode().group_by().sum()`. + +It reaches exactly one cell of the matrix: `sum` is the only aggregation with a non-null +identity, every other one yields null over nothing either way, and `count` yields 0 either +way. + +The same philosophy sets the `strict` default. `apply_interp` verifies that `x` is +non-decreasing, `agg_slices` rejects inverted or NaN range bounds at expression +construction, and `agg_lists` raises on lists longer than `list_length` rather than +silently truncating them — all **on by default**, since each covers a failure that is +otherwise silent. Pass `strict=False` for the lenient readings (numpy's unsorted-`x` +behaviour, empty selections from inverted ranges, `list_length` as a deliberate window). +NaN, which announces itself in the output, propagates rather than raising either way. Errors raised from inside a row describe the offending row by its shape rather than its position: Polars hands a plugin one chunk at a time, so any row number counted there would diff --git a/polars_list_utils/__init__.py b/polars_list_utils/__init__.py index d270ba7..87da9af 100644 --- a/polars_list_utils/__init__.py +++ b/polars_list_utils/__init__.py @@ -43,6 +43,7 @@ Window = Literal["hann", "hanning", "blackman"] Scaling = Literal["amplitude", "power", "psd"] Aggregation = Literal["sum", "mean", "median", "std", "min", "max", "delta", "count"] +Empty = Literal["null", "zero"] BinaryOp = Literal[ "add", "sub", "mul", "div", "and", "or", "gt", "ge", "lt", "le", "eq", "ne" ] @@ -57,6 +58,7 @@ __all__ = [ "Aggregation", "BinaryOp", + "Empty", "IntoExprColumn", "Range", "RangeBound", @@ -101,7 +103,7 @@ def apply_interp( y_column: IntoExprColumn, xp_column: IntoExprColumn, *, - strict: bool = False, + strict: bool = True, ) -> pl.Expr: """Interpolate `(x, y)` data onto the `xp` coordinates. @@ -109,12 +111,12 @@ def apply_interp( values outside the data range clamped to the first/last `y` value. `x` must be sorted in increasing order. - With `strict=True` a row whose `x` values descend raises - (duplicates stay legal, as in numpy). The default keeps numpy's - silent behaviour: unsorted `x` interpolates garbage without - complaint. A NaN in `x` is a legitimate float either way and - propagates into the output rather than raising — the check exists - for the failure you cannot see, not the one you can. + With `strict=True` (the default) a row whose `x` values descend + raises; duplicates stay legal, as in numpy. Pass `strict=False` for + numpy's silent behaviour, where unsorted `x` interpolates garbage + without complaint. A NaN in `x` is a legitimate float either way + and propagates into the output rather than raising — the check + exists for the failure you cannot see, not the one you can. Returns a `List[f64]` column of interpolated y values, one per `xp` coordinate. Raises if x and y lists differ in length. @@ -235,7 +237,8 @@ def agg_slices( aggregation: Aggregation, slices_include: list[Range] | None = None, slices_exclude: list[Range] | None = None, - strict: bool = False, + empty: Empty = "null", + strict: bool = True, ) -> pl.Expr: """Aggregate values whose paired index falls within the given ranges. @@ -251,20 +254,31 @@ def agg_slices( are skipped (pairwise with their index), while NaN is a legitimate float value — it poisons `sum`/`mean`/`std`, is skipped by `min`/`max` unless all values are NaN, and sorts as the largest - value for `median`. An empty selection yields null (`count`: 0, - `sum`: 0.0). NaN indices never match any range; mismatched list - lengths raise. + value for `median`. NaN indices never match any range; mismatched + list lengths raise. `std` is the sample standard deviation (ddof=1, polars' default) and yields null for selections with fewer than two values. + An empty selection — no element's index fell in range — yields + null, except `count`, which yields 0. `empty` decides the one + aggregation where the two conventions disagree: + + - `"null"` (default): summing nothing is unknown, not zero. An + empty selection usually means a misconfigured range, and a real + `0.0` disguises that as a measurement — the dangerous direction + for a feature feeding a threshold. + - `"zero"`: polars' convention, the identity element of addition. + Equivalent to `.fill_null(0.0)` on the default, and the setting + to use when this must agree with `explode().group_by().sum()`. + Range bounds accept anything `float()` converts, and the explicit form may be given as tuples or lists (so ranges loaded from JSON work as-is). - With `strict=True`, inverted (`lo > hi`) or NaN range bounds raise - at expression construction; the default keeps them as legitimately - empty selections. + With `strict=True` (the default), inverted (`lo > hi`) or NaN range + bounds raise at expression construction. Pass `strict=False` to + keep them as legitimately empty selections. """ return _plugin( "agg_slices", @@ -272,11 +286,14 @@ def agg_slices( aggregation=aggregation, slices_include=_normalize_ranges(slices_include, "slices_include", strict), slices_exclude=_normalize_ranges(slices_exclude, "slices_exclude", strict), + empty=empty, ) _AGGS: dict[str, Callable[[pl.Expr], pl.Expr]] = { - "sum": pl.Expr.sum, + # polars' own sum of nothing is 0.0; `empty="null"` restores the + # "no values, no answer" reading by guarding on the non-null count. + "sum": lambda e: pl.when(e.count() > 0).then(e.sum()), "mean": pl.Expr.mean, "median": pl.Expr.median, "std": pl.Expr.std, @@ -286,13 +303,19 @@ def agg_slices( "count": lambda e: e.count().cast(pl.Float64), } +_AGGS_EMPTY_ZERO: dict[str, Callable[[pl.Expr], pl.Expr]] = { + **_AGGS, + "sum": pl.Expr.sum, +} + def agg_lists( list_column: IntoExprColumn, *, list_length: int, aggregation: Aggregation, - strict: bool = False, + empty: Empty = "null", + strict: bool = True, ) -> pl.Expr: """Aggregate a list-type column elementwise over rows. @@ -302,15 +325,17 @@ def agg_lists( group. Lists shorter than `list_length` contribute nulls at the missing positions. - Lists longer than `list_length` are silently truncated to it — a - deliberate window for some callers. With `strict=True` such rows - raise instead. + With `strict=True` (the default), lists longer than `list_length` + raise. Pass `strict=False` to keep the silent truncation, which + some callers want as a deliberate window. Implemented in pure polars (no plugin), so missing data follows the polars convention exactly as in :func:`agg_slices`: null elements - are skipped (`count` counts non-null values, an all-null position - sums to 0.0), NaN propagates per-kernel, and `std` is the sample - standard deviation (ddof=1). + are skipped (`count` counts non-null values), NaN propagates + per-kernel, and `std` is the sample standard deviation (ddof=1). + `empty` decides what an all-null position sums to — null by + default, `0.0` under `empty="zero"` — exactly as in + :func:`agg_slices`. """ if list_length < 1: raise ValueError(f"list_length must be at least 1, got {list_length}") @@ -319,7 +344,7 @@ def agg_lists( # Pure polars cannot raise from inside an expression, so the # length check rides along as an elementwise pass-through plugin. col = _plugin("check_list_len", [col], list_length=list_length) - agg = _AGGS[aggregation] + agg = (_AGGS if empty == "null" else _AGGS_EMPTY_ZERO)[aggregation] return pl.concat_list( agg(col.list.slice(offset=n, length=1).list.first()) for n in range(list_length) diff --git a/src/agg_slices.rs b/src/agg_slices.rs index e6f58c5..1519600 100644 --- a/src/agg_slices.rs +++ b/src/agg_slices.rs @@ -1,5 +1,5 @@ use crate::slice::{Range, Slice}; -use crate::statistic::Aggregation; +use crate::statistic::{Aggregation, Empty}; use crate::util::ListInputs; use polars::prelude::*; use pyo3_polars::derive::polars_expr; @@ -10,6 +10,7 @@ struct AggSlicesKwargs { aggregation: Aggregation, slices_include: Option>, slices_exclude: Option>, + empty: Empty, } /// Aggregate the values whose paired index falls within the requested @@ -18,8 +19,8 @@ struct AggSlicesKwargs { /// Missing data follows the polars convention: null elements are /// skipped (pairwise with their index), while NaN is a legitimate /// float value that flows into the aggregation. An empty selection -/// yields null (`count`: 0, `sum`: 0.0). NaN indices never match any -/// range. +/// yields null (`count` reading 0, and `sum` following the `empty` +/// setting). NaN indices never match any range. #[polars_expr(output_type=Float64)] fn agg_slices( inputs: &[Series], @@ -48,8 +49,8 @@ fn agg_slices( // Polars missing-data convention: null elements are skipped // (pairwise with their index), NaN values flow through. The - // aggregations turn an empty selection into None (count: 0, - // sum: 0.0). + // aggregations turn an empty selection into None, except + // count, which reads 0. let selected = values .iter() .zip(indices) @@ -58,7 +59,7 @@ fn agg_slices( .map(|(value, _index)| value) .collect::>(); - Ok(kwargs.aggregation.apply(&selected)) + Ok(kwargs.aggregation.apply(&selected, kwargs.empty)) }) .collect::>()?; diff --git a/src/statistic.rs b/src/statistic.rs index 66db9d2..1049001 100644 --- a/src/statistic.rs +++ b/src/statistic.rs @@ -16,12 +16,31 @@ pub(crate) enum Aggregation { Count, } +/// What an empty selection yields for `sum`, the only aggregation with +/// a non-null identity element. +/// +/// Every other aggregation yields null over nothing under either +/// setting, and `count` yields 0 under either, so this switch reaches +/// exactly one cell of the matrix. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Deserialize)] +#[serde(rename_all = "lowercase")] +pub(crate) enum Empty { + /// Summing nothing is unknown, not zero: the safer reading when an + /// empty selection usually means a misconfigured range. + #[default] + Null, + /// Polars' convention: the identity element of addition. + Zero, +} + impl Aggregation { pub(crate) fn apply( &self, values: &[f64], + empty: Empty, ) -> Option { match self { + Self::Sum if values.is_empty() && empty == Empty::Zero => Some(0.0), Self::Sum => values.sum(), Self::Mean => values.mean(), Self::Median => values.median(), @@ -86,10 +105,21 @@ pub(crate) trait Statistic { self.values().is_empty() } - /// Sum from the `+0.0` identity, so an empty selection and an - /// all-`-0.0` one both read `+0.0`, as polars does. (Rust's - /// `Iterator::sum` folds from IEEE's `-0.0` identity instead.) + /// Sum of the selected values, or null when nothing was selected. + /// + /// A deliberate deviation from polars, which returns its identity + /// element `0.0` there: summing nothing is not the same as summing + /// to zero, and reporting a real zero for absent data is the more + /// dangerous of the two errors. `count` keeps polars' `0`, which is + /// the semantically correct answer to "how many values". + /// + /// A non-empty selection folds from `+0.0`, so an all-`-0.0` one + /// reads `+0.0` as polars does; Rust's `Iterator::sum` would fold + /// from IEEE's `-0.0` identity instead. fn sum(&self) -> Option { + if self.is_empty() { + return None; + } Some(self.values().iter().fold(0.0, |acc, &v| acc + v)) } diff --git a/tests/test_kwargs_serialization.py b/tests/test_kwargs_serialization.py index b727018..9cdab66 100644 --- a/tests/test_kwargs_serialization.py +++ b/tests/test_kwargs_serialization.py @@ -118,7 +118,7 @@ def test_infinite_and_nan_bounds_still_work(df_ramp): "v", "i", aggregation="count", slices_include=[(-INF, INF)] ).alias("all"), polist.agg_slices( - "v", "i", aggregation="count", slices_include=[(NAN, NAN)] + "v", "i", aggregation="count", slices_include=[(NAN, NAN)], strict=False ).alias("none"), ) assert out["all"][0] == 11.0 diff --git a/tests/test_plugins.py b/tests/test_plugins.py index 75b987e..a32b657 100644 --- a/tests/test_plugins.py +++ b/tests/test_plugins.py @@ -249,27 +249,62 @@ def test_agg_slices_empty_selection_is_null(df_ramp): assert out["cnt"][0] == 0.0 -def test_agg_slices_empty_selection_sum_is_positive_zero(df_ramp): - # Sum's empty identity is +0.0 (polars itself is split between - # +0.0 for Series.sum and -0.0 for list.sum of an empty list). +def test_agg_slices_empty_selection_sum_follows_the_empty_flag(df_ramp): + # Summing nothing is unknown by default; "zero" restores polars' + # identity element for callers who want the two to agree. out = df_ramp.with_columns( polist.agg_slices( "v", "i", aggregation="sum", slices_include=[(100.0, 200.0)] - ).alias("s") + ).alias("null"), + polist.agg_slices( + "v", "i", aggregation="sum", slices_include=[(100.0, 200.0)], empty="zero" + ).alias("zero"), ) - assert out["s"][0] == 0.0 - assert math.copysign(1.0, out["s"][0]) == 1.0 + assert out["null"][0] is None + assert out["zero"][0] == 0.0 + # polars is itself split between +0.0 (Series.sum) and -0.0 + # (list.sum of an empty list); the library standardizes on +0.0. + assert math.copysign(1.0, out["zero"][0]) == 1.0 -def test_agg_slices_all_null_selection_sums_to_zero(): +def test_agg_slices_all_null_selection_follows_the_empty_flag(): df = pl.DataFrame( {"v": [[None, None]], "i": [[0.0, 1.0]]}, schema={"v": pl.List(pl.Float64), "i": pl.List(pl.Float64)}, ) out = df.with_columns( - polist.agg_slices("v", "i", aggregation="sum").alias("s") + polist.agg_slices("v", "i", aggregation="sum").alias("null"), + polist.agg_slices("v", "i", aggregation="sum", empty="zero").alias("zero"), + ) + assert out["null"][0] is None + assert out["zero"][0] == 0.0 + + +def test_empty_zero_reproduces_the_explode_equivalence(): + # Any list aggregation can be rewritten as explode + group_by + agg. + # That equivalence holds for every aggregation under the default, + # and for `sum` only under empty="zero" -- the one documented + # departure from polars' vertical convention. + data = [1.0, 2.0, 3.0] + df = pl.DataFrame( + {"v": [data], "i": [[0.0, 1.0, 2.0]]}, + schema={"v": pl.List(pl.Float64), "i": pl.List(pl.Float64)}, ) - assert out["s"][0] == 0.0 + empty_range: list[polist.Range] = [(100.0, 200.0)] + exploded = pl.Series([], dtype=pl.Float64) # the same empty selection + + got_zero = df.select( + polist.agg_slices( + "v", "i", aggregation="sum", slices_include=empty_range, empty="zero" + ).alias("r") + )["r"][0] + got_null = df.select( + polist.agg_slices( + "v", "i", aggregation="sum", slices_include=empty_range + ).alias("r") + )["r"][0] + assert got_zero == exploded.sum() # agrees with polars + assert got_null is None # deliberately does not def test_agg_slices_matches_polars_vertical_semantics(): @@ -390,16 +425,21 @@ def test_agg_lists_elementwise_over_groups(): assert out["mean"][1].to_list() == pytest.approx([7.0, 8.0, 9.0]) -def test_agg_lists_sum_of_all_null_position_is_zero(): - # Vertical sum over an all-null position is polars' identity, 0.0. +def test_agg_lists_sum_of_all_null_position_follows_the_empty_flag(): + # The same switch as agg_slices, so both aggregating functions give + # one answer to "what is the sum of no values". df = pl.DataFrame( {"g": [1, 1], "v": [[1.0, None], [2.0, None]]}, schema={"g": pl.Int64, "v": pl.List(pl.Float64)}, ) out = df.group_by("g").agg( - polist.agg_lists("v", list_length=2, aggregation="sum").alias("s") + polist.agg_lists("v", list_length=2, aggregation="sum").alias("null"), + polist.agg_lists( + "v", list_length=2, aggregation="sum", empty="zero" + ).alias("zero"), ) - assert out["s"][0].to_list() == [3.0, 0.0] + assert out["null"][0].to_list() == [3.0, None] + assert out["zero"][0].to_list() == [3.0, 0.0] def test_agg_lists_invalid_length_raises(): @@ -407,6 +447,32 @@ def test_agg_lists_invalid_length_raises(): polist.agg_lists("v", list_length=0, aggregation="mean") +@pytest.mark.parametrize("empty", ["null", "zero"]) +def test_both_aggregators_agree_on_the_sum_of_nothing(empty): + # The two functions must give one answer to "what is the sum of no + # values", under either setting. + horizontal = pl.DataFrame( + {"v": [[None, None]], "i": [[0.0, 1.0]]}, + schema={"v": pl.List(pl.Float64), "i": pl.List(pl.Float64)}, + ).select( + polist.agg_slices("v", "i", aggregation="sum", empty=empty).alias("r") + )["r"][0] + vertical = ( + pl.DataFrame( + {"g": [1, 1], "v": [[None], [None]]}, + schema={"g": pl.Int64, "v": pl.List(pl.Float64)}, + ) + .group_by("g") + .agg( + polist.agg_lists( + "v", list_length=1, aggregation="sum", empty=empty + ).alias("r") + )["r"][0] + .to_list()[0] + ) + assert horizontal == vertical + + def test_agg_lists_and_agg_slices_share_missing_data_semantics(): # The same values with a null and a NaN, aggregated horizontally by # agg_slices and vertically by agg_lists, must agree for every diff --git a/tests/test_strict.py b/tests/test_strict.py index cbf6536..a214d86 100644 --- a/tests/test_strict.py +++ b/tests/test_strict.py @@ -26,9 +26,24 @@ def test_interp_strict_raises_on_unsorted_x(): ) -def test_interp_default_silently_interpolates_garbage(): +def test_strict_is_the_default(): + # Feature-extraction-friendly defaults: the silent failures are + # errors unless the caller opts out. + with pytest.raises(polars.exceptions.PolarsError, match="non-decreasing"): + _interp_frame().with_columns( + polist.apply_interp("x", "y", pl.lit([0.5, 1.5])).alias("r") + ) + with pytest.raises(ValueError, match="lo <= hi"): + polist.agg_slices("v", "i", aggregation="mean", slices_include=[(2.0, 0.0)]) + with pytest.raises(polars.exceptions.PolarsError, match="list_length"): + pl.DataFrame({"g": [1], "v": [[1.0, 2.0, 3.0]]}).group_by("g").agg( + polist.agg_lists("v", list_length=2, aggregation="mean") + ) + + +def test_interp_lenient_silently_interpolates_garbage(): out = _interp_frame().with_columns( - polist.apply_interp("x", "y", pl.lit([0.5, 1.5])).alias("r") + polist.apply_interp("x", "y", pl.lit([0.5, 1.5]), strict=False).alias("r") ) assert out["r"][0].to_list() == [20.0, 20.0] @@ -88,24 +103,24 @@ def test_agg_slices_strict_raises_on_nan_bound(): ) -def test_agg_slices_default_nan_bound_is_empty_selection(): - # The negative half: without strict, a NaN bound is a legitimately +def test_agg_slices_lenient_nan_bound_is_empty_selection(): + # The negative half: with strict off, a NaN bound is a legitimately # empty selection rather than an error. df = pl.DataFrame({"v": [[1.0, 2.0]], "i": [[0.0, 1.0]]}) out = df.with_columns( polist.agg_slices( - "v", "i", aggregation="mean", slices_include=[(NAN, 1.0)] + "v", "i", aggregation="mean", slices_include=[(NAN, 1.0)], strict=False ).alias("mean"), polist.agg_slices( - "v", "i", aggregation="count", slices_include=[(NAN, 1.0)] + "v", "i", aggregation="count", slices_include=[(NAN, 1.0)], strict=False ).alias("cnt"), polist.agg_slices( - "v", "i", aggregation="sum", slices_include=[(NAN, 1.0)] + "v", "i", aggregation="sum", slices_include=[(NAN, 1.0)], strict=False ).alias("sum"), ) assert out["mean"][0] is None assert out["cnt"][0] == 0.0 - assert out["sum"][0] == 0.0 + assert out["sum"][0] is None def test_agg_slices_strict_accepts_list_shaped_ranges(): @@ -130,14 +145,14 @@ def test_agg_slices_strict_accepts_list_shaped_ranges(): ) -def test_agg_slices_default_inverted_range_is_empty_selection(): +def test_agg_slices_lenient_inverted_range_is_empty_selection(): df = pl.DataFrame({"v": [[1.0, 2.0]], "i": [[0.0, 1.0]]}) out = df.with_columns( polist.agg_slices( - "v", "i", aggregation="mean", slices_include=[(2.0, 0.0)] + "v", "i", aggregation="mean", slices_include=[(2.0, 0.0)], strict=False ).alias("mean"), polist.agg_slices( - "v", "i", aggregation="count", slices_include=[(2.0, 0.0)] + "v", "i", aggregation="count", slices_include=[(2.0, 0.0)], strict=False ).alias("cnt"), ) assert out["mean"][0] is None @@ -154,11 +169,13 @@ def test_agg_lists_strict_raises_on_overlong_list(): ) -def test_agg_lists_default_truncates_silently(): +def test_agg_lists_lenient_truncates_silently(): # list_length as a deliberate window: third elements are dropped. df = pl.DataFrame({"g": [1, 1], "v": [[1.0, 2.0, 30.0], [3.0, 4.0, 50.0]]}) out = df.group_by("g").agg( - polist.agg_lists("v", list_length=2, aggregation="mean").alias("r") + polist.agg_lists( + "v", list_length=2, aggregation="mean", strict=False + ).alias("r") ) assert out["r"][0].to_list() == pytest.approx([2.0, 3.0]) From 3882639da9c20fe11cdd9dfd473b9f94fe34ddd5 Mon Sep 17 00:00:00 2001 From: Travis Hammond Date: Tue, 1 Sep 2026 12:48:27 +0200 Subject: [PATCH 05/17] fix: close review gaps in the empty flag and the strict defaults The new docs overclaimed and two of the strict checks were miscalibrated. empty="zero" was documented as equivalent to .fill_null(0.0). It is not, and the difference runs in the unsafe direction: agg_slices yields null both for an empty selection, which the flag governs, and for a null row, which it deliberately does not. fill_null cannot tell the two apart, so following that advice launders a missing measurement into a real zero -- the exact failure the flag exists to prevent. The claim of agreement with explode/group_by is also formulation-dependent, so it now names the formulation, and the test that carried the name performs the round trip instead of asserting 0.0 == 0.0. An empty list row is a null row today rather than an empty selection, so neither empty nor count's zero reaches it. The spec changes that in the breaking pass; until then the docs say so and a test pins it, so the divergence from agg_lists is deliberate rather than discovered. agg_lists selected its aggregation table by testing empty against the string "null" and falling through otherwise, so any typo, or the None a JSON null decodes to, silently chose the zero convention -- picking the dangerous reading precisely when the caller misspelled the safe one. Both aggregating functions now validate aggregation and empty at expression construction and fail the same way. apply_interp's monotonicity check compared adjacent pairs, and every IEEE comparison with NaN is false, so a NaN between two elements hid the descent: [0, 10, NaN, 1, 2] passed and interpolated finite, plausible, wrong numbers. It now compares each coordinate against the last non-NaN one, so NaN stays legal and propagates while the descent around it raises. agg_lists' strict check raised on any row longer than list_length, but truncating a null-padded tail discards nothing, since every aggregation skips nulls. Fixed-capacity buffers and widened ragged data look like that, and the only remedy was disabling the check for the rows that do lose data. It now raises only when something past the window is non-null. Co-Authored-By: Claude Fable 5 --- README.md | 10 ++- polars_list_utils/__init__.py | 46 ++++++++++-- src/agg_lists.rs | 29 +++++--- src/apply_interp.rs | 36 ++++++--- src/statistic.rs | 4 +- tests/test_error_messages.py | 2 +- tests/test_plugins.py | 133 ++++++++++++++++++++++++++++++---- tests/test_strict.py | 38 ++++++++++ 8 files changed, 254 insertions(+), 44 deletions(-) diff --git a/README.md b/README.md index 79d82c3..f4f751a 100644 --- a/README.md +++ b/README.md @@ -107,13 +107,19 @@ misconfigured range, and a real `0.0` disguises that as a measurement which then through a threshold. So `agg_slices` and `agg_lists` take an `empty` argument: - `empty="null"` (default) — summing nothing is unknown, not zero. -- `empty="zero"` — Polars' convention. Equivalent to `.fill_null(0.0)` on the default, and - what you want when the result must agree with `explode().group_by().sum()`. +- `empty="zero"` — Polars' convention, matching + `explode().group_by().agg(col.filter(...).sum())`. It reaches exactly one cell of the matrix: `sum` is the only aggregation with a non-null identity, every other one yields null over nothing either way, and `count` yields 0 either way. +The switch governs an empty *selection* — the lists held elements, but none was selected. +A null row, and a row whose lists are themselves empty, stay null under both settings for +every aggregation. So `empty="zero"` is not the same as `.fill_null(0.0)`: that cannot tell +a missing row from an empty selection, and would turn a missing measurement into a real +zero, which is the error this flag exists to avoid. + The same philosophy sets the `strict` default. `apply_interp` verifies that `x` is non-decreasing, `agg_slices` rejects inverted or NaN range bounds at expression construction, and `agg_lists` raises on lists longer than `list_length` rather than diff --git a/polars_list_utils/__init__.py b/polars_list_utils/__init__.py index 87da9af..ba2d8ac 100644 --- a/polars_list_utils/__init__.py +++ b/polars_list_utils/__init__.py @@ -44,6 +44,9 @@ Scaling = Literal["amplitude", "power", "psd"] Aggregation = Literal["sum", "mean", "median", "std", "min", "max", "delta", "count"] Empty = Literal["null", "zero"] + +_AGGREGATIONS = ("sum", "mean", "median", "std", "min", "max", "delta", "count") +_EMPTY = ("null", "zero") BinaryOp = Literal[ "add", "sub", "mul", "div", "and", "or", "gt", "ge", "lt", "le", "eq", "ne" ] @@ -190,6 +193,28 @@ def apply_fft( ) +def _check_choice( + value: object, + allowed: tuple[str, ...], + param: str, + func: str, +) -> None: + """Reject an out-of-vocabulary keyword at expression construction. + + Both aggregating functions validate here rather than leaving it to + the plugin, so a typo fails the same way in each. It matters most + for `empty`, whose two settings differ in exactly the direction + this library cares about: silently falling through to `"zero"` + would turn "no data" into a real 0.0, the failure the flag exists + to prevent. + """ + if value not in allowed: + raise ValueError( + f"{func}: {param} must be one of {', '.join(map(repr, allowed))}, " + f"got {value!r}" + ) + + def _normalize_ranges( ranges: list[Range] | None, param: str, strict: bool ) -> list[Range] | None: @@ -260,17 +285,22 @@ def agg_slices( `std` is the sample standard deviation (ddof=1, polars' default) and yields null for selections with fewer than two values. - An empty selection — no element's index fell in range — yields - null, except `count`, which yields 0. `empty` decides the one - aggregation where the two conventions disagree: + An empty selection — the lists held elements, but no index fell in + range — yields null, except `count`, which yields 0. `empty` + decides the one aggregation where the two conventions disagree: - `"null"` (default): summing nothing is unknown, not zero. An empty selection usually means a misconfigured range, and a real `0.0` disguises that as a measurement — the dangerous direction for a feature feeding a threshold. - - `"zero"`: polars' convention, the identity element of addition. - Equivalent to `.fill_null(0.0)` on the default, and the setting - to use when this must agree with `explode().group_by().sum()`. + - `"zero"`: polars' convention, the identity element of addition, + matching `explode().group_by().agg(col.filter(...).sum())`. + + `empty` governs the empty *selection* only. A null row, and a row + whose lists are themselves empty, stay null under both settings for + every aggregation including `count` — so `empty="zero"` is not the + same as `.fill_null(0.0)`, which cannot tell a missing row from an + empty selection and would launder the former into a real zero. Range bounds accept anything `float()` converts, and the explicit form may be given as tuples or lists (so ranges loaded from JSON @@ -280,6 +310,8 @@ def agg_slices( bounds raise at expression construction. Pass `strict=False` to keep them as legitimately empty selections. """ + _check_choice(aggregation, _AGGREGATIONS, "aggregation", "agg_slices") + _check_choice(empty, _EMPTY, "empty", "agg_slices") return _plugin( "agg_slices", [value_column, index_column], @@ -339,6 +371,8 @@ def agg_lists( """ if list_length < 1: raise ValueError(f"list_length must be at least 1, got {list_length}") + _check_choice(aggregation, _AGGREGATIONS, "aggregation", "agg_lists") + _check_choice(empty, _EMPTY, "empty", "agg_lists") col = polars_func_arg_into_col_expr(list_column) if strict: # Pure polars cannot raise from inside an expression, so the diff --git a/src/agg_lists.rs b/src/agg_lists.rs index 4a0c4cf..10f9e55 100644 --- a/src/agg_lists.rs +++ b/src/agg_lists.rs @@ -15,9 +15,15 @@ fn passthrough_output(input_fields: &[Field]) -> PolarsResult { Ok(input_fields[0].clone()) } -/// Identity pass-through that raises when any row's list is longer -/// than `list_length` (which `agg_lists` would otherwise silently -/// truncate to). Null rows are data and pass through. +/// Identity pass-through that raises when a row would lose data to +/// `agg_lists`' silent truncation at `list_length`. +/// +/// Being merely longer than the window is not enough: every +/// aggregation skips null elements, so truncating a null-padded tail +/// is provably lossless, and raising on it would be a false positive +/// whose only remedy is turning the check off entirely. Fixed-capacity +/// buffers and ragged data widened to a common length are ordinary +/// inputs here. Null rows are data and pass through. /// /// The message names no row index: the plugin is handed one chunk at a /// time, so any position named here would be chunk-local and point at @@ -28,13 +34,18 @@ fn check_list_len( kwargs: CheckListLenKwargs, ) -> PolarsResult { let s = &inputs[0]; + let width = kwargs.list_length; for row in s.list()?.into_iter().flatten() { - polars_ensure!( - row.len() <= kwargs.list_length, - ComputeError: - "agg_lists: a row has list length {} but list_length={} (strict=true)", - row.len(), kwargs.list_length - ); + if row.len() > width { + let tail = row.slice(width as i64, row.len() - width); + polars_ensure!( + tail.null_count() == tail.len(), + ComputeError: + "agg_lists: a row has list length {} with non-null values past \ + list_length={} (strict=true)", + row.len(), width + ); + } } Ok(s.clone()) } diff --git a/src/apply_interp.rs b/src/apply_interp.rs index 3023e8e..be0c428 100644 --- a/src/apply_interp.rs +++ b/src/apply_interp.rs @@ -30,18 +30,30 @@ fn apply_interp( ComputeError: "apply_interp: x and y lists differ in length ({} vs {})", cols[0].len(), cols[1].len() ); - // Only a strictly descending pair is a violation; a NaN pair is - // incomparable, and NaN propagates per the transform family's - // rule. No row index: the plugin is handed one chunk at a time, - // so any position named here would be chunk-local. - if kwargs.strict - && let Some(w) = cols[0].windows(2).find(|w| w[0] > w[1]) - { - polars_bail!(ComputeError: - "apply_interp: x_column must be non-decreasing with strict=true, \ - found {} followed by {}", - w[0], w[1] - ); + // Compare each coordinate against the last non-NaN one rather + // than its immediate neighbour. NaN stays a legitimate float + // that propagates into the output, but it must not hide a + // descent: every IEEE comparison with NaN is false, so an + // adjacent-pair test reads [0, 10, NaN, 1] as sorted. + // + // No row index in the message: the plugin is handed one chunk + // at a time, so any position named here would be chunk-local. + if kwargs.strict { + let mut last: Option = None; + for &x in cols[0] { + if x.is_nan() { + continue; + } + if let Some(prev) = last + && x < prev + { + polars_bail!(ComputeError: + "apply_interp: x_column must be non-decreasing with \ + strict=true, found {prev} followed by {x}" + ); + } + last = Some(x); + } } Ok(Some(interp_slice( cols[0], diff --git a/src/statistic.rs b/src/statistic.rs index 1049001..7b9d3c3 100644 --- a/src/statistic.rs +++ b/src/statistic.rs @@ -21,7 +21,9 @@ pub(crate) enum Aggregation { /// /// Every other aggregation yields null over nothing under either /// setting, and `count` yields 0 under either, so this switch reaches -/// exactly one cell of the matrix. +/// exactly one cell of the matrix. It applies to a selection that came +/// out empty; a null row or an empty list never reaches this code, +/// having already short-circuited to a null output row. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Deserialize)] #[serde(rename_all = "lowercase")] pub(crate) enum Empty { diff --git a/tests/test_error_messages.py b/tests/test_error_messages.py index 1908f1f..e6580dd 100644 --- a/tests/test_error_messages.py +++ b/tests/test_error_messages.py @@ -75,7 +75,7 @@ def test_agg_lists_strict_message_names_no_row(): polist.agg_lists("v", list_length=2, aggregation="mean", strict=True) ) ) - assert "list length 3 but list_length=2" in msg + assert "list length 3 with non-null values past list_length=2" in msg assert "row 1" not in msg diff --git a/tests/test_plugins.py b/tests/test_plugins.py index a32b657..d138825 100644 --- a/tests/test_plugins.py +++ b/tests/test_plugins.py @@ -280,31 +280,138 @@ def test_agg_slices_all_null_selection_follows_the_empty_flag(): assert out["zero"][0] == 0.0 -def test_empty_zero_reproduces_the_explode_equivalence(): +@pytest.mark.parametrize( + ("lo", "hi", "selects_anything"), [(0.0, 1.0, True), (100.0, 200.0, False)] +) +def test_empty_zero_reproduces_the_explode_round_trip(lo, hi, selects_anything): # Any list aggregation can be rewritten as explode + group_by + agg. - # That equivalence holds for every aggregation under the default, - # and for `sum` only under empty="zero" -- the one documented - # departure from polars' vertical convention. - data = [1.0, 2.0, 3.0] + # This performs that round trip rather than describing it, for both + # a range that selects values and one that selects none -- the + # latter is the cell where empty="zero" earns its keep. df = pl.DataFrame( - {"v": [data], "i": [[0.0, 1.0, 2.0]]}, + {"v": [[1.0, 2.0, 3.0]], "i": [[0.0, 1.0, 2.0]]}, schema={"v": pl.List(pl.Float64), "i": pl.List(pl.Float64)}, ) - empty_range: list[polist.Range] = [(100.0, 200.0)] - exploded = pl.Series([], dtype=pl.Float64) # the same empty selection + rng: list[polist.Range] = [(lo, hi)] + + exploded = ( + df.with_row_index() + .explode(["v", "i"]) + .group_by("index") + .agg( + pl.col("v") + .filter(pl.col("i").is_between(lo, hi)) + .sum() + .alias("r") + ) + )["r"][0] got_zero = df.select( polist.agg_slices( - "v", "i", aggregation="sum", slices_include=empty_range, empty="zero" + "v", "i", aggregation="sum", slices_include=rng, empty="zero" ).alias("r") )["r"][0] got_null = df.select( + polist.agg_slices("v", "i", aggregation="sum", slices_include=rng).alias("r") + )["r"][0] + + assert got_zero == exploded + if selects_anything: + assert got_null == exploded # the two settings agree here + else: + assert got_null is None # and deliberately part company here + + +@pytest.mark.parametrize("agg", ["mean", "median", "std", "min", "max", "delta", "count"]) +@pytest.mark.parametrize("include", [[(0.0, 2.0)], [(100.0, 200.0)]]) +def test_empty_flag_touches_only_sum(agg, include): + # Every aggregation but sum must be byte-identical under both + # settings, whether or not the selection came out empty. + df = pl.DataFrame( + {"v": [[1.0, 2.0, 3.0]], "i": [[0.0, 1.0, 2.0]]}, + schema={"v": pl.List(pl.Float64), "i": pl.List(pl.Float64)}, + ) + out = df.select( + polist.agg_slices( + "v", "i", aggregation=agg, slices_include=include + ).alias("null"), + polist.agg_slices( + "v", "i", aggregation=agg, slices_include=include, empty="zero" + ).alias("zero"), + ) + assert out["null"][0] == out["zero"][0] + + +@pytest.mark.parametrize("empty", ["null", "zero"]) +def test_empty_does_not_change_a_non_empty_sum(empty): + # Guards the `values.is_empty()` half of the kernel's condition: + # empty="zero" must not flatten a selection that has values. + df = pl.DataFrame( + {"v": [[1.0, 2.0, 3.0]], "i": [[0.0, 1.0, 2.0]]}, + schema={"v": pl.List(pl.Float64), "i": pl.List(pl.Float64)}, + ) + out = df.select( polist.agg_slices( - "v", "i", aggregation="sum", slices_include=empty_range + "v", "i", aggregation="sum", slices_include=[(0.0, 2.0)], empty=empty ).alias("r") - )["r"][0] - assert got_zero == exploded.sum() # agrees with polars - assert got_null is None # deliberately does not + ) + assert out["r"][0] == 6.0 + + +def test_empty_zero_is_not_fill_null(): + # A null row and an empty selection both read null by default, but + # they are different facts: empty="zero" resolves only the second, + # while fill_null(0.0) cannot tell them apart and launders a + # missing measurement into a real zero. + df = pl.DataFrame( + {"v": [[1.0, 2.0], None], "i": [[0.0, 1.0], None]}, + schema={"v": pl.List(pl.Float64), "i": pl.List(pl.Float64)}, + ) + rng: list[polist.Range] = [(100.0, 200.0)] + out = df.select( + polist.agg_slices( + "v", "i", aggregation="sum", slices_include=rng, empty="zero" + ).alias("zero"), + polist.agg_slices("v", "i", aggregation="sum", slices_include=rng) + .fill_null(0.0) + .alias("filled"), + ) + assert out["zero"].to_list() == [0.0, None] + assert out["filled"].to_list() == [0.0, 0.0] + + +def test_empty_list_rows_are_null_rows_under_both_settings(): + # Today an empty list is a null row rather than an empty selection, + # so `empty` (and count's 0) cannot reach it. Spec section 7 changes + # this in the breaking pass; pin it until then so the divergence + # from agg_lists is deliberate rather than discovered. + df = pl.DataFrame( + {"v": [[]], "i": [[]]}, + schema={"v": pl.List(pl.Float64), "i": pl.List(pl.Float64)}, + ) + out = df.select( + polist.agg_slices("v", "i", aggregation="sum", empty="zero").alias("sum"), + polist.agg_slices("v", "i", aggregation="count", empty="zero").alias("count"), + ) + assert out["sum"][0] is None + assert out["count"][0] is None + + +@pytest.mark.parametrize("bad", ["nul", "Null", "zeroes", None, 0, ""]) +def test_both_aggregators_reject_a_bogus_empty(bad): + # Falling through to "zero" on a typo would silently pick the very + # reading the flag exists to avoid. + with pytest.raises(ValueError, match="empty must be"): + polist.agg_slices("v", "i", aggregation="sum", empty=bad) + with pytest.raises(ValueError, match="empty must be"): + polist.agg_lists("v", list_length=1, aggregation="sum", empty=bad) + + +def test_both_aggregators_reject_a_bogus_aggregation(): + with pytest.raises(ValueError, match="aggregation must be"): + polist.agg_slices("v", "i", aggregation="total") # ty: ignore[invalid-argument-type] + with pytest.raises(ValueError, match="aggregation must be"): + polist.agg_lists("v", list_length=1, aggregation="total") # ty: ignore[invalid-argument-type] def test_agg_slices_matches_polars_vertical_semantics(): diff --git a/tests/test_strict.py b/tests/test_strict.py index a214d86..e380e8c 100644 --- a/tests/test_strict.py +++ b/tests/test_strict.py @@ -69,6 +69,29 @@ def test_interp_strict_ignores_lone_nan_x(): assert out["r"][0].to_list() == [1.0] +def test_interp_strict_sees_a_descent_masked_by_nan(): + # Every IEEE comparison with NaN is false, so an adjacent-pair test + # reads [0, 10, NaN, 1, 2] as sorted and silently interpolates + # finite, plausible, wrong numbers -- exactly the failure the check + # exists for. + df = pl.DataFrame( + {"x": [[0.0, 10.0, NAN, 1.0, 2.0]], "y": [[0.0, 1.0, 2.0, 3.0, 4.0]]} + ) + with pytest.raises(polars.exceptions.PolarsError, match="non-decreasing"): + df.with_columns( + polist.apply_interp("x", "y", pl.lit([0.5]), strict=True).alias("r") + ) + + +def test_interp_strict_allows_nan_in_genuinely_sorted_x(): + # NaN itself is not a violation; it propagates into the output. + df = pl.DataFrame({"x": [[0.0, NAN, 2.0]], "y": [[0.0, 1.0, 2.0]]}) + out = df.with_columns( + polist.apply_interp("x", "y", pl.lit([0.5]), strict=True).alias("r") + ) + assert math.isnan(out["r"][0].to_list()[0]) + + def test_interp_strict_accepts_duplicate_x(): # Duplicates stay legal, as in numpy. df = pl.DataFrame({"x": [[0.0, 1.0, 1.0, 2.0]], "y": [[0.0, 10.0, 20.0, 30.0]]}) @@ -169,6 +192,21 @@ def test_agg_lists_strict_raises_on_overlong_list(): ) +def test_agg_lists_strict_allows_null_padded_rows(): + # Truncating a null-padded tail discards nothing -- every + # aggregation skips nulls anyway -- so raising would be a false + # positive whose only remedy is disabling the check outright. + # Fixed-capacity buffers and widened ragged data look like this. + df = pl.DataFrame( + {"g": [1, 1], "v": [[1.0, 2.0, None, None], [3.0, 4.0, None, None]]}, + schema={"g": pl.Int64, "v": pl.List(pl.Float64)}, + ) + out = df.group_by("g").agg( + polist.agg_lists("v", list_length=2, aggregation="mean", strict=True).alias("r") + ) + assert out["r"][0].to_list() == pytest.approx([2.0, 3.0]) + + def test_agg_lists_lenient_truncates_silently(): # list_length as a deliberate window: third elements are dropped. df = pl.DataFrame({"g": [1, 1], "v": [[1.0, 2.0, 30.0], [3.0, 4.0, 50.0]]}) From 66f80465bdcef21cc03e0979cdcc3e617f7a330b Mon Sep 17 00:00:00 2001 From: Travis Hammond Date: Tue, 1 Sep 2026 14:39:58 +0200 Subject: [PATCH 06/17] chore: target 2.0.0 and keep the requirements spec out of the repo The spec is a local working document, not a shipped artefact, so it is untracked and ignored rather than committed alongside the code. The version goes straight to 2.0.0: everything on this branch lands as a single breaking release, so there is no intermediate number to occupy. Co-Authored-By: Claude Fable 5 --- .gitignore | 3 +++ Cargo.lock | 2 +- Cargo.toml | 2 +- 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index 36104b3..deea853 100644 --- a/.gitignore +++ b/.gitignore @@ -9,6 +9,9 @@ target # Claude .claude/ +# Design/requirements working documents, kept local +polars_list_utils_spec.md + # These are backup files generated by rustfmt **/*.rs.bk diff --git a/Cargo.lock b/Cargo.lock index ea4b3e7..bb70d95 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1646,7 +1646,7 @@ dependencies = [ [[package]] name = "polars-list-utils" -version = "1.1.0" +version = "2.0.0" dependencies = [ "butterworth", "interp", diff --git a/Cargo.toml b/Cargo.toml index eb8d5cb..ca8389c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "polars-list-utils" -version = "1.1.0" +version = "2.0.0" edition = "2024" [lib] From c911efd7c3513fd09b7741293a4a85c168d352ae Mon Sep 17 00:00:00 2001 From: Travis Hammond Date: Tue, 1 Sep 2026 15:06:58 +0200 Subject: [PATCH 07/17] feat: accept Array columns everywhere, alongside List Kernels only ever see List. An Array is normalized on the way in and rebuilt on the way out, so container handling lives entirely at the entry/exit boundary and every function keeps one kernel. polars casts between the two losslessly -- inner dtype, null rows, inner nulls all survive -- which is what makes the round trip free. The container algebra: length-preserving functions return what they were given; apply_fft narrows Array(w) to Array(w/2 + 1), being the only width-transforming function; apply_interp follows xp, since the query coordinates decide the output's shape. zip_binary and cum_agg_runs return an Array only when both inputs are Arrays, and their widths must agree -- a mismatch the schema already knows about, so it raises before execution rather than per row. A mixed List/Array pair is allowed and gives a List. Two functions sit outside that. agg_slices reduces to a scalar, so it has no container to propagate. agg_lists always returns a List: it reduces many rows to one, which an elementwise plugin cannot do, so it must stay an expression composition -- and an expression is built before any schema is resolved, so the container it would have to choose is not knowable there. It gains a pass-through plugin to normalize Array to List, since the .list namespace rejects Array outright; that plugin also absorbed the strict width check and now validates list_length against an Array's width. Zero-width Arrays are rejected rather than supported. polars-arrow panics when slicing a zero-width FixedSizeList, inside the FFI boundary before any of this code runs, so the case is not ours to support. Every function resolves its output type at plan time, which is what turns that panic into a typed error -- agg_slices now resolves its Float64 through a function rather than declaring it statically, purely so the check runs. Zero-length Lists are unaffected. The three tests that pinned "Array raises" become support tests. The original PR 0 regression is kept honest by asserting that the calls which once aborted the process now complete. Co-Authored-By: Claude Fable 5 --- README.md | 11 ++ polars_list_utils/__init__.py | 39 +++-- src/agg_lists.rs | 78 +++++++--- src/agg_slices.rs | 16 +- src/apply_butterworth.rs | 15 +- src/apply_fft.rs | 18 ++- src/apply_interp.rs | 14 +- src/cum_agg_runs.rs | 49 ++++-- src/util.rs | 191 ++++++++++++++++++----- src/zip_binary.rs | 57 ++++--- tests/test_array_container.py | 279 ++++++++++++++++++++++++++++++++++ tests/test_cum_agg_runs.py | 12 +- tests/test_plugins.py | 20 ++- tests/test_zip_binary.py | 11 +- 14 files changed, 670 insertions(+), 140 deletions(-) create mode 100644 tests/test_array_container.py diff --git a/README.md b/README.md index f4f751a..241fb13 100644 --- a/README.md +++ b/README.md @@ -99,6 +99,17 @@ meaningfully transformed. `zip_binary` and `cum_agg_runs` mirror the correspondi and cumulative Polars operations per element instead, and treat empty lists as valid (empty in, empty out). +### List and Array + +Every function accepts both `List` and `Array` columns. A length-preserving function +returns the container it was given; `apply_fft` narrows an `Array(w)` to `Array(w // 2 + 1)`, +and `apply_interp` follows its `xp` argument, since that decides the output's shape. +`zip_binary` returns an `Array` only when both operands are `Array`s of equal width — a +mismatch is caught before execution — and a mixed pair gives a `List`. Two exceptions: +`agg_slices` reduces to a scalar, and `agg_lists` always returns a `List`, being an +expression composition that never sees its input's schema. Zero-width `Array`s are +rejected, because Polars panics when slicing them; zero-length `List`s are fine. + ### Two conventions, one switch Polars answers "what is the sum of no values" with `0.0`, the identity element of addition. diff --git a/polars_list_utils/__init__.py b/polars_list_utils/__init__.py index ba2d8ac..aaf5bc3 100644 --- a/polars_list_utils/__init__.py +++ b/polars_list_utils/__init__.py @@ -1,10 +1,9 @@ -"""Polars expression plugins for signal processing on List columns. +"""Polars expression plugins for signal processing on list columns. -Six Rust plugins operating elementwise on List columns, plus one -pure-polars helper. The four original functions take `List[f64]` -(anything numeric is cast to it); :func:`zip_binary` and -:func:`cum_agg_runs` instead take float or Boolean inners as their -operands require, and preserve `Float32`: +Six Rust plugins operating elementwise, plus one pure-polars helper. The +four original functions take `List[f64]` (anything numeric is cast to +it); :func:`zip_binary` and :func:`cum_agg_runs` instead take float or +Boolean inners as their operands require, and preserve `Float32`: - :func:`apply_interp`: interpolate (x, y) data onto new x coordinates. - :func:`apply_butterworth`: zero-phase Butterworth filtering. @@ -14,6 +13,18 @@ - :func:`zip_binary`: element-wise binary ops between two list columns. - :func:`cum_agg_runs`: cumulative aggregation within gated runs. +Both `List` and `Array` columns are accepted everywhere, and a +length-preserving function returns the container it was given — +`apply_fft` narrows an `Array(w)` to `Array(w // 2 + 1)`, and +`apply_interp` follows its `xp` argument, since that is what decides the +output's shape. `zip_binary` returns an `Array` only when both operands +are `Array`s of equal width, which the schema settles before execution; +a mixed pair is allowed and gives a `List`. Two exceptions: +:func:`agg_slices` reduces to a scalar, and :func:`agg_lists` always +returns a `List`, being an expression composition that never sees its +input's schema. Zero-width `Array`s are rejected — polars panics when +slicing them — while zero-length `List`s are fine. + The plugins accept length-1 literal list columns (e.g. `pl.lit(...)`) for any input and broadcast them. Null rows produce null output rows. @@ -373,11 +384,17 @@ def agg_lists( raise ValueError(f"list_length must be at least 1, got {list_length}") _check_choice(aggregation, _AGGREGATIONS, "aggregation", "agg_lists") _check_choice(empty, _EMPTY, "empty", "agg_lists") - col = polars_func_arg_into_col_expr(list_column) - if strict: - # Pure polars cannot raise from inside an expression, so the - # length check rides along as an elementwise pass-through plugin. - col = _plugin("check_list_len", [col], list_length=list_length) + # The pass-through plugin normalizes Array to List (the `.list` + # namespace rejects Array, and an expression cannot branch on a dtype + # it has not resolved yet) and, under strict, raises on rows that + # would lose data — which pure polars cannot do from inside an + # expression. + col = _plugin( + "prepare_agg_lists", + [polars_func_arg_into_col_expr(list_column)], + list_length=list_length, + strict=strict, + ) agg = (_AGGS if empty == "null" else _AGGS_EMPTY_ZERO)[aggregation] return pl.concat_list( agg(col.list.slice(offset=n, length=1).list.first()) diff --git a/src/agg_lists.rs b/src/agg_lists.rs index 10f9e55..7a86a8c 100644 --- a/src/agg_lists.rs +++ b/src/agg_lists.rs @@ -1,39 +1,71 @@ -//! `agg_lists` itself is pure polars (see the Python wrapper); this -//! module only holds the plugin backing its `strict` check, because a -//! pure-polars expression has no way to raise at run time. +//! `agg_lists` itself is an expression composition over polars' own +//! vertical aggregations (see the Python wrapper) — it has to be, since +//! it reduces many rows to one and an elementwise plugin cannot. This +//! module holds the one plugin it needs: a pass-through that normalizes +//! the container and, under `strict`, validates the row width. +//! +//! Normalization is not optional. The `.list` namespace rejects `Array` +//! columns outright, so without this the composition would only accept +//! `List`; and the wrapper cannot branch on dtype itself, because an +//! expression is built before any schema is resolved. +use crate::util::{Container, as_list}; use polars::prelude::*; use pyo3_polars::derive::polars_expr; use serde::Deserialize; #[derive(Deserialize)] -struct CheckListLenKwargs { +struct PrepareKwargs { list_length: usize, + strict: bool, } -fn passthrough_output(input_fields: &[Field]) -> PolarsResult { - Ok(input_fields[0].clone()) +fn prepare_output( + input_fields: &[Field], + kwargs: PrepareKwargs, +) -> PolarsResult { + let (container, inner) = + Container::split(input_fields[0].dtype(), "agg_lists: list_column")?; + // An Array's width is in the schema, so a mismatch against + // list_length is decidable before execution (validation, not + // inference: the signature does not fork on container). + if let Container::Array(width) = container { + polars_ensure!( + width == kwargs.list_length, + ComputeError: + "agg_lists: list_length={} does not match the Array width {width}", + kwargs.list_length + ); + } + Ok(Field::new( + PlSmallStr::from(""), + DataType::List(Box::new(inner)), + )) } -/// Identity pass-through that raises when a row would lose data to -/// `agg_lists`' silent truncation at `list_length`. +/// Normalize `list_column` to a `List` column, and under `strict` raise +/// when a row would lose data to `agg_lists`' truncation at +/// `list_length`. /// -/// Being merely longer than the window is not enough: every -/// aggregation skips null elements, so truncating a null-padded tail -/// is provably lossless, and raising on it would be a false positive -/// whose only remedy is turning the check off entirely. Fixed-capacity -/// buffers and ragged data widened to a common length are ordinary -/// inputs here. Null rows are data and pass through. +/// Being merely longer than the window is not enough: every aggregation +/// skips null elements, so truncating a null-padded tail is provably +/// lossless, and raising on it would be a false positive whose only +/// remedy is turning the check off entirely. Fixed-capacity buffers and +/// ragged data widened to a common length are ordinary inputs here. +/// Null rows are data and pass through. /// /// The message names no row index: the plugin is handed one chunk at a -/// time, so any position named here would be chunk-local and point at -/// a different, innocent row of the frame. -#[polars_expr(output_type_func=passthrough_output)] -fn check_list_len( +/// time, so any position named here would be chunk-local and point at a +/// different, innocent row of the frame. +#[polars_expr(output_type_func_with_kwargs=prepare_output)] +fn prepare_agg_lists( inputs: &[Series], - kwargs: CheckListLenKwargs, + kwargs: PrepareKwargs, ) -> PolarsResult { - let s = &inputs[0]; + let s = as_list(&inputs[0])?; + if !kwargs.strict { + return Ok(s); + } let width = kwargs.list_length; for row in s.list()?.into_iter().flatten() { if row.len() > width { @@ -42,10 +74,10 @@ fn check_list_len( tail.null_count() == tail.len(), ComputeError: "agg_lists: a row has list length {} with non-null values past \ - list_length={} (strict=true)", - row.len(), width + list_length={width} (strict=true)", + row.len() ); } } - Ok(s.clone()) + Ok(s) } diff --git a/src/agg_slices.rs b/src/agg_slices.rs index 1519600..ed3f556 100644 --- a/src/agg_slices.rs +++ b/src/agg_slices.rs @@ -1,6 +1,6 @@ use crate::slice::{Range, Slice}; use crate::statistic::{Aggregation, Empty}; -use crate::util::ListInputs; +use crate::util::{Container, ListInputs}; use polars::prelude::*; use pyo3_polars::derive::polars_expr; use serde::Deserialize; @@ -13,6 +13,18 @@ struct AggSlicesKwargs { empty: Empty, } +/// Always a `Float64` scalar — there is no container to propagate — but +/// resolved through a function rather than declared statically, so the +/// inputs are validated at plan time. That is what keeps a zero-width +/// `Array` from reaching the FFI boundary, where polars would panic +/// rather than raise. +fn agg_slices_output(input_fields: &[Field]) -> PolarsResult { + for (field, label) in input_fields.iter().zip(["value_column", "index_column"]) { + Container::split(field.dtype(), &format!("agg_slices: {label}"))?; + } + Ok(Field::new(PlSmallStr::from(""), DataType::Float64)) +} + /// Aggregate the values whose paired index falls within the requested /// index ranges. /// @@ -21,7 +33,7 @@ struct AggSlicesKwargs { /// float value that flows into the aggregation. An empty selection /// yields null (`count` reading 0, and `sum` following the `empty` /// setting). NaN indices never match any range. -#[polars_expr(output_type=Float64)] +#[polars_expr(output_type_func=agg_slices_output)] fn agg_slices( inputs: &[Series], kwargs: AggSlicesKwargs, diff --git a/src/apply_butterworth.rs b/src/apply_butterworth.rs index bfafa2d..b79cbe3 100644 --- a/src/apply_butterworth.rs +++ b/src/apply_butterworth.rs @@ -1,9 +1,14 @@ use crate::filter::Filter; -use crate::util::{apply_list_transform, list_f64_output}; +use crate::util::{Container, apply_list_transform, transform_output}; use polars::prelude::*; use pyo3_polars::derive::polars_expr; use serde::Deserialize; +/// Length-preserving, so the output keeps the signal's container. +fn butterworth_output(input_fields: &[Field]) -> PolarsResult { + transform_output(input_fields, 0, "apply_butterworth: signal_column", |w| w) +} + #[derive(Deserialize)] struct ButterworthKwargs { sample_rate: f64, @@ -27,24 +32,26 @@ struct ButterworthKwargs { /// magnitude response, cutoff at -6 dB. Bandpass doubles the design /// order, as in scipy.signal.butter. Rows with no more than 3x the /// effective order samples (the reflection padding) yield null. -#[polars_expr(output_type_func=list_f64_output)] +#[polars_expr(output_type_func=butterworth_output)] fn apply_butterworth( inputs: &[Series], kwargs: ButterworthKwargs, ) -> PolarsResult { + let (out, _) = + Container::split(inputs[0].dtype(), "apply_butterworth: signal_column")?; let filter = match (kwargs.min_freq, kwargs.max_freq) { (Some(lo), Some(hi)) => Filter::Bandpass(lo, hi), (Some(lo), None) => Filter::Highpass(lo), (None, Some(hi)) => Filter::Lowpass(hi), (None, None) => { - return apply_list_transform(inputs, |cols| Ok(Some(cols[0].to_vec()))); + return apply_list_transform(inputs, out, |cols| Ok(Some(cols[0].to_vec()))); } }; let order = kwargs.filter_order.unwrap_or(4); let built = filter.build(kwargs.sample_rate, order)?; let min_samples = filter.min_samples(order); - apply_list_transform(inputs, move |cols| { + apply_list_transform(inputs, out, move |cols| { if cols[0].len() < min_samples { return Ok(None); } diff --git a/src/apply_fft.rs b/src/apply_fft.rs index 05d778a..c490467 100644 --- a/src/apply_fft.rs +++ b/src/apply_fft.rs @@ -1,9 +1,16 @@ use crate::transform::{Scaling, Window, compute_fft}; -use crate::util::{apply_list_transform, list_f64_output}; +use crate::util::{Container, apply_list_transform, transform_output}; use polars::prelude::*; use pyo3_polars::derive::polars_expr; use serde::Deserialize; +/// The only width-transforming function: a one-sided real FFT of `n` +/// samples emits `n/2 + 1` bins, so an `Array(w)` narrows to +/// `Array(w/2 + 1)`. +fn fft_output(input_fields: &[Field]) -> PolarsResult { + transform_output(input_fields, 0, "apply_fft: signal_column", |w| w / 2 + 1) +} + #[derive(Deserialize)] struct FftKwargs { sample_rate: f64, @@ -20,7 +27,7 @@ struct FftKwargs { /// /// Emits N/2 + 1 values from DC to Nyquist. Rows whose signal is not a /// power-of-two length, or contains non-finite values, yield null. -#[polars_expr(output_type_func=list_f64_output)] +#[polars_expr(output_type_func=fft_output)] fn apply_fft( inputs: &[Series], kwargs: FftKwargs, @@ -32,8 +39,13 @@ fn apply_fft( ); let window = kwargs.window.unwrap_or_default(); let scaling = kwargs.scaling.unwrap_or_default(); + let (input, _) = Container::split(inputs[0].dtype(), "apply_fft: signal_column")?; + let out = match input { + Container::List => Container::List, + Container::Array(w) => Container::Array(w / 2 + 1), + }; - apply_list_transform(inputs, |cols| { + apply_list_transform(inputs, out, |cols| { Ok(compute_fft(cols[0], kwargs.sample_rate, window, scaling)) }) } diff --git a/src/apply_interp.rs b/src/apply_interp.rs index be0c428..24a69ba 100644 --- a/src/apply_interp.rs +++ b/src/apply_interp.rs @@ -1,4 +1,4 @@ -use crate::util::{apply_list_transform, list_f64_output}; +use crate::util::{Container, apply_list_transform, transform_output}; use interp::{InterpMode, interp_slice}; use polars::prelude::*; use pyo3_polars::derive::polars_expr; @@ -10,6 +10,13 @@ struct InterpKwargs { strict: bool, } +/// Shape follows `xp`: the output holds one value per query coordinate, +/// so the query column decides both length and container. The `x`/`y` +/// containers are irrelevant. +fn interp_output(input_fields: &[Field]) -> PolarsResult { + transform_output(input_fields, 2, "apply_interp: xp_column", |w| w) +} + /// Interpolate `y` values from `(x, y)` data onto new `xp` coordinates, /// as in numpy.interp; `xp` values outside the data range clamp to the /// first/last `y` value. @@ -19,12 +26,13 @@ struct InterpKwargs { /// silently interpolates garbage, as in numpy. NaN in `x` is a /// legitimate float that propagates into the output either way: the /// check exists for the silent failure, not the visible one. -#[polars_expr(output_type_func=list_f64_output)] +#[polars_expr(output_type_func=interp_output)] fn apply_interp( inputs: &[Series], kwargs: InterpKwargs, ) -> PolarsResult { - apply_list_transform(inputs, |cols| { + let (out, _) = Container::split(inputs[2].dtype(), "apply_interp: xp_column")?; + apply_list_transform(inputs, out, |cols| { polars_ensure!( cols[0].len() == cols[1].len(), ComputeError: "apply_interp: x and y lists differ in length ({} vs {})", diff --git a/src/cum_agg_runs.rs b/src/cum_agg_runs.rs index 58f6a2e..a423799 100644 --- a/src/cum_agg_runs.rs +++ b/src/cum_agg_runs.rs @@ -1,6 +1,6 @@ use crate::statistic::{Aggregation, Statistic, max_fold, midpoint, min_fold}; use crate::util::{ - TypedListInput, broadcast_len, build_float_list, build_u32_list, tot_cmp, + Container, TypedListInput, broadcast_len, build_float_list, build_u32_list, tot_cmp, }; use polars::prelude::*; use pyo3_polars::derive::polars_expr; @@ -169,20 +169,19 @@ fn cum_agg_runs_output( ComputeError: "cum_agg_runs expects exactly 2 inputs, got {}", input_fields.len() ); - let value_inner = match input_fields[0].dtype() { - DataType::List(inner) - if matches!(**inner, DataType::Float32 | DataType::Float64) => - { - *inner.clone() - } - dt => polars_bail!(ComputeError: - "cum_agg_runs: value_column must be a List column with Float32 or Float64 \ - inner dtype, got {dt}"), - }; + let (values, value_inner) = + Container::split(input_fields[0].dtype(), "cum_agg_runs: value_column")?; polars_ensure!( - matches!(input_fields[1].dtype(), DataType::List(inner) if **inner == DataType::Boolean), - ComputeError: "cum_agg_runs: gate_column must be a List column with Boolean inner \ + matches!(value_inner, DataType::Float32 | DataType::Float64), + ComputeError: "cum_agg_runs: value_column must have a Float32 or Float64 inner \ dtype, got {}", + input_fields[0].dtype() + ); + let (gates, gate_inner) = + Container::split(input_fields[1].dtype(), "cum_agg_runs: gate_column")?; + polars_ensure!( + gate_inner == DataType::Boolean, + ComputeError: "cum_agg_runs: gate_column must have a Boolean inner dtype, got {}", input_fields[1].dtype() ); // count emits UInt32, polars' count dtype; everything else keeps @@ -193,10 +192,27 @@ fn cum_agg_runs_output( }; Ok(Field::new( PlSmallStr::from(""), - DataType::List(Box::new(inner)), + run_container(values, gates)?.dtype(inner), )) } +/// The scan is length-preserving, so the output keeps the value +/// column's container. Two `Array` inputs must agree on width — visible +/// in the schema, so it raises at plan time. +fn run_container( + values: Container, + gates: Container, +) -> PolarsResult { + if let (Container::Array(vw), Container::Array(gw)) = (values, gates) { + polars_ensure!( + vw == gw, + ComputeError: + "cum_agg_runs: value and gate Arrays must have equal widths, got {vw} and {gw}" + ); + } + Ok(values) +} + /// Cumulative aggregation within gated runs of a list column. /// /// A null row in either input yields a null row; empty lists scan to @@ -213,6 +229,7 @@ fn cum_agg_runs( ); let values = TypedListInput::::float(&inputs[0], "cum_agg_runs: value_column")?; let gates = TypedListInput::::boolean(&inputs[1], "cum_agg_runs: gate_column")?; + let out = run_container(values.container(), gates.container())?; let len = broadcast_len(&[values.n_rows(), gates.n_rows()])?; let mut rows: Vec>>> = Vec::with_capacity(len); @@ -238,8 +255,8 @@ fn cum_agg_runs( }) }) .collect(); - Ok(build_u32_list(rows)) + build_u32_list(rows, out) } else { - build_float_list(rows, values.inner_dtype()) + build_float_list(rows, values.inner_dtype(), out) } } diff --git a/src/util.rs b/src/util.rs index 8200911..2ecf170 100644 --- a/src/util.rs +++ b/src/util.rs @@ -16,11 +16,102 @@ pub(crate) fn tot_cmp( } } -/// Output dtype for list-returning expressions. -pub(crate) fn list_f64_output(_: &[Field]) -> PolarsResult { +/// The container an input arrived in, so an output can be restored to +/// the same shape. +/// +/// Kernels only ever see `List`: `Array` is normalized on the way in and +/// rebuilt on the way out, which keeps one kernel per function with the +/// container handled entirely at the boundary. polars casts between the +/// two losslessly, preserving inner dtype, null rows, inner nulls and +/// zero widths. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum Container { + List, + Array(usize), +} + +impl Container { + /// Split a `List`/`Array` dtype into its container and inner dtype. + /// + /// Zero-width `Array`s are rejected here rather than supported: + /// polars-arrow panics when slicing a zero-width `FixedSizeList` + /// ("the offset of the new Buffer cannot exceed the existing + /// length"), which happens inside the plugin FFI boundary before any + /// of this code runs. Since every function resolves its output type + /// through this call at plan time, raising here stops the query + /// before execution — a typed error instead of a panic. A + /// zero-length `List` remains perfectly valid; only the fixed-width + /// container is affected. + pub(crate) fn split( + dtype: &DataType, + label: &str, + ) -> PolarsResult<(Self, DataType)> { + match dtype { + DataType::List(inner) => Ok((Self::List, (**inner).clone())), + DataType::Array(_, 0) => polars_bail!( + ComputeError: + "{label}: zero-width Arrays are not supported (polars panics when \ + slicing them); use a List column for zero-length rows" + ), + DataType::Array(inner, width) => Ok((Self::Array(*width), (**inner).clone())), + dt => polars_bail!( + ComputeError: "{label} must be a List or Array column, got {dt}" + ), + } + } + + /// This container holding elements of `inner`. + pub(crate) fn dtype( + self, + inner: DataType, + ) -> DataType { + match self { + Self::List => DataType::List(Box::new(inner)), + Self::Array(w) => DataType::Array(Box::new(inner), w), + } + } +} + +/// Normalize an `Array` column to the equivalent `List` column, so +/// kernels see a single container. `List` columns pass through. +pub(crate) fn as_list(s: &Series) -> PolarsResult { + match s.dtype() { + DataType::Array(inner, _) => s.cast(&DataType::List(inner.clone())), + _ => Ok(s.clone()), + } +} + +/// Restore a computed `List` series to `container`. +pub(crate) fn restore( + s: Series, + container: Container, +) -> PolarsResult { + match container { + Container::List => Ok(s), + Container::Array(width) => { + let (_, inner) = Container::split(s.dtype(), "output")?; + s.cast(&Container::Array(width).dtype(inner)) + } + } +} + +/// Output dtype for the transforms: `Float64` elements in the container +/// `input_fields[from]` arrived in, whose width `width_of` maps (the FFT +/// halves it; everything else preserves it). +pub(crate) fn transform_output( + input_fields: &[Field], + from: usize, + label: &str, + width_of: impl Fn(usize) -> usize, +) -> PolarsResult { + let (container, _) = Container::split(input_fields[from].dtype(), label)?; + let out = match container { + Container::List => Container::List, + Container::Array(w) => Container::Array(width_of(w)), + }; Ok(Field::new( PlSmallStr::from(""), - DataType::List(Box::new(DataType::Float64)), + out.dtype(DataType::Float64), )) } @@ -97,12 +188,14 @@ impl ListInputs { let cas = inputs .iter() .map(|s| { + let (_, inner) = Container::split(s.dtype(), "input")?; polars_ensure!( - matches!(s.dtype(), DataType::List(inner) if inner.is_primitive_numeric()), - ComputeError: "expected List column with numeric inner dtype, got {}", + inner.is_primitive_numeric(), + ComputeError: + "expected List or Array column with numeric inner dtype, got {}", s.dtype() ); - s.cast(&list_f64) + as_list(s)?.cast(&list_f64) }) .collect::>>()?; @@ -145,9 +238,12 @@ impl ListInputs { /// /// The closure receives one `&[f64]` per input column (length-1 /// literals broadcast) and returns `Ok(None)` to emit a null row, or -/// `Err` to abort the whole query. Null input rows stay null. +/// `Err` to abort the whole query. Null input rows stay null. The result +/// is restored to `out_container`, which the caller derives from +/// whichever input governs the output shape. pub(crate) fn apply_list_transform( inputs: &[Series], + out_container: Container, f: F, ) -> PolarsResult where @@ -169,7 +265,7 @@ where } } - Ok(builder.finish().into_series()) + restore(builder.finish().into_series(), out_container) } /// One `List` input for the dtype-preserving entry layer used by @@ -182,24 +278,25 @@ where pub(crate) struct TypedListInput { rows: Vec>>>, inner_dtype: DataType, + container: Container, } impl TypedListInput { - /// Extract a float List column; inner dtype must be `Float32` or - /// `Float64` (anything else raises, ints included). `label` names + /// Extract a float List/Array column; inner dtype must be `Float32` + /// or `Float64` (anything else raises, ints included). `label` names /// the function and parameter for error messages. pub(crate) fn float( s: &Series, label: &str, ) -> PolarsResult { - let inner_dtype = list_inner_dtype(s, label)?; + let (container, inner_dtype) = Container::split(s.dtype(), label)?; polars_ensure!( matches!(inner_dtype, DataType::Float32 | DataType::Float64), ComputeError: - "{label} must be a List column with Float32 or Float64 inner dtype, got {}", + "{label} must have a Float32 or Float64 inner dtype, got {}", s.dtype() ); - let cast = s.cast(&DataType::List(Box::new(DataType::Float64)))?; + let cast = as_list(s)?.cast(&DataType::List(Box::new(DataType::Float64)))?; let mut rows = Vec::with_capacity(cast.len()); for opt in cast.list()?.into_iter() { rows.push(match opt { @@ -207,30 +304,39 @@ impl TypedListInput { Some(row) => Some(row.f64()?.into_iter().collect()), }); } - Ok(Self { rows, inner_dtype }) + Ok(Self { + rows, + inner_dtype, + container, + }) } } impl TypedListInput { - /// Extract a Boolean List column; any other inner dtype raises. + /// Extract a Boolean List/Array column; any other inner dtype raises. pub(crate) fn boolean( s: &Series, label: &str, ) -> PolarsResult { - let inner_dtype = list_inner_dtype(s, label)?; + let (container, inner_dtype) = Container::split(s.dtype(), label)?; polars_ensure!( inner_dtype == DataType::Boolean, - ComputeError: "{label} must be a List column with Boolean inner dtype, got {}", + ComputeError: "{label} must have a Boolean inner dtype, got {}", s.dtype() ); - let mut rows = Vec::with_capacity(s.len()); - for opt in s.list()?.into_iter() { + let normalized = as_list(s)?; + let mut rows = Vec::with_capacity(normalized.len()); + for opt in normalized.list()?.into_iter() { rows.push(match opt { None => None, Some(row) => Some(row.bool()?.into_iter().collect()), }); } - Ok(Self { rows, inner_dtype }) + Ok(Self { + rows, + inner_dtype, + container, + }) } } @@ -244,6 +350,10 @@ impl TypedListInput { &self.inner_dtype } + pub(crate) fn container(&self) -> Container { + self.container + } + /// Row `i` with length-1 literal broadcasting; `None` is a null row. pub(crate) fn row( &self, @@ -253,22 +363,12 @@ impl TypedListInput { } } -/// The inner dtype of a `List` column, or a clear error. -fn list_inner_dtype( - s: &Series, - label: &str, -) -> PolarsResult { - match s.dtype() { - DataType::List(inner) => Ok(*inner.clone()), - dt => polars_bail!(ComputeError: "{label} must be a List column, got {dt}"), - } -} - /// Build a `List` float series from per-row optional elements, casting /// the `f64`-computed values to `inner` once at the exit boundary. pub(crate) fn build_float_list( rows: Vec>>>, inner: &DataType, + container: Container, ) -> PolarsResult { let mut builder = ListPrimitiveChunkedBuilder::::new( PlSmallStr::from(""), @@ -283,15 +383,19 @@ pub(crate) fn build_float_list( } } let s = builder.finish().into_series(); - if *inner == DataType::Float64 { - Ok(s) + let s = if *inner == DataType::Float64 { + s } else { - s.cast(&DataType::List(Box::new(inner.clone()))) - } + s.cast(&DataType::List(Box::new(inner.clone())))? + }; + restore(s, container) } -/// Build a `List(Boolean)` series from per-row optional elements. -pub(crate) fn build_bool_list(rows: Vec>>>) -> Series { +/// Build a Boolean series from per-row optional elements. +pub(crate) fn build_bool_list( + rows: Vec>>>, + container: Container, +) -> PolarsResult { let mut builder = ListBooleanChunkedBuilder::new(PlSmallStr::from(""), rows.len(), rows.len() * 8); for row in rows { @@ -300,12 +404,15 @@ pub(crate) fn build_bool_list(rows: Vec>>>) -> Series { None => builder.append_null(), } } - builder.finish().into_series() + restore(builder.finish().into_series(), container) } -/// Build a `List(UInt32)` series from per-row optional elements -/// (the `count` aggregation's dtype, as in polars). -pub(crate) fn build_u32_list(rows: Vec>>>) -> Series { +/// Build a `UInt32` series from per-row optional elements (the `count` +/// aggregation's dtype, as in polars). +pub(crate) fn build_u32_list( + rows: Vec>>>, + container: Container, +) -> PolarsResult { let mut builder = ListPrimitiveChunkedBuilder::::new( PlSmallStr::from(""), rows.len(), @@ -318,5 +425,5 @@ pub(crate) fn build_u32_list(rows: Vec>>>) -> Series { None => builder.append_null(), } } - builder.finish().into_series() + restore(builder.finish().into_series(), container) } diff --git a/src/zip_binary.rs b/src/zip_binary.rs index c37c214..303d06c 100644 --- a/src/zip_binary.rs +++ b/src/zip_binary.rs @@ -1,5 +1,5 @@ use crate::util::{ - TypedListInput, broadcast_len, build_bool_list, build_float_list, tot_cmp, + Container, TypedListInput, broadcast_len, build_bool_list, build_float_list, tot_cmp, }; use polars::prelude::*; use pyo3_polars::derive::polars_expr; @@ -138,29 +138,47 @@ fn arithmetic_inner( } } -/// Plan-time dtype validation for one operand. +/// Plan-time container and dtype validation for one operand. fn validate_operand( dtype: &DataType, op: Op, label: &str, -) -> PolarsResult { - let DataType::List(inner) = dtype else { - polars_bail!(ComputeError: "{label} must be a List column, got {dtype}"); - }; +) -> PolarsResult<(Container, DataType)> { + let (container, inner) = Container::split(dtype, label)?; if op.takes_booleans() { polars_ensure!( - **inner == DataType::Boolean, - ComputeError: "{label} must be a List column with Boolean inner dtype for \ - 'and'/'or', got {dtype}" + inner == DataType::Boolean, + ComputeError: "{label} must have a Boolean inner dtype for 'and'/'or', \ + got {dtype}" ); } else { polars_ensure!( - matches!(**inner, DataType::Float32 | DataType::Float64), - ComputeError: - "{label} must be a List column with Float32 or Float64 inner dtype, got {dtype}" + matches!(inner, DataType::Float32 | DataType::Float64), + ComputeError: "{label} must have a Float32 or Float64 inner dtype, got {dtype}" ); } - Ok(*inner.clone()) + Ok((container, inner)) +} + +/// The output container: `Array` only when both operands are `Array`, in +/// which case their widths must agree — a mismatch the schema already +/// knows about, so it raises at plan time. A mixed `List`/`Array` pair is +/// allowed and yields `List`, with the length checked per row. +fn zip_container( + left: Container, + right: Container, +) -> PolarsResult { + match (left, right) { + (Container::Array(lw), Container::Array(rw)) => { + polars_ensure!( + lw == rw, + ComputeError: + "zip_binary: Array operands must have equal widths, got {lw} and {rw}" + ); + Ok(Container::Array(lw)) + } + _ => Ok(Container::List), + } } fn zip_binary_output( @@ -172,12 +190,12 @@ fn zip_binary_output( ComputeError: "zip_binary expects exactly 2 inputs, got {}", input_fields.len() ); - let left = validate_operand( + let (lc, left) = validate_operand( input_fields[0].dtype(), kwargs.op, "zip_binary: left_column", )?; - let right = validate_operand( + let (rc, right) = validate_operand( input_fields[1].dtype(), kwargs.op, "zip_binary: right_column", @@ -189,7 +207,7 @@ fn zip_binary_output( }; Ok(Field::new( PlSmallStr::from(""), - DataType::List(Box::new(inner)), + zip_container(lc, rc)?.dtype(inner), )) } @@ -217,15 +235,17 @@ fn zip_binary( TypedListInput::::boolean(&inputs[0], "zip_binary: left_column")?; let right = TypedListInput::::boolean(&inputs[1], "zip_binary: right_column")?; + let out = zip_container(left.container(), right.container())?; let len = broadcast_len(&[left.n_rows(), right.n_rows()])?; let rows = (0..len) .map(|i| zip_row(left.row(i), right.row(i), |a, b| kleene(op, a, b))) .collect::>>()?; - return Ok(build_bool_list(rows)); + return build_bool_list(rows, out); } let left = TypedListInput::::float(&inputs[0], "zip_binary: left_column")?; let right = TypedListInput::::float(&inputs[1], "zip_binary: right_column")?; + let out = zip_container(left.container(), right.container())?; let len = broadcast_len(&[left.n_rows(), right.n_rows()])?; if op.is_arithmetic() { @@ -235,11 +255,12 @@ fn zip_binary( build_float_list( rows, &arithmetic_inner(left.inner_dtype(), right.inner_dtype()), + out, ) } else { let rows = (0..len) .map(|i| zip_row(left.row(i), right.row(i), |a, b| comparison(op, a, b))) .collect::>>()?; - Ok(build_bool_list(rows)) + build_bool_list(rows, out) } } diff --git a/tests/test_array_container.py b/tests/test_array_container.py new file mode 100644 index 0000000..0037a66 --- /dev/null +++ b/tests/test_array_container.py @@ -0,0 +1,279 @@ +"""Array support. + +The kernels only ever see List: Array is normalized on the way in and +rebuilt on the way out. So the oracle for every value question is the +List result, and what needs testing on top is the container algebra -- +which container comes out, and which structural mismatches the schema +can catch before execution.""" +import polars as pl +import polars.exceptions +import pytest +from polars.datatypes import DataTypeClass + +import polars_list_utils as polist + +FS = 200.0 +# A power of two for the FFT, and longer than the order-4 Butterworth +# reflection padding (3 * 4 + 1) so the filter produces a real row. +SIGNAL = [float(i % 5) for i in range(32)] + + +def as_array(df: pl.DataFrame, **cols: DataTypeClass) -> pl.DataFrame: + """The same frame with the named List columns retyped as Arrays, + each at the width of its first row.""" + return df.with_columns( + pl.col(name).cast(pl.Array(inner, len(df[name][0]))) + for name, inner in cols.items() + ) + + +# ------------------------------------------------- values are container-blind + +def test_transforms_agree_between_containers(): + lists = pl.DataFrame({"s": [SIGNAL]}, schema={"s": pl.List(pl.Float64)}) + arrays = as_array(lists, s=pl.Float64) + for expr in [ + polist.apply_fft("s", sample_rate=FS, window="hanning", scaling="amplitude"), + polist.apply_butterworth("s", sample_rate=FS, max_freq=50.0), + ]: + got = arrays.select(expr.alias("r"))["r"][0].to_list() + expected = lists.select(expr.alias("r"))["r"][0].to_list() + assert got == pytest.approx(expected) + + +def test_aggregations_agree_between_containers(): + lists = pl.DataFrame( + {"v": [[1.0, None, 3.0, 4.0]], "i": [[0.0, 1.0, 2.0, 3.0]]}, + schema={"v": pl.List(pl.Float64), "i": pl.List(pl.Float64)}, + ) + arrays = as_array(lists, v=pl.Float64, i=pl.Float64) + for agg in ["sum", "mean", "median", "std", "min", "max", "delta", "count"]: + expr = polist.agg_slices("v", "i", aggregation=agg) # ty: ignore[invalid-argument-type] + assert arrays.select(expr.alias("r"))["r"][0] == pytest.approx( + lists.select(expr.alias("r"))["r"][0] + ), agg + + +def test_elementwise_and_scan_agree_between_containers(): + lists = pl.DataFrame( + {"v": [[1.0, 2.0, 3.0, 4.0]], "g": [[True, True, False, True]]}, + schema={"v": pl.List(pl.Float64), "g": pl.List(pl.Boolean)}, + ) + arrays = as_array(lists, v=pl.Float64, g=pl.Boolean) + for expr in [ + polist.zip_binary("v", "v", op="mul"), + polist.cum_agg_runs("v", "g", aggregation="sum"), + polist.cum_agg_runs("v", "g", aggregation="count"), + ]: + assert ( + arrays.select(expr.alias("r"))["r"][0].to_list() + == lists.select(expr.alias("r"))["r"][0].to_list() + ) + + +# ------------------------------------------------------ container propagation + +def test_length_preserving_transforms_keep_the_container(): + df = as_array( + pl.DataFrame({"s": [SIGNAL]}, schema={"s": pl.List(pl.Float64)}), s=pl.Float64 + ) + out = df.select( + polist.apply_butterworth("s", sample_rate=FS, max_freq=50.0).alias("r") + ) + assert out["r"].dtype == pl.Array(pl.Float64, len(SIGNAL)) + assert len(out["r"][0].to_list()) == len(SIGNAL) + + +def test_fft_halves_the_array_width(): + # The one width-transforming function: n samples -> n/2 + 1 bins. + df = as_array( + pl.DataFrame({"s": [SIGNAL]}, schema={"s": pl.List(pl.Float64)}), s=pl.Float64 + ) + half = len(SIGNAL) // 2 + 1 + out = df.select(polist.apply_fft("s", sample_rate=FS).alias("r")) + assert out["r"].dtype == pl.Array(pl.Float64, half) + assert len(out["r"][0].to_list()) == half + + +def test_interp_container_follows_xp(): + # The output holds one value per query coordinate, so xp decides the + # shape and the x/y containers are irrelevant. + df = pl.DataFrame( + {"x": [[0.0, 1.0, 2.0]], "y": [[0.0, 10.0, 20.0]]}, + schema={"x": pl.List(pl.Float64), "y": pl.List(pl.Float64)}, + ) + xp_array = pl.Series("xp", [[0.5, 1.5]], dtype=pl.Array(pl.Float64, 2)) + out = df.select(polist.apply_interp("x", "y", pl.lit(xp_array)).alias("r")) + assert out["r"].dtype == pl.Array(pl.Float64, 2) + assert out["r"][0].to_list() == pytest.approx([5.0, 15.0]) + + # List xp over Array x/y gives List back. + arrays = as_array(df, x=pl.Float64, y=pl.Float64) + out = arrays.select(polist.apply_interp("x", "y", pl.lit([0.5])).alias("r")) + assert out["r"].dtype == pl.List(pl.Float64) + + +def test_agg_slices_is_scalar_regardless_of_container(): + df = as_array( + pl.DataFrame({"v": [[1.0, 2.0]], "i": [[0.0, 1.0]]}, + schema={"v": pl.List(pl.Float64), "i": pl.List(pl.Float64)}), + v=pl.Float64, i=pl.Float64, + ) + out = df.select(polist.agg_slices("v", "i", aggregation="mean").alias("r")) + assert out["r"].dtype == pl.Float64 + + +def test_agg_lists_returns_list_even_for_array_input(): + # Documented exception to "output container follows input": agg_lists + # is an expression composition, not a plugin -- it reduces rows, which + # an elementwise plugin cannot -- so it never sees the input schema at + # the point where the output container would have to be chosen. + df = pl.DataFrame( + {"g": [1, 1], "a": [[1.0, 2.0], [3.0, 4.0]]}, + schema={"g": pl.Int64, "a": pl.Array(pl.Float64, 2)}, + ) + out = df.group_by("g").agg( + polist.agg_lists("a", list_length=2, aggregation="mean").alias("r") + ) + assert out["r"].dtype == pl.List(pl.Float64) + assert out["r"][0].to_list() == pytest.approx([2.0, 3.0]) + + +@pytest.mark.parametrize("inner", [pl.Float32, pl.Float64]) +def test_inner_dtype_survives_the_container_round_trip(inner): + df = pl.DataFrame( + {"v": [[1.0, 2.0]], "g": [[True, True]]}, + schema={"v": pl.Array(inner, 2), "g": pl.Array(pl.Boolean, 2)}, + ) + assert df.select(polist.zip_binary("v", "v", op="add").alias("r"))[ + "r" + ].dtype == pl.Array(inner, 2) + assert df.select( + polist.cum_agg_runs("v", "g", aggregation="sum").alias("r") + )["r"].dtype == pl.Array(inner, 2) + # count keeps polars' count dtype, in the input's container. + assert df.select( + polist.cum_agg_runs("v", "g", aggregation="count").alias("r") + )["r"].dtype == pl.Array(pl.UInt32, 2) + + +def test_comparison_on_arrays_yields_boolean_array(): + df = pl.DataFrame({"a": [[1.0, 2.0]]}, schema={"a": pl.Array(pl.Float64, 2)}) + out = df.select(polist.zip_binary("a", "a", op="le").alias("r")) + assert out["r"].dtype == pl.Array(pl.Boolean, 2) + + +# --------------------------------------------------------- mixed and mismatched + +def test_mixed_containers_yield_a_list(): + df = pl.DataFrame( + {"a": [[1.0, 2.0]], "b": [[3.0, 4.0]]}, + schema={"a": pl.Array(pl.Float64, 2), "b": pl.List(pl.Float64)}, + ) + out = df.select(polist.zip_binary("a", "b", op="add").alias("r")) + assert out["r"].dtype == pl.List(pl.Float64) + assert out["r"][0].to_list() == [4.0, 6.0] + + +def test_array_width_mismatch_raises_at_plan_time(): + df = pl.LazyFrame( + {"a": [[1.0, 2.0]], "b": [[1.0, 2.0, 3.0]]}, + schema={"a": pl.Array(pl.Float64, 2), "b": pl.Array(pl.Float64, 3)}, + ) + with pytest.raises(polars.exceptions.PolarsError, match="equal widths"): + df.select(polist.zip_binary("a", "b", op="add")).collect_schema() + + +def test_cum_agg_runs_array_width_mismatch_raises_at_plan_time(): + df = pl.LazyFrame( + {"v": [[1.0, 2.0]], "g": [[True, True, True]]}, + schema={"v": pl.Array(pl.Float64, 2), "g": pl.Array(pl.Boolean, 3)}, + ) + with pytest.raises(polars.exceptions.PolarsError, match="equal widths"): + df.select(polist.cum_agg_runs("v", "g", aggregation="sum")).collect_schema() + + +def test_agg_lists_validates_list_length_against_array_width(): + df = pl.DataFrame( + {"g": [1], "a": [[1.0, 2.0]]}, + schema={"g": pl.Int64, "a": pl.Array(pl.Float64, 2)}, + ) + with pytest.raises(polars.exceptions.PolarsError, match="does not match the Array width"): + df.group_by("g").agg(polist.agg_lists("a", list_length=3, aggregation="mean")) + + +def test_integer_arrays_raise_for_the_dtype_strict_functions(): + df = pl.DataFrame({"a": [[1, 2]]}, schema={"a": pl.Array(pl.Int64, 2)}) + with pytest.raises(polars.exceptions.PolarsError, match="Float32 or Float64"): + df.select(polist.zip_binary("a", "a", op="add")) + + +# ------------------------------------------------------------- zero-width edge + +@pytest.mark.parametrize( + "build", + [ + lambda d: d.select(polist.agg_slices("a", "a", aggregation="count")), + lambda d: d.select(polist.zip_binary("a", "a", op="add")), + lambda d: d.select(polist.apply_fft("a", sample_rate=FS)), + ], +) +def test_zero_width_arrays_raise_instead_of_panicking(build): + # polars-arrow panics when slicing a zero-width FixedSizeList, inside + # the FFI boundary before any of this library's code runs. Every + # function resolves its output type at plan time, so the query is + # stopped there -- a typed error, never a panic. A zero-length List + # stays perfectly valid. + df = pl.DataFrame({"a": [[]]}, schema={"a": pl.Array(pl.Float64, 0)}) + with pytest.raises(polars.exceptions.PolarsError) as excinfo: + build(df) + message = str(excinfo.value) + assert "zero-width Arrays are not supported" in message + # "the plugin panicked" is how pyo3-polars reports an unwound panic. + assert "the plugin panicked" not in message + + +def test_zero_length_lists_are_unaffected(): + df = pl.DataFrame({"a": [[]]}, schema={"a": pl.List(pl.Float64)}) + out = df.select(polist.zip_binary("a", "a", op="add").alias("r")) + assert out["r"][0].to_list() == [] + + +# ------------------------------------------------------- nulls and broadcasting + +def test_null_rows_and_inner_nulls_behave_as_for_lists(): + df = pl.DataFrame( + {"a": [[1.0, 2.0], None]}, schema={"a": pl.Array(pl.Float64, 2)} + ) + out = df.select(polist.zip_binary("a", "a", op="add").alias("r")) + assert out["r"].to_list() == [[2.0, 4.0], None] + + inner_null = pl.DataFrame( + {"a": [[1.0, None]], "i": [[0.0, 1.0]]}, + schema={"a": pl.Array(pl.Float64, 2), "i": pl.Array(pl.Float64, 2)}, + ) + assert inner_null.select( + polist.agg_slices("a", "i", aggregation="count").alias("r") + )["r"][0] == 1.0 + + +def test_array_literal_broadcasts(): + df = pl.DataFrame( + {"a": [[1.0, 2.0], [3.0, 4.0]]}, schema={"a": pl.Array(pl.Float64, 2)} + ) + template = pl.Series("t", [[10.0, 20.0]], dtype=pl.Array(pl.Float64, 2)) + out = df.select(polist.zip_binary("a", pl.lit(template), op="add").alias("r")) + assert out["r"].dtype == pl.Array(pl.Float64, 2) + assert out["r"].to_list() == [[11.0, 22.0], [13.0, 24.0]] + + +@pytest.mark.parametrize("engine", ["in-memory", "streaming"]) +def test_arrays_survive_lazy_and_streaming(engine): + df = pl.DataFrame( + {"a": [[float(i), float(i) + 1.0] for i in range(2000)]}, + schema={"a": pl.Array(pl.Float64, 2)}, + ) + expr = polist.zip_binary("a", "a", op="add").alias("r") + eager = df.select(expr)["r"].to_list() + lazy = df.lazy().select(expr).collect(engine=engine)["r"].to_list() # ty: ignore[not-subscriptable] + assert lazy == eager diff --git a/tests/test_cum_agg_runs.py b/tests/test_cum_agg_runs.py index 93c2966..74741b0 100644 --- a/tests/test_cum_agg_runs.py +++ b/tests/test_cum_agg_runs.py @@ -467,15 +467,17 @@ def test_run_state_resets_after_a_nan_poisoned_run(): assert math.isnan(means[2]) and means[4] == 7.0 -def test_array_input_raises_cleanly(): - # PR 0 exists because an Array input used to abort the process; the - # new plugins must also produce a catchable typed error. +def test_array_input_is_supported(): + # Array used to abort the process, then raised; it is now supported. + # A mixed Array value column and List gate is allowed, and the output + # follows the value column. See tests/test_array_container.py. df = pl.DataFrame( {"v": [[1.0, 2.0]], "g": [[True, True]]}, schema={"v": pl.Array(pl.Float64, 2), "g": pl.List(pl.Boolean)}, ) - with pytest.raises(polars.exceptions.PolarsError, match="List column"): - df.select(polist.cum_agg_runs("v", "g", aggregation="sum")) + out = df.select(polist.cum_agg_runs("v", "g", aggregation="sum").alias("r")) + assert out["r"].dtype == pl.Array(pl.Float64, 2) + assert out["r"][0].to_list() == [1.0, 3.0] def test_float32_accumulates_in_f64_and_rounds_once(): diff --git a/tests/test_plugins.py b/tests/test_plugins.py index d138825..4f4624a 100644 --- a/tests/test_plugins.py +++ b/tests/test_plugins.py @@ -628,19 +628,23 @@ def test_integer_lists_are_accepted(): assert df["yp"][0].to_list() == pytest.approx([15.0]) -def test_array_input_raises_instead_of_aborting(): +def test_array_input_is_supported_not_aborting(): # Without the polars `dtype-array` feature, an Array input panicked in - # polars' arrow-FFI import before any plugin check ran -- a non-unwinding - # panic that aborted the whole process. It must be a catchable typed - # error (until Array support proper lands). + # polars' arrow-FFI import before any plugin check ran -- a + # non-unwinding panic that aborted the whole process. It is now a + # supported container; see tests/test_array_container.py for the + # container algebra. This keeps the original regression honest: the + # calls that used to abort must complete. df = pl.DataFrame( {"v": [[1.0, 2.0], [3.0, 4.0]]}, schema={"v": pl.Array(pl.Float64, 2)}, ) - with pytest.raises(polars.exceptions.PolarsError, match="expected List column"): - df.with_columns(polist.agg_slices("v", "v", aggregation="mean").alias("a")) - with pytest.raises(polars.exceptions.PolarsError, match="expected List column"): - df.with_columns(polist.apply_fft("v", sample_rate=FS).alias("a")) + out = df.with_columns( + polist.agg_slices("v", "v", aggregation="mean").alias("a"), + polist.apply_fft("v", sample_rate=FS).alias("f"), + ) + assert out["a"].to_list() == pytest.approx([1.5, 3.5]) + assert out["f"].dtype == pl.Array(pl.Float64, 2) def test_broadcasting_literal_over_rows(): diff --git a/tests/test_zip_binary.py b/tests/test_zip_binary.py index 5c071df..4a3842d 100644 --- a/tests/test_zip_binary.py +++ b/tests/test_zip_binary.py @@ -258,15 +258,16 @@ def test_empty_lists_zip_to_empty_on_every_path(op, dtype): assert out["r"][0].to_list() == [] -def test_array_input_raises_cleanly(): - # PR 0 exists because an Array input used to abort the process; the - # new plugins must also produce a catchable typed error. +def test_array_input_is_supported(): + # Array used to abort the process, then raised; it is now supported. + # The container algebra lives in tests/test_array_container.py. df = pl.DataFrame( {"a": [[1.0, 2.0]], "b": [[3.0, 4.0]]}, schema={"a": pl.Array(pl.Float64, 2), "b": pl.Array(pl.Float64, 2)}, ) - with pytest.raises(polars.exceptions.PolarsError, match="List column"): - df.select(polist.zip_binary("a", "b", op="add")) + out = df.select(polist.zip_binary("a", "b", op="add").alias("r")) + assert out["r"].dtype == pl.Array(pl.Float64, 2) + assert out["r"][0].to_list() == [4.0, 6.0] @pytest.mark.parametrize("engine", ["in-memory", "streaming"]) From ae7b81ce5293b0367de51280bb35d7b6a61975ae Mon Sep 17 00:00:00 2001 From: Travis Hammond Date: Tue, 1 Sep 2026 15:48:28 +0200 Subject: [PATCH 08/17] fix: close the remaining panic routes into the FFI boundary Adversarial review of the Array pass found three ways to reach a polars-arrow panic that the plan-time guard was meant to have closed. apply_interp resolved its output from xp alone, so nothing inspected x or y. A zero-width Array in either of those positions sailed past the guard and panicked, while the identical column in the xp slot raised cleanly -- the failure mode depended purely on argument position. Every function must resolve every container input, not only the one that governs its output shape, and the tests now cover all eleven positions rather than three functions. The guard also matched only the outermost dimension, so Array(Float64, (2, 0)) -- which is Array(Array(Float64, 0), 2), outer width an innocent 2 -- walked straight through to the same panic. It now rejects a zero width at any nesting depth. The third route was not about Arrays at all. polars gates the narrow integer dtypes behind cargo features, and reconstructs the Series across the FFI boundary before any plugin code runs, so List(Int8) panicked there exactly as Array once did -- and no plan-time check can help, because the panic precedes our code. The features are now enabled, which is the same one-line remedy the Array abort took. Two structural checks the spec requires were also missing: apply_interp's x/y pair and agg_slices' value/index pair are decidable from the schema when both sides are Arrays, but only zip_binary and cum_agg_runs were enforcing theirs. All four now share one helper, so the asymmetry cannot come back quietly. Docs corrected where the Array pass had outdated them: agg_lists no longer claims to be pure polars, and now states that an Array input must match list_length exactly rather than being null-padded like a short List. The mixed-container rules are documented as the deliberate asymmetry they are -- zip_binary needs both operands to be Arrays, while cum_agg_runs follows its value column, its gate being a mask rather than an operand. Co-Authored-By: Claude Fable 5 --- Cargo.lock | 2 + Cargo.toml | 13 +++- polars_list_utils/__init__.py | 65 ++++++++++------ src/agg_slices.rs | 14 +++- src/apply_interp.rs | 12 ++- src/cum_agg_runs.rs | 18 +++-- src/util.rs | 37 +++++++-- src/zip_binary.rs | 23 +++--- tests/test_array_container.py | 140 ++++++++++++++++++++++++++++++++++ 9 files changed, 268 insertions(+), 56 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index bb70d95..8c055b7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1412,9 +1412,11 @@ dependencies = [ "polars-compute", "polars-core", "polars-error", + "polars-io", "polars-lazy", "polars-ops", "polars-plan", + "polars-time", "polars-utils", "version_check", ] diff --git a/Cargo.toml b/Cargo.toml index ca8389c..79cc1d1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,7 +10,18 @@ crate-type = ["cdylib"] [dependencies] pyo3 = { version = "0.26", features = ["extension-module", "abi3-py310"] } pyo3-polars = { version = "0.25", features = ["derive"] } -polars = { version = "0.52", default-features = false, features = ["dtype-array"] } +# Every dtype a caller can put in a list column must be compiled in: +# polars reconstructs the Series across the plugin's arrow-FFI boundary +# before any of our code runs, and a dtype whose feature is missing +# panics there rather than erroring (the same root cause as the Array +# process-abort these features were first enabled for). +polars = { version = "0.52", default-features = false, features = [ + "dtype-array", + "dtype-i8", + "dtype-i16", + "dtype-u8", + "dtype-u16", +] } serde = { version = "1", features = ["derive"] } realfft = "3.4.0" butterworth = "0.1.0" diff --git a/polars_list_utils/__init__.py b/polars_list_utils/__init__.py index aaf5bc3..52361ae 100644 --- a/polars_list_utils/__init__.py +++ b/polars_list_utils/__init__.py @@ -1,9 +1,11 @@ """Polars expression plugins for signal processing on list columns. -Six Rust plugins operating elementwise, plus one pure-polars helper. The -four original functions take `List[f64]` (anything numeric is cast to -it); :func:`zip_binary` and :func:`cum_agg_runs` instead take float or -Boolean inners as their operands require, and preserve `Float32`: +Six elementwise Rust plugins, plus :func:`agg_lists`, which composes +polars' own vertical aggregations (with a small pass-through plugin to +prepare its input). The four original functions take `List[f64]` +(anything numeric is cast to it); :func:`zip_binary` and +:func:`cum_agg_runs` instead take float or Boolean inners as their +operands require, and preserve `Float32`: - :func:`apply_interp`: interpolate (x, y) data onto new x coordinates. - :func:`apply_butterworth`: zero-phase Butterworth filtering. @@ -18,11 +20,13 @@ `apply_fft` narrows an `Array(w)` to `Array(w // 2 + 1)`, and `apply_interp` follows its `xp` argument, since that is what decides the output's shape. `zip_binary` returns an `Array` only when both operands -are `Array`s of equal width, which the schema settles before execution; -a mixed pair is allowed and gives a `List`. Two exceptions: -:func:`agg_slices` reduces to a scalar, and :func:`agg_lists` always -returns a `List`, being an expression composition that never sees its -input's schema. Zero-width `Array`s are rejected — polars panics when +are `Array`s, since both are operands; `cum_agg_runs` follows its value +column alone, its gate being a mask rather than an operand. Paired +inputs that must agree element-for-element have their widths checked +before execution when both are `Array`s, and per row otherwise. Two +exceptions to container propagation: :func:`agg_slices` reduces to a +scalar, and :func:`agg_lists` always returns a `List`, being an +expression composition that never sees its input's schema. Zero-width `Array`s are rejected — polars panics when slicing them — while zero-length `List`s are fine. The plugins accept length-1 literal list columns (e.g. `pl.lit(...)`) @@ -364,21 +368,34 @@ def agg_lists( For use in a group_by/aggregation context: element n of the result is `aggregation` applied vertically to element n of every list in - the group, giving one `List[f64]` of length `list_length` per - group. Lists shorter than `list_length` contribute nulls at the - missing positions. - - With `strict=True` (the default), lists longer than `list_length` - raise. Pass `strict=False` to keep the silent truncation, which - some callers want as a deliberate window. - - Implemented in pure polars (no plugin), so missing data follows the - polars convention exactly as in :func:`agg_slices`: null elements - are skipped (`count` counts non-null values), NaN propagates - per-kernel, and `std` is the sample standard deviation (ddof=1). - `empty` decides what an all-null position sums to — null by - default, `0.0` under `empty="zero"` — exactly as in - :func:`agg_slices`. + the group, giving one `List` of length `list_length` per group. + The result is always a `List`, even for `Array` input: this is an + expression composition rather than a plugin, because it reduces many + rows to one, and an expression is built before any schema is + resolved — so the container to return is not knowable here. + + `List` rows shorter than `list_length` contribute nulls at the + missing positions. An `Array` input must instead match `list_length` + exactly, since its width is fixed for every row and a mismatch is + therefore a configuration error rather than data; it raises, under + `strict` or not. + + With `strict=True` (the default), `List` rows carrying non-null data + past `list_length` raise. Pass `strict=False` to keep the silent + truncation, which some callers want as a deliberate window. A + null-padded tail is never a violation: every aggregation skips + nulls, so truncating it loses nothing. + + The aggregations are polars' own, so missing data follows the polars + convention exactly as in :func:`agg_slices`: null elements are + skipped (`count` counts non-null values), NaN propagates per-kernel, + and `std` is the sample standard deviation (ddof=1). `empty` decides + what an all-null position sums to — null by default, `0.0` under + `empty="zero"` — exactly as in :func:`agg_slices`. + + A small pass-through plugin runs first, to normalize the container + and carry the checks above; polars' `.list` namespace rejects + `Array` outright, and an expression cannot raise from inside itself. """ if list_length < 1: raise ValueError(f"list_length must be at least 1, got {list_length}") diff --git a/src/agg_slices.rs b/src/agg_slices.rs index ed3f556..6832d2b 100644 --- a/src/agg_slices.rs +++ b/src/agg_slices.rs @@ -1,6 +1,6 @@ use crate::slice::{Range, Slice}; use crate::statistic::{Aggregation, Empty}; -use crate::util::{Container, ListInputs}; +use crate::util::{Container, ListInputs, ensure_equal_widths}; use polars::prelude::*; use pyo3_polars::derive::polars_expr; use serde::Deserialize; @@ -19,9 +19,15 @@ struct AggSlicesKwargs { /// `Array` from reaching the FFI boundary, where polars would panic /// rather than raise. fn agg_slices_output(input_fields: &[Field]) -> PolarsResult { - for (field, label) in input_fields.iter().zip(["value_column", "index_column"]) { - Container::split(field.dtype(), &format!("agg_slices: {label}"))?; - } + let (values, _) = + Container::split(input_fields[0].dtype(), "agg_slices: value_column")?; + let (indices, _) = + Container::split(input_fields[1].dtype(), "agg_slices: index_column")?; + ensure_equal_widths( + values, + indices, + "agg_slices: value and index Arrays must have equal widths", + )?; Ok(Field::new(PlSmallStr::from(""), DataType::Float64)) } diff --git a/src/apply_interp.rs b/src/apply_interp.rs index 24a69ba..3184125 100644 --- a/src/apply_interp.rs +++ b/src/apply_interp.rs @@ -1,4 +1,6 @@ -use crate::util::{Container, apply_list_transform, transform_output}; +use crate::util::{ + Container, apply_list_transform, ensure_equal_widths, transform_output, +}; use interp::{InterpMode, interp_slice}; use polars::prelude::*; use pyo3_polars::derive::polars_expr; @@ -12,8 +14,14 @@ struct InterpKwargs { /// Shape follows `xp`: the output holds one value per query coordinate, /// so the query column decides both length and container. The `x`/`y` -/// containers are irrelevant. +/// containers do not reach the output — but they are still split here, +/// because that is the only thing standing between a zero-width `Array` +/// and the polars-arrow panic, and because an `x`/`y` width mismatch is +/// decidable from the schema. fn interp_output(input_fields: &[Field]) -> PolarsResult { + let (x, _) = Container::split(input_fields[0].dtype(), "apply_interp: x_column")?; + let (y, _) = Container::split(input_fields[1].dtype(), "apply_interp: y_column")?; + ensure_equal_widths(x, y, "apply_interp: x and y Arrays must have equal widths")?; transform_output(input_fields, 2, "apply_interp: xp_column", |w| w) } diff --git a/src/cum_agg_runs.rs b/src/cum_agg_runs.rs index a423799..8f8e1ee 100644 --- a/src/cum_agg_runs.rs +++ b/src/cum_agg_runs.rs @@ -1,6 +1,7 @@ use crate::statistic::{Aggregation, Statistic, max_fold, midpoint, min_fold}; use crate::util::{ - Container, TypedListInput, broadcast_len, build_float_list, build_u32_list, tot_cmp, + Container, TypedListInput, broadcast_len, build_float_list, build_u32_list, + ensure_equal_widths, tot_cmp, }; use polars::prelude::*; use pyo3_polars::derive::polars_expr; @@ -203,13 +204,14 @@ fn run_container( values: Container, gates: Container, ) -> PolarsResult { - if let (Container::Array(vw), Container::Array(gw)) = (values, gates) { - polars_ensure!( - vw == gw, - ComputeError: - "cum_agg_runs: value and gate Arrays must have equal widths, got {vw} and {gw}" - ); - } + ensure_equal_widths( + values, + gates, + "cum_agg_runs: value and gate Arrays must have equal widths", + )?; + // The gate is a mask, not a symmetric operand, so the value column + // alone decides the output shape -- unlike zip_binary, where both + // sides are operands and an Array result needs both to be Arrays. Ok(values) } diff --git a/src/util.rs b/src/util.rs index 2ecf170..1fddad8 100644 --- a/src/util.rs +++ b/src/util.rs @@ -46,13 +46,16 @@ impl Container { dtype: &DataType, label: &str, ) -> PolarsResult<(Self, DataType)> { + // Nested too: `Array(Float64, (2, 0))` is `Array(Array(Float64, + // 0), 2)`, whose outer width is a perfectly innocent 2. + polars_ensure!( + !has_zero_width_array(dtype), + ComputeError: + "{label}: zero-width Arrays are not supported (polars panics when \ + slicing them); use a List column for zero-length rows, got {dtype}" + ); match dtype { DataType::List(inner) => Ok((Self::List, (**inner).clone())), - DataType::Array(_, 0) => polars_bail!( - ComputeError: - "{label}: zero-width Arrays are not supported (polars panics when \ - slicing them); use a List column for zero-length rows" - ), DataType::Array(inner, width) => Ok((Self::Array(*width), (**inner).clone())), dt => polars_bail!( ComputeError: "{label} must be a List or Array column, got {dt}" @@ -72,6 +75,30 @@ impl Container { } } +/// Whether a zero-width `Array` appears at any nesting depth. +fn has_zero_width_array(dtype: &DataType) -> bool { + match dtype { + DataType::Array(_, 0) => true, + DataType::Array(inner, _) | DataType::List(inner) => has_zero_width_array(inner), + _ => false, + } +} + +/// Equal-width validation for a pair of inputs that must agree +/// element-for-element. Two `Array`s carry their widths in the schema, +/// so a mismatch raises before execution; any other pairing is checked +/// per row inside the kernel. +pub(crate) fn ensure_equal_widths( + left: Container, + right: Container, + message: &str, +) -> PolarsResult<()> { + if let (Container::Array(lw), Container::Array(rw)) = (left, right) { + polars_ensure!(lw == rw, ComputeError: "{message}, got {lw} and {rw}"); + } + Ok(()) +} + /// Normalize an `Array` column to the equivalent `List` column, so /// kernels see a single container. `List` columns pass through. pub(crate) fn as_list(s: &Series) -> PolarsResult { diff --git a/src/zip_binary.rs b/src/zip_binary.rs index 303d06c..d621016 100644 --- a/src/zip_binary.rs +++ b/src/zip_binary.rs @@ -1,5 +1,6 @@ use crate::util::{ - Container, TypedListInput, broadcast_len, build_bool_list, build_float_list, tot_cmp, + Container, TypedListInput, broadcast_len, build_bool_list, build_float_list, + ensure_equal_widths, tot_cmp, }; use polars::prelude::*; use pyo3_polars::derive::polars_expr; @@ -168,17 +169,15 @@ fn zip_container( left: Container, right: Container, ) -> PolarsResult { - match (left, right) { - (Container::Array(lw), Container::Array(rw)) => { - polars_ensure!( - lw == rw, - ComputeError: - "zip_binary: Array operands must have equal widths, got {lw} and {rw}" - ); - Ok(Container::Array(lw)) - } - _ => Ok(Container::List), - } + ensure_equal_widths( + left, + right, + "zip_binary: Array operands must have equal widths", + )?; + Ok(match (left, right) { + (Container::Array(w), Container::Array(_)) => Container::Array(w), + _ => Container::List, + }) } fn zip_binary_output( diff --git a/tests/test_array_container.py b/tests/test_array_container.py index 0037a66..2e63026 100644 --- a/tests/test_array_container.py +++ b/tests/test_array_container.py @@ -210,6 +210,146 @@ def test_integer_arrays_raise_for_the_dtype_strict_functions(): # ------------------------------------------------------------- zero-width edge +ZERO_WIDTH_POSITIONS = { + # Every position that accepts a container needs its own case: the + # guard lives in each function's plan-time output resolution, so a + # position left out of that resolution is a position that panics. + "interp_x": lambda z, ok: pl.DataFrame( + {"z": z, "a": ok} + ).select(polist.apply_interp("z", "a", "a")), + "interp_y": lambda z, ok: pl.DataFrame( + {"z": z, "a": ok} + ).select(polist.apply_interp("a", "z", "a")), + "interp_xp": lambda z, ok: pl.DataFrame( + {"z": z, "a": ok} + ).select(polist.apply_interp("a", "a", "z")), + "fft": lambda z, ok: pl.DataFrame({"z": z}).select( + polist.apply_fft("z", sample_rate=FS) + ), + "butterworth": lambda z, ok: pl.DataFrame({"z": z}).select( + polist.apply_butterworth("z", sample_rate=FS, max_freq=50.0) + ), + "agg_slices_value": lambda z, ok: pl.DataFrame({"z": z, "a": ok}).select( + polist.agg_slices("z", "a", aggregation="count") + ), + "agg_slices_index": lambda z, ok: pl.DataFrame({"z": z, "a": ok}).select( + polist.agg_slices("a", "z", aggregation="count") + ), + "zip_left": lambda z, ok: pl.DataFrame({"z": z, "a": ok}).select( + polist.zip_binary("z", "a", op="add") + ), + "zip_right": lambda z, ok: pl.DataFrame({"z": z, "a": ok}).select( + polist.zip_binary("a", "z", op="add") + ), + "cum_value": lambda z, ok: pl.DataFrame( + {"z": z, "g": pl.Series("g", [[True, True]], dtype=pl.Array(pl.Boolean, 2))} + ).select(polist.cum_agg_runs("z", "g", aggregation="sum")), + "cum_gate": lambda z, ok: pl.DataFrame({"a": ok, "z": z}).select( + polist.cum_agg_runs("a", "z", aggregation="sum") + ), + "agg_lists": lambda z, ok: pl.DataFrame({"g": [1], "z": z}) + .group_by("g") + .agg(polist.agg_lists("z", list_length=1, aggregation="mean")), +} + + +@pytest.mark.parametrize("position", list(ZERO_WIDTH_POSITIONS)) +def test_zero_width_arrays_raise_in_every_position(position): + z = pl.Series("z", [[]], dtype=pl.Array(pl.Float64, 0)) + ok = pl.Series("a", [[1.0, 2.0]], dtype=pl.Array(pl.Float64, 2)) + with pytest.raises(polars.exceptions.PolarsError) as excinfo: + ZERO_WIDTH_POSITIONS[position](z, ok) + message = str(excinfo.value) + assert "zero-width Arrays are not supported" in message, message + assert "the plugin panicked" not in message, message + + +@pytest.mark.parametrize( + "dtype", + [ + pl.Array(pl.Float64, (2, 0)), # Array(Array(Float64, 0), 2) + pl.List(pl.Array(pl.Float64, 0)), + ], +) +def test_nested_zero_width_arrays_raise(dtype): + # The outer width is innocent, so a top-level-only check waves these + # through and the inner zero width panics inside polars-arrow. + df = pl.DataFrame({"a": pl.Series("a", [[[], []]], dtype=dtype)}) + with pytest.raises(polars.exceptions.PolarsError) as excinfo: + df.select(polist.apply_fft("a", sample_rate=FS)) + message = str(excinfo.value) + assert "zero-width Arrays are not supported" in message, message + assert "the plugin panicked" not in message, message + + +@pytest.mark.parametrize( + ("build", "match"), + [ + ( + lambda: polist.apply_interp("x", "y", pl.lit([0.5])), + "x and y Arrays must have equal widths", + ), + ( + lambda: polist.agg_slices("x", "y", aggregation="mean"), + "value and index Arrays must have equal widths", + ), + ], +) +def test_paired_array_widths_are_checked_at_plan_time(build, match): + # All four paired equal-length requirements resolve from the schema + # when both sides are Arrays, so none of them waits for rows. + lf = pl.LazyFrame( + {"x": [[1.0, 2.0, 3.0]], "y": [[1.0, 2.0]]}, + schema={"x": pl.Array(pl.Float64, 3), "y": pl.Array(pl.Float64, 2)}, + ) + with pytest.raises(polars.exceptions.PolarsError, match=match): + lf.select(build()).collect_schema() + + +def test_mixed_container_length_mismatch_is_a_runtime_error(): + # A List against an Array cannot be settled from the schema, so the + # per-row check has to catch it. + df = pl.DataFrame( + {"a": [[1.0, 2.0]], "b": [[1.0, 2.0, 3.0]]}, + schema={"a": pl.Array(pl.Float64, 2), "b": pl.List(pl.Float64)}, + ) + with pytest.raises(polars.exceptions.PolarsError, match="differ in length"): + df.select(polist.zip_binary("a", "b", op="add")) + + +def test_cum_agg_runs_follows_the_value_column_not_both(): + # Deliberately unlike zip_binary: the gate is a mask, not an operand, + # so a List gate does not demote an Array value column. + array_value = pl.DataFrame( + {"v": [[1.0, 2.0]], "g": [[True, True]]}, + schema={"v": pl.Array(pl.Float64, 2), "g": pl.List(pl.Boolean)}, + ) + assert array_value.select( + polist.cum_agg_runs("v", "g", aggregation="sum").alias("r") + )["r"].dtype == pl.Array(pl.Float64, 2) + + list_value = pl.DataFrame( + {"v": [[1.0, 2.0]], "g": [[True, True]]}, + schema={"v": pl.List(pl.Float64), "g": pl.Array(pl.Boolean, 2)}, + ) + assert list_value.select( + polist.cum_agg_runs("v", "g", aggregation="sum").alias("r") + )["r"].dtype == pl.List(pl.Float64) + + +@pytest.mark.parametrize( + "inner", [pl.Int8, pl.Int16, pl.UInt8, pl.UInt16, pl.Int32, pl.Int64] +) +def test_narrow_integer_inners_do_not_panic(inner): + # polars gates the narrow integer dtypes behind cargo features, and + # reconstructs the Series across the FFI boundary before any of this + # library's code runs -- so a missing feature panics there, exactly + # as the Array dtype once did. + df = pl.DataFrame({"a": [[1, 2, 3, 4]]}, schema={"a": pl.List(inner)}) + out = df.select(polist.agg_slices("a", "a", aggregation="sum").alias("r")) + assert out["r"][0] == 10.0 + + @pytest.mark.parametrize( "build", [ From 732569349bcd6803fc1e7f08a18340e1ed4e4e48 Mon Sep 17 00:00:00 2001 From: Travis Hammond Date: Tue, 1 Sep 2026 16:00:49 +0200 Subject: [PATCH 09/17] fix: compile polars with dtype-full so no input dtype can panic Probing the whole dtype space turned up five more inners that panic rather than error: Categorical, Enum, Decimal, Struct and Int128. The cause is the one PR 0 identified for Array -- polars reconstructs the Series across the plugin's arrow-FFI boundary before any of our code runs, so a dtype whose cargo feature is absent panics somewhere no check of ours can reach. Hand-picking the four narrow integer features fixed the case the review reported and left the other five. dtype-full fixes the class, and stops the set growing silently the next time polars gates a new dtype. Rejecting a dtype is this library's job; being able to see it is polars'. The wheel grows to about 25 MB, which is the price of the doctrine that plugins never panic. Int128 turns out to be accepted rather than rejected once visible, being primitive and numeric -- so the test asserts the actual invariant, that no dtype produces a panic, rather than presuming which ones we reject. Co-Authored-By: Claude Fable 5 --- Cargo.lock | 182 ++++++++++++++++++++++++++++++++++ Cargo.toml | 20 ++-- README.md | 15 ++- tests/test_array_container.py | 27 +++++ 4 files changed, 227 insertions(+), 17 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8c055b7..9c3eaf8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8,6 +8,19 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "getrandom 0.3.4", + "once_cell", + "version_check", + "zerocopy", +] + [[package]] name = "aho-corasick" version = "1.1.4" @@ -554,6 +567,15 @@ dependencies = [ "miniz_oxide", ] +[[package]] +name = "float-cmp" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b09cf3155332e944990140d967ff5eceb70df778b34f77d8075db46e4704e6d8" +dependencies = [ + "num-traits", +] + [[package]] name = "fnv" version = "1.0.7" @@ -755,6 +777,16 @@ dependencies = [ "tracing", ] +[[package]] +name = "halfbrown" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c7ed2f2edad8a14c8186b847909a41fbb9c3eafa44f88bd891114ed5019da09" +dependencies = [ + "hashbrown 0.16.1", + "serde", +] + [[package]] name = "hashbrown" version = "0.15.5" @@ -792,6 +824,12 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + [[package]] name = "home" version = "0.5.12" @@ -1416,6 +1454,7 @@ dependencies = [ "polars-lazy", "polars-ops", "polars-plan", + "polars-sql", "polars-time", "polars-utils", "version_check", @@ -1427,6 +1466,7 @@ version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "33d3fe43f8702cf7899ff3d516c2e5f7dc84ee6f6a3007e1a831a0ff87940704" dependencies = [ + "atoi_simd", "bitflags", "bytemuck", "chrono", @@ -1437,6 +1477,7 @@ dependencies = [ "getrandom 0.2.17", "getrandom 0.3.4", "hashbrown 0.16.1", + "itoa", "lz4", "num-traits", "polars-arrow-format", @@ -1496,6 +1537,7 @@ dependencies = [ "boxcar", "bytemuck", "chrono", + "chrono-tz", "either", "hashbrown 0.16.1", "indexmap", @@ -1566,10 +1608,12 @@ dependencies = [ "polars-ops", "polars-plan", "polars-row", + "polars-time", "polars-utils", "rand 0.9.5", "rayon", "recursive", + "regex", "version_check", ] @@ -1593,6 +1637,7 @@ dependencies = [ "atoi_simd", "blake3", "bytes", + "chrono", "fast-float2", "fs4", "futures", @@ -1609,8 +1654,10 @@ dependencies = [ "polars-compute", "polars-core", "polars-error", + "polars-json", "polars-parquet", "polars-schema", + "polars-time", "polars-utils", "rayon", "regex", @@ -1622,6 +1669,27 @@ dependencies = [ "tokio", ] +[[package]] +name = "polars-json" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3dd2126daebf58da564fc5840cd55eb8eb2479d24dfced0a1aea2178a9b33b12" +dependencies = [ + "chrono", + "fallible-streaming-iterator", + "hashbrown 0.16.1", + "indexmap", + "itoa", + "num-traits", + "polars-arrow", + "polars-compute", + "polars-error", + "polars-utils", + "ryu", + "simd-json", + "streaming-iterator", +] + [[package]] name = "polars-lazy" version = "0.52.0" @@ -1641,6 +1709,7 @@ dependencies = [ "polars-ops", "polars-plan", "polars-stream", + "polars-time", "polars-utils", "rayon", "version_check", @@ -1673,6 +1742,7 @@ dependencies = [ "polars-io", "polars-ops", "polars-plan", + "polars-time", "polars-utils", "rayon", "recursive", @@ -1685,9 +1755,13 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7b59c80a019ef0e6f09b4416d2647076a52839305c9eb11919e8298ec667f853" dependencies = [ "argminmax", + "base64", "bytemuck", + "chrono", + "chrono-tz", "either", "hashbrown 0.16.1", + "hex", "indexmap", "libm", "memchr", @@ -1703,6 +1777,8 @@ dependencies = [ "regex-syntax", "serde", "strum_macros", + "unicode-normalization", + "unicode-reverse", "version_check", ] @@ -1748,6 +1824,8 @@ dependencies = [ "bitflags", "bytemuck", "bytes", + "chrono", + "chrono-tz", "either", "hashbrown 0.16.1", "memmap2", @@ -1765,6 +1843,7 @@ dependencies = [ "pyo3", "rayon", "recursive", + "regex", "serde", "sha2", "slotmap", @@ -1800,6 +1879,27 @@ dependencies = [ "version_check", ] +[[package]] +name = "polars-sql" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3e67ac1cbb0c972a57af3be12f19aa9803898863fe95c33cdd39df05f5738a75" +dependencies = [ + "bitflags", + "hex", + "polars-core", + "polars-error", + "polars-lazy", + "polars-ops", + "polars-plan", + "polars-time", + "polars-utils", + "rand 0.9.5", + "regex", + "serde", + "sqlparser", +] + [[package]] name = "polars-stream" version = "0.52.0" @@ -1849,6 +1949,7 @@ dependencies = [ "atoi_simd", "bytemuck", "chrono", + "chrono-tz", "now", "num-traits", "polars-arrow", @@ -2265,6 +2366,26 @@ dependencies = [ "bitflags", ] +[[package]] +name = "ref-cast" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e440fb4e4b4147295338efb76001ab9e4efc0e5839df2c47fc5ac2381d365c3" +dependencies = [ + "ref-cast-impl", +] + +[[package]] +name = "ref-cast-impl" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92ecd8964f8453721699a1ed72037b0db49ce2f5a5138486ee89bed6f67cdf3a" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + [[package]] name = "regex" version = "1.13.1" @@ -2630,6 +2751,22 @@ version = "0.3.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3a219298ac11a56ea9a6d2120044824d6f01aeb034955e7af7bc16858527deea" +[[package]] +name = "simd-json" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4255126f310d2ba20048db6321c81ab376f6a6735608bf11f0785c41f01f64e3" +dependencies = [ + "ahash", + "halfbrown", + "once_cell", + "ref-cast", + "serde", + "serde_json", + "simdutf8", + "value-trait", +] + [[package]] name = "simdutf8" version = "0.1.5" @@ -2674,6 +2811,15 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "sqlparser" +version = "0.53.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05a528114c392209b3264855ad491fcce534b94a38771b0a0b97a79379275ce8" +dependencies = [ + "log", +] + [[package]] name = "stable_deref_trait" version = "1.2.1" @@ -3005,6 +3151,30 @@ version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" +[[package]] +name = "unicode-normalization" +version = "0.1.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5fd4f6878c9cb28d874b009da9e8d183b5abc80117c40bbd187a1fde336be6e8" +dependencies = [ + "tinyvec", +] + +[[package]] +name = "unicode-reverse" +version = "1.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b6f4888ebc23094adfb574fdca9fdc891826287a6397d2cd28802ffd6f20c76" +dependencies = [ + "unicode-segmentation", +] + +[[package]] +name = "unicode-segmentation" +version = "1.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8" + [[package]] name = "unindent" version = "0.2.4" @@ -3053,6 +3223,18 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "value-trait" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e80f0c733af0720a501b3905d22e2f97662d8eacfe082a75ed7ffb5ab08cb59" +dependencies = [ + "float-cmp", + "halfbrown", + "itoa", + "ryu", +] + [[package]] name = "version_check" version = "0.9.5" diff --git a/Cargo.toml b/Cargo.toml index 79cc1d1..688c8f5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,18 +10,14 @@ crate-type = ["cdylib"] [dependencies] pyo3 = { version = "0.26", features = ["extension-module", "abi3-py310"] } pyo3-polars = { version = "0.25", features = ["derive"] } -# Every dtype a caller can put in a list column must be compiled in: -# polars reconstructs the Series across the plugin's arrow-FFI boundary -# before any of our code runs, and a dtype whose feature is missing -# panics there rather than erroring (the same root cause as the Array -# process-abort these features were first enabled for). -polars = { version = "0.52", default-features = false, features = [ - "dtype-array", - "dtype-i8", - "dtype-i16", - "dtype-u8", - "dtype-u16", -] } +# Every dtype a caller can put in a column must be compiled in, hence +# dtype-full rather than a hand-picked list. polars reconstructs the +# Series across the plugin's arrow-FFI boundary before any of our code +# runs, so a dtype whose feature is missing panics there instead of +# reaching our error path -- the same root cause as the Array +# process-abort, and unreachable by any check we could add. Rejecting a +# dtype is our job; being able to see it is polars'. +polars = { version = "0.52", default-features = false, features = ["dtype-full"] } serde = { version = "1", features = ["derive"] } realfft = "3.4.0" butterworth = "0.1.0" diff --git a/README.md b/README.md index 241fb13..e210b5f 100644 --- a/README.md +++ b/README.md @@ -104,11 +104,16 @@ and cumulative Polars operations per element instead, and treat empty lists as v Every function accepts both `List` and `Array` columns. A length-preserving function returns the container it was given; `apply_fft` narrows an `Array(w)` to `Array(w // 2 + 1)`, and `apply_interp` follows its `xp` argument, since that decides the output's shape. -`zip_binary` returns an `Array` only when both operands are `Array`s of equal width — a -mismatch is caught before execution — and a mixed pair gives a `List`. Two exceptions: -`agg_slices` reduces to a scalar, and `agg_lists` always returns a `List`, being an -expression composition that never sees its input's schema. Zero-width `Array`s are -rejected, because Polars panics when slicing them; zero-length `List`s are fine. +`zip_binary` returns an `Array` only when both operands are `Array`s, since both are +operands; `cum_agg_runs` follows its value column alone, its gate being a mask rather than +an operand. Where two inputs must agree element-for-element, their widths are checked +before execution if both are `Array`s, and per row otherwise. Two exceptions to container +propagation: `agg_slices` reduces to a scalar, and `agg_lists` always returns a `List`, +being an expression composition that never sees its input's schema (an `Array` input to it +must match `list_length` exactly, rather than being null-padded like a short `List`). + +Zero-width `Array`s are rejected at any nesting depth, because Polars panics when slicing +them; zero-length `List`s are fine. ### Two conventions, one switch diff --git a/tests/test_array_container.py b/tests/test_array_container.py index 2e63026..3901f77 100644 --- a/tests/test_array_container.py +++ b/tests/test_array_container.py @@ -350,6 +350,33 @@ def test_narrow_integer_inners_do_not_panic(inner): assert out["r"][0] == 10.0 +@pytest.mark.parametrize( + ("inner", "values"), + [ + (pl.Categorical, [["a", "b"]]), + (pl.Enum(["a", "b"]), [["a", "b"]]), + (pl.Decimal(10, 2), [[1, 2]]), + (pl.Struct({"x": pl.Int64}), [[{"x": 1}]]), + (pl.Int128, [[1, 2]]), + (pl.String, [["a", "b"]]), + (pl.Datetime, [[1, 2]]), + ], +) +def test_exotic_inner_dtypes_never_panic(inner, values): + # Rejecting a dtype is this library's job; being able to *see* it is + # polars'. A feature-gated dtype panics during the FFI + # reconstruction, before any check of ours can run, so the whole + # gated set has to be compiled in. What each dtype then does -- a + # clean rejection, or acceptance if it is numeric after all -- is + # decided by our own dtype rules; the invariant here is only that + # the answer is never a panic. + df = pl.DataFrame({"a": values}, schema={"a": pl.List(inner)}) + try: + df.select(polist.agg_slices("a", "a", aggregation="count")) + except polars.exceptions.PolarsError as exc: + assert "the plugin panicked" not in str(exc), str(exc) + + @pytest.mark.parametrize( "build", [ From a7d51c00de035ca55d0f90b8b130b6d1a0440f7d Mon Sep 17 00:00:00 2001 From: Travis Hammond Date: Wed, 2 Sep 2026 10:19:10 +0200 Subject: [PATCH 10/17] feat!: dtype preservation and the error-policy pass The last piece of the 2.0 contract: the library stops converting invalid data into missing data, and stops flattening every dtype to Float64. Dtypes. Float32 in gives Float32 out for all seven functions now, not just the two that shipped with it: kernels compute in f64 and round once at the output boundary, so Float32 results are deterministic rounded-once values. apply_interp's output follows its y column, values determining the value dtype. count emits UInt32 -- polars' count dtype -- from agg_slices and agg_lists as it already did from cum_agg_runs, the one exception to preservation. Integer inners raise at plan time naming the parameter, instead of being silently cast; integer support would multiply the vetting matrix and is deliberately out of scope. Error policy. Wrong shape is structural, not data, so it raises where it used to yield a null row: a non-power-of-two FFT length (or one below 2, where the windows sum to zero), a Butterworth signal shorter than its reflection padding, and an empty list into any transform -- the limiting case of too-short, and a raise policy that softened as the input got worse would be indefensible. For an Array the width sits in the schema and these raise at plan time; for mixed-length List data the escape hatch is pre-filtering with list.len(). Non-finite samples go the other way: they are legitimate floats that now propagate -- one NaN yields an all-NaN spectrum instead of a null row, because converting invalid to missing erased exactly the distinction the null/NaN protocol depends on. Internal kernel failures raise rather than nulling, and butterworth validates sample_rate itself so the message names the parameter. Empty lists split by family. Transforms raise, as above. Aggregations treat them as the empty selection they are -- count 0, sum per the empty flag -- which closes the one disagreement between agg_slices and agg_lists. apply_interp splits precisely: an empty x/y is no interpolant and raises, an empty xp is zero query points and returns an empty list. The FFT scaling vocabulary becomes scipy's: power -> spectrum, psd -> density, with the old names removed rather than aliased -- the ambiguous name disappearing is the point, "power" having already been misread as A-squared in practice. amplitude_squared is added as exactly that quantity, the square of the amplitude scaling with the one-sided doubling squared along with it. Window names are untouched. Tests pinned to the old behavior are flipped deliberately per the plan, and a dtype-preservation suite covers the new contract per function, aggregation and container. Co-Authored-By: Claude Fable 5 --- README.md | 41 ++++--- polars_list_utils/__init__.py | 111 +++++++++++------ src/agg_lists.rs | 4 +- src/agg_slices.rs | 62 +++++++--- src/apply_butterworth.rs | 118 ++++++++++++++---- src/apply_fft.rs | 65 +++++++--- src/apply_interp.rs | 111 +++++++++-------- src/transform.rs | 86 ++++++++------ src/util.rs | 108 ++++++++++------- tests/test_array_container.py | 20 ++-- tests/test_dtype_preservation.py | 197 +++++++++++++++++++++++++++++++ tests/test_plugins.py | 156 ++++++++++++++++++------ 12 files changed, 800 insertions(+), 279 deletions(-) create mode 100644 tests/test_dtype_preservation.py diff --git a/README.md b/README.md index e210b5f..22a261e 100644 --- a/README.md +++ b/README.md @@ -31,12 +31,15 @@ boost. - `polist.apply_fft` - Applies a (real) Fast Fourier Transform (FFT) to a List-type column of signal data, producing the `N/2 + 1` amplitudes from DC to Nyquist. - - Can pre-process the signals with a windowing function (e.g. Hann, Blackman). - - Can scale the amplitudes following the scipy conventions: `amplitude` reads a tone's peak - amplitude, `power` its mean-square (`scipy.signal.periodogram(scaling="spectrum")`), and - `psd` its power density (`scaling="density"`, integrates to the signal's mean-square power). - - The length of each signal must be a power of two, and the corresponding frequency axis - is given by `[i * sample_rate / N for i in range(N // 2 + 1)]`. + - Can pre-process the signals with a windowing function (e.g. Hann, Blackman), applied + in periodic (sym=False) form as in `scipy.signal.periodogram`. + - Scales the amplitudes under scipy's names: `amplitude` reads a tone's peak amplitude A, + `amplitude_squared` reads A² (the square of `amplitude`, a convenience), `spectrum` its + mean-square A²/2 (`scipy.signal.periodogram(scaling="spectrum")`), and `density` its + power density (`scaling="density"`, integrates to the signal's mean-square power). + - The length of each signal must be a power of two of at least 2 — a violation raises — + and the corresponding frequency axis is + `[i * sample_rate / N for i in range(N // 2 + 1)]`. - `polist.agg_slices` - Computes an aggregation of a range of y-values defined by some x-values for List-type columns. @@ -75,9 +78,17 @@ boost. - `Float32` values stay `Float32`; `count` emits `UInt32` (Polars' count dtype). The six plugin functions accept a length-1 literal list column (e.g. `pl.lit([...])`) for -any input and broadcast it across rows. Invalid configuration (bad cutoffs, unknown window -names, mismatched list lengths) raises an error instead of silently returning nulls, while -per-row data problems (e.g. a signal length that is not a power of two) yield null rows. +any input and broadcast it across rows. Anything structural raises instead of silently +returning nulls: invalid configuration (bad cutoffs, unknown window names), mismatched +paired lengths, integer inner dtypes, and per-row shapes a transform cannot process — a +non-power-of-two FFT length, a Butterworth signal shorter than its reflection padding, an +empty list. For mixed-length List data the escape hatch is pre-filtering with `list.len()`. +Only genuinely missing data nulls: a null row, or a null element inside a transform's +signal. + +`Float32` columns stay `Float32` — kernels compute in `f64` and round once at the output +boundary — with one exception: `count` emits `UInt32`, Polars' count dtype, everywhere it +appears. Both aggregation functions support `sum`, `mean`, `median`, `std`, `min`, `max`, `delta`, and `count`, and both handle missing data like the vertical aggregations of Polars itself. @@ -93,11 +104,13 @@ agree to rounding rather than bit-for-bit.) - `std` is the sample standard deviation (ddof=1) and yields null for fewer than two values, as in Polars. -The three transforms instead turn any row with missing data (a null row, an empty list, or -a list containing nulls) into a null output row: a signal with missing samples cannot be -meaningfully transformed. `zip_binary` and `cum_agg_runs` mirror the corresponding scalar -and cumulative Polars operations per element instead, and treat empty lists as valid -(empty in, empty out). +The three transforms instead follow a three-line rule: missing data nulls (a null row or +a null element yields a null row — a signal with missing samples cannot be meaningfully +transformed); invalid values propagate (NaN and ±inf flow through the arithmetic — one NaN +yields an all-NaN spectrum, staying visibly invalid rather than becoming missing); wrong +shape raises. For the aggregations an empty list is simply an empty selection, and +`zip_binary` and `cum_agg_runs` mirror the corresponding scalar and cumulative Polars +operations per element, treating empty lists as valid (empty in, empty out). ### List and Array diff --git a/polars_list_utils/__init__.py b/polars_list_utils/__init__.py index 52361ae..263c728 100644 --- a/polars_list_utils/__init__.py +++ b/polars_list_utils/__init__.py @@ -2,10 +2,12 @@ Six elementwise Rust plugins, plus :func:`agg_lists`, which composes polars' own vertical aggregations (with a small pass-through plugin to -prepare its input). The four original functions take `List[f64]` -(anything numeric is cast to it); :func:`zip_binary` and -:func:`cum_agg_runs` instead take float or Boolean inners as their -operands require, and preserve `Float32`: +prepare its input). Value-like inputs take `Float32` or `Float64` +inners — anything else, integers included, raises at plan time naming +the parameter — and `Float32` is preserved: kernels compute in `f64` +and round once at the output boundary. The one exception to dtype +preservation is `count`, which emits `UInt32`, polars' count dtype. +Gates and the Boolean operands of `zip_binary` take Boolean inners: - :func:`apply_interp`: interpolate (x, y) data onto new x coordinates. - :func:`apply_butterworth`: zero-phase Butterworth filtering. @@ -26,22 +28,27 @@ before execution when both are `Array`s, and per row otherwise. Two exceptions to container propagation: :func:`agg_slices` reduces to a scalar, and :func:`agg_lists` always returns a `List`, being an -expression composition that never sees its input's schema. Zero-width `Array`s are rejected — polars panics when -slicing them — while zero-length `List`s are fine. +expression composition that never sees its input's schema. Zero-width +`Array`s are rejected — polars panics when slicing them — while +zero-length `List`s are fine. The plugins accept length-1 literal list columns (e.g. `pl.lit(...)`) for any input and broadcast them. Null rows produce null output rows. -Missing data follows three deliberate family regimes: - -- Transforms (`apply_interp`, `apply_butterworth`, `apply_fft`) null - out rows whose lists are empty or contain null elements (a signal - with missing samples cannot be transformed). -- Aggregations (`agg_slices`, `agg_lists`) skip missing data - polars-style; NaN is a legitimate float and flows through. -- `zip_binary` and `cum_agg_runs` mirror polars' scalar and `cum_*` - operations per element: nulls propagate per those ops' rules, and - empty lists are valid (empty in, empty out). +Data conditions follow three deliberate family regimes. A signal is +atomic, so for the transforms (`apply_interp`, `apply_butterworth`, +`apply_fft`) the rule is three lines: **missing data nulls** — a null +row or any null element yields a null row, a signal with missing +samples being untransformable; **invalid values propagate** — NaN and +±inf are legitimate floats that flow through the arithmetic and stay +visible in the output; **wrong shape raises** — a length the transform +cannot process, including the empty list, is structural, not data. +The aggregations (`agg_slices`, `agg_lists`) instead skip missing data +polars-style, and treat an empty list as an empty selection. And +`zip_binary`/`cum_agg_runs` mirror polars' scalar and `cum_*` +operations per element: nulls propagate per those ops' rules, and +empty lists are valid (empty in, empty out — total functions over zero +elements). """ import math @@ -56,7 +63,7 @@ IntoExprColumn = str | pl.Expr | pl.Series Window = Literal["hann", "hanning", "blackman"] -Scaling = Literal["amplitude", "power", "psd"] +Scaling = Literal["amplitude", "amplitude_squared", "spectrum", "density"] Aggregation = Literal["sum", "mean", "median", "std", "min", "max", "delta", "count"] Empty = Literal["null", "zero"] @@ -129,6 +136,11 @@ def apply_interp( values outside the data range clamped to the first/last `y` value. `x` must be sorted in increasing order. + Empty splits precisely: an empty `x`/`y` is no interpolant and + raises, while an empty `xp` is zero query points — a valid question + whose answer is an empty list. The output's container follows `xp` + and its inner dtype follows `y` (values determine the value dtype). + With `strict=True` (the default) a row whose `x` values descend raises; duplicates stay legal, as in numpy. Pass `strict=False` for numpy's silent behaviour, where unsorted `x` interpolates garbage @@ -163,8 +175,13 @@ def apply_butterworth( Filtering is bidirectional (like `scipy.signal.filtfilt`): zero phase lag, squared magnitude response, and the cutoff sits at -6 dB rather than -3 dB. Bandpass doubles the design order, as in - `scipy.signal.butter`. Rows with no more than 3x the effective - order samples (the reflection padding) yield null. + `scipy.signal.butter`. A row with no more than 3x the effective + order samples (the reflection padding) raises — wrong shape is + structural, not data; for an `Array` the width is in the schema and + the raise lands at plan time, and for mixed-length `List` data the + escape hatch is pre-filtering with `list.len()`. Non-finite samples + propagate and contaminate the filtered output rather than nulling + it. """ return _plugin( "apply_butterworth", @@ -187,17 +204,30 @@ def apply_fft( Emits `N/2 + 1` values from DC to Nyquist, so the frequency axis is `[i * sample_rate / N for i in range(N // 2 + 1)]`. Signal lengths - must be powers of two; rows violating that (or containing non-finite - values) yield null. + must be powers of two of at least 2 — a violation raises (for an + `Array`, already at plan time, the width being in the schema; for + mixed-length `List` data, pre-filter with `list.len()`). Non-finite + samples are legitimate floats and propagate: every bin sums all + samples, so one NaN yields an all-NaN spectrum rather than a null + row — invalid data stays visibly invalid instead of becoming + missing. Guard downstream thresholds with `is_finite()`. - Scaling conventions match scipy for a tone of peak amplitude A: + Scaling names are scipy's, for a tone of peak amplitude A: - `"amplitude"`: peak-amplitude spectrum, the tone reads A. - - `"power"`: power spectrum, the tone reads its mean-square A²/2 + - `"amplitude_squared"`: `amplitude`² per bin, the tone reads A². + A convenience with no external anchor — elementwise squaring of a + list column is exactly the boilerplate this library removes. + - `"spectrum"`: power spectrum, the tone reads its mean-square A²/2 (`scipy.signal.periodogram(scaling="spectrum")`). - - `"psd"`: power spectral density in unit²/Hz; integrates to the - signal's mean-square power (`scaling="density"`). - - `None`: raw FFT magnitudes. + - `"density"`: power spectral density in unit²/Hz; integrates to + the signal's mean-square power (`scaling="density"`). + - `None`: raw FFT magnitudes of the windowed signal. + + The one-sided doubling happens in each scaling's own domain and + skips DC and Nyquist, so `spectrum != amplitude²/2` at exactly those + two bins — matching scipy. Windows are periodic (sym=False), as in + `periodogram`. """ return _plugin( "apply_fft", @@ -284,8 +314,9 @@ def agg_slices( Selects the elements of `value_column` whose corresponding element in `index_column` lies in any `slices_include` range and no - `slices_exclude` range, then aggregates them into a `Float64`. - Omitting `slices_include` includes everything. + `slices_exclude` range, then aggregates them into a scalar that + keeps the value column's inner dtype (`count` alone emits `UInt32`, + polars' count dtype). Omitting `slices_include` includes everything. Missing data follows the polars convention, so results follow polars' vertical aggregations over the same values (semantically — @@ -300,9 +331,10 @@ def agg_slices( `std` is the sample standard deviation (ddof=1, polars' default) and yields null for selections with fewer than two values. - An empty selection — the lists held elements, but no index fell in - range — yields null, except `count`, which yields 0. `empty` - decides the one aggregation where the two conventions disagree: + An empty selection — no index fell in range, or the lists were + themselves empty, the limiting case of the same condition — yields + null, except `count`, which yields 0. `empty` decides the one + aggregation where the two conventions disagree: - `"null"` (default): summing nothing is unknown, not zero. An empty selection usually means a misconfigured range, and a real @@ -311,11 +343,11 @@ def agg_slices( - `"zero"`: polars' convention, the identity element of addition, matching `explode().group_by().agg(col.filter(...).sum())`. - `empty` governs the empty *selection* only. A null row, and a row - whose lists are themselves empty, stay null under both settings for - every aggregation including `count` — so `empty="zero"` is not the - same as `.fill_null(0.0)`, which cannot tell a missing row from an - empty selection and would launder the former into a real zero. + `empty` governs the empty *selection* only. A null row stays null + under both settings for every aggregation including `count` — so + `empty="zero"` is not the same as `.fill_null(0.0)`, which cannot + tell a missing row from an empty selection and would launder the + former into a real zero. Range bounds accept anything `float()` converts, and the explicit form may be given as tuples or lists (so ranges loaded from JSON @@ -347,7 +379,7 @@ def agg_slices( "min": pl.Expr.min, "max": pl.Expr.max, "delta": lambda e: e.max() - e.min(), - "count": lambda e: e.count().cast(pl.Float64), + "count": pl.Expr.count, # UInt32, polars' count dtype } _AGGS_EMPTY_ZERO: dict[str, Callable[[pl.Expr], pl.Expr]] = { @@ -388,8 +420,9 @@ def agg_lists( The aggregations are polars' own, so missing data follows the polars convention exactly as in :func:`agg_slices`: null elements are - skipped (`count` counts non-null values), NaN propagates per-kernel, - and `std` is the sample standard deviation (ddof=1). `empty` decides + skipped (`count` counts non-null values and emits `UInt32`, polars' + count dtype), NaN propagates per-kernel, and `std` is the sample + standard deviation (ddof=1). `empty` decides what an all-null position sums to — null by default, `0.0` under `empty="zero"` — exactly as in :func:`agg_slices`. diff --git a/src/agg_lists.rs b/src/agg_lists.rs index 7a86a8c..1667a5b 100644 --- a/src/agg_lists.rs +++ b/src/agg_lists.rs @@ -9,7 +9,7 @@ //! `List`; and the wrapper cannot branch on dtype itself, because an //! expression is built before any schema is resolved. -use crate::util::{Container, as_list}; +use crate::util::{Container, as_list, require_float_inner}; use polars::prelude::*; use pyo3_polars::derive::polars_expr; use serde::Deserialize; @@ -25,7 +25,7 @@ fn prepare_output( kwargs: PrepareKwargs, ) -> PolarsResult { let (container, inner) = - Container::split(input_fields[0].dtype(), "agg_lists: list_column")?; + require_float_inner(input_fields[0].dtype(), "agg_lists: list_column")?; // An Array's width is in the schema, so a mismatch against // list_length is decidable before execution (validation, not // inference: the signature does not fork on container). diff --git a/src/agg_slices.rs b/src/agg_slices.rs index 6832d2b..c42ab6a 100644 --- a/src/agg_slices.rs +++ b/src/agg_slices.rs @@ -1,6 +1,6 @@ use crate::slice::{Range, Slice}; use crate::statistic::{Aggregation, Empty}; -use crate::util::{Container, ListInputs, ensure_equal_widths}; +use crate::util::{ListInputs, ensure_equal_widths, require_float_inner}; use polars::prelude::*; use pyo3_polars::derive::polars_expr; use serde::Deserialize; @@ -13,22 +13,42 @@ struct AggSlicesKwargs { empty: Empty, } -/// Always a `Float64` scalar — there is no container to propagate — but -/// resolved through a function rather than declared statically, so the -/// inputs are validated at plan time. That is what keeps a zero-width -/// `Array` from reaching the FFI boundary, where polars would panic -/// rather than raise. -fn agg_slices_output(input_fields: &[Field]) -> PolarsResult { - let (values, _) = - Container::split(input_fields[0].dtype(), "agg_slices: value_column")?; +/// The output dtype for the requested aggregation: the value column's +/// inner dtype preserved, except `count`, which emits `UInt32` (polars' +/// count dtype, the one exception to inner-dtype preservation). +fn scalar_dtype( + aggregation: Aggregation, + value_inner: &DataType, +) -> DataType { + match aggregation { + Aggregation::Count => DataType::UInt32, + _ => value_inner.clone(), + } +} + +/// A scalar — there is no container to propagate — but resolved through +/// a function rather than declared statically, so the inputs are +/// validated at plan time. That is what keeps a zero-width `Array` from +/// reaching the FFI boundary, where polars would panic rather than +/// raise, and what makes an integer inner a plan-time error naming the +/// parameter. +fn agg_slices_output( + input_fields: &[Field], + kwargs: AggSlicesKwargs, +) -> PolarsResult { + let (values, value_inner) = + require_float_inner(input_fields[0].dtype(), "agg_slices: value_column")?; let (indices, _) = - Container::split(input_fields[1].dtype(), "agg_slices: index_column")?; + require_float_inner(input_fields[1].dtype(), "agg_slices: index_column")?; ensure_equal_widths( values, indices, "agg_slices: value and index Arrays must have equal widths", )?; - Ok(Field::new(PlSmallStr::from(""), DataType::Float64)) + Ok(Field::new( + PlSmallStr::from(""), + scalar_dtype(kwargs.aggregation, &value_inner), + )) } /// Aggregate the values whose paired index falls within the requested @@ -36,14 +56,17 @@ fn agg_slices_output(input_fields: &[Field]) -> PolarsResult { /// /// Missing data follows the polars convention: null elements are /// skipped (pairwise with their index), while NaN is a legitimate -/// float value that flows into the aggregation. An empty selection -/// yields null (`count` reading 0, and `sum` following the `empty` +/// float value that flows into the aggregation. An empty list is an +/// empty selection, exactly like a non-empty list whose indices all +/// miss: null (`count` reading 0, `sum` following the `empty` /// setting). NaN indices never match any range. -#[polars_expr(output_type_func=agg_slices_output)] +#[polars_expr(output_type_func_with_kwargs=agg_slices_output)] fn agg_slices( inputs: &[Series], kwargs: AggSlicesKwargs, ) -> PolarsResult { + let (_, value_inner) = + require_float_inner(inputs[0].dtype(), "agg_slices: value_column")?; let li = ListInputs::with_inner_nulls(inputs)?; // No include ranges means include everything; excludes then carve it up. @@ -81,5 +104,14 @@ fn agg_slices( }) .collect::>()?; - Ok(out.into_series()) + // Computed in f64; rounded once to the declared dtype at the exit + // boundary (doctrine 4). Counts are exact integers well below 2^53, + // so their detour through f64 is lossless. + let target = scalar_dtype(kwargs.aggregation, &value_inner); + let s = out.into_series(); + if target == DataType::Float64 { + Ok(s) + } else { + s.cast(&target) + } } diff --git a/src/apply_butterworth.rs b/src/apply_butterworth.rs index b79cbe3..44911c8 100644 --- a/src/apply_butterworth.rs +++ b/src/apply_butterworth.rs @@ -1,14 +1,9 @@ use crate::filter::Filter; -use crate::util::{Container, apply_list_transform, transform_output}; +use crate::util::{Container, apply_list_transform, require_float_inner}; use polars::prelude::*; use pyo3_polars::derive::polars_expr; use serde::Deserialize; -/// Length-preserving, so the output keeps the signal's container. -fn butterworth_output(input_fields: &[Field]) -> PolarsResult { - transform_output(input_fields, 0, "apply_butterworth: signal_column", |w| w) -} - #[derive(Deserialize)] struct ButterworthKwargs { sample_rate: f64, @@ -17,6 +12,49 @@ struct ButterworthKwargs { filter_order: Option, } +impl ButterworthKwargs { + fn filter(&self) -> Option { + match (self.min_freq, self.max_freq) { + (Some(lo), Some(hi)) => Some(Filter::Bandpass(lo, hi)), + (Some(lo), None) => Some(Filter::Highpass(lo)), + (None, Some(hi)) => Some(Filter::Lowpass(hi)), + (None, None) => None, + } + } + + fn order(&self) -> usize { + self.filter_order.unwrap_or(4) + } +} + +/// Length-preserving, so the output keeps the signal's container and +/// inner dtype. The kwargs are available here, which lets an `Array` +/// too short for the reflection padding raise at plan time — the width +/// is in the schema, and a fixed width means every row would fail. +fn butterworth_output( + input_fields: &[Field], + kwargs: ButterworthKwargs, +) -> PolarsResult { + let (container, inner) = + require_float_inner(input_fields[0].dtype(), "apply_butterworth: signal_column")?; + polars_ensure!( + kwargs.sample_rate > 0.0, + ComputeError: "apply_butterworth: sample_rate must be positive, got {}", + kwargs.sample_rate + ); + if let (Container::Array(w), Some(filter)) = (container, kwargs.filter()) { + let min_samples = filter.min_samples(kwargs.order()); + polars_ensure!( + w >= min_samples, + ComputeError: + "apply_butterworth: signal_column has Array width {w} but the order-{} \ + filter needs at least {min_samples} samples (reflection padding)", + kwargs.order() + ); + } + Ok(Field::new(PlSmallStr::from(""), container.dtype(inner))) +} + /// Apply a zero-phase Butterworth filter to a list column. /// /// # Arguments (via kwargs) @@ -26,35 +64,63 @@ struct ButterworthKwargs { /// - `filter_order`: Filter order (default 4) /// /// If both cutoffs are specified, a bandpass filter is applied; with one, -/// a highpass or lowpass; with neither, the input passes through. +/// a highpass or lowpass; with neither, the input passes through (the +/// family entry contract still applies: null-element rows null out, and +/// an empty row raises — a bypass that skipped it would make row +/// semantics depend on parameter values). /// /// Filtering is bidirectional (filtfilt): zero phase lag, squared /// magnitude response, cutoff at -6 dB. Bandpass doubles the design -/// order, as in scipy.signal.butter. Rows with no more than 3x the -/// effective order samples (the reflection padding) yield null. -#[polars_expr(output_type_func=butterworth_output)] +/// order, as in scipy.signal.butter. A row with no more than 3x the +/// effective order samples (the reflection padding) raises — wrong shape +/// is structural, not data. Non-finite samples are legitimate floats and +/// contaminate the filtered output rather than nulling it. +#[polars_expr(output_type_func_with_kwargs=butterworth_output)] fn apply_butterworth( inputs: &[Series], kwargs: ButterworthKwargs, ) -> PolarsResult { - let (out, _) = - Container::split(inputs[0].dtype(), "apply_butterworth: signal_column")?; - let filter = match (kwargs.min_freq, kwargs.max_freq) { - (Some(lo), Some(hi)) => Filter::Bandpass(lo, hi), - (Some(lo), None) => Filter::Highpass(lo), - (None, Some(hi)) => Filter::Lowpass(hi), - (None, None) => { - return apply_list_transform(inputs, out, |cols| Ok(Some(cols[0].to_vec()))); - } + polars_ensure!( + kwargs.sample_rate > 0.0, + ComputeError: "apply_butterworth: sample_rate must be positive, got {}", + kwargs.sample_rate + ); + let (out, inner) = + require_float_inner(inputs[0].dtype(), "apply_butterworth: signal_column")?; + let Some(filter) = kwargs.filter() else { + return apply_list_transform( + inputs, + "apply_butterworth", + &[false], + out, + &inner, + |cols| Ok(Some(cols[0].to_vec())), + ); }; - let order = kwargs.filter_order.unwrap_or(4); + let order = kwargs.order(); let built = filter.build(kwargs.sample_rate, order)?; let min_samples = filter.min_samples(order); - apply_list_transform(inputs, out, move |cols| { - if cols[0].len() < min_samples { - return Ok(None); - } - Ok(built.bidirectional(&cols[0].to_vec()).ok()) - }) + apply_list_transform( + inputs, + "apply_butterworth", + &[false], + out, + &inner, + move |cols| { + polars_ensure!( + cols[0].len() >= min_samples, + ComputeError: + "apply_butterworth: a row's signal has {} samples but the order-{order} \ + filter needs at least {min_samples} (reflection padding; pre-filter \ + with list.len() to drop such rows)", + cols[0].len() + ); + // An internal filter failure is a bug surfacing, not + // missing data: raise rather than null. + built.bidirectional(&cols[0].to_vec()).map(Some).map_err( + |e| polars_err!(ComputeError: "apply_butterworth: filter failed: {e:?}"), + ) + }, + ) } diff --git a/src/apply_fft.rs b/src/apply_fft.rs index c490467..8f6c554 100644 --- a/src/apply_fft.rs +++ b/src/apply_fft.rs @@ -1,16 +1,9 @@ use crate::transform::{Scaling, Window, compute_fft}; -use crate::util::{Container, apply_list_transform, transform_output}; +use crate::util::{Container, apply_list_transform, require_float_inner}; use polars::prelude::*; use pyo3_polars::derive::polars_expr; use serde::Deserialize; -/// The only width-transforming function: a one-sided real FFT of `n` -/// samples emits `n/2 + 1` bins, so an `Array(w)` narrows to -/// `Array(w/2 + 1)`. -fn fft_output(input_fields: &[Field]) -> PolarsResult { - transform_output(input_fields, 0, "apply_fft: signal_column", |w| w / 2 + 1) -} - #[derive(Deserialize)] struct FftKwargs { sample_rate: f64, @@ -18,15 +11,49 @@ struct FftKwargs { scaling: Option, } +/// The FFT needs a power-of-two signal length of at least 2 (a +/// length-1 signal *is* a power of two, but the windows sum to zero on +/// it and the scaled spectrum reads NaN; a one-sample spectrum is +/// meaningless anyway). For an `Array` the width sits in the schema, so +/// a violation raises here at plan time; a `List` is checked per row. +fn fft_width_ok(width: usize) -> bool { + width >= 2 && width.is_power_of_two() +} + +/// The only width-transforming function: a one-sided real FFT of `n` +/// samples emits `n/2 + 1` bins, so an `Array(w)` narrows to +/// `Array(w/2 + 1)`. The inner dtype follows the signal. +fn fft_output(input_fields: &[Field]) -> PolarsResult { + let (container, inner) = + require_float_inner(input_fields[0].dtype(), "apply_fft: signal_column")?; + let out = match container { + Container::List => Container::List, + Container::Array(w) => { + polars_ensure!( + fft_width_ok(w), + ComputeError: + "apply_fft: signal_column has Array width {w}; the FFT needs a \ + power-of-two length of at least 2" + ); + Container::Array(w / 2 + 1) + } + }; + Ok(Field::new(PlSmallStr::from(""), out.dtype(inner))) +} + /// Apply a one-sided FFT to a list column of time-domain signals. /// /// # Arguments (via kwargs) /// - `sample_rate`: Sampling rate in Hz -/// - `window`: Optional window function ("hann"/"hanning", "blackman") -/// - `scaling`: Optional scaling ("amplitude", "power", "psd") +/// - `window`: Optional window function ("hanning"/"hann", "blackman") +/// - `scaling`: Optional scaling ("amplitude", "amplitude_squared", +/// "spectrum", "density") /// -/// Emits N/2 + 1 values from DC to Nyquist. Rows whose signal is not a -/// power-of-two length, or contains non-finite values, yield null. +/// Emits N/2 + 1 values from DC to Nyquist. A signal length that is not +/// a power of two of at least 2 raises — wrong shape is structural, not +/// data. Non-finite samples are legitimate floats and propagate: every +/// bin sums all samples, so one NaN yields an all-NaN spectrum rather +/// than a null row that would launder invalid data into missing data. #[polars_expr(output_type_func=fft_output)] fn apply_fft( inputs: &[Series], @@ -39,13 +66,21 @@ fn apply_fft( ); let window = kwargs.window.unwrap_or_default(); let scaling = kwargs.scaling.unwrap_or_default(); - let (input, _) = Container::split(inputs[0].dtype(), "apply_fft: signal_column")?; + let (input, inner) = + require_float_inner(inputs[0].dtype(), "apply_fft: signal_column")?; let out = match input { Container::List => Container::List, Container::Array(w) => Container::Array(w / 2 + 1), }; - apply_list_transform(inputs, out, |cols| { - Ok(compute_fft(cols[0], kwargs.sample_rate, window, scaling)) + apply_list_transform(inputs, "apply_fft", &[false], out, &inner, |cols| { + polars_ensure!( + fft_width_ok(cols[0].len()), + ComputeError: + "apply_fft: a row's signal has {} samples; the FFT needs a power-of-two \ + length of at least 2 (pre-filter with list.len() to drop such rows)", + cols[0].len() + ); + compute_fft(cols[0], kwargs.sample_rate, window, scaling).map(Some) }) } diff --git a/src/apply_interp.rs b/src/apply_interp.rs index 3184125..a06e385 100644 --- a/src/apply_interp.rs +++ b/src/apply_interp.rs @@ -1,6 +1,4 @@ -use crate::util::{ - Container, apply_list_transform, ensure_equal_widths, transform_output, -}; +use crate::util::{apply_list_transform, ensure_equal_widths, require_float_inner}; use interp::{InterpMode, interp_slice}; use polars::prelude::*; use pyo3_polars::derive::polars_expr; @@ -13,22 +11,29 @@ struct InterpKwargs { } /// Shape follows `xp`: the output holds one value per query coordinate, -/// so the query column decides both length and container. The `x`/`y` -/// containers do not reach the output — but they are still split here, -/// because that is the only thing standing between a zero-width `Array` -/// and the polars-arrow panic, and because an `x`/`y` width mismatch is -/// decidable from the schema. +/// so the query column decides both length and container. The inner +/// dtype follows `y` — values determine the value dtype. `x` and `y` +/// are still split here, because that is what rejects zero-width +/// `Array`s and non-float inners before execution, and an `x`/`y` +/// width mismatch is decidable from the schema. fn interp_output(input_fields: &[Field]) -> PolarsResult { - let (x, _) = Container::split(input_fields[0].dtype(), "apply_interp: x_column")?; - let (y, _) = Container::split(input_fields[1].dtype(), "apply_interp: y_column")?; + let (x, _) = require_float_inner(input_fields[0].dtype(), "apply_interp: x_column")?; + let (y, y_inner) = + require_float_inner(input_fields[1].dtype(), "apply_interp: y_column")?; ensure_equal_widths(x, y, "apply_interp: x and y Arrays must have equal widths")?; - transform_output(input_fields, 2, "apply_interp: xp_column", |w| w) + let (xp, _) = + require_float_inner(input_fields[2].dtype(), "apply_interp: xp_column")?; + Ok(Field::new(PlSmallStr::from(""), xp.dtype(y_inner))) } /// Interpolate `y` values from `(x, y)` data onto new `xp` coordinates, /// as in numpy.interp; `xp` values outside the data range clamp to the /// first/last `y` value. /// +/// Empty splits precisely: an empty `x`/`y` is no interpolant and +/// raises, while an empty `xp` is zero query points — a valid question +/// whose answer is an empty list. +/// /// `x` must be non-decreasing for the result to be meaningful; with /// `strict` the first descending pair raises, without it unsorted `x` /// silently interpolates garbage, as in numpy. NaN in `x` is a @@ -39,43 +44,53 @@ fn apply_interp( inputs: &[Series], kwargs: InterpKwargs, ) -> PolarsResult { - let (out, _) = Container::split(inputs[2].dtype(), "apply_interp: xp_column")?; - apply_list_transform(inputs, out, |cols| { - polars_ensure!( - cols[0].len() == cols[1].len(), - ComputeError: "apply_interp: x and y lists differ in length ({} vs {})", - cols[0].len(), cols[1].len() - ); - // Compare each coordinate against the last non-NaN one rather - // than its immediate neighbour. NaN stays a legitimate float - // that propagates into the output, but it must not hide a - // descent: every IEEE comparison with NaN is false, so an - // adjacent-pair test reads [0, 10, NaN, 1] as sorted. - // - // No row index in the message: the plugin is handed one chunk - // at a time, so any position named here would be chunk-local. - if kwargs.strict { - let mut last: Option = None; - for &x in cols[0] { - if x.is_nan() { - continue; - } - if let Some(prev) = last - && x < prev - { - polars_bail!(ComputeError: - "apply_interp: x_column must be non-decreasing with \ - strict=true, found {prev} followed by {x}" - ); + let (xp, _) = require_float_inner(inputs[2].dtype(), "apply_interp: xp_column")?; + let (_, y_inner) = require_float_inner(inputs[1].dtype(), "apply_interp: y_column")?; + apply_list_transform( + inputs, + "apply_interp", + &[false, false, true], + xp, + &y_inner, + |cols| { + polars_ensure!( + cols[0].len() == cols[1].len(), + ComputeError: "apply_interp: x and y lists differ in length ({} vs {})", + cols[0].len(), cols[1].len() + ); + // Compare each coordinate against the last non-NaN one + // rather than its immediate neighbour. NaN stays a + // legitimate float that propagates into the output, but it + // must not hide a descent: every IEEE comparison with NaN + // is false, so an adjacent-pair test reads [0, 10, NaN, 1] + // as sorted. + // + // No row index in the message: the plugin is handed one + // chunk at a time, so any position named here would be + // chunk-local. + if kwargs.strict { + let mut last: Option = None; + for &x in cols[0] { + if x.is_nan() { + continue; + } + if let Some(prev) = last + && x < prev + { + polars_bail!(ComputeError: + "apply_interp: x_column must be non-decreasing with \ + strict=true, found {prev} followed by {x}" + ); + } + last = Some(x); } - last = Some(x); } - } - Ok(Some(interp_slice( - cols[0], - cols[1], - cols[2], - &InterpMode::FirstLast, - ))) - }) + Ok(Some(interp_slice( + cols[0], + cols[1], + cols[2], + &InterpMode::FirstLast, + ))) + }, + ) } diff --git a/src/transform.rs b/src/transform.rs index 4140235..57b875e 100644 --- a/src/transform.rs +++ b/src/transform.rs @@ -10,6 +10,7 @@ use std::cell::RefCell; +use polars::prelude::*; use realfft::RealFftPlanner; use serde::Deserialize; @@ -54,39 +55,61 @@ impl Window { } } -/// Scaling of the one-sided spectrum. +/// Scaling of the one-sided spectrum, in scipy's vocabulary. +/// +/// The names follow `scipy.signal.periodogram` because the previous set +/// produced a real misreading: `"power"` was read as A² when it yields +/// the mean-square A²/2. `"spectrum"` and `"density"` are scipy's +/// unambiguous names for those quantities; `"amplitude_squared"` is the +/// convenience the misreading was reaching for, defined as the square +/// of the `"amplitude"` scaling. The old names are removed, not +/// aliased — the ambiguous name disappearing is the point. #[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize)] -#[serde(rename_all = "lowercase")] +#[serde(rename_all = "snake_case")] pub(crate) enum Scaling { - /// Raw FFT magnitudes. + /// Raw FFT magnitudes of the windowed signal. #[default] None, /// Peak-amplitude spectrum: a tone of amplitude A reads A. Amplitude, - /// Power spectrum: a tone of amplitude A reads its mean-square A²/2. - Power, + /// `amplitude`² per bin (one-sided doubling squared along with it): + /// a tone of amplitude A reads A². A convenience with no external + /// anchor — elementwise squaring of a list column is exactly the + /// boilerplate this library exists to remove. + AmplitudeSquared, + /// Power spectrum: a tone of amplitude A reads its mean-square + /// A²/2, as `scipy.signal.periodogram(scaling="spectrum")`. + Spectrum, /// Power spectral density in [unit²/Hz]: power divided by the /// window's equivalent noise bandwidth. Integrates to the signal's - /// mean-square power (Parseval). - Psd, + /// mean-square power (Parseval), as `scaling="density"`. + Density, } /// Raw one-sided FFT magnitudes (N/2 + 1 values from DC to Nyquist), /// using a thread-local planner to cache FFT plans across rows. -fn fft_magnitudes(samples: &[f64]) -> Option> { +/// +/// A failed plan or process call raises: it is a bug surfacing, not +/// missing data, and converting it to a null row would launder an +/// internal error into "missing signal". +fn fft_magnitudes(samples: &[f64]) -> PolarsResult> { FFT_PLANNER.with(|cell| { let fft = cell.borrow_mut().plan_fft_forward(samples.len()); let mut input = samples.to_vec(); let mut spectrum = fft.make_output_vec(); - fft.process(&mut input, &mut spectrum).ok()?; - Some(spectrum.iter().map(|c| c.norm()).collect()) + fft.process(&mut input, &mut spectrum).map_err( + |e| polars_err!(ComputeError: "apply_fft: FFT kernel failed: {e}"), + )?; + Ok(spectrum.iter().map(|c| c.norm()).collect()) }) } /// One-sided spectrum of `signal` with the given window and scaling. /// -/// Returns `None` for signals the FFT pipeline cannot process: empty, -/// non-power-of-two length, or containing non-finite values. +/// The caller guarantees a power-of-two length of at least 2. Non-finite +/// samples are legitimate floats and flow through the arithmetic: every +/// bin sums all samples, so one NaN yields an all-NaN spectrum — invalid +/// data stays visibly invalid rather than becoming missing. /// /// Scaling normalizes by S1 = Σw[n]; interior bins are doubled **in the /// power domain** to account for the discarded negative frequencies, @@ -96,11 +119,7 @@ pub(crate) fn compute_fft( sample_rate: f64, window: Window, scaling: Scaling, -) -> Option> { - if !signal.len().is_power_of_two() || signal.iter().any(|v| !v.is_finite()) { - return None; - } - +) -> PolarsResult> { let coeffs = window.coefficients(signal.len()); let windowed: Vec = signal.iter().zip(&coeffs).map(|(v, c)| v * c).collect(); let magnitudes = fft_magnitudes(&windowed)?; @@ -111,21 +130,20 @@ pub(crate) fn compute_fft( let enbw_hz = sample_rate * s2 / (s1 * s1); let last = magnitudes.len() - 1; - Some( - magnitudes - .iter() - .enumerate() - .map(|(i, &mag)| { - // One-sided doubling: DC (i=0) and Nyquist (i=last) have - // no negative-frequency counterpart and are not doubled. - let sf = if i == 0 || i == last { 1.0 } else { 2.0 }; - match scaling { - Scaling::None => mag, - Scaling::Amplitude => sf * mag / s1, - Scaling::Power => sf * (mag / s1).powi(2), - Scaling::Psd => sf * (mag / s1).powi(2) / enbw_hz, - } - }) - .collect(), - ) + Ok(magnitudes + .iter() + .enumerate() + .map(|(i, &mag)| { + // One-sided doubling: DC (i=0) and Nyquist (i=last) have + // no negative-frequency counterpart and are not doubled. + let sf = if i == 0 || i == last { 1.0 } else { 2.0 }; + match scaling { + Scaling::None => mag, + Scaling::Amplitude => sf * mag / s1, + Scaling::AmplitudeSquared => (sf * mag / s1).powi(2), + Scaling::Spectrum => sf * (mag / s1).powi(2), + Scaling::Density => sf * (mag / s1).powi(2) / enbw_hz, + } + }) + .collect()) } diff --git a/src/util.rs b/src/util.rs index 1fddad8..c474e81 100644 --- a/src/util.rs +++ b/src/util.rs @@ -122,50 +122,46 @@ pub(crate) fn restore( } } -/// Output dtype for the transforms: `Float64` elements in the container -/// `input_fields[from]` arrived in, whose width `width_of` maps (the FFT -/// halves it; everything else preserves it). -pub(crate) fn transform_output( - input_fields: &[Field], - from: usize, +/// Split a `List`/`Array` dtype and require a float inner: `Float32` or +/// `Float64`, anything else (ints included) raising with the parameter's +/// name. The schema knows the inner dtype, so every output-type function +/// routes its inputs through this and the rejection lands at plan time. +pub(crate) fn require_float_inner( + dtype: &DataType, label: &str, - width_of: impl Fn(usize) -> usize, -) -> PolarsResult { - let (container, _) = Container::split(input_fields[from].dtype(), label)?; - let out = match container { - Container::List => Container::List, - Container::Array(w) => Container::Array(width_of(w)), - }; - Ok(Field::new( - PlSmallStr::from(""), - out.dtype(DataType::Float64), - )) +) -> PolarsResult<(Container, DataType)> { + let (container, inner) = Container::split(dtype, label)?; + polars_ensure!( + matches!(inner, DataType::Float32 | DataType::Float64), + ComputeError: "{label} must have a Float32 or Float64 inner dtype, got {dtype}" + ); + Ok((container, inner)) } /// Extract one row's list as a `Vec`. /// -/// Returns `None` (-> null output row) if the list is empty or contains -/// null elements: dropping inner nulls would silently misalign this list -/// against its paired columns. +/// Returns `None` (-> null output row) if the list contains null +/// elements: dropping inner nulls would silently misalign this list +/// against its paired columns. Empty lists are kept — whether an empty +/// list is valid is the caller's per-input policy, not an extraction +/// question. fn extract_row(s: &Series) -> Option> { - let values: Option> = s.f64().ok()?.into_iter().collect(); - values.filter(|v| !v.is_empty()) + s.f64().ok()?.into_iter().collect() } -/// Extract one row's list with inner nulls preserved. -/// -/// Returns `None` (-> null output row) only if the list is empty. +/// Extract one row's list with inner nulls preserved. An empty list is +/// an empty selection, not a null row. fn extract_row_nullable(s: &Series) -> Option>> { - let values: Vec> = s.f64().ok()?.into_iter().collect(); - (!values.is_empty()).then_some(values) + Some(s.f64().ok()?.into_iter().collect()) } /// Pre-extracted list-column inputs with automatic literal broadcasting. /// /// Polars expression plugins receive `&[Series]` as-is from the engine, /// without broadcasting length-1 literals. This struct handles that -/// uniformly: all inputs must be `List` columns with a numeric inner -/// dtype (cast to f64), of either length N or length 1 (broadcast to N). +/// uniformly: all inputs must be `List`/`Array` columns with a float +/// inner dtype (computed in f64), of either length N or length 1 +/// (broadcast to N). /// /// [`ListInputs::new`] nulls out rows whose lists contain null elements /// (for transforms, where a missing sample invalidates the signal); @@ -211,17 +207,14 @@ impl ListInputs { ComputeError: "expected at least one input" ); + // A runtime backstop: each function's output-type function has + // already rejected non-float inners at plan time with the + // parameter's name. let list_f64 = DataType::List(Box::new(DataType::Float64)); let cas = inputs .iter() .map(|s| { - let (_, inner) = Container::split(s.dtype(), "input")?; - polars_ensure!( - inner.is_primitive_numeric(), - ComputeError: - "expected List or Array column with numeric inner dtype, got {}", - s.dtype() - ); + require_float_inner(s.dtype(), "input")?; as_list(s)?.cast(&list_f64) }) .collect::>>()?; @@ -265,12 +258,24 @@ impl ListInputs { /// /// The closure receives one `&[f64]` per input column (length-1 /// literals broadcast) and returns `Ok(None)` to emit a null row, or -/// `Err` to abort the whole query. Null input rows stay null. The result -/// is restored to `out_container`, which the caller derives from -/// whichever input governs the output shape. +/// `Err` to abort the whole query. Null input rows stay null. +/// +/// `allow_empty` is the per-input empty-list policy, one flag per +/// input. An empty list where it is disallowed raises: it is the +/// limiting case of too-short, and a raise policy that softened as the +/// input got worse would be indefensible. (`apply_interp` allows an +/// empty `xp` — zero query points is a valid question with an empty +/// answer — while an empty `x`/`y` is no interpolant at all.) +/// +/// The result is computed in `f64` and restored to `out_container` and +/// `out_inner` at the exit boundary, both derived by the caller from +/// whichever input governs the output. pub(crate) fn apply_list_transform( inputs: &[Series], + name: &'static str, + allow_empty: &[bool], out_container: Container, + out_inner: &DataType, f: F, ) -> PolarsResult where @@ -286,13 +291,34 @@ where ); for i in 0..li.len() { - match li.row(i).map(|slices| f(&slices)).transpose()?.flatten() { + let row = match li.row(i) { + None => { + builder.append_null(); + continue; + } + Some(slices) => slices, + }; + for (slice, &allowed) in row.iter().zip(allow_empty) { + polars_ensure!( + allowed || !slice.is_empty(), + ComputeError: + "{name}: a row's list is empty; a transform needs at least one \ + sample (pre-filter with list.len() to drop such rows)" + ); + } + match f(&row)? { Some(result) => builder.append_slice(&result), None => builder.append_null(), } } - restore(builder.finish().into_series(), out_container) + let s = builder.finish().into_series(); + let s = if *out_inner == DataType::Float64 { + s + } else { + s.cast(&DataType::List(Box::new(out_inner.clone())))? + }; + restore(s, out_container) } /// One `List` input for the dtype-preserving entry layer used by diff --git a/tests/test_array_container.py b/tests/test_array_container.py index 3901f77..2aa5dff 100644 --- a/tests/test_array_container.py +++ b/tests/test_array_container.py @@ -340,14 +340,20 @@ def test_cum_agg_runs_follows_the_value_column_not_both(): @pytest.mark.parametrize( "inner", [pl.Int8, pl.Int16, pl.UInt8, pl.UInt16, pl.Int32, pl.Int64] ) -def test_narrow_integer_inners_do_not_panic(inner): - # polars gates the narrow integer dtypes behind cargo features, and - # reconstructs the Series across the FFI boundary before any of this - # library's code runs -- so a missing feature panics there, exactly - # as the Array dtype once did. +def test_integer_inners_are_rejected_not_panicked(inner): + # Two layers guard this. polars gates the narrow integer dtypes + # behind cargo features and reconstructs the Series across the FFI + # boundary before any of this library's code runs -- a missing + # feature panics there, exactly as the Array dtype once did, which + # is why the crate compiles with dtype-full. Once visible, integers + # are then rejected by our own dtype rules: a clean typed error at + # plan time, never a panic. df = pl.DataFrame({"a": [[1, 2, 3, 4]]}, schema={"a": pl.List(inner)}) - out = df.select(polist.agg_slices("a", "a", aggregation="sum").alias("r")) - assert out["r"][0] == 10.0 + with pytest.raises(polars.exceptions.PolarsError) as excinfo: + df.select(polist.agg_slices("a", "a", aggregation="sum").alias("r")) + message = str(excinfo.value) + assert "Float32 or Float64" in message + assert "the plugin panicked" not in message @pytest.mark.parametrize( diff --git a/tests/test_dtype_preservation.py b/tests/test_dtype_preservation.py new file mode 100644 index 0000000..dff429f --- /dev/null +++ b/tests/test_dtype_preservation.py @@ -0,0 +1,197 @@ +"""Inner-dtype preservation across every function. + +Dtype preservation is an interface property, not a computation +constraint: kernels compute in f64 and round once at the exit boundary, +so Float32 results are deterministic rounded-once values. The single +exception is `count`, which emits UInt32 -- polars' count dtype -- +vertically and cumulatively.""" +import polars as pl +import polars.exceptions +import pytest + +import polars_list_utils as polist + +FS = 200.0 +SIGNAL32 = [float(i % 5) for i in range(32)] + + +def frame(inner: object) -> pl.DataFrame: + return pl.DataFrame( + {"s": [SIGNAL32], "i": [[float(i) for i in range(32)]]}, + schema={"s": pl.List(inner), "i": pl.List(inner)}, + ) + + +# ------------------------------------------------------ Float32 preservation + +@pytest.mark.parametrize("inner", [pl.Float32, pl.Float64]) +def test_transforms_preserve_the_inner_dtype(inner): + df = frame(inner) + fft = df.select(polist.apply_fft("s", sample_rate=FS).alias("r")) + assert fft["r"].dtype == pl.List(inner) + filt = df.select( + polist.apply_butterworth("s", sample_rate=FS, max_freq=50.0).alias("r") + ) + assert filt["r"].dtype == pl.List(inner) + + +def test_interp_inner_dtype_follows_y(): + # Values determine the value dtype: x and xp are coordinates. + df = pl.DataFrame( + {"x": [[0.0, 1.0, 2.0]], "y": [[0.0, 10.0, 20.0]]}, + schema={"x": pl.List(pl.Float64), "y": pl.List(pl.Float32)}, + ) + out = df.select(polist.apply_interp("x", "y", pl.lit([0.5, 1.5])).alias("r")) + assert out["r"].dtype == pl.List(pl.Float32) + assert out["r"][0].to_list() == pytest.approx([5.0, 15.0]) + + # And the other way round: f32 coordinates, f64 values -> f64 out. + df = pl.DataFrame( + {"x": [[0.0, 1.0, 2.0]], "y": [[0.0, 10.0, 20.0]]}, + schema={"x": pl.List(pl.Float32), "y": pl.List(pl.Float64)}, + ) + out = df.select(polist.apply_interp("x", "y", pl.lit([0.5])).alias("r")) + assert out["r"].dtype == pl.List(pl.Float64) + + +@pytest.mark.parametrize("inner", [pl.Float32, pl.Float64]) +@pytest.mark.parametrize( + "agg", ["sum", "mean", "median", "std", "min", "max", "delta"] +) +def test_agg_slices_scalar_keeps_the_value_dtype(inner, agg): + df = frame(inner) + out = df.select( + polist.agg_slices("s", "i", aggregation=agg).alias("r") + ) + assert out["r"].dtype == inner, agg + + +def test_agg_slices_float32_values_are_rounded_once(): + # Computed in f64, cast at the boundary -- the value is the f64 + # answer rounded, not an f32-accumulated one. + df = frame(pl.Float32) + out = df.select(polist.agg_slices("s", "i", aggregation="mean").alias("r")) + expected = sum(SIGNAL32) / len(SIGNAL32) + assert out["r"][0] == pytest.approx(expected) + + +# ------------------------------------------------------------- count dtype + +@pytest.mark.parametrize("inner", [pl.Float32, pl.Float64]) +def test_agg_slices_count_emits_uint32(inner): + df = frame(inner) + out = df.select(polist.agg_slices("s", "i", aggregation="count").alias("r")) + assert out["r"].dtype == pl.UInt32 + assert out["r"][0] == 32 + + +def test_agg_lists_count_emits_uint32(): + df = pl.DataFrame( + {"g": [1, 1], "v": [[1.0, None], [2.0, 3.0]]}, + schema={"g": pl.Int64, "v": pl.List(pl.Float64)}, + ) + out = df.group_by("g").agg( + polist.agg_lists("v", list_length=2, aggregation="count").alias("r") + ) + assert out["r"].dtype == pl.List(pl.UInt32) + assert out["r"][0].to_list() == [2, 1] + + +def test_both_aggregators_agree_on_count_dtype(): + # One exception to preservation, applied uniformly: UInt32 from the + # slice reduction, the group aggregation, and the scan alike. + horizontal = pl.DataFrame({"v": [[1.0, 2.0]], "i": [[0.0, 1.0]]}).select( + polist.agg_slices("v", "i", aggregation="count").alias("r") + )["r"] + vertical = ( + pl.DataFrame({"g": [1], "v": [[1.0, 2.0]]}) + .group_by("g") + .agg(polist.agg_lists("v", list_length=2, aggregation="count").alias("r"))["r"] + ) + scan = pl.DataFrame({"v": [[1.0, 2.0]]}).select( + polist.cum_agg_runs("v", pl.lit([True, True]), aggregation="count").alias("r") + )["r"] + assert horizontal.dtype == pl.UInt32 + assert vertical.dtype == pl.List(pl.UInt32) + assert scan.dtype == pl.List(pl.UInt32) + + +# ------------------------------------------------------- integer rejection + +def test_agg_lists_integer_inner_raises(): + df = pl.DataFrame( + {"g": [1], "v": [[1, 2]]}, schema={"g": pl.Int64, "v": pl.List(pl.Int64)} + ) + with pytest.raises(polars.exceptions.PolarsError, match="Float32 or Float64"): + df.group_by("g").agg( + polist.agg_lists("v", list_length=2, aggregation="mean") + ) + + +@pytest.mark.parametrize( + "build", + [ + lambda d: d.select(polist.apply_fft("s", sample_rate=FS)), + lambda d: d.select( + polist.apply_butterworth("s", sample_rate=FS, max_freq=50.0) + ), + lambda d: d.select(polist.agg_slices("s", "s", aggregation="mean")), + ], +) +def test_integer_rejection_lands_at_plan_time(build): + lf = pl.LazyFrame({"s": [[1, 2, 3, 4]]}, schema={"s": pl.List(pl.Int64)}) + with pytest.raises(polars.exceptions.PolarsError, match="Float32 or Float64"): + build(lf).collect_schema() + + +# --------------------------------------------- plan-time shape checks (Array) + +def test_fft_non_power_of_two_array_width_raises_at_plan_time(): + lf = pl.LazyFrame( + {"s": [[1.0] * 6]}, schema={"s": pl.Array(pl.Float64, 6)} + ) + with pytest.raises(polars.exceptions.PolarsError, match="power-of-two"): + lf.select(polist.apply_fft("s", sample_rate=FS)).collect_schema() + + +def test_butterworth_short_array_width_raises_at_plan_time(): + # A fixed width shorter than the reflection padding means every row + # would fail; the schema knows it, so the query never starts. + lf = pl.LazyFrame( + {"s": [[1.0] * 8]}, schema={"s": pl.Array(pl.Float64, 8)} + ) + with pytest.raises(polars.exceptions.PolarsError, match="reflection padding"): + lf.select( + polist.apply_butterworth("s", sample_rate=FS, max_freq=50.0) + ).collect_schema() + + +def test_butterworth_negative_sample_rate_raises(): + df = pl.DataFrame({"s": [SIGNAL32]}) + with pytest.raises(polars.exceptions.PolarsError, match="sample_rate"): + df.select(polist.apply_butterworth("s", sample_rate=-1.0, max_freq=50.0)) + + +# ---------------------------------------------------------- interp empties + +def test_interp_empty_xp_returns_an_empty_list(): + # Zero query points is a valid question with an empty answer. + df = pl.DataFrame({"x": [[0.0, 1.0]], "y": [[0.0, 10.0]]}) + out = df.select( + polist.apply_interp( + "x", "y", pl.lit([], dtype=pl.List(pl.Float64)) + ).alias("r") + ) + assert out["r"][0].to_list() == [] + + +def test_interp_empty_x_or_y_raises(): + # No interpolant at all is structural, not data. + empty = pl.DataFrame( + {"e": [[]], "v": [[0.0, 1.0]]}, + schema={"e": pl.List(pl.Float64), "v": pl.List(pl.Float64)}, + ) + with pytest.raises(polars.exceptions.PolarsError, match="empty"): + empty.select(polist.apply_interp("e", "e", pl.lit([0.5]))) + with pytest.raises(polars.exceptions.PolarsError, match="empty"): + empty.select(polist.apply_interp("v", "e", pl.lit([0.5]))) diff --git a/tests/test_plugins.py b/tests/test_plugins.py index 4f4624a..4aed30f 100644 --- a/tests/test_plugins.py +++ b/tests/test_plugins.py @@ -42,30 +42,82 @@ def test_fft_hann_alias_matches_hanning(): assert df["a"][0].to_list() == df["b"][0].to_list() -def test_fft_power_reads_mean_square(): +def test_fft_spectrum_reads_mean_square(): # scipy 'spectrum' convention: a tone of amplitude A reads A^2 / 2. df = pl.DataFrame({"s": [sine(25.0, amplitude=2.0)]}).with_columns( - polist.apply_fft("s", sample_rate=FS, scaling="power").alias("p") + polist.apply_fft("s", sample_rate=FS, scaling="spectrum").alias("p") ) assert df["p"][0].to_list()[peak_bin(25.0)] == pytest.approx(2.0, rel=1e-9) -def test_fft_psd_satisfies_parseval(): +def test_fft_amplitude_squared_is_amplitude_squared(): + # Defined as amplitude^2 per bin, one-sided doubling squared along + # with it -- the quantity "power" used to be misread as. + df = pl.DataFrame({"s": [sine(25.0, amplitude=3.0)]}).with_columns( + polist.apply_fft( + "s", sample_rate=FS, window="hann", scaling="amplitude" + ).alias("a"), + polist.apply_fft( + "s", sample_rate=FS, window="hann", scaling="amplitude_squared" + ).alias("a2"), + ) + amp, amp2 = df["a"][0].to_list(), df["a2"][0].to_list() + assert amp2 == pytest.approx([v * v for v in amp], rel=1e-12) + assert amp2[peak_bin(25.0)] == pytest.approx(9.0, rel=1e-9) + + +def test_fft_density_satisfies_parseval(): # For any signal, sum(PSD) * df == mean(x^2) exactly (no window). rng = np.random.default_rng(42) signal = list(rng.standard_normal(N)) df = pl.DataFrame({"s": [signal]}).with_columns( - polist.apply_fft("s", sample_rate=FS, scaling="psd").alias("psd") + polist.apply_fft("s", sample_rate=FS, scaling="density").alias("psd") ) integral = sum(df["psd"][0].to_list()) * FS / N assert integral == pytest.approx(float(np.mean(np.square(signal))), rel=1e-9) -def test_fft_non_power_of_two_yields_null(): - df = pl.DataFrame({"s": [[1.0] * 100]}).with_columns( - polist.apply_fft("s", sample_rate=FS).alias("a") +def test_fft_non_power_of_two_raises(): + # Wrong shape is structural, not data: previously a null row. The + # escape hatch for mixed-length List data is list.len() pre-filtering. + with pytest.raises(polars.exceptions.PolarsError, match="power-of-two"): + pl.DataFrame({"s": [[1.0] * 100]}).with_columns( + polist.apply_fft("s", sample_rate=FS).alias("a") + ) + + +def test_fft_length_one_raises(): + # A length-1 signal IS a power of two, but the windows sum to zero + # on it; the raise is uniform across windows. + with pytest.raises(polars.exceptions.PolarsError, match="at least 2"): + pl.DataFrame({"s": [[5.0]]}).with_columns( + polist.apply_fft("s", sample_rate=FS, window="hann").alias("a") + ) + + +def test_fft_nan_row_yields_all_nan_spectrum(): + # Invalid data propagates instead of becoming missing data: every + # bin sums all samples, so one NaN contaminates the whole spectrum. + signal = sine(25.0) + signal[100] = float("nan") + df = pl.DataFrame({"s": [signal]}).with_columns( + polist.apply_fft("s", sample_rate=FS, scaling="amplitude").alias("a") ) - assert df["a"][0] is None + spectrum = df["a"][0].to_list() + assert spectrum is not None + assert all(math.isnan(v) for v in spectrum) + + +def test_transform_empty_list_raises(): + # The limiting case of too-short: a raise policy that softened as + # the input got worse would be indefensible. + df = pl.DataFrame({"s": [[]]}, schema={"s": pl.List(pl.Float64)}) + with pytest.raises(polars.exceptions.PolarsError, match="empty"): + df.with_columns(polist.apply_fft("s", sample_rate=FS).alias("a")) + with pytest.raises(polars.exceptions.PolarsError, match="empty"): + df.with_columns( + polist.apply_butterworth("s", sample_rate=FS, max_freq=50.0).alias("a") + ) def test_fft_unknown_window_raises(): @@ -80,12 +132,14 @@ def test_fft_unknown_window_raises(): def test_fft_unknown_scaling_raises(): + # "power" was the old name; after the rename to scipy's vocabulary + # it must be invalid, not quietly aliased. with pytest.raises(polars.exceptions.PolarsError): pl.DataFrame({"s": [sine(25.0)]}).with_columns( polist.apply_fft( "s", sample_rate=FS, - scaling="density", # ty: ignore[invalid-argument-type] + scaling="power", # ty: ignore[invalid-argument-type] ).alias("a") ) @@ -101,8 +155,12 @@ def test_fft_windowed_psd_enbw_identity(): # For any signal and window, psd == power / enbw_hz elementwise; # the periodic hann window's ENBW is exactly 1.5 bins. df = pl.DataFrame({"s": [sine(25.0)]}).with_columns( - polist.apply_fft("s", sample_rate=FS, window="hann", scaling="power").alias("pow"), - polist.apply_fft("s", sample_rate=FS, window="hann", scaling="psd").alias("psd"), + polist.apply_fft( + "s", sample_rate=FS, window="hann", scaling="spectrum" + ).alias("pow"), + polist.apply_fft( + "s", sample_rate=FS, window="hann", scaling="density" + ).alias("psd"), ) enbw_hz = 1.5 * FS / N power, psd = df["pow"][0].to_list(), df["psd"][0].to_list() @@ -146,26 +204,35 @@ def test_butterworth_invalid_cutoff_raises(): ) -def test_butterworth_short_list_yields_null(): - # 12 samples equals the order-4 reflection padding, which used to - # panic inside the butterworth crate; must now be a null row. - df = pl.DataFrame({"s": [[1.0] * 12, sine(5.0)]}).with_columns( +def test_butterworth_short_list_raises(): + # 12 samples equals the order-4 reflection padding, which once + # panicked inside the butterworth crate and then yielded a null row; + # wrong shape is structural, so it now raises. + with pytest.raises(polars.exceptions.PolarsError, match="reflection padding"): + pl.DataFrame({"s": [[1.0] * 12, sine(5.0)]}).with_columns( + polist.apply_butterworth("s", sample_rate=FS, max_freq=20.0).alias("f") + ) + + +def test_butterworth_short_rows_can_be_prefiltered(): + # The documented escape hatch for mixed-length List data. + df = pl.DataFrame({"s": [[1.0] * 12, sine(5.0)]}) + out = df.filter(pl.col("s").list.len() >= 13).with_columns( polist.apply_butterworth("s", sample_rate=FS, max_freq=20.0).alias("f") ) - assert df["f"][0] is None - assert df["f"][1] is not None + assert out.height == 1 + assert out["f"][0] is not None -def test_butterworth_bandpass_short_list_yields_null(): +def test_butterworth_bandpass_short_list_raises(): # Bandpass doubles the design order, so order 4 pads with 24 samples; - # a 24-sample row is exactly the panic length and must yield null. - df = pl.DataFrame({"s": [[1.0] * 24, sine(25.0, n=32)]}).with_columns( - polist.apply_butterworth( - "s", sample_rate=FS, min_freq=10.0, max_freq=40.0 - ).alias("f") - ) - assert df["f"][0] is None - assert df["f"][1] is not None + # a 24-sample row is exactly the old panic length and must raise. + with pytest.raises(polars.exceptions.PolarsError, match="reflection padding"): + pl.DataFrame({"s": [[1.0] * 24, sine(25.0, n=32)]}).with_columns( + polist.apply_butterworth( + "s", sample_rate=FS, min_freq=10.0, max_freq=40.0 + ).alias("f") + ) def test_butterworth_order_zero_raises(): @@ -380,21 +447,27 @@ def test_empty_zero_is_not_fill_null(): assert out["filled"].to_list() == [0.0, 0.0] -def test_empty_list_rows_are_null_rows_under_both_settings(): - # Today an empty list is a null row rather than an empty selection, - # so `empty` (and count's 0) cannot reach it. Spec section 7 changes - # this in the breaking pass; pin it until then so the divergence - # from agg_lists is deliberate rather than discovered. +def test_empty_list_rows_are_empty_selections(): + # An empty list is the limiting case of "no index fell in range": + # count reads 0, sum follows `empty`, everything else is null. This + # is what makes agg_slices and agg_lists give one answer to the same + # question, and Array(w=0) needs no special case only because it is + # rejected outright. (Previously an empty list was a null row, which + # neither `empty` nor count's 0 could reach.) df = pl.DataFrame( {"v": [[]], "i": [[]]}, schema={"v": pl.List(pl.Float64), "i": pl.List(pl.Float64)}, ) out = df.select( - polist.agg_slices("v", "i", aggregation="sum", empty="zero").alias("sum"), - polist.agg_slices("v", "i", aggregation="count", empty="zero").alias("count"), + polist.agg_slices("v", "i", aggregation="sum").alias("sum_null"), + polist.agg_slices("v", "i", aggregation="sum", empty="zero").alias("sum_zero"), + polist.agg_slices("v", "i", aggregation="count").alias("count"), + polist.agg_slices("v", "i", aggregation="mean").alias("mean"), ) - assert out["sum"][0] is None - assert out["count"][0] is None + assert out["sum_null"][0] is None + assert out["sum_zero"][0] == 0.0 + assert out["count"][0] == 0 + assert out["mean"][0] is None @pytest.mark.parametrize("bad", ["nul", "Null", "zeroes", None, 0, ""]) @@ -621,11 +694,18 @@ def test_null_and_inner_null_rows_yield_null(): assert df["a"][2] is None # inner nulls must not silently misalign -def test_integer_lists_are_accepted(): - df = pl.DataFrame({"x": [[0, 1, 2]], "y": [[0, 10, 20]]}).with_columns( +def test_integer_lists_raise(): + # Previously accepted and silently cast to f64; integer support is + # deliberately out of scope, so the rejection is loud, at plan time, + # and names the parameter. + df = pl.DataFrame({"x": [[0, 1, 2]], "y": [[0, 10, 20]]}) + with pytest.raises(polars.exceptions.PolarsError, match="Float32 or Float64"): + df.with_columns(polist.apply_interp("x", "y", pl.lit([1.5])).alias("yp")) + lf = df.lazy().with_columns( polist.apply_interp("x", "y", pl.lit([1.5])).alias("yp") ) - assert df["yp"][0].to_list() == pytest.approx([15.0]) + with pytest.raises(polars.exceptions.PolarsError, match="x_column"): + lf.collect_schema() def test_array_input_is_supported_not_aborting(): From 831398a006cca23823b4cfeee5c0047b1fb64d1c Mon Sep 17 00:00:00 2001 From: Travis Hammond Date: Wed, 2 Sep 2026 10:49:53 +0200 Subject: [PATCH 11/17] fix: close the review gaps in the semantics pass One missed sibling edit and five smaller findings from the adversarial review of the dtype and error-policy pass. The README's empty-flag section still described the pre-2.0 rule that a row with empty lists stays null -- three paragraphs after the same file correctly said an empty list is now an empty selection, and one file away from the docstring this pass had already corrected. The paragraph now matches both. apply_interp's closing Returns line likewise still promised List[f64], contradicting the dtype rule two paragraphs above it in its own docstring. Butterworth parameter errors now all carry the function name and the offending keyword: the cutoff range is checked in the plugin rather than left to the butterworth crate, whose errors name neither and leak a transitive dependency's Debug enum into a user-facing message. A NaN cutoff fails the range check and raises too. min_samples saturates instead of wrapping, so an absurd filter_order can no longer arithmetically defeat the plan-time Array guard in a release build, and filter_order >= 1 is validated at plan time alongside sample_rate. The escape-hatch docs now say the list.len() pre-filter must sit upstream of the transform: in a lazy query the optimizer pushes a later filter below the plugin, so offending rows are dropped before validation rather than raising, and the eager and lazy outcomes of the post-filter form differ. A streaming variant of the pre-filter test pins the direction that works. Two coverage holes: the amplitude_squared test ran on a zero-mean tone whose DC and Nyquist bins sit at 1e-15, where the default absolute tolerance would accept a wrong one-sided factor at exactly the two bins the doubling skips -- the signal now carries a DC offset and the two bins are asserted explicitly, along with the documented spectrum-vs-amplitude-squared wrinkle at DC. And the rename test only probed "power": "psd" must be rejected too, and now is. Co-Authored-By: Claude Fable 5 --- README.md | 14 +++++----- polars_list_utils/__init__.py | 15 ++++++++--- src/apply_butterworth.rs | 5 ++++ src/filter.rs | 51 +++++++++++++++++++++++++++-------- tests/test_plugins.py | 44 +++++++++++++++++++++++++----- 5 files changed, 101 insertions(+), 28 deletions(-) diff --git a/README.md b/README.md index 22a261e..c37b832 100644 --- a/README.md +++ b/README.md @@ -82,7 +82,9 @@ any input and broadcast it across rows. Anything structural raises instead of si returning nulls: invalid configuration (bad cutoffs, unknown window names), mismatched paired lengths, integer inner dtypes, and per-row shapes a transform cannot process — a non-power-of-two FFT length, a Butterworth signal shorter than its reflection padding, an -empty list. For mixed-length List data the escape hatch is pre-filtering with `list.len()`. +empty list. For mixed-length List data the escape hatch is pre-filtering with `list.len()` +*before* the transform — in a lazy query the optimizer pushes a later filter below the +transform, so offending rows are dropped before validation rather than raising. Only genuinely missing data nulls: a null row, or a null element inside a transform's signal. @@ -143,11 +145,11 @@ It reaches exactly one cell of the matrix: `sum` is the only aggregation with a identity, every other one yields null over nothing either way, and `count` yields 0 either way. -The switch governs an empty *selection* — the lists held elements, but none was selected. -A null row, and a row whose lists are themselves empty, stay null under both settings for -every aggregation. So `empty="zero"` is not the same as `.fill_null(0.0)`: that cannot tell -a missing row from an empty selection, and would turn a missing measurement into a real -zero, which is the error this flag exists to avoid. +The switch governs an empty *selection* — no index fell in range, or the lists were +themselves empty, the limiting case of the same condition. A null row stays null under both +settings for every aggregation, `count` included. So `empty="zero"` is not the same as +`.fill_null(0.0)`: that cannot tell a missing row from an empty selection, and would turn a +missing measurement into a real zero, which is the error this flag exists to avoid. The same philosophy sets the `strict` default. `apply_interp` verifies that `x` is non-decreasing, `agg_slices` rejects inverted or NaN range bounds at expression diff --git a/polars_list_utils/__init__.py b/polars_list_utils/__init__.py index 263c728..b5f38e1 100644 --- a/polars_list_utils/__init__.py +++ b/polars_list_utils/__init__.py @@ -148,8 +148,9 @@ def apply_interp( and propagates into the output rather than raising — the check exists for the failure you cannot see, not the one you can. - Returns a `List[f64]` column of interpolated y values, one per - `xp` coordinate. Raises if x and y lists differ in length. + Returns one interpolated y value per `xp` coordinate, in the + container and inner dtype the paragraph above describes. Raises if + x and y lists differ in length. """ return _plugin( "apply_interp", @@ -179,7 +180,10 @@ def apply_butterworth( order samples (the reflection padding) raises — wrong shape is structural, not data; for an `Array` the width is in the schema and the raise lands at plan time, and for mixed-length `List` data the - escape hatch is pre-filtering with `list.len()`. Non-finite samples + escape hatch is pre-filtering with `list.len()` before this + transform (in a lazy query the optimizer pushes a later filter + below it, dropping offending rows before validation rather than + raising). Non-finite samples propagate and contaminate the filtered output rather than nulling it. """ @@ -206,7 +210,10 @@ def apply_fft( `[i * sample_rate / N for i in range(N // 2 + 1)]`. Signal lengths must be powers of two of at least 2 — a violation raises (for an `Array`, already at plan time, the width being in the schema; for - mixed-length `List` data, pre-filter with `list.len()`). Non-finite + mixed-length `List` data, pre-filter with `list.len()` *before* + this transform — in a lazy query, a filter written after it is + pushed below it by the optimizer, so offending rows are dropped + before validation rather than raising). Non-finite samples are legitimate floats and propagate: every bin sums all samples, so one NaN yields an all-NaN spectrum rather than a null row — invalid data stays visibly invalid instead of becoming diff --git a/src/apply_butterworth.rs b/src/apply_butterworth.rs index 44911c8..6bf8d2b 100644 --- a/src/apply_butterworth.rs +++ b/src/apply_butterworth.rs @@ -42,6 +42,11 @@ fn butterworth_output( ComputeError: "apply_butterworth: sample_rate must be positive, got {}", kwargs.sample_rate ); + polars_ensure!( + kwargs.order() >= 1, + ComputeError: "apply_butterworth: filter_order must be at least 1, got {}", + kwargs.order() + ); if let (Container::Array(w), Some(filter)) = (container, kwargs.filter()) { let min_samples = filter.min_samples(kwargs.order()); polars_ensure!( diff --git a/src/filter.rs b/src/filter.rs index 3aa4374..48dbb77 100644 --- a/src/filter.rs +++ b/src/filter.rs @@ -28,7 +28,7 @@ impl Filter { order: usize, ) -> usize { match self { - Filter::Bandpass(..) => 2 * order, + Filter::Bandpass(..) => order.saturating_mul(2), _ => order, } } @@ -36,18 +36,28 @@ impl Filter { /// Minimum sample count accepted by `bidirectional`, which reflects /// `3 * effective_order` samples of padding onto each end of the /// signal (and panics on exactly that length, hence the `+ 1`). + /// + /// Saturating: release builds wrap on overflow, so an absurd + /// `filter_order` could otherwise wrap this to a tiny number and + /// arithmetically defeat the very guard that keeps the crate from + /// indexing past the end of the data. pub(crate) fn min_samples( self, order: usize, ) -> usize { - 3 * self.effective_order(order) + 1 + self.effective_order(order) + .saturating_mul(3) + .saturating_add(1) } /// Validate parameters and design the Butterworth filter once. /// - /// Cutoffs outside `(0, sample_rate / 2)` are rejected by the - /// butterworth crate; order and cutoff ordering are checked here - /// because the crate panics instead of erroring on them. + /// All parameter errors carry the function name and the offending + /// keyword. The cutoff range is checked here rather than left to + /// the butterworth crate, whose errors name neither and would leak + /// a transitive dependency's Debug enum into a user-facing message + /// (the `map_err` below is an unreachable-in-practice backstop). A + /// NaN cutoff fails the range check and raises too. pub(crate) fn build( self, sample_rate: f64, @@ -55,15 +65,34 @@ impl Filter { ) -> PolarsResult { polars_ensure!( order >= 1, - ComputeError: "filter_order must be at least 1, got {order}" + ComputeError: "apply_butterworth: filter_order must be at least 1, got {order}" ); - if let Filter::Bandpass(low, high) = self { + let nyquist = sample_rate / 2.0; + let in_range = |name: &str, freq: f64| -> PolarsResult<()> { polars_ensure!( - low < high, - ComputeError: "bandpass requires min_freq < max_freq, got {low} >= {high}" + freq > 0.0 && freq < nyquist, + ComputeError: + "apply_butterworth: {name} must be in (0, sample_rate/2 = {nyquist}), \ + got {freq}" ); + Ok(()) + }; + match self { + Filter::Highpass(freq) => in_range("min_freq", freq)?, + Filter::Lowpass(freq) => in_range("max_freq", freq)?, + Filter::Bandpass(low, high) => { + in_range("min_freq", low)?; + in_range("max_freq", high)?; + polars_ensure!( + low < high, + ComputeError: + "apply_butterworth: bandpass requires min_freq < max_freq, \ + got {low} >= {high}" + ); + } } - ButterworthFilter::new(order, sample_rate, self.to_cutoff()) - .map_err(|e| polars_err!(ComputeError: "failed to create filter: {e:?}")) + ButterworthFilter::new(order, sample_rate, self.to_cutoff()).map_err( + |e| polars_err!(ComputeError: "apply_butterworth: failed to create filter: {e:?}"), + ) } } diff --git a/tests/test_plugins.py b/tests/test_plugins.py index 4aed30f..91a5e50 100644 --- a/tests/test_plugins.py +++ b/tests/test_plugins.py @@ -52,18 +52,35 @@ def test_fft_spectrum_reads_mean_square(): def test_fft_amplitude_squared_is_amplitude_squared(): # Defined as amplitude^2 per bin, one-sided doubling squared along - # with it -- the quantity "power" used to be misread as. - df = pl.DataFrame({"s": [sine(25.0, amplitude=3.0)]}).with_columns( + # with it -- the quantity "power" used to be misread as. The DC + # offset matters: a zero-mean tone leaves the DC and Nyquist bins + # near 1e-15, where any default absolute tolerance would accept a + # wrong one-sided factor at exactly the two bins the doubling + # skips. + signal = [2.0 + v for v in sine(25.0, amplitude=3.0)] + df = pl.DataFrame({"s": [signal]}).with_columns( polist.apply_fft( "s", sample_rate=FS, window="hann", scaling="amplitude" ).alias("a"), polist.apply_fft( "s", sample_rate=FS, window="hann", scaling="amplitude_squared" ).alias("a2"), + polist.apply_fft( + "s", sample_rate=FS, window="hann", scaling="spectrum" + ).alias("sp"), ) amp, amp2 = df["a"][0].to_list(), df["a2"][0].to_list() assert amp2 == pytest.approx([v * v for v in amp], rel=1e-12) + assert amp2[0] == pytest.approx(amp[0] ** 2, rel=1e-12) + assert amp2[-1] == pytest.approx(amp[-1] ** 2, rel=1e-12) assert amp2[peak_bin(25.0)] == pytest.approx(9.0, rel=1e-9) + assert amp2[0] == pytest.approx(4.0, rel=1e-9) # the DC offset, squared + # And the documented wrinkle: spectrum == amplitude^2 / 2 at + # interior bins, but NOT at DC, where the doubling is absent from + # both scalings and the halving therefore has nothing to cancel. + sp = df["sp"][0].to_list() + assert sp[peak_bin(25.0)] == pytest.approx(amp2[peak_bin(25.0)] / 2.0, rel=1e-9) + assert sp[0] == pytest.approx(amp2[0], rel=1e-9) def test_fft_density_satisfies_parseval(): @@ -131,15 +148,16 @@ def test_fft_unknown_window_raises(): ) -def test_fft_unknown_scaling_raises(): - # "power" was the old name; after the rename to scipy's vocabulary - # it must be invalid, not quietly aliased. +@pytest.mark.parametrize("old_name", ["power", "psd"]) +def test_removed_scaling_names_are_rejected(old_name): + # The old names must be invalid, not quietly aliased -- the + # ambiguous name disappearing is the point of the rename. with pytest.raises(polars.exceptions.PolarsError): pl.DataFrame({"s": [sine(25.0)]}).with_columns( polist.apply_fft( "s", sample_rate=FS, - scaling="power", # ty: ignore[invalid-argument-type] + scaling=old_name, ).alias("a") ) @@ -215,7 +233,8 @@ def test_butterworth_short_list_raises(): def test_butterworth_short_rows_can_be_prefiltered(): - # The documented escape hatch for mixed-length List data. + # The documented escape hatch for mixed-length List data: filter + # upstream of the transform. Pinned in eager and streaming alike. df = pl.DataFrame({"s": [[1.0] * 12, sine(5.0)]}) out = df.filter(pl.col("s").list.len() >= 13).with_columns( polist.apply_butterworth("s", sample_rate=FS, max_freq=20.0).alias("f") @@ -223,6 +242,17 @@ def test_butterworth_short_rows_can_be_prefiltered(): assert out.height == 1 assert out["f"][0] is not None + streamed = ( + df.lazy() + .filter(pl.col("s").list.len() >= 13) + .with_columns( + polist.apply_butterworth("s", sample_rate=FS, max_freq=20.0).alias("f") + ) + .collect(engine="streaming") + ) + assert streamed.height == 1 + assert streamed["f"][0] is not None + def test_butterworth_bandpass_short_list_raises(): # Bandpass doubles the design order, so order 4 pads with 24 samples; From 98da6b1431c9278085faa13b1286e90dca83d16c Mon Sep 17 00:00:00 2001 From: Travis Hammond Date: Wed, 2 Sep 2026 10:50:21 +0200 Subject: [PATCH 12/17] chore: silence ty on the streaming collect union type The house pattern for collect(engine=...), whose stub returns InProcessQuery | DataFrame. Co-Authored-By: Claude Fable 5 --- tests/test_plugins.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_plugins.py b/tests/test_plugins.py index 91a5e50..b253301 100644 --- a/tests/test_plugins.py +++ b/tests/test_plugins.py @@ -250,8 +250,8 @@ def test_butterworth_short_rows_can_be_prefiltered(): ) .collect(engine="streaming") ) - assert streamed.height == 1 - assert streamed["f"][0] is not None + assert streamed.height == 1 # ty: ignore[unresolved-attribute] + assert streamed["f"][0] is not None # ty: ignore[not-subscriptable] def test_butterworth_bandpass_short_list_raises(): From 4d48ce1bceea3df6bebe097959cad9df5eb310dc Mon Sep 17 00:00:00 2001 From: Travis Hammond Date: Wed, 2 Sep 2026 11:26:07 +0200 Subject: [PATCH 13/17] chore: keep the prefade migration notes out of git alongside the spec Co-Authored-By: Claude Fable 5 --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index deea853..7af0ed3 100644 --- a/.gitignore +++ b/.gitignore @@ -11,6 +11,7 @@ target # Design/requirements working documents, kept local polars_list_utils_spec.md +prefade_migration_todos.md # These are backup files generated by rustfmt **/*.rs.bk From 4a612e90c5d6e785250d974bbb00df2874c00c36 Mon Sep 17 00:00:00 2001 From: Travis Hammond Date: Wed, 2 Sep 2026 14:01:30 +0200 Subject: [PATCH 14/17] fix: survive modern polars, and close the whole-release review findings The whole-release cross-check found the ground shifting under the suite: polars fixed the very bugs this library had documented itself as deviating from. On >= 1.40 the cum_min/cum_max fold-seed leak is gone (1.44 agrees exactly with this library's prefix rule), 1.44 flipped polars' own min/max tie-break from first-wins to last-wins, and 1.44 widened the Float32 cum_sum accumulator. Eight tests asserted polars' old behavior via live oracles, so CI -- which installs the newest polars, not the lock -- would have failed its gating job and a 2.0.0 tag would have published nothing. The library's semantics do not move: the prefix rule, first-wins ties, and round-once Float32 were always version-independent choices, and they stay. The tests now assert that contract directly -- the oracle computes min/max with its own first-wins fold instead of leaning on a polars whose tie-break drifted, and the two remaining native comparisons branch on the polars version. The suite passes 304/304 on both 1.30 (the lock) and 1.44.1 (what CI installs). The docs that described the deviations as facts about polars now scope them to the versions where they were true. The dev extra records that the tests need polars >= 1.25 for the streaming engine, though the library itself still runs on the declared 1.23 floor. Behavioral fixes from the same review: Statistic::mean now routes through the +0.0-seeded sum, so an all-negative-zero mean reads +0.0 everywhere (agg_slices had kept -0.0 while cum_agg_runs and polars said +0.0); broadcast-length errors carry the calling function's name instead of being the one anonymous message reachable from six plugins; and a typo'd range boundary mode fails at construction naming the parameter instead of surfacing serde's untagged-enum internals. Coverage the review showed to be missing or toothless, now present: even-count median VALUES for agg_slices against polars at the extremes; the kwargs rebuild exercised on slices_exclude alone (a shared object in both kwargs stops aliasing once include is rebuilt, so those tests could not catch a missing exclude rebuild); a +-inf FFT row (a mix of inf and NaN bins, not all-NaN -- the spec bullet said otherwise and now matches the docs); null elements pinned across all six transform entry paths including the butterworth no-cutoff pass-through; agg_lists Float32 preservation; and an agg_slices f32 fixture built around 2^24, where step-wise f32 accumulation and f64-rounded-once actually part company. Docs: README gains the requested authorship disclaimer and a 2.0.0 changelog section; the "six plugin functions" broadcast paragraph is rewritten to say what actually broadcasts and why the count was misleading; the agg_lists bullet no longer claims "pure Polars (no plugin)" and now states that it aggregates in the column's own dtype (overflowing near the Float32 limit where the round-once kernels stay finite); the strict bullet no longer overstates the agg_lists check; and the inverted +-/-+ signs in the seed-leak note are replaced by explicit per-function values. Co-Authored-By: Claude Fable 5 --- README.md | 68 +++++++++++++++++----- polars_list_utils/__init__.py | 31 ++++++---- pyproject.toml | 4 ++ src/agg_slices.rs | 2 +- src/cum_agg_runs.rs | 2 +- src/statistic.rs | 23 +++++--- src/transform.rs | 7 ++- src/util.rs | 32 +++++++---- src/zip_binary.rs | 4 +- tests/test_cum_agg_runs.py | 77 ++++++++++++++++++------- tests/test_dtype_preservation.py | 37 ++++++++++-- tests/test_kwargs_serialization.py | 35 ++++++++++++ tests/test_plugins.py | 90 ++++++++++++++++++++++++++++++ uv.lock | 2 + 14 files changed, 340 insertions(+), 74 deletions(-) diff --git a/README.md b/README.md index c37b832..9d5be3c 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,13 @@ than their Numpy counterparts (although they might be in some cases), the integration with Polars gave my larger processing pipeline a significant speed boost. +## A disclaimer about authorship + +This library started life as a hand-crafted solution to a problem at work. The time +pressures of life and work eventually led me to succumb to taking shortcuts and let AI +take the wheel. This repo is therefore no longer "artisanal": it is AI-generated, reviewed +to the best of my time and abilities, and fully dog-fooded by me at work. + ## Features - `polist.apply_interp` @@ -50,8 +57,12 @@ boost. - `polist.agg_lists` - Applies element-wise list-aggregations to a List-type column in a GroupBy context. - - This is implemented in pure Polars (no plugin), and lives here so that both aggregation - functions offer the same aggregations with the same missing-data behaviour. + - Composed from Polars' own vertical aggregations, with a small pass-through plugin + that normalizes `Array` to `List` and carries the `strict`/`list_length` checks; it + lives here so that both aggregation functions offer the same aggregations with the + same missing-data behaviour. One consequence: it aggregates in the column's own + dtype, as Polars does, rather than in `f64` -- so near the `Float32` range limit it + can overflow where the round-once plugin kernels stay finite. - `polist.zip_binary` - Applies a binary operation element-wise between two List-type columns, per element @@ -70,15 +81,20 @@ boost. so an all-`True` gate over `Float64` tracks `cum_sum`/`cum_min`/`cum_max`/`cum_count`; null values emit null but keep the running state, except `count`, which emits the running count as `cum_count` does. - - Two documented departures from the native cumulatives: `cum_min`/`cum_max` fold from - ∓`f64::MAX` and never replace that seed on a tie, so they report ±1.8e308 where the - prefix is all-NaN or opens with an infinity of the same sign, while this function - reports the true prefix extreme; and `Float32` values accumulate in `f64` and round - once at the boundary, where Polars rounds every step. + - This function has always emitted the true prefix extreme -- NaN for an all-NaN + prefix, or a leading infinity -- where Polars up to 1.43 leaked its `cum_min`/ + `cum_max` fold seed (`+f64::MAX` for `cum_min`, `-f64::MAX` for `cum_max`; briefly + `±inf`) at those positions. Polars >= 1.44 agrees with this function. `Float32` + values accumulate in `f64` and round once at the boundary, so a Polars version that + accumulates step-wise in `Float32` can differ in the last bits and near the range + limit. - `Float32` values stay `Float32`; `count` emits `UInt32` (Polars' count dtype). -The six plugin functions accept a length-1 literal list column (e.g. `pl.lit([...])`) for -any input and broadcast it across rows. Anything structural raises instead of silently +Every function accepts a length-1 literal list column (e.g. `pl.lit([...])`) for any +input and broadcasts it across rows, provided at least one other input is a real column -- +as in Polars generally, an expression built only from length-1 literals produces a single +row rather than broadcasting, so the claim is moot for the single-input transforms and for +`agg_lists`, which aggregates rows rather than mapping them. Anything structural raises instead of silently returning nulls: invalid configuration (bad cutoffs, unknown window names), mismatched paired lengths, integer inner dtypes, and per-row shapes a transform cannot process — a non-power-of-two FFT length, a Butterworth signal shorter than its reflection padding, an @@ -88,8 +104,9 @@ transform, so offending rows are dropped before validation rather than raising. Only genuinely missing data nulls: a null row, or a null element inside a transform's signal. -`Float32` columns stay `Float32` — kernels compute in `f64` and round once at the output -boundary — with one exception: `count` emits `UInt32`, Polars' count dtype, everywhere it +`Float32` columns stay `Float32`; the plugin kernels compute in `f64` and round once at +the output boundary, while `agg_lists` aggregates in the column's own dtype as Polars +does. The one dtype exception: `count` emits `UInt32`, Polars' count dtype, everywhere it appears. Both aggregation functions support `sum`, `mean`, `median`, `std`, `min`, `max`, `delta`, @@ -153,8 +170,9 @@ missing measurement into a real zero, which is the error this flag exists to avo The same philosophy sets the `strict` default. `apply_interp` verifies that `x` is non-decreasing, `agg_slices` rejects inverted or NaN range bounds at expression -construction, and `agg_lists` raises on lists longer than `list_length` rather than -silently truncating them — all **on by default**, since each covers a failure that is +construction, and `agg_lists` raises on lists carrying non-null data past +`list_length` rather than silently truncating it (a null-padded tail is lossless and +never raises) — all **on by default**, since each covers a failure that is otherwise silent. Pass `strict=False` for the lenient readings (numpy's unsorted-`x` behaviour, empty selections from inverted ranges, `list_length` as a deliberate window). NaN, which announces itself in the output, propagates rather than raising either way. @@ -247,3 +265,27 @@ uvx ruff check cargo clippy --release -- -D warnings cargo fmt ``` + +## Changelog + +### 2.0.0 + +One bundled breaking release. New: the `sum` aggregation everywhere; `zip_binary` +(element-wise arithmetic, Kleene `and`/`or`, and total-ordered comparisons between two +list columns); `cum_agg_runs` (cumulative aggregation within gate-delimited runs, matching +Polars' `cum_*` semantics); `Array` columns accepted by every function, with the container +preserved; and an `empty` switch on the aggregations choosing what the sum of nothing is +-- `"null"` (default) or Polars' `"zero"`. + +Breaking: `strict` checks are on by default (unsorted interpolation axes, invalid ranges, +and lossy `agg_lists` truncation now raise; pass `strict=False` for the old leniency). +Wrong shape raises instead of yielding null rows (non-power-of-two FFT lengths, too-short +Butterworth signals, empty lists into transforms -- pre-filter with `list.len()` upstream), +while NaN/±inf now propagate into visibly invalid output instead of becoming null: guard +thresholds with `is_finite()`. `Float32` columns stay `Float32`, `count` returns `UInt32`, +and integer inner dtypes raise. FFT scalings use scipy's names -- `power` -> `spectrum`, +`psd` -> `density`, plus new `amplitude_squared`; window names are unchanged. The +even-count median now matches Polars exactly (last-ulp changes on stored features). + +Fixed: `Array` inputs no longer abort the process; no input dtype can panic; aliased or +JSON-shaped range arguments no longer crash `agg_slices`. diff --git a/polars_list_utils/__init__.py b/polars_list_utils/__init__.py index b5f38e1..3ec84d3 100644 --- a/polars_list_utils/__init__.py +++ b/polars_list_utils/__init__.py @@ -33,7 +33,10 @@ zero-length `List`s are fine. The plugins accept length-1 literal list columns (e.g. `pl.lit(...)`) -for any input and broadcast them. Null rows produce null output rows. +for any input and broadcast them, provided at least one other input is +a real column — as in polars generally, an expression built only from +length-1 literals produces a single row rather than broadcasting. Null +rows produce null output rows. Data conditions follow three deliberate family regimes. A signal is atomic, so for the transforms (`apply_interp`, `apply_butterworth`, @@ -288,6 +291,14 @@ def bound(b: Any) -> tuple[float, str] | float: # Explicit ((value, mode)) form, as a tuple or (from JSON) a list. if isinstance(b, (tuple, list)): value, mode = b + # Validated here so a typo'd mode fails with the parameter's + # name, instead of surfacing serde's untagged-enum internals + # from the plugin boundary. + if mode not in ("closed", "open"): + raise ValueError( + f"agg_slices: boundary mode in {param} must be 'closed' " + f"or 'open', got {mode!r}" + ) return (float(value), str(mode)) return float(b) @@ -508,17 +519,17 @@ def cum_agg_runs( vertical `aggregation` over the run's elements so far — so with an all-`True` gate over `Float64` values, `sum`/`min`/`max`/`count` track polars' `cum_sum`/`cum_min`/`cum_max`/`cum_count`, with two - documented exceptions: + caveats: - - `cum_min`/`cum_max` fold from ∓`f64::MAX` and never replace that - seed on a tie, so they emit ±1.7976931348623157e+308 wherever the - prefix is all-NaN or opens with an infinity of the same sign. - This function follows the prefix rule and emits the true prefix - minimum or maximum (NaN, or that infinity). + - This function has always followed the prefix rule and emitted the + true prefix extreme — NaN for an all-NaN prefix, or a leading + infinity. polars up to 1.43 instead leaked its fold seed at those + positions (`+f64::MAX` from `cum_min`, `-f64::MAX` from + `cum_max`; briefly `±inf`); polars 1.44 fixed it and agrees. - `Float32` values accumulate in `f64` and round once at the output - boundary, while polars rounds every step, so the two can differ - in the last bits and around the `Float32` range limit (where this - function can stay finite as polars overflows to infinity). + boundary, so a polars version that accumulates `cum_*` step-wise + in `Float32` can differ in the last bits and around the `Float32` + range limit. Floating-point accumulation order differs from polars generally, so treat the correspondence as one of semantics — missing data, NaN, diff --git a/pyproject.toml b/pyproject.toml index 8ad121d..0826399 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,10 @@ dev = [ "pytest >= 8.0", "numpy >= 2.2.3", "matplotlib >= 3.4.0", + # The library itself runs on the declared floor (polars >= 1.23), + # but the test suite drives collect(engine="streaming"), which + # polars added in 1.25. + "polars >= 1.25", ] [build-system] diff --git a/src/agg_slices.rs b/src/agg_slices.rs index c42ab6a..ce495e9 100644 --- a/src/agg_slices.rs +++ b/src/agg_slices.rs @@ -67,7 +67,7 @@ fn agg_slices( ) -> PolarsResult { let (_, value_inner) = require_float_inner(inputs[0].dtype(), "agg_slices: value_column")?; - let li = ListInputs::with_inner_nulls(inputs)?; + let li = ListInputs::with_inner_nulls(inputs, "agg_slices")?; // No include ranges means include everything; excludes then carve it up. let mut slice = match kwargs.slices_include { diff --git a/src/cum_agg_runs.rs b/src/cum_agg_runs.rs index 8f8e1ee..8d58b69 100644 --- a/src/cum_agg_runs.rs +++ b/src/cum_agg_runs.rs @@ -232,7 +232,7 @@ fn cum_agg_runs( let values = TypedListInput::::float(&inputs[0], "cum_agg_runs: value_column")?; let gates = TypedListInput::::boolean(&inputs[1], "cum_agg_runs: gate_column")?; let out = run_container(values.container(), gates.container())?; - let len = broadcast_len(&[values.n_rows(), gates.n_rows()])?; + let len = broadcast_len(&[values.n_rows(), gates.n_rows()], "cum_agg_runs")?; let mut rows: Vec>>> = Vec::with_capacity(len); for i in 0..len { diff --git a/src/statistic.rs b/src/statistic.rs index 7b9d3c3..6df5b8a 100644 --- a/src/statistic.rs +++ b/src/statistic.rs @@ -21,9 +21,11 @@ pub(crate) enum Aggregation { /// /// Every other aggregation yields null over nothing under either /// setting, and `count` yields 0 under either, so this switch reaches -/// exactly one cell of the matrix. It applies to a selection that came -/// out empty; a null row or an empty list never reaches this code, -/// having already short-circuited to a null output row. +/// exactly one cell of the matrix. It applies to any selection that +/// came out empty — including a row whose lists were themselves empty, +/// the limiting case of the same condition. Only a null row never +/// reaches this code, having already short-circuited to a null output +/// row. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Deserialize)] #[serde(rename_all = "lowercase")] pub(crate) enum Empty { @@ -56,8 +58,11 @@ impl Aggregation { } /// Running minimum that skips NaN and keeps the **first** of values -/// that compare equal, as polars does. `f64::min` would return the -/// later one, which flips the sign of a `-0.0`/`+0.0` tie. +/// that compare equal — matching polars through 1.43 and this +/// library's own stable-sorted `median`; polars 1.44 changed its own +/// tie-break to last-wins, and this library stays fixed rather than +/// tracking it (visible only in the sign of a `-0.0`/`+0.0` tie). +/// `f64::min` would return the later value. pub(crate) fn min_fold( acc: f64, x: f64, @@ -143,10 +148,10 @@ pub(crate) trait Statistic { } fn mean(&self) -> Option { - if self.is_empty() { - return None; - } - Some(self.values().iter().sum::() / self.values().len() as f64) + // Via `sum`, so the mean of an all-`-0.0` selection reads + // `+0.0` like everywhere else in the library -- a bare + // `Iterator::sum` folds from IEEE's `-0.0` identity. + Some(self.sum()? / self.values().len() as f64) } /// Sample standard deviation (ddof = 1), matching polars' default. diff --git a/src/transform.rs b/src/transform.rs index 57b875e..d942ef9 100644 --- a/src/transform.rs +++ b/src/transform.rs @@ -111,9 +111,10 @@ fn fft_magnitudes(samples: &[f64]) -> PolarsResult> { /// bin sums all samples, so one NaN yields an all-NaN spectrum — invalid /// data stays visibly invalid rather than becoming missing. /// -/// Scaling normalizes by S1 = Σw[n]; interior bins are doubled **in the -/// power domain** to account for the discarded negative frequencies, -/// while DC and Nyquist (which have no mirror bin) are not. +/// Scaling normalizes by S1 = Σw[n]; interior bins are doubled **in +/// each scaling's own domain** (×2 amplitude, ×2 power) to account for +/// the discarded negative frequencies, while DC and Nyquist (which +/// have no mirror bin) are not. pub(crate) fn compute_fft( signal: &[f64], sample_rate: f64, diff --git a/src/util.rs b/src/util.rs index c474e81..602ef2e 100644 --- a/src/util.rs +++ b/src/util.rs @@ -173,25 +173,35 @@ pub(crate) struct ListInputs { } impl ListInputs { - pub(crate) fn new(inputs: &[Series]) -> PolarsResult { - Self::build(inputs, extract_row) + pub(crate) fn new( + inputs: &[Series], + name: &str, + ) -> PolarsResult { + Self::build(inputs, name, extract_row) } } impl ListInputs> { - pub(crate) fn with_inner_nulls(inputs: &[Series]) -> PolarsResult { - Self::build(inputs, extract_row_nullable) + pub(crate) fn with_inner_nulls( + inputs: &[Series], + name: &str, + ) -> PolarsResult { + Self::build(inputs, name, extract_row_nullable) } } /// The broadcast length across inputs: length-1 inputs are literals, -/// all others must agree. -pub(crate) fn broadcast_len(lens: &[usize]) -> PolarsResult { +/// all others must agree. `name` is the calling function, so the error +/// is attributable in a pipeline with several plugin calls. +pub(crate) fn broadcast_len( + lens: &[usize], + name: &str, +) -> PolarsResult { let len = lens.iter().copied().filter(|&l| l != 1).max().unwrap_or(1); for &l in lens { polars_ensure!( l == len || l == 1, - ComputeError: "all inputs must have length {len} or 1 (got {l})" + ComputeError: "{name}: all inputs must have length {len} or 1 (got {l})" ); } Ok(len) @@ -200,6 +210,7 @@ pub(crate) fn broadcast_len(lens: &[usize]) -> PolarsResult { impl ListInputs { fn build( inputs: &[Series], + name: &str, extract: impl Fn(&Series) -> Option>, ) -> PolarsResult { polars_ensure!( @@ -220,7 +231,7 @@ impl ListInputs { .collect::>>()?; let lens = cas.iter().map(|s| s.len()).collect::>(); - let len = broadcast_len(&lens)?; + let len = broadcast_len(&lens, name)?; let columns = cas .iter() @@ -241,7 +252,8 @@ impl ListInputs { } /// All column values for row `i`, with length-1 inputs broadcast. - /// Returns `None` if any column is null (or empty) at this row. + /// Returns `None` if any column is null at this row; an empty list + /// is `Some(&[])`, judged by the caller's `allow_empty` policy. pub(crate) fn row( &self, i: usize, @@ -281,7 +293,7 @@ pub(crate) fn apply_list_transform( where F: Fn(&[&[f64]]) -> PolarsResult>>, { - let li = ListInputs::new(inputs)?; + let li = ListInputs::new(inputs, name)?; let mut builder = ListPrimitiveChunkedBuilder::::new( PlSmallStr::from(""), diff --git a/src/zip_binary.rs b/src/zip_binary.rs index d621016..225c64c 100644 --- a/src/zip_binary.rs +++ b/src/zip_binary.rs @@ -235,7 +235,7 @@ fn zip_binary( let right = TypedListInput::::boolean(&inputs[1], "zip_binary: right_column")?; let out = zip_container(left.container(), right.container())?; - let len = broadcast_len(&[left.n_rows(), right.n_rows()])?; + let len = broadcast_len(&[left.n_rows(), right.n_rows()], "zip_binary")?; let rows = (0..len) .map(|i| zip_row(left.row(i), right.row(i), |a, b| kleene(op, a, b))) .collect::>>()?; @@ -245,7 +245,7 @@ fn zip_binary( let left = TypedListInput::::float(&inputs[0], "zip_binary: left_column")?; let right = TypedListInput::::float(&inputs[1], "zip_binary: right_column")?; let out = zip_container(left.container(), right.container())?; - let len = broadcast_len(&[left.n_rows(), right.n_rows()])?; + let len = broadcast_len(&[left.n_rows(), right.n_rows()], "zip_binary")?; if op.is_arithmetic() { let rows = (0..len) diff --git a/tests/test_cum_agg_runs.py b/tests/test_cum_agg_runs.py index 74741b0..f82cbf5 100644 --- a/tests/test_cum_agg_runs.py +++ b/tests/test_cum_agg_runs.py @@ -12,6 +12,13 @@ NAN = float("nan") +# polars changed two behaviors this suite touches: 1.44 fixed the +# cum_min/cum_max fold-seed leak (agreeing with this library's prefix +# rule) and flipped its min/max tie-break from first-wins to last-wins +# (this library stays first-wins). Tests that compare against native +# polars branch on this. +PL_VERSION = tuple(int(p) for p in pl.__version__.split(".")[:2]) + AGGS = ["sum", "mean", "median", "std", "min", "max", "delta", "count"] INF = float("inf") @@ -85,9 +92,24 @@ def two_pass_std(prefix: list[float]) -> float | None: return math.sqrt(variance) if variance == variance and variance >= 0 else NAN +def first_wins_extreme(prefix: list[float], is_better) -> float: + """NaN-skipping min/max that keeps the FIRST of equal values — + the library's fixed tie-break (and polars' own through 1.43; 1.44 + flipped polars to last-wins, so the oracle cannot lean on live + polars for these two without inheriting a moving tie-break).""" + acc = float("nan") + for x in prefix: + if math.isnan(acc): + acc = x + elif not math.isnan(x) and is_better(x, acc): + acc = x + return acc + + def vertical(agg: str, prefix: list[float]) -> float | int | None: """The vertical aggregation over a run prefix's non-null values, - with polars itself as the kernel (the deference doctrine).""" + with polars itself as the kernel wherever polars' answer is stable + across the supported versions.""" s = pl.Series(prefix, dtype=pl.Float64) if agg == "sum": return s.sum() @@ -95,9 +117,14 @@ def vertical(agg: str, prefix: list[float]) -> float | int | None: return len(prefix) # prefix holds only non-null values if agg == "std": return two_pass_std(prefix) + if agg == "min": + return first_wins_extreme(prefix, lambda x, acc: x < acc) + if agg == "max": + return first_wins_extreme(prefix, lambda x, acc: x > acc) if agg == "delta": - mx, mn = s.max(), s.min() - return None if mx is None else mx - mn # ty: ignore[unsupported-operator] + mx = first_wins_extreme(prefix, lambda x, acc: x > acc) + mn = first_wins_extreme(prefix, lambda x, acc: x < acc) + return mx - mn return getattr(s, agg)() @@ -251,8 +278,11 @@ def test_median_matches_polars_at_extremes(): @pytest.mark.parametrize("values", [[-0.0, 0.0], [0.0, -0.0], [-0.0, -0.0]]) def test_signed_zero_ties_keep_the_first_value(values): - # polars' min/max keep the earlier of two values that compare equal; - # Rust's f64::min/max keep the later, which flips the sign of zero. + # The library's fixed tie-break: min/max/median keep the earlier of + # two values that compare equal, so a signed-zero tie keeps its + # first sign. (polars agreed through 1.43; 1.44 flipped its own + # min/max to last-wins. Rust's f64::min/max would also keep the + # later value.) The expectation is therefore stated directly. df = pl.DataFrame({"v": [values]}, schema={"v": pl.List(pl.Float64)}) for agg in ["min", "max", "median"]: got = df.select( @@ -260,25 +290,26 @@ def test_signed_zero_ties_keep_the_first_value(values): "v", pl.lit([True] * len(values)), aggregation=agg # ty: ignore[invalid-argument-type] ).alias("r") )["r"][0].to_list() - expected = [ - getattr(pl.Series(values[:k], dtype=pl.Float64), agg)() - for k in range(1, len(values) + 1) - ] + # Every prefix's extreme (and two-element median) is its first + # element: the values all compare equal. + expected = [values[0]] * len(values) assert_row_matches(got, expected, f"{agg} {values}") -def test_leading_infinity_also_deviates_from_native_cum_seed_leak(): - # The native seed leak is not NaN-specific: cum_min/cum_max fold - # from -+f64::MAX and never replace the seed on a tie, so a leading - # infinity of the same sign leaks it too. The prefix rule gives the - # true infinity instead. +def test_leading_infinity_follows_the_prefix_rule(): + # The prefix rule gives the true leading infinity. Native polars + # through 1.43 leaked its fold seed there (the seed leak was + # tie-blind, not NaN-specific); 1.44 fixed it and agrees. df = pl.DataFrame({"v": [[INF, 1.0]]}, schema={"v": pl.List(pl.Float64)}) got = df.select( polist.cum_agg_runs("v", pl.lit([True, True]), aggregation="min").alias("r") )["r"][0].to_list() + assert got == [INF, 1.0] native = pl.Series([INF, 1.0], dtype=pl.Float64).cum_min().to_list() - assert got[0] == INF - assert native[0] == pytest.approx(1.7976931348623157e308) + if PL_VERSION < (1, 44): + assert native[0] == pytest.approx(1.7976931348623157e308) + else: + assert native == got def test_leading_nan_follows_prefix_rule_not_polars_seed_leak(): @@ -490,11 +521,17 @@ def test_float32_accumulates_in_f64_and_rounds_once(): got = df.select( polist.cum_agg_runs("v", pl.lit([True] * 3), aggregation="sum").alias("r") )["r"][0].to_list() - native = pl.Series(values, dtype=pl.Float32).cum_sum().to_list() - # polars overflows at the second step and never recovers; the f64 - # accumulator stays in range and rounds back to a finite Float32. - assert math.isinf(native[1]) and math.isinf(native[2]) + # This library's contract, version-independent: the f64 running + # value exceeds Float32 range only at the middle position, so the + # once-rounded output recovers to a finite value. assert math.isinf(got[1]) and got[2] == pytest.approx(3e38, rel=1e-6) + native = pl.Series(values, dtype=pl.Float32).cum_sum().to_list() + if PL_VERSION < (1, 44): + # polars accumulated step-wise in Float32: overflow was sticky. + assert math.isinf(native[1]) and math.isinf(native[2]) + else: + # polars widened its accumulator and agrees. + assert math.isinf(native[1]) and native[2] == pytest.approx(3e38, rel=1e-6) @pytest.mark.parametrize("engine", ["in-memory", "streaming"]) diff --git a/tests/test_dtype_preservation.py b/tests/test_dtype_preservation.py index dff429f..166cb1c 100644 --- a/tests/test_dtype_preservation.py +++ b/tests/test_dtype_preservation.py @@ -67,12 +67,20 @@ def test_agg_slices_scalar_keeps_the_value_dtype(inner, agg): def test_agg_slices_float32_values_are_rounded_once(): - # Computed in f64, cast at the boundary -- the value is the f64 - # answer rounded, not an f32-accumulated one. - df = frame(pl.Float32) + # A fixture that discriminates: adding 1.0 to 2^24 in Float32 loses + # the 1.0 (the ulp at 2^24 is 2), so step-wise f32 accumulation and + # f64-then-round-once give different answers. + values = [16777216.0] + [1.0] * 8 # f64 sum 16777224 -> mean 1864136.0 + idx = [float(i) for i in range(len(values))] + df = pl.DataFrame( + {"s": [values], "i": [idx]}, + schema={"s": pl.List(pl.Float32), "i": pl.List(pl.Float32)}, + ) out = df.select(polist.agg_slices("s", "i", aggregation="mean").alias("r")) - expected = sum(SIGNAL32) / len(SIGNAL32) - assert out["r"][0] == pytest.approx(expected) + rounded_once = sum(values) / len(values) # exactly representable + stepwise_f32 = 16777216.0 / len(values) # what f32 accumulation gives + assert out["r"][0] == rounded_once + assert out["r"][0] != pytest.approx(stepwise_f32, rel=1e-9) # ------------------------------------------------------------- count dtype @@ -85,6 +93,25 @@ def test_agg_slices_count_emits_uint32(inner): assert out["r"][0] == 32 +@pytest.mark.parametrize("inner", [pl.Float32, pl.Float64]) +@pytest.mark.parametrize("empty", ["null", "zero"]) +def test_agg_lists_preserves_the_inner_dtype(inner, empty): + # agg_lists composes polars' own aggregations, which preserve the + # inner dtype natively -- but the declared output dtype flows from + # the pass-through plugin, so hardcoding it there must fail here. + df = pl.DataFrame( + {"g": [1, 1], "v": [[1.0, 2.0], [3.0, 4.0]]}, + schema={"g": pl.Int64, "v": pl.List(inner)}, + ) + for agg in ["sum", "mean", "median", "std", "min", "max", "delta"]: + out = df.group_by("g").agg( + polist.agg_lists( + "v", list_length=2, aggregation=agg, empty=empty # ty: ignore[invalid-argument-type] + ).alias("r") + ) + assert out["r"].dtype == pl.List(inner), (agg, empty) + + def test_agg_lists_count_emits_uint32(): df = pl.DataFrame( {"g": [1, 1], "v": [[1.0, None], [2.0, 3.0]]}, diff --git a/tests/test_kwargs_serialization.py b/tests/test_kwargs_serialization.py index 9cdab66..8497aff 100644 --- a/tests/test_kwargs_serialization.py +++ b/tests/test_kwargs_serialization.py @@ -70,6 +70,41 @@ def test_same_list_object_for_both_kwargs(df_ramp): assert out["r"][0] == 0.0 +def test_identical_bound_tuples_in_exclude_alone(df_ramp): + # The rebuild must run on slices_exclude independently: a shared + # object in BOTH kwargs stops aliasing once include is rebuilt, so + # only an exclude-only case can catch a missing exclude rebuild. + out = df_ramp.with_columns( + polist.agg_slices( + "v", "i", + aggregation="count", + slices_exclude=[((1.0, "closed"), (1.0, "closed"))], + ).alias("r") + ) + assert out["r"][0] == 10.0 + + +def test_numpy_bounds_in_exclude_alone(df_ramp): + out = df_ramp.with_columns( + polist.agg_slices( + "v", "i", + aggregation="count", + slices_exclude=[(np.float64(2.0), np.float64(4.0))], + ).alias("r") + ) + assert out["r"][0] == 8.0 + + +def test_bogus_boundary_mode_fails_with_the_parameter_name(df_ramp): + # A typo mode must not surface serde untagged-enum internals. + with pytest.raises(ValueError, match="boundary mode in slices_include"): + polist.agg_slices( + "v", "i", + aggregation="count", + slices_include=[((0.0, "clsoed"), (1.0, "open"))], # ty: ignore[invalid-argument-type] + ) + + def test_duplicate_simple_ranges(df_ramp): out = df_ramp.with_columns( polist.agg_slices( diff --git a/tests/test_plugins.py b/tests/test_plugins.py index b253301..c6f68da 100644 --- a/tests/test_plugins.py +++ b/tests/test_plugins.py @@ -125,6 +125,60 @@ def test_fft_nan_row_yields_all_nan_spectrum(): assert all(math.isnan(v) for v in spectrum) +def test_fft_inf_row_yields_nonfinite_spectrum(): + # An infinity is invalid data too, but unlike NaN it does not poison + # every bin: inf*0 window products and inf-inf bin sums are NaN + # while the rest stay infinite -- a visibly invalid mix, never null. + signal = sine(25.0) + signal[100] = float("inf") + df = pl.DataFrame({"s": [signal]}).with_columns( + polist.apply_fft("s", sample_rate=FS, scaling="amplitude").alias("a") + ) + spectrum = df["a"][0].to_list() + assert spectrum is not None + assert all(not math.isfinite(v) for v in spectrum) + assert any(math.isinf(v) for v in spectrum) + + +def test_butterworth_nan_contaminates_not_nulls(): + # Same rule for the filter: invalid propagates, it does not vanish. + signal = sine(5.0) + signal[50] = float("nan") + df = pl.DataFrame({"s": [signal]}).with_columns( + polist.apply_butterworth("s", sample_rate=FS, max_freq=20.0).alias("f") + ) + out = df["f"][0].to_list() + assert out is not None + assert any(math.isnan(v) for v in out) + + +def test_agg_slices_even_count_median_matches_polars(): + # The interpolated even-count median is a headline 2.0 change; pin + # its VALUES against polars, including the extremes where the old + # (lo + hi) / 2 midpoint overflowed or mis-signed. + for values in ( + [1.0, 2.0, 3.0, 4.0], + [0.1, 0.2, 0.3, 0.4], + [1e308, 1e308], + [-1e308, 1e308], + [1.0, float("-inf")], + [float("inf"), float("inf")], + ): + idx = [float(i) for i in range(len(values))] + df = pl.DataFrame( + {"v": [values], "i": [idx]}, + schema={"v": pl.List(pl.Float64), "i": pl.List(pl.Float64)}, + ) + got = df.select( + polist.agg_slices("v", "i", aggregation="median").alias("r") + )["r"][0] + expected = pl.Series(values, dtype=pl.Float64).median() + if isinstance(expected, float) and math.isnan(expected): + assert got is not None and math.isnan(got), values + else: + assert got == expected, values + + def test_transform_empty_list_raises(): # The limiting case of too-short: a raise policy that softened as # the input got worse would be indefensible. @@ -377,6 +431,7 @@ def test_agg_slices_all_null_selection_follows_the_empty_flag(): assert out["zero"][0] == 0.0 +@pytest.mark.filterwarnings("ignore:In Polars 2.0") # explode(empty_as_null) @pytest.mark.parametrize( ("lo", "hi", "selects_anything"), [(0.0, 1.0, True), (100.0, 200.0, False)] ) @@ -712,6 +767,41 @@ def test_agg_lists_and_agg_slices_share_missing_data_semantics(): # ----------------------------------------------------------- null semantics +NULL_ELEMENT_TRANSFORMS = { + # Every null-element position across the transform family: a single + # entry-path change (dropping nulls, unwrap_or) must fail loudly. + "fft": lambda good, bad: pl.DataFrame({"s": [good, bad]}).select( + polist.apply_fft("s", sample_rate=FS).alias("r") + ), + "butterworth": lambda good, bad: pl.DataFrame({"s": [good, bad]}).select( + polist.apply_butterworth("s", sample_rate=FS, max_freq=20.0).alias("r") + ), + # The no-cutoff pass-through is not a no-op: the family entry + # contract still applies (spec 9.20). + "butterworth_passthrough": lambda good, bad: pl.DataFrame( + {"s": [good, bad]} + ).select(polist.apply_butterworth("s", sample_rate=FS).alias("r")), + "interp_x": lambda good, bad: pl.DataFrame( + {"x": [good, bad], "y": [good, good]} + ).select(polist.apply_interp("x", "y", pl.lit([0.5]), strict=False).alias("r")), + "interp_y": lambda good, bad: pl.DataFrame( + {"x": [good, good], "y": [good, bad]} + ).select(polist.apply_interp("x", "y", pl.lit([0.5])).alias("r")), + "interp_xp": lambda good, bad: pl.DataFrame( + {"x": [good, good], "y": [good, good], "p": [good, bad]} + ).select(polist.apply_interp("x", "y", "p").alias("r")), +} + + +@pytest.mark.parametrize("position", list(NULL_ELEMENT_TRANSFORMS)) +def test_null_elements_null_the_row_across_the_family(position): + good = [float(v) for v in range(16)] + bad = [0.0, None] + [float(v) for v in range(14)] + out = NULL_ELEMENT_TRANSFORMS[position](good, bad)["r"] + assert out[0] is not None, position + assert out[1] is None, position # inner nulls must not silently misalign + + def test_null_and_inner_null_rows_yield_null(): df = pl.DataFrame( {"s": [sine(25.0), None, [1.0, None, 3.0] + [0.0] * 509]}, diff --git a/uv.lock b/uv.lock index b9c9f4e..7fe8a12 100644 --- a/uv.lock +++ b/uv.lock @@ -482,6 +482,7 @@ dev = [ { name = "matplotlib" }, { name = "maturin" }, { name = "numpy" }, + { name = "polars" }, { name = "pytest" }, ] @@ -491,6 +492,7 @@ requires-dist = [ { name = "maturin", marker = "extra == 'dev'", specifier = ">=1.0,<2.0" }, { name = "numpy", marker = "extra == 'dev'", specifier = ">=2.2.3" }, { name = "polars", specifier = ">=1.23.0" }, + { name = "polars", marker = "extra == 'dev'", specifier = ">=1.25" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, ] provides-extras = ["dev"] From c0f0cee1bbff9140bd4e0b75d2401c97518cb786 Mon Sep 17 00:00:00 2001 From: Travis Hammond Date: Wed, 2 Sep 2026 14:32:08 +0200 Subject: [PATCH 15/17] chore: drop the local working documents from .gitignore They stay untracked and uncommitted, and git status simply shows them -- preferred over a permanent ignore entry for files that will move elsewhere. Co-Authored-By: Claude Fable 5 --- .gitignore | 4 ---- 1 file changed, 4 deletions(-) diff --git a/.gitignore b/.gitignore index 7af0ed3..36104b3 100644 --- a/.gitignore +++ b/.gitignore @@ -9,10 +9,6 @@ target # Claude .claude/ -# Design/requirements working documents, kept local -polars_list_utils_spec.md -prefade_migration_todos.md - # These are backup files generated by rustfmt **/*.rs.bk From f3c82cb709b7b3f56ccd15bb699b79d4dfbc299c Mon Sep 17 00:00:00 2001 From: Travis Hammond Date: Wed, 2 Sep 2026 15:05:29 +0200 Subject: [PATCH 16/17] chore: declare the MIT license A LICENSE file plus the pyproject and Cargo metadata, so PyPI stops showing the package as all-rights-reserved and cargo stops warning about a missing license during the sdist build. The wheel metadata carries License-File plus the MIT classifier; the license file ships in every wheel and the sdist. Co-Authored-By: Claude Fable 5 --- Cargo.toml | 1 + LICENSE | 21 +++++++++++++++++++++ pyproject.toml | 2 ++ 3 files changed, 24 insertions(+) create mode 100644 LICENSE diff --git a/Cargo.toml b/Cargo.toml index 688c8f5..fdd9b44 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,6 +2,7 @@ name = "polars-list-utils" version = "2.0.0" edition = "2024" +license = "MIT" [lib] name = "polars_list_utils" diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..5186eed --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025-2026 Travis Hammond + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/pyproject.toml b/pyproject.toml index 0826399..6a2e4b8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,8 +3,10 @@ name = "polars-list-utils" authors = [{ name="Travis Hammond", email="dashdeckers@gmail.com" }] description = "Signal-processing utilities for List columns in Polars" readme = "README.md" +license = { file = "LICENSE" } dynamic = ["version"] classifiers = [ + "License :: OSI Approved :: MIT License", "Programming Language :: Rust", "Programming Language :: Python :: 3", "Programming Language :: Python :: Implementation :: CPython", From 05818d83860e338e0101919e2ea0b3a707ba4812 Mon Sep 17 00:00:00 2001 From: Travis Hammond Date: Wed, 2 Sep 2026 15:40:34 +0200 Subject: [PATCH 17/17] fix: run the plan-purity subprocess outside the repo root python -c prepends its working directory to sys.path, and at the repo root the source directory shadows the installed package. Locally the shadow is harmless -- maturin develop drops the compiled module into the source tree -- but on CI, where the package is pip-installed into site-packages, the shadow has no _internal module and the import dies. The subprocess now runs from the tests directory, which cannot shadow anything. Only this test was affected: it is the suite's one subprocess, and pytest itself puts tests/ on the path, not the root. Co-Authored-By: Claude Fable 5 --- tests/test_plan_purity.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/test_plan_purity.py b/tests/test_plan_purity.py index 35de9e5..fb9d836 100644 --- a/tests/test_plan_purity.py +++ b/tests/test_plan_purity.py @@ -70,6 +70,12 @@ def test_pipeline_streams_without_in_memory_fallback(): text=True, env=env, timeout=120, + # Run from the tests directory, not the repo root: `python -c` + # prepends its CWD to sys.path, and at the repo root the source + # directory shadows the installed package -- which locally holds + # the maturin-develop .pyd but on CI has no compiled module at + # all, failing the import. + cwd=os.path.dirname(os.path.abspath(__file__)), ) log = proc.stderr + proc.stdout assert proc.returncode == 0, log