diff --git a/.github/workflows/ci_image_build.yml b/.github/workflows/ci_image_build.yml index 27139f26..07b42580 100644 --- a/.github/workflows/ci_image_build.yml +++ b/.github/workflows/ci_image_build.yml @@ -5,6 +5,7 @@ on: paths: - 'docker/ci/**' - '.github/workflows/ci_image_build.yml' + - 'tools/dependencies.sh' workflow_dispatch: inputs: tag-base: diff --git a/Cargo.lock b/Cargo.lock index 1bea6cd3..46c3c655 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -20,6 +20,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "anyhow" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224afbd727c3d6e4b90103ece64b8d1b67fbb1973b1046c2281eed3f3803f800" + [[package]] name = "approx" version = "0.4.0" @@ -250,7 +256,7 @@ checksum = "29e724a68d9319343bb3328c9cc2dfde263f4b3142ee1059a9980580171c954b" dependencies = [ "atty", "bitflags", - "clap_lex", + "clap_lex 0.2.4", "indexmap", "strsim", "termcolor", @@ -258,6 +264,41 @@ dependencies = [ "textwrap", ] +[[package]] +name = "clap" +version = "4.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f13b9c79b5d1dd500d20ef541215a6423c75829ef43117e1b4d17fd8af0b5d76" +dependencies = [ + "bitflags", + "clap_derive", + "clap_lex 0.3.1", + "once_cell", +] + +[[package]] +name = "clap-cargo" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eca953650a7350560b61db95a0ab1d9c6f7b74d146a9e08fb258b834f3cf7e2c" +dependencies = [ + "clap 4.1.4", + "doc-comment", +] + +[[package]] +name = "clap_derive" +version = "4.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "684a277d672e91966334af371f1a7b5833f9aa00b07c84e92fbce95e00208ce8" +dependencies = [ + "heck", + "proc-macro-error", + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "clap_lex" version = "0.2.4" @@ -267,6 +308,15 @@ dependencies = [ "os_str_bytes", ] +[[package]] +name = "clap_lex" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "783fe232adfca04f90f56201b26d79682d4cd2625e0bc7290b95123afe558ade" +dependencies = [ + "os_str_bytes", +] + [[package]] name = "colored" version = "2.0.0" @@ -298,6 +348,12 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fb4a24b1aaf0fd0ce8b45161144d6f42cd91677fd5940fd431183eb023b3a2b8" +[[package]] +name = "core-foundation-sys" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5827cebf4670468b8772dd191856768aedcb1b0278a04f989f7766351917b9dc" + [[package]] name = "cortex-m" version = "0.7.6" @@ -406,22 +462,6 @@ dependencies = [ "typenum", ] -[[package]] -name = "cstr_core" -version = "0.2.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dd98742e4fdca832d40cab219dc2e3048de17d873248f83f17df47c1bea70956" -dependencies = [ - "cty", - "memchr", -] - -[[package]] -name = "cty" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b365fabc795046672053e29c954733ec3b05e4be654ab130fe8f1f94d7051f35" - [[package]] name = "digest" version = "0.10.3" @@ -453,6 +493,12 @@ dependencies = [ "winapi", ] +[[package]] +name = "doc-comment" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fea41bba32d969b513997752735605054bc0dfa92b4c56bf1189f2e174be7a10" + [[package]] name = "either" version = "1.7.0" @@ -696,6 +742,12 @@ dependencies = [ "stable_deref_trait", ] +[[package]] +name = "heck" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8" + [[package]] name = "hermit-abi" version = "0.1.19" @@ -814,9 +866,9 @@ checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" [[package]] name = "libc" -version = "0.2.138" +version = "0.2.139" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "db6d7e329c562c5dfab7a46a2afabc8b987ab9a4834c9d1ca04dc54c1546cef8" +checksum = "201de327520df007757c1f0adce6e827fe8562fbc28bfd9c15571c66ca1f5f79" [[package]] name = "libloading" @@ -967,6 +1019,15 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "ntapi" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc51db7b362b205941f71232e56c625156eb9a929f8cf74a428fd5bc094a4afc" +dependencies = [ + "winapi", +] + [[package]] name = "nu-ansi-term" version = "0.46.0" @@ -1031,9 +1092,9 @@ dependencies = [ [[package]] name = "once_cell" -version = "1.16.0" +version = "1.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "86f0b0d4bf799edbc74508c1e8bf170ff5f41238e5f8225603ca7caaae2b7860" +checksum = "b7e5500299e16ebb147ae15a00a942af264cf3688f47923b8fc2cd5858f23ad3" [[package]] name = "ordered-float" @@ -1092,6 +1153,16 @@ version = "1.0.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9423e2b32f7a043629287a536f21951e8c6a82482d0acb1eeebfc90bc2225b22" +[[package]] +name = "pathsearch" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da983bc5e582ab17179c190b4b66c7d76c5943a69c6d34df2a2b6bf8a2977b05" +dependencies = [ + "anyhow", + "libc", +] + [[package]] name = "peeking_take_while" version = "0.1.2" @@ -1160,22 +1231,19 @@ dependencies = [ [[package]] name = "pgx" -version = "0.6.1" +version = "0.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "708693428ee16491645c1c2795f6777a78b9b1d08ab2e481a82b7977c9814b56" +checksum = "fc91f19f84e7c1ba7b25953b042bd487b6e1bbec4c3af09f61a6ac31207ff776" dependencies = [ "atomic-traits", "bitflags", "bitvec", - "cstr_core", - "eyre", "heapless", "libc", "once_cell", "pgx-macros", "pgx-pg-sys", - "pgx-utils", - "quote", + "pgx-sql-entity-graph", "seahash", "seq-macro", "serde", @@ -1185,31 +1253,31 @@ dependencies = [ "time", "tracing", "tracing-error", - "uuid 1.2.2", + "uuid 1.3.0", ] [[package]] name = "pgx-macros" -version = "0.6.1" +version = "0.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0050ca15df7dbfe718f9006d2b9a38d2a00a40934d8154450cdde6323b58b690" +checksum = "1ebfde3c33353d42c2fbcc76bea758b37018b33b1391c93d6402546569914e94" dependencies = [ - "pgx-utils", + "pgx-sql-entity-graph", "proc-macro2", "quote", "syn", - "unescape", ] [[package]] name = "pgx-pg-config" -version = "0.6.1" +version = "0.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "000bba0f67f2aa20e971a6127a8971bd85a6b724d3e95d9066f64816de3b4e67" +checksum = "e97c27bab88fdb7b94e549b02267ab9595bd9d1043718d6d72bc2d34cf1e3952" dependencies = [ "dirs", "eyre", "owo-colors", + "pathsearch", "serde", "serde_derive", "serde_json", @@ -1219,9 +1287,9 @@ dependencies = [ [[package]] name = "pgx-pg-sys" -version = "0.6.1" +version = "0.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6bd3fd6e1bcbfed67da7ac2dc71986a9f1396aeaec3d449d0697eb54909b34a" +checksum = "6b79c48c564bed305d202b852321603107e5f3ac31f25ea2cc4031475f38d0b3" dependencies = [ "bindgen", "eyre", @@ -1230,21 +1298,42 @@ dependencies = [ "once_cell", "pgx-macros", "pgx-pg-config", - "pgx-utils", + "pgx-sql-entity-graph", "proc-macro2", "quote", - "rayon", + "serde", "shlex", "sptr", "syn", ] +[[package]] +name = "pgx-sql-entity-graph" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "573a8d8c23be24c39f7b7fbbc7e15d95aa0327acd61ba95c9c9f237fec51f205" +dependencies = [ + "convert_case", + "eyre", + "petgraph", + "proc-macro2", + "quote", + "regex", + "seq-macro", + "syn", + "tracing", + "tracing-error", + "tracing-subscriber", + "unescape", +] + [[package]] name = "pgx-tests" -version = "0.6.1" +version = "0.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44db59e5a5473f9192cff05a284677b8b491bffcf85dcc7b892339f084ad2f2b" +checksum = "fc09f25ae560bc4e3308022999416966beda5b60d2957b9ab92bffaf2d6a86c3" dependencies = [ + "clap-cargo", "eyre", "libc", "once_cell", @@ -1252,40 +1341,15 @@ dependencies = [ "pgx", "pgx-macros", "pgx-pg-config", - "pgx-utils", "postgres", "regex", "serde", "serde_json", + "sysinfo", "thiserror", "time", ] -[[package]] -name = "pgx-utils" -version = "0.6.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "62cdc413efcd90a1e94c7f09dea24ec1ecfa1275350cbe378a8776eb7b5448f9" -dependencies = [ - "atty", - "convert_case", - "cstr_core", - "eyre", - "petgraph", - "proc-macro2", - "quote", - "regex", - "seq-macro", - "serde", - "serde_derive", - "serde_json", - "syn", - "tracing", - "tracing-error", - "tracing-subscriber", - "unescape", -] - [[package]] name = "phf" version = "0.11.1" @@ -1377,11 +1441,35 @@ version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eb9f9e6e233e5c4a35559a617bf40a4ec447db2e84c20b55a6f83167b7e57872" +[[package]] +name = "proc-macro-error" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da25490ff9892aab3fcf7c36f08cfb902dd3e71ca0f9f9517bea02a73a5ce38c" +dependencies = [ + "proc-macro-error-attr", + "proc-macro2", + "quote", + "syn", + "version_check", +] + +[[package]] +name = "proc-macro-error-attr" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1be40180e52ecc98ad80b184934baf3d0d29f979574e439af5a55274b35f869" +dependencies = [ + "proc-macro2", + "quote", + "version_check", +] + [[package]] name = "proc-macro2" -version = "1.0.47" +version = "1.0.51" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ea3d908b0e36316caf9e9e2c4625cdde190a7e6f440d794667ed17a1855e725" +checksum = "5d727cae5b39d21da60fa540906919ad737832fe0b1c165da3a34d6548c849d6" dependencies = [ "unicode-ident", ] @@ -1422,9 +1510,9 @@ dependencies = [ [[package]] name = "quote" -version = "1.0.21" +version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bbe448f377a7d6961e30f5955f9b8d106c3f5e449d493ee1b125c1d43c2b5179" +checksum = "8856d8364d252a14d474036ea1358d63c9e6965c8e5c1885c18f73d70bff9c7b" dependencies = [ "proc-macro2", ] @@ -1526,9 +1614,9 @@ dependencies = [ [[package]] name = "regex" -version = "1.7.0" +version = "1.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e076559ef8e241f2ae3479e36f97bd5741c0330689e217ad51ce2c76808b868a" +checksum = "48aaa5748ba571fb95cd2c85c09f629215d3a6ece942baa100950af03a34f733" dependencies = [ "aho-corasick", "memchr", @@ -1689,9 +1777,9 @@ checksum = "0772c5c30e1a0d91f6834f8e545c69281c099dfa9a3ac58d96a9fd629c8d4898" [[package]] name = "serde" -version = "1.0.149" +version = "1.0.152" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "256b9932320c590e707b94576e3cc1f7c9024d0ee6612dfbcf1cb106cbe8e055" +checksum = "bb7d1f0d3021d347a83e556fc4683dea2ea09d87bccdf88ff5c12545d89d5efb" dependencies = [ "serde_derive", ] @@ -1708,9 +1796,9 @@ dependencies = [ [[package]] name = "serde_derive" -version = "1.0.149" +version = "1.0.152" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b4eae9b04cbffdfd550eb462ed33bc6a1b68c935127d008b27444d08380f94e4" +checksum = "af487d118eecd09402d70a5d72551860e788df87b464af30e5ea6a38c75c541e" dependencies = [ "proc-macro2", "quote", @@ -1719,9 +1807,9 @@ dependencies = [ [[package]] name = "serde_json" -version = "1.0.89" +version = "1.0.93" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "020ff22c755c2ed3f8cf162dbb41a7268d934702f3ed3631656ea597e08fc3db" +checksum = "cad406b69c91885b5107daf2c29572f6c8cdb3c66826821e286c533490c0bc76" dependencies = [ "itoa", "ryu", @@ -1839,7 +1927,7 @@ name = "sql-doctester" version = "0.1.0" dependencies = [ "bytecount", - "clap", + "clap 3.2.17", "colored", "postgres", "pulldown-cmark", @@ -1903,15 +1991,30 @@ checksum = "6bdef32e8150c2a081110b42772ffe7d7c9032b606bc226c8260fd97e0976601" [[package]] name = "syn" -version = "1.0.105" +version = "1.0.107" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "60b9b43d45702de4c839cb9b51d9f529c5dd26a4aff255b42b1ebc03e88ee908" +checksum = "1f4064b5b16e03ae50984a5a8ed5d4f8803e6bc1fd170a3cda91a1be4b18e3f5" dependencies = [ "proc-macro2", "quote", "unicode-ident", ] +[[package]] +name = "sysinfo" +version = "0.27.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "975fe381e0ecba475d4acff52466906d95b153a40324956552e027b2a9eaa89e" +dependencies = [ + "cfg-if", + "core-foundation-sys", + "libc", + "ntapi", + "once_cell", + "rayon", + "winapi", +] + [[package]] name = "tap" version = "1.0.1" @@ -2051,7 +2154,7 @@ dependencies = [ [[package]] name = "timescaledb_toolkit" -version = "1.14.0-dev" +version = "1.15.0" dependencies = [ "aggregate_builder", "approx 0.4.0", @@ -2071,6 +2174,7 @@ dependencies = [ "pest_derive", "pgx", "pgx-macros", + "pgx-sql-entity-graph", "pgx-tests", "rand", "rand_chacha", @@ -2160,9 +2264,9 @@ dependencies = [ [[package]] name = "toml" -version = "0.5.9" +version = "0.5.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8d82e1a7758622a465f8cee077614c73484dac5b836c02ff6a40d5d1010324d7" +checksum = "f4f7f0dd8d50a853a531c426359045b1998f04219d88799810762cd4ad314234" dependencies = [ "serde", ] @@ -2390,7 +2494,7 @@ name = "update-tester" version = "0.3.0" dependencies = [ "bytecount", - "clap", + "clap 3.2.17", "colored", "control_file_reader", "postgres", @@ -2424,9 +2528,9 @@ dependencies = [ [[package]] name = "uuid" -version = "1.2.2" +version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "422ee0de9031b5b948b97a8fc04e3aa35230001a722ddd27943e0be31564ce4c" +checksum = "1674845326ee10d37ca60470760d4288a6f80f304007d92e5c53bab78c9cfd79" dependencies = [ "getrandom", ] diff --git a/Changelog.md b/Changelog.md index cb2b8714..4daafe23 100644 --- a/Changelog.md +++ b/Changelog.md @@ -4,7 +4,27 @@ This changelog should be updated as part of a PR if the work is worth noting (most of them should be). If unsure, always add an entry here for any PR targeted for the next release. It's easier to remove than add an entry at final review time for the next release. -## Next Release (Date TBD) +## [1.15.0](https://github.com/timescale/timescaledb-toolkit/releases/tag/1.15.0) (2023-03-08) + +#### New experimental features + +#### Bug fixes +- [#715](https://github.com/timescale/timescaledb-toolkit/pull/715): Fix out-of-bounds indexing error in `state_agg` rollup + +#### Stabilized features +- [#722](https://github.com/timescale/timescaledb-toolkit/pull/722): Stabilize heartbeat aggregate. +- [#724](https://github.com/timescale/timescaledb-toolkit/pull/724): Stabilize integral and interpolated_integral for time-weighted-average. +- [#723](https://github.com/timescale/timescaledb-toolkit/pull/723): Stabilized `state_agg` + +#### Other notable changes +- [#716](https://github.com/timescale/timescaledb-toolkit/issues/716): Add arrow operator support for counter aggregate and time-weighted aggregate interpolated accessors. +- [#716](https://github.com/timescale/timescaledb-toolkit/issues/716): Remove experimental versions of interpolated accessors for counter aggregate and time-weighted aggregates. The stable versions introduced in 1.14.0 should be used instead. +- [#723](https://github.com/timescale/timescaledb-toolkit/pull/723): Added `state_at` function for `state_agg` +- [#709](https://github.com/timescale/timescaledb-toolkit/pull/709): Updated pgx version to 0.7.1 + +**Full Changelog**: https://github.com/timescale/timescaledb-toolkit/compare/1.14.0...1.15.0 + +## [1.14.0](https://github.com/timescale/timescaledb-toolkit/releases/tag/1.14.0) (2023-02-09) #### New experimental features @@ -24,13 +44,11 @@ This changelog should be updated as part of a PR if the work is worth noting (mo - [#699](https://github.com/timescale/timescaledb-toolkit/pull/699): `interpolated_duration_in`/`duration_in`/`interpolated_state_periods`/`state_periods` have the first two arguments swapped: now the aggregate is first and the state is second - [#699](https://github.com/timescale/timescaledb-toolkit/pull/699): `into_values`/`into_int_values` now returns a table with intervals instead of microseconds -#### Shout-outs - -**Full Changelog**: [TODO] +**Full Changelog**: https://github.com/timescale/timescaledb-toolkit/compare/1.13.1...1.14.0 ## [1.13.1](https://github.com/timescale/timescaledb-toolkit/releases/tag/1.13.1) (2023-01-03) -**Full Changelog**: +**Full Changelog**: https://github.com/timescale/timescaledb-toolkit/compare/1.13.0...1.13.1 - [#664](https://github.com/timescale/timescaledb-toolkit/pull/664) Support PostgreSQL 15. @@ -77,6 +95,8 @@ Users can use the new experimental function `toolkit_experimental.to_text(timeve - [#646](https://github.com/timescale/timescaledb-toolkit/pull/646): Added experimental support for PostgreSQL 15. - [#621](https://github.com/timescale/timescaledb-toolkit/pull/621): Rocky Linux 9 support +**Full Changelog**: https://github.com/timescale/timescaledb-toolkit/compare/1.12.1...1.13.0 + ## [1.12.1](https://github.com/timescale/timescaledb-toolkit/releases/tag/1.12.1) (2022-11-17) #### Bug fixes diff --git a/Readme.md b/Readme.md index a28fddc9..4ddfed27 100644 --- a/Readme.md +++ b/Readme.md @@ -38,7 +38,7 @@ As for other platforms: patches welcome! ### 🔧 Tools Setup ### -Building the extension requires valid [rust](https://www.rust-lang.org/) (we build and test on 1.64), [rustfmt](https://github.com/rust-lang/rustfmt), and clang installs, along with the postgres headers for whichever version of postgres you are running, and pgx. +Building the extension requires valid [rust](https://www.rust-lang.org/) (we build and test on 1.65), [rustfmt](https://github.com/rust-lang/rustfmt), and clang installs, along with the postgres headers for whichever version of postgres you are running, and pgx. We recommend installing rust using the [official instructions](https://www.rust-lang.org/tools/install): ```bash curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh @@ -51,7 +51,7 @@ sudo apt-get install make gcc pkg-config clang postgresql-server-dev-14 libssl-d Next you need [cargo-pgx](https://github.com/tcdi/pgx), which can be installed with ```bash -cargo install --version '=0.6.1' --force cargo-pgx +cargo install --version '=0.7.1' --force cargo-pgx ``` You must reinstall cargo-pgx whenever you update your Rust compiler, since cargo-pgx needs to be built with the same compiler as Toolkit. diff --git a/bors.toml b/bors.toml index 78f331cc..0a1d3761 100644 --- a/bors.toml +++ b/bors.toml @@ -3,6 +3,7 @@ status = [ "Test Postgres (12)", "Test Postgres (13)", "Test Postgres (14)", + "Test Postgres (15)", ] delete-merged-branches = true timeout_sec = 3600 # 60 min diff --git a/docker/ci/setup.sh b/docker/ci/setup.sh index fabdb4a7..15d6545f 100755 --- a/docker/ci/setup.sh +++ b/docker/ci/setup.sh @@ -78,6 +78,8 @@ if $privileged; then ;; 8) + yum -qy install dnf-plugins-core + dnf config-manager --enable powertools dnf -qy module disable postgresql # fpm suddenly requires newer public_suffix that requires newer ruby # https://github.com/jordansissel/fpm/issues/1923 ¯\_(ツ)_/¯ @@ -86,6 +88,8 @@ if $privileged; then ;; 9) + yum -qy install dnf-plugins-core + dnf config-manager --enable crb dnf -qy install ruby-devel rubygems ;; diff --git a/docs/release.md b/docs/release.md new file mode 100644 index 00000000..77339d10 --- /dev/null +++ b/docs/release.md @@ -0,0 +1,287 @@ +# Release and build procedures + +We build the timescaledb_toolkit extension using Cargo, but we have many +higher-level tasks in need of automation: + +- Build, lint, and test with particular flags in multiple environments +- Extract SQL examples from documentation and test +- Test upgrades +- Installation +- Publish a release +- Make a container image to run all the above on + +The rest of this document elaborates on each of those. But first.. + +## Dependency management + +Ideally, all dependencies would be specified in just one place. But that's +not quite feasible. Cargo.toml files capture the crate dependencies. + +The rest are needed by the six shell scripts used to solve the above list +of problems. We configure those in `tools/dependencies.sh`. + +## Build, lint, and test + +`tools/build` is the relatively simple shell script that owns the cargo flags +for running clippy, running tests, installing, and testing upgrades. +The latter two are arguably out of place here. + +Testing upgrades is now handled by `testbin` (below), but the version here was +useful for Mac. That has now degraded as it would need to support a third +pgx... + +Installing is only relevant for local development. + +## Extract SQL examples from documentation and test + +`tools/sql-doctester` is a Rust program which extracts example SQL programs +and expected output from documentation, runs the programs, and asserts their +output matches what was expected. + +The intent here is merely to prevent sample code from bitrotting, but some +functionality is currently only tested here. + +## Test upgrades + +We include in each release a set of scripts to upgrade an installation from a +set of previous versions. We test these upgrades by installing a supported +old version, materializing some data, running the upgrade script, and +asserting the extension can still load the old data. + +`tools/update-tester` is a Rust program which loads tests from `tests/update` +to implement the materialize and verify steps. It needs to know which version +each function was stabilized in, and we store that information in +`extension/src/stabilization_info.rs` (also used by post-install, see below). + +`tools/testbin` is a shell script that uses `update-tester` to test upgrades +between released binaries (deb and rpm). + +## Installation + +Installation is a two-step process currently duplicated in three places. +The two steps are: + +1. `cargo pgx install --release` OR `cargo pgx package` +2. `tools/post-install` + +These steps are repeated in: + +1. `Readme.md` +2. `tools/build` +3. `toolkit/package-deb.sh` and `toolkit/package-rpm.sh` + +`Readme.md` could simply recommend running `tools/build install`. + +`package-deb.sh` and `package-rpm.sh` could run `tools/build package` (which +doesn't yet exist). + +`cargo pgx install` installs the extension into the directory specified by +`pg_config`. `cargo pgx package` installs into a directory under +`$CARGO_TARGET_DIR` where we pick it up and pack it into deb and rpm packages. + +`tools/post-install` performs miscellaneous install-time procedures: +- finalize control file +- rename `timescaledb_toolkit.so` to include the version number +- generate update scripts + +`tools/post-install` needs to know which version each function was stabilized +in, and we store that information in `extension/src/stabilization_info.rs`. + +## Publish a release + +`tools/release` automates all the steps of our release process. We run it via +github action (`.github/workflows/release.yml`). + +`tools/release` creates and pushes a release branch and tag, runs tests, +starts a package build, prepares the `main` branch for the next release, and +creates an issue so we don't forget some tasks not yet automated. + +The package build happens in a different repository for reasons described in +comments at the top of `tools/release`. + +Over in that repository, we have `.github/workflows/toolkit-package.yml` which +runs `toolkit/docker-run-package.sh` to build packages and +`toolkit/upload-packages.sh` to upload them to PackageCloud. + +`toolkit/docker-run-package.sh` runs `toolkit/package-deb.sh` and +`toolkit/package-rpm.sh` in various container images to build packags for +those platforms. Which platforms we build for is controlled in the `yml` +action file. + +### Usage: + +1. https://github.com/timescale/timescaledb-toolkit/actions/workflows/release.yml +2. Click "Run workflow" +3. Fill out the form +4. Be sure to replace `-n` with `-push`! + +We can replace this last one with a checkbox: +- unchecked: run with neither `-n` nor `-push` +- checked: run with `-push` + +The script has three modes: + +- `-n`: print what would be done without doing anything +- `-push`: do everything including pushing to Github and PackageCloud +- neither: do all the work (branch, edit, test, package, upgrade-test), but don't push anywhere + +The third mode is the most useful but it is not available from the Github +action. Very sad. We need to fix that. + +### Resume after failure + +Up until the packaging step, just rerun the release action after the problem +is resolved. + +If packages have been published, the choices are: +- do the rest of what the script does manually +- increment the patch revision (1.3.X) and start another release + +An obvious improvement would be to teach `tools/release` to resume at a +specific step, something like `tools/release --start-at-step 7`. It would +need to verify that the previous steps were actually done and bail out if not. + +Once packaging is no longer asynchronous in the other repository, +`tools/release` can simply be taught to figure out which steps are done all on +its own, without an operator having to tell it where to resume. + +### Debugging + +We run `tools/release` with the shell's `-x` option so it prints each command +it runs. We redirect the standard error stream to the standard output because +Docker will otherwise separate them such that error messages may appear far +from related output. + +So, when something goes wrong, it is easy to pinpoint exactly which part of +the script failed and how. + +Things that can go wrong: + +#### Transient network hiccough + +This can happen at almost any stage. A simple retry might be the easiest way +to see if the issue is transient. If it's not, options are limited: +- wait +- complain, then wait + +#### cargo install cargo-edit + +- Is crates.io down? +- Has cargo-edit vanished from crates.io? + +#### Install gh + +- The version we use is gone. Find the latest and figure out whether all our + usage has been invalidated by incompatible changes. Be careful! +- Or, just squirrel away a copy of the old binary and keep rolling, until the + underlying APIs it uses break. +- The checksum doesn't match. Did they break it? Why would they do such a + thing? Were they hacked? Probably should go ahead and update at this point. + +#### `extension/timescaledb_toolkit.control` problems + +`tools/release` edits this file, so it is very careful that the file looks the +way it expects. It is and should remain very picky. + +If we've made some unexpected edits, it will complain. If the edits were +erroneous, fix them; else, you have to teach `tools/release` what you've done. + +One of the things it checks is the `upgradeable_from` line. Most importantly, +it expects that patch releases are upgradeable from the previous version in +the same minor version (e.g. 1.3.1 is upgradeable from 1.3.0). + +#### `Changelog.md` problems + +`tools/release` ensures the version being released has Changelog.md entries. + +It also requires some particular boiler-plate text at the top to know where to +make its edits. The boiler-plate is arbitrary text for intended for +consumption by the development team. If we change that text, `tools/release` +needs to know about it. + +#### Tests fail + +Oh boy! + +Test output is logged. + +`tools/build test-extension` shouldn't fail since it already passed when the +release commit was merged to master. + +You're not trying to release a commit that didn't pass CI, are you? + +But, the upgrade tests are being run for the first time! So those might +break. We should run `tools/release --no-push` nightly. In the mean time... +to the debugger! + +#### git push fails + +We've had branch permission problems before... + +Is the authentication token working? + +#### `gh` fails + +Is GitHub API struggling? + +Is the authentication token working? + +Has the packaging action in the `release-build-scripts` repository +gone missing? + +## Make a container image to run all the above on + +`.github/workflows/toolkit-image.yml` configures the GitHub action which +builds all our supported container images. + +One image is special: debian-11-amd64. This is the one we run all our GitHub +actions on. + +`docker/ci/Dockerfile` is the entry-point and it runs `docker/ci/setup.sh` to +do the work: + +- Create the build user +- Install necessary build tools and libraries +- Install postgresql and timescaledb +- Install `gh` github command-line tool used by `tools/release` +- Install Rust and PGX +- Pre-fetch toolkit's crate dependencies to minimize work done at CI time + +## Maintenance tasks + +So, we've automated build and release! ONCE AND FOR ALL. Right? + +As the great Balki Bartokomous often said: +of course not; don't be ridiculous. + +These are the sorts of things we have to do from time to time: + +- Update Rust. It moves pretty fast. +- Update PGX. It moves even faster. +- Update other crates. `cargo audit` and `cargo update` are our friends. +- Update OS versions. Labels such as `rockylinux:9` eventually point to + something different or disappear entirely. The former actually surprised us + once already. + +### Things we update blindly + +We install the latest version of these every time, so they may change in +surprising ways at inopportune times. + +- fpm: It's a Ruby script with lots of dependencies and we install the latest + version and it bit us on the ass once already. We use it because someone + set it up for us a long time ago and no one has had the chance to sit down + and figure out how to write an RPM spec file. Shouldn't take more than a + few hours, just haven't done it... + +- postgresql: We install the latest version of a fixed set of major versions, + so this should be very unlikely to break on us. Listed for completeness. + +- timescaledb: We test with their master branch nightly, so we should be + ahead of this one. + +### Unknown Unknowns + +lol + +They're inevitable. You just need a good nose for debugging. diff --git a/docs/state_agg.md b/docs/state_agg.md index 1126fbf4..1aa9ee8c 100644 --- a/docs/state_agg.md +++ b/docs/state_agg.md @@ -30,13 +30,17 @@ INSERT INTO states_test_4 VALUES ('2020-01-01 00:01:00+00', 2), ('2020-01-01 00:01:03+00', 51351), ('2020-01-01 00:02:00+00', -9); -CREATE TABLE states_test_5(ts TIMESTAMPTZ, state BIGINT); +CREATE TABLE states_test_5(ts TIMESTAMPTZ, state BIGINT); -- states_test with integer states INSERT INTO states_test_5 VALUES ('2020-01-01 00:00:00+00', 4), ('2020-01-01 00:00:11+00', 51351), ('2020-01-01 00:01:00+00', 2), ('2020-01-01 00:02:03+00', 51351), ('2020-01-01 00:02:05+00', -9); +CREATE TABLE states_test_6(ts TIMESTAMPTZ, state BIGINT); -- states_test_3 with integer states +INSERT INTO states_test_6 VALUES + ('2019-12-31 00:00:11+00', 456789), + ('2019-12-31 00:01:00+00', 4); ``` ## Functions @@ -79,7 +83,7 @@ FROM states_test; #### duration_in for a range ```SQL -SELECT toolkit_experimental.duration_in(toolkit_experimental.state_agg(ts, state), 'OK', '2020-01-01 00:01:00+00', '2 days') FROM states_test; +SELECT duration_in(state_agg(ts, state), 'OK', '2020-01-01 00:01:00+00', '2 days') FROM states_test; ``` ```output duration_in @@ -87,7 +91,7 @@ SELECT toolkit_experimental.duration_in(toolkit_experimental.state_agg(ts, state 00:00:57 ``` ```SQL -SELECT toolkit_experimental.duration_in(toolkit_experimental.state_agg(ts, state), 'OK', '2020-01-01 00:01:00+00', NULL) FROM states_test; +SELECT duration_in(state_agg(ts, state), 'OK', '2020-01-01 00:01:00+00', NULL) FROM states_test; ``` ```output duration_in @@ -95,7 +99,7 @@ SELECT toolkit_experimental.duration_in(toolkit_experimental.state_agg(ts, state 00:00:57 ``` ```SQL -SELECT toolkit_experimental.duration_in(toolkit_experimental.state_agg(ts, state), 'OK', '2020-01-01 00:01:00+00') FROM states_test; +SELECT duration_in(state_agg(ts, state), 'OK', '2020-01-01 00:01:00+00') FROM states_test; ``` ```output duration_in @@ -103,7 +107,7 @@ SELECT toolkit_experimental.duration_in(toolkit_experimental.state_agg(ts, state 00:00:57 ``` ```SQL -SELECT toolkit_experimental.duration_in(toolkit_experimental.state_agg(ts, state), 51351, '2020-01-01 00:01:00+00', '2 days') FROM states_test_4; +SELECT duration_in(state_agg(ts, state), 51351, '2020-01-01 00:01:00+00', '2 days') FROM states_test_4; ``` ```output duration_in @@ -111,7 +115,7 @@ SELECT toolkit_experimental.duration_in(toolkit_experimental.state_agg(ts, state 00:00:57 ``` ```SQL -SELECT toolkit_experimental.duration_in(toolkit_experimental.state_agg(ts, state), 51351, '2020-01-01 00:01:00+00', NULL) FROM states_test_4; +SELECT duration_in(state_agg(ts, state), 51351, '2020-01-01 00:01:00+00', NULL) FROM states_test_4; ``` ```output duration_in @@ -120,7 +124,7 @@ SELECT toolkit_experimental.duration_in(toolkit_experimental.state_agg(ts, state ``` ```SQL -SELECT toolkit_experimental.duration_in(toolkit_experimental.state_agg(ts, state), 'OK', '2020-01-01 00:00:15+00', '30 seconds') FROM states_test; +SELECT duration_in(state_agg(ts, state), 'OK', '2020-01-01 00:00:15+00', '30 seconds') FROM states_test; ``` ```output duration_in @@ -129,7 +133,7 @@ SELECT toolkit_experimental.duration_in(toolkit_experimental.state_agg(ts, state ``` ```SQL -SELECT toolkit_experimental.duration_in(toolkit_experimental.state_agg(ts, state), 'OK', '2020-01-01 00:00:15+00', '1 minute 1 second') FROM states_test; +SELECT duration_in(state_agg(ts, state), 51351, '2020-01-01 00:00:15+00', '1 minute 1 second') FROM states_test_4; ``` ```output duration_in @@ -137,6 +141,33 @@ SELECT toolkit_experimental.duration_in(toolkit_experimental.state_agg(ts, state 00:00:58 ``` +```SQL +SELECT duration_in(state_agg(ts, state), 'OK', '2020-01-01 00:00:15+00', '1 minute 1 second') FROM states_test; +``` +```output + duration_in +------------- + 00:00:58 +``` + +```SQL +SELECT (SELECT state_agg(ts, state) FROM states_test) -> duration_in('OK'::text, '2020-01-01 00:00:15+00', '1 minute 1 second'); +``` +```output + ?column? +------------- + 00:00:58 +``` + +```SQL +SELECT (SELECT state_agg(ts, state) FROM states_test) -> duration_in('OK'); +``` +```output + ?column? +------------- + 00:01:46 +``` + ### into_values ```SQL @@ -153,8 +184,8 @@ SELECT state, duration FROM toolkit_experimental.into_values( STOP | 00:00:00 ``` ```SQL -SELECT state, duration FROM toolkit_experimental.into_int_values( - (SELECT toolkit_experimental.state_agg(ts, state) FROM states_test_4)) +SELECT state, duration FROM into_int_values( + (SELECT state_agg(ts, state) FROM states_test_4)) ORDER BY state, duration; ``` ```output @@ -165,12 +196,36 @@ SELECT state, duration FROM toolkit_experimental.into_int_values( 4 | 00:00:11 51351 | 00:01:46 ``` +```SQL +SELECT (state_agg(ts, state) -> into_values()).* FROM states_test ORDER BY state; +``` +```output + state | duration +-------+---------- + ERROR | 00:00:03 + OK | 00:01:46 + START | 00:00:11 + STOP | 00:00:00 +``` ### state_timeline ```SQL -SELECT state, start_time, end_time FROM toolkit_experimental.state_timeline( - (SELECT toolkit_experimental.state_agg(ts, state) FROM states_test)) +SELECT (state_agg(ts, state) -> state_timeline()).* FROM states_test; +``` +```output + state | start_time | end_time +-------+------------------------+------------------------ + START | 2020-01-01 00:00:00+00 | 2020-01-01 00:00:11+00 + OK | 2020-01-01 00:00:11+00 | 2020-01-01 00:01:00+00 + ERROR | 2020-01-01 00:01:00+00 | 2020-01-01 00:01:03+00 + OK | 2020-01-01 00:01:03+00 | 2020-01-01 00:02:00+00 + STOP | 2020-01-01 00:02:00+00 | 2020-01-01 00:02:00+00 +``` + +```SQL +SELECT state, start_time, end_time FROM state_timeline( + (SELECT state_agg(ts, state) FROM states_test)) ORDER BY start_time; ``` ```output @@ -184,8 +239,8 @@ ERROR | 2020-01-01 00:01:00+00 | 2020-01-01 00:01:03+00 ``` ```SQL -SELECT state, start_time, end_time FROM toolkit_experimental.state_int_timeline( - (SELECT toolkit_experimental.state_agg(ts, state) FROM states_test_4)) +SELECT state, start_time, end_time FROM state_int_timeline( + (SELECT state_agg(ts, state) FROM states_test_4)) ORDER BY start_time; ``` ```output @@ -200,8 +255,8 @@ state | start_time | end_time ```SQL -SELECT state, start_time, end_time FROM toolkit_experimental.state_timeline( - (SELECT toolkit_experimental.state_agg(ts, state) FROM states_test_2)) +SELECT state, start_time, end_time FROM state_timeline( + (SELECT state_agg(ts, state) FROM states_test_2)) ORDER BY start_time; ``` ```output @@ -212,12 +267,100 @@ START | 2019-12-31 00:00:00+00 | 2019-12-31 00:00:11+00 STOP | 2019-12-31 00:02:00+00 | 2019-12-31 00:02:00+00 ``` +### state_in + +```SQL +SELECT state_at( + (SELECT state_agg(ts, state) FROM states_test), + '2020-01-01 00:01:02+00' +); +``` +```output + state_at +---------- + ERROR +``` +```SQL +SELECT state_at_int( + (SELECT state_agg(ts, state) FROM states_test_5), + '2020-01-01 00:01:02+00' +); +``` +```output + state_at +---------- + 2 +``` +```SQL +SELECT state_at( + (SELECT state_agg(ts, state) FROM states_test), + '2020-01-01 00:01:00+00' +); +``` +```output + state_at +---------- + ERROR +``` +```SQL +SELECT state_at( + (SELECT state_agg(ts, state) FROM states_test), + '2020-01-01 00:00:05+00' +); +``` +```output + state_at +---------- + START +``` +```SQL +SELECT state_at( + (SELECT state_agg(ts, state) FROM states_test), + '2020-01-01 00:00:00+00' +); +``` +```output + state_at +---------- + START +``` +```SQL +SELECT state_at( + (SELECT state_agg(ts, state) FROM states_test), + '2019-12-31 23:59:59.999999+00' +); +``` +```output + state_at +---------- + +``` +```SQL +SELECT state_at( + (SELECT state_agg(ts, state) FROM states_test), + '2025-01-01 00:00:00+00' +); +``` +```output + state_at +---------- + STOP +``` +```SQL +SELECT (SELECT state_agg(ts, state) FROM states_test) -> state_at('2025-01-01 00:00:00+00'); +``` +```output + ?column? +---------- + STOP +``` + ## state_periods ```SQL SELECT start_time, end_time -FROM toolkit_experimental.state_periods( - (SELECT toolkit_experimental.state_agg(ts, state) FROM states_test), +FROM state_periods( + (SELECT state_agg(ts, state) FROM states_test), 'OK' ) ORDER BY start_time; @@ -229,10 +372,20 @@ start_time | end_time 2020-01-01 00:01:03+00 | 2020-01-01 00:02:00+00 ``` +```SQL +SELECT ((SELECT state_agg(ts, state) FROM states_test) -> state_periods('OK')).*; +``` +```output + start_time | end_time +------------------------+------------------------ + 2020-01-01 00:00:11+00 | 2020-01-01 00:01:00+00 + 2020-01-01 00:01:03+00 | 2020-01-01 00:02:00+00 +``` + ```SQL SELECT start_time, end_time -FROM toolkit_experimental.state_periods( - (SELECT toolkit_experimental.state_agg(ts, state) FROM states_test_4), +FROM state_periods( + (SELECT state_agg(ts, state) FROM states_test_4), 51351 ) ORDER BY start_time; @@ -246,8 +399,8 @@ start_time | end_time ```SQL SELECT start_time, end_time -FROM toolkit_experimental.state_periods( - (SELECT toolkit_experimental.state_agg(ts, state) FROM states_test), +FROM state_periods( + (SELECT state_agg(ts, state) FROM states_test), 'ANYTHING' ) ORDER BY start_time; @@ -260,10 +413,10 @@ start_time | end_time ## interpolated_state_timeline ```SQL -SELECT state, start_time, end_time FROM toolkit_experimental.interpolated_state_timeline( - (SELECT toolkit_experimental.state_agg(ts, state) FROM states_test), +SELECT state, start_time, end_time FROM interpolated_state_timeline( + (SELECT state_agg(ts, state) FROM states_test), '2019-12-31', '1 days', - (SELECT toolkit_experimental.state_agg(ts, state) FROM states_test_3) + (SELECT state_agg(ts, state) FROM states_test_3) ) ORDER BY start_time; ``` @@ -278,10 +431,45 @@ ERROR | 2020-01-01 00:01:00+00 | 2020-01-01 00:01:03+00 ``` ```SQL -SELECT state, start_time, end_time FROM toolkit_experimental.interpolated_state_timeline( - (SELECT toolkit_experimental.state_agg(ts, state) FROM states_test), +SELECT ((SELECT state_agg(ts, state) FROM states_test) -> interpolated_state_timeline( + '2019-12-31', '1 days', + (SELECT state_agg(ts, state) FROM states_test_3) +)).* +ORDER BY start_time; +``` +```output + state | start_time | end_time +-------+------------------------+------------------------ + START | 2019-12-31 00:00:00+00 | 2020-01-01 00:00:11+00 + OK | 2020-01-01 00:00:11+00 | 2020-01-01 00:01:00+00 + ERROR | 2020-01-01 00:01:00+00 | 2020-01-01 00:01:03+00 + OK | 2020-01-01 00:01:03+00 | 2020-01-01 00:02:00+00 + STOP | 2020-01-01 00:02:00+00 | 2020-01-01 00:02:00+00 +``` + +```SQL +SELECT state, start_time, end_time FROM interpolated_state_int_timeline( + (SELECT state_agg(ts, state) FROM states_test_5), + '2019-12-31', '1 days', + (SELECT state_agg(ts, state) FROM states_test_6) +) +ORDER BY start_time; +``` +```output +state | start_time | end_time +------+------------------------+----------------------- + 4 | 2019-12-31 00:00:00+00 | 2020-01-01 00:00:11+00 +51351 | 2020-01-01 00:00:11+00 | 2020-01-01 00:01:00+00 + 2 | 2020-01-01 00:01:00+00 | 2020-01-01 00:02:03+00 +51351 | 2020-01-01 00:02:03+00 | 2020-01-01 00:02:05+00 + -9 | 2020-01-01 00:02:05+00 | 2020-01-01 00:02:05+00 +``` + +```SQL +SELECT state, start_time, end_time FROM interpolated_state_timeline( + (SELECT state_agg(ts, state) FROM states_test), '2019-12-31', '5 days', - (SELECT toolkit_experimental.state_agg(ts, state) FROM states_test_3) + (SELECT state_agg(ts, state) FROM states_test_3) ) ORDER BY start_time; ``` @@ -296,10 +484,10 @@ ERROR | 2020-01-01 00:01:00+00 | 2020-01-01 00:01:03+00 ``` ```SQL -SELECT state, start_time, end_time FROM toolkit_experimental.interpolated_state_timeline( - (SELECT toolkit_experimental.state_agg(ts, state) FROM states_test), +SELECT state, start_time, end_time FROM interpolated_state_timeline( + (SELECT state_agg(ts, state) FROM states_test), '2019-12-31', '1 days', - (SELECT toolkit_experimental.state_agg(ts, state) FROM states_test_2) + (SELECT state_agg(ts, state) FROM states_test_2) ) ORDER BY start_time; ``` @@ -315,10 +503,10 @@ ERROR | 2020-01-01 00:01:00+00 | 2020-01-01 00:01:03+00 ``` ```SQL -SELECT state, start_time, end_time FROM toolkit_experimental.interpolated_state_timeline( - (SELECT toolkit_experimental.state_agg(ts, state) FROM states_test), +SELECT state, start_time, end_time FROM interpolated_state_timeline( + (SELECT state_agg(ts, state) FROM states_test), '2019-12-31', '5 days', - (SELECT toolkit_experimental.state_agg(ts, state) FROM states_test_2) + (SELECT state_agg(ts, state) FROM states_test_2) ) ORDER BY start_time; ``` @@ -333,15 +521,24 @@ ERROR | 2020-01-01 00:01:00+00 | 2020-01-01 00:01:03+00 STOP | 2020-01-01 00:02:00+00 | 2020-01-05 00:00:00+00 ``` +```SQL +SELECT (state_agg(ts, state) -> state_periods('OK')).* FROM states_test; +``` +```output + start_time | end_time +------------------------+------------------------ + 2020-01-01 00:00:11+00 | 2020-01-01 00:01:00+00 + 2020-01-01 00:01:03+00 | 2020-01-01 00:02:00+00 +``` ## interpolated_state_periods ```SQL -SELECT start_time, end_time FROM toolkit_experimental.interpolated_state_periods( - (SELECT toolkit_experimental.state_agg(ts, state) FROM states_test), +SELECT start_time, end_time FROM interpolated_state_periods( + (SELECT state_agg(ts, state) FROM states_test), 'OK', '2019-12-31', '1 days', - (SELECT toolkit_experimental.state_agg(ts, state) FROM states_test_3) + (SELECT state_agg(ts, state) FROM states_test_3) ) ORDER BY start_time; ``` @@ -353,11 +550,41 @@ start_time | end_time ``` ```SQL -SELECT start_time, end_time FROM toolkit_experimental.interpolated_state_periods( - (SELECT toolkit_experimental.state_agg(ts, state) FROM states_test), +SELECT ((SELECT state_agg(ts, state) FROM states_test) -> interpolated_state_periods( + 'OK', + '2019-12-31', '1 days', + (SELECT state_agg(ts, state) FROM states_test_3) +)).* +ORDER BY start_time; +``` +```output + start_time | end_time +------------------------+------------------------ + 2020-01-01 00:00:11+00 | 2020-01-01 00:01:00+00 + 2020-01-01 00:01:03+00 | 2020-01-01 00:02:00+00 +``` + +```SQL +SELECT start_time, end_time FROM interpolated_state_periods( + (SELECT state_agg(ts, state) FROM states_test), 'START', '2019-12-31', '5 days', - (SELECT toolkit_experimental.state_agg(ts, state) FROM states_test_3) + (SELECT state_agg(ts, state) FROM states_test_3) +) +ORDER BY start_time; +``` +```output +start_time | end_time +-----------------------+----------------------- +2019-12-31 00:00:00+00 | 2020-01-01 00:00:11+00 +``` + +```SQL +SELECT start_time, end_time FROM interpolated_state_periods( + (SELECT state_agg(ts, state) FROM states_test_5), + 4, + '2019-12-31', '5 days', + (SELECT state_agg(ts, state) FROM states_test_6) ) ORDER BY start_time; ``` @@ -368,11 +595,11 @@ start_time | end_time ``` ```SQL -SELECT start_time, end_time FROM toolkit_experimental.interpolated_state_periods( - (SELECT toolkit_experimental.state_agg(ts, state) FROM states_test), +SELECT start_time, end_time FROM interpolated_state_periods( + (SELECT state_agg(ts, state) FROM states_test), 'STOP', '2019-12-31', '1 days', - (SELECT toolkit_experimental.state_agg(ts, state) FROM states_test_2) + (SELECT state_agg(ts, state) FROM states_test_2) ) ORDER BY start_time; ``` @@ -384,11 +611,11 @@ start_time | end_time ``` ```SQL -SELECT start_time, end_time FROM toolkit_experimental.interpolated_state_periods( - (SELECT toolkit_experimental.state_agg(ts, state) FROM states_test), +SELECT start_time, end_time FROM interpolated_state_periods( + (SELECT state_agg(ts, state) FROM states_test), 'STOP', '2019-12-31', '5 days', - (SELECT toolkit_experimental.state_agg(ts, state) FROM states_test_2) + (SELECT state_agg(ts, state) FROM states_test_2) ) ORDER BY start_time; ``` @@ -399,7 +626,7 @@ start_time | end_time 2020-01-01 00:02:00+00 | 2020-01-05 00:00:00+00 ``` -## rolllup +## rollup ```SQL WITH buckets AS (SELECT @@ -440,11 +667,11 @@ FROM buckets; ```SQL WITH buckets AS (SELECT date_trunc('minute', ts) as dt, - toolkit_experimental.state_agg(ts, state) AS sa + state_agg(ts, state) AS sa FROM states_test GROUP BY date_trunc('minute', ts)) -SELECT toolkit_experimental.state_timeline( - toolkit_experimental.rollup(buckets.sa) +SELECT state_timeline( + rollup(buckets.sa) ) FROM buckets; ``` @@ -461,12 +688,12 @@ FROM buckets; ```SQL WITH buckets AS (SELECT date_trunc('minute', ts) as dt, - toolkit_experimental.state_agg(ts, state) AS sa + state_agg(ts, state) AS sa FROM states_test GROUP BY date_trunc('minute', ts) HAVING date_trunc('minute', ts) != '2020-01-01 00:01:00+00'::timestamptz) -SELECT toolkit_experimental.state_timeline( - toolkit_experimental.rollup(buckets.sa) +SELECT state_timeline( + rollup(buckets.sa) ) FROM buckets; ``` @@ -481,12 +708,12 @@ FROM buckets; ```SQL WITH buckets AS (SELECT date_trunc('minute', ts) as dt, - toolkit_experimental.state_agg(ts, state) AS sa + state_agg(ts, state) AS sa FROM states_test_5 GROUP BY date_trunc('minute', ts) HAVING date_trunc('minute', ts) != '2020-01-01 00:01:00+00'::timestamptz) -SELECT toolkit_experimental.state_int_timeline( - toolkit_experimental.rollup(buckets.sa) +SELECT state_int_timeline( + rollup(buckets.sa) ) FROM buckets; ``` @@ -497,3 +724,56 @@ FROM buckets; (51351,"2020-01-01 00:00:11+00","2020-01-01 00:02:05+00") (-9,"2020-01-01 00:02:05+00","2020-01-01 00:02:05+00") ``` + +## With continuous aggregate + +```SQL ,non-transactional,ignore-output +CREATE TABLE email_status ( + id BIGINT, + ts TIMESTAMPTZ, + status TEXT +); +SELECT create_hypertable('email_status','ts'); + +INSERT INTO email_status("ts", "id", "status") +VALUES +('2022-01-11 11:51:12',1,'draft'), +('2022-01-11 11:53:23',1,'queued'), +('2022-01-11 11:57:46',1,'sending'), +('2022-01-11 11:57:50',1,'sent'), +('2022-01-11 11:52:12',2,'draft'), +('2022-01-11 11:58:23',2,'queued'), +('2022-01-11 12:00:46',2,'sending'), +('2022-01-11 12:01:03',2,'bounced'); +``` + +```SQL ,non-transactional,ignore-output +CREATE MATERIALIZED VIEW sa +WITH (timescaledb.continuous) AS +SELECT time_bucket('1 minute'::interval, ts) AS bucket, + id, + state_agg(ts, status) AS agg +FROM email_status +GROUP BY bucket, id; +``` + +```SQL +SELECT rollup(agg) -> duration_in('draft') FROM sa WHERE id = 1; +``` +```output + ?column? +---------- + 00:02:11 +``` + +```SQL +SELECT (state_timeline(rollup(agg))).* FROM sa WHERE id = 2; +``` +```output + state | start_time | end_time +---------+------------------------+------------------------ + draft | 2022-01-11 11:52:12+00 | 2022-01-11 11:58:23+00 + queued | 2022-01-11 11:58:23+00 | 2022-01-11 12:00:46+00 + sending | 2022-01-11 12:00:46+00 | 2022-01-11 12:01:03+00 + bounced | 2022-01-11 12:01:03+00 | 2022-01-11 12:01:03+00 +``` diff --git a/docs/test_caggs.md b/docs/test_caggs.md index 81ff2ea5..f00358d0 100644 --- a/docs/test_caggs.md +++ b/docs/test_caggs.md @@ -23,7 +23,8 @@ AS SELECT hyperloglog(64, value1) as hll, counter_agg(time, value1) as counter, stats_agg(value1, value2) as stats, - timevector(time, value2) as tvec + timevector(time, value2) as tvec, + heartbeat_agg(time, time_bucket('7 day'::interval, time), '1w', '55m') as hb FROM test GROUP BY time_bucket('7 day'::interval, time); ``` @@ -114,3 +115,22 @@ ORDER BY week; 2020-07-20 00:00:00+00 | 168 2020-07-27 00:00:00+00 | 59 ``` + +```SQL +SELECT week, uptime(hb), interpolated_uptime(hb, LAG(hb) OVER (ORDER BY week)) +FROM weekly_aggs +WHERE week > '2020-06-01' +ORDER BY week; +``` +```output + week | uptime | interpolated_uptime +------------------------+-----------------+--------------------- + 2020-06-08 00:00:00+00 | 6 days 10:00:00 | 6 days 10:00:00 + 2020-06-15 00:00:00+00 | 6 days 10:00:00 | 6 days 10:00:00 + 2020-06-22 00:00:00+00 | 6 days 10:00:00 | 6 days 10:00:00 + 2020-06-29 00:00:00+00 | 6 days 10:00:00 | 6 days 10:00:00 + 2020-07-06 00:00:00+00 | 6 days 10:00:00 | 6 days 10:00:00 + 2020-07-13 00:00:00+00 | 6 days 10:00:00 | 6 days 10:00:00 + 2020-07-20 00:00:00+00 | 6 days 10:00:00 | 6 days 10:00:00 + 2020-07-27 00:00:00+00 | 2 days 06:05:00 | 2 days 06:05:00 +``` diff --git a/extension/Cargo.toml b/extension/Cargo.toml index 643a1f79..14c52529 100644 --- a/extension/Cargo.toml +++ b/extension/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "timescaledb_toolkit" -version = "1.14.0-dev" +version = "1.15.0" edition = "2021" [lib] @@ -17,8 +17,9 @@ pg_test = ["approx"] [dependencies] # Keep synchronized with `cargo install --version N.N.N cargo-pgx` in Readme.md and docker/ci/Dockerfile # Also `pgx-tests` down below in `dev-dependencies`. -pgx = "=0.6.1" -pgx-macros = "=0.6.1" +pgx = "=0.7.1" +pgx-macros = "=0.7.1" +pgx-sql-entity-graph = "=0.7.1" encodings = {path="../crates/encodings"} flat_serialize = {path="../crates/flat_serialize/flat_serialize"} flat_serialize_macro = {path="../crates/flat_serialize/flat_serialize_macro"} @@ -55,5 +56,5 @@ spfunc = "0.1.0" statrs = "0.15.0" [dev-dependencies] -pgx-tests = "=0.6.1" +pgx-tests = "=0.7.1" approx = "0.4.0" diff --git a/extension/src/accessors.rs b/extension/src/accessors.rs index 40af1908..8703240e 100644 --- a/extension/src/accessors.rs +++ b/extension/src/accessors.rs @@ -97,8 +97,36 @@ accessor! { open_time() } accessor! { high_time() } accessor! { low_time() } accessor! { close_time() } +accessor! { live_ranges() } +accessor! { dead_ranges() } +accessor! { uptime() } +accessor! { downtime() } +accessor! { into_values() } +accessor! { into_int_values() } +accessor! { state_timeline() } +accessor! { state_int_timeline() } // The rest are more complex, with String or other challenges. Leaving alone for now. +pg_type! { + #[derive(Debug)] + struct AccessorLiveAt { + time: u64, + } +} + +ron_inout_funcs!(AccessorLiveAt); + +#[pg_extern(immutable, parallel_safe, name = "live_at")] +pub fn accessor_live_at(ts: crate::raw::TimestampTz) -> AccessorLiveAt<'static> { + unsafe { + flatten! { + AccessorLiveAt { + time: ts.0.value() as u64, + } + } + } +} + pg_type! { #[derive(Debug)] struct AccessorStdDev<'input> { @@ -526,32 +554,31 @@ pub fn accessor_unnest() -> AccessorUnnest<'static> { } } -#[pg_schema] -pub mod toolkit_experimental { - use super::*; - - pg_type! { - #[derive(Debug)] - struct AccessorIntegral<'input> { - len: u32, - bytes: [u8; self.len], - } +pg_type! { + #[derive(Debug)] + struct AccessorIntegral<'input> { + len: u32, + bytes: [u8; self.len], } +} - // FIXME string IO - ron_inout_funcs!(AccessorIntegral); +// FIXME string IO +ron_inout_funcs!(AccessorIntegral); - #[pg_extern(immutable, parallel_safe, name = "integral")] - pub fn accessor_integral(unit: default!(&str, "'second'")) -> AccessorIntegral<'static> { - unsafe { - flatten! { - AccessorIntegral { - len: unit.len().try_into().unwrap(), - bytes: unit.as_bytes().into(), - } +#[pg_extern(immutable, parallel_safe, name = "integral")] +pub fn accessor_integral(unit: default!(&str, "'second'")) -> AccessorIntegral<'static> { + unsafe { + flatten! { + AccessorIntegral { + len: unit.len().try_into().unwrap(), + bytes: unit.as_bytes().into(), } } } +} +#[pg_schema] +pub mod toolkit_experimental { + use super::*; pg_type! { #[derive(Debug)] diff --git a/extension/src/aggregate_builder_tests.rs b/extension/src/aggregate_builder_tests.rs index 38bd43e8..2dbbdcec 100644 --- a/extension/src/aggregate_builder_tests.rs +++ b/extension/src/aggregate_builder_tests.rs @@ -87,24 +87,26 @@ mod tests { #[pg_test] fn test_anything_in_experimental_and_returns_first() { - Spi::execute(|client| { + Spi::connect(|mut client| { let output = client - .select( + .update( "SELECT toolkit_experimental.anything(val) \ FROM (VALUES ('foo'), ('bar'), ('baz')) as v(val)", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(output.as_deref(), Some("foo")); }) } #[pg_test] fn test_anything_has_correct_fn_names_and_def() { - Spi::execute(|client| { - let spec = get_aggregate_spec(&client, "anything"); + Spi::connect(|mut client| { + let spec = get_aggregate_spec(&mut client, "anything"); // output is // fn kind (`a`), volatility, parallel-safety, num args, final fn modify (is this right?) // transition type (`internal`) @@ -131,8 +133,8 @@ mod tests { #[pg_test] fn test_cagg_anything_has_correct_fn_names_and_def() { - Spi::execute(|client| { - let spec = get_aggregate_spec(&client, "cagg_anything"); + Spi::connect(|mut client| { + let spec = get_aggregate_spec(&mut client, "cagg_anything"); // output is // fn kind (`a`), volatility, parallel-safety, num args, final fn modify (is this right?) // transition type (`internal`) @@ -159,8 +161,8 @@ mod tests { #[pg_test] fn test_parallel_anything_has_correct_fn_names_and_def() { - Spi::execute(|client| { - let spec = get_aggregate_spec(&client, "parallel_anything"); + Spi::connect(|mut client| { + let spec = get_aggregate_spec(&mut client, "parallel_anything"); // output is // fn kind (`a`), volatility, parallel-safety, num args, final fn modify (is this right?) // transition type (`internal`) @@ -188,9 +190,9 @@ mod tests { // It gets annoying, and segfaulty to handle many arguments from the Spi. // For simplicity, we just return a single string representing the tuple // and use string-comparison. - fn get_aggregate_spec(client: &spi::SpiClient, aggregate_name: &str) -> String { + fn get_aggregate_spec(client: &mut spi::SpiClient, aggregate_name: &str) -> String { client - .select( + .update( &format!( r#"SELECT ( prokind, @@ -213,8 +215,10 @@ mod tests { None, None, ) + .unwrap() .first() .get_one::() + .unwrap() .expect("no aggregate found") } } diff --git a/extension/src/aggregate_utils.rs b/extension/src/aggregate_utils.rs index 602cc77d..211ea9c7 100644 --- a/extension/src/aggregate_utils.rs +++ b/extension/src/aggregate_utils.rs @@ -4,7 +4,7 @@ use pgx::pg_sys; // TODO move to func_utils once there are enough function to warrant one pub unsafe fn get_collation(fcinfo: pg_sys::FunctionCallInfo) -> Option { - if (*fcinfo).fncollation == 0 { + if (*fcinfo).fncollation == pg_sys::Oid::INVALID { None } else { Some((*fcinfo).fncollation) @@ -13,7 +13,7 @@ pub unsafe fn get_collation(fcinfo: pg_sys::FunctionCallInfo) -> Option Option { if fcinfo.is_null() { - Some(100) // TODO: default OID, there should be a constant for this + Some(unsafe { pg_sys::Oid::from_u32_unchecked(100) }) // TODO: default OID, there should be a constant for this } else { unsafe { get_collation(fcinfo) } } diff --git a/extension/src/asap.rs b/extension/src/asap.rs index dab699d1..d558cd80 100644 --- a/extension/src/asap.rs +++ b/extension/src/asap.rs @@ -202,9 +202,9 @@ mod tests { fn test_against_reference() { // Test our ASAP implementation against the reference implementation at http://www.futuredata.io.s3-website-us-west-2.amazonaws.com/asap/ // The sample data is the first 100 points of the second sample data set. Note that the dates are not important for this test. - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); - let mut result = client.select( + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); + let mut result = client.update( " SELECT value FROM unnest( @@ -223,46 +223,46 @@ mod tests { ) AS v(i, val) )) s", None, - None); + None).unwrap(); assert_relative_eq!( - result.next().unwrap()[1].value::().unwrap() as f32, + result.next().unwrap()[1].value::().unwrap().unwrap() as f32, 10.39 ); assert_relative_eq!( - result.next().unwrap()[1].value::().unwrap() as f32, + result.next().unwrap()[1].value::().unwrap().unwrap() as f32, 9.29 ); assert_relative_eq!( - result.next().unwrap()[1].value::().unwrap() as f32, + result.next().unwrap()[1].value::().unwrap().unwrap() as f32, 7.54 ); assert_relative_eq!( - result.next().unwrap()[1].value::().unwrap() as f32, + result.next().unwrap()[1].value::().unwrap().unwrap() as f32, 7.8 ); assert_relative_eq!( - result.next().unwrap()[1].value::().unwrap() as f32, + result.next().unwrap()[1].value::().unwrap().unwrap() as f32, 10.34 ); assert_relative_eq!( - result.next().unwrap()[1].value::().unwrap() as f32, + result.next().unwrap()[1].value::().unwrap().unwrap() as f32, 11.01 ); assert_relative_eq!( - result.next().unwrap()[1].value::().unwrap() as f32, + result.next().unwrap()[1].value::().unwrap().unwrap() as f32, 10.54 ); assert_relative_eq!( - result.next().unwrap()[1].value::().unwrap() as f32, + result.next().unwrap()[1].value::().unwrap().unwrap() as f32, 8.01 ); assert_relative_eq!( - result.next().unwrap()[1].value::().unwrap() as f32, + result.next().unwrap()[1].value::().unwrap().unwrap() as f32, 8.99 ); assert_relative_eq!( - result.next().unwrap()[1].value::().unwrap() as f32, + result.next().unwrap()[1].value::().unwrap().unwrap() as f32, 8.73 ); assert!(result.next().is_none()); @@ -271,8 +271,8 @@ mod tests { #[pg_test] fn test_asap_equivalence() { - Spi::execute(|client| { - let mut value_result = client.select( + Spi::connect(|mut client| { + let mut value_result = client.update( " SELECT time::text, value FROM unnest( @@ -291,9 +291,9 @@ mod tests { ) AS v(i, val) )) s", None, - None); + None).unwrap(); - let mut tvec_result = client.select( + let mut tvec_result = client.update( " SELECT time::text, value FROM unnest( @@ -314,13 +314,13 @@ mod tests { ), 10) ))", None, - None); + None).unwrap(); for _ in 0..10 { let v = value_result.next().unwrap(); let t = tvec_result.next().unwrap(); assert_eq!(v[1].value::<&str>(), t[1].value::<&str>()); - assert_eq!(v[2].value::(), t[2].value::()); + assert_eq!(v[2].value::().unwrap(), t[2].value::().unwrap()); } assert!(value_result.next().is_none()); assert!(tvec_result.next().is_none()); diff --git a/extension/src/candlestick.rs b/extension/src/candlestick.rs index 6db5fb9f..e11ca16d 100644 --- a/extension/src/candlestick.rs +++ b/extension/src/candlestick.rs @@ -531,22 +531,29 @@ mod tests { macro_rules! select_one { ($client:expr, $stmt:expr, $type:ty) => { - $client.select($stmt, None, None).first().get_one::<$type>() + $client + .update($stmt, None, None) + .unwrap() + .first() + .get_one::<$type>() + .unwrap() }; } macro_rules! select_two { ($client:expr, $stmt:expr, $type1:ty, $type2:ty) => { $client - .select($stmt, None, None) + .update($stmt, None, None) + .unwrap() .first() .get_two::<$type1, $type2>() + .unwrap() }; } #[pg_test] fn candlestick_single_point() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); let stmt = r#"SELECT candlestick(ts, open, high, low, close, volume)::text FROM ( @@ -569,8 +576,8 @@ mod tests { #[pg_test] fn candlestick_agg_single_point() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); let stmt = r#"SELECT candlestick_agg(ts, price, volume)::text FROM ( @@ -593,8 +600,8 @@ mod tests { #[pg_test] fn candlestick_accessors() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); for ohlc in ["open", "high", "low", "close"] { let stmt = format!( @@ -648,8 +655,8 @@ mod tests { #[pg_test] fn candlestick_agg_accessors() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); for ohlc in ["open", "high", "low", "close"] { let stmt = format!( @@ -686,8 +693,8 @@ mod tests { #[pg_test] fn candlestick_agg_extreme_values() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // timestamptz low and high val according to https://www.postgresql.org/docs/14/datatype-datetime.html for extreme_time in &["4713-01-01 00:00:00+00 BC", "294276-12-31 23:59:59+00"] { @@ -744,7 +751,7 @@ mod tests { #[pg_test] fn candlestick_null_inputs() { - Spi::execute(|client| { + Spi::connect(|mut client| { for (t, o, h, l, c, v) in &[ ("NULL", "NULL", "NULL", "NULL", "NULL", "NULL"), ("NULL", "1.0", "1.0", "1.0", "1.0", "1.0"), @@ -753,7 +760,7 @@ mod tests { ("now()", "1.0", "1.0", "NULL", "1.0", "1.0"), ("now()", "1.0", "1.0", "1.0", "NULL", "1.0"), ] { - let stmt = format!("SELECT candlestick({t}, {o}, {h}, {l}, {c}, {v})"); + let stmt = format!("SELECT candlestick({t}, {o}, {h}, {l}, {c}, {v})::TEXT"); let output = select_one!(client, &stmt, String); assert_eq!(output, None); } @@ -762,18 +769,18 @@ mod tests { #[pg_test] fn candlestick_agg_null_inputs() { - Spi::execute(|client| { + Spi::connect(|mut client| { for (ts, price, vol) in &[ ("NULL", "NULL", "NULL"), ("NULL", "1.0", "1.0"), ("now()", "NULL", "1.0"), ] { - let stmt = format!("SELECT candlestick_agg({ts}, {price}, {vol})"); + let stmt = format!("SELECT candlestick_agg({ts}, {price}, {vol})::text"); let output = select_one!(client, &stmt, String); assert_eq!(output, None); } - client.select("SET timezone TO 'UTC'", None, None); + client.update("SET timezone TO 'UTC'", None, None).unwrap(); let expected = "(\ version:1,\ @@ -796,8 +803,8 @@ mod tests { #[pg_test] fn candlestick_as_constructor() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); let stmt = r#"SELECT candlestick(ts, open, high, low, close, volume)::text @@ -806,7 +813,7 @@ mod tests { ('2022-08-02 00:00:00+00'::timestamptz, 9.0, 12.0, 3.0, 6.0, 1.0) ) AS v(ts, open, high, low, close, volume)"#; - let mut candlesticks = client.select(stmt, None, None); + let mut candlesticks = client.update(stmt, None, None).unwrap(); let expected = "(\ version:1,\ @@ -817,7 +824,10 @@ mod tests { volume:Transaction(vol:1,vwap:0)\ )"; - assert_eq!(Some(expected), candlesticks.next().unwrap()[1].value()); + assert_eq!( + Some(expected), + candlesticks.next().unwrap()[1].value().unwrap() + ); let expected = "(\ version:1,\ @@ -828,14 +838,17 @@ mod tests { volume:Transaction(vol:1,vwap:7)\ )"; - assert_eq!(Some(expected), candlesticks.next().unwrap()[1].value()); + assert_eq!( + Some(expected), + candlesticks.next().unwrap()[1].value().unwrap() + ); }); } #[pg_test] fn candlestick_agg_constant() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); let stmt = r#"SELECT date_trunc('day', ts)::text, @@ -864,8 +877,8 @@ mod tests { #[pg_test] fn candlestick_agg_strictly_increasing() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); let stmt = r#"SELECT date_trunc('day', ts)::text, @@ -894,8 +907,8 @@ mod tests { #[pg_test] fn candlestick_agg_strictly_decreasing() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); let stmt = r#"SELECT date_trunc('day', ts)::text, @@ -924,8 +937,8 @@ mod tests { #[pg_test] fn candlestick_agg_oscillating() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); let stmt = r#"SELECT date_trunc('day', ts)::text, @@ -961,8 +974,8 @@ mod tests { #[pg_test] fn candlestick_rollup() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); let stmt = r#"WITH t AS ( SELECT @@ -992,8 +1005,8 @@ mod tests { #[pg_test] fn candlestick_agg_rollup() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); let stmt = r#"WITH t AS ( SELECT diff --git a/extension/src/counter_agg.rs b/extension/src/counter_agg.rs index 2e760f55..9fdf7ae8 100644 --- a/extension/src/counter_agg.rs +++ b/extension/src/counter_agg.rs @@ -29,6 +29,10 @@ use crate::raw::tstzrange; use crate::raw::bytea; +mod accessors; + +use accessors::{CounterInterpolatedDeltaAccessor, CounterInterpolatedRateAccessor}; + // pg_type! can't handle generics so use a type alias to specify the type for `stats` type PgTypeHackStatsSummary2D = StatsSummary2D; // TODO wrap FlatSummary a la GaugeSummary - requires serialization version bump @@ -619,20 +623,6 @@ fn counter_agg_extrapolated_delta<'a>(summary: CounterSummary<'a>, method: &str) } } -// Public facing interpolated_delta -extension_sql!( - "\n\ - CREATE FUNCTION toolkit_experimental.interpolated_delta(summary countersummary,\n\ - start timestamptz,\n\ - duration interval,\n\ - prev countersummary,\n\ - next countersummary) RETURNS DOUBLE PRECISION\n\ - AS $$\n\ - SELECT interpolated_delta(summary,start,duration,prev,next) $$ LANGUAGE SQL IMMUTABLE PARALLEL SAFE;\n\ -", - name = "experimental_interpolated_delta", requires=[counter_agg_interpolated_delta] -); - #[pg_extern(name = "interpolated_delta", immutable, parallel_safe)] fn counter_agg_interpolated_delta<'a>( summary: CounterSummary<'a>, @@ -648,6 +638,32 @@ fn counter_agg_interpolated_delta<'a>( .delta() } +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_counter_interpolated_delta<'a>( + sketch: CounterSummary<'a>, + accessor: CounterInterpolatedDeltaAccessor<'a>, +) -> f64 { + let prev = if accessor.flags & 1 == 1 { + Some(accessor.prev.clone().into()) + } else { + None + }; + let next = if accessor.flags & 2 == 2 { + Some(accessor.next.clone().into()) + } else { + None + }; + + counter_agg_interpolated_delta( + sketch, + accessor.timestamp.into(), + accessor.interval.into(), + prev, + next, + ) +} + #[pg_operator(immutable, parallel_safe)] #[opname(->)] pub fn arrow_counter_agg_extrapolated_rate<'a>( @@ -668,21 +684,6 @@ fn counter_agg_extrapolated_rate<'a>(summary: CounterSummary<'a>, method: &str) } } -// Public facing interpolated_rate -extension_sql!( - "\n\ - CREATE FUNCTION toolkit_experimental.interpolated_rate(summary countersummary,\n\ - start timestamptz,\n\ - duration interval,\n\ - prev countersummary,\n\ - next countersummary) RETURNS DOUBLE PRECISION\n\ - AS $$\n\ - SELECT interpolated_rate(summary,start,duration,prev,next) $$ LANGUAGE SQL IMMUTABLE PARALLEL SAFE;\n\ -", - name = "experimental_interpolated_rate", - requires = [counter_agg_interpolated_rate] -); - #[pg_extern(name = "interpolated_rate", immutable, parallel_safe)] fn counter_agg_interpolated_rate<'a>( summary: CounterSummary<'a>, @@ -698,6 +699,32 @@ fn counter_agg_interpolated_rate<'a>( .rate() } +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_counter_interpolated_rate<'a>( + sketch: CounterSummary<'a>, + accessor: CounterInterpolatedRateAccessor<'a>, +) -> Option { + let prev = if accessor.flags & 1 == 1 { + Some(accessor.prev.clone().into()) + } else { + None + }; + let next = if accessor.flags & 2 == 2 { + Some(accessor.next.clone().into()) + } else { + None + }; + + counter_agg_interpolated_rate( + sketch, + accessor.timestamp.into(), + accessor.interval.into(), + prev, + next, + ) +} + #[pg_operator(immutable, parallel_safe)] #[opname(->)] pub fn arrow_counter_agg_num_elements<'a>( @@ -885,19 +912,23 @@ mod tests { macro_rules! select_one { ($client:expr, $stmt:expr, $type:ty) => { $client - .select($stmt, None, None) + .update($stmt, None, None) + .unwrap() .first() .get_one::<$type>() .unwrap() + .unwrap() }; } macro_rules! select_and_check_one { ($client:expr, $stmt:expr, $type:ty) => {{ let (a, b) = $client - .select($stmt, None, None) + .update($stmt, None, None) + .unwrap() .first() - .get_two::<$type, $type>(); + .get_two::<$type, $type>() + .unwrap(); assert_eq!(a, b); a.unwrap() }}; @@ -923,16 +954,19 @@ mod tests { #[pg_test] fn test_counter_aggregate() { - Spi::execute(|client| { + Spi::connect(|mut client| { // set search_path after defining our table so we don't pollute the wrong schema + client.update("SET timezone TO 'UTC'", None, None).unwrap(); let stmt = "SELECT format('toolkit_experimental, %s',current_setting('search_path'))"; let search_path = select_one!(client, stmt, String); - client.select( - &format!("SET LOCAL search_path TO {}", search_path), - None, - None, - ); - make_test_table(&client, "test"); + client + .update( + &format!("SET LOCAL search_path TO {}", search_path), + None, + None, + ) + .unwrap(); + make_test_table(&mut client, "test"); // NULL bounds are equivalent to none provided let stmt = "SELECT counter_agg(ts, val) FROM test"; @@ -976,7 +1010,7 @@ mod tests { assert_relative_eq!(select_and_check_one!(client, stmt, f64), 20.0 / 120.0); let stmt = "INSERT INTO test VALUES('2020-01-01 00:02:00+00', 10.0), ('2020-01-01 00:03:00+00', 20.0), ('2020-01-01 00:04:00+00', 10.0)"; - client.select(stmt, None, None); + client.update(stmt, None, None).unwrap(); let stmt = "SELECT \ slope(counter_agg(ts, val)), \ @@ -997,15 +1031,14 @@ mod tests { assert_relative_eq!(select_and_check_one!(client, stmt, f64), 1.0); let stmt = "SELECT \ - counter_zero_time(counter_agg(ts, val)), \ - counter_agg(ts, val)->counter_zero_time() \ + counter_zero_time(counter_agg(ts, val))::TEXT, \ + (counter_agg(ts, val)->counter_zero_time())::TEXT \ FROM test"; - let zp = select_and_check_one!(client, stmt, i64); - let real_zp = select_one!(client, "SELECT '2019-12-31 23:59:00+00'::timestamptz", i64); - assert_eq!(zp, real_zp); + let zp = select_and_check_one!(client, stmt, String); + assert_eq!(&zp, "2019-12-31 23:59:00+00"); let stmt = "INSERT INTO test VALUES('2020-01-01 00:08:00+00', 30.0), ('2020-01-01 00:10:00+00', 30.0), ('2020-01-01 00:10:30+00', 10.0), ('2020-01-01 00:20:00+00', 40.0)"; - client.select(stmt, None, None); + client.update(stmt, None, None).unwrap(); let stmt = "SELECT \ num_elements(counter_agg(ts, val)), \ @@ -1039,13 +1072,15 @@ mod tests { #[pg_test] fn test_counter_io() { - Spi::execute(|client| { - client.select( - "CREATE TABLE test(ts timestamptz, val DOUBLE PRECISION)", - None, - None, - ); - client.select("SET TIME ZONE 'UTC'", None, None); + Spi::connect(|mut client| { + client + .update( + "CREATE TABLE test(ts timestamptz, val DOUBLE PRECISION)", + None, + None, + ) + .unwrap(); + client.update("SET TIME ZONE 'UTC'", None, None).unwrap(); let stmt = "INSERT INTO test VALUES\ ('2020-01-01 00:00:00+00', 10.0),\ ('2020-01-01 00:01:00+00', 20.0),\ @@ -1056,7 +1091,7 @@ mod tests { ('2020-01-01 00:06:00+00', 10.0),\ ('2020-01-01 00:07:00+00', 30.0),\ ('2020-01-01 00:08:00+00', 10.0)"; - client.select(stmt, None, None); + client.update(stmt, None, None).unwrap(); let expected = "(\ version:1,\ @@ -1186,8 +1221,8 @@ mod tests { #[pg_test] fn delta_after_counter_decrease() { - Spi::execute(|client| { - decrease(&client); + Spi::connect(|mut client| { + decrease(&mut client); let stmt = "SELECT delta(counter_agg(ts, val)) FROM test"; // 10 after 30 means there was a reset so we add 30 + 10 = 40. // Delta from 30 to 40 => 10 @@ -1197,8 +1232,8 @@ mod tests { #[pg_test] fn delta_after_counter_increase() { - Spi::execute(|client| { - increase(&client); + Spi::connect(|mut client| { + increase(&mut client); let stmt = "SELECT delta(counter_agg(ts, val)) FROM test"; assert_eq!(20.0, select_one!(client, stmt, f64)); }); @@ -1206,8 +1241,8 @@ mod tests { #[pg_test] fn delta_after_counter_decrease_then_increase_to_same_value() { - Spi::execute(|client| { - decrease_then_increase_to_same_value(&client); + Spi::connect(|mut client| { + decrease_then_increase_to_same_value(&mut client); let stmt = "SELECT delta(counter_agg(ts, val)) FROM test"; // 10 after 30 means there was a reset so we add 30 + 10 + 30 = 70. // Delta from 30 to 70 => 30 @@ -1217,8 +1252,8 @@ mod tests { #[pg_test] fn delta_after_counter_increase_then_decrease_to_same_value() { - Spi::execute(|client| { - increase_then_decrease_to_same_value(&client); + Spi::connect(|mut client| { + increase_then_decrease_to_same_value(&mut client); let stmt = "SELECT delta(counter_agg(ts, val)) FROM test"; // In this case, counter goes 10, 30, 40 (reset + 10). // Delta from 10 to 40 => 30 @@ -1228,8 +1263,8 @@ mod tests { #[pg_test] fn idelta_left_after_counter_decrease() { - Spi::execute(|client| { - decrease(&client); + Spi::connect(|mut client| { + decrease(&mut client); let stmt = "SELECT idelta_left(counter_agg(ts, val)) FROM test"; assert_eq!(10.0, select_one!(client, stmt, f64)); }); @@ -1237,8 +1272,8 @@ mod tests { #[pg_test] fn idelta_left_after_counter_increase() { - Spi::execute(|client| { - increase(&client); + Spi::connect(|mut client| { + increase(&mut client); let stmt = "SELECT idelta_left(counter_agg(ts, val)) FROM test"; assert_eq!(20.0, select_one!(client, stmt, f64)); }); @@ -1246,8 +1281,8 @@ mod tests { #[pg_test] fn idelta_left_after_counter_increase_then_decrease_to_same_value() { - Spi::execute(|client| { - increase_then_decrease_to_same_value(&client); + Spi::connect(|mut client| { + increase_then_decrease_to_same_value(&mut client); let stmt = "SELECT idelta_left(counter_agg(ts, val)) FROM test"; assert_eq!(20.0, select_one!(client, stmt, f64)); }); @@ -1255,8 +1290,8 @@ mod tests { #[pg_test] fn idelta_left_after_counter_decrease_then_increase_to_same_value() { - Spi::execute(|client| { - decrease_then_increase_to_same_value(&client); + Spi::connect(|mut client| { + decrease_then_increase_to_same_value(&mut client); let stmt = "SELECT idelta_left(counter_agg(ts, val)) FROM test"; assert_eq!(10.0, select_one!(client, stmt, f64)); @@ -1265,8 +1300,8 @@ mod tests { #[pg_test] fn idelta_right_after_counter_decrease() { - Spi::execute(|client| { - decrease(&client); + Spi::connect(|mut client| { + decrease(&mut client); let stmt = "SELECT idelta_right(counter_agg(ts, val)) FROM test"; assert_eq!(10.0, select_one!(client, stmt, f64)); }); @@ -1274,8 +1309,8 @@ mod tests { #[pg_test] fn idelta_right_after_counter_increase() { - Spi::execute(|client| { - increase(&client); + Spi::connect(|mut client| { + increase(&mut client); let stmt = "SELECT idelta_right(counter_agg(ts, val)) FROM test"; assert_eq!(20.0, select_one!(client, stmt, f64)); }); @@ -1283,8 +1318,8 @@ mod tests { #[pg_test] fn idelta_right_after_counter_increase_then_decrease_to_same_value() { - Spi::execute(|client| { - increase_then_decrease_to_same_value(&client); + Spi::connect(|mut client| { + increase_then_decrease_to_same_value(&mut client); let stmt = "SELECT idelta_right(counter_agg(ts, val)) FROM test"; assert_eq!(10.0, select_one!(client, stmt, f64)); }); @@ -1292,8 +1327,8 @@ mod tests { #[pg_test] fn idelta_right_after_counter_decrease_then_increase_to_same_value() { - Spi::execute(|client| { - decrease_then_increase_to_same_value(&client); + Spi::connect(|mut client| { + decrease_then_increase_to_same_value(&mut client); let stmt = "SELECT idelta_right(counter_agg(ts, val)) FROM test"; assert_eq!(20.0, select_one!(client, stmt, f64)); }); @@ -1301,14 +1336,15 @@ mod tests { #[pg_test] fn counter_agg_interpolation() { - Spi::execute(|client| { - client.select( + Spi::connect(|mut client| { + client.update( "CREATE TABLE test(time timestamptz, value double precision, bucket timestamptz)", None, None, - ); - client.select( - r#"INSERT INTO test VALUES + ).unwrap(); + client + .update( + r#"INSERT INTO test VALUES ('2020-1-1 10:00'::timestamptz, 10.0, '2020-1-1'::timestamptz), ('2020-1-1 12:00'::timestamptz, 40.0, '2020-1-1'::timestamptz), ('2020-1-1 16:00'::timestamptz, 20.0, '2020-1-1'::timestamptz), @@ -1318,13 +1354,15 @@ mod tests { ('2020-1-3 4:00'::timestamptz, 30.0, '2020-1-3'::timestamptz), ('2020-1-3 12:00'::timestamptz, 0.0, '2020-1-3'::timestamptz), ('2020-1-3 16:00'::timestamptz, 35.0, '2020-1-3'::timestamptz)"#, - None, - None, - ); - - let mut deltas = client.select( - r#"SELECT - toolkit_experimental.interpolated_delta( + None, + None, + ) + .unwrap(); + + let mut deltas = client + .update( + r#"SELECT + interpolated_delta( agg, bucket, '1 day'::interval, @@ -1336,26 +1374,33 @@ mod tests { GROUP BY bucket ) s ORDER BY bucket"#, - None, - None, - ); + None, + None, + ) + .unwrap(); // Day 1, start at 10, interpolated end of day is 10 (after reset), reset at 40 and 20 assert_eq!( - deltas.next().unwrap()[1].value(), + deltas.next().unwrap()[1].value().unwrap(), Some(10. + 40. + 20. - 10.) ); // Day 2, interpolated start is 10, interpolated end is 27.5, reset at 50 - assert_eq!(deltas.next().unwrap()[1].value(), Some(27.5 + 50. - 10.)); + assert_eq!( + deltas.next().unwrap()[1].value().unwrap(), + Some(27.5 + 50. - 10.) + ); // Day 3, interpolated start is 27.5, end is 35, reset at 30 - assert_eq!(deltas.next().unwrap()[1].value(), Some(35. + 30. - 27.5)); + assert_eq!( + deltas.next().unwrap()[1].value().unwrap(), + Some(35. + 30. - 27.5) + ); assert!(deltas.next().is_none()); - // test that the non experimental version also returns the same result - let mut deltas = client.select( - r#"SELECT - interpolated_delta( - agg, + // test that the arrow version also returns the same result + let mut deltas = client + .update( + r#"SELECT + agg -> interpolated_delta( bucket, '1 day'::interval, LAG(agg) OVER (ORDER BY bucket), @@ -1366,24 +1411,32 @@ mod tests { GROUP BY bucket ) s ORDER BY bucket"#, - None, - None, - ); + None, + None, + ) + .unwrap(); // Day 1, start at 10, interpolated end of day is 10 (after reset), reset at 40 and 20 assert_eq!( - deltas.next().unwrap()[1].value(), + deltas.next().unwrap()[1].value().unwrap(), Some(10. + 40. + 20. - 10.) ); // Day 2, interpolated start is 10, interpolated end is 27.5, reset at 50 - assert_eq!(deltas.next().unwrap()[1].value(), Some(27.5 + 50. - 10.)); + assert_eq!( + deltas.next().unwrap()[1].value().unwrap(), + Some(27.5 + 50. - 10.) + ); // Day 3, interpolated start is 27.5, end is 35, reset at 30 - assert_eq!(deltas.next().unwrap()[1].value(), Some(35. + 30. - 27.5)); + assert_eq!( + deltas.next().unwrap()[1].value().unwrap(), + Some(35. + 30. - 27.5) + ); assert!(deltas.next().is_none()); - let mut rates = client.select( - r#"SELECT - toolkit_experimental.interpolated_rate( + let mut rates = client + .update( + r#"SELECT + interpolated_rate( agg, bucket, '1 day'::interval, @@ -1395,32 +1448,33 @@ mod tests { GROUP BY bucket ) s ORDER BY bucket"#, - None, - None, - ); + None, + None, + ) + .unwrap(); // Day 1, 14 hours (rate is per second) assert_eq!( - rates.next().unwrap()[1].value(), + rates.next().unwrap()[1].value().unwrap(), Some((10. + 40. + 20. - 10.) / (14. * 60. * 60.)) ); // Day 2, 24 hours assert_eq!( - rates.next().unwrap()[1].value(), + rates.next().unwrap()[1].value().unwrap(), Some((27.5 + 50. - 10.) / (24. * 60. * 60.)) ); // Day 3, 16 hours assert_eq!( - rates.next().unwrap()[1].value(), + rates.next().unwrap()[1].value().unwrap(), Some((35. + 30. - 27.5) / (16. * 60. * 60.)) ); - - // test that the non experimental version also returns the same result assert!(rates.next().is_none()); - let mut rates = client.select( - r#"SELECT - interpolated_rate( - agg, + + // test that the arrow operator version also returns the same result + let mut rates = client + .update( + r#"SELECT + agg -> interpolated_rate( bucket, '1 day'::interval, LAG(agg) OVER (ORDER BY bucket), @@ -1431,23 +1485,24 @@ mod tests { GROUP BY bucket ) s ORDER BY bucket"#, - None, - None, - ); + None, + None, + ) + .unwrap(); // Day 1, 14 hours (rate is per second) assert_eq!( - rates.next().unwrap()[1].value(), + rates.next().unwrap()[1].value().unwrap(), Some((10. + 40. + 20. - 10.) / (14. * 60. * 60.)) ); // Day 2, 24 hours assert_eq!( - rates.next().unwrap()[1].value(), + rates.next().unwrap()[1].value().unwrap(), Some((27.5 + 50. - 10.) / (24. * 60. * 60.)) ); // Day 3, 16 hours assert_eq!( - rates.next().unwrap()[1].value(), + rates.next().unwrap()[1].value().unwrap(), Some((35. + 30. - 27.5) / (16. * 60. * 60.)) ); assert!(rates.next().is_none()); @@ -1456,27 +1511,30 @@ mod tests { #[pg_test] fn interpolated_delta_with_aligned_point() { - Spi::execute(|client| { - client.select( + Spi::connect(|mut client| { + client.update( "CREATE TABLE test(time timestamptz, value double precision, bucket timestamptz)", None, None, - ); - client.select( - r#"INSERT INTO test VALUES + ).unwrap(); + client + .update( + r#"INSERT INTO test VALUES ('2020-1-1 10:00'::timestamptz, 10.0, '2020-1-1'::timestamptz), ('2020-1-1 12:00'::timestamptz, 40.0, '2020-1-1'::timestamptz), ('2020-1-1 16:00'::timestamptz, 20.0, '2020-1-1'::timestamptz), ('2020-1-2 0:00'::timestamptz, 15.0, '2020-1-2'::timestamptz), ('2020-1-2 12:00'::timestamptz, 50.0, '2020-1-2'::timestamptz), ('2020-1-2 20:00'::timestamptz, 25.0, '2020-1-2'::timestamptz)"#, - None, - None, - ); - - let mut deltas = client.select( - r#"SELECT - toolkit_experimental.interpolated_delta( + None, + None, + ) + .unwrap(); + + let mut deltas = client + .update( + r#"SELECT + interpolated_delta( agg, bucket, '1 day'::interval, @@ -1488,24 +1546,28 @@ mod tests { GROUP BY bucket ) s ORDER BY bucket"#, - None, - None, - ); + None, + None, + ) + .unwrap(); // Day 1, start at 10, interpolated end of day is 15 (after reset), reset at 40 and 20 assert_eq!( - deltas.next().unwrap()[1].value(), + deltas.next().unwrap()[1].value().unwrap(), Some(15. + 40. + 20. - 10.) ); // Day 2, start is 15, end is 25, reset at 50 - assert_eq!(deltas.next().unwrap()[1].value(), Some(25. + 50. - 15.)); + assert_eq!( + deltas.next().unwrap()[1].value().unwrap(), + Some(25. + 50. - 15.) + ); assert!(deltas.next().is_none()); }); } #[pg_test] fn irate_left_arrow_match() { - Spi::execute(|client| { - make_test_table(&client, "test"); + Spi::connect(|mut client| { + make_test_table(&mut client, "test"); assert_relative_eq!( select_and_check_one!( @@ -1523,8 +1585,8 @@ mod tests { #[pg_test] fn irate_right_arrow_match() { - Spi::execute(|client| { - make_test_table(&client, "test"); + Spi::connect(|mut client| { + make_test_table(&mut client, "test"); assert_relative_eq!( select_and_check_one!( @@ -1542,8 +1604,8 @@ mod tests { #[pg_test] fn idelta_left_arrow_match() { - Spi::execute(|client| { - make_test_table(&client, "test"); + Spi::connect(|mut client| { + make_test_table(&mut client, "test"); assert_relative_eq!( select_and_check_one!( @@ -1561,8 +1623,8 @@ mod tests { #[pg_test] fn idelta_right_arrow_match() { - Spi::execute(|client| { - make_test_table(&client, "test"); + Spi::connect(|mut client| { + make_test_table(&mut client, "test"); assert_relative_eq!( select_and_check_one!( @@ -1580,15 +1642,15 @@ mod tests { #[pg_test] fn num_resets_arrow_match() { - Spi::execute(|client| { - make_test_table(&client, "test"); + Spi::connect(|mut client| { + make_test_table(&mut client, "test"); assert_relative_eq!( select_and_check_one!( client, "SELECT \ - num_resets(counter_agg(ts, val)), \ - counter_agg(ts, val) -> num_resets() \ + num_resets(counter_agg(ts, val))::float, \ + (counter_agg(ts, val) -> num_resets())::float \ FROM test", f64 ), @@ -1599,8 +1661,8 @@ mod tests { #[pg_test] fn first_and_last_val() { - Spi::execute(|client| { - make_test_table(&client, "test"); + Spi::connect(|mut client| { + make_test_table(&mut client, "test"); assert_relative_eq!( select_one!( @@ -1628,8 +1690,8 @@ mod tests { #[pg_test] fn first_and_last_val_arrow_match() { - Spi::execute(|client| { - make_test_table(&client, "test"); + Spi::connect(|mut client| { + make_test_table(&mut client, "test"); assert_relative_eq!( select_and_check_one!( @@ -1659,9 +1721,9 @@ mod tests { #[pg_test] fn first_and_last_time() { - Spi::execute(|client| { - make_test_table(&client, "test"); - client.select("SET TIME ZONE 'UTC'", None, None); + Spi::connect(|mut client| { + make_test_table(&mut client, "test"); + client.update("SET TIME ZONE 'UTC'", None, None).unwrap(); assert_eq!( select_one!( @@ -1689,9 +1751,9 @@ mod tests { #[pg_test] fn first_and_last_time_arrow_match() { - Spi::execute(|client| { - make_test_table(&client, "test"); - client.select("SET TIME ZONE 'UTC'", None, None); + Spi::connect(|mut client| { + make_test_table(&mut client, "test"); + client.update("SET TIME ZONE 'UTC'", None, None).unwrap(); assert_eq!( select_and_check_one!( @@ -1721,7 +1783,7 @@ mod tests { // #[pg_test] // fn test_combine_aggregate(){ - // Spi::execute(|client| { + // Spi::connect(|mut client| { // }); // } @@ -1729,85 +1791,103 @@ mod tests { #[cfg(any(test, feature = "pg_test"))] pub(crate) mod testing { - pub fn decrease(client: &pgx::SpiClient) { - client.select( - "CREATE TABLE test(ts timestamptz, val DOUBLE PRECISION)", - None, - None, - ); - client.select("SET TIME ZONE 'UTC'", None, None); - client.select( - r#"INSERT INTO test VALUES + pub fn decrease(client: &mut pgx::spi::SpiClient) { + client + .update( + "CREATE TABLE test(ts timestamptz, val DOUBLE PRECISION)", + None, + None, + ) + .unwrap(); + client.update("SET TIME ZONE 'UTC'", None, None).unwrap(); + client + .update( + r#"INSERT INTO test VALUES ('2020-01-01 00:00:00+00', 30.0), ('2020-01-01 00:07:00+00', 10.0)"#, - None, - None, - ); + None, + None, + ) + .unwrap(); } - pub fn increase(client: &pgx::SpiClient) { - client.select( - "CREATE TABLE test(ts timestamptz, val DOUBLE PRECISION)", - None, - None, - ); - client.select("SET TIME ZONE 'UTC'", None, None); - client.select( - r#"INSERT INTO test VALUES + pub fn increase(client: &mut pgx::spi::SpiClient) { + client + .update( + "CREATE TABLE test(ts timestamptz, val DOUBLE PRECISION)", + None, + None, + ) + .unwrap(); + client.update("SET TIME ZONE 'UTC'", None, None).unwrap(); + client + .update( + r#"INSERT INTO test VALUES ('2020-01-01 00:00:00+00', 10.0), ('2020-01-01 00:07:00+00', 30.0)"#, - None, - None, - ); + None, + None, + ) + .unwrap(); } - pub fn decrease_then_increase_to_same_value(client: &pgx::SpiClient) { - client.select( - "CREATE TABLE test(ts timestamptz, val DOUBLE PRECISION)", - None, - None, - ); - client.select("SET TIME ZONE 'UTC'", None, None); - client.select( - r#"INSERT INTO test VALUES + pub fn decrease_then_increase_to_same_value(client: &mut pgx::spi::SpiClient) { + client + .update( + "CREATE TABLE test(ts timestamptz, val DOUBLE PRECISION)", + None, + None, + ) + .unwrap(); + client.update("SET TIME ZONE 'UTC'", None, None).unwrap(); + client + .update( + r#"INSERT INTO test VALUES ('2020-01-01 00:00:00+00', 30.0), ('2020-01-01 00:07:00+00', 10.0), ('2020-01-01 00:08:00+00', 30.0)"#, - None, - None, - ); + None, + None, + ) + .unwrap(); } - pub fn increase_then_decrease_to_same_value(client: &pgx::SpiClient) { - client.select( - "CREATE TABLE test(ts timestamptz, val DOUBLE PRECISION)", - None, - None, - ); - client.select("SET TIME ZONE 'UTC'", None, None); - client.select( - r#"INSERT INTO test VALUES + pub fn increase_then_decrease_to_same_value(client: &mut pgx::spi::SpiClient) { + client + .update( + "CREATE TABLE test(ts timestamptz, val DOUBLE PRECISION)", + None, + None, + ) + .unwrap(); + client.update("SET TIME ZONE 'UTC'", None, None).unwrap(); + client + .update( + r#"INSERT INTO test VALUES ('2020-01-01 00:00:00+00', 10.0), ('2020-01-01 00:07:00+00', 30.0), ('2020-01-01 00:08:00+00', 10.0)"#, - None, - None, - ); + None, + None, + ) + .unwrap(); } - pub fn make_test_table(client: &pgx::SpiClient, name: &str) { - client.select( - &format!( - "CREATE TABLE {}(ts timestamptz, val DOUBLE PRECISION)", - name - ), - None, - None, - ); - client.select( + pub fn make_test_table(client: &mut pgx::spi::SpiClient, name: &str) { + client + .update( + &format!( + "CREATE TABLE {}(ts timestamptz, val DOUBLE PRECISION)", + name + ), + None, + None, + ) + .unwrap(); + client.update( &format!("INSERT INTO {} VALUES('2020-01-01 00:00:00+00', 10.0), ('2020-01-01 00:01:00+00', 20.0)", name), None, None, - ); + ).unwrap(); } } diff --git a/extension/src/counter_agg/accessors.rs b/extension/src/counter_agg/accessors.rs new file mode 100644 index 00000000..54ec7041 --- /dev/null +++ b/extension/src/counter_agg/accessors.rs @@ -0,0 +1,91 @@ +use pgx::*; + +use crate::{ + counter_agg::{CounterSummary, CounterSummaryData, MetricSummary}, + datum_utils::interval_to_ms, + pg_type, ron_inout_funcs, +}; + +use tspoint::TSPoint; + +pg_type! { + struct CounterInterpolatedRateAccessor { + timestamp : i64, + interval : i64, + prev : CounterSummaryData, + next : CounterSummaryData, + flags : u64, + } +} + +ron_inout_funcs!(CounterInterpolatedRateAccessor); + +#[pg_extern(immutable, parallel_safe, name = "interpolated_rate")] +fn counter_interpolated_rate_accessor<'a>( + start: crate::raw::TimestampTz, + duration: crate::raw::Interval, + prev: Option>, + next: Option>, +) -> CounterInterpolatedRateAccessor<'static> { + fn empty_summary<'b>() -> Option> { + let tmp = TSPoint { ts: 0, val: 0.0 }; + let tmp = MetricSummary::new(&tmp, None); + let tmp = CounterSummary::from_internal_counter_summary(tmp); + Some(tmp) + } + + let flags = u64::from(prev.is_some()) + if next.is_some() { 2 } else { 0 }; + let prev = prev.or_else(empty_summary).unwrap().0; + let next = next.or_else(empty_summary).unwrap().0; + let interval = interval_to_ms(&start, &duration); + crate::build! { + CounterInterpolatedRateAccessor { + timestamp : start.into(), + interval, + prev, + next, + flags, + } + } +} + +pg_type! { + struct CounterInterpolatedDeltaAccessor { + timestamp : i64, + interval : i64, + prev : CounterSummaryData, + next : CounterSummaryData, + flags : u64, + } +} + +ron_inout_funcs!(CounterInterpolatedDeltaAccessor); + +#[pg_extern(immutable, parallel_safe, name = "interpolated_delta")] +fn counter_interpolated_delta_accessor<'a>( + start: crate::raw::TimestampTz, + duration: crate::raw::Interval, + prev: Option>, + next: Option>, +) -> CounterInterpolatedDeltaAccessor<'static> { + fn empty_summary<'b>() -> Option> { + let tmp = TSPoint { ts: 0, val: 0.0 }; + let tmp = MetricSummary::new(&tmp, None); + let tmp = CounterSummary::from_internal_counter_summary(tmp); + Some(tmp) + } + + let flags = u64::from(prev.is_some()) + if next.is_some() { 2 } else { 0 }; + let prev = prev.or_else(empty_summary).unwrap().0; + let next = next.or_else(empty_summary).unwrap().0; + let interval = interval_to_ms(&start, &duration); + crate::build! { + CounterInterpolatedDeltaAccessor { + timestamp : start.into(), + interval, + prev, + next, + flags, + } + } +} diff --git a/extension/src/countminsketch.rs b/extension/src/countminsketch.rs index 6240fdf0..da1147c1 100644 --- a/extension/src/countminsketch.rs +++ b/extension/src/countminsketch.rs @@ -139,32 +139,40 @@ mod tests { #[pg_test] fn test_countminsketch() { - Spi::execute(|client| { - client.select("CREATE TABLE test (data TEXT)", None, None); - client.select("INSERT INTO test SELECT generate_series(1, 100)::TEXT UNION ALL SELECT generate_series(1, 50)::TEXT", None, None); + Spi::connect(|mut client| { + client + .update("CREATE TABLE test (data TEXT)", None, None) + .unwrap(); + client.update("INSERT INTO test SELECT generate_series(1, 100)::TEXT UNION ALL SELECT generate_series(1, 50)::TEXT", None, None).unwrap(); let sanity = client - .select("SELECT COUNT(*) FROM test", None, None) + .update("SELECT COUNT(*) FROM test", None, None) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(Some(150), sanity); - client.select( - "CREATE VIEW sketch AS \ + client + .update( + "CREATE VIEW sketch AS \ SELECT toolkit_experimental.count_min_sketch(data, 0.01, 0.01) \ FROM test", - None, - None, - ); + None, + None, + ) + .unwrap(); let sanity = client - .select("SELECT COUNT(*) FROM sketch", None, None) + .update("SELECT COUNT(*) FROM sketch", None, None) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert!(sanity.unwrap_or(0) > 0); let (col1, col2, col3) = client - .select( + .update( "SELECT \ toolkit_experimental.approx_count('1', count_min_sketch), \ toolkit_experimental.approx_count('51', count_min_sketch), \ @@ -173,8 +181,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_three::(); + .get_three::() + .unwrap(); // 0.01 => error param to the sketch, 150 => number of items added to the sketch let err_margin = 0.01 * 150.0; @@ -193,9 +203,9 @@ mod tests { #[pg_test] fn test_countminsketch_combine() { - Spi::execute(|client| { + Spi::connect(|mut client| { let combined = client - .select( + .update( "SELECT toolkit_experimental.approx_count('1', toolkit_experimental.count_min_sketch(v::text, 0.01, 0.01)) FROM (SELECT * FROM generate_series(1, 100) v \ UNION ALL \ @@ -203,8 +213,8 @@ mod tests { None, None, ) - .first() - .get_one::(); + .unwrap().first() + .get_one::().unwrap(); let expected = 2; // 0.01 => error param to the sketch, 200 => number of items added to the sketch @@ -221,18 +231,20 @@ mod tests { #[pg_test] fn countminsketch_io_test() { - Spi::execute(|client| { - client.select("CREATE TABLE io_test (value TEXT)", None, None); - client.select("INSERT INTO io_test VALUES ('lorem'), ('ipsum'), ('dolor'), ('sit'), ('amet'), ('consectetur'), ('adipiscing'), ('elit')", None, None); + Spi::connect(|mut client| { + client + .update("CREATE TABLE io_test (value TEXT)", None, None) + .unwrap(); + client.update("INSERT INTO io_test VALUES ('lorem'), ('ipsum'), ('dolor'), ('sit'), ('amet'), ('consectetur'), ('adipiscing'), ('elit')", None, None).unwrap(); let sketch = client - .select( + .update( "SELECT toolkit_experimental.count_min_sketch(value, 0.5, 0.01)::text FROM io_test", None, None, ) - .first() - .get_one::(); + .unwrap().first() + .get_one::().unwrap(); let expected = "(\ version:1,\ @@ -253,30 +265,32 @@ mod tests { #[pg_test] fn test_cms_null_input_yields_null_output() { - Spi::execute(|client| { + Spi::connect(|mut client| { let output = client - .select( + .update( "SELECT toolkit_experimental.count_min_sketch(NULL::TEXT, 0.1, 0.1)::TEXT", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(output, None) }) } #[pg_test] fn test_approx_count_null_input_yields_null_output() { - Spi::execute(|client| { + Spi::connect(|mut client| { let output = client - .select( + .update( "SELECT toolkit_experimental.approx_count('1'::text, NULL::toolkit_experimental.countminsketch)", None, None, ) - .first() - .get_one::(); + .unwrap().first() + .get_one::().unwrap(); assert_eq!(output, None) }) } diff --git a/extension/src/datum_utils.rs b/extension/src/datum_utils.rs index f445b4c9..ea0f328e 100644 --- a/extension/src/datum_utils.rs +++ b/extension/src/datum_utils.rs @@ -62,7 +62,7 @@ pub struct TextSerializableDatumWriter { impl TextSerializableDatumWriter { pub fn from_oid(typoid: Oid) -> Self { - let mut type_output = 0; + let mut type_output = pg_sys::Oid::INVALID; let mut typ_is_varlena = false; let mut flinfo = unsafe { std::mem::MaybeUninit::zeroed().assume_init() }; @@ -81,13 +81,13 @@ impl TextSerializableDatumWriter { pub struct DatumFromSerializedTextReader { flinfo: pg_sys::FmgrInfo, - typ_io_param: u32, + typ_io_param: pg_sys::Oid, } impl DatumFromSerializedTextReader { pub fn from_oid(typoid: Oid) -> Self { - let mut type_input = 0; - let mut typ_io_param = 0; + let mut type_input = pg_sys::Oid::INVALID; + let mut typ_io_param = pg_sys::oids::Oid::INVALID; let mut flinfo = unsafe { std::mem::MaybeUninit::zeroed().assume_init() }; unsafe { pg_sys::getTypeInputInfo(typoid, &mut type_input, &mut typ_io_param); @@ -101,7 +101,7 @@ impl DatumFromSerializedTextReader { } pub fn read_datum(&mut self, datum_str: &str) -> Datum { - let cstr = pgx::cstr_core::CString::new(datum_str).unwrap(); // TODO: error handling + let cstr = std::ffi::CString::new(datum_str).unwrap(); // TODO: error handling let cstr_ptr = cstr.as_ptr() as *mut std::os::raw::c_char; unsafe { pg_sys::InputFunctionCall(&mut self.flinfo, cstr_ptr, self.typ_io_param, -1) } } @@ -115,7 +115,7 @@ impl Serialize for TextSerializeableDatum { S: serde::Serializer, { let chars = unsafe { pg_sys::OutputFunctionCall(self.1, self.0) }; - let cstr = unsafe { pgx::cstr_core::CStr::from_ptr(chars) }; + let cstr = unsafe { std::ffi::CStr::from_ptr(chars) }; serializer.serialize_str(cstr.to_str().unwrap()) } } @@ -237,7 +237,7 @@ impl Serialize for DatumHashBuilder { where S: serde::Serializer, { - let collation = if self.collation == 0 { + let collation = if self.collation == pg_sys::oids::Oid::INVALID { None } else { Some(PgCollationId(self.collation)) @@ -296,7 +296,7 @@ impl<'a> Serialize for DatumStore<'a> { let mut writer = TextSerializableDatumWriter::from_oid(self.type_oid.0); let count = self.iter().count(); let mut seq = serializer.serialize_seq(Some(count + 1))?; - seq.serialize_element(&self.type_oid.0)?; + seq.serialize_element(&self.type_oid.0.as_u32())?; for element in self.iter() { seq.serialize_element(&writer.make_serializable(element))?; } @@ -322,7 +322,8 @@ impl<'a, 'de> Deserialize<'de> for DatumStore<'a> { where A: SeqAccess<'de>, { - let oid = seq.next_element::().unwrap().unwrap(); // TODO: error handling + let oid = + unsafe { Oid::from_u32_unchecked(seq.next_element::().unwrap().unwrap()) }; // TODO: error handling // TODO separate human-readable and binary forms let mut reader = DatumFromSerializedTextReader::from_oid(oid); @@ -685,10 +686,10 @@ mod tests { #[pg_test] fn test_value_datum_store() { - Spi::execute(|client| { - let test = client.select("SELECT toolkit_experimental.datum_test_agg(r.data)::TEXT FROM (SELECT generate_series(10, 100, 10) as data) r", None, None) - .first() - .get_one::().unwrap(); + Spi::connect(|mut client| { + let test = client.update("SELECT toolkit_experimental.datum_test_agg(r.data)::TEXT FROM (SELECT generate_series(10, 100, 10) as data) r", None, None) + .unwrap().first() + .get_one::().unwrap().unwrap(); let expected = "(version:1,datums:[23,\"10\",\"20\",\"30\",\"40\",\"50\",\"60\",\"70\",\"80\",\"90\",\"100\"])"; assert_eq!(test, expected); }); @@ -696,10 +697,10 @@ mod tests { #[pg_test] fn test_varlena_datum_store() { - Spi::execute(|client| { - let test = client.select("SELECT toolkit_experimental.datum_test_agg(r.data)::TEXT FROM (SELECT generate_series(10, 100, 10)::TEXT as data) r", None, None) - .first() - .get_one::().unwrap(); + Spi::connect(|mut client| { + let test = client.update("SELECT toolkit_experimental.datum_test_agg(r.data)::TEXT FROM (SELECT generate_series(10, 100, 10)::TEXT as data) r", None, None) + .unwrap().first() + .get_one::().unwrap().unwrap(); let expected = "(version:1,datums:[25,\"10\",\"20\",\"30\",\"40\",\"50\",\"60\",\"70\",\"80\",\"90\",\"100\"])"; assert_eq!(test, expected); }); @@ -707,10 +708,10 @@ mod tests { #[pg_test] fn test_byref_datum_store() { - Spi::execute(|client| { - let test = client.select("SELECT toolkit_experimental.datum_test_agg(r.data)::TEXT FROM (SELECT (generate_series(10, 100, 10)::TEXT || ' seconds')::INTERVAL as data) r", None, None) - .first() - .get_one::().unwrap(); + Spi::connect(|mut client| { + let test = client.update("SELECT toolkit_experimental.datum_test_agg(r.data)::TEXT FROM (SELECT (generate_series(10, 100, 10)::TEXT || ' seconds')::INTERVAL as data) r", None, None) + .unwrap().first() + .get_one::().unwrap().unwrap(); let expected = "(version:1,datums:[1186,\"00:00:10\",\"00:00:20\",\"00:00:30\",\"00:00:40\",\"00:00:50\",\"00:01:00\",\"00:01:10\",\"00:01:20\",\"00:01:30\",\"00:01:40\"])"; assert_eq!(test, expected); }); diff --git a/extension/src/duration.rs b/extension/src/duration.rs index cc252070..c2886cb9 100644 --- a/extension/src/duration.rs +++ b/extension/src/duration.rs @@ -2,6 +2,8 @@ //! PostgreSQL parses duration units. Currently units longer than an hour are unsupported since //! the length of days varies when in a timezone with daylight savings time. +use core::fmt::{self, Formatter}; + // Canonical PostgreSQL units: https://github.com/postgres/postgres/blob/b76fb6c2a99eb7d49f96e56599fef1ffc1c134c9/src/include/utils/datetime.h#L48-L60 #[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)] pub enum DurationUnit { @@ -14,7 +16,7 @@ pub enum DurationUnit { } impl DurationUnit { - fn microseconds(self) -> u32 { + pub fn microseconds(self) -> u32 { match self { Self::Microsec => 1, Self::Millisec => 1000, @@ -46,6 +48,18 @@ impl DurationUnit { } } +impl fmt::Display for DurationUnit { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + match self { + DurationUnit::Microsec => write!(f, "microsecond"), + DurationUnit::Millisec => write!(f, "millisecond"), + DurationUnit::Second => write!(f, "second"), + DurationUnit::Minute => write!(f, "minute"), + DurationUnit::Hour => write!(f, "hour"), + } + } +} + #[cfg(test)] mod test { use super::*; diff --git a/extension/src/frequency.rs b/extension/src/frequency.rs index 3d26a9df..89078eec 100644 --- a/extension/src/frequency.rs +++ b/extension/src/frequency.rs @@ -444,7 +444,7 @@ pub mod toolkit_experimental { build! { SpaceSavingAggregate { - type_oid: trans.type_oid() as _, + type_oid: trans.type_oid().into(), num_values: trans.entries.len() as _, values_seen: trans.total_vals, freq_param: trans.freq_param, @@ -466,14 +466,14 @@ pub mod toolkit_experimental { let mut trans = if agg.topn == 0 { SpaceSavingTransState::freq_agg_from_type_id( agg.freq_param, - agg.type_oid, + unsafe { Oid::from_u32_unchecked(agg.type_oid) }, collation, ) } else { SpaceSavingTransState::topn_agg_from_type_id( agg.freq_param, agg.topn as u32, - agg.type_oid, + unsafe { Oid::from_u32_unchecked(agg.type_oid) }, collation, ) }; @@ -851,7 +851,7 @@ pub fn rollup_agg_trans_inner( in_aggregate_context(fcinfo, || { let trans = (&value, &fcinfo).into(); if let Some(state) = state { - Some(SpaceSavingTransState::combine(&*state, &trans).into()) + Some(SpaceSavingTransState::combine(&state, &trans).into()) } else { Some(trans.into()) } @@ -881,7 +881,7 @@ pub fn rollup_agg_bigint_trans_inner( in_aggregate_context(fcinfo, || { let trans = (&value, &fcinfo).into(); if let Some(state) = state { - Some(SpaceSavingTransState::combine(&*state, &trans).into()) + Some(SpaceSavingTransState::combine(&state, &trans).into()) } else { Some(trans.into()) } @@ -911,7 +911,7 @@ pub fn rollup_agg_text_trans_inner( in_aggregate_context(fcinfo, || { let trans = (&value, &fcinfo).into(); if let Some(state) = state { - Some(SpaceSavingTransState::combine(&*state, &trans).into()) + Some(SpaceSavingTransState::combine(&state, &trans).into()) } else { Some(trans.into()) } @@ -1287,14 +1287,19 @@ pub fn freq_iter<'a>( ), > { unsafe { - if ty.oid() != agg.type_oid { + if ty.oid().as_u32() != agg.type_oid { pgx::error!("mischatched types") } let counts = agg.counts.slice().iter().zip(agg.overcounts.slice().iter()); TableIterator::new(agg.datums.clone().into_iter().zip(counts).map_while( move |(value, (&count, &overcount))| { let total = agg.values_seen as f64; - let value = AnyElement::from_polymorphic_datum(value, false, agg.type_oid).unwrap(); + let value = AnyElement::from_polymorphic_datum( + value, + false, + Oid::from_u32_unchecked(agg.type_oid), + ) + .unwrap(); let min_freq = (count - overcount) as f64 / total; let max_freq = count as f64 / total; Some((value, min_freq, max_freq)) @@ -1394,7 +1399,7 @@ pub fn topn( ty: Option, ) -> SetOfIterator { // If called with a NULL, assume type matches - if ty.is_some() && ty.unwrap().oid() != agg.type_oid { + if ty.is_some() && ty.unwrap().oid().as_u32() != agg.type_oid { pgx::error!("mischatched types") } @@ -1418,7 +1423,7 @@ pub fn topn( ) // TODO Shouldn't failure to convert to AnyElement cause error, not early stop? .map_while(move |value| unsafe { - AnyElement::from_polymorphic_datum(value, false, type_oid) + AnyElement::from_polymorphic_datum(value, false, Oid::from_u32_unchecked(type_oid)) }), ) } @@ -1527,7 +1532,7 @@ pub fn max_frequency(agg: SpaceSavingAggregate<'_>, value: AnyElement) -> f64 { match agg .datums .iter() - .position(|datum| value == (datum, agg.type_oid).into()) + .position(|datum| value == (datum, unsafe { Oid::from_u32_unchecked(agg.type_oid) }).into()) { Some(idx) => agg.counts.slice()[idx] as f64 / agg.values_seen as f64, None => 0., @@ -1540,7 +1545,7 @@ pub fn min_frequency(agg: SpaceSavingAggregate<'_>, value: AnyElement) -> f64 { match agg .datums .iter() - .position(|datum| value == (datum, agg.type_oid).into()) + .position(|datum| value == (datum, unsafe { Oid::from_u32_unchecked(agg.type_oid) }).into()) { Some(idx) => { (agg.counts.slice()[idx] - agg.overcounts.slice()[idx]) as f64 / agg.values_seen as f64 @@ -1681,39 +1686,45 @@ mod tests { #[pg_test] fn test_freq_aggregate() { - Spi::execute(|client| { + Spi::connect(|mut client| { // using the search path trick for this test to make it easier to stabilize later on let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); - client.select("SET TIMEZONE to UTC", None, None); - client.select( - "CREATE TABLE test (data INTEGER, time TIMESTAMPTZ)", - None, - None, - ); + client.update("SET TIMEZONE to UTC", None, None).unwrap(); + client + .update( + "CREATE TABLE test (data INTEGER, time TIMESTAMPTZ)", + None, + None, + ) + .unwrap(); for i in (0..100).rev() { - client.select(&format!("INSERT INTO test SELECT i, '2020-1-1'::TIMESTAMPTZ + ('{} days, ' || i::TEXT || ' seconds')::INTERVAL FROM generate_series({}, 99, 1) i", 100 - i, i), None, None); + client.update(&format!("INSERT INTO test SELECT i, '2020-1-1'::TIMESTAMPTZ + ('{} days, ' || i::TEXT || ' seconds')::INTERVAL FROM generate_series({}, 99, 1) i", 100 - i, i), None, None).unwrap(); } - let test = client.select("SELECT freq_agg(0.015, s.data)::TEXT FROM (SELECT data FROM test ORDER BY time) s", None, None) - .first() - .get_one::().unwrap(); + let test = client.update("SELECT freq_agg(0.015, s.data)::TEXT FROM (SELECT data FROM test ORDER BY time) s", None, None) + .unwrap().first() + .get_one::().unwrap().unwrap(); let expected = "(version:1,num_values:67,topn:0,values_seen:5050,freq_param:0.015,counts:[100,99,98,97,96,95,94,93,92,91,90,89,88,87,86,85,84,83,82,81,80,79,78,77,76,75,74,73,72,71,70,69,68,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67],overcounts:[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66],datums:[99,98,97,96,95,94,93,92,91,90,89,88,87,86,85,84,83,82,81,80,79,78,77,76,75,74,73,72,71,70,69,68,67,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66])"; assert_eq!(test, expected); - let test = client.select("SELECT raw_freq_agg(0.015, s.data)::TEXT FROM (SELECT data FROM test ORDER BY time) s", None, None) - .first() - .get_one::().unwrap(); + let test = client.update("SELECT raw_freq_agg(0.015, s.data)::TEXT FROM (SELECT data FROM test ORDER BY time) s", None, None) + .unwrap().first() + .get_one::().unwrap().unwrap(); let expected = "(version:1,type_oid:23,num_values:67,values_seen:5050,freq_param:0.015,topn:0,counts:[100,99,98,97,96,95,94,93,92,91,90,89,88,87,86,85,84,83,82,81,80,79,78,77,76,75,74,73,72,71,70,69,68,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67,67],overcounts:[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66],datums:[23,\"99\",\"98\",\"97\",\"96\",\"95\",\"94\",\"93\",\"92\",\"91\",\"90\",\"89\",\"88\",\"87\",\"86\",\"85\",\"84\",\"83\",\"82\",\"81\",\"80\",\"79\",\"78\",\"77\",\"76\",\"75\",\"74\",\"73\",\"72\",\"71\",\"70\",\"69\",\"68\",\"67\",\"33\",\"34\",\"35\",\"36\",\"37\",\"38\",\"39\",\"40\",\"41\",\"42\",\"43\",\"44\",\"45\",\"46\",\"47\",\"48\",\"49\",\"50\",\"51\",\"52\",\"53\",\"54\",\"55\",\"56\",\"57\",\"58\",\"59\",\"60\",\"61\",\"62\",\"63\",\"64\",\"65\",\"66\"])"; assert_eq!(test, expected); }); @@ -1721,33 +1732,39 @@ mod tests { #[pg_test] fn test_topn_aggregate() { - Spi::execute(|client| { + Spi::connect(|mut client| { // using the search path trick for this test to make it easier to stabilize later on let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); - client.select("SET TIMEZONE to UTC", None, None); - client.select( - "CREATE TABLE test (data INTEGER, time TIMESTAMPTZ)", - None, - None, - ); + client.update("SET TIMEZONE to UTC", None, None).unwrap(); + client + .update( + "CREATE TABLE test (data INTEGER, time TIMESTAMPTZ)", + None, + None, + ) + .unwrap(); for i in (0..200).rev() { - client.select(&format!("INSERT INTO test SELECT i, '2020-1-1'::TIMESTAMPTZ + ('{} days, ' || i::TEXT || ' seconds')::INTERVAL FROM generate_series({}, 199, 1) i", 200 - i, i), None, None); + client.update(&format!("INSERT INTO test SELECT i, '2020-1-1'::TIMESTAMPTZ + ('{} days, ' || i::TEXT || ' seconds')::INTERVAL FROM generate_series({}, 199, 1) i", 200 - i, i), None, None).unwrap(); } - let test = client.select("SELECT topn_agg(10, s.data)::TEXT FROM (SELECT data FROM test ORDER BY time) s", None, None) - .first() - .get_one::().unwrap(); + let test = client.update("SELECT topn_agg(10, s.data)::TEXT FROM (SELECT data FROM test ORDER BY time) s", None, None) + .unwrap().first() + .get_one::().unwrap().unwrap(); let expected = "(version:1,num_values:110,topn:10,values_seen:20100,freq_param:1.1,counts:[200,199,198,197,196,195,194,193,192,191,190,189,188,187,186,185,184,183,182,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181,181],overcounts:[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180,180],datums:[199,198,197,196,195,194,193,192,191,190,189,188,187,186,185,184,183,182,181,90,91,92,93,94,95,96,97,98,99,100,101,102,103,104,105,106,107,108,109,110,111,112,113,114,115,116,117,118,119,120,121,122,123,124,125,126,127,128,129,130,131,132,133,134,135,136,137,138,139,140,141,142,143,144,145,146,147,148,149,150,151,152,153,154,155,156,157,158,159,160,161,162,163,164,165,166,167,168,169,170,171,172,173,174,175,176,177,178,179,180])"; assert_eq!(test, expected); }); @@ -1967,84 +1984,96 @@ mod tests { } // Setup environment and create table 'test' with some aggregates in table 'aggs' - fn setup_with_test_table(client: &SpiClient) { + fn setup_with_test_table(client: &mut pgx::spi::SpiClient) { // using the search path trick for this test to make it easier to stabilize later on let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); - client.select("SET TIMEZONE to UTC", None, None); - client.select( - "CREATE TABLE test (data INTEGER, time TIMESTAMPTZ)", - None, - None, - ); + client.update("SET TIMEZONE to UTC", None, None).unwrap(); + client + .update( + "CREATE TABLE test (data INTEGER, time TIMESTAMPTZ)", + None, + None, + ) + .unwrap(); for i in (0..20).rev() { - client.select(&format!("INSERT INTO test SELECT i, '2020-1-1'::TIMESTAMPTZ + ('{} days, ' || i::TEXT || ' seconds')::INTERVAL FROM generate_series({}, 19, 1) i", 10 - i, i), None, None); + client.update(&format!("INSERT INTO test SELECT i, '2020-1-1'::TIMESTAMPTZ + ('{} days, ' || i::TEXT || ' seconds')::INTERVAL FROM generate_series({}, 19, 1) i", 10 - i, i), None, None).unwrap(); } - client.select( - "CREATE TABLE aggs (name TEXT, agg SPACESAVINGBIGINTAGGREGATE)", - None, - None, - ); - client.select("INSERT INTO aggs SELECT 'topn_default', topn_agg(5, s.data) FROM (SELECT data FROM test ORDER BY time) s", None, None); - client.select("INSERT INTO aggs SELECT 'topn_1.5', topn_agg(5, 1.5, s.data) FROM (SELECT data FROM test ORDER BY time) s", None, None); - client.select("INSERT INTO aggs SELECT 'topn_2', topn_agg(5, 2, s.data) FROM (SELECT data FROM test ORDER BY time) s", None, None); - client.select("INSERT INTO aggs SELECT 'freq_8', freq_agg(0.08, s.data) FROM (SELECT data FROM test ORDER BY time) s", None, None); - client.select("INSERT INTO aggs SELECT 'freq_5', freq_agg(0.05, s.data) FROM (SELECT data FROM test ORDER BY time) s", None, None); - client.select("INSERT INTO aggs SELECT 'freq_2', freq_agg(0.02, s.data) FROM (SELECT data FROM test ORDER BY time) s", None, None); + client + .update( + "CREATE TABLE aggs (name TEXT, agg SPACESAVINGBIGINTAGGREGATE)", + None, + None, + ) + .unwrap(); + client.update("INSERT INTO aggs SELECT 'topn_default', topn_agg(5, s.data) FROM (SELECT data FROM test ORDER BY time) s", None, None).unwrap(); + client.update("INSERT INTO aggs SELECT 'topn_1.5', topn_agg(5, 1.5, s.data) FROM (SELECT data FROM test ORDER BY time) s", None, None).unwrap(); + client.update("INSERT INTO aggs SELECT 'topn_2', topn_agg(5, 2, s.data) FROM (SELECT data FROM test ORDER BY time) s", None, None).unwrap(); + client.update("INSERT INTO aggs SELECT 'freq_8', freq_agg(0.08, s.data) FROM (SELECT data FROM test ORDER BY time) s", None, None).unwrap(); + client.update("INSERT INTO aggs SELECT 'freq_5', freq_agg(0.05, s.data) FROM (SELECT data FROM test ORDER BY time) s", None, None).unwrap(); + client.update("INSERT INTO aggs SELECT 'freq_2', freq_agg(0.02, s.data) FROM (SELECT data FROM test ORDER BY time) s", None, None).unwrap(); } // API tests #[pg_test] fn test_topn() { - Spi::execute(|client| { - setup_with_test_table(&client); + Spi::connect(|mut client| { + setup_with_test_table(&mut client); // simple tests let rows = client - .select( + .update( "SELECT topn(agg) FROM aggs WHERE name = 'topn_default'", None, None, ) + .unwrap() .count(); assert_eq!(rows, 5); let rows = client - .select( + .update( "SELECT topn(agg, 5) FROM aggs WHERE name = 'freq_5'", None, None, ) + .unwrap() .count(); assert_eq!(rows, 5); // can limit below topn_agg value let rows = client - .select( + .update( "SELECT topn(agg, 3) FROM aggs WHERE name = 'topn_default'", None, None, ) + .unwrap() .count(); assert_eq!(rows, 3); // only 4 rows with freq >= 0.08 let rows = client - .select( + .update( "SELECT topn(agg, 5) FROM aggs WHERE name = 'freq_8'", None, None, ) + .unwrap() .count(); assert_eq!(rows, 4); }); @@ -2054,44 +2083,47 @@ mod tests { error = "data is not skewed enough to find top 0 parameters with a skew of 1.5, try reducing the skew factor" )] fn topn_on_underskewed_topn_agg() { - Spi::execute(|client| { - setup_with_test_table(&client); + Spi::connect(|mut client| { + setup_with_test_table(&mut client); client - .select( + .update( "SELECT topn(agg, 0::int) FROM aggs WHERE name = 'topn_1.5'", None, None, ) + .unwrap() .count(); }); } #[pg_test(error = "requested N (8) exceeds creation parameter of topn aggregate (5)")] fn topn_high_n_on_topn_agg() { - Spi::execute(|client| { - setup_with_test_table(&client); + Spi::connect(|mut client| { + setup_with_test_table(&mut client); client - .select( + .update( "SELECT topn(agg, 8) FROM aggs WHERE name = 'topn_default'", None, None, ) + .unwrap() .count(); }); } #[pg_test(error = "frequency aggregates require a N parameter to topn")] fn topn_requires_n_for_freq_agg() { - Spi::execute(|client| { - setup_with_test_table(&client); + Spi::connect(|mut client| { + setup_with_test_table(&mut client); assert_eq!( 0, client - .select( + .update( "SELECT topn(agg) FROM aggs WHERE name = 'freq_2'", None, None ) + .unwrap() .count(), ); }); @@ -2099,31 +2131,34 @@ mod tests { #[pg_test] fn test_into_values() { - Spi::execute(|client| { - setup_with_test_table(&client); + Spi::connect(|mut client| { + setup_with_test_table(&mut client); let rows = client - .select( + .update( "SELECT into_values(agg) FROM aggs WHERE name = 'freq_8'", None, None, ) + .unwrap() .count(); assert_eq!(rows, 13); let rows = client - .select( + .update( "SELECT into_values(agg) FROM aggs WHERE name = 'freq_5'", None, None, ) + .unwrap() .count(); assert_eq!(rows, 20); let rows = client - .select( + .update( "SELECT into_values(agg) FROM aggs WHERE name = 'freq_2'", None, None, ) + .unwrap() .count(); assert_eq!(rows, 20); }); @@ -2131,66 +2166,65 @@ mod tests { #[pg_test] fn test_frequency_getters() { - Spi::execute(|client| { - setup_with_test_table(&client); + Spi::connect(|mut client| { + setup_with_test_table(&mut client); // simple tests - let (min, max) = client.select("SELECT min_frequency(agg, 3), max_frequency(agg, 3) FROM aggs WHERE name = 'freq_2'", None, None) - .first() - .get_two::(); + let (min, max) = client.update("SELECT min_frequency(agg, 3), max_frequency(agg, 3) FROM aggs WHERE name = 'freq_2'", None, None) + .unwrap().first() + .get_two::().unwrap(); assert_eq!(min.unwrap(), 0.01904761904761905); assert_eq!(max.unwrap(), 0.01904761904761905); - let (min, max) = client.select("SELECT min_frequency(agg, 11), max_frequency(agg, 11) FROM aggs WHERE name = 'topn_default'", None, None) - .first() - .get_two::(); + let (min, max) = client.update("SELECT min_frequency(agg, 11), max_frequency(agg, 11) FROM aggs WHERE name = 'topn_default'", None, None) + .unwrap().first() + .get_two::().unwrap(); assert_eq!(min.unwrap(), 0.05714285714285714); assert_eq!(max.unwrap(), 0.05714285714285714); // missing value - let (min, max) = client.select("SELECT min_frequency(agg, 3), max_frequency(agg, 3) FROM aggs WHERE name = 'freq_8'", None, None) - .first() - .get_two::(); + let (min, max) = client.update("SELECT min_frequency(agg, 3), max_frequency(agg, 3) FROM aggs WHERE name = 'freq_8'", None, None) + .unwrap().first() + .get_two::().unwrap(); assert_eq!(min.unwrap(), 0.); assert_eq!(max.unwrap(), 0.); - let (min, max) = client.select("SELECT min_frequency(agg, 20), max_frequency(agg, 20) FROM aggs WHERE name = 'topn_2'", None, None) - .first() - .get_two::(); + let (min, max) = client.update("SELECT min_frequency(agg, 20), max_frequency(agg, 20) FROM aggs WHERE name = 'topn_2'", None, None) + .unwrap().first() + .get_two::().unwrap(); assert_eq!(min.unwrap(), 0.); assert_eq!(max.unwrap(), 0.); // noisy value - let (min, max) = client.select("SELECT min_frequency(agg, 8), max_frequency(agg, 8) FROM aggs WHERE name = 'topn_1.5'", None, None) - .first() - .get_two::(); + let (min, max) = client.update("SELECT min_frequency(agg, 8), max_frequency(agg, 8) FROM aggs WHERE name = 'topn_1.5'", None, None) + .unwrap().first() + .get_two::().unwrap(); assert_eq!(min.unwrap(), 0.004761904761904762); assert_eq!(max.unwrap(), 0.05238095238095238); }); } #[pg_test] - #[ignore] fn test_rollups() { - Spi::execute(|client| { - client.select( + Spi::connect(|mut client| { + client.update( "CREATE TABLE test (raw_data DOUBLE PRECISION, int_data INTEGER, text_data TEXT, bucket INTEGER)", None, None, - ); + ).unwrap(); - // Generate an array of 1000 values by taking the probability curve for a + // Generate an array of 10000 values by taking the probability curve for a // zeta curve with an s of 1.1 for the top 5 values, then adding smaller // amounts of the next 5 most common values, and finally filling with unique values. - let mut vals = vec![1; 95]; - vals.append(&mut vec![2; 45]); - vals.append(&mut vec![3; 39]); - vals.append(&mut vec![4; 21]); - vals.append(&mut vec![5; 17]); + let mut vals = vec![1; 945]; + vals.append(&mut vec![2; 441]); + vals.append(&mut vec![3; 283]); + vals.append(&mut vec![4; 206]); + vals.append(&mut vec![5; 161]); for v in 6..=10 { - vals.append(&mut vec![v, 10]); + vals.append(&mut vec![v, 125]); } - for v in 0..(1000 - 95 - 45 - 39 - 21 - 17 - (5 * 10)) { + for v in 0..(10000 - 945 - 441 - 283 - 206 - 161 - (5 * 125)) { vals.push(11 + v); } vals.shuffle(&mut thread_rng()); @@ -2201,45 +2235,45 @@ mod tests { "INSERT INTO test SELECT {}, {}::INT, {}::TEXT, FLOOR(RANDOM() * 10)", v, v, v ); - client.select(&cmd, None, None); + client.update(&cmd, None, None).unwrap(); } // No matter how the values are batched into subaggregates, we should always // see the same top 5 values - let mut result = client.select( + let mut result = client.update( "WITH aggs AS (SELECT bucket, toolkit_experimental.raw_topn_agg(5, raw_data) as raw_agg FROM test GROUP BY bucket) SELECT toolkit_experimental.topn(toolkit_experimental.rollup(raw_agg), NULL::DOUBLE PRECISION)::TEXT from aggs", None, None - ); - assert_eq!(result.next().unwrap()[1].value(), Some("1")); - assert_eq!(result.next().unwrap()[1].value(), Some("2")); - assert_eq!(result.next().unwrap()[1].value(), Some("3")); - assert_eq!(result.next().unwrap()[1].value(), Some("4")); - assert_eq!(result.next().unwrap()[1].value(), Some("5")); + ).unwrap(); + assert_eq!(result.next().unwrap()[1].value().unwrap(), Some("1")); + assert_eq!(result.next().unwrap()[1].value().unwrap(), Some("2")); + assert_eq!(result.next().unwrap()[1].value().unwrap(), Some("3")); + assert_eq!(result.next().unwrap()[1].value().unwrap(), Some("4")); + assert_eq!(result.next().unwrap()[1].value().unwrap(), Some("5")); assert!(result.next().is_none()); - let mut result = client.select( + let mut result = client.update( "WITH aggs AS (SELECT bucket, toolkit_experimental.topn_agg(5, int_data) as int_agg FROM test GROUP BY bucket) SELECT toolkit_experimental.topn(toolkit_experimental.rollup(int_agg))::TEXT from aggs", None, None - ); - assert_eq!(result.next().unwrap()[1].value(), Some("1")); - assert_eq!(result.next().unwrap()[1].value(), Some("2")); - assert_eq!(result.next().unwrap()[1].value(), Some("3")); - assert_eq!(result.next().unwrap()[1].value(), Some("4")); - assert_eq!(result.next().unwrap()[1].value(), Some("5")); + ).unwrap(); + assert_eq!(result.next().unwrap()[1].value().unwrap(), Some("1")); + assert_eq!(result.next().unwrap()[1].value().unwrap(), Some("2")); + assert_eq!(result.next().unwrap()[1].value().unwrap(), Some("3")); + assert_eq!(result.next().unwrap()[1].value().unwrap(), Some("4")); + assert_eq!(result.next().unwrap()[1].value().unwrap(), Some("5")); assert!(result.next().is_none()); - let mut result = client.select( + let mut result = client.update( "WITH aggs AS (SELECT bucket, toolkit_experimental.topn_agg(5, text_data) as text_agg FROM test GROUP BY bucket) SELECT toolkit_experimental.topn(toolkit_experimental.rollup(text_agg))::TEXT from aggs", None, None - ); - assert_eq!(result.next().unwrap()[1].value(), Some("1")); - assert_eq!(result.next().unwrap()[1].value(), Some("2")); - assert_eq!(result.next().unwrap()[1].value(), Some("3")); - assert_eq!(result.next().unwrap()[1].value(), Some("4")); - assert_eq!(result.next().unwrap()[1].value(), Some("5")); + ).unwrap(); + assert_eq!(result.next().unwrap()[1].value().unwrap(), Some("1")); + assert_eq!(result.next().unwrap()[1].value().unwrap(), Some("2")); + assert_eq!(result.next().unwrap()[1].value().unwrap(), Some("3")); + assert_eq!(result.next().unwrap()[1].value().unwrap(), Some("4")); + assert_eq!(result.next().unwrap()[1].value().unwrap(), Some("5")); assert!(result.next().is_none()); }); } diff --git a/extension/src/gauge_agg.rs b/extension/src/gauge_agg.rs index 59351001..7e499414 100644 --- a/extension/src/gauge_agg.rs +++ b/extension/src/gauge_agg.rs @@ -698,22 +698,26 @@ mod tests { macro_rules! select_one { ($client:expr, $stmt:expr, $type:ty) => { $client - .select($stmt, None, None) + .update($stmt, None, None) + .unwrap() .first() .get_one::<$type>() .unwrap() + .unwrap() }; } #[pg_test] fn round_trip() { - Spi::execute(|client| { - client.select( - "CREATE TABLE test(ts timestamptz, val DOUBLE PRECISION)", - None, - None, - ); - client.select("SET TIME ZONE 'UTC'", None, None); + Spi::connect(|mut client| { + client + .update( + "CREATE TABLE test(ts timestamptz, val DOUBLE PRECISION)", + None, + None, + ) + .unwrap(); + client.update("SET TIME ZONE 'UTC'", None, None).unwrap(); let stmt = "INSERT INTO test VALUES\ ('2020-01-01 00:00:00+00', 10.0),\ ('2020-01-01 00:01:00+00', 20.0),\ @@ -724,7 +728,7 @@ mod tests { ('2020-01-01 00:06:00+00', 10.0),\ ('2020-01-01 00:07:00+00', 30.0),\ ('2020-01-01 00:08:00+00', 10.0)"; - client.select(stmt, None, None); + client.update(stmt, None, None).unwrap(); let expected = "(\ version:1,\ @@ -784,8 +788,8 @@ mod tests { #[pg_test] fn delta_after_gauge_decrease() { - Spi::execute(|client| { - decrease(&client); + Spi::connect(|mut client| { + decrease(&mut client); let stmt = "SELECT toolkit_experimental.delta(toolkit_experimental.gauge_agg(ts, val)) FROM test"; assert_eq!(-20.0, select_one!(client, stmt, f64)); }); @@ -793,8 +797,8 @@ mod tests { #[pg_test] fn delta_after_gauge_increase() { - Spi::execute(|client| { - increase(&client); + Spi::connect(|mut client| { + increase(&mut client); let stmt = "SELECT toolkit_experimental.delta(toolkit_experimental.gauge_agg(ts, val)) FROM test"; assert_eq!(20.0, select_one!(client, stmt, f64)); }); @@ -802,8 +806,8 @@ mod tests { #[pg_test] fn delta_after_gauge_decrease_then_increase_to_same_value() { - Spi::execute(|client| { - decrease_then_increase_to_same_value(&client); + Spi::connect(|mut client| { + decrease_then_increase_to_same_value(&mut client); let stmt = "SELECT toolkit_experimental.delta(toolkit_experimental.gauge_agg(ts, val)) FROM test"; assert_eq!(0.0, select_one!(client, stmt, f64)); }); @@ -811,8 +815,8 @@ mod tests { #[pg_test] fn delta_after_gauge_increase_then_decrease_to_same_value() { - Spi::execute(|client| { - increase_then_decrease_to_same_value(&client); + Spi::connect(|mut client| { + increase_then_decrease_to_same_value(&mut client); let stmt = "SELECT toolkit_experimental.delta(toolkit_experimental.gauge_agg(ts, val)) FROM test"; assert_eq!(0.0, select_one!(client, stmt, f64)); }); @@ -820,8 +824,8 @@ mod tests { #[pg_test] fn idelta_left_after_gauge_decrease() { - Spi::execute(|client| { - decrease(&client); + Spi::connect(|mut client| { + decrease(&mut client); let stmt = "SELECT toolkit_experimental.idelta_left(toolkit_experimental.gauge_agg(ts, val)) FROM test"; assert_eq!(10.0, select_one!(client, stmt, f64)); }); @@ -829,8 +833,8 @@ mod tests { #[pg_test] fn idelta_left_after_gauge_increase() { - Spi::execute(|client| { - increase(&client); + Spi::connect(|mut client| { + increase(&mut client); let stmt = "SELECT toolkit_experimental.idelta_left(toolkit_experimental.gauge_agg(ts, val)) FROM test"; assert_eq!(20.0, select_one!(client, stmt, f64)); }); @@ -838,8 +842,8 @@ mod tests { #[pg_test] fn idelta_left_after_gauge_increase_then_decrease_to_same_value() { - Spi::execute(|client| { - increase_then_decrease_to_same_value(&client); + Spi::connect(|mut client| { + increase_then_decrease_to_same_value(&mut client); let stmt = "SELECT toolkit_experimental.idelta_left(toolkit_experimental.gauge_agg(ts, val)) FROM test"; assert_eq!(20.0, select_one!(client, stmt, f64)); }); @@ -847,8 +851,8 @@ mod tests { #[pg_test] fn idelta_left_after_gauge_decrease_then_increase_to_same_value() { - Spi::execute(|client| { - decrease_then_increase_to_same_value(&client); + Spi::connect(|mut client| { + decrease_then_increase_to_same_value(&mut client); let stmt = "SELECT toolkit_experimental.idelta_left(toolkit_experimental.gauge_agg(ts, val)) FROM test"; assert_eq!(10.0, select_one!(client, stmt, f64)); }); @@ -856,8 +860,8 @@ mod tests { #[pg_test] fn idelta_right_after_gauge_decrease() { - Spi::execute(|client| { - decrease(&client); + Spi::connect(|mut client| { + decrease(&mut client); let stmt = "SELECT toolkit_experimental.idelta_right(toolkit_experimental.gauge_agg(ts, val)) FROM test"; assert_eq!(10.0, select_one!(client, stmt, f64)); }); @@ -865,8 +869,8 @@ mod tests { #[pg_test] fn idelta_right_after_gauge_increase() { - Spi::execute(|client| { - increase(&client); + Spi::connect(|mut client| { + increase(&mut client); let stmt = "SELECT toolkit_experimental.idelta_right(toolkit_experimental.gauge_agg(ts, val)) FROM test"; assert_eq!(20.0, select_one!(client, stmt, f64)); }); @@ -874,8 +878,8 @@ mod tests { #[pg_test] fn idelta_right_after_gauge_increase_then_decrease_to_same_value() { - Spi::execute(|client| { - increase_then_decrease_to_same_value(&client); + Spi::connect(|mut client| { + increase_then_decrease_to_same_value(&mut client); let stmt = "SELECT toolkit_experimental.idelta_right(toolkit_experimental.gauge_agg(ts, val)) FROM test"; assert_eq!(10.0, select_one!(client, stmt, f64)); }); @@ -883,8 +887,8 @@ mod tests { #[pg_test] fn idelta_right_after_gauge_decrease_then_increase_to_same_value() { - Spi::execute(|client| { - decrease_then_increase_to_same_value(&client); + Spi::connect(|mut client| { + decrease_then_increase_to_same_value(&mut client); let stmt = "SELECT toolkit_experimental.idelta_right(toolkit_experimental.gauge_agg(ts, val)) FROM test"; assert_eq!(20.0, select_one!(client, stmt, f64)); }); @@ -910,28 +914,46 @@ mod tests { #[pg_test] fn rollup() { - Spi::execute(|client| { - client.select( - "CREATE TABLE test(ts timestamptz, val DOUBLE PRECISION)", - None, - None, - ); + Spi::connect(|mut client| { + // needed so that GaugeSummary type can be resolved + let sp = client + .update( + "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", + None, + None, + ) + .unwrap() + .first() + .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) + .unwrap(); + + client + .update( + "CREATE TABLE test(ts timestamptz, val DOUBLE PRECISION)", + None, + None, + ) + .unwrap(); // This tests GaugeSummary::single_value - the old first == last // check erroneously saw 21.0 == 21.0 and called it a single value. let stmt = "INSERT INTO test VALUES('2020-01-01 00:00:00+00', 10.0), ('2020-01-01 00:01:00+00', 21.0), ('2020-01-01 00:01:00+00', 22.0), ('2020-01-01 00:01:00+00', 21.0)"; - client.select(stmt, None, None); + client.update(stmt, None, None).unwrap(); let stmt = "INSERT INTO test VALUES('2020-01-01 00:02:00+00', 10.0), ('2020-01-01 00:03:00+00', 20.0), ('2020-01-01 00:04:00+00', 10.0)"; - client.select(stmt, None, None); + client.update(stmt, None, None).unwrap(); let stmt = "INSERT INTO test VALUES('2020-01-01 00:08:00+00', 30.0), ('2020-01-01 00:10:00+00', 30.0), ('2020-01-01 00:10:30+00', 10.0), ('2020-01-01 00:20:00+00', 40.0)"; - client.select(stmt, None, None); + client.update(stmt, None, None).unwrap(); //combine function works as expected - let stmt = "SELECT toolkit_experimental.gauge_agg(ts, val) FROM test"; + let stmt = "SELECT gauge_agg(ts, val) FROM test"; let a = select_one!(client, stmt, GaugeSummary); - let stmt = "WITH t as (SELECT date_trunc('minute', ts), toolkit_experimental.gauge_agg(ts, val) as agg FROM test group by 1 ) SELECT toolkit_experimental.rollup(agg) FROM t"; + let stmt = "WITH t as (SELECT date_trunc('minute', ts), gauge_agg(ts, val) as agg FROM test group by 1 ) SELECT rollup(agg) FROM t"; let b = select_one!(client, stmt, GaugeSummary); assert_close_enough(&a.into(), &b.into()); }); @@ -939,14 +961,15 @@ mod tests { #[pg_test] fn gauge_agg_interpolation() { - Spi::execute(|client| { - client.select( + Spi::connect(|mut client| { + client.update( "CREATE TABLE test(time timestamptz, value double precision, bucket timestamptz)", None, None, - ); - client.select( - r#"INSERT INTO test VALUES + ).unwrap(); + client + .update( + r#"INSERT INTO test VALUES ('2020-1-1 10:00'::timestamptz, 10.0, '2020-1-1'::timestamptz), ('2020-1-1 12:00'::timestamptz, 40.0, '2020-1-1'::timestamptz), ('2020-1-1 16:00'::timestamptz, 20.0, '2020-1-1'::timestamptz), @@ -956,12 +979,14 @@ mod tests { ('2020-1-3 4:00'::timestamptz, 30.0, '2020-1-3'::timestamptz), ('2020-1-3 12:00'::timestamptz, 0.0, '2020-1-3'::timestamptz), ('2020-1-3 16:00'::timestamptz, 35.0, '2020-1-3'::timestamptz)"#, - None, - None, - ); + None, + None, + ) + .unwrap(); - let mut deltas = client.select( - r#"SELECT + let mut deltas = client + .update( + r#"SELECT toolkit_experimental.interpolated_delta( agg, bucket, @@ -974,19 +999,21 @@ mod tests { GROUP BY bucket ) s ORDER BY bucket"#, - None, - None, - ); + None, + None, + ) + .unwrap(); // Day 1, start at 10, interpolated end of day is 16 - assert_eq!(deltas.next().unwrap()[1].value(), Some(16. - 10.)); + assert_eq!(deltas.next().unwrap()[1].value().unwrap(), Some(16. - 10.)); // Day 2, interpolated start is 16, interpolated end is 27.5 - assert_eq!(deltas.next().unwrap()[1].value(), Some(27.5 - 16.)); + assert_eq!(deltas.next().unwrap()[1].value().unwrap(), Some(27.5 - 16.)); // Day 3, interpolated start is 27.5, end is 35 - assert_eq!(deltas.next().unwrap()[1].value(), Some(35. - 27.5)); + assert_eq!(deltas.next().unwrap()[1].value().unwrap(), Some(35. - 27.5)); - let mut rates = client.select( - r#"SELECT + let mut rates = client + .update( + r#"SELECT toolkit_experimental.interpolated_rate( agg, bucket, @@ -999,23 +1026,24 @@ mod tests { GROUP BY bucket ) s ORDER BY bucket"#, - None, - None, - ); + None, + None, + ) + .unwrap(); // Day 1, 14 hours (rate is per second) assert_eq!( - rates.next().unwrap()[1].value(), + rates.next().unwrap()[1].value().unwrap(), Some((16. - 10.) / (14. * 60. * 60.)) ); // Day 2, 24 hours assert_eq!( - rates.next().unwrap()[1].value(), + rates.next().unwrap()[1].value().unwrap(), Some((27.5 - 16.) / (24. * 60. * 60.)) ); // Day 3, 16 hours assert_eq!( - rates.next().unwrap()[1].value(), + rates.next().unwrap()[1].value().unwrap(), Some((35. - 27.5) / (16. * 60. * 60.)) ); }); @@ -1023,26 +1051,29 @@ mod tests { #[pg_test] fn guage_agg_interpolated_delta_with_aligned_point() { - Spi::execute(|client| { - client.select( + Spi::connect(|mut client| { + client.update( "CREATE TABLE test(time timestamptz, value double precision, bucket timestamptz)", None, None, - ); - client.select( - r#"INSERT INTO test VALUES + ).unwrap(); + client + .update( + r#"INSERT INTO test VALUES ('2020-1-1 10:00'::timestamptz, 10.0, '2020-1-1'::timestamptz), ('2020-1-1 12:00'::timestamptz, 40.0, '2020-1-1'::timestamptz), ('2020-1-1 16:00'::timestamptz, 20.0, '2020-1-1'::timestamptz), ('2020-1-2 0:00'::timestamptz, 15.0, '2020-1-2'::timestamptz), ('2020-1-2 12:00'::timestamptz, 50.0, '2020-1-2'::timestamptz), ('2020-1-2 20:00'::timestamptz, 25.0, '2020-1-2'::timestamptz)"#, - None, - None, - ); + None, + None, + ) + .unwrap(); - let mut deltas = client.select( - r#"SELECT + let mut deltas = client + .update( + r#"SELECT toolkit_experimental.interpolated_delta( agg, bucket, @@ -1055,34 +1086,55 @@ mod tests { GROUP BY bucket ) s ORDER BY bucket"#, - None, - None, - ); + None, + None, + ) + .unwrap(); // Day 1, start at 10, interpolated end of day is 15 (after reset) - assert_eq!(deltas.next().unwrap()[1].value(), Some(15. - 10.)); + assert_eq!(deltas.next().unwrap()[1].value().unwrap(), Some(15. - 10.)); // Day 2, start is 15, end is 25 - assert_eq!(deltas.next().unwrap()[1].value(), Some(25. - 15.)); + assert_eq!(deltas.next().unwrap()[1].value().unwrap(), Some(25. - 15.)); assert!(deltas.next().is_none()); }); } #[pg_test] fn no_results_on_null_input() { - Spi::execute(|client| { - client.select( - "CREATE TABLE test(ts timestamptz, val DOUBLE PRECISION)", - None, - None, - ); + Spi::connect(|mut client| { + // needed so that GaugeSummary type can be resolved + let sp = client + .update( + "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", + None, + None, + ) + .unwrap() + .first() + .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) + .unwrap(); + + client + .update( + "CREATE TABLE test(ts timestamptz, val DOUBLE PRECISION)", + None, + None, + ) + .unwrap(); let stmt = "INSERT INTO test VALUES (NULL, NULL)"; - client.select(stmt, None, None); + client.update(stmt, None, None).unwrap(); let stmt = "SELECT toolkit_experimental.gauge_agg(ts, val) FROM test"; assert!(client - .select(stmt, None, None) + .update(stmt, None, None) + .unwrap() .first() .get_one::() + .unwrap() .is_none()); }); } diff --git a/extension/src/heartbeat_agg.rs b/extension/src/heartbeat_agg.rs index 8d46fcd5..1ef63413 100644 --- a/extension/src/heartbeat_agg.rs +++ b/extension/src/heartbeat_agg.rs @@ -2,6 +2,9 @@ use pgx::iter::TableIterator; use pgx::*; use crate::{ + accessors::{ + AccessorDeadRanges, AccessorDowntime, AccessorLiveAt, AccessorLiveRanges, AccessorUptime, + }, aggregate_utils::in_aggregate_context, datum_utils::interval_to_ms, flatten, @@ -13,7 +16,12 @@ use crate::{ use std::cmp::{max, min}; -use toolkit_experimental::HeartbeatAggData; +mod accessors; + +use accessors::{ + HeartbeatInterpolateAccessor, HeartbeatInterpolatedDowntimeAccessor, + HeartbeatInterpolatedUptimeAccessor, +}; const BUFFER_SIZE: usize = 1000; // How many values to absorb before consolidating @@ -180,195 +188,263 @@ impl HeartbeatTransState { } } -#[pg_schema] -mod toolkit_experimental { - use super::*; - - pg_type! { - #[derive(Debug)] - struct HeartbeatAgg<'input> - { - start_time : i64, - end_time : i64, - last_seen : i64, - interval_len : i64, - num_intervals : u64, - interval_starts : [i64; self.num_intervals], - interval_ends : [i64; self.num_intervals], - } +pg_type! { + #[derive(Debug)] + struct HeartbeatAgg<'input> + { + start_time : i64, + end_time : i64, + last_seen : i64, + interval_len : i64, + num_intervals : u64, + interval_starts : [i64; self.num_intervals], + interval_ends : [i64; self.num_intervals], } +} - ron_inout_funcs!(HeartbeatAgg); +ron_inout_funcs!(HeartbeatAgg); - impl HeartbeatAgg<'_> { - fn sum_live_intervals(self) -> i64 { - let starts = self.interval_starts.as_slice(); - let ends = self.interval_ends.as_slice(); - let mut sum = 0; - for i in 0..self.num_intervals as usize { - sum += ends[i] - starts[i]; - } - sum +impl HeartbeatAgg<'_> { + fn sum_live_intervals(self) -> i64 { + let starts = self.interval_starts.as_slice(); + let ends = self.interval_ends.as_slice(); + let mut sum = 0; + for i in 0..self.num_intervals as usize { + sum += ends[i] - starts[i]; } + sum + } - fn interpolate_start(&mut self, pred: &Self) { - // only allow interpolation of non-overlapping ranges - assert!(pred.end_time <= self.start_time); - let pred_end = pred.last_seen + self.interval_len; - - if pred_end <= self.start_time { - return; - } + fn interpolate_start(&mut self, pred: &Self) { + // only allow interpolation of non-overlapping ranges + assert!(pred.end_time <= self.start_time); + let pred_end = pred.last_seen + self.interval_len; - // If first range already covers (start_time, pred_end) return - if self - .interval_starts - .as_slice() - .first() - .filter(|v| **v == self.start_time) - .is_some() - && self - .interval_ends - .as_slice() - .first() - .filter(|v| **v >= pred_end) - .is_some() - { - return; - } + if pred_end <= self.start_time { + return; + } - if self - .interval_starts + // If first range already covers (start_time, pred_end) return + if self + .interval_starts + .as_slice() + .first() + .filter(|v| **v == self.start_time) + .is_some() + && self + .interval_ends .as_slice() .first() - .filter(|v| **v <= pred_end) + .filter(|v| **v >= pred_end) .is_some() - { - self.interval_starts.as_owned()[0] = self.start_time; - } else { - let start = self.start_time; - self.interval_starts.as_owned().insert(0, start); - self.interval_ends.as_owned().insert(0, pred_end); - self.num_intervals += 1; - } + { + return; } - } - #[pg_extern] - pub fn live_ranges( - agg: HeartbeatAgg<'static>, - ) -> TableIterator<'static, (name!(start, TimestampTz), name!(end, TimestampTz))> { - let starts = agg.interval_starts.clone(); - let ends = agg.interval_ends.clone(); - TableIterator::new( - starts - .into_iter() - .map(|x| x.into()) - .zip(ends.into_iter().map(|x| x.into())), - ) + if self + .interval_starts + .as_slice() + .first() + .filter(|v| **v <= pred_end) + .is_some() + { + self.interval_starts.as_owned()[0] = self.start_time; + } else { + let start = self.start_time; + self.interval_starts.as_owned().insert(0, start); + self.interval_ends.as_owned().insert(0, pred_end); + self.num_intervals += 1; + } } +} - #[pg_extern] - pub fn dead_ranges( - agg: HeartbeatAgg<'static>, - ) -> TableIterator<'static, (name!(start, TimestampTz), name!(end, TimestampTz))> { - if agg.num_intervals == 0 { - return TableIterator::new(std::iter::once(( - agg.start_time.into(), - agg.end_time.into(), - ))); - } +#[pg_extern] +pub fn live_ranges( + agg: HeartbeatAgg<'static>, +) -> TableIterator<'static, (name!(start, TimestampTz), name!(end, TimestampTz))> { + let starts = agg.interval_starts.clone(); + let ends = agg.interval_ends.clone(); + TableIterator::new( + starts + .into_iter() + .map(|x| x.into()) + .zip(ends.into_iter().map(|x| x.into())), + ) +} - // Dead ranges are the opposite of the intervals stored in the aggregate - let mut starts = agg.interval_ends.clone().into_vec(); - let mut ends = agg.interval_starts.clone().into_vec(); +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_heartbeat_agg_live_ranges( + sketch: HeartbeatAgg<'static>, + _accessor: AccessorLiveRanges<'static>, +) -> TableIterator<'static, (name!(start, TimestampTz), name!(end, TimestampTz))> { + live_ranges(sketch) +} - // Fix the first point depending on whether the aggregate starts in a live or dead range - if ends[0] == agg.start_time { - ends.remove(0); - } else { - starts.insert(0, agg.start_time); - } +#[pg_extern] +pub fn dead_ranges( + agg: HeartbeatAgg<'static>, +) -> TableIterator<'static, (name!(start, TimestampTz), name!(end, TimestampTz))> { + if agg.num_intervals == 0 { + return TableIterator::new(std::iter::once(( + agg.start_time.into(), + agg.end_time.into(), + ))); + } - // Fix the last point depending on whether the aggregate starts in a live or dead range - if *starts.last().unwrap() == agg.end_time { - starts.pop(); - } else { - ends.push(agg.end_time); - } + // Dead ranges are the opposite of the intervals stored in the aggregate + let mut starts = agg.interval_ends.clone().into_vec(); + let mut ends = agg.interval_starts.clone().into_vec(); - TableIterator::new( - starts - .into_iter() - .map(|x| x.into()) - .zip(ends.into_iter().map(|x| x.into())), - ) + // Fix the first point depending on whether the aggregate starts in a live or dead range + if ends[0] == agg.start_time { + ends.remove(0); + } else { + starts.insert(0, agg.start_time); } - #[pg_extern] - pub fn uptime(agg: HeartbeatAgg<'static>) -> Interval { - agg.sum_live_intervals().into() + // Fix the last point depending on whether the aggregate starts in a live or dead range + if *starts.last().unwrap() == agg.end_time { + starts.pop(); + } else { + ends.push(agg.end_time); } - #[pg_extern] - pub fn interpolated_uptime( - agg: HeartbeatAgg<'static>, - pred: Option>, - ) -> Interval { - uptime(interpolate_heartbeat_agg(agg, pred)) - } + TableIterator::new( + starts + .into_iter() + .map(|x| x.into()) + .zip(ends.into_iter().map(|x| x.into())), + ) +} - #[pg_extern] - pub fn downtime(agg: HeartbeatAgg<'static>) -> Interval { - (agg.end_time - agg.start_time - agg.sum_live_intervals()).into() - } +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_heartbeat_agg_dead_ranges( + sketch: HeartbeatAgg<'static>, + _accessor: AccessorDeadRanges<'static>, +) -> TableIterator<'static, (name!(start, TimestampTz), name!(end, TimestampTz))> { + dead_ranges(sketch) +} + +#[pg_extern] +pub fn uptime(agg: HeartbeatAgg<'static>) -> Interval { + agg.sum_live_intervals().into() +} + +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_heartbeat_agg_uptime( + sketch: HeartbeatAgg<'static>, + _accessor: AccessorUptime<'static>, +) -> Interval { + uptime(sketch) +} + +#[pg_extern] +pub fn interpolated_uptime( + agg: HeartbeatAgg<'static>, + pred: Option>, +) -> Interval { + uptime(interpolate_heartbeat_agg(agg, pred)) +} + +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_heartbeat_agg_interpolated_uptime( + sketch: HeartbeatAgg<'static>, + accessor: HeartbeatInterpolatedUptimeAccessor<'static>, +) -> Interval { + interpolated_uptime(sketch, accessor.pred()) +} + +#[pg_extern] +pub fn downtime(agg: HeartbeatAgg<'static>) -> Interval { + (agg.end_time - agg.start_time - agg.sum_live_intervals()).into() +} + +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_heartbeat_agg_downtime( + sketch: HeartbeatAgg<'static>, + _accessor: AccessorDowntime<'static>, +) -> Interval { + downtime(sketch) +} + +#[pg_extern] +pub fn interpolated_downtime( + agg: HeartbeatAgg<'static>, + pred: Option>, +) -> Interval { + downtime(interpolate_heartbeat_agg(agg, pred)) +} + +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_heartbeat_agg_interpolated_downtime( + sketch: HeartbeatAgg<'static>, + accessor: HeartbeatInterpolatedDowntimeAccessor<'static>, +) -> Interval { + interpolated_downtime(sketch, accessor.pred()) +} - #[pg_extern] - pub fn interpolated_downtime( - agg: HeartbeatAgg<'static>, - pred: Option>, - ) -> Interval { - downtime(interpolate_heartbeat_agg(agg, pred)) +#[pg_extern] +pub fn live_at(agg: HeartbeatAgg<'static>, test: TimestampTz) -> bool { + if agg.num_intervals == 0 { + return false; } - #[pg_extern] - pub fn live_at(agg: HeartbeatAgg<'static>, test: TimestampTz) -> bool { - if agg.num_intervals == 0 { + let test = i64::from(test); + let mut start_iter = agg.interval_starts.iter().enumerate().peekable(); + while let Some((idx, val)) = start_iter.next() { + if test < val { + // Only possible if test shows up before first interval return false; } - - let test = i64::from(test); - let mut start_iter = agg.interval_starts.iter().enumerate().peekable(); - while let Some((idx, val)) = start_iter.next() { - if test < val { - // Only possible if test shows up before first interval - return false; - } - if let Some((_, next_val)) = start_iter.peek() { - if test < *next_val { - return test < agg.interval_ends.as_slice()[idx]; - } + if let Some((_, next_val)) = start_iter.peek() { + if test < *next_val { + return test < agg.interval_ends.as_slice()[idx]; } } - // Fall out the loop if test > start of last interval - return test < *agg.interval_ends.as_slice().last().unwrap(); } + // Fall out the loop if test > start of last interval + return test < *agg.interval_ends.as_slice().last().unwrap(); +} - #[pg_extern(name = "interpolate")] - fn interpolate_heartbeat_agg( - agg: HeartbeatAgg<'static>, - pred: Option>, - ) -> HeartbeatAgg<'static> { - let mut r = agg.clone(); - if let Some(pred) = pred { - r.interpolate_start(&pred); - } - r +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_heartbeat_agg_live_at( + sketch: HeartbeatAgg<'static>, + accessor: AccessorLiveAt<'static>, +) -> bool { + let ts = TimestampTz(accessor.time.into()); + live_at(sketch, ts) +} + +#[pg_extern(name = "interpolate")] +fn interpolate_heartbeat_agg( + agg: HeartbeatAgg<'static>, + pred: Option>, +) -> HeartbeatAgg<'static> { + let mut r = agg.clone(); + if let Some(pred) = pred { + r.interpolate_start(&pred); } + r } -impl From> for HeartbeatTransState { - fn from(agg: toolkit_experimental::HeartbeatAgg<'static>) -> Self { +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_heartbeat_agg_interpolate( + sketch: HeartbeatAgg<'static>, + accessor: HeartbeatInterpolateAccessor<'static>, +) -> HeartbeatAgg<'static> { + interpolate_heartbeat_agg(sketch, accessor.pred()) +} + +impl From> for HeartbeatTransState { + fn from(agg: HeartbeatAgg<'static>) -> Self { HeartbeatTransState { start: agg.start_time, end: agg.end_time, @@ -384,7 +460,7 @@ impl From> for HeartbeatTransState { } } -#[pg_extern(schema = "toolkit_experimental", immutable, parallel_safe)] +#[pg_extern(immutable, parallel_safe)] pub fn heartbeat_trans( state: Internal, heartbeat: TimestampTz, @@ -425,17 +501,17 @@ pub fn heartbeat_trans_inner( } } -#[pg_extern(schema = "toolkit_experimental", immutable, parallel_safe)] +#[pg_extern(immutable, parallel_safe)] pub fn heartbeat_final( state: Internal, fcinfo: pg_sys::FunctionCallInfo, -) -> Option> { +) -> Option> { heartbeat_final_inner(unsafe { state.to_inner() }, fcinfo) } pub fn heartbeat_final_inner( state: Option>, fcinfo: pg_sys::FunctionCallInfo, -) -> Option> { +) -> Option> { unsafe { in_aggregate_context(fcinfo, || { state.map(|mut s| { @@ -464,17 +540,17 @@ pub fn heartbeat_final_inner( } } -#[pg_extern(schema = "toolkit_experimental", immutable, parallel_safe)] +#[pg_extern(immutable, parallel_safe)] pub fn heartbeat_rollup_trans( state: Internal, - value: Option>, + value: Option>, fcinfo: pg_sys::FunctionCallInfo, ) -> Option { heartbeat_rollup_trans_inner(unsafe { state.to_inner() }, value, fcinfo).internal() } pub fn heartbeat_rollup_trans_inner( state: Option>, - value: Option>, + value: Option>, fcinfo: pg_sys::FunctionCallInfo, ) -> Option> { unsafe { @@ -491,12 +567,12 @@ pub fn heartbeat_rollup_trans_inner( extension_sql!( "\n\ - CREATE AGGREGATE toolkit_experimental.heartbeat_agg(\n\ + CREATE AGGREGATE heartbeat_agg(\n\ heartbeat TIMESTAMPTZ, agg_start TIMESTAMPTZ, agg_duration INTERVAL, heartbeat_liveness INTERVAL\n\ ) (\n\ - sfunc = toolkit_experimental.heartbeat_trans,\n\ + sfunc = heartbeat_trans,\n\ stype = internal,\n\ - finalfunc = toolkit_experimental.heartbeat_final\n\ + finalfunc = heartbeat_final\n\ );\n\ ", name = "heartbeat_agg", @@ -508,12 +584,12 @@ extension_sql!( extension_sql!( "\n\ - CREATE AGGREGATE toolkit_experimental.rollup(\n\ - toolkit_experimental.HeartbeatAgg\n\ + CREATE AGGREGATE rollup(\n\ + HeartbeatAgg\n\ ) (\n\ - sfunc = toolkit_experimental.heartbeat_rollup_trans,\n\ + sfunc = heartbeat_rollup_trans,\n\ stype = internal,\n\ - finalfunc = toolkit_experimental.heartbeat_final\n\ + finalfunc = heartbeat_final\n\ );\n\ ", name = "heartbeat_agg_rollup", @@ -582,13 +658,16 @@ mod tests { #[pg_test] pub fn test_heartbeat_agg() { - Spi::execute(|client| { - client.select("SET TIMEZONE to UTC", None, None); + Spi::connect(|mut client| { + client.update("SET TIMEZONE to UTC", None, None).unwrap(); - client.select("CREATE TABLE liveness(heartbeat TIMESTAMPTZ)", None, None); + client + .update("CREATE TABLE liveness(heartbeat TIMESTAMPTZ)", None, None) + .unwrap(); - client.select( - "INSERT INTO liveness VALUES + client + .update( + "INSERT INTO liveness VALUES ('01-01-2020 0:2:20 UTC'), ('01-01-2020 0:10 UTC'), ('01-01-2020 0:17 UTC'), @@ -610,80 +689,230 @@ mod tests { ('01-01-2020 1:57 UTC'), ('01-01-2020 1:59:50 UTC') ", - None, - None, - ); + None, + None, + ) + .unwrap(); - let mut result = client.select( - "SELECT toolkit_experimental.live_ranges(toolkit_experimental.heartbeat_agg(heartbeat, '01-01-2020 UTC', '2h', '10m'))::TEXT - FROM liveness", None, None); + let mut result = client.update( + "SELECT live_ranges(heartbeat_agg(heartbeat, '01-01-2020 UTC', '2h', '10m'))::TEXT + FROM liveness", None, None).unwrap(); + let mut arrow_result = client.update( + "SELECT (heartbeat_agg(heartbeat, '01-01-2020 UTC', '2h', '10m') -> live_ranges())::TEXT + FROM liveness", None, None).unwrap(); + + let test = arrow_result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + test + ); + assert_eq!( + test, "(\"2020-01-01 00:02:20+00\",\"2020-01-01 00:27:00+00\")" ); + let test = arrow_result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(); + assert_eq!( + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + test + ); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + test, "(\"2020-01-01 00:30:00+00\",\"2020-01-01 00:50:00+00\")" ); + let test = arrow_result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(); + assert_eq!( + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + test + ); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + test, "(\"2020-01-01 00:50:30+00\",\"2020-01-01 01:38:00+00\")" ); + let test = arrow_result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(); + assert_eq!( + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + test + ); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + test, "(\"2020-01-01 01:38:01+00\",\"2020-01-01 02:00:00+00\")" ); assert!(result.next().is_none()); + assert!(arrow_result.next().is_none()); + + let mut result = client.update( + "SELECT dead_ranges(heartbeat_agg(heartbeat, '01-01-2020 UTC', '2h', '10m'))::TEXT + FROM liveness", None, None).unwrap(); - let mut result = client.select( - "SELECT toolkit_experimental.dead_ranges(toolkit_experimental.heartbeat_agg(heartbeat, '01-01-2020 UTC', '2h', '10m'))::TEXT - FROM liveness", None, None); + let mut arrow_result = client.update( + "SELECT (heartbeat_agg(heartbeat, '01-01-2020 UTC', '2h', '10m') -> dead_ranges())::TEXT + FROM liveness", None, None).unwrap(); + let test = arrow_result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(); + assert_eq!( + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + test + ); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + test, "(\"2020-01-01 00:00:00+00\",\"2020-01-01 00:02:20+00\")" ); + let test = arrow_result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(); + assert_eq!( + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + test + ); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + test, "(\"2020-01-01 00:27:00+00\",\"2020-01-01 00:30:00+00\")" ); + let test = arrow_result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + test + ); + assert_eq!( + test, "(\"2020-01-01 00:50:00+00\",\"2020-01-01 00:50:30+00\")" ); + let test = arrow_result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + test + ); + assert_eq!( + test, "(\"2020-01-01 01:38:00+00\",\"2020-01-01 01:38:01+00\")" ); assert!(result.next().is_none()); + assert!(arrow_result.next().is_none()); - let result = client.select( - "SELECT toolkit_experimental.uptime(toolkit_experimental.heartbeat_agg(heartbeat, '01-01-2020 UTC', '2h', '10m'))::TEXT - FROM liveness", None, None).first().get_one::().unwrap(); + let result = client + .update( + "SELECT uptime(heartbeat_agg(heartbeat, '01-01-2020 UTC', '2h', '10m'))::TEXT + FROM liveness", + None, + None, + ) + .unwrap() + .first() + .get_one::() + .unwrap() + .unwrap(); + assert_eq!("01:54:09", result); + + let result = client.update( + "SELECT (heartbeat_agg(heartbeat, '01-01-2020 UTC', '2h', '10m') -> uptime())::TEXT + FROM liveness", None, None).unwrap().first().get_one::().unwrap().unwrap(); assert_eq!("01:54:09", result); - let result = client.select( - "SELECT toolkit_experimental.downtime(toolkit_experimental.heartbeat_agg(heartbeat, '01-01-2020 UTC', '2h', '10m'))::TEXT - FROM liveness", None, None).first().get_one::().unwrap(); + let result = client + .update( + "SELECT downtime(heartbeat_agg(heartbeat, '01-01-2020 UTC', '2h', '10m'))::TEXT + FROM liveness", + None, + None, + ) + .unwrap() + .first() + .get_one::() + .unwrap() + .unwrap(); + assert_eq!("00:05:51", result); + + let result = client.update( + "SELECT (heartbeat_agg(heartbeat, '01-01-2020 UTC', '2h', '10m') -> downtime())::TEXT + FROM liveness", None, None).unwrap().first().get_one::().unwrap().unwrap(); assert_eq!("00:05:51", result); let (result1, result2, result3) = - client.select( - "WITH agg AS (SELECT toolkit_experimental.heartbeat_agg(heartbeat, '01-01-2020 UTC', '2h', '10m') AS agg FROM liveness) - SELECT toolkit_experimental.live_at(agg, '01-01-2020 00:01:00 UTC')::TEXT, - toolkit_experimental.live_at(agg, '01-01-2020 00:05:00 UTC')::TEXT, - toolkit_experimental.live_at(agg, '01-01-2020 00:30:00 UTC')::TEXT FROM agg", None, None) - .first() - .get_three::(); + client.update( + "WITH agg AS (SELECT heartbeat_agg(heartbeat, '01-01-2020 UTC', '2h', '10m') AS agg FROM liveness) + SELECT live_at(agg, '01-01-2020 00:01:00 UTC')::TEXT, + live_at(agg, '01-01-2020 00:05:00 UTC')::TEXT, + live_at(agg, '01-01-2020 00:30:00 UTC')::TEXT FROM agg", None, None) + .unwrap().first() + .get_three::().unwrap(); let (result4, result5) = - client.select( - "WITH agg AS (SELECT toolkit_experimental.heartbeat_agg(heartbeat, '01-01-2020 UTC', '2h', '10m') AS agg FROM liveness) - SELECT toolkit_experimental.live_at(agg, '01-01-2020 01:38:00 UTC')::TEXT, - toolkit_experimental.live_at(agg, '01-01-2020 02:01:00 UTC')::TEXT FROM agg", None, None) - .first() - .get_two::(); + client.update( + "WITH agg AS (SELECT heartbeat_agg(heartbeat, '01-01-2020 UTC', '2h', '10m') AS agg FROM liveness) + SELECT live_at(agg, '01-01-2020 01:38:00 UTC')::TEXT, + live_at(agg, '01-01-2020 02:01:00 UTC')::TEXT FROM agg", None, None) + .unwrap().first() + .get_two::().unwrap(); + + assert_eq!(result1.unwrap(), "false"); // outside ranges + assert_eq!(result2.unwrap(), "true"); // inside ranges + assert_eq!(result3.unwrap(), "true"); // first point of range + assert_eq!(result4.unwrap(), "false"); // last point of range + assert_eq!(result5.unwrap(), "false"); // outside aggregate + + let (result1, result2, result3) = + client.update( + "WITH agg AS (SELECT heartbeat_agg(heartbeat, '01-01-2020 UTC', '2h', '10m') AS agg FROM liveness) + SELECT (agg -> live_at('01-01-2020 00:01:00 UTC'))::TEXT, + (agg -> live_at('01-01-2020 00:05:00 UTC'))::TEXT, + (agg -> live_at('01-01-2020 00:30:00 UTC'))::TEXT FROM agg", None, None) + .unwrap().first() + .get_three::().unwrap(); + + let (result4, result5) = + client.update( + "WITH agg AS (SELECT heartbeat_agg(heartbeat, '01-01-2020 UTC', '2h', '10m') AS agg FROM liveness) + SELECT (agg -> live_at('01-01-2020 01:38:00 UTC'))::TEXT, + (agg -> live_at('01-01-2020 02:01:00 UTC'))::TEXT FROM agg", None, None) + .unwrap().first() + .get_two::().unwrap(); assert_eq!(result1.unwrap(), "false"); // outside ranges assert_eq!(result2.unwrap(), "true"); // inside ranges @@ -695,16 +924,18 @@ mod tests { #[pg_test] pub fn test_heartbeat_rollup() { - Spi::execute(|client| { - client.select("SET TIMEZONE to UTC", None, None); + Spi::connect(|mut client| { + client.update("SET TIMEZONE to UTC", None, None).unwrap(); - client.select( - "CREATE TABLE heartbeats(time timestamptz, batch timestamptz)", - None, - None, - ); + client + .update( + "CREATE TABLE heartbeats(time timestamptz, batch timestamptz)", + None, + None, + ) + .unwrap(); - client.select( + client.update( "INSERT INTO heartbeats VALUES ('01-01-2020 3:02:20 UTC'::timestamptz, '01-01-2020 3:00:00 UTC'::timestamptz), ('01-01-2020 3:03:10 UTC'::timestamptz, '01-01-2020 3:00:00 UTC'::timestamptz), @@ -720,20 +951,22 @@ mod tests { ('01-01-2020 23:39:00 UTC'::timestamptz, '01-01-2020 23:00:00 UTC'::timestamptz)", None, None, - ); + ).unwrap(); let result = client - .select( + .update( "WITH aggs AS ( - SELECT toolkit_experimental.heartbeat_agg(time, batch, '1h', '1m') + SELECT heartbeat_agg(time, batch, '1h', '1m') FROM heartbeats GROUP BY batch - ) SELECT toolkit_experimental.rollup(heartbeat_agg)::TEXT FROM aggs", + ) SELECT rollup(heartbeat_agg)::TEXT FROM aggs", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() .unwrap(); assert_eq!("(version:1,start_time:631162800000000,end_time:631238400000000,last_seen:631237140000000,interval_len:60000000,num_intervals:7,interval_starts:[631162940000000,631178360000000,631179560000000,631180760000000,631184350000000,631236860000000,631237040000000],interval_ends:[631163107000000,631178420000000,631179620000000,631180870000000,631184410000000,631236920000000,631237200000000])", result); }) @@ -741,17 +974,16 @@ mod tests { #[pg_test] pub fn test_heartbeat_combining_rollup() { - Spi::execute(|client| { - client.select("SET TIMEZONE to UTC", None, None); + Spi::connect(|mut client| { + client.update("SET TIMEZONE to UTC", None, None).unwrap(); - client.select( - "CREATE TABLE aggs(agg toolkit_experimental.heartbeatagg)", - None, - None, - ); + client + .update("CREATE TABLE aggs(agg heartbeatagg)", None, None) + .unwrap(); - client.select( - "INSERT INTO aggs SELECT toolkit_experimental.heartbeat_agg(hb, '01-01-2020 UTC', '1h', '10m') + client + .update( + "INSERT INTO aggs SELECT heartbeat_agg(hb, '01-01-2020 UTC', '1h', '10m') FROM (VALUES ('01-01-2020 0:2:20 UTC'::timestamptz), ('01-01-2020 0:10 UTC'::timestamptz), @@ -763,10 +995,12 @@ mod tests { ) AS _(hb)", None, None, - ); + ) + .unwrap(); - client.select( - "INSERT INTO aggs SELECT toolkit_experimental.heartbeat_agg(hb, '01-01-2020 0:30 UTC', '1h', '10m') + client + .update( + "INSERT INTO aggs SELECT heartbeat_agg(hb, '01-01-2020 0:30 UTC', '1h', '10m') FROM (VALUES ('01-01-2020 0:35 UTC'::timestamptz), ('01-01-2020 0:40 UTC'::timestamptz), @@ -776,10 +1010,12 @@ mod tests { ) AS _(hb)", None, None, - ); + ) + .unwrap(); - client.select( - "INSERT INTO aggs SELECT toolkit_experimental.heartbeat_agg(hb, '01-01-2020 1:00 UTC', '1h', '10m') + client + .update( + "INSERT INTO aggs SELECT heartbeat_agg(hb, '01-01-2020 1:00 UTC', '1h', '10m') FROM (VALUES ('01-01-2020 1:00 UTC'::timestamptz), ('01-01-2020 1:28 UTC'::timestamptz), @@ -792,29 +1028,44 @@ mod tests { ) AS _(hb)", None, None, - ); + ) + .unwrap(); - let mut result = client.select( - "SELECT toolkit_experimental.dead_ranges(toolkit_experimental.rollup(agg))::TEXT + let mut result = client + .update( + "SELECT dead_ranges(rollup(agg))::TEXT FROM aggs", - None, - None, - ); + None, + None, + ) + .unwrap(); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), "(\"2020-01-01 00:00:00+00\",\"2020-01-01 00:02:20+00\")" ); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), "(\"2020-01-01 00:27:00+00\",\"2020-01-01 00:30:00+00\")" ); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), "(\"2020-01-01 00:50:00+00\",\"2020-01-01 00:50:30+00\")" ); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), "(\"2020-01-01 01:38:00+00\",\"2020-01-01 01:38:01+00\")" ); assert!(result.next().is_none()); @@ -823,17 +1074,20 @@ mod tests { #[pg_test] pub fn test_heartbeat_agg_interpolation() { - Spi::execute(|client| { - client.select("SET TIMEZONE to UTC", None, None); + Spi::connect(|mut client| { + client.update("SET TIMEZONE to UTC", None, None).unwrap(); - client.select( - "CREATE TABLE liveness(heartbeat TIMESTAMPTZ, start TIMESTAMPTZ)", - None, - None, - ); + client + .update( + "CREATE TABLE liveness(heartbeat TIMESTAMPTZ, start TIMESTAMPTZ)", + None, + None, + ) + .unwrap(); - client.select( - "INSERT INTO liveness VALUES + client + .update( + "INSERT INTO liveness VALUES ('01-01-2020 0:2:20 UTC', '01-01-2020 0:0 UTC'), ('01-01-2020 0:10 UTC', '01-01-2020 0:0 UTC'), ('01-01-2020 0:17 UTC', '01-01-2020 0:0 UTC'), @@ -855,103 +1109,473 @@ mod tests { ('01-01-2020 1:57 UTC', '01-01-2020 1:30 UTC'), ('01-01-2020 1:59:50 UTC', '01-01-2020 1:30 UTC') ", - None, - None, - ); + None, + None, + ) + .unwrap(); - let mut result = client.select( - "WITH s AS ( + let mut result = client + .update( + "WITH s AS ( SELECT start, - toolkit_experimental.heartbeat_agg(heartbeat, start, '30m', '10m') AS agg + heartbeat_agg(heartbeat, start, '30m', '10m') AS agg FROM liveness GROUP BY start), t AS ( SELECT start, - toolkit_experimental.interpolate(agg, LAG (agg) OVER (ORDER BY start)) AS agg + interpolate(agg, LAG (agg) OVER (ORDER BY start)) AS agg FROM s) - SELECT toolkit_experimental.downtime(agg)::TEXT FROM t;", None, None); + SELECT downtime(agg)::TEXT FROM t;", + None, + None, + ) + .unwrap(); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), "00:05:20" ); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), "00:00:30" ); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), "00:00:00" ); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), "00:00:01" ); assert!(result.next().is_none()); - let mut result = client.select( - "WITH s AS ( + let mut result = client + .update( + "WITH s AS ( + SELECT start, + heartbeat_agg(heartbeat, start, '30m', '10m') AS agg + FROM liveness + GROUP BY start), + t AS ( + SELECT start, + interpolate(agg, LAG (agg) OVER (ORDER BY start)) AS agg + FROM s) + SELECT live_ranges(agg)::TEXT FROM t;", + None, + None, + ) + .unwrap(); + assert_eq!( + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + "(\"2020-01-01 00:02:20+00\",\"2020-01-01 00:27:00+00\")" + ); + assert_eq!( + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + "(\"2020-01-01 00:30:00+00\",\"2020-01-01 00:50:00+00\")" + ); + assert_eq!( + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + "(\"2020-01-01 00:50:30+00\",\"2020-01-01 01:00:00+00\")" + ); + assert_eq!( + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + "(\"2020-01-01 01:00:00+00\",\"2020-01-01 01:30:00+00\")" + ); + assert_eq!( + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + "(\"2020-01-01 01:30:00+00\",\"2020-01-01 01:38:00+00\")" + ); + assert_eq!( + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + "(\"2020-01-01 01:38:01+00\",\"2020-01-01 02:00:00+00\")" + ); + assert!(result.next().is_none()); + + let mut result = client + .update( + "WITH s AS ( SELECT start, - toolkit_experimental.heartbeat_agg(heartbeat, start, '30m', '10m') AS agg + heartbeat_agg(heartbeat, start, '30m', '10m') AS agg FROM liveness GROUP BY start), t AS ( SELECT start, - toolkit_experimental.interpolate(agg, LAG (agg) OVER (ORDER BY start)) AS agg + agg -> interpolate(LAG (agg) OVER (ORDER BY start)) AS agg FROM s) - SELECT toolkit_experimental.live_ranges(agg)::TEXT FROM t;", None, None); + SELECT live_ranges(agg)::TEXT FROM t;", + None, + None, + ) + .unwrap(); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), "(\"2020-01-01 00:02:20+00\",\"2020-01-01 00:27:00+00\")" ); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), "(\"2020-01-01 00:30:00+00\",\"2020-01-01 00:50:00+00\")" ); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), "(\"2020-01-01 00:50:30+00\",\"2020-01-01 01:00:00+00\")" ); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), "(\"2020-01-01 01:00:00+00\",\"2020-01-01 01:30:00+00\")" ); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), "(\"2020-01-01 01:30:00+00\",\"2020-01-01 01:38:00+00\")" ); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), "(\"2020-01-01 01:38:01+00\",\"2020-01-01 02:00:00+00\")" ); assert!(result.next().is_none()); - let mut result = client.select( - "WITH s AS ( + let mut result = client + .update( + "WITH s AS ( SELECT start, - toolkit_experimental.heartbeat_agg(heartbeat, start, '30m', '10m') AS agg + heartbeat_agg(heartbeat, start, '30m', '10m') AS agg FROM liveness GROUP BY start) - SELECT toolkit_experimental.interpolated_uptime(agg, LAG (agg) OVER (ORDER BY start))::TEXT - FROM s", None, None); + SELECT interpolated_uptime(agg, LAG (agg) OVER (ORDER BY start))::TEXT + FROM s", + None, + None, + ) + .unwrap(); + + assert_eq!( + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + "00:24:40" + ); + assert_eq!( + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + "00:29:30" + ); + assert_eq!( + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + "00:30:00" + ); + assert_eq!( + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + "00:29:59" + ); + assert!(result.next().is_none()); + + let mut result = client + .update( + "WITH s AS ( + SELECT start, + heartbeat_agg(heartbeat, start, '30m', '10m') AS agg + FROM liveness + GROUP BY start) + SELECT (agg -> interpolated_uptime(LAG (agg) OVER (ORDER BY start)))::TEXT + FROM s", + None, + None, + ) + .unwrap(); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), "00:24:40" ); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), "00:29:30" ); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), "00:30:00" ); assert_eq!( - result.next().unwrap()[1].value::().unwrap(), + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), "00:29:59" ); assert!(result.next().is_none()); + + let mut result = client + .update( + "WITH s AS ( + SELECT start, + heartbeat_agg(heartbeat, start, '30m', '10m') AS agg + FROM liveness + GROUP BY start) + SELECT interpolated_downtime(agg, LAG (agg) OVER (ORDER BY start))::TEXT + FROM s", + None, + None, + ) + .unwrap(); + + assert_eq!( + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + "00:05:20" + ); + assert_eq!( + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + "00:00:30" + ); + assert_eq!( + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + "00:00:00" + ); + assert_eq!( + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + "00:00:01" + ); + assert!(result.next().is_none()); + + let mut result = client + .update( + "WITH s AS ( + SELECT start, + heartbeat_agg(heartbeat, start, '30m', '10m') AS agg + FROM liveness + GROUP BY start) + SELECT (agg -> interpolated_downtime(LAG (agg) OVER (ORDER BY start)))::TEXT + FROM s", + None, + None, + ) + .unwrap(); + + assert_eq!( + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + "00:05:20" + ); + assert_eq!( + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + "00:00:30" + ); + assert_eq!( + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + "00:00:00" + ); + assert_eq!( + result.next().unwrap()[1] + .value::() + .unwrap() + .unwrap(), + "00:00:01" + ); + assert!(result.next().is_none()); }) } + + #[pg_test] + fn test_heartbeat_agg_text_io() { + Spi::connect(|mut client| { + client.update("SET TIMEZONE to UTC", None, None).unwrap(); + + client + .update("CREATE TABLE liveness(heartbeat TIMESTAMPTZ)", None, None) + .unwrap(); + + client + .update( + "INSERT INTO liveness VALUES + ('01-01-2020 0:2:20 UTC'), + ('01-01-2020 0:10 UTC'), + ('01-01-2020 0:17 UTC') + ", + None, + None, + ) + .unwrap(); + + let output = client + .update( + "SELECT heartbeat_agg(heartbeat, '01-01-2020', '30m', '5m')::TEXT + FROM liveness;", + None, + None, + ) + .unwrap() + .first() + .get_one::() + .unwrap(); + + let expected = "(version:1,start_time:631152000000000,end_time:631153800000000,last_seen:631153020000000,interval_len:300000000,num_intervals:3,interval_starts:[631152140000000,631152600000000,631153020000000],interval_ends:[631152440000000,631152900000000,631153320000000])"; + + assert_eq!(output, Some(expected.into())); + + let estimate = client + .update( + &format!("SELECT uptime('{}'::heartbeatagg)::TEXT", expected), + None, + None, + ) + .unwrap() + .first() + .get_one::() + .unwrap(); + assert_eq!(estimate.unwrap().as_str(), "00:15:00"); + }); + } + + #[pg_test] + fn test_heartbeat_agg_byte_io() { + use std::ptr; + + // Create a heartbeat agg from 0 to 250 with intervals from 40-50, 60-85, and 100-110 + let state = heartbeat_trans_inner( + None, + 40.into(), + 0.into(), + 250.into(), + 10.into(), + ptr::null_mut(), + ); + let state = heartbeat_trans_inner( + state, + 60.into(), + 0.into(), + 250.into(), + 10.into(), + ptr::null_mut(), + ); + let state = heartbeat_trans_inner( + state, + 65.into(), + 0.into(), + 250.into(), + 10.into(), + ptr::null_mut(), + ); + let state = heartbeat_trans_inner( + state, + 75.into(), + 0.into(), + 250.into(), + 10.into(), + ptr::null_mut(), + ); + let state = heartbeat_trans_inner( + state, + 100.into(), + 0.into(), + 250.into(), + 10.into(), + ptr::null_mut(), + ); + + let agg = heartbeat_final_inner(state, ptr::null_mut()) + .expect("failed to build finalized heartbeat_agg"); + let serial = agg.to_pg_bytes(); + + let expected = [ + 128, 1, 0, 0, // header + 1, // version + 0, 0, 0, // padding + 0, 0, 0, 0, 0, 0, 0, 0, // start_time + 250, 0, 0, 0, 0, 0, 0, 0, // end_time + 100, 0, 0, 0, 0, 0, 0, 0, // last_seen + 10, 0, 0, 0, 0, 0, 0, 0, // interval_len + 3, 0, 0, 0, 0, 0, 0, 0, // num_intervals + 40, 0, 0, 0, 0, 0, 0, 0, // interval_starts[0] + 60, 0, 0, 0, 0, 0, 0, 0, // interval_starts[1] + 100, 0, 0, 0, 0, 0, 0, 0, // interval_starts[2] + 50, 0, 0, 0, 0, 0, 0, 0, // interval_ends[0] + 85, 0, 0, 0, 0, 0, 0, 0, // interval_ends[1] + 110, 0, 0, 0, 0, 0, 0, 0, // interval_ends[2] + ]; + assert_eq!(serial, expected); + } } diff --git a/extension/src/heartbeat_agg/accessors.rs b/extension/src/heartbeat_agg/accessors.rs new file mode 100644 index 00000000..e60b7c95 --- /dev/null +++ b/extension/src/heartbeat_agg/accessors.rs @@ -0,0 +1,123 @@ +use pgx::*; + +use crate::{ + flatten, + heartbeat_agg::{HeartbeatAgg, HeartbeatAggData}, + pg_type, ron_inout_funcs, +}; + +fn empty_agg<'a>() -> HeartbeatAgg<'a> { + unsafe { + flatten!(HeartbeatAgg { + start_time: 0, + end_time: 0, + last_seen: 0, + interval_len: 0, + num_intervals: 0, + interval_starts: vec!().into(), + interval_ends: vec!().into(), + }) + } +} + +pg_type! { + struct HeartbeatInterpolatedUptimeAccessor<'input> { + has_prev : u64, + prev : HeartbeatAggData<'input>, + } +} + +ron_inout_funcs!(HeartbeatInterpolatedUptimeAccessor); + +#[pg_extern(immutable, parallel_safe, name = "interpolated_uptime")] +fn heartbeat_agg_interpolated_uptime_accessor<'a>( + prev: Option>, +) -> HeartbeatInterpolatedUptimeAccessor<'a> { + let has_prev = u64::from(prev.is_some()); + let prev = prev.unwrap_or_else(empty_agg).0; + + crate::build! { + HeartbeatInterpolatedUptimeAccessor { + has_prev, + prev, + } + } +} + +impl<'a> HeartbeatInterpolatedUptimeAccessor<'a> { + pub fn pred(&self) -> Option> { + if self.has_prev == 0 { + None + } else { + Some(self.prev.clone().into()) + } + } +} + +pg_type! { + struct HeartbeatInterpolatedDowntimeAccessor<'input> { + has_prev : u64, + prev : HeartbeatAggData<'input>, + } +} + +ron_inout_funcs!(HeartbeatInterpolatedDowntimeAccessor); + +#[pg_extern(immutable, parallel_safe, name = "interpolated_downtime")] +fn heartbeat_agg_interpolated_downtime_accessor<'a>( + prev: Option>, +) -> HeartbeatInterpolatedDowntimeAccessor<'a> { + let has_prev = u64::from(prev.is_some()); + let prev = prev.unwrap_or_else(empty_agg).0; + + crate::build! { + HeartbeatInterpolatedDowntimeAccessor { + has_prev, + prev, + } + } +} + +impl<'a> HeartbeatInterpolatedDowntimeAccessor<'a> { + pub fn pred(&self) -> Option> { + if self.has_prev == 0 { + None + } else { + Some(self.prev.clone().into()) + } + } +} + +pg_type! { + struct HeartbeatInterpolateAccessor<'input> { + has_prev : u64, + prev : HeartbeatAggData<'input>, + } +} + +ron_inout_funcs!(HeartbeatInterpolateAccessor); + +#[pg_extern(immutable, parallel_safe, name = "interpolate")] +fn heartbeat_agg_interpolate_accessor<'a>( + prev: Option>, +) -> HeartbeatInterpolateAccessor<'a> { + let has_prev = u64::from(prev.is_some()); + let prev = prev.unwrap_or_else(empty_agg).0; + + crate::build! { + HeartbeatInterpolateAccessor { + has_prev, + prev, + } + } +} + +impl<'a> HeartbeatInterpolateAccessor<'a> { + pub fn pred(&self) -> Option> { + if self.has_prev == 0 { + None + } else { + Some(self.prev.clone().into()) + } + } +} diff --git a/extension/src/hyperloglog.rs b/extension/src/hyperloglog.rs index 36898646..31d97f51 100644 --- a/extension/src/hyperloglog.rs +++ b/extension/src/hyperloglog.rs @@ -482,17 +482,19 @@ mod tests { #[pg_test] fn test_hll_aggregate() { - Spi::execute(|client| { + Spi::connect(|mut client| { let text = client - .select( + .update( "SELECT \ hyperloglog(32, v::float)::TEXT \ FROM generate_series(1, 100) v", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); let expected = "(\ version:1,\ @@ -509,7 +511,7 @@ mod tests { assert_eq!(text.unwrap(), expected); let (count, arrow_count) = client - .select( + .update( "SELECT \ distinct_count(\ hyperloglog(32, v::float)\ @@ -519,19 +521,23 @@ mod tests { None, None, ) + .unwrap() .first() - .get_two::(); + .get_two::() + .unwrap(); assert_eq!(count, Some(132)); assert_eq!(count, arrow_count); let count2 = client - .select( + .update( &format!("SELECT distinct_count('{}')", expected), None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(count2, count); }); } @@ -539,17 +545,19 @@ mod tests { #[pg_test] // Should have same results as test_hll_distinct_aggregate running with the same number of buckets fn test_approx_count_distinct_aggregate() { - Spi::execute(|client| { + Spi::connect(|mut client| { let text = client - .select( + .update( "SELECT \ toolkit_experimental.approx_count_distinct(v::float)::TEXT \ FROM generate_series(1, 100) v", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); let expected = "(\ version:1,\ @@ -584,7 +592,7 @@ mod tests { assert_eq!(text.unwrap(), expected); let (count, arrow_count) = client - .select( + .update( "SELECT \ distinct_count(\ toolkit_experimental.approx_count_distinct(v::float)\ @@ -594,19 +602,23 @@ mod tests { None, None, ) + .unwrap() .first() - .get_two::(); + .get_two::() + .unwrap(); assert_eq!(count, Some(100)); assert_eq!(count, arrow_count); let count2 = client - .select( + .update( &format!("SELECT distinct_count('{}')", expected), None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(count2, count); }); } @@ -616,8 +628,10 @@ mod tests { unsafe { // Unable to build the hyperloglog through hyperloglog_trans, as that requires a valid fcinfo to determine OIDs. - // FIXME: use named constant for default correlation oid - let hasher = DatumHashBuilder::from_type_id(pg_sys::TEXTOID, Some(100)); + let hasher = DatumHashBuilder::from_type_id( + pg_sys::TEXTOID, + Some(crate::serialization::collations::DEFAULT_COLLATION_OID), + ); let mut control = HyperLogLogTrans { logger: HLL::new(6, hasher), }; @@ -644,7 +658,11 @@ mod tests { 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 12, 0, 0, 0, 0, 0, 0, 0, 136, 136, 9, 7, 8, 74, 76, 47, 200, 231, 53, 25, 3, 0, 0, 0, 0, 0, 0, 0, 6, 9, 0, 0, 0, 1, ]; - bincode::serialize_into(&mut expected, &PgCollationId(100)).unwrap(); + bincode::serialize_into( + &mut expected, + &PgCollationId(crate::serialization::collations::DEFAULT_COLLATION_OID), + ) + .unwrap(); assert_eq!(buffer, expected); let expected = pgx::varlena::rust_byte_slice_to_bytea(&expected); @@ -670,7 +688,11 @@ mod tests { 49, 2, 8, 65, 131, 24, 32, 133, 12, 50, 66, 12, 48, 197, 12, 81, 130, 255, 58, 6, 255, 255, 255, 255, 255, 255, 255, 3, 9, 0, 0, 0, 1, ]; - bincode::serialize_into(&mut expected, &PgCollationId(100)).unwrap(); + bincode::serialize_into( + &mut expected, + &PgCollationId(crate::serialization::collations::DEFAULT_COLLATION_OID), + ) + .unwrap(); assert_eq!(buffer, expected); let expected = pgx::varlena::rust_byte_slice_to_bytea(&expected); @@ -683,16 +705,18 @@ mod tests { #[pg_test] fn test_hll_aggregate_int() { - Spi::execute(|client| { + Spi::connect(|mut client| { let text = client - .select( + .update( "SELECT hyperloglog(32, v::int)::TEXT FROM generate_series(1, 100) v", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); let expected = "(\ version:1,\ @@ -709,7 +733,7 @@ mod tests { assert_eq!(text.unwrap(), expected); let count = client - .select( + .update( "SELECT \ distinct_count(\ hyperloglog(32, v::int)\ @@ -717,39 +741,48 @@ mod tests { None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(count, Some(96)); let count2 = client - .select( + .update( &format!("SELECT distinct_count('{}')", expected), None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(count2, count); }); } #[pg_test] fn test_hll_aggregate_text() { - Spi::execute(|client| { + Spi::connect(|mut client| { use crate::serialization::PgCollationId; let text = client - .select( + .update( "SELECT \ hyperloglog(32, v::text)::TEXT \ FROM generate_series(1, 100) v", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); - let default_collation = ron::to_string(&PgCollationId(100)).unwrap(); + let default_collation = ron::to_string(&PgCollationId( + crate::serialization::collations::DEFAULT_COLLATION_OID, + )) + .unwrap(); let expected = format!( "(\ version:1,\ @@ -768,48 +801,54 @@ mod tests { assert_eq!(text.unwrap(), expected); let count = client - .select( + .update( "SELECT distinct_count(\ hyperloglog(32, v::text)\ ) FROM generate_series(1, 100) v", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(count, Some(111)); let count2 = client - .select( + .update( &format!("SELECT distinct_count('{}')", expected), None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(count2, count); }); } #[pg_test] fn test_hll_union_text() { - Spi::execute(|client| { + Spi::connect(|mut client| { { // self-union should be a nop let expected = client - .select( + .update( "SELECT \ hyperloglog(32, v::text)::TEXT \ FROM generate_series(1, 100) v", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() .unwrap(); let text = client - .select( + .update( "SELECT rollup(logs)::text \ FROM (\ (SELECT hyperloglog(32, v::text) logs \ @@ -821,8 +860,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(text.unwrap(), expected); } @@ -837,7 +878,12 @@ mod tests { (SELECT hyperloglog(32, v::text) \ FROM generate_series(50, 150) v)\ ) q"; - let count = client.select(query, None, None).first().get_one::(); + let count = client + .update(query, None, None) + .unwrap() + .first() + .get_one::() + .unwrap(); assert_eq!(count, Some(153)); } @@ -846,11 +892,13 @@ mod tests { #[pg_test] fn test_hll_null_input_yields_null_output() { - Spi::execute(|client| { + Spi::connect(|mut client| { let output = client - .select("SELECT hyperloglog(32, null::int)::TEXT", None, None) + .update("SELECT hyperloglog(32, null::int)::TEXT", None, None) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(output, None) }) } @@ -859,42 +907,48 @@ mod tests { error = "Invalid value for size 2. Size must be between 16 and 262144, though less than 1024 not recommended" )] fn test_hll_error_too_small() { - Spi::execute(|client| { + Spi::connect(|mut client| { let output = client - .select("SELECT hyperloglog(2, 'foo'::text)::TEXT", None, None) + .update("SELECT hyperloglog(2, 'foo'::text)::TEXT", None, None) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(output, None) }) } #[pg_test] fn test_hll_size_min() { - Spi::execute(|client| { + Spi::connect(|mut client| { let output = client - .select("SELECT hyperloglog(16, 'foo'::text)::TEXT", None, None) + .update("SELECT hyperloglog(16, 'foo'::text)::TEXT", None, None) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert!(output.is_some()) }) } #[pg_test] fn test_hll_size_max() { - Spi::execute(|client| { + Spi::connect(|mut client| { let output = client - .select("SELECT hyperloglog(262144, 'foo'::text)::TEXT", None, None) + .update("SELECT hyperloglog(262144, 'foo'::text)::TEXT", None, None) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert!(output.is_some()) }) } #[pg_test] fn stderror_arrow_match() { - Spi::execute(|client| { + Spi::connect(|mut client| { let (count, arrow_count) = client - .select( + .update( "SELECT \ stderror(\ hyperloglog(32, v::float)\ @@ -904,9 +958,11 @@ mod tests { None, None, ) + .unwrap() .first() - .get_two::(); - assert_eq!(Some(-788581389), count); + .get_two::() + .unwrap(); + assert_eq!(Some(0.18384776310850234), count); assert_eq!(count, arrow_count); }); } @@ -915,7 +971,7 @@ mod tests { fn bias_correct_values_accurate() { const NUM_BIAS_TRIALS: usize = 5; const MAX_TRIAL_ERROR: f64 = 0.05; - Spi::execute(|client| { + Spi::connect(|mut client| { // This should match THRESHOLD_DATA_VEC from b=12-18 let thresholds = vec![3100, 6500, 11500, 20000, 50000, 120000, 350000]; let rand_precision: Uniform = Uniform::new_inclusive(12, 18); @@ -932,9 +988,11 @@ mod tests { ); let estimate = client - .select(&query, None, None) + .update(&query, None, None) + .unwrap() .first() .get_one::() + .unwrap() .unwrap(); let error = (estimate as f64 / cardinality as f64).abs() - 1.; @@ -947,11 +1005,13 @@ mod tests { error = "Invalid value for size 262145. Size must be between 16 and 262144, though less than 1024 not recommended" )] fn test_hll_error_too_large() { - Spi::execute(|client| { + Spi::connect(|mut client| { let output = client - .select("SELECT hyperloglog(262145, 'foo'::text)::TEXT", None, None) + .update("SELECT hyperloglog(262145, 'foo'::text)::TEXT", None, None) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(output, None) }) } diff --git a/extension/src/lib.rs b/extension/src/lib.rs index 2af617f5..cfb205d2 100644 --- a/extension/src/lib.rs +++ b/extension/src/lib.rs @@ -24,6 +24,7 @@ pub mod lttb; pub mod nmost; pub mod range; pub mod saturation; +pub(crate) mod serialization; pub mod state_aggregate; pub mod stats_agg; pub mod tdigest; @@ -38,7 +39,6 @@ mod duration; mod palloc; mod pg_any_element; mod raw; -mod serialization; mod stabilization_info; mod stabilization_tests; mod type_builder; diff --git a/extension/src/lttb.rs b/extension/src/lttb.rs index 8af90d85..53295f3b 100644 --- a/extension/src/lttb.rs +++ b/extension/src/lttb.rs @@ -435,63 +435,74 @@ mod tests { #[pg_test] fn test_lttb_equivalence() { - Spi::execute(|client| { - client.select( - "CREATE TABLE test(time TIMESTAMPTZ, value DOUBLE PRECISION);", - None, - None, - ); - client.select( + Spi::connect(|mut client| { + client + .update( + "CREATE TABLE test(time TIMESTAMPTZ, value DOUBLE PRECISION);", + None, + None, + ) + .unwrap(); + client.update( "INSERT INTO test SELECT time, value - FROM toolkit_experimental.generate_periodic_normal_series('2020-01-01 UTC'::timestamptz, NULL);", None, None); - - client.select( - "CREATE TABLE results1(time TIMESTAMPTZ, value DOUBLE PRECISION);", - None, - None, - ); - client.select( - "INSERT INTO results1 + FROM toolkit_experimental.generate_periodic_normal_series('2020-01-01 UTC'::timestamptz, NULL);", None, None).unwrap(); + + client + .update( + "CREATE TABLE results1(time TIMESTAMPTZ, value DOUBLE PRECISION);", + None, + None, + ) + .unwrap(); + client + .update( + "INSERT INTO results1 SELECT time, value FROM unnest( (SELECT lttb(time, value, 100) FROM test) );", - None, - None, - ); - - client.select( - "CREATE TABLE results2(time TIMESTAMPTZ, value DOUBLE PRECISION);", - None, - None, - ); - client.select( - "INSERT INTO results2 + None, + None, + ) + .unwrap(); + + client + .update( + "CREATE TABLE results2(time TIMESTAMPTZ, value DOUBLE PRECISION);", + None, + None, + ) + .unwrap(); + client + .update( + "INSERT INTO results2 SELECT time, value FROM unnest( (SELECT lttb( (SELECT timevector(time, value) FROM test), 100) ) );", - None, - None, - ); + None, + None, + ) + .unwrap(); let delta = client - .select("SELECT count(*) FROM results1 r1 FULL OUTER JOIN results2 r2 ON r1 = r2 WHERE r1 IS NULL OR r2 IS NULL;" , None, None) - .first() - .get_one::(); + .update("SELECT count(*) FROM results1 r1 FULL OUTER JOIN results2 r2 ON r1 = r2 WHERE r1 IS NULL OR r2 IS NULL;" , None, None) + .unwrap().first() + .get_one::().unwrap(); assert_eq!(delta.unwrap(), 0); }) } #[pg_test] fn test_lttb_result() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); - let mut result = client.select( - r#"SELECT unnest(lttb(ts, val, 5))::TEXT + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); + let mut result = client + .update( + r#"SELECT unnest(lttb(ts, val, 5))::TEXT FROM (VALUES ('2020-1-1'::timestamptz, 10), ('2020-1-2'::timestamptz, 21), @@ -505,28 +516,29 @@ mod tests { ('2020-1-10'::timestamptz, 27), ('2020-1-11'::timestamptz, 14) ) AS v(ts, val)"#, - None, - None, - ); + None, + None, + ) + .unwrap(); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-01 00:00:00+00\",10)") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-04 00:00:00+00\",32)") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-05 00:00:00+00\",12)") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-08 00:00:00+00\",29)") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-11 00:00:00+00\",14)") ); assert!(result.next().is_none()); @@ -535,10 +547,11 @@ mod tests { #[pg_test] fn test_gp_lttb() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); - let mut result = client.select( - r#"SELECT unnest(toolkit_experimental.gp_lttb(ts, val, 7))::TEXT + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); + let mut result = client + .update( + r#"SELECT unnest(toolkit_experimental.gp_lttb(ts, val, 7))::TEXT FROM (VALUES ('2020-1-1'::timestamptz, 10), ('2020-1-2'::timestamptz, 21), @@ -552,36 +565,37 @@ mod tests { ('2020-3-10'::timestamptz, 27), ('2020-3-11'::timestamptz, 14) ) AS v(ts, val)"#, - None, - None, - ); + None, + None, + ) + .unwrap(); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-01 00:00:00+00\",10)") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-04 00:00:00+00\",32)") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-05 00:00:00+00\",12)") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-02-06 00:00:00+00\",14)") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-03-07 00:00:00+00\",18)") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-03-08 00:00:00+00\",29)") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-03-11 00:00:00+00\",14)") ); assert!(result.next().is_none()); @@ -590,10 +604,11 @@ mod tests { #[pg_test] fn test_gp_lttb_with_gap() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); - let mut result = client.select( - r#"SELECT unnest(toolkit_experimental.gp_lttb(ts, val, '36hr', 5))::TEXT + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); + let mut result = client + .update( + r#"SELECT unnest(toolkit_experimental.gp_lttb(ts, val, '36hr', 5))::TEXT FROM (VALUES ('2020-1-1'::timestamptz, 10), ('2020-1-2'::timestamptz, 21), @@ -605,45 +620,46 @@ mod tests { ('2020-3-10'::timestamptz, 27), ('2020-3-11'::timestamptz, 14) ) AS v(ts, val)"#, - None, - None, - ); + None, + None, + ) + .unwrap(); // This should include everything, despite target resolution of 5 assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-01 00:00:00+00\",10)") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-02 00:00:00+00\",21)") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-04 00:00:00+00\",32)") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-05 00:00:00+00\",12)") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-02-06 00:00:00+00\",14)") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-03-07 00:00:00+00\",18)") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-03-08 00:00:00+00\",29)") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-03-10 00:00:00+00\",27)") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-03-11 00:00:00+00\",14)") ); assert!(result.next().is_none()); diff --git a/extension/src/nmost/max_by_float.rs b/extension/src/nmost/max_by_float.rs index 711db4b7..c9fea174 100644 --- a/extension/src/nmost/max_by_float.rs +++ b/extension/src/nmost/max_by_float.rs @@ -150,68 +150,72 @@ mod tests { #[pg_test] fn max_by_float_correctness() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); - client.select( - "CREATE TABLE data(val DOUBLE PRECISION, category INT)", - None, - None, - ); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); + client + .update( + "CREATE TABLE data(val DOUBLE PRECISION, category INT)", + None, + None, + ) + .unwrap(); for i in 0..100 { let i = (i * 83) % 100; // mess with the ordering just a little - client.select( - &format!("INSERT INTO data VALUES ({}.0/128, {})", i, i % 4), - None, - None, - ); + client + .update( + &format!("INSERT INTO data VALUES ({}.0/128, {})", i, i % 4), + None, + None, + ) + .unwrap(); } // Test into_values let mut result = - client.select("SELECT toolkit_experimental.into_values(toolkit_experimental.max_n_by(val, data, 3), NULL::data)::TEXT from data", + client.update("SELECT toolkit_experimental.into_values(toolkit_experimental.max_n_by(val, data, 3), NULL::data)::TEXT from data", None, None, - ); + ).unwrap(); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(0.7734375,\"(0.7734375,3)\")") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(0.765625,\"(0.765625,2)\")") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(0.7578125,\"(0.7578125,1)\")") ); assert!(result.next().is_none()); // Test rollup let mut result = - client.select( + client.update( "WITH aggs as (SELECT category, toolkit_experimental.max_n_by(val, data, 5) as agg from data GROUP BY category) SELECT toolkit_experimental.into_values(toolkit_experimental.rollup(agg), NULL::data)::TEXT FROM aggs", None, None, - ); + ).unwrap(); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(0.7734375,\"(0.7734375,3)\")") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(0.765625,\"(0.765625,2)\")") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(0.7578125,\"(0.7578125,1)\")") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(0.75,\"(0.75,0)\")") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(0.7421875,\"(0.7421875,3)\")") ); assert!(result.next().is_none()); diff --git a/extension/src/nmost/max_by_int.rs b/extension/src/nmost/max_by_int.rs index 29649b8a..505a2f1b 100644 --- a/extension/src/nmost/max_by_int.rs +++ b/extension/src/nmost/max_by_int.rs @@ -149,42 +149,70 @@ mod tests { #[pg_test] fn max_by_int_correctness() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); - client.select("CREATE TABLE data(val INT8, category INT)", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); + client + .update("CREATE TABLE data(val INT8, category INT)", None, None) + .unwrap(); for i in 0..100 { let i = (i * 83) % 100; // mess with the ordering just a little - client.select( - &format!("INSERT INTO data VALUES ({}, {})", i, i % 4), - None, - None, - ); + client + .update( + &format!("INSERT INTO data VALUES ({}, {})", i, i % 4), + None, + None, + ) + .unwrap(); } // Test into_values let mut result = - client.select("SELECT toolkit_experimental.into_values(toolkit_experimental.max_n_by(val, data, 3), NULL::data)::TEXT from data", + client.update("SELECT toolkit_experimental.into_values(toolkit_experimental.max_n_by(val, data, 3), NULL::data)::TEXT from data", None, None, - ); - assert_eq!(result.next().unwrap()[1].value(), Some("(99,\"(99,3)\")")); - assert_eq!(result.next().unwrap()[1].value(), Some("(98,\"(98,2)\")")); - assert_eq!(result.next().unwrap()[1].value(), Some("(97,\"(97,1)\")")); + ).unwrap(); + assert_eq!( + result.next().unwrap()[1].value().unwrap(), + Some("(99,\"(99,3)\")") + ); + assert_eq!( + result.next().unwrap()[1].value().unwrap(), + Some("(98,\"(98,2)\")") + ); + assert_eq!( + result.next().unwrap()[1].value().unwrap(), + Some("(97,\"(97,1)\")") + ); assert!(result.next().is_none()); // Test rollup let mut result = - client.select( + client.update( "WITH aggs as (SELECT category, toolkit_experimental.max_n_by(val, data, 5) as agg from data GROUP BY category) SELECT toolkit_experimental.into_values(toolkit_experimental.rollup(agg), NULL::data)::TEXT FROM aggs", None, None, - ); - assert_eq!(result.next().unwrap()[1].value(), Some("(99,\"(99,3)\")")); - assert_eq!(result.next().unwrap()[1].value(), Some("(98,\"(98,2)\")")); - assert_eq!(result.next().unwrap()[1].value(), Some("(97,\"(97,1)\")")); - assert_eq!(result.next().unwrap()[1].value(), Some("(96,\"(96,0)\")")); - assert_eq!(result.next().unwrap()[1].value(), Some("(95,\"(95,3)\")")); + ).unwrap(); + assert_eq!( + result.next().unwrap()[1].value().unwrap(), + Some("(99,\"(99,3)\")") + ); + assert_eq!( + result.next().unwrap()[1].value().unwrap(), + Some("(98,\"(98,2)\")") + ); + assert_eq!( + result.next().unwrap()[1].value().unwrap(), + Some("(97,\"(97,1)\")") + ); + assert_eq!( + result.next().unwrap()[1].value().unwrap(), + Some("(96,\"(96,0)\")") + ); + assert_eq!( + result.next().unwrap()[1].value().unwrap(), + Some("(95,\"(95,3)\")") + ); assert!(result.next().is_none()); }) } diff --git a/extension/src/nmost/max_by_time.rs b/extension/src/nmost/max_by_time.rs index 18d2db98..ccee7021 100644 --- a/extension/src/nmost/max_by_time.rs +++ b/extension/src/nmost/max_by_time.rs @@ -156,68 +156,70 @@ mod tests { #[pg_test] fn max_by_time_correctness() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); - client.select( - "CREATE TABLE data(val TIMESTAMPTZ, category INT)", - None, - None, - ); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); + client + .update( + "CREATE TABLE data(val TIMESTAMPTZ, category INT)", + None, + None, + ) + .unwrap(); for i in 0..100 { let i = (i * 83) % 100; // mess with the ordering just a little - client.select( + client.update( &format!("INSERT INTO data VALUES ('2020-1-1 UTC'::timestamptz + {} * '1d'::interval, {})", i, i % 4), None, None, - ); + ).unwrap(); } // Test into_values let mut result = - client.select("SELECT toolkit_experimental.into_values(toolkit_experimental.max_n_by(val, data, 3), NULL::data)::TEXT from data", + client.update("SELECT toolkit_experimental.into_values(toolkit_experimental.max_n_by(val, data, 3), NULL::data)::TEXT from data", None, None, - ); + ).unwrap(); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-04-09 00:00:00+00\",\"(\"\"2020-04-09 00:00:00+00\"\",3)\")") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-04-08 00:00:00+00\",\"(\"\"2020-04-08 00:00:00+00\"\",2)\")") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-04-07 00:00:00+00\",\"(\"\"2020-04-07 00:00:00+00\"\",1)\")") ); assert!(result.next().is_none()); // Test rollup let mut result = - client.select( + client.update( "WITH aggs as (SELECT category, toolkit_experimental.max_n_by(val, data, 5) as agg from data GROUP BY category) SELECT toolkit_experimental.into_values(toolkit_experimental.rollup(agg), NULL::data)::TEXT FROM aggs", None, None, - ); + ).unwrap(); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-04-09 00:00:00+00\",\"(\"\"2020-04-09 00:00:00+00\"\",3)\")") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-04-08 00:00:00+00\",\"(\"\"2020-04-08 00:00:00+00\"\",2)\")") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-04-07 00:00:00+00\",\"(\"\"2020-04-07 00:00:00+00\"\",1)\")") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-04-06 00:00:00+00\",\"(\"\"2020-04-06 00:00:00+00\"\",0)\")") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-04-05 00:00:00+00\",\"(\"\"2020-04-05 00:00:00+00\"\",3)\")") ); assert!(result.next().is_none()); diff --git a/extension/src/nmost/max_float.rs b/extension/src/nmost/max_float.rs index 8e953ca7..5498b35e 100644 --- a/extension/src/nmost/max_float.rs +++ b/extension/src/nmost/max_float.rs @@ -186,29 +186,33 @@ mod tests { #[pg_test] fn max_float_correctness() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); - client.select( - "CREATE TABLE data(val DOUBLE PRECISION, category INT)", - None, - None, - ); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); + client + .update( + "CREATE TABLE data(val DOUBLE PRECISION, category INT)", + None, + None, + ) + .unwrap(); for i in 0..100 { let i = (i * 83) % 100; // mess with the ordering just a little - client.select( - &format!("INSERT INTO data VALUES ({}.0/128, {})", i, i % 4), - None, - None, - ); + client + .update( + &format!("INSERT INTO data VALUES ({}.0/128, {})", i, i % 4), + None, + None, + ) + .unwrap(); } // Test into_array let result = - client.select("SELECT toolkit_experimental.into_array(toolkit_experimental.max_n(val, 5)) from data", + client.update("SELECT toolkit_experimental.into_array(toolkit_experimental.max_n(val, 5)) from data", None, None, - ).first().get_one::>(); + ).unwrap().first().get_one::>().unwrap(); assert_eq!( result.unwrap(), vec![99. / 128., 98. / 128., 97. / 128., 96. / 128., 95. / 128.] @@ -216,21 +220,27 @@ mod tests { // Test into_values let mut result = - client.select("SELECT toolkit_experimental.into_values(toolkit_experimental.max_n(val, 3))::TEXT from data", + client.update("SELECT toolkit_experimental.into_values(toolkit_experimental.max_n(val, 3))::TEXT from data", None, None, - ); - assert_eq!(result.next().unwrap()[1].value(), Some("0.7734375")); - assert_eq!(result.next().unwrap()[1].value(), Some("0.765625")); - assert_eq!(result.next().unwrap()[1].value(), Some("0.7578125")); + ).unwrap(); + assert_eq!( + result.next().unwrap()[1].value().unwrap(), + Some("0.7734375") + ); + assert_eq!(result.next().unwrap()[1].value().unwrap(), Some("0.765625")); + assert_eq!( + result.next().unwrap()[1].value().unwrap(), + Some("0.7578125") + ); assert!(result.next().is_none()); // Test rollup let result = - client.select( + client.update( "WITH aggs as (SELECT category, toolkit_experimental.max_n(val, 5) as agg from data GROUP BY category) SELECT toolkit_experimental.into_array(toolkit_experimental.rollup(agg)) FROM aggs", None, None, - ).first().get_one::>(); + ).unwrap().first().get_one::>().unwrap(); assert_eq!( result.unwrap(), vec![99. / 128., 98. / 128., 97. / 128., 96. / 128., 95. / 128.] diff --git a/extension/src/nmost/max_int.rs b/extension/src/nmost/max_int.rs index 95bf2ce9..da47eef9 100644 --- a/extension/src/nmost/max_int.rs +++ b/extension/src/nmost/max_int.rs @@ -180,44 +180,48 @@ mod tests { #[pg_test] fn max_int_correctness() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); - client.select("CREATE TABLE data(val INT8, category INT)", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); + client + .update("CREATE TABLE data(val INT8, category INT)", None, None) + .unwrap(); for i in 0..100 { let i = (i * 83) % 100; // mess with the ordering just a little - client.select( - &format!("INSERT INTO data VALUES ({}, {})", i, i % 4), - None, - None, - ); + client + .update( + &format!("INSERT INTO data VALUES ({}, {})", i, i % 4), + None, + None, + ) + .unwrap(); } // Test into_array let result = - client.select("SELECT toolkit_experimental.into_array(toolkit_experimental.max_n(val, 5)) from data", + client.update("SELECT toolkit_experimental.into_array(toolkit_experimental.max_n(val, 5)) from data", None, None, - ).first().get_one::>(); + ).unwrap().first().get_one::>().unwrap(); assert_eq!(result.unwrap(), vec![99, 98, 97, 96, 95]); // Test into_values let mut result = - client.select("SELECT toolkit_experimental.into_values(toolkit_experimental.max_n(val, 3))::TEXT from data", + client.update("SELECT toolkit_experimental.into_values(toolkit_experimental.max_n(val, 3))::TEXT from data", None, None, - ); - assert_eq!(result.next().unwrap()[1].value(), Some("99")); - assert_eq!(result.next().unwrap()[1].value(), Some("98")); - assert_eq!(result.next().unwrap()[1].value(), Some("97")); + ).unwrap(); + assert_eq!(result.next().unwrap()[1].value().unwrap(), Some("99")); + assert_eq!(result.next().unwrap()[1].value().unwrap(), Some("98")); + assert_eq!(result.next().unwrap()[1].value().unwrap(), Some("97")); assert!(result.next().is_none()); // Test rollup let result = - client.select( + client.update( "WITH aggs as (SELECT category, toolkit_experimental.max_n(val, 5) as agg from data GROUP BY category) SELECT toolkit_experimental.into_array(toolkit_experimental.rollup(agg)) FROM aggs", None, None, - ).first().get_one::>(); + ).unwrap().first().get_one::>().unwrap(); assert_eq!(result.unwrap(), vec![99, 98, 97, 96, 95]); }) } diff --git a/extension/src/nmost/max_time.rs b/extension/src/nmost/max_time.rs index 17c50b5e..99ff0fd1 100644 --- a/extension/src/nmost/max_time.rs +++ b/extension/src/nmost/max_time.rs @@ -194,57 +194,59 @@ mod tests { #[pg_test] fn max_time_correctness() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); - client.select( - "CREATE TABLE data(val TIMESTAMPTZ, category INT)", - None, - None, - ); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); + client + .update( + "CREATE TABLE data(val TIMESTAMPTZ, category INT)", + None, + None, + ) + .unwrap(); for i in 0..100 { let i = (i * 83) % 100; // mess with the ordering just a little - client.select( + client.update( &format!("INSERT INTO data VALUES ('2020-1-1 UTC'::timestamptz + {} * '1d'::interval, {})", i, i % 4), None, None, - ); + ).unwrap(); } // Test into_array let result = - client.select("SELECT toolkit_experimental.into_array(toolkit_experimental.max_n(val, 5))::TEXT from data", + client.update("SELECT toolkit_experimental.into_array(toolkit_experimental.max_n(val, 5))::TEXT from data", None, None, - ).first().get_one::<&str>(); + ).unwrap().first().get_one::<&str>().unwrap(); assert_eq!(result.unwrap(), "{\"2020-04-09 00:00:00+00\",\"2020-04-08 00:00:00+00\",\"2020-04-07 00:00:00+00\",\"2020-04-06 00:00:00+00\",\"2020-04-05 00:00:00+00\"}"); // Test into_values let mut result = - client.select("SELECT toolkit_experimental.into_values(toolkit_experimental.max_n(val, 3))::TEXT from data", + client.update("SELECT toolkit_experimental.into_values(toolkit_experimental.max_n(val, 3))::TEXT from data", None, None, - ); + ).unwrap(); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("2020-04-09 00:00:00+00") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("2020-04-08 00:00:00+00") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("2020-04-07 00:00:00+00") ); assert!(result.next().is_none()); // Test rollup let result = - client.select( + client.update( "WITH aggs as (SELECT category, toolkit_experimental.max_n(val, 5) as agg from data GROUP BY category) SELECT toolkit_experimental.into_array(toolkit_experimental.rollup(agg))::TEXT FROM aggs", None, None, - ).first().get_one::<&str>(); + ).unwrap().first().get_one::<&str>().unwrap(); assert_eq!(result.unwrap(), "{\"2020-04-09 00:00:00+00\",\"2020-04-08 00:00:00+00\",\"2020-04-07 00:00:00+00\",\"2020-04-06 00:00:00+00\",\"2020-04-05 00:00:00+00\"}"); }) } diff --git a/extension/src/nmost/min_by_float.rs b/extension/src/nmost/min_by_float.rs index d1fe0af0..c422698e 100644 --- a/extension/src/nmost/min_by_float.rs +++ b/extension/src/nmost/min_by_float.rs @@ -149,62 +149,72 @@ mod tests { #[pg_test] fn min_by_float_correctness() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); - client.select( - "CREATE TABLE data(val DOUBLE PRECISION, category INT)", - None, - None, - ); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); + client + .update( + "CREATE TABLE data(val DOUBLE PRECISION, category INT)", + None, + None, + ) + .unwrap(); for i in 0..100 { let i = (i * 83) % 100; // mess with the ordering just a little - client.select( - &format!("INSERT INTO data VALUES ({}.0/128, {})", i, i % 4), - None, - None, - ); + client + .update( + &format!("INSERT INTO data VALUES ({}.0/128, {})", i, i % 4), + None, + None, + ) + .unwrap(); } // Test into_values let mut result = - client.select("SELECT toolkit_experimental.into_values(toolkit_experimental.min_n_by(val, data, 3), NULL::data)::TEXT from data", + client.update("SELECT toolkit_experimental.into_values(toolkit_experimental.min_n_by(val, data, 3), NULL::data)::TEXT from data", None, None, - ); - assert_eq!(result.next().unwrap()[1].value(), Some("(0,\"(0,0)\")")); + ).unwrap(); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), + Some("(0,\"(0,0)\")") + ); + assert_eq!( + result.next().unwrap()[1].value().unwrap(), Some("(0.0078125,\"(0.0078125,1)\")") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(0.015625,\"(0.015625,2)\")") ); assert!(result.next().is_none()); // Test rollup let mut result = - client.select( + client.update( "WITH aggs as (SELECT category, toolkit_experimental.min_n_by(val, data, 5) as agg from data GROUP BY category) SELECT toolkit_experimental.into_values(toolkit_experimental.rollup(agg), NULL::data)::TEXT FROM aggs", None, None, - ); - assert_eq!(result.next().unwrap()[1].value(), Some("(0,\"(0,0)\")")); + ).unwrap(); + assert_eq!( + result.next().unwrap()[1].value().unwrap(), + Some("(0,\"(0,0)\")") + ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(0.0078125,\"(0.0078125,1)\")") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(0.015625,\"(0.015625,2)\")") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(0.0234375,\"(0.0234375,3)\")") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(0.03125,\"(0.03125,0)\")") ); assert!(result.next().is_none()); diff --git a/extension/src/nmost/min_by_int.rs b/extension/src/nmost/min_by_int.rs index 18536289..026633f2 100644 --- a/extension/src/nmost/min_by_int.rs +++ b/extension/src/nmost/min_by_int.rs @@ -136,42 +136,70 @@ mod tests { #[pg_test] fn min_by_int_correctness() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); - client.select("CREATE TABLE data(val INT8, category INT)", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); + client + .update("CREATE TABLE data(val INT8, category INT)", None, None) + .unwrap(); for i in 0..100 { let i = (i * 83) % 100; // mess with the ordering just a little - client.select( - &format!("INSERT INTO data VALUES ({}, {})", i, i % 4), - None, - None, - ); + client + .update( + &format!("INSERT INTO data VALUES ({}, {})", i, i % 4), + None, + None, + ) + .unwrap(); } // Test into_values let mut result = - client.select("SELECT toolkit_experimental.into_values(toolkit_experimental.min_n_by(val, data, 3), NULL::data)::TEXT from data", + client.update("SELECT toolkit_experimental.into_values(toolkit_experimental.min_n_by(val, data, 3), NULL::data)::TEXT from data", None, None, - ); - assert_eq!(result.next().unwrap()[1].value(), Some("(0,\"(0,0)\")")); - assert_eq!(result.next().unwrap()[1].value(), Some("(1,\"(1,1)\")")); - assert_eq!(result.next().unwrap()[1].value(), Some("(2,\"(2,2)\")")); + ).unwrap(); + assert_eq!( + result.next().unwrap()[1].value().unwrap(), + Some("(0,\"(0,0)\")") + ); + assert_eq!( + result.next().unwrap()[1].value().unwrap(), + Some("(1,\"(1,1)\")") + ); + assert_eq!( + result.next().unwrap()[1].value().unwrap(), + Some("(2,\"(2,2)\")") + ); assert!(result.next().is_none()); // Test rollup let mut result = - client.select( + client.update( "WITH aggs as (SELECT category, toolkit_experimental.min_n_by(val, data, 5) as agg from data GROUP BY category) SELECT toolkit_experimental.into_values(toolkit_experimental.rollup(agg), NULL::data)::TEXT FROM aggs", None, None, - ); - assert_eq!(result.next().unwrap()[1].value(), Some("(0,\"(0,0)\")")); - assert_eq!(result.next().unwrap()[1].value(), Some("(1,\"(1,1)\")")); - assert_eq!(result.next().unwrap()[1].value(), Some("(2,\"(2,2)\")")); - assert_eq!(result.next().unwrap()[1].value(), Some("(3,\"(3,3)\")")); - assert_eq!(result.next().unwrap()[1].value(), Some("(4,\"(4,0)\")")); + ).unwrap(); + assert_eq!( + result.next().unwrap()[1].value().unwrap(), + Some("(0,\"(0,0)\")") + ); + assert_eq!( + result.next().unwrap()[1].value().unwrap(), + Some("(1,\"(1,1)\")") + ); + assert_eq!( + result.next().unwrap()[1].value().unwrap(), + Some("(2,\"(2,2)\")") + ); + assert_eq!( + result.next().unwrap()[1].value().unwrap(), + Some("(3,\"(3,3)\")") + ); + assert_eq!( + result.next().unwrap()[1].value().unwrap(), + Some("(4,\"(4,0)\")") + ); assert!(result.next().is_none()); }) } diff --git a/extension/src/nmost/min_by_time.rs b/extension/src/nmost/min_by_time.rs index fa1dce43..5fca2f03 100644 --- a/extension/src/nmost/min_by_time.rs +++ b/extension/src/nmost/min_by_time.rs @@ -143,68 +143,70 @@ mod tests { #[pg_test] fn min_by_time_correctness() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); - client.select( - "CREATE TABLE data(val TIMESTAMPTZ, category INT)", - None, - None, - ); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); + client + .update( + "CREATE TABLE data(val TIMESTAMPTZ, category INT)", + None, + None, + ) + .unwrap(); for i in 0..100 { let i = (i * 83) % 100; // mess with the ordering just a little - client.select( + client.update( &format!("INSERT INTO data VALUES ('2020-1-1 UTC'::timestamptz + {} * '1d'::interval, {})", i, i % 4), None, None, - ); + ).unwrap(); } // Test into_values let mut result = - client.select("SELECT toolkit_experimental.into_values(toolkit_experimental.min_n_by(val, data, 3), NULL::data)::TEXT from data", + client.update("SELECT toolkit_experimental.into_values(toolkit_experimental.min_n_by(val, data, 3), NULL::data)::TEXT from data", None, None, - ); + ).unwrap(); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-01 00:00:00+00\",\"(\"\"2020-01-01 00:00:00+00\"\",0)\")") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-02 00:00:00+00\",\"(\"\"2020-01-02 00:00:00+00\"\",1)\")") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-03 00:00:00+00\",\"(\"\"2020-01-03 00:00:00+00\"\",2)\")") ); assert!(result.next().is_none()); // Test rollup let mut result = - client.select( + client.update( "WITH aggs as (SELECT category, toolkit_experimental.min_n_by(val, data, 5) as agg from data GROUP BY category) SELECT toolkit_experimental.into_values(toolkit_experimental.rollup(agg), NULL::data)::TEXT FROM aggs", None, None, - ); + ).unwrap(); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-01 00:00:00+00\",\"(\"\"2020-01-01 00:00:00+00\"\",0)\")") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-02 00:00:00+00\",\"(\"\"2020-01-02 00:00:00+00\"\",1)\")") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-03 00:00:00+00\",\"(\"\"2020-01-03 00:00:00+00\"\",2)\")") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-04 00:00:00+00\",\"(\"\"2020-01-04 00:00:00+00\"\",3)\")") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-05 00:00:00+00\",\"(\"\"2020-01-05 00:00:00+00\"\",0)\")") ); assert!(result.next().is_none()); diff --git a/extension/src/nmost/min_float.rs b/extension/src/nmost/min_float.rs index 7f074413..a538aa66 100644 --- a/extension/src/nmost/min_float.rs +++ b/extension/src/nmost/min_float.rs @@ -185,29 +185,33 @@ mod tests { #[pg_test] fn min_float_correctness() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); - client.select( - "CREATE TABLE data(val DOUBLE PRECISION, category INT)", - None, - None, - ); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); + client + .update( + "CREATE TABLE data(val DOUBLE PRECISION, category INT)", + None, + None, + ) + .unwrap(); for i in 0..100 { let i = (i * 83) % 100; // mess with the ordering just a little - client.select( - &format!("INSERT INTO data VALUES ({}.0/128, {})", i, i % 4), - None, - None, - ); + client + .update( + &format!("INSERT INTO data VALUES ({}.0/128, {})", i, i % 4), + None, + None, + ) + .unwrap(); } // Test into_array let result = - client.select("SELECT toolkit_experimental.into_array(toolkit_experimental.min_n(val, 5)) from data", + client.update("SELECT toolkit_experimental.into_array(toolkit_experimental.min_n(val, 5)) from data", None, None, - ).first().get_one::>(); + ).unwrap().first().get_one::>().unwrap(); assert_eq!( result.unwrap(), vec![0. / 128., 1. / 128., 2. / 128., 3. / 128., 4. / 128.] @@ -215,23 +219,26 @@ mod tests { // Test into_values let mut result = - client.select("SELECT toolkit_experimental.into_values(toolkit_experimental.min_n(val, 3))::TEXT from data", + client.update("SELECT toolkit_experimental.into_values(toolkit_experimental.min_n(val, 3))::TEXT from data", None, None, - ); - assert_eq!(result.next().unwrap()[1].value(), Some("0")); - assert_eq!(result.next().unwrap()[1].value(), Some("0.0078125")); - assert_eq!(result.next().unwrap()[1].value(), Some("0.015625")); + ).unwrap(); + assert_eq!(result.next().unwrap()[1].value().unwrap(), Some("0")); + assert_eq!( + result.next().unwrap()[1].value().unwrap(), + Some("0.0078125") + ); + assert_eq!(result.next().unwrap()[1].value().unwrap(), Some("0.015625")); assert!(result.next().is_none()); // Test rollup let result = - client.select( + client.update( "WITH aggs as (SELECT category, toolkit_experimental.min_n(val, 5) as agg from data GROUP BY category) SELECT toolkit_experimental.into_array(toolkit_experimental.rollup(agg)) FROM aggs", None, None, - ).first().get_one::>(); + ).unwrap().first().get_one::>(); assert_eq!( - result.unwrap(), + result.unwrap().unwrap(), vec![0. / 128., 1. / 128., 2. / 128., 3. / 128., 4. / 128.] ); }) diff --git a/extension/src/nmost/min_int.rs b/extension/src/nmost/min_int.rs index 8ea53b7c..bc3da09b 100644 --- a/extension/src/nmost/min_int.rs +++ b/extension/src/nmost/min_int.rs @@ -172,44 +172,48 @@ mod tests { #[pg_test] fn min_int_correctness() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); - client.select("CREATE TABLE data(val INT8, category INT)", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); + client + .update("CREATE TABLE data(val INT8, category INT)", None, None) + .unwrap(); for i in 0..100 { let i = (i * 83) % 100; // mess with the ordering just a little - client.select( - &format!("INSERT INTO data VALUES ({}, {})", i, i % 4), - None, - None, - ); + client + .update( + &format!("INSERT INTO data VALUES ({}, {})", i, i % 4), + None, + None, + ) + .unwrap(); } // Test into_array let result = - client.select("SELECT toolkit_experimental.into_array(toolkit_experimental.min_n(val, 5)) from data", + client.update("SELECT toolkit_experimental.into_array(toolkit_experimental.min_n(val, 5)) from data", None, None, - ).first().get_one::>(); + ).unwrap().first().get_one::>().unwrap(); assert_eq!(result.unwrap(), vec![0, 1, 2, 3, 4]); // Test into_values let mut result = - client.select("SELECT toolkit_experimental.into_values(toolkit_experimental.min_n(val, 3))::TEXT from data", + client.update("SELECT toolkit_experimental.into_values(toolkit_experimental.min_n(val, 3))::TEXT from data", None, None, - ); - assert_eq!(result.next().unwrap()[1].value(), Some("0")); - assert_eq!(result.next().unwrap()[1].value(), Some("1")); - assert_eq!(result.next().unwrap()[1].value(), Some("2")); + ).unwrap(); + assert_eq!(result.next().unwrap()[1].value().unwrap(), Some("0")); + assert_eq!(result.next().unwrap()[1].value().unwrap(), Some("1")); + assert_eq!(result.next().unwrap()[1].value().unwrap(), Some("2")); assert!(result.next().is_none()); // Test rollup let result = - client.select( + client.update( "WITH aggs as (SELECT category, toolkit_experimental.min_n(val, 5) as agg from data GROUP BY category) SELECT toolkit_experimental.into_array(toolkit_experimental.rollup(agg)) FROM aggs", None, None, - ).first().get_one::>(); + ).unwrap().first().get_one::>().unwrap(); assert_eq!(result.unwrap(), vec![0, 1, 2, 3, 4]); }) } diff --git a/extension/src/nmost/min_time.rs b/extension/src/nmost/min_time.rs index 464e9a73..af188302 100644 --- a/extension/src/nmost/min_time.rs +++ b/extension/src/nmost/min_time.rs @@ -185,57 +185,59 @@ mod tests { #[pg_test] fn min_time_correctness() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); - client.select( - "CREATE TABLE data(val TIMESTAMPTZ, category INT)", - None, - None, - ); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); + client + .update( + "CREATE TABLE data(val TIMESTAMPTZ, category INT)", + None, + None, + ) + .unwrap(); for i in 0..100 { let i = (i * 83) % 100; // mess with the ordering just a little - client.select( + client.update( &format!("INSERT INTO data VALUES ('2020-1-1 UTC'::timestamptz + {} * '1d'::interval, {})", i, i % 4), None, None, - ); + ).unwrap(); } // Test into_array let result = - client.select("SELECT toolkit_experimental.into_array(toolkit_experimental.min_n(val, 5))::TEXT from data", + client.update("SELECT toolkit_experimental.into_array(toolkit_experimental.min_n(val, 5))::TEXT from data", None, None, - ).first().get_one::<&str>(); + ).unwrap().first().get_one::<&str>().unwrap(); assert_eq!(result.unwrap(), "{\"2020-01-01 00:00:00+00\",\"2020-01-02 00:00:00+00\",\"2020-01-03 00:00:00+00\",\"2020-01-04 00:00:00+00\",\"2020-01-05 00:00:00+00\"}"); // Test into_values let mut result = - client.select("SELECT toolkit_experimental.into_values(toolkit_experimental.min_n(val, 3))::TEXT from data", + client.update("SELECT toolkit_experimental.into_values(toolkit_experimental.min_n(val, 3))::TEXT from data", None, None, - ); + ).unwrap(); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("2020-01-01 00:00:00+00") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("2020-01-02 00:00:00+00") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("2020-01-03 00:00:00+00") ); assert!(result.next().is_none()); // Test rollup let result = - client.select( + client.update( "WITH aggs as (SELECT category, toolkit_experimental.min_n(val, 5) as agg from data GROUP BY category) SELECT toolkit_experimental.into_array(toolkit_experimental.rollup(agg))::TEXT FROM aggs", None, None, - ).first().get_one::<&str>(); + ).unwrap().first().get_one::<&str>().unwrap(); assert_eq!(result.unwrap(), "{\"2020-01-01 00:00:00+00\",\"2020-01-02 00:00:00+00\",\"2020-01-03 00:00:00+00\",\"2020-01-04 00:00:00+00\",\"2020-01-05 00:00:00+00\"}"); }) } diff --git a/extension/src/raw.rs b/extension/src/raw.rs index 40daf8fd..be86084a 100644 --- a/extension/src/raw.rs +++ b/extension/src/raw.rs @@ -1,10 +1,8 @@ #![allow(non_camel_case_types)] -use pgx::{ - utils::sql_entity_graph::metadata::{ - ArgumentError, Returns, ReturnsError, SqlMapping, SqlTranslatable, - }, - *, +use pgx::*; +use pgx_sql_entity_graph::metadata::{ + ArgumentError, Returns, ReturnsError, SqlMapping, SqlTranslatable, }; extension_sql!( diff --git a/extension/src/serialization.rs b/extension/src/serialization.rs index 2c1de62a..fc9e6106 100644 --- a/extension/src/serialization.rs +++ b/extension/src/serialization.rs @@ -6,10 +6,10 @@ use std::{ os::raw::{c_char, c_int}, }; -use pgx::cstr_core::CStr; use pgx::pg_sys; +use std::ffi::CStr; -mod collations; +pub(crate) mod collations; mod functions; mod types; diff --git a/extension/src/serialization/collations.rs b/extension/src/serialization/collations.rs index 97918c53..e2283901 100644 --- a/extension/src/serialization/collations.rs +++ b/extension/src/serialization/collations.rs @@ -41,7 +41,7 @@ impl PgCollationId { #[allow(non_upper_case_globals)] const Anum_pg_collation_oid: u32 = 1; // https://github.com/postgres/postgres/blob/e955bd4b6c2bcdbd253837f6cf4c7520b98e69d4/src/include/catalog/pg_collation.dat -const DEFAULT_COLLATION_OID: u32 = 100; +pub(crate) const DEFAULT_COLLATION_OID: Oid = unsafe { pg_sys::Oid::from_u32_unchecked(100) }; #[allow(non_camel_case_types)] #[derive(Copy, Clone)] @@ -168,7 +168,7 @@ impl<'de> Deserialize<'de> for PgCollationId { let collation = >::deserialize(deserializer)?; let (namespace, name) = match collation { - None => return Ok(Self(0)), + None => return Ok(Self(pg_sys::Oid::INVALID)), Some(qualified_name) => qualified_name, }; @@ -229,7 +229,7 @@ impl<'de> Deserialize<'de> for PgCollationId { // The default collation doesn't necessarily exist in the // collations catalog, so check that specially if name == &**DEFAULT_COLLATION_NAME { - return Ok(PgCollationId(100)); + return Ok(PgCollationId(DEFAULT_COLLATION_OID)); } return Err(D::Error::custom(format!( "invalid collation {:?}.{:?}", @@ -254,22 +254,33 @@ unsafe fn get_struct(tuple: pg_sys::HeapTuple) -> *mut T { mod tests { use super::PgCollationId; - use pgx::{pg_guard, pg_sys, pg_test}; + use pgx::{pg_sys, pg_test}; + + const COLLATION_ID_950: PgCollationId = + PgCollationId(unsafe { pg_sys::Oid::from_u32_unchecked(950) }); + const COLLATION_ID_951: PgCollationId = + PgCollationId(unsafe { pg_sys::Oid::from_u32_unchecked(951) }); // TODO is there a way we can test more of this without making it flaky? #[pg_test] fn test_pg_collation_id_serialize_default_collation_ron() { - let serialized = ron::to_string(&PgCollationId(100)).unwrap(); + let serialized = ron::to_string(&PgCollationId( + crate::serialization::collations::DEFAULT_COLLATION_OID, + )) + .unwrap(); let deserialized: PgCollationId = ron::from_str(&serialized).unwrap(); - assert_ne!(deserialized.0, 0); - let serialized = ron::to_string(&PgCollationId(100)).unwrap(); + assert_ne!(deserialized.0, pg_sys::Oid::INVALID); + let serialized = ron::to_string(&PgCollationId( + crate::serialization::collations::DEFAULT_COLLATION_OID, + )) + .unwrap(); let deserialized2: PgCollationId = ron::from_str(&serialized).unwrap(); assert_eq!(deserialized2.0, deserialized.0); } #[pg_test] fn test_pg_collation_id_serialize_c_collation() { - let serialized = bincode::serialize(&PgCollationId(950)).unwrap(); + let serialized = bincode::serialize(&COLLATION_ID_950).unwrap(); assert_eq!( serialized, vec![ @@ -278,21 +289,21 @@ mod tests { ] ); let deserialized: PgCollationId = bincode::deserialize(&serialized).unwrap(); - assert_eq!(deserialized.0, 950); + assert_eq!(deserialized.0, COLLATION_ID_950.0); } // TODO this test may be too flaky depending on what the default collation actually is #[pg_test] fn test_pg_collation_id_serialize_c_collation_ron() { - let serialized = ron::to_string(&PgCollationId(950)).unwrap(); + let serialized = ron::to_string(&COLLATION_ID_950).unwrap(); assert_eq!(&*serialized, "Some((\"pg_catalog\",\"C\"))",); let deserialized: PgCollationId = ron::from_str(&serialized).unwrap(); - assert_eq!(deserialized.0, 950); + assert_eq!(deserialized.0, COLLATION_ID_950.0); } #[pg_test] fn test_pg_collation_id_serialize_posix_collation() { - let serialized = bincode::serialize(&PgCollationId(951)).unwrap(); + let serialized = bincode::serialize(&COLLATION_ID_951).unwrap(); assert_eq!( serialized, vec![ @@ -301,15 +312,15 @@ mod tests { ] ); let deserialized: PgCollationId = bincode::deserialize(&serialized).unwrap(); - assert_eq!(deserialized.0, 951); + assert_eq!(deserialized.0, COLLATION_ID_951.0); } // TODO this test may be too flaky depending on what the default collation actually is #[pg_test] fn test_pg_collation_id_serialize_posix_collation_ron() { - let serialized = ron::to_string(&PgCollationId(951)).unwrap(); + let serialized = ron::to_string(&COLLATION_ID_951).unwrap(); assert_eq!(&*serialized, "Some((\"pg_catalog\",\"POSIX\"))",); let deserialized: PgCollationId = ron::from_str(&serialized).unwrap(); - assert_eq!(deserialized.0, 951); + assert_eq!(deserialized.0, COLLATION_ID_951.0); } } diff --git a/extension/src/serialization/functions.rs b/extension/src/serialization/functions.rs index 750d36dd..8ec6fa6d 100644 --- a/extension/src/serialization/functions.rs +++ b/extension/src/serialization/functions.rs @@ -64,6 +64,6 @@ impl<'de> Deserialize<'de> for PgProcId { ) }; - Ok(Self(oid.value() as _)) + Ok(Self(unsafe { Oid::from_u32_unchecked(oid.value() as _) })) } } diff --git a/extension/src/serialization/types.rs b/extension/src/serialization/types.rs index 7f3bf698..d53f98f9 100644 --- a/extension/src/serialization/types.rs +++ b/extension/src/serialization/types.rs @@ -17,17 +17,17 @@ use pgx::*; /// that these types can be stored more compactly if desired. #[derive(Debug, Clone, Copy)] #[repr(transparent)] -pub struct ShortTypeId(pub u32); +pub struct ShortTypeId(pub Oid); impl_flat_serializable!(ShortTypeId); -impl From for ShortTypeId { - fn from(id: u32) -> Self { +impl From for ShortTypeId { + fn from(id: Oid) -> Self { Self(id) } } -impl From for u32 { +impl From for Oid { fn from(id: ShortTypeId) -> Self { id.0 } @@ -321,8 +321,7 @@ mod tests { use super::{PgTypId, ShortTypeId}; use pgx::{ - pg_guard, - pg_sys::{self, BOOLOID, CHAROID, CIRCLEOID}, + pg_sys::{BOOLOID, CHAROID, CIRCLEOID}, pg_test, }; diff --git a/extension/src/stabilization_info.rs b/extension/src/stabilization_info.rs index 5259d894..be03cb02 100644 --- a/extension/src/stabilization_info.rs +++ b/extension/src/stabilization_info.rs @@ -11,6 +11,183 @@ crate::functions_stabilized_at! { STABLE_FUNCTIONS + "1.15.0" => { + arrow_counter_interpolated_delta(countersummary,counterinterpolateddeltaaccessor), + arrow_counter_interpolated_rate(countersummary,counterinterpolatedrateaccessor), + arrow_time_weighted_average_interpolated_average(timeweightsummary,timeweightinterpolatedaverageaccessor), + counterinterpolateddeltaaccessor_in(cstring), + counterinterpolateddeltaaccessor_out(counterinterpolateddeltaaccessor), + counterinterpolatedrateaccessor_in(cstring), + counterinterpolatedrateaccessor_out(counterinterpolatedrateaccessor), + interpolated_average(timestamp with time zone,interval,timeweightsummary,timeweightsummary), + interpolated_delta(timestamp with time zone,interval,countersummary,countersummary), + interpolated_rate(timestamp with time zone,interval,countersummary,countersummary), + timeweightinterpolatedaverageaccessor_in(cstring), + timeweightinterpolatedaverageaccessor_out(timeweightinterpolatedaverageaccessor), + accessorintegral_in(cstring), + accessorintegral_out(accessorintegral), + arrow_time_weighted_average_integral(timeweightsummary,accessorintegral), + arrow_time_weighted_average_interpolated_integral(timeweightsummary,timeweightinterpolatedintegralaccessor), + integral(text), + integral(timeweightsummary,text), + interpolated_integral(timestamp with time zone,interval,timeweightsummary,timeweightsummary,text), + interpolated_integral(timeweightsummary,timestamp with time zone, interval,timeweightsummary,timeweightsummary,text), + timeweightinterpolatedintegralaccessor_in(cstring), + timeweightinterpolatedintegralaccessor_out(timeweightinterpolatedintegralaccessor), + dead_ranges(heartbeatagg), + downtime(heartbeatagg), + heartbeat_agg(timestamp with time zone,timestamp with time zone,interval,interval), + heartbeat_final(internal), + heartbeat_rollup_trans(internal,heartbeatagg), + heartbeat_trans(internal,timestamp with time zone,timestamp with time zone,interval,interval), + heartbeatagg_in(cstring), + heartbeatagg_out(heartbeatagg), + interpolate(heartbeatagg,heartbeatagg), + interpolated_downtime(heartbeatagg,heartbeatagg), + interpolated_uptime(heartbeatagg,heartbeatagg), + live_at(heartbeatagg,timestamp with time zone), + live_ranges(heartbeatagg), + rollup(heartbeatagg), + uptime(heartbeatagg), + accessordeadranges_in(cstring), + accessordeadranges_out(accessordeadranges), + accessordowntime_in(cstring), + accessordowntime_out(accessordowntime), + accessorliveat_in(cstring), + accessorliveat_out(accessorliveat), + accessorliveranges_in(cstring), + accessorliveranges_out(accessorliveranges), + accessoruptime_in(cstring), + accessoruptime_out(accessoruptime), + arrow_heartbeat_agg_dead_ranges(heartbeatagg,accessordeadranges), + arrow_heartbeat_agg_downtime(heartbeatagg,accessordowntime), + arrow_heartbeat_agg_live_at(heartbeatagg,accessorliveat), + arrow_heartbeat_agg_live_ranges(heartbeatagg,accessorliveranges), + arrow_heartbeat_agg_uptime(heartbeatagg,accessoruptime), + dead_ranges(), + downtime(), + live_at(timestamp with time zone), + live_ranges(), + uptime(), + arrow_heartbeat_agg_interpolate(heartbeatagg,heartbeatinterpolateaccessor), + arrow_heartbeat_agg_interpolated_downtime(heartbeatagg,heartbeatinterpolateddowntimeaccessor), + arrow_heartbeat_agg_interpolated_uptime(heartbeatagg,heartbeatinterpolateduptimeaccessor), + heartbeatinterpolateaccessor_in(cstring), + heartbeatinterpolateaccessor_out(heartbeatinterpolateaccessor), + heartbeatinterpolateddowntimeaccessor_in(cstring), + heartbeatinterpolateddowntimeaccessor_out(heartbeatinterpolateddowntimeaccessor), + heartbeatinterpolateduptimeaccessor_in(cstring), + heartbeatinterpolateduptimeaccessor_out(heartbeatinterpolateduptimeaccessor), + interpolate(heartbeatagg), + interpolated_downtime(heartbeatagg), + interpolated_uptime(heartbeatagg), + duration_in(stateagg,bigint), + duration_in(stateagg,bigint,timestamp with time zone,interval), + duration_in(stateagg,text), + duration_in(stateagg,text,timestamp with time zone,interval), + interpolated_duration_in(stateagg,bigint,timestamp with time zone,interval,stateagg), + interpolated_duration_in(stateagg,text,timestamp with time zone,interval,stateagg), + interpolated_state_periods(stateagg,bigint,timestamp with time zone,interval,stateagg), + interpolated_state_periods(stateagg,text,timestamp with time zone,interval,stateagg), + interpolated_state_timeline(stateagg,timestamp with time zone,interval,stateagg), + interpolated_state_int_timeline(stateagg,timestamp with time zone,interval,stateagg), + into_int_values(stateagg), + into_values(stateagg), + rollup(stateagg), + state_agg(timestamp with time zone,bigint), + state_agg(timestamp with time zone,text), + state_agg_combine_fn_outer(internal,internal), + state_agg_deserialize_fn_outer(bytea,internal), + state_agg_finally_fn_outer(internal), + state_agg_int_trans(internal,timestamp with time zone,bigint), + state_agg_rollup_final(internal), + state_agg_rollup_trans(internal,stateagg), + state_agg_serialize_fn_outer(internal), + state_agg_transition_fn_outer(internal,timestamp with time zone,text), + state_at(stateagg,timestamp with time zone), + state_at_int(stateagg,timestamp with time zone), + state_int_timeline(stateagg), + state_periods(stateagg,bigint), + state_periods(stateagg,text), + state_timeline(stateagg), + stateagg_in(cstring), + stateagg_out(stateagg), + state_agg_rollup_combine(internal,internal), + state_agg_rollup_deserialize(bytea,internal), + state_agg_rollup_serialize(internal), + accessordurationin_in(cstring), + accessordurationin_out(accessordurationin), + accessordurationinint_in(cstring), + accessordurationinint_out(accessordurationinint), + accessordurationinrange_in(cstring), + accessordurationinrange_out(accessordurationinrange), + accessordurationinrangeint_in(cstring), + accessordurationinrangeint_out(accessordurationinrangeint), + accessorinterpolateddurationin_in(cstring), + accessorinterpolateddurationin_out(accessorinterpolateddurationin), + accessorinterpolateddurationinint_in(cstring), + accessorinterpolateddurationinint_out(accessorinterpolateddurationinint), + accessorinterpolatedstateinttimeline_in(cstring), + accessorinterpolatedstateinttimeline_out(accessorinterpolatedstateinttimeline), + accessorinterpolatedstateperiods_in(cstring), + accessorinterpolatedstateperiods_out(accessorinterpolatedstateperiods), + accessorinterpolatedstateperiodsint_in(cstring), + accessorinterpolatedstateperiodsint_out(accessorinterpolatedstateperiodsint), + accessorinterpolatedstatetimeline_in(cstring), + accessorinterpolatedstatetimeline_out(accessorinterpolatedstatetimeline), + accessorintointvalues_in(cstring), + accessorintointvalues_out(accessorintointvalues), + accessorintovalues_in(cstring), + accessorintovalues_out(accessorintovalues), + accessorstateat_in(cstring), + accessorstateat_out(accessorstateat), + accessorstateatint_in(cstring), + accessorstateatint_out(accessorstateatint), + accessorstateinttimeline_in(cstring), + accessorstateinttimeline_out(accessorstateinttimeline), + accessorstateperiods_in(cstring), + accessorstateperiods_out(accessorstateperiods), + accessorstateperiodsint_in(cstring), + accessorstateperiodsint_out(accessorstateperiodsint), + accessorstatetimeline_in(cstring), + accessorstatetimeline_out(accessorstatetimeline), + arrow_state_agg_duration_in_int(stateagg,accessordurationinint), + arrow_state_agg_duration_in_range_int(stateagg,accessordurationinrangeint), + arrow_state_agg_duration_in_range_string(stateagg,accessordurationinrange), + arrow_state_agg_duration_in_string(stateagg,accessordurationin), + arrow_state_agg_interpolated_duration_in_int(stateagg,accessorinterpolateddurationinint), + arrow_state_agg_interpolated_duration_in_string(stateagg,accessorinterpolateddurationin), + arrow_state_agg_interpolated_state_int_timeline(stateagg,accessorinterpolatedstateinttimeline), + arrow_state_agg_interpolated_state_periods_int(stateagg,accessorinterpolatedstateperiodsint), + arrow_state_agg_interpolated_state_periods_string(stateagg,accessorinterpolatedstateperiods), + arrow_state_agg_interpolated_state_timeline(stateagg,accessorinterpolatedstatetimeline), + arrow_state_agg_into_int_values(stateagg,accessorintointvalues), + arrow_state_agg_into_values(stateagg,accessorintovalues), + arrow_state_agg_state_at_int(stateagg,accessorstateatint), + arrow_state_agg_state_at_string(stateagg,accessorstateat), + arrow_state_agg_state_int_timeline(stateagg,accessorstateinttimeline), + arrow_state_agg_state_periods_int(stateagg,accessorstateperiodsint), + arrow_state_agg_state_periods_string(stateagg,accessorstateperiods), + arrow_state_agg_state_timeline(stateagg,accessorstatetimeline), + duration_in(bigint), + duration_in(bigint,timestamp with time zone,interval), + duration_in(text), + duration_in(text,timestamp with time zone,interval), + interpolated_duration_in(bigint,timestamp with time zone,interval,stateagg), + interpolated_duration_in(text,timestamp with time zone,interval,stateagg), + interpolated_state_int_timeline(timestamp with time zone,interval,stateagg), + interpolated_state_periods(bigint,timestamp with time zone,interval,stateagg), + interpolated_state_periods(text,timestamp with time zone,interval,stateagg), + interpolated_state_timeline(timestamp with time zone,interval,stateagg), + into_int_values(), + into_values(), + state_at(timestamp with time zone), + state_at_int(timestamp with time zone), + state_int_timeline(), + state_periods(bigint), + state_periods(text), + state_timeline(), + } "1.14.0" => { interpolated_average(timeweightsummary,timestamp with time zone,interval,timeweightsummary,timeweightsummary), interpolated_delta(countersummary,timestamp with time zone,interval,countersummary,countersummary), @@ -492,6 +669,41 @@ crate::functions_stabilized_at! { crate::types_stabilized_at! { STABLE_TYPES + "1.15.0" => { + counterinterpolateddeltaaccessor, + counterinterpolatedrateaccessor, + timeweightinterpolatedaverageaccessor, + timeweightinterpolatedintegralaccessor, + accessorintegral, + heartbeatagg, + accessordeadranges, + accessordowntime, + accessorliveat, + accessorliveranges, + accessoruptime, + heartbeatinterpolateaccessor, + heartbeatinterpolateddowntimeaccessor, + heartbeatinterpolateduptimeaccessor, + stateagg, + accessordurationin, + accessordurationinint, + accessordurationinrange, + accessordurationinrangeint, + accessorinterpolateddurationin, + accessorinterpolateddurationinint, + accessorinterpolatedstateinttimeline, + accessorinterpolatedstateperiods, + accessorinterpolatedstateperiodsint, + accessorinterpolatedstatetimeline, + accessorintointvalues, + accessorintovalues, + accessorstateat, + accessorstateatint, + accessorstateinttimeline, + accessorstateperiods, + accessorstateperiodsint, + accessorstatetimeline, + } "1.14.0" => { candlestick, accessorclose, @@ -581,6 +793,39 @@ crate::types_stabilized_at! { crate::operators_stabilized_at! { STABLE_OPERATORS + "1.15.0" => { + "->"(countersummary,counterinterpolateddeltaaccessor), + "->"(countersummary,counterinterpolatedrateaccessor), + "->"(timeweightsummary,timeweightinterpolatedaverageaccessor), + "->"(timeweightsummary,timeweightinterpolatedintegralaccessor), + "->"(timeweightsummary,accessorintegral), + "->"(heartbeatagg,accessordeadranges), + "->"(heartbeatagg,accessordowntime), + "->"(heartbeatagg,accessorliveat), + "->"(heartbeatagg,accessorliveranges), + "->"(heartbeatagg,accessoruptime), + "->"(heartbeatagg,heartbeatinterpolateaccessor), + "->"(heartbeatagg,heartbeatinterpolateddowntimeaccessor), + "->"(heartbeatagg,heartbeatinterpolateduptimeaccessor), + "->"(stateagg,accessordurationin), + "->"(stateagg,accessordurationinint), + "->"(stateagg,accessordurationinrange), + "->"(stateagg,accessordurationinrangeint), + "->"(stateagg,accessorinterpolateddurationin), + "->"(stateagg,accessorinterpolateddurationinint), + "->"(stateagg,accessorinterpolatedstateinttimeline), + "->"(stateagg,accessorinterpolatedstateperiods), + "->"(stateagg,accessorinterpolatedstateperiodsint), + "->"(stateagg,accessorinterpolatedstatetimeline), + "->"(stateagg,accessorintointvalues), + "->"(stateagg,accessorintovalues), + "->"(stateagg,accessorstateat), + "->"(stateagg,accessorstateatint), + "->"(stateagg,accessorstateinttimeline), + "->"(stateagg,accessorstateperiods), + "->"(stateagg,accessorstateperiodsint), + "->"(stateagg,accessorstatetimeline), + } "1.14.0" => { "->"(candlestick,accessorclose), "->"(candlestick,accessorclosetime), diff --git a/extension/src/stabilization_tests.rs b/extension/src/stabilization_tests.rs index d9183082..c66a03ab 100644 --- a/extension/src/stabilization_tests.rs +++ b/extension/src/stabilization_tests.rs @@ -12,12 +12,12 @@ mod tests { // Test that any new features are added to the the experimental schema #[pg_test] fn test_schema_qualification() { - Spi::execute(|client| { + Spi::connect(|mut client| { let stable_functions: HashSet = stable_functions(); let stable_types: HashSet = stable_types(); let stable_operators: HashSet = stable_operators(); let unexpected_features: Vec<_> = client - .select( + .update( "SELECT pg_catalog.pg_describe_object(classid, objid, 0) \ FROM pg_catalog.pg_extension e, pg_catalog.pg_depend d \ WHERE e.extname='timescaledb_toolkit' \ @@ -28,8 +28,14 @@ mod tests { None, None, ) + .unwrap() .filter_map(|row| { - let val: String = row.by_ordinal(1).unwrap().value().unwrap(); + let val: String = row + .get_datum_by_ordinal(1) + .unwrap() + .value() + .unwrap() + .unwrap(); if let Some(schema) = val.strip_prefix("schema ") { // the only schemas we should define are diff --git a/extension/src/state_aggregate.rs b/extension/src/state_aggregate.rs index 3ab5ed4f..0a56d735 100644 --- a/extension/src/state_aggregate.rs +++ b/extension/src/state_aggregate.rs @@ -15,6 +15,9 @@ use flat_serialize::*; use flat_serialize_macro::FlatSerializable; use crate::{ + accessors::{ + AccessorIntoIntValues, AccessorIntoValues, AccessorStateIntTimeline, AccessorStateTimeline, + }, flatten, palloc::{Inner, Internal}, pg_type, @@ -22,8 +25,10 @@ use crate::{ ron_inout_funcs, }; -use toolkit_experimental::{CompactStateAgg, StateAgg}; +use toolkit_experimental::CompactStateAgg; +mod accessors; +use accessors::*; pub mod rollup; /// The data of a state. @@ -46,11 +51,18 @@ impl MaterializedState { Self::String(s) => StateEntry::from_existing_str(states, s), } } - fn try_existing_entry(&self, states: &str) -> Option { - Some(match self { - Self::Integer(i) => StateEntry { a: i64::MAX, b: *i }, - Self::String(s) => StateEntry::try_from_existing_str(states, s)?, - }) + + fn into_string(self) -> String { + match self { + Self::String(str) => str, + _ => panic!("MaterializedState::into_string called with non-string"), + } + } + fn into_integer(self) -> i64 { + match self { + Self::Integer(int) => int, + _ => panic!("MaterializedState::into_integer called with non-integer"), + } } } @@ -64,6 +76,7 @@ pub struct StateEntry { b: i64, } impl StateEntry { + #[cfg(test)] // only used by tests fn from_integer(int: i64) -> Self { Self { a: i64::MAX, @@ -104,16 +117,23 @@ impl StateEntry { if self.a == i64::MAX { MaterializedState::Integer(self.b) } else { - MaterializedState::String(states[self.a as usize..self.b as usize].to_string()) + MaterializedState::String( + states + .get(self.a as usize..self.b as usize) + .expect("tried to materialize out-of-bounds state") + .to_string(), + ) } } fn as_str(self, states: &str) -> &str { assert!(self.a != i64::MAX, "Tried to get non-string state"); - &states[self.a as usize..self.b as usize] + states + .get(self.a as usize..self.b as usize) + .expect("tried to stringify out-of-bounds state") } - fn as_integer(self) -> i64 { + fn into_integer(self) -> i64 { assert!(self.a == i64::MAX, "Tried to get non-integer state"); self.b } @@ -141,13 +161,6 @@ pub mod toolkit_experimental { } } - pg_type! { - #[derive(Debug)] - struct StateAgg<'input> { - compact_state_agg: CompactStateAggData<'input>, - } - } - impl CompactStateAgg<'_> { pub(super) fn empty(compact: bool, integer_states: bool) -> Self { unsafe { @@ -235,9 +248,11 @@ pub mod toolkit_experimental { } pub fn get(&self, state: StateEntry) -> Option { - let materialized = state.materialize(self.states_as_str()); + self.get_materialized(&state.materialize(self.states_as_str())) + } + pub(super) fn get_materialized(&self, state: &MaterializedState) -> Option { for record in self.durations.iter() { - if record.state.materialize(self.states_as_str()) == materialized { + if record.state.materialize(self.states_as_str()) == *state { return Some(record.duration); } } @@ -395,36 +410,47 @@ pub mod toolkit_experimental { } } - impl<'input> StateAgg<'input> { - pub fn new(compact_state_agg: CompactStateAgg) -> Self { - unsafe { - flatten!(StateAgg { - compact_state_agg: compact_state_agg.0, - }) - } - } + ron_inout_funcs!(CompactStateAgg); +} +use toolkit_experimental::*; - pub fn as_compact_state_agg(self) -> CompactStateAgg<'input> { - unsafe { self.0.compact_state_agg.flatten() } +pg_type! { + #[derive(Debug)] + struct StateAgg<'input> { + compact_state_agg: CompactStateAggData<'input>, + } +} +impl<'input> StateAgg<'input> { + pub fn new(compact_state_agg: CompactStateAgg) -> Self { + unsafe { + flatten!(StateAgg { + compact_state_agg: compact_state_agg.0, + }) } + } - pub fn assert_int<'a>(&self) { - assert!( - self.0.compact_state_agg.integer_states, - "State must have integer values for this function" - ); - } - pub fn assert_str<'a>(&self) { - assert!( - !self.0.compact_state_agg.integer_states, - "State must have string values for this function" - ); - } + pub fn empty(integer_states: bool) -> Self { + Self::new(CompactStateAgg::empty(false, integer_states)) } - ron_inout_funcs!(CompactStateAgg); - ron_inout_funcs!(StateAgg); + pub fn as_compact_state_agg(self) -> toolkit_experimental::CompactStateAgg<'input> { + unsafe { self.0.compact_state_agg.flatten() } + } + + pub fn assert_int<'a>(&self) { + assert!( + self.0.compact_state_agg.integer_states, + "State must have integer values for this function" + ); + } + pub fn assert_str<'a>(&self) { + assert!( + !self.0.compact_state_agg.integer_states, + "State must have string values for this function" + ); + } } +ron_inout_funcs!(StateAgg); fn state_trans_inner( state: Option, @@ -546,7 +572,7 @@ fn compact_state_agg_int_trans( } #[aggregate] -impl toolkit_experimental::state_agg { +impl state_agg { type State = CompactStateAggTransState; const PARALLEL_SAFE: bool = true; @@ -617,17 +643,17 @@ impl toolkit_experimental::state_agg { } extension_sql!( - "CREATE AGGREGATE toolkit_experimental.state_agg( + "CREATE AGGREGATE state_agg( ts timestamptz, value bigint ) ( stype = internal, - sfunc = toolkit_experimental.state_agg_int_trans, - finalfunc = toolkit_experimental.state_agg_finally_fn_outer, + sfunc = state_agg_int_trans, + finalfunc = state_agg_finally_fn_outer, parallel = safe, - serialfunc = toolkit_experimental.state_agg_serialize_fn_outer, - deserialfunc = toolkit_experimental.state_agg_deserialize_fn_outer, - combinefunc = toolkit_experimental.state_agg_combine_fn_outer + serialfunc = state_agg_serialize_fn_outer, + deserialfunc = state_agg_deserialize_fn_outer, + combinefunc = state_agg_combine_fn_outer );", name = "state_agg_bigint", requires = [ @@ -638,7 +664,7 @@ extension_sql!( state_agg_combine_fn_outer ], ); -#[pg_extern(immutable, parallel_safe, schema = "toolkit_experimental")] +#[pg_extern(immutable, parallel_safe)] fn state_agg_int_trans( __inner: pgx::Internal, ts: TimestampTz, @@ -737,7 +763,7 @@ impl CompactStateAggTransState { fn duration_in_inner<'a>( aggregate: Option>, - state: Option, + state: MaterializedState, range: Option<(i64, Option)>, // start and interval ) -> crate::raw::Interval { let time: i64 = if let Some((start, interval)) = range { @@ -748,13 +774,12 @@ fn duration_in_inner<'a>( i64::MAX }; assert!(end >= start, "End time must be after start time"); - if let (Some(state), Some(agg)) = (state, aggregate) { + if let Some(agg) = aggregate { assert!( !agg.0.compact, "unreachable: interval specified for compact aggregate" ); - let state = state.materialize(agg.states_as_str()); let mut total = 0; for tis in agg.combined_durations.iter() { let tis_start_time = i64::max(tis.start_time, start); @@ -774,7 +799,9 @@ fn duration_in_inner<'a>( 0 } } else { - state.and_then(|state| aggregate?.get(state)).unwrap_or(0) + aggregate + .and_then(|aggregate| aggregate.get_materialized(&state)) + .unwrap_or(0) }; time.into() } @@ -784,9 +811,7 @@ pub fn duration_in<'a>(agg: Option>, state: String) -> crate if let Some(ref agg) = agg { agg.assert_str() }; - let state = agg - .as_ref() - .and_then(|agg| StateEntry::try_from_existing_str(agg.states_as_str(), &state)); + let state = MaterializedState::String(state); duration_in_inner(agg, state, None) } @@ -800,15 +825,10 @@ pub fn duration_in_int<'a>(agg: Option>, state: i64) -> crat if let Some(ref agg) = agg { agg.assert_int() }; - duration_in_inner(agg, Some(StateEntry::from_integer(state)), None) + duration_in_inner(agg, MaterializedState::Integer(state), None) } -#[pg_extern( - immutable, - parallel_safe, - name = "duration_in", - schema = "toolkit_experimental" -)] +#[pg_extern(immutable, parallel_safe, name = "duration_in")] pub fn duration_in_tl<'a>(agg: Option>, state: String) -> crate::raw::Interval { if let Some(ref agg) = agg { agg.assert_str() @@ -816,29 +836,39 @@ pub fn duration_in_tl<'a>(agg: Option>, state: String) -> crate::ra duration_in(agg.map(StateAgg::as_compact_state_agg), state) } -#[pg_extern( - immutable, - parallel_safe, - name = "duration_in", - schema = "toolkit_experimental" -)] +#[pg_extern(immutable, parallel_safe, name = "duration_in")] pub fn duration_in_tl_int<'a>(agg: Option>, state: i64) -> crate::raw::Interval { if let Some(ref agg) = agg { agg.assert_int() }; duration_in_inner( agg.map(StateAgg::as_compact_state_agg), - Some(StateEntry::from_integer(state)), + MaterializedState::Integer(state), None, ) } +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_state_agg_duration_in_string<'a>( + agg: StateAgg<'a>, + accessor: AccessorDurationIn<'a>, +) -> crate::raw::Interval { + let state = MaterializedState::String( + String::from_utf8_lossy(accessor.state_bytes.as_slice()).to_string(), + ); + duration_in_inner(Some(agg.as_compact_state_agg()), state, None) +} +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_state_agg_duration_in_int<'a>( + agg: StateAgg<'a>, + accessor: AccessorDurationInInt<'a>, +) -> crate::raw::Interval { + let state = MaterializedState::Integer(accessor.state); + duration_in_inner(Some(agg.as_compact_state_agg()), state, None) +} -#[pg_extern( - immutable, - parallel_safe, - name = "duration_in", - schema = "toolkit_experimental" -)] +#[pg_extern(immutable, parallel_safe, name = "duration_in")] pub fn duration_in_range<'a>( agg: Option>, state: String, @@ -851,18 +881,14 @@ pub fn duration_in_range<'a>( let agg = agg.map(StateAgg::as_compact_state_agg); let interval = interval.map(|interval| crate::datum_utils::interval_to_ms(&start, &interval)); let start = start.into(); - let state = agg - .as_ref() - .and_then(|agg| StateEntry::try_from_existing_str(agg.states_as_str(), &state)); - duration_in_inner(agg, state, Some((start, interval))) + duration_in_inner( + agg, + MaterializedState::String(state), + Some((start, interval)), + ) } -#[pg_extern( - immutable, - parallel_safe, - name = "duration_in", - schema = "toolkit_experimental" -)] +#[pg_extern(immutable, parallel_safe, name = "duration_in")] pub fn duration_in_range_int<'a>( agg: Option>, state: i64, @@ -876,16 +902,58 @@ pub fn duration_in_range_int<'a>( let start = start.into(); duration_in_inner( agg.map(StateAgg::as_compact_state_agg), - Some(StateEntry::from_integer(state)), + MaterializedState::Integer(state), Some((start, interval)), ) } +/// Used to indicate no interval was specified. The interval cannot be negative anyways, so this +/// value will never be a valid argument. +const NO_INTERVAL_MARKER: i64 = i64::MIN; +fn range_tuple(start: i64, interval: i64) -> (i64, Option) { + ( + start, + if interval == NO_INTERVAL_MARKER { + None + } else { + Some(interval) + }, + ) +} +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_state_agg_duration_in_range_string<'a>( + agg: StateAgg<'a>, + accessor: AccessorDurationInRange<'a>, +) -> crate::raw::Interval { + let state = MaterializedState::String( + String::from_utf8_lossy(accessor.state_bytes.as_slice()).to_string(), + ); + duration_in_inner( + Some(agg.as_compact_state_agg()), + state, + Some(range_tuple(accessor.start, accessor.interval)), + ) +} +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_state_agg_duration_in_range_int<'a>( + agg: StateAgg<'a>, + accessor: AccessorDurationInRangeInt<'a>, +) -> crate::raw::Interval { + let state = MaterializedState::Integer(accessor.state); + duration_in_inner( + Some(agg.as_compact_state_agg()), + state, + Some(range_tuple(accessor.start, accessor.interval)), + ) +} + fn interpolated_duration_in_inner<'a>( aggregate: Option>, - state: Option, - start: TimestampTz, - interval: crate::raw::Interval, + state: MaterializedState, + start: i64, + interval: i64, prev: Option>, ) -> crate::raw::Interval { match aggregate { @@ -893,8 +961,6 @@ fn interpolated_duration_in_inner<'a>( "when interpolating data between grouped data, all groups must contain some data" ), Some(aggregate) => { - let interval = crate::datum_utils::interval_to_ms(&start, &interval); - let start = start.into(); if let Some(ref prev) = prev { assert!( start >= prev.0.last_time, @@ -915,9 +981,7 @@ fn interpolated_duration_in_inner<'a>( Some((start, Some(interval))) }; let new_agg = aggregate.interpolate(start, interval, prev); - let state_entry = - state.and_then(|state| state.try_existing_entry(new_agg.states_as_str())); - duration_in_inner(Some(new_agg), state_entry, range) + duration_in_inner(Some(new_agg), state, range) } } } @@ -932,21 +996,17 @@ pub fn interpolated_duration_in<'a>( if let Some(ref agg) = agg { agg.assert_str() }; + let interval = crate::datum_utils::interval_to_ms(&start, &interval); interpolated_duration_in_inner( agg, - Some(MaterializedState::String(state)), - start, + MaterializedState::String(state), + start.into(), interval, prev, ) } -#[pg_extern( - immutable, - parallel_safe, - name = "interpolated_duration_in", - schema = "toolkit_experimental" -)] +#[pg_extern(immutable, parallel_safe, name = "interpolated_duration_in")] pub fn interpolated_duration_in_tl<'a>( agg: Option>, state: String, @@ -982,21 +1042,17 @@ pub fn interpolated_duration_in_int<'a>( if let Some(ref agg) = agg { agg.assert_int() }; + let interval = crate::datum_utils::interval_to_ms(&start, &interval); interpolated_duration_in_inner( agg, - Some(MaterializedState::Integer(state)), - start, + MaterializedState::Integer(state), + start.into(), interval, prev, ) } -#[pg_extern( - immutable, - parallel_safe, - name = "interpolated_duration_in", - schema = "toolkit_experimental" -)] +#[pg_extern(immutable, parallel_safe, name = "interpolated_duration_in")] pub fn interpolated_duration_in_tl_int<'a>( agg: Option>, state: i64, @@ -1015,6 +1071,46 @@ pub fn interpolated_duration_in_tl_int<'a>( prev.map(StateAgg::as_compact_state_agg), ) } +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_state_agg_interpolated_duration_in_string<'a>( + agg: Option>, + accessor: AccessorInterpolatedDurationIn<'a>, +) -> crate::raw::Interval { + let state = MaterializedState::String( + String::from_utf8_lossy(accessor.state_bytes.as_slice()).to_string(), + ); + interpolated_duration_in_inner( + agg.map(StateAgg::as_compact_state_agg), + state, + accessor.start, + accessor.interval, + if accessor.prev_present { + Some(unsafe { accessor.prev.flatten() }.as_compact_state_agg()) + } else { + None + }, + ) +} +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_state_agg_interpolated_duration_in_int<'a>( + agg: Option>, + accessor: AccessorInterpolatedDurationInInt<'a>, +) -> crate::raw::Interval { + let state = MaterializedState::Integer(accessor.state); + interpolated_duration_in_inner( + agg.map(StateAgg::as_compact_state_agg), + state, + accessor.start, + accessor.interval, + if accessor.prev_present { + Some(unsafe { accessor.prev.flatten() }.as_compact_state_agg()) + } else { + None + }, + ) +} fn duration_in_bad_args_inner() -> ! { panic!("The start and interval parameters cannot be used for duration_in with a compact state aggregate") @@ -1085,17 +1181,12 @@ pub fn into_int_values<'a>( agg.durations .clone() .into_iter() - .map(move |record| (record.state.as_integer(), record.duration.into())) + .map(move |record| (record.state.into_integer(), record.duration.into())) .collect::>() .into_iter(), // make map panic now instead of at iteration time ) } -#[pg_extern( - immutable, - parallel_safe, - name = "into_values", - schema = "toolkit_experimental" -)] +#[pg_extern(immutable, parallel_safe, name = "into_values")] pub fn into_values_tl<'a>( agg: StateAgg<'a>, ) -> TableIterator< @@ -1108,12 +1199,7 @@ pub fn into_values_tl<'a>( agg.assert_str(); into_values(agg.as_compact_state_agg()) } -#[pg_extern( - immutable, - parallel_safe, - name = "into_int_values", - schema = "toolkit_experimental" -)] +#[pg_extern(immutable, parallel_safe, name = "into_int_values")] pub fn into_values_tl_int<'a>( agg: StateAgg<'a>, ) -> TableIterator< @@ -1126,6 +1212,34 @@ pub fn into_values_tl_int<'a>( agg.assert_int(); into_int_values(agg.as_compact_state_agg()) } +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_state_agg_into_values<'a>( + agg: StateAgg<'a>, + _accessor: AccessorIntoValues<'a>, +) -> TableIterator< + 'a, + ( + pgx::name!(state, String), + pgx::name!(duration, crate::raw::Interval), + ), +> { + into_values_tl(agg) +} +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_state_agg_into_int_values<'a>( + agg: StateAgg<'a>, + _accessor: AccessorIntoIntValues<'a>, +) -> TableIterator< + 'a, + ( + pgx::name!(state, i64), + pgx::name!(duration, crate::raw::Interval), + ), +> { + into_values_tl_int(agg) +} fn state_timeline_inner<'a>( agg: CompactStateAgg<'a>, @@ -1175,7 +1289,7 @@ fn state_int_timeline_inner<'a>( .into_iter() .map(move |record| { ( - record.state.as_integer(), + record.state.into_integer(), TimestampTz::from(record.start_time), TimestampTz::from(record.end_time), ) @@ -1185,7 +1299,7 @@ fn state_int_timeline_inner<'a>( ) } -#[pg_extern(immutable, parallel_safe, schema = "toolkit_experimental")] +#[pg_extern(immutable, parallel_safe)] pub fn state_timeline<'a>( agg: StateAgg<'a>, ) -> TableIterator< @@ -1199,7 +1313,7 @@ pub fn state_timeline<'a>( agg.assert_str(); state_timeline_inner(agg.as_compact_state_agg()) } -#[pg_extern(immutable, parallel_safe, schema = "toolkit_experimental")] +#[pg_extern(immutable, parallel_safe)] pub fn state_int_timeline<'a>( agg: StateAgg<'a>, ) -> TableIterator< @@ -1214,11 +1328,41 @@ pub fn state_int_timeline<'a>( state_int_timeline_inner(agg.as_compact_state_agg()) } -#[pg_extern(immutable, parallel_safe, schema = "toolkit_experimental")] -pub fn interpolated_state_timeline<'a>( +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_state_agg_state_timeline<'a>( + agg: StateAgg<'a>, + _accessor: AccessorStateTimeline<'a>, +) -> TableIterator< + 'a, + ( + pgx::name!(state, String), + pgx::name!(start_time, TimestampTz), + pgx::name!(end_time, TimestampTz), + ), +> { + state_timeline(agg) +} +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_state_agg_state_int_timeline<'a>( + agg: StateAgg<'a>, + _accessor: AccessorStateIntTimeline<'a>, +) -> TableIterator< + 'a, + ( + pgx::name!(state, i64), + pgx::name!(start_time, TimestampTz), + pgx::name!(end_time, TimestampTz), + ), +> { + state_int_timeline(agg) +} + +fn interpolated_state_timeline_inner<'a>( agg: Option>, - start: TimestampTz, - interval: crate::raw::Interval, + start: i64, + interval: i64, prev: Option>, ) -> TableIterator< 'a, @@ -1235,25 +1379,21 @@ pub fn interpolated_state_timeline<'a>( None => pgx::error!( "when interpolating data between grouped data, all groups must contain some data" ), - Some(agg) => { - let interval = crate::datum_utils::interval_to_ms(&start, &interval); - TableIterator::new( - state_timeline_inner(agg.as_compact_state_agg().interpolate( - start.into(), - interval, - prev.map(StateAgg::as_compact_state_agg), - )) - .collect::>() - .into_iter(), - ) - } + Some(agg) => TableIterator::new( + state_timeline_inner(agg.as_compact_state_agg().interpolate( + start, + interval, + prev.map(StateAgg::as_compact_state_agg), + )) + .collect::>() + .into_iter(), + ), } } -#[pg_extern(immutable, parallel_safe, schema = "toolkit_experimental")] -pub fn interpolated_int_state_timeline<'a>( +fn interpolated_state_int_timeline_inner<'a>( agg: Option>, - start: TimestampTz, - interval: crate::raw::Interval, + start: i64, + interval: i64, prev: Option>, ) -> TableIterator< 'a, @@ -1270,77 +1410,151 @@ pub fn interpolated_int_state_timeline<'a>( None => pgx::error!( "when interpolating data between grouped data, all groups must contain some data" ), - Some(agg) => { - let interval = crate::datum_utils::interval_to_ms(&start, &interval); - TableIterator::new( - state_int_timeline_inner(agg.as_compact_state_agg().interpolate( - start.into(), - interval, - prev.map(StateAgg::as_compact_state_agg), - )) - .collect::>() - .into_iter(), - ) - } + Some(agg) => TableIterator::new( + state_int_timeline_inner(agg.as_compact_state_agg().interpolate( + start, + interval, + prev.map(StateAgg::as_compact_state_agg), + )) + .collect::>() + .into_iter(), + ), } } - -fn state_periods_inner<'a>( - state: MaterializedState, - agg: CompactStateAgg<'a>, +#[pg_extern(immutable, parallel_safe)] +pub fn interpolated_state_timeline<'a>( + agg: Option>, + start: TimestampTz, + interval: crate::raw::Interval, + prev: Option>, ) -> TableIterator< 'a, ( + pgx::name!(state, String), pgx::name!(start_time, TimestampTz), pgx::name!(end_time, TimestampTz), ), > { - assert!( - !agg.compact, - "state_periods can only be called on a compact_state_agg built from state_agg" - ); - let states: String = agg.states_as_str().to_owned(); - TableIterator::new( - agg.combined_durations - .clone() - .into_iter() - .filter_map(move |record| { - if record.state.materialize(&states) == state { - Some(( - TimestampTz::from(record.start_time), - TimestampTz::from(record.end_time), - )) - } else { - None - } - }), - ) + let interval = crate::datum_utils::interval_to_ms(&start, &interval); + interpolated_state_timeline_inner(agg, start.into(), interval, prev) } - -#[pg_extern(immutable, parallel_safe, schema = "toolkit_experimental")] -pub fn state_periods<'a>( - agg: StateAgg<'a>, - state: String, +#[pg_extern(immutable, parallel_safe)] +pub fn interpolated_state_int_timeline<'a>( + agg: Option>, + start: TimestampTz, + interval: crate::raw::Interval, + prev: Option>, ) -> TableIterator< 'a, ( + pgx::name!(state, i64), pgx::name!(start_time, TimestampTz), pgx::name!(end_time, TimestampTz), ), > { - agg.assert_str(); - let agg = agg.as_compact_state_agg(); - state_periods_inner(MaterializedState::String(state), agg) + let interval = crate::datum_utils::interval_to_ms(&start, &interval); + interpolated_state_int_timeline_inner(agg, start.into(), interval, prev) } -#[pg_extern( - immutable, - parallel_safe, - schema = "toolkit_experimental", - name = "state_periods" -)] -pub fn state_int_periods<'a>( - agg: StateAgg<'a>, - state: i64, +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_state_agg_interpolated_state_timeline<'a>( + agg: Option>, + accessor: AccessorInterpolatedStateTimeline<'a>, +) -> TableIterator< + 'a, + ( + pgx::name!(state, String), + pgx::name!(start_time, TimestampTz), + pgx::name!(end_time, TimestampTz), + ), +> { + interpolated_state_timeline_inner( + agg, + accessor.start, + accessor.interval, + if accessor.prev_present { + Some(unsafe { accessor.prev.flatten() }) + } else { + None + }, + ) +} +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_state_agg_interpolated_state_int_timeline<'a>( + agg: Option>, + accessor: AccessorInterpolatedStateIntTimeline<'a>, +) -> TableIterator< + 'a, + ( + pgx::name!(state, i64), + pgx::name!(start_time, TimestampTz), + pgx::name!(end_time, TimestampTz), + ), +> { + interpolated_state_int_timeline_inner( + agg, + accessor.start, + accessor.interval, + if accessor.prev_present { + Some(unsafe { accessor.prev.flatten() }) + } else { + None + }, + ) +} + +fn state_periods_inner<'a>( + agg: CompactStateAgg<'a>, + state: MaterializedState, +) -> TableIterator< + 'a, + ( + pgx::name!(start_time, TimestampTz), + pgx::name!(end_time, TimestampTz), + ), +> { + assert!( + !agg.compact, + "state_periods can only be called on a compact_state_agg built from state_agg" + ); + let states: String = agg.states_as_str().to_owned(); + TableIterator::new( + agg.combined_durations + .clone() + .into_iter() + .filter_map(move |record| { + if record.state.materialize(&states) == state { + Some(( + TimestampTz::from(record.start_time), + TimestampTz::from(record.end_time), + )) + } else { + None + } + }), + ) +} + +#[pg_extern(immutable, parallel_safe)] +pub fn state_periods<'a>( + agg: StateAgg<'a>, + state: String, +) -> TableIterator< + 'a, + ( + pgx::name!(start_time, TimestampTz), + pgx::name!(end_time, TimestampTz), + ), +> { + agg.assert_str(); + let agg = agg.as_compact_state_agg(); + state_periods_inner(agg, MaterializedState::String(state)) +} +#[pg_extern(immutable, parallel_safe, name = "state_periods")] +pub fn state_int_periods<'a>( + agg: StateAgg<'a>, + state: i64, ) -> TableIterator< 'a, ( @@ -1350,17 +1564,50 @@ pub fn state_int_periods<'a>( > { agg.assert_int(); state_periods_inner( - MaterializedState::Integer(state), agg.as_compact_state_agg(), + MaterializedState::Integer(state), ) } +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_state_agg_state_periods_string<'a>( + agg: StateAgg<'a>, + accessor: AccessorStatePeriods<'a>, +) -> TableIterator< + 'a, + ( + pgx::name!(start_time, TimestampTz), + pgx::name!(end_time, TimestampTz), + ), +> { + let state = MaterializedState::String( + String::from_utf8_lossy(accessor.state_bytes.as_slice()).to_string(), + ); + state_periods_inner(agg.as_compact_state_agg(), state) +} +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_state_agg_state_periods_int<'a>( + agg: StateAgg<'a>, + accessor: AccessorStatePeriodsInt<'a>, +) -> TableIterator< + 'a, + ( + pgx::name!(start_time, TimestampTz), + pgx::name!(end_time, TimestampTz), + ), +> { + let state = MaterializedState::Integer(accessor.state); + state_periods_inner(agg.as_compact_state_agg(), state) +} + fn interpolated_state_periods_inner<'a>( - aggregate: Option>, + aggregate: Option>, state: MaterializedState, - start: TimestampTz, - interval: crate::raw::Interval, - prev: Option>, + start: i64, + interval: i64, + prev: Option>, ) -> TableIterator< 'a, ( @@ -1372,24 +1619,14 @@ fn interpolated_state_periods_inner<'a>( None => pgx::error!( "when interpolating data between grouped data, all groups must contain some data" ), - Some(aggregate) => { - let interval = crate::datum_utils::interval_to_ms(&start, &interval); - TableIterator::new( - state_periods_inner( - state, - aggregate.as_compact_state_agg().interpolate( - start.into(), - interval, - prev.map(StateAgg::as_compact_state_agg), - ), - ) + Some(aggregate) => TableIterator::new( + state_periods_inner(aggregate.interpolate(start, interval, prev), state) .collect::>() .into_iter(), - ) - } + ), } } -#[pg_extern(immutable, parallel_safe, schema = "toolkit_experimental")] +#[pg_extern(immutable, parallel_safe)] pub fn interpolated_state_periods<'a>( agg: Option>, state: String, @@ -1406,14 +1643,16 @@ pub fn interpolated_state_periods<'a>( if let Some(ref agg) = agg { agg.assert_str() }; - interpolated_state_periods_inner(agg, MaterializedState::String(state), start, interval, prev) + let interval = crate::datum_utils::interval_to_ms(&start, &interval); + interpolated_state_periods_inner( + agg.map(StateAgg::as_compact_state_agg), + MaterializedState::String(state), + start.into(), + interval, + prev.map(StateAgg::as_compact_state_agg), + ) } -#[pg_extern( - immutable, - parallel_safe, - schema = "toolkit_experimental", - name = "interpolated_state_periods" -)] +#[pg_extern(immutable, parallel_safe, name = "interpolated_state_periods")] pub fn interpolated_state_periods_int<'a>( agg: Option>, state: i64, @@ -1430,14 +1669,115 @@ pub fn interpolated_state_periods_int<'a>( if let Some(ref agg) = agg { agg.assert_int() }; + let interval = crate::datum_utils::interval_to_ms(&start, &interval); interpolated_state_periods_inner( - agg, + agg.map(StateAgg::as_compact_state_agg), MaterializedState::Integer(state), - start, + start.into(), interval, - prev, + prev.map(StateAgg::as_compact_state_agg), ) } +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_state_agg_interpolated_state_periods_string<'a>( + agg: Option>, + accessor: AccessorInterpolatedStatePeriods<'a>, +) -> TableIterator< + 'a, + ( + pgx::name!(start_time, TimestampTz), + pgx::name!(end_time, TimestampTz), + ), +> { + let state = MaterializedState::String( + String::from_utf8_lossy(accessor.state_bytes.as_slice()).to_string(), + ); + interpolated_state_periods_inner( + agg.map(StateAgg::as_compact_state_agg), + state, + accessor.start, + accessor.interval, + if accessor.prev_present { + Some(unsafe { accessor.prev.flatten() }.as_compact_state_agg()) + } else { + None + }, + ) +} +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_state_agg_interpolated_state_periods_int<'a>( + agg: Option>, + accessor: AccessorInterpolatedStatePeriodsInt<'a>, +) -> TableIterator< + 'a, + ( + pgx::name!(start_time, TimestampTz), + pgx::name!(end_time, TimestampTz), + ), +> { + let state = MaterializedState::Integer(accessor.state); + interpolated_state_periods_inner( + agg.map(StateAgg::as_compact_state_agg), + state, + accessor.start, + accessor.interval, + if accessor.prev_present { + Some(unsafe { accessor.prev.flatten() }.as_compact_state_agg()) + } else { + None + }, + ) +} + +fn state_at_inner<'a>(agg: StateAgg<'a>, point: i64) -> Option { + let agg = agg.as_compact_state_agg(); + let point: i64 = point.into(); + if agg.combined_durations.is_empty() { + return None; + } + + // binary search to find the first time at or after the start time + let slice = agg.combined_durations.as_slice(); + let idx = match slice.binary_search_by(|tis| tis.start_time.cmp(&point)) { + Ok(idx) => idx, + Err(idx) => idx.checked_sub(1)?, // return NULL if before first item + }; + let tis = slice.get(idx).expect("binary search index out-of-bounds"); + + Some(tis.state.materialize(agg.states_as_str())) +} + +#[pg_extern(immutable, parallel_safe, name = "state_at")] +fn state_at<'a>(agg: StateAgg<'a>, point: TimestampTz) -> Option { + agg.assert_str(); + state_at_inner(agg, point.into()).map(MaterializedState::into_string) +} + +#[pg_extern(immutable, parallel_safe, name = "state_at_int")] +fn state_at_int<'a>(agg: StateAgg<'a>, point: TimestampTz) -> Option { + agg.assert_int(); + state_at_inner(agg, point.into()).map(MaterializedState::into_integer) +} + +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_state_agg_state_at_string<'a>( + agg: StateAgg<'a>, + accessor: AccessorStateAt<'a>, +) -> Option { + state_at_inner(agg, accessor.time).map(MaterializedState::into_string) +} + +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_state_agg_state_at_int<'a>( + agg: StateAgg<'a>, + accessor: AccessorStateAtInt<'a>, +) -> Option { + state_at_inner(agg, accessor.time).map(MaterializedState::into_integer) +} #[derive(Clone, Debug, Deserialize, Eq, FlatSerializable, PartialEq, Serialize)] #[repr(C)] @@ -1512,18 +1852,22 @@ mod tests { macro_rules! select_one { ($client:expr, $stmt:expr, $type:ty) => { $client - .select($stmt, None, None) + .update($stmt, None, None) + .unwrap() .first() .get_one::<$type>() .unwrap() + .unwrap() }; } #[pg_test] #[should_panic = "The start and interval parameters cannot be used for duration_in with"] fn duration_in_misuse_error() { - Spi::execute(|client| { - client.select("CREATE TABLE test(ts timestamptz, state TEXT)", None, None); + Spi::connect(|mut client| { + client + .update("CREATE TABLE test(ts timestamptz, state TEXT)", None, None) + .unwrap(); assert_eq!( "365 days 00:02:00", select_one!( @@ -1537,16 +1881,20 @@ mod tests { #[pg_test] fn one_state_one_change() { - Spi::execute(|client| { - client.select("CREATE TABLE test(ts timestamptz, state TEXT)", None, None); - client.select( - r#"INSERT INTO test VALUES + Spi::connect(|mut client| { + client + .update("CREATE TABLE test(ts timestamptz, state TEXT)", None, None) + .unwrap(); + client + .update( + r#"INSERT INTO test VALUES ('2020-01-01 00:00:00+00', 'one'), ('2020-12-31 00:02:00+00', 'end') "#, - None, - None, - ); + None, + None, + ) + .unwrap(); assert_eq!( "365 days 00:02:00", select_one!( @@ -1559,7 +1907,7 @@ mod tests { "365 days 00:02:00", select_one!( client, - "SELECT toolkit_experimental.duration_in(toolkit_experimental.state_agg(ts, state), 'one')::TEXT FROM test", + "SELECT duration_in(state_agg(ts, state), 'one')::TEXT FROM test", &str ) ); @@ -1568,17 +1916,21 @@ mod tests { #[pg_test] fn two_states_two_changes() { - Spi::execute(|client| { - client.select("CREATE TABLE test(ts timestamptz, state TEXT)", None, None); - client.select( - r#"INSERT INTO test VALUES + Spi::connect(|mut client| { + client + .update("CREATE TABLE test(ts timestamptz, state TEXT)", None, None) + .unwrap(); + client + .update( + r#"INSERT INTO test VALUES ('2020-01-01 00:00:00+00', 'one'), ('2020-01-01 00:01:00+00', 'two'), ('2020-12-31 00:02:00+00', 'end') "#, - None, - None, - ); + None, + None, + ) + .unwrap(); assert_eq!( "00:01:00", @@ -1601,18 +1953,22 @@ mod tests { #[pg_test] fn two_states_three_changes() { - Spi::execute(|client| { - client.select("CREATE TABLE test(ts timestamptz, state TEXT)", None, None); - client.select( - r#"INSERT INTO test VALUES + Spi::connect(|mut client| { + client + .update("CREATE TABLE test(ts timestamptz, state TEXT)", None, None) + .unwrap(); + client + .update( + r#"INSERT INTO test VALUES ('2020-01-01 00:00:00+00', 'one'), ('2020-01-01 00:01:00+00', 'two'), ('2020-01-01 00:02:00+00', 'one'), ('2020-12-31 00:02:00+00', 'end') "#, - None, - None, - ); + None, + None, + ) + .unwrap(); assert_eq!( "365 days 00:01:00", @@ -1635,7 +1991,7 @@ mod tests { "365 days 00:01:00", select_one!( client, - "SELECT toolkit_experimental.duration_in(toolkit_experimental.state_agg(ts, state), 'one')::TEXT FROM test", + "SELECT duration_in(state_agg(ts, state), 'one')::TEXT FROM test", &str ) ); @@ -1643,7 +1999,7 @@ mod tests { "00:01:00", select_one!( client, - "SELECT toolkit_experimental.duration_in(toolkit_experimental.state_agg(ts, state), 'two')::TEXT FROM test", + "SELECT duration_in(state_agg(ts, state), 'two')::TEXT FROM test", &str ) ); @@ -1652,18 +2008,22 @@ mod tests { #[pg_test] fn out_of_order_times() { - Spi::execute(|client| { - client.select("CREATE TABLE test(ts timestamptz, state TEXT)", None, None); - client.select( - r#"INSERT INTO test VALUES + Spi::connect(|mut client| { + client + .update("CREATE TABLE test(ts timestamptz, state TEXT)", None, None) + .unwrap(); + client + .update( + r#"INSERT INTO test VALUES ('2020-01-01 00:00:00+00', 'one'), ('2020-01-01 00:02:00+00', 'one'), ('2020-01-01 00:01:00+00', 'two'), ('2020-12-31 00:02:00+00', 'end') "#, - None, - None, - ); + None, + None, + ) + .unwrap(); assert_eq!( "365 days 00:01:00", @@ -1688,18 +2048,22 @@ mod tests { fn same_state_twice() { // TODO Do we care? Could be that states are recorded not only when they change but // also at regular intervals even when they don't? - Spi::execute(|client| { - client.select("CREATE TABLE test(ts timestamptz, state TEXT)", None, None); - client.select( - r#"INSERT INTO test VALUES + Spi::connect(|mut client| { + client + .update("CREATE TABLE test(ts timestamptz, state TEXT)", None, None) + .unwrap(); + client + .update( + r#"INSERT INTO test VALUES ('2020-01-01 00:00:00+00', 'one'), ('2020-01-01 00:01:00+00', 'one'), ('2020-01-01 00:02:00+00', 'two'), ('2020-12-31 00:02:00+00', 'end') "#, - None, - None, - ); + None, + None, + ) + .unwrap(); assert_eq!( "00:02:00", select_one!( @@ -1721,17 +2085,21 @@ mod tests { #[pg_test] fn duration_in_two_states_two_changes() { - Spi::execute(|client| { - client.select("CREATE TABLE test(ts timestamptz, state TEXT)", None, None); - client.select( - r#"INSERT INTO test VALUES + Spi::connect(|mut client| { + client + .update("CREATE TABLE test(ts timestamptz, state TEXT)", None, None) + .unwrap(); + client + .update( + r#"INSERT INTO test VALUES ('2020-01-01 00:00:00+00', 'one'), ('2020-01-01 00:01:00+00', 'two'), ('2020-12-31 00:02:00+00', 'end') "#, - None, - None, - ); + None, + None, + ) + .unwrap(); assert_eq!( "00:01:00", select_one!( @@ -1753,17 +2121,21 @@ mod tests { #[pg_test] fn same_state_twice_last() { - Spi::execute(|client| { - client.select("CREATE TABLE test(ts timestamptz, state TEXT)", None, None); - client.select( - r#"INSERT INTO test VALUES + Spi::connect(|mut client| { + client + .update("CREATE TABLE test(ts timestamptz, state TEXT)", None, None) + .unwrap(); + client + .update( + r#"INSERT INTO test VALUES ('2020-01-01 00:00:00+00', 'one'), ('2020-01-01 00:01:00+00', 'two'), ('2020-01-01 00:02:00+00', 'two') "#, - None, - None, - ); + None, + None, + ) + .unwrap(); assert_eq!( "00:01:00", select_one!( @@ -1778,9 +2150,11 @@ mod tests { #[pg_test] fn combine_using_muchos_data() { compact_state_agg::counters::reset(); - Spi::execute(|client| { - client.select("CREATE TABLE test(ts timestamptz, state TEXT)", None, None); - client.select( + Spi::connect(|mut client| { + client + .update("CREATE TABLE test(ts timestamptz, state TEXT)", None, None) + .unwrap(); + client.update( r#" insert into test values ('2020-01-01 00:00:00+00', 'one'); insert into test select '2020-01-02 UTC'::timestamptz + make_interval(days=>v), 'two' from generate_series(1,300000) v; @@ -1789,7 +2163,7 @@ insert into test select '2020-01-02 UTC'::timestamptz + make_interval(days=>v), "#, None, None, - ); + ).unwrap(); assert_eq!( "2 days", select_one!( @@ -1813,16 +2187,20 @@ insert into test select '2020-01-02 UTC'::timestamptz + make_interval(days=>v), #[allow(dead_code)] fn combine_using_settings() { compact_state_agg::counters::reset(); - Spi::execute(|client| { - client.select("CREATE TABLE test(ts timestamptz, state TEXT)", None, None); - client.select( - r#"INSERT INTO test VALUES + Spi::connect(|mut client| { + client + .update("CREATE TABLE test(ts timestamptz, state TEXT)", None, None) + .unwrap(); + client + .update( + r#"INSERT INTO test VALUES ('2020-01-01 00:00:00+00', 'one'), ('2020-01-03 00:00:00+00', 'two') "#, - None, - None, - ); + None, + None, + ) + .unwrap(); assert_eq!( "2 days", select_one!( @@ -1854,19 +2232,23 @@ SELECT toolkit_experimental.duration_in('one', toolkit_experimental.compact_stat // the sample query from the ticket #[pg_test] fn sample_query() { - Spi::execute(|client| { - client.select("CREATE TABLE test(ts timestamptz, state TEXT)", None, None); - client.select( - r#"INSERT INTO test VALUES + Spi::connect(|mut client| { + client + .update("CREATE TABLE test(ts timestamptz, state TEXT)", None, None) + .unwrap(); + client + .update( + r#"INSERT INTO test VALUES ('2020-01-01 00:00:00+00', 'START'), ('2020-01-01 00:01:00+00', 'ERROR'), ('2020-01-01 00:02:00+00', 'STOPPED')"#, - None, - None, - ); + None, + None, + ) + .unwrap(); assert_eq!( client - .select( + .update( r#"SELECT toolkit_experimental.duration_in(states, 'ERROR')::TEXT as error, toolkit_experimental.duration_in(states, 'START')::TEXT as start, toolkit_experimental.duration_in(states, 'STOPPED')::TEXT as stopped @@ -1874,22 +2256,24 @@ SELECT toolkit_experimental.duration_in('one', toolkit_experimental.compact_stat None, None, ) - .first() - .get_three::<&str, &str, &str>(), + .unwrap().first() + .get_three::<&str, &str, &str>().unwrap(), (Some("00:01:00"), Some("00:01:00"), Some("00:00:00")) ); assert_eq!( client - .select( - r#"SELECT toolkit_experimental.duration_in(states, 'ERROR')::TEXT as error, - toolkit_experimental.duration_in(states, 'START')::TEXT as start, - toolkit_experimental.duration_in(states, 'STOPPED')::TEXT as stopped - FROM (SELECT toolkit_experimental.state_agg(ts, state) as states FROM test) as foo"#, + .update( + r#"SELECT duration_in(states, 'ERROR')::TEXT as error, + duration_in(states, 'START')::TEXT as start, + duration_in(states, 'STOPPED')::TEXT as stopped + FROM (SELECT state_agg(ts, state) as states FROM test) as foo"#, None, None, ) + .unwrap() .first() - .get_three::<&str, &str, &str>(), + .get_three::<&str, &str, &str>() + .unwrap(), (Some("00:01:00"), Some("00:01:00"), Some("00:00:00")) ); }) @@ -1897,16 +2281,19 @@ SELECT toolkit_experimental.duration_in('one', toolkit_experimental.compact_stat #[pg_test] fn interpolated_duration() { - Spi::execute(|client| { - client.select( - "SET TIME ZONE 'UTC'; + Spi::connect(|mut client| { + client + .update( + "SET TIME ZONE 'UTC'; CREATE TABLE inttest(time TIMESTAMPTZ, state TEXT, bucket INT); CREATE TABLE inttest2(time TIMESTAMPTZ, state BIGINT, bucket INT);", - None, - None, - ); - client.select( - r#"INSERT INTO inttest VALUES + None, + None, + ) + .unwrap(); + client + .update( + r#"INSERT INTO inttest VALUES ('2020-1-1 10:00'::timestamptz, 'one', 1), ('2020-1-1 12:00'::timestamptz, 'two', 1), ('2020-1-1 16:00'::timestamptz, 'three', 1), @@ -1926,12 +2313,13 @@ SELECT toolkit_experimental.duration_in('one', toolkit_experimental.compact_stat ('2020-1-3 10:00'::timestamptz, 10001, 3), ('2020-1-3 12:00'::timestamptz, 10002, 3), ('2020-1-3 16:00'::timestamptz, 10003, 3);"#, - None, - None, - ); + None, + None, + ) + .unwrap(); // Interpolate time spent in state "three" each day - let mut durations = client.select( + let mut durations = client.update( r#"SELECT toolkit_experimental.interpolated_duration_in( agg, @@ -1946,42 +2334,94 @@ SELECT toolkit_experimental.duration_in('one', toolkit_experimental.compact_stat ORDER BY bucket"#, None, None, - ); + ).unwrap(); // Day 1, in "three" from "16:00" to end of day - assert_eq!(durations.next().unwrap()[1].value(), Some("08:00:00")); + assert_eq!( + durations.next().unwrap()[1].value().unwrap(), + Some("08:00:00") + ); // Day 2, in "three" from start of day to "2:00" and "20:00" to end of day - assert_eq!(durations.next().unwrap()[1].value(), Some("06:00:00")); + assert_eq!( + durations.next().unwrap()[1].value().unwrap(), + Some("06:00:00") + ); // Day 3, in "three" from start of day to end - assert_eq!(durations.next().unwrap()[1].value(), Some("18:00:00")); + assert_eq!( + durations.next().unwrap()[1].value().unwrap(), + Some("18:00:00") + ); assert!(durations.next().is_none()); - let mut durations = client.select( + let mut durations = client.update( r#"SELECT - toolkit_experimental.interpolated_duration_in( + interpolated_duration_in( agg, 'three', '2019-12-31 0:00'::timestamptz + (bucket * '1 day'::interval), '1 day'::interval, LAG(agg) OVER (ORDER BY bucket) )::TEXT FROM ( - SELECT bucket, toolkit_experimental.state_agg(time, state) as agg + SELECT bucket, state_agg(time, state) as agg FROM inttest GROUP BY bucket ) s ORDER BY bucket"#, None, None, + ).unwrap(); + + // Day 1, in "three" from "16:00" to end of day + assert_eq!( + durations.next().unwrap()[1].value().unwrap(), + Some("08:00:00") ); + // Day 2, in "three" from start of day to "2:00" and "20:00" to end of day + assert_eq!( + durations.next().unwrap()[1].value().unwrap(), + Some("06:00:00") + ); + // Day 3, in "three" from start of day to end + assert_eq!( + durations.next().unwrap()[1].value().unwrap(), + Some("18:00:00") + ); + assert!(durations.next().is_none()); + + let mut durations = client.update( + r#"SELECT + interpolated_duration_in( + agg, + 10003, + '2019-12-31 0:00'::timestamptz + (bucket * '1 day'::interval), '1 day'::interval, + LAG(agg) OVER (ORDER BY bucket) + )::TEXT FROM ( + SELECT bucket, state_agg(time, state) as agg + FROM inttest2 + GROUP BY bucket + ) s + ORDER BY bucket"#, + None, + None, + ).unwrap(); // Day 1, in "three" from "16:00" to end of day - assert_eq!(durations.next().unwrap()[1].value(), Some("08:00:00")); + assert_eq!( + durations.next().unwrap()[1].value().unwrap(), + Some("08:00:00") + ); // Day 2, in "three" from start of day to "2:00" and "20:00" to end of day - assert_eq!(durations.next().unwrap()[1].value(), Some("06:00:00")); + assert_eq!( + durations.next().unwrap()[1].value().unwrap(), + Some("06:00:00") + ); // Day 3, in "three" from start of day to end - assert_eq!(durations.next().unwrap()[1].value(), Some("18:00:00")); + assert_eq!( + durations.next().unwrap()[1].value().unwrap(), + Some("18:00:00") + ); assert!(durations.next().is_none()); - let mut durations = client.select( + let mut durations = client.update( r#"SELECT toolkit_experimental.interpolated_duration_in( agg, @@ -1996,14 +2436,23 @@ SELECT toolkit_experimental.duration_in('one', toolkit_experimental.compact_stat ORDER BY bucket"#, None, None, - ); + ).unwrap(); // Day 1, in "three" from "16:00" to end of day - assert_eq!(durations.next().unwrap()[1].value(), Some("08:00:00")); + assert_eq!( + durations.next().unwrap()[1].value().unwrap(), + Some("08:00:00") + ); // Day 2, in "three" from start of day to "2:00" and "20:00" to end of day - assert_eq!(durations.next().unwrap()[1].value(), Some("06:00:00")); + assert_eq!( + durations.next().unwrap()[1].value().unwrap(), + Some("06:00:00") + ); // Day 3, in "three" from start of day to end - assert_eq!(durations.next().unwrap()[1].value(), Some("18:00:00")); + assert_eq!( + durations.next().unwrap()[1].value().unwrap(), + Some("18:00:00") + ); assert!(durations.next().is_none()); }); } @@ -2012,50 +2461,63 @@ SELECT toolkit_experimental.duration_in('one', toolkit_experimental.compact_stat error = "state cannot be both String(\"ERROR\") and String(\"START\") at 631152000000000" )] fn two_states_at_one_time() { - Spi::execute(|client| { - client.select("CREATE TABLE test(ts timestamptz, state TEXT)", None, None); - client.select( - r#"INSERT INTO test VALUES + Spi::connect(|mut client| { + client + .update("CREATE TABLE test(ts timestamptz, state TEXT)", None, None) + .unwrap(); + client + .update( + r#"INSERT INTO test VALUES ('2020-01-01 00:00:00+00', 'START'), ('2020-01-01 00:00:00+00', 'ERROR')"#, - None, - None, - ); - client.select( - "SELECT toolkit_experimental.duration_in(toolkit_experimental.compact_state_agg(ts, state), 'one') FROM test", - None, - None, - ); - client.select( - "SELECT toolkit_experimental.duration_in(toolkit_experimental.state_agg(ts, state), 'one') FROM test", - None, - None, - ); + None, + None, + ) + .unwrap(); + client + .update( + "SELECT toolkit_experimental.duration_in(toolkit_experimental.compact_state_agg(ts, state), 'one') FROM test", + None, + None, + ) + .unwrap(); + client + .update( + "SELECT duration_in(state_agg(ts, state), 'one') FROM test", + None, + None, + ) + .unwrap(); }) } #[pg_test] fn interpolate_introduces_state() { - Spi::execute(|client| { - client.select( - "CREATE TABLE states(time TIMESTAMPTZ, state TEXT, bucket INT)", - None, - None, - ); - client.select( - r#"INSERT INTO states VALUES + Spi::connect(|mut client| { + client + .update( + "CREATE TABLE states(time TIMESTAMPTZ, state TEXT, bucket INT)", + None, + None, + ) + .unwrap(); + client + .update( + r#"INSERT INTO states VALUES ('2020-1-1 10:00', 'starting', 1), ('2020-1-1 10:30', 'running', 1), ('2020-1-2 16:00', 'error', 2), ('2020-1-3 18:30', 'starting', 3), ('2020-1-3 19:30', 'running', 3), ('2020-1-4 12:00', 'stopping', 4)"#, - None, - None, - ); + None, + None, + ) + .unwrap(); - let mut durations = client.select( - r#"SELECT + let mut durations = client + .update( + r#"SELECT toolkit_experimental.interpolated_duration_in( agg, 'running', @@ -2067,36 +2529,243 @@ SELECT toolkit_experimental.duration_in('one', toolkit_experimental.compact_stat GROUP BY bucket ) s ORDER BY bucket"#, - None, - None, - ); + None, + None, + ) + .unwrap(); - assert_eq!(durations.next().unwrap()[1].value(), Some("13:30:00")); - assert_eq!(durations.next().unwrap()[1].value(), Some("16:00:00")); - assert_eq!(durations.next().unwrap()[1].value(), Some("04:30:00")); - assert_eq!(durations.next().unwrap()[1].value(), Some("12:00:00")); + assert_eq!( + durations.next().unwrap()[1].value().unwrap(), + Some("13:30:00") + ); + assert_eq!( + durations.next().unwrap()[1].value().unwrap(), + Some("16:00:00") + ); + assert_eq!( + durations.next().unwrap()[1].value().unwrap(), + Some("04:30:00") + ); + assert_eq!( + durations.next().unwrap()[1].value().unwrap(), + Some("12:00:00") + ); - let mut durations = client.select( - r#"SELECT - toolkit_experimental.interpolated_duration_in( + let mut durations = client + .update( + r#"SELECT + interpolated_duration_in( agg, 'running', '2019-12-31 0:00'::timestamptz + (bucket * '1 day'::interval), '1 day'::interval, LAG(agg) OVER (ORDER BY bucket) )::TEXT FROM ( - SELECT bucket, toolkit_experimental.state_agg(time, state) as agg + SELECT bucket, state_agg(time, state) as agg FROM states GROUP BY bucket ) s ORDER BY bucket"#, - None, - None, + None, + None, + ) + .unwrap(); + + assert_eq!( + durations.next().unwrap()[1].value().unwrap(), + Some("13:30:00") ); + assert_eq!( + durations.next().unwrap()[1].value().unwrap(), + Some("16:00:00") + ); + assert_eq!( + durations.next().unwrap()[1].value().unwrap(), + Some("04:30:00") + ); + assert_eq!( + durations.next().unwrap()[1].value().unwrap(), + Some("12:00:00") + ); + + let mut durations = client + .update( + r#"SELECT + (agg -> interpolated_duration_in( + 'running', + '2019-12-31 0:00'::timestamptz + (bucket * '1 day'::interval), '1 day'::interval, + LAG(agg) OVER (ORDER BY bucket) + ))::TEXT FROM ( + SELECT bucket, state_agg(time, state) as agg + FROM states + GROUP BY bucket + ) s + ORDER BY bucket"#, + None, + None, + ) + .unwrap(); - assert_eq!(durations.next().unwrap()[1].value(), Some("13:30:00")); - assert_eq!(durations.next().unwrap()[1].value(), Some("16:00:00")); - assert_eq!(durations.next().unwrap()[1].value(), Some("04:30:00")); - assert_eq!(durations.next().unwrap()[1].value(), Some("12:00:00")); + assert_eq!( + durations.next().unwrap()[1].value().unwrap(), + Some("13:30:00") + ); + assert_eq!( + durations.next().unwrap()[1].value().unwrap(), + Some("16:00:00") + ); + assert_eq!( + durations.next().unwrap()[1].value().unwrap(), + Some("04:30:00") + ); + assert_eq!( + durations.next().unwrap()[1].value().unwrap(), + Some("12:00:00") + ); }) } + + #[pg_test] + fn text_serialization() { + Spi::connect(|mut client| { + client + .update( + "SET TIME ZONE 'UTC'; + CREATE TABLE states(ts TIMESTAMPTZ, state TEXT); + CREATE TABLE states_int(ts TIMESTAMPTZ, state BIGINT);", + None, + None, + ) + .unwrap(); + // only a single entry so ordering is consistent between runs + client + .update( + r#"INSERT INTO states VALUES + ('2020-1-1 10:00', 'starting'); + INSERT INTO states_int VALUES + ('2020-1-1 10:00', -67876545);"#, + None, + None, + ) + .unwrap(); + + let agg_text = select_one!( + client, + "SELECT state_agg(ts, state)::TEXT FROM states", + &str + ); + let expected = "(version:1,compact_state_agg:(version:1,states_len:8,durations_len:1,durations:[(duration:0,state:(a:0,b:8))],combined_durations_len:1,combined_durations:[(start_time:631188000000000,end_time:631188000000000,state:(a:0,b:8))],first_time:631188000000000,last_time:631188000000000,first_state:0,last_state:0,states:[115,116,97,114,116,105,110,103],compact:false,integer_states:false))"; + assert_eq!(agg_text, expected); + + let agg_text = select_one!( + client, + "SELECT state_agg(ts, state)::TEXT FROM states_int", + &str + ); + let expected = "(version:1,compact_state_agg:(version:1,states_len:0,durations_len:1,durations:[(duration:0,state:(a:9223372036854775807,b:-67876545))],combined_durations_len:1,combined_durations:[(start_time:631188000000000,end_time:631188000000000,state:(a:9223372036854775807,b:-67876545))],first_time:631188000000000,last_time:631188000000000,first_state:0,last_state:0,states:[],compact:false,integer_states:true))"; + assert_eq!(agg_text, expected); + }); + } + + #[pg_test] + fn combine() { + assert_eq!(state_agg::combine(None, None), None); + + let mut trans_state_2 = CompactStateAggTransState::new(true); + trans_state_2.record(MaterializedState::Integer(444), 10005000); + let mut trans_state_1 = CompactStateAggTransState::new(true); + trans_state_1.record(MaterializedState::Integer(333), 10000000); + let trans_state = state_agg::combine(Some(&trans_state_1), Some(&trans_state_2)).unwrap(); + let trans_state = state_agg::combine(Some(&trans_state), None).unwrap(); + let trans_state = state_agg::combine(None, Some(&trans_state)).unwrap(); + assert_eq!( + trans_state, + CompactStateAggTransState { + records: vec![ + Record { + state: MaterializedState::Integer(333), + time: 10000000 + }, + Record { + state: MaterializedState::Integer(444), + time: 10005000 + } + ], + integer_states: true, + } + ); + } + + #[pg_test] + fn binary_serialization_integer() { + let mut trans_state = CompactStateAggTransState::new(true); + // only inserting one state since to avoid random ordering + trans_state.record(MaterializedState::Integer(22), 99); + let agg = state_agg::finally(Some(&mut trans_state)).unwrap(); + + // dis: duration i64, state entry (i64, i64) + let expected = [ + 232, 1, 0, 0, // header + 1, // version + 0, 0, 0, // padding + // inner compact_state_agg: + 200, 1, 0, 0, // header + 1, // version + 0, 0, 0, // padding + 0, 0, 0, 0, 0, 0, 0, 0, // states_len (empty since integer states) + 1, 0, 0, 0, 0, 0, 0, 0, // durations_len + 0, 0, 0, 0, 0, 0, 0, 0, // state 1: duration + 255, 255, 255, 255, 255, 255, 255, 127, // state 1: a + 22, 0, 0, 0, 0, 0, 0, 0, // state 1: b + 1, 0, 0, 0, 0, 0, 0, 0, // combined_durations_len + 99, 0, 0, 0, 0, 0, 0, 0, // state 1: start time + 99, 0, 0, 0, 0, 0, 0, 0, // state 1: end time + 255, 255, 255, 255, 255, 255, 255, 127, // state 1: a + 22, 0, 0, 0, 0, 0, 0, 0, // state 1: b + 99, 0, 0, 0, 0, 0, 0, 0, // first_time + 99, 0, 0, 0, 0, 0, 0, 0, // last_time + 0, 0, 0, 0, // first_state (index) + 0, 0, 0, 0, // last_state (index) + // states array is empty + 0, // compact (false) + 1, // integer states (true) + ]; + assert_eq!(agg.to_pg_bytes(), expected); + } + + #[pg_test] + fn binary_serialization_string() { + let mut trans_state = CompactStateAggTransState::new(false); + // only inserting one state since to avoid random ordering + trans_state.record(MaterializedState::String("ABC".to_string()), 99); + let agg = state_agg::finally(Some(&mut trans_state)).unwrap(); + + // dis: duration i64, state entry (i64, i64) + let expected = [ + 244, 1, 0, 0, // header + 1, // version + 0, 0, 0, // padding + // inner compact_state_agg: + 212, 1, 0, 0, // header + 1, // version + 0, 0, 0, // padding + 3, 0, 0, 0, 0, 0, 0, 0, // states_len + 1, 0, 0, 0, 0, 0, 0, 0, // durations_len + 0, 0, 0, 0, 0, 0, 0, 0, // state 1: duration + 0, 0, 0, 0, 0, 0, 0, 0, // state 1: a + 3, 0, 0, 0, 0, 0, 0, 0, // state 1: b + 1, 0, 0, 0, 0, 0, 0, 0, // combined_durations_len + 99, 0, 0, 0, 0, 0, 0, 0, // state 1: start time + 99, 0, 0, 0, 0, 0, 0, 0, // state 1: end time + 0, 0, 0, 0, 0, 0, 0, 0, // state 1: a + 3, 0, 0, 0, 0, 0, 0, 0, // state 1: b + 99, 0, 0, 0, 0, 0, 0, 0, // first_time + 99, 0, 0, 0, 0, 0, 0, 0, // last_time + 0, 0, 0, 0, // first_state (index) + 0, 0, 0, 0, // last_state (index) + 65, 66, 67, // states array + 0, // compact (false) + 0, // integer states (false) + ]; + assert_eq!(agg.to_pg_bytes(), expected); + } } diff --git a/extension/src/state_aggregate/accessors.rs b/extension/src/state_aggregate/accessors.rs new file mode 100644 index 00000000..7a85df7e --- /dev/null +++ b/extension/src/state_aggregate/accessors.rs @@ -0,0 +1,335 @@ +use crate::{ + datum_utils::interval_to_ms, + raw::{Interval, TimestampTz}, + state_aggregate::*, +}; + +pg_type! { + struct AccessorInterpolatedStateTimeline<'input> { + start: i64, + interval: i64, + prev: StateAggData<'input>, + prev_present: bool, + } +} +ron_inout_funcs!(AccessorInterpolatedStateTimeline); + +#[pg_extern(immutable, parallel_safe, name = "interpolated_state_timeline")] +fn accessor_state_agg_interpolated_interpolated_state_timeline<'a>( + start: TimestampTz, + interval: Interval, + prev: Option>, +) -> AccessorInterpolatedStateTimeline<'a> { + crate::build! { + AccessorInterpolatedStateTimeline { + interval: interval_to_ms(&start, &interval), + start: start.into(), + prev_present: prev.is_some(), + prev: prev.unwrap_or_else(|| StateAgg::empty(false)).0, + } + } +} + +pg_type! { + struct AccessorInterpolatedStateIntTimeline<'input> { + start: i64, + interval: i64, + prev: StateAggData<'input>, + prev_present: bool, + } +} +ron_inout_funcs!(AccessorInterpolatedStateIntTimeline); + +#[pg_extern(immutable, parallel_safe, name = "interpolated_state_int_timeline")] +fn accessor_state_agg_interpolated_interpolated_state_int_timeline<'a>( + start: TimestampTz, + interval: Interval, + prev: Option>, +) -> AccessorInterpolatedStateIntTimeline<'a> { + crate::build! { + AccessorInterpolatedStateIntTimeline { + interval: interval_to_ms(&start, &interval), + start: start.into(), + prev_present: prev.is_some(), + prev: prev.unwrap_or_else(|| StateAgg::empty(false)).0, + } + } +} + +// weird ordering is needed for alignment +pg_type! { + struct AccessorInterpolatedDurationIn<'input> { + start: i64, + interval: i64, + state_len: u32, + padding_2: [u8; 4], + prev: StateAggData<'input>, + state_bytes: [u8; self.state_len], + prev_present: bool, + } +} +ron_inout_funcs!(AccessorInterpolatedDurationIn); +pg_type! { + struct AccessorInterpolatedDurationInInt<'input> { + start: i64, + interval: i64, + state: i64, + prev_present: bool, + padding_2: [u8; 7], + prev: StateAggData<'input>, + } +} +ron_inout_funcs!(AccessorInterpolatedDurationInInt); + +#[pg_extern(immutable, parallel_safe, name = "interpolated_duration_in")] +fn accessor_state_agg_interpolated_interpolated_duration_in<'a>( + state: String, + start: TimestampTz, + interval: Interval, + prev: Option>, +) -> AccessorInterpolatedDurationIn<'a> { + crate::build! { + AccessorInterpolatedDurationIn { + state_len: state.len().try_into().unwrap(), + state_bytes: state.into_bytes().into(), + interval: interval_to_ms(&start, &interval), + start: start.into(), + prev_present: prev.is_some(), + prev: prev.unwrap_or_else(|| StateAgg::empty(false)).0, + padding_2: Default::default(), + } + } +} +#[pg_extern(immutable, parallel_safe, name = "interpolated_duration_in")] +fn accessor_state_agg_interpolated_interpolated_duration_in_int<'a>( + state: i64, + start: TimestampTz, + interval: Interval, + prev: Option>, +) -> AccessorInterpolatedDurationInInt<'a> { + crate::build! { + AccessorInterpolatedDurationInInt { + state, + interval: interval_to_ms(&start, &interval), + start: start.into(), + prev_present: prev.is_some(), + prev: prev.unwrap_or_else(|| StateAgg::empty(false)).0, + padding_2: Default::default(), + } + } +} + +// weird ordering is needed for alignment +pg_type! { + struct AccessorInterpolatedStatePeriods<'input> { + start: i64, + interval: i64, + state_len: u32, + padding_2: [u8; 4], + prev: StateAggData<'input>, + state_bytes: [u8; self.state_len], + prev_present: bool, + } +} +ron_inout_funcs!(AccessorInterpolatedStatePeriods); +pg_type! { + struct AccessorInterpolatedStatePeriodsInt<'input> { + start: i64, + interval: i64, + state: i64, + prev_present: bool, + padding_2: [u8; 7], + prev: StateAggData<'input>, + } +} +ron_inout_funcs!(AccessorInterpolatedStatePeriodsInt); + +#[pg_extern(immutable, parallel_safe, name = "interpolated_state_periods")] +fn accessor_state_agg_interpolated_interpolated_state_periods<'a>( + state: String, + start: TimestampTz, + interval: Interval, + prev: Option>, +) -> AccessorInterpolatedStatePeriods<'a> { + crate::build! { + AccessorInterpolatedStatePeriods { + state_len: state.len().try_into().unwrap(), + state_bytes: state.into_bytes().into(), + interval: interval_to_ms(&start, &interval), + start: start.into(), + prev_present: prev.is_some(), + prev: prev.unwrap_or_else(|| StateAgg::empty(false)).0, + padding_2: Default::default(), + } + } +} +#[pg_extern(immutable, parallel_safe, name = "interpolated_state_periods")] +fn accessor_state_agg_interpolated_interpolated_state_periods_int<'a>( + state: i64, + start: TimestampTz, + interval: Interval, + prev: Option>, +) -> AccessorInterpolatedStatePeriodsInt<'a> { + crate::build! { + AccessorInterpolatedStatePeriodsInt { + state, + interval: interval_to_ms(&start, &interval), + start: start.into(), + prev_present: prev.is_some(), + prev: prev.unwrap_or_else(|| StateAgg::empty(false)).0, + padding_2: Default::default(), + } + } +} + +pg_type! { + struct AccessorDurationIn<'input> { + state_len: u32, + state_bytes: [u8; self.state_len], + } +} +ron_inout_funcs!(AccessorDurationIn); +pg_type! { + struct AccessorDurationInInt { + state: i64, + } +} +ron_inout_funcs!(AccessorDurationInInt); + +#[pg_extern(immutable, parallel_safe, name = "duration_in")] +fn accessor_state_agg_duration_in<'a>(state: String) -> AccessorDurationIn<'a> { + crate::build! { + AccessorDurationIn { + state_len: state.len().try_into().unwrap(), + state_bytes: state.into_bytes().into(), + } + } +} +#[pg_extern(immutable, parallel_safe, name = "duration_in")] +fn accessor_state_agg_duration_in_int<'a>(state: i64) -> AccessorDurationInInt<'a> { + crate::build! { + AccessorDurationInInt { + state, + } + } +} + +pg_type! { + struct AccessorStatePeriods<'input> { + state_len: u32, + state_bytes: [u8; self.state_len], + } +} +ron_inout_funcs!(AccessorStatePeriods); +pg_type! { + struct AccessorStatePeriodsInt { + state: i64, + } +} +ron_inout_funcs!(AccessorStatePeriodsInt); + +#[pg_extern(immutable, parallel_safe, name = "state_periods")] +fn accessor_state_agg_state_periods<'a>(state: String) -> AccessorStatePeriods<'a> { + crate::build! { + AccessorStatePeriods { + state_len: state.len().try_into().unwrap(), + state_bytes: state.into_bytes().into(), + } + } +} +#[pg_extern(immutable, parallel_safe, name = "state_periods")] +fn accessor_state_agg_state_periods_int<'a>(state: i64) -> AccessorStatePeriodsInt<'a> { + crate::build! { + AccessorStatePeriodsInt { + state, + } + } +} + +pg_type! { + struct AccessorDurationInRange<'input> { + state_len: u32, + padding_2: [u8; 4], + start: i64, + interval: i64, + state_bytes: [u8; self.state_len], + } +} +ron_inout_funcs!(AccessorDurationInRange); +pg_type! { + struct AccessorDurationInRangeInt { + state: i64, + start: i64, + interval: i64, + } +} +ron_inout_funcs!(AccessorDurationInRangeInt); + +#[pg_extern(immutable, parallel_safe, name = "duration_in")] +fn accessor_state_agg_duration_in_range<'a>( + state: String, + start: TimestampTz, + interval: default!(Option, "NULL"), +) -> AccessorDurationInRange<'a> { + let interval = interval + .map(|interval| crate::datum_utils::interval_to_ms(&start, &interval)) + .unwrap_or(NO_INTERVAL_MARKER); + let start = start.into(); + crate::build! { + AccessorDurationInRange { + state_len: state.len().try_into().unwrap(), + state_bytes: state.into_bytes().into(), + padding_2: [0; 4], + start, interval + } + } +} +#[pg_extern(immutable, parallel_safe, name = "duration_in")] +fn accessor_state_agg_duration_in_range_int<'a>( + state: i64, + start: TimestampTz, + interval: default!(Option, "NULL"), +) -> AccessorDurationInRangeInt<'a> { + let interval = interval + .map(|interval| crate::datum_utils::interval_to_ms(&start, &interval)) + .unwrap_or(NO_INTERVAL_MARKER); + let start = start.into(); + crate::build! { + AccessorDurationInRangeInt { + state, + start, interval + } + } +} + +pg_type! { + struct AccessorStateAt { + time: i64, + } +} +ron_inout_funcs!(AccessorStateAt); + +#[pg_extern(immutable, parallel_safe, name = "state_at")] +fn accessor_state_agg_state_at<'a>(time: TimestampTz) -> AccessorStateAt<'a> { + crate::build! { + AccessorStateAt { + time: time.into(), + } + } +} + +pg_type! { + struct AccessorStateAtInt { + time: i64, + } +} +ron_inout_funcs!(AccessorStateAtInt); + +#[pg_extern(immutable, parallel_safe, name = "state_at_int")] +fn accessor_state_agg_state_at_int<'a>(time: TimestampTz) -> AccessorStateAtInt<'a> { + crate::build! { + AccessorStateAtInt { + time: time.into(), + } + } +} diff --git a/extension/src/state_aggregate/rollup.rs b/extension/src/state_aggregate/rollup.rs index d7f18bc8..9899aa36 100644 --- a/extension/src/state_aggregate/rollup.rs +++ b/extension/src/state_aggregate/rollup.rs @@ -12,40 +12,40 @@ extension_sql!( sfunc = toolkit_experimental.compact_state_agg_rollup_trans, stype = internal, finalfunc = toolkit_experimental.compact_state_agg_rollup_final, - combinefunc = toolkit_experimental.compact_state_agg_rollup_combine, - serialfunc = toolkit_experimental.compact_state_agg_rollup_serialize, - deserialfunc = toolkit_experimental.compact_state_agg_rollup_deserialize, + combinefunc = state_agg_rollup_combine, + serialfunc = state_agg_rollup_serialize, + deserialfunc = state_agg_rollup_deserialize, parallel = restricted );", name = "compact_state_agg_rollup", requires = [ compact_state_agg_rollup_trans, compact_state_agg_rollup_final, - compact_state_agg_rollup_combine, - compact_state_agg_rollup_serialize, - compact_state_agg_rollup_deserialize, + state_agg_rollup_combine, + state_agg_rollup_serialize, + state_agg_rollup_deserialize, CompactStateAgg, ], ); extension_sql!( - "CREATE AGGREGATE toolkit_experimental.rollup( - value toolkit_experimental.StateAgg + "CREATE AGGREGATE rollup( + value StateAgg ) ( - sfunc = toolkit_experimental.state_agg_rollup_trans, + sfunc = state_agg_rollup_trans, stype = internal, - finalfunc = toolkit_experimental.state_agg_rollup_final, - combinefunc = toolkit_experimental.compact_state_agg_rollup_combine, - serialfunc = toolkit_experimental.compact_state_agg_rollup_serialize, - deserialfunc = toolkit_experimental.compact_state_agg_rollup_deserialize, + finalfunc = state_agg_rollup_final, + combinefunc = state_agg_rollup_combine, + serialfunc = state_agg_rollup_serialize, + deserialfunc = state_agg_rollup_deserialize, parallel = restricted );", name = "state_agg_rollup", requires = [ state_agg_rollup_trans, state_agg_rollup_final, - compact_state_agg_rollup_combine, - compact_state_agg_rollup_serialize, - compact_state_agg_rollup_deserialize, + state_agg_rollup_combine, + state_agg_rollup_serialize, + state_agg_rollup_deserialize, StateAgg, ], ); @@ -56,7 +56,7 @@ pub struct RollupTransState { compact: bool, } -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] struct OwnedCompactStateAgg { durations: Vec, combined_durations: Vec, @@ -80,14 +80,22 @@ impl OwnedCompactStateAgg { "can't merge aggs with different state types" ); - let (earlier, later) = match self.first_time.cmp(&other.first_time) { + let (earlier, later) = match self.cmp(&other) { Ordering::Less => (self, other), Ordering::Greater => (other, self), - Ordering::Equal => panic!("can't merge overlapping aggregates (same start time)"), + Ordering::Equal => panic!( + "can't merge overlapping aggregates (same start time: {})", + self.first_time + ), }; + assert!( earlier.last_time <= later.first_time, - "can't merge overlapping aggregates" + "can't merge overlapping aggregates (earlier={}-{}, later={}-{})", + earlier.first_time, + earlier.last_time, + later.first_time, + later.last_time, ); assert_ne!( later.durations.len(), @@ -97,7 +105,7 @@ impl OwnedCompactStateAgg { assert_ne!( earlier.durations.len(), 0, - "later aggregate must be non-empty" + "earlier aggregate must be non-empty" ); let later_states = @@ -108,25 +116,37 @@ impl OwnedCompactStateAgg { let earlier_len = earlier.combined_durations.len(); - let mut added_entries = 0; - for dis in later.durations.iter() { - let merged_duration_to_update = merged_durations.iter_mut().find(|merged_dis| { - merged_dis.state.materialize(&merged_states) == dis.state.materialize(&later_states) - }); - if let Some(merged_duration_to_update) = merged_duration_to_update { - merged_duration_to_update.duration += dis.duration; - } else { - let state = dis - .state - .materialize(&later_states) - .entry(&mut merged_states); - merged_durations.push(DurationInState { - state, - duration: dis.duration, - }); - added_entries += 1; + let mut merged_last_state = None; + for (later_idx, dis) in later.durations.iter().enumerate() { + let materialized_dis = dis.state.materialize(&later_states); + let merged_duration_info = + merged_durations + .iter_mut() + .enumerate() + .find(|(_, merged_dis)| { + merged_dis.state.materialize(&merged_states) == materialized_dis + }); + + let merged_idx = + if let Some((merged_idx, merged_duration_to_update)) = merged_duration_info { + merged_duration_to_update.duration += dis.duration; + merged_idx + } else { + let state = materialized_dis.entry(&mut merged_states); + merged_durations.push(DurationInState { + state, + duration: dis.duration, + }); + merged_durations.len() - 1 + }; + + if later_idx == later.last_state as usize { + // this is the last state + merged_last_state = Some(merged_idx); }; } + let merged_last_state = + merged_last_state.expect("later last_state not in later.durations") as u32; let mut combined_durations = earlier .combined_durations @@ -142,22 +162,36 @@ impl OwnedCompactStateAgg { let gap = later.first_time - earlier.last_time; assert!(gap >= 0); - merged_durations[earlier.last_state as usize].duration += gap; + merged_durations + .get_mut(earlier.last_state as usize) + .expect("earlier.last_state doesn't point to a state") + .duration += gap; // ensure combined_durations covers the whole range of time if !earlier.compact { - if combined_durations[earlier_len - 1] + if combined_durations + .get_mut(earlier_len - 1) + .expect("invalid combined_durations: nothing at end of earlier") .state .materialize(&merged_states) - == combined_durations[earlier_len] + == combined_durations + .get(earlier_len) + .expect("invalid combined_durations: nothing at start of earlier") .state .materialize(&merged_states) { - combined_durations[earlier_len - 1].end_time = - combined_durations.remove(earlier_len).end_time; + combined_durations + .get_mut(earlier_len - 1) + .expect("invalid combined_durations (nothing at earlier_len - 1, equal)") + .end_time = combined_durations.remove(earlier_len).end_time; } else { - combined_durations[earlier_len - 1].end_time = - combined_durations[earlier_len].start_time; + combined_durations + .get_mut(earlier_len - 1) + .expect("invalid combined_durations (nothing at earlier_len - 1, not equal)") + .end_time = combined_durations + .get(earlier_len) + .expect("invalid combined_durations (nothing at earlier_len, not equal)") + .start_time; } } @@ -169,8 +203,8 @@ impl OwnedCompactStateAgg { first_time: earlier.first_time, last_time: later.last_time, - first_state: earlier.first_state, - last_state: added_entries + later.last_state, + first_state: earlier.first_state, // indexes into earlier durations are same for merged_durations + last_state: merged_last_state, // these values are always the same for both compact: earlier.compact, @@ -216,8 +250,23 @@ impl<'a> From> for OwnedCompactStateAgg { } } +impl PartialOrd for OwnedCompactStateAgg { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for OwnedCompactStateAgg { + fn cmp(&self, other: &Self) -> Ordering { + // compare using first time (OwnedCompactStateAgg::merge will handle any overlap) + self.first_time.cmp(&other.first_time) + } +} + impl RollupTransState { fn merge(&mut self) { + // OwnedCompactStateAgg::merge can't merge overlapping aggregates + self.values.sort(); self.values = self .values .drain(..) @@ -260,7 +309,7 @@ pub fn compact_state_agg_rollup_trans_inner<'a>( } } -#[pg_extern(immutable, parallel_safe, schema = "toolkit_experimental")] +#[pg_extern(immutable, parallel_safe)] pub fn state_agg_rollup_trans<'a>( state: Internal, next: Option>, @@ -300,7 +349,7 @@ fn compact_state_agg_rollup_final_inner<'a>( } } -#[pg_extern(immutable, parallel_safe, schema = "toolkit_experimental")] +#[pg_extern(immutable, parallel_safe)] fn state_agg_rollup_final<'a>( state: Internal, fcinfo: pg_sys::FunctionCallInfo, @@ -326,36 +375,35 @@ fn state_agg_rollup_final_inner<'a>( } } -#[pg_extern(immutable, parallel_safe, strict, schema = "toolkit_experimental")] -pub fn compact_state_agg_rollup_serialize(state: Internal) -> bytea { +#[pg_extern(immutable, parallel_safe, strict)] +pub fn state_agg_rollup_serialize(state: Internal) -> bytea { let mut state: Inner = unsafe { state.to_inner().unwrap() }; state.merge(); crate::do_serialize!(state) } -#[pg_extern(strict, immutable, parallel_safe, schema = "toolkit_experimental")] -pub fn compact_state_agg_rollup_deserialize(bytes: bytea, _internal: Internal) -> Option { - compact_state_agg_rollup_deserialize_inner(bytes).internal() +#[pg_extern(strict, immutable, parallel_safe)] +pub fn state_agg_rollup_deserialize(bytes: bytea, _internal: Internal) -> Option { + state_agg_rollup_deserialize_inner(bytes).internal() } -pub fn compact_state_agg_rollup_deserialize_inner(bytes: bytea) -> Inner { +pub fn state_agg_rollup_deserialize_inner(bytes: bytea) -> Inner { let t: RollupTransState = crate::do_deserialize!(bytes, RollupTransState); t.into() } -#[pg_extern(immutable, parallel_safe, schema = "toolkit_experimental")] -pub fn compact_state_agg_rollup_combine( +#[pg_extern(immutable, parallel_safe)] +pub fn state_agg_rollup_combine( state1: Internal, state2: Internal, fcinfo: pg_sys::FunctionCallInfo, ) -> Option { unsafe { - compact_state_agg_rollup_combine_inner(state1.to_inner(), state2.to_inner(), fcinfo) - .internal() + state_agg_rollup_combine_inner(state1.to_inner(), state2.to_inner(), fcinfo).internal() } } #[allow(clippy::redundant_clone)] // clone is needed so we don't mutate shared memory -pub fn compact_state_agg_rollup_combine_inner( +pub fn state_agg_rollup_combine_inner( state1: Option>, state2: Option>, fcinfo: pg_sys::FunctionCallInfo, @@ -417,4 +465,108 @@ mod tests { r2.merge(r1); } + + #[test] + fn merges_compact_aggs_correctly() { + let s1 = OwnedCompactStateAgg { + durations: vec![ + DurationInState { + duration: 500, + state: StateEntry::from_integer(555_2), + }, + DurationInState { + duration: 400, + state: StateEntry::from_integer(555_1), + }, + ], + combined_durations: vec![], + first_time: 100, + last_time: 1000, + first_state: 1, + last_state: 0, + states: vec![], + compact: true, + integer_states: true, + }; + let s2 = OwnedCompactStateAgg { + durations: vec![ + DurationInState { + duration: 500, + state: StateEntry::from_integer(555_2), + }, + DurationInState { + duration: 400, + state: StateEntry::from_integer(555_1), + }, + ], + combined_durations: vec![], + first_time: 1000 + 12345, + last_time: 1900 + 12345, + first_state: 1, + last_state: 0, + states: vec![], + compact: true, + integer_states: true, + }; + let s3 = OwnedCompactStateAgg { + durations: vec![ + DurationInState { + duration: 500, + state: StateEntry::from_integer(555_2), + }, + DurationInState { + duration: 400, + state: StateEntry::from_integer(555_1), + }, + ], + combined_durations: vec![], + first_time: 1900 + 12345, + last_time: 1900 + 12345 + 900, + first_state: 1, + last_state: 0, + states: vec![], + compact: true, + integer_states: true, + }; + let expected = OwnedCompactStateAgg { + durations: vec![ + DurationInState { + duration: 500 * 3 + 12345, + state: StateEntry::from_integer(555_2), + }, + DurationInState { + duration: 400 * 3, + state: StateEntry::from_integer(555_1), + }, + ], + combined_durations: vec![], + first_time: 100, + last_time: 1900 + 12345 + 900, + first_state: 1, + last_state: 0, + states: vec![], + compact: true, + integer_states: true, + }; + let merged = s1.clone().merge(s2.clone().merge(s3.clone())); + assert_eq!(merged, expected); + let merged = s3.clone().merge(s2.clone().merge(s1.clone())); + assert_eq!(merged, expected); + + let mut trans_state = RollupTransState { + values: vec![s1.clone(), s2.clone(), s3.clone()], + compact: true, + }; + trans_state.merge(); + assert_eq!(trans_state.values.len(), 1); + assert_eq!(trans_state.values[0], expected.clone()); + + let mut trans_state = RollupTransState { + values: vec![s3.clone(), s1.clone(), s2.clone()], + compact: true, + }; + trans_state.merge(); + assert_eq!(trans_state.values.len(), 1); + assert_eq!(trans_state.values[0], expected.clone()); + } } diff --git a/extension/src/stats_agg.rs b/extension/src/stats_agg.rs index 513f02fa..389a1aa7 100644 --- a/extension/src/stats_agg.rs +++ b/extension/src/stats_agg.rs @@ -1408,10 +1408,11 @@ pub fn as_method(method: &str) -> Option { // macro_rules! select_one { // ($client:expr, $stmt:expr, $type:ty) => { // $client -// .select($stmt, None, None) +// .update($stmt, None, None) // .first() // .get_one::<$type>() // .unwrap() +// .unwrap() // }; // } @@ -1434,7 +1435,7 @@ pub fn as_method(method: &str) -> Option { // // #[pg_test] // // fn test_combine_aggregate(){ -// // Spi::execute(|client| { +// // Spi::connect(|mut client| { // // }); // // } @@ -1458,48 +1459,60 @@ mod tests { #[pg_test] fn test_stats_agg_text_io() { - Spi::execute(|client| { - client.select( - "CREATE TABLE test_table (test_x DOUBLE PRECISION, test_y DOUBLE PRECISION)", - None, - None, - ); + Spi::connect(|mut client| { + client + .update( + "CREATE TABLE test_table (test_x DOUBLE PRECISION, test_y DOUBLE PRECISION)", + None, + None, + ) + .unwrap(); let test = client - .select( + .update( "SELECT stats_agg(test_y, test_x)::TEXT FROM test_table", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert!(test.is_none()); - client.select("INSERT INTO test_table VALUES (10, 10);", None, None); + client + .update("INSERT INTO test_table VALUES (10, 10);", None, None) + .unwrap(); let test = client - .select( + .update( "SELECT stats_agg(test_y, test_x)::TEXT FROM test_table", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() .unwrap(); assert_eq!( test, "(version:1,n:1,sx:10,sx2:0,sx3:0,sx4:0,sy:10,sy2:0,sy3:0,sy4:0,sxy:0)" ); - client.select("INSERT INTO test_table VALUES (20, 20);", None, None); + client + .update("INSERT INTO test_table VALUES (20, 20);", None, None) + .unwrap(); let test = client - .select( + .update( "SELECT stats_agg(test_y, test_x)::TEXT FROM test_table", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() .unwrap(); let expected = "(version:1,n:2,sx:30,sx2:50,sx3:0,sx4:1250,sy:30,sy2:50,sy3:0,sy4:1250,sxy:50)"; @@ -1508,55 +1521,61 @@ mod tests { // Test a few functions to see that the text serialized object behave the same as the constructed one assert_eq!( client - .select( + .update( "SELECT skewness_x(stats_agg(test_y, test_x)) FROM test_table", None, None ) + .unwrap() .first() .get_one::(), client - .select( + .update( &format!("SELECT skewness_x('{}'::StatsSummary2D)", expected), None, None ) + .unwrap() .first() .get_one::() ); assert_eq!( client - .select( + .update( "SELECT kurtosis_y(stats_agg(test_y, test_x)) FROM test_table", None, None ) + .unwrap() .first() .get_one::(), client - .select( + .update( &format!("SELECT kurtosis_y('{}'::StatsSummary2D)", expected), None, None ) + .unwrap() .first() .get_one::() ); assert_eq!( client - .select( + .update( "SELECT covariance(stats_agg(test_y, test_x)) FROM test_table", None, None ) + .unwrap() .first() .get_one::(), client - .select( + .update( &format!("SELECT covariance('{}'::StatsSummary2D)", expected), None, None ) + .unwrap() .first() .get_one::() ); @@ -1564,38 +1583,48 @@ mod tests { // Test text round trip assert_eq!( client - .select( + .update( &format!("SELECT '{}'::StatsSummary2D::TEXT", expected), None, None ) + .unwrap() .first() .get_one::() + .unwrap() .unwrap(), expected ); - client.select("INSERT INTO test_table VALUES ('NaN', 30);", None, None); + client + .update("INSERT INTO test_table VALUES ('NaN', 30);", None, None) + .unwrap(); let test = client - .select( + .update( "SELECT stats_agg(test_y, test_x)::TEXT FROM test_table", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() .unwrap(); assert_eq!(test, "(version:1,n:3,sx:NaN,sx2:NaN,sx3:NaN,sx4:NaN,sy:60,sy2:200,sy3:0,sy4:20000,sxy:NaN)"); - client.select("INSERT INTO test_table VALUES (40, 'Inf');", None, None); + client + .update("INSERT INTO test_table VALUES (40, 'Inf');", None, None) + .unwrap(); let test = client - .select( + .update( "SELECT stats_agg(test_y, test_x)::TEXT FROM test_table", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() .unwrap(); assert_eq!(test, "(version:1,n:4,sx:NaN,sx2:NaN,sx3:NaN,sx4:NaN,sy:inf,sy2:NaN,sy3:NaN,sy4:NaN,sxy:NaN)"); }); @@ -1705,22 +1734,27 @@ mod tests { #[allow(clippy::float_cmp)] fn check_agg_equivalence( state: &TestState, - client: &SpiClient, + client: &mut pgx::spi::SpiClient, pg_cmd: &str, tk_cmd: &str, allowed_diff: f64, do_moving_agg: bool, ) { - let (pg_result, pg_moving_agg_result) = client - .select(pg_cmd, None, None) - .first() - .get_two::(); + warning!("pg_cmd={} ; tk_cmd={}", pg_cmd, tk_cmd); + let pg_row = client.update(pg_cmd, None, None).unwrap().first(); + let (pg_result, pg_moving_agg_result) = if do_moving_agg { + pg_row.get_two::().unwrap() + } else { + (pg_row.get_one::().unwrap(), None) + }; let pg_result = pg_result.unwrap(); let (tk_result, arrow_result, tk_moving_agg_result) = client - .select(tk_cmd, None, None) + .update(tk_cmd, None, None) + .unwrap() .first() - .get_three::(); + .get_three::() + .unwrap(); let (tk_result, arrow_result) = (tk_result.unwrap(), arrow_result.unwrap()); assert_eq!(tk_result, arrow_result, "Arrow didn't match in {}", tk_cmd); @@ -1761,7 +1795,7 @@ mod tests { } fn pg1d_aggx(agg: &str) -> String { - format!("SELECT {agg}(test_x), (SELECT {agg}(test_x) OVER (ORDER BY test_x ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING) FROM test_table LIMIT 1 OFFSET 3) FROM test_table", agg = agg) + format!("SELECT {agg}(test_x)::float, (SELECT {agg}(test_x) OVER (ORDER BY test_x ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING) FROM test_table LIMIT 1 OFFSET 3)::float FROM test_table", agg = agg) } fn pg1d_aggy(agg: &str) -> String { @@ -1769,15 +1803,15 @@ mod tests { } fn pg2d_agg(agg: &str) -> String { - format!("SELECT {agg}(test_y, test_x), (SELECT {agg}(test_y, test_x) OVER (ORDER BY test_x ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING) FROM test_table LIMIT 1 OFFSET 3) FROM test_table", agg = agg) + format!("SELECT {agg}(test_y, test_x)::float, (SELECT {agg}(test_y, test_x) OVER (ORDER BY test_x ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING) FROM test_table LIMIT 1 OFFSET 3)::float FROM test_table", agg = agg) } fn tk1d_agg(agg: &str) -> String { format!( "SELECT \ - {agg}(stats_agg(test_x)), \ - stats_agg(test_x)->{agg}(), \ - {agg}((SELECT stats_agg(test_x) OVER (ORDER BY test_x ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING) FROM test_table LIMIT 1 OFFSET 3)) \ + {agg}(stats_agg(test_x))::float, \ + (stats_agg(test_x)->{agg}())::float, \ + {agg}((SELECT stats_agg(test_x) OVER (ORDER BY test_x ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING) FROM test_table LIMIT 1 OFFSET 3))::float \ FROM test_table", agg = agg ) @@ -1798,9 +1832,9 @@ mod tests { fn tk2d_agg(agg: &str) -> String { format!( "SELECT \ - {agg}(stats_agg(test_y, test_x)), \ - stats_agg(test_y, test_x)->{agg}(), \ - {agg}((SELECT stats_agg(test_y, test_x) OVER (ORDER BY test_x ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING) FROM test_table LIMIT 1 OFFSET 3)) \ + {agg}(stats_agg(test_y, test_x))::float, \ + (stats_agg(test_y, test_x)->{agg}())::float, \ + {agg}((SELECT stats_agg(test_y, test_x) OVER (ORDER BY test_x ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING) FROM test_table LIMIT 1 OFFSET 3))::float \ FROM test_table", agg = agg ) @@ -1827,32 +1861,36 @@ mod tests { } fn test_aggs(state: &mut TestState) { - Spi::execute(|client| { - client.select( - "CREATE TABLE test_table (test_x DOUBLE PRECISION, test_y DOUBLE PRECISION)", - None, - None, - ); + Spi::connect(|mut client| { + client + .update( + "CREATE TABLE test_table (test_x DOUBLE PRECISION, test_y DOUBLE PRECISION)", + None, + None, + ) + .unwrap(); - client.select( - &format!( - "INSERT INTO test_table VALUES {}", - state - .x_values - .iter() - .zip(state.y_values.iter()) - .map(|(x, y)| "(".to_string() - + &x.to_string() - + "," - + &y.to_string() - + ")" - + ",") - .collect::() - .trim_end_matches(',') - ), - None, - None, - ); + client + .update( + &format!( + "INSERT INTO test_table VALUES {}", + state + .x_values + .iter() + .zip(state.y_values.iter()) + .map(|(x, y)| "(".to_string() + + &x.to_string() + + "," + + &y.to_string() + + ")" + + ",") + .collect::() + .trim_end_matches(',') + ), + None, + None, + ) + .unwrap(); // Definitions for allowed errors for different aggregates const NONE: f64 = 0.; // Exact match @@ -1863,7 +1901,7 @@ mod tests { check_agg_equivalence( state, - &client, + &mut client, &pg1d_aggx("avg"), &tk1d_agg("average"), NONE, @@ -1871,7 +1909,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg1d_aggx("sum"), &tk1d_agg("sum"), NONE, @@ -1879,7 +1917,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg1d_aggx("count"), &tk1d_agg("num_vals"), NONE, @@ -1887,7 +1925,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg1d_aggx("stddev"), &tk1d_agg("stddev"), EPS2, @@ -1895,7 +1933,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg1d_aggx("stddev_pop"), &tk1d_agg_arg("stddev", "population"), EPS2, @@ -1903,7 +1941,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg1d_aggx("stddev_samp"), &tk1d_agg_arg("stddev", "sample"), EPS2, @@ -1911,7 +1949,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg1d_aggx("variance"), &tk1d_agg("variance"), EPS3, @@ -1919,7 +1957,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg1d_aggx("var_pop"), &tk1d_agg_arg("variance", "population"), EPS3, @@ -1927,7 +1965,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg1d_aggx("var_samp"), &tk1d_agg_arg("variance", "sample"), EPS3, @@ -1936,7 +1974,7 @@ mod tests { check_agg_equivalence( state, - &client, + &mut client, &pg2d_agg("regr_avgx"), &tk2d_agg("average_x"), NONE, @@ -1944,7 +1982,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg2d_agg("regr_avgy"), &tk2d_agg("average_y"), NONE, @@ -1952,7 +1990,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg1d_aggx("sum"), &tk2d_agg("sum_x"), NONE, @@ -1960,7 +1998,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg1d_aggy("sum"), &tk2d_agg("sum_y"), NONE, @@ -1968,7 +2006,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg1d_aggx("stddev"), &tk2d_agg("stddev_x"), EPS2, @@ -1976,7 +2014,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg1d_aggy("stddev"), &tk2d_agg("stddev_y"), EPS2, @@ -1984,7 +2022,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg1d_aggx("stddev_pop"), &tk2d_agg_arg("stddev_x", "population"), EPS2, @@ -1992,7 +2030,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg1d_aggy("stddev_pop"), &tk2d_agg_arg("stddev_y", "population"), EPS2, @@ -2000,7 +2038,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg1d_aggx("stddev_samp"), &tk2d_agg_arg("stddev_x", "sample"), EPS2, @@ -2008,7 +2046,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg1d_aggy("stddev_samp"), &tk2d_agg_arg("stddev_y", "sample"), EPS2, @@ -2016,7 +2054,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg1d_aggx("variance"), &tk2d_agg("variance_x"), EPS3, @@ -2024,7 +2062,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg1d_aggy("variance"), &tk2d_agg("variance_y"), EPS3, @@ -2032,7 +2070,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg1d_aggx("var_pop"), &tk2d_agg_arg("variance_x", "population"), EPS3, @@ -2040,7 +2078,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg1d_aggy("var_pop"), &tk2d_agg_arg("variance_y", "population"), EPS3, @@ -2048,7 +2086,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg1d_aggx("var_samp"), &tk2d_agg_arg("variance_x", "sample"), EPS3, @@ -2056,7 +2094,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg1d_aggy("var_samp"), &tk2d_agg_arg("variance_y", "sample"), EPS3, @@ -2064,7 +2102,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg2d_agg("regr_count"), &tk2d_agg("num_vals"), NONE, @@ -2073,7 +2111,7 @@ mod tests { check_agg_equivalence( state, - &client, + &mut client, &pg2d_agg("regr_slope"), &tk2d_agg("slope"), EPS1, @@ -2081,7 +2119,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg2d_agg("corr"), &tk2d_agg("corr"), EPS1, @@ -2089,7 +2127,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg2d_agg("regr_intercept"), &tk2d_agg("intercept"), EPS1, @@ -2100,15 +2138,17 @@ mod tests { { let query = tk2d_agg("x_intercept"); let (result, arrow_result) = client - .select(&query, None, None) + .update(&query, None, None) + .unwrap() .first() - .get_two::(); + .get_two::() + .unwrap(); assert_eq!(result, arrow_result, "Arrow didn't match in {}", query); } check_agg_equivalence( state, - &client, + &mut client, &pg2d_agg("regr_r2"), &tk2d_agg("determination_coeff"), EPS1, @@ -2116,7 +2156,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg2d_agg("covar_pop"), &tk2d_agg_arg("covariance", "population"), BILLIONTH, @@ -2124,7 +2164,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg2d_agg("covar_samp"), &tk2d_agg_arg("covariance", "sample"), BILLIONTH, @@ -2134,7 +2174,7 @@ mod tests { // Skewness and kurtosis don't have aggregate functions in postgres, but we can compute them check_agg_equivalence( state, - &client, + &mut client, &pg_moment_pop_query(3, "test_x"), &tk1d_agg_arg("skewness", "population"), BILLIONTH, @@ -2142,7 +2182,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg_moment_pop_query(3, "test_x"), &tk2d_agg_arg("skewness_x", "population"), BILLIONTH, @@ -2150,7 +2190,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg_moment_pop_query(3, "test_y"), &tk2d_agg_arg("skewness_y", "population"), BILLIONTH, @@ -2158,7 +2198,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg_moment_pop_query(4, "test_x"), &tk1d_agg_arg("kurtosis", "population"), BILLIONTH, @@ -2166,7 +2206,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg_moment_pop_query(4, "test_x"), &tk2d_agg_arg("kurtosis_x", "population"), BILLIONTH, @@ -2174,7 +2214,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg_moment_pop_query(4, "test_y"), &tk2d_agg_arg("kurtosis_y", "population"), BILLIONTH, @@ -2183,7 +2223,7 @@ mod tests { check_agg_equivalence( state, - &client, + &mut client, &pg_moment_samp_query(3, "test_x"), &tk1d_agg_arg("skewness", "sample"), BILLIONTH, @@ -2191,7 +2231,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg_moment_samp_query(3, "test_x"), &tk2d_agg_arg("skewness_x", "sample"), BILLIONTH, @@ -2199,7 +2239,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg_moment_samp_query(3, "test_y"), &tk2d_agg_arg("skewness_y", "sample"), BILLIONTH, @@ -2207,7 +2247,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg_moment_samp_query(4, "test_x"), &tk1d_agg_arg("kurtosis", "sample"), BILLIONTH, @@ -2215,7 +2255,7 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg_moment_samp_query(4, "test_x"), &tk2d_agg_arg("kurtosis_x", "sample"), BILLIONTH, @@ -2223,22 +2263,23 @@ mod tests { ); check_agg_equivalence( state, - &client, + &mut client, &pg_moment_samp_query(4, "test_y"), &tk2d_agg_arg("kurtosis_y", "sample"), BILLIONTH, false, ); - client.select("DROP TABLE test_table", None, None); + client.update("DROP TABLE test_table", None, None).unwrap(); }); } #[pg_test] fn stats_agg_rolling() { - Spi::execute(|client| { - client.select( - " + Spi::connect(|mut client| { + client + .update( + " SET timezone TO 'UTC'; CREATE TABLE prices(ts TIMESTAMPTZ, price FLOAT); INSERT INTO prices ( @@ -2254,25 +2295,30 @@ INSERT INTO prices ( dates ); ", - None, - None, - ); + None, + None, + ) + .unwrap(); - let mut vals = client.select( + let mut vals = client.update( "SELECT stddev(data.stats_agg) FROM (SELECT stats_agg(price) OVER (ORDER BY ts RANGE '50 minutes' PRECEDING) FROM prices) data", None, None - ); - assert!(vals.next().unwrap()[1].value::().unwrap().is_nan()); - assert!(vals.next().unwrap()[1].value::().is_some()); - assert!(vals.next().unwrap()[1].value::().is_some()); - - let mut vals = client.select( + ).unwrap(); + assert!(vals.next().unwrap()[1] + .value::() + .unwrap() + .unwrap() + .is_nan()); + assert!(vals.next().unwrap()[1].value::().unwrap().is_some()); + assert!(vals.next().unwrap()[1].value::().unwrap().is_some()); + + let mut vals = client.update( "SELECT slope(data.stats_agg) FROM (SELECT stats_agg((EXTRACT(minutes FROM ts)), price) OVER (ORDER BY ts RANGE '50 minutes' PRECEDING) FROM prices) data;", None, None - ); - assert!(vals.next().unwrap()[1].value::().is_none()); // trendline is zero initially - assert!(vals.next().unwrap()[1].value::().is_some()); - assert!(vals.next().unwrap()[1].value::().is_some()); + ).unwrap(); + assert!(vals.next().unwrap()[1].value::().unwrap().is_none()); // trendline is zero initially + assert!(vals.next().unwrap()[1].value::().unwrap().is_some()); + assert!(vals.next().unwrap()[1].value::().unwrap().is_some()); }); } } diff --git a/extension/src/tdigest.rs b/extension/src/tdigest.rs index e1e235c6..3671d459 100644 --- a/extension/src/tdigest.rs +++ b/extension/src/tdigest.rs @@ -130,7 +130,7 @@ impl<'input> InOutFuncs for TDigest<'input> { } } - fn input(input: &pgx::cstr_core::CStr) -> TDigest<'input> + fn input(input: &std::ffi::CStr) -> TDigest<'input> where Self: Sized, { @@ -441,29 +441,37 @@ mod tests { #[pg_test] fn test_tdigest_aggregate() { - Spi::execute(|client| { - client.select("CREATE TABLE test (data DOUBLE PRECISION)", None, None); - client.select( - "INSERT INTO test SELECT generate_series(0.01, 100, 0.01)", - None, - None, - ); + Spi::connect(|mut client| { + client + .update("CREATE TABLE test (data DOUBLE PRECISION)", None, None) + .unwrap(); + client + .update( + "INSERT INTO test SELECT generate_series(0.01, 100, 0.01)", + None, + None, + ) + .unwrap(); let sanity = client - .select("SELECT COUNT(*) FROM test", None, None) + .update("SELECT COUNT(*) FROM test", None, None) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(10000, sanity.unwrap()); - client.select( - "CREATE VIEW digest AS \ + client + .update( + "CREATE VIEW digest AS \ SELECT tdigest(100, data) FROM test", - None, - None, - ); + None, + None, + ) + .unwrap(); let (min, max, count) = client - .select( + .update( "SELECT \ min_val(tdigest), \ max_val(tdigest), \ @@ -472,15 +480,17 @@ mod tests { None, None, ) + .unwrap() .first() - .get_three::(); + .get_three::() + .unwrap(); apx_eql(min.unwrap(), 0.01, 0.000001); apx_eql(max.unwrap(), 100.0, 0.000001); apx_eql(count.unwrap(), 10000.0, 0.000001); let (min2, max2, count2) = client - .select( + .update( "SELECT \ tdigest->min_val(), \ tdigest->max_val(), \ @@ -489,15 +499,17 @@ mod tests { None, None, ) + .unwrap() .first() - .get_three::(); + .get_three::() + .unwrap(); assert_eq!(min2, min); assert_eq!(max2, max); assert_eq!(count2, count); let (mean, mean2) = client - .select( + .update( "SELECT \ mean(tdigest), \ tdigest -> mean() @@ -505,8 +517,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_two::(); + .get_two::() + .unwrap(); apx_eql(mean.unwrap(), 50.005, 0.0001); assert_eq!(mean, mean2); @@ -516,7 +530,7 @@ mod tests { let quantile = value / 100.0; let (est_val, est_quant) = client - .select( + .update( &format!( "SELECT approx_percentile({}, tdigest), \ @@ -527,8 +541,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_two::(); + .get_two::() + .unwrap(); if i == 0 { pct_eql(est_val.unwrap(), 0.01, 1.0); @@ -539,7 +555,7 @@ mod tests { } let (est_val2, est_quant2) = client - .select( + .update( &format!( "SELECT tdigest->approx_percentile({}), \ @@ -550,8 +566,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_two::(); + .get_two::() + .unwrap(); assert_eq!(est_val2, est_val); assert_eq!(est_quant2, est_quant); } @@ -560,9 +578,9 @@ mod tests { #[pg_test] fn test_tdigest_small_count() { - Spi::execute(|client| { + Spi::connect(|mut client| { let estimate = client - .select( + .update( "SELECT \ approx_percentile(\ 0.99, \ @@ -571,8 +589,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_one(); + .get_one() + .unwrap(); assert_eq!(estimate, Some(99.5)); }); @@ -593,30 +613,34 @@ mod tests { #[pg_test] fn test_tdigest_io() { - Spi::execute(|client| { + Spi::connect(|mut client| { let output = client - .select( + .update( "SELECT \ tdigest(100, data)::text \ FROM generate_series(1, 100) data;", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); let expected = "(version:1,buckets:88,max_buckets:100,count:100,sum:5050,min:1,max:100,centroids:[(mean:1,weight:1),(mean:2,weight:1),(mean:3,weight:1),(mean:4,weight:1),(mean:5,weight:1),(mean:6,weight:1),(mean:7,weight:1),(mean:8,weight:1),(mean:9,weight:1),(mean:10,weight:1),(mean:11,weight:1),(mean:12,weight:1),(mean:13,weight:1),(mean:14,weight:1),(mean:15,weight:1),(mean:16,weight:1),(mean:17,weight:1),(mean:18,weight:1),(mean:19,weight:1),(mean:20,weight:1),(mean:21,weight:1),(mean:22,weight:1),(mean:23,weight:1),(mean:24,weight:1),(mean:25,weight:1),(mean:26,weight:1),(mean:27,weight:1),(mean:28,weight:1),(mean:29,weight:1),(mean:30,weight:1),(mean:31,weight:1),(mean:32,weight:1),(mean:33,weight:1),(mean:34,weight:1),(mean:35,weight:1),(mean:36,weight:1),(mean:37,weight:1),(mean:38,weight:1),(mean:39,weight:1),(mean:40,weight:1),(mean:41,weight:1),(mean:42,weight:1),(mean:43,weight:1),(mean:44,weight:1),(mean:45,weight:1),(mean:46,weight:1),(mean:47,weight:1),(mean:48,weight:1),(mean:49,weight:1),(mean:50,weight:1),(mean:51,weight:1),(mean:52.5,weight:2),(mean:54.5,weight:2),(mean:56.5,weight:2),(mean:58.5,weight:2),(mean:60.5,weight:2),(mean:62.5,weight:2),(mean:64,weight:1),(mean:65.5,weight:2),(mean:67.5,weight:2),(mean:69,weight:1),(mean:70.5,weight:2),(mean:72,weight:1),(mean:73.5,weight:2),(mean:75,weight:1),(mean:76,weight:1),(mean:77.5,weight:2),(mean:79,weight:1),(mean:80,weight:1),(mean:81.5,weight:2),(mean:83,weight:1),(mean:84,weight:1),(mean:85,weight:1),(mean:86,weight:1),(mean:87,weight:1),(mean:88,weight:1),(mean:89,weight:1),(mean:90,weight:1),(mean:91,weight:1),(mean:92,weight:1),(mean:93,weight:1),(mean:94,weight:1),(mean:95,weight:1),(mean:96,weight:1),(mean:97,weight:1),(mean:98,weight:1),(mean:99,weight:1),(mean:100,weight:1)])"; assert_eq!(output, Some(expected.into())); let estimate = client - .select( + .update( &format!("SELECT approx_percentile(0.90, '{}'::tdigest)", expected), None, None, ) + .unwrap() .first() - .get_one(); + .get_one() + .unwrap(); assert_eq!(estimate, Some(90.5)); }); } @@ -655,66 +679,80 @@ mod tests { #[pg_test] fn test_tdigest_compound_agg() { - Spi::execute(|client| { - client.select( - "CREATE TABLE new_test (device INTEGER, value DOUBLE PRECISION)", - None, - None, - ); - client.select("INSERT INTO new_test SELECT dev, dev - v FROM generate_series(1,10) dev, generate_series(0, 1.0, 0.01) v", None, None); + Spi::connect(|mut client| { + client + .update( + "CREATE TABLE new_test (device INTEGER, value DOUBLE PRECISION)", + None, + None, + ) + .unwrap(); + client.update("INSERT INTO new_test SELECT dev, dev - v FROM generate_series(1,10) dev, generate_series(0, 1.0, 0.01) v", None, None).unwrap(); let sanity = client - .select("SELECT COUNT(*) FROM new_test", None, None) + .update("SELECT COUNT(*) FROM new_test", None, None) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(Some(1010), sanity); - client.select( - "CREATE VIEW digests AS \ + client + .update( + "CREATE VIEW digests AS \ SELECT device, tdigest(20, value) \ FROM new_test \ GROUP BY device", - None, - None, - ); + None, + None, + ) + .unwrap(); - client.select( - "CREATE VIEW composite AS \ + client + .update( + "CREATE VIEW composite AS \ SELECT tdigest(tdigest) \ FROM digests", - None, - None, - ); + None, + None, + ) + .unwrap(); - client.select( - "CREATE VIEW base AS \ + client + .update( + "CREATE VIEW base AS \ SELECT tdigest(20, value) \ FROM new_test", - None, - None, - ); + None, + None, + ) + .unwrap(); let value = client - .select( + .update( "SELECT \ approx_percentile(0.9, tdigest) \ FROM base", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); let test_value = client - .select( + .update( "SELECT \ approx_percentile(0.9, tdigest) \ FROM composite", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); apx_eql(test_value.unwrap(), value.unwrap(), 0.1); apx_eql(test_value.unwrap(), 9.0, 0.1); diff --git a/extension/src/time_vector.rs b/extension/src/time_vector.rs index f8d16710..e64905ef 100644 --- a/extension/src/time_vector.rs +++ b/extension/src/time_vector.rs @@ -528,48 +528,54 @@ mod tests { #[pg_test] pub fn test_unnest() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); - client.select( - "CREATE TABLE data(time TIMESTAMPTZ, value DOUBLE PRECISION)", - None, - None, - ); - client.select( - r#"INSERT INTO data VALUES + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); + client + .update( + "CREATE TABLE data(time TIMESTAMPTZ, value DOUBLE PRECISION)", + None, + None, + ) + .unwrap(); + client + .update( + r#"INSERT INTO data VALUES ('2020-1-1', 30.0), ('2020-1-2', 45.0), ('2020-1-3', NULL), ('2020-1-4', 55.5), ('2020-1-5', 10.0)"#, - None, - None, - ); + None, + None, + ) + .unwrap(); - let mut unnest = client.select( - "SELECT unnest(timevector(time, value))::TEXT FROM data", - None, - None, - ); + let mut unnest = client + .update( + "SELECT unnest(timevector(time, value))::TEXT FROM data", + None, + None, + ) + .unwrap(); assert_eq!( - unnest.next().unwrap()[1].value(), + unnest.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-01 00:00:00+00\",30)") ); assert_eq!( - unnest.next().unwrap()[1].value(), + unnest.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-02 00:00:00+00\",45)") ); assert_eq!( - unnest.next().unwrap()[1].value(), + unnest.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-03 00:00:00+00\",NaN)") ); assert_eq!( - unnest.next().unwrap()[1].value(), + unnest.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-04 00:00:00+00\",55.5)") ); assert_eq!( - unnest.next().unwrap()[1].value(), + unnest.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-05 00:00:00+00\",10)") ); assert!(unnest.next().is_none()); @@ -578,43 +584,49 @@ mod tests { #[pg_test] pub fn test_format_timevector() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); - client.select( - "CREATE TABLE data(time TIMESTAMPTZ, value DOUBLE PRECISION)", - None, - None, - ); - client.select( - r#"INSERT INTO data VALUES + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); + client + .update( + "CREATE TABLE data(time TIMESTAMPTZ, value DOUBLE PRECISION)", + None, + None, + ) + .unwrap(); + client + .update( + r#"INSERT INTO data VALUES ('2020-1-1', 30.0), ('2020-1-2', 45.0), ('2020-1-3', NULL), ('2020-1-4', 55.5), ('2020-1-5', 10.0)"#, - None, - None, - ); + None, + None, + ) + .unwrap(); let test_plotly_template = client - .select( + .update( "SELECT toolkit_experimental.to_plotly(timevector(time, value)) FROM data", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() .unwrap(); assert_eq!(test_plotly_template, "{\"times\": [\"2020-01-01 00:00:00+00\",\"2020-01-02 00:00:00+00\",\"2020-01-03 00:00:00+00\",\"2020-01-04 00:00:00+00\",\"2020-01-05 00:00:00+00\"], \"vals\": [\"30\",\"45\",\"null\",\"55.5\",\"10\"]}" ); - let test_paired_timevals_template = client.select( + let test_paired_timevals_template = client.update( "SELECT toolkit_experimental.to_text(timevector(time, value),'{{TIMEVALS}}') FROM data", None, None, - ).first() - .get_one::() + ).unwrap().first() + .get_one::().unwrap() .unwrap(); assert_eq!( @@ -622,23 +634,23 @@ mod tests { ); let test_user_supplied_template = client - .select( + .update( "SELECT toolkit_experimental.to_text(timevector(time,value), '{\"times\": {{ TIMES }}, \"vals\": {{ VALUES }}}') FROM data", None, None, ) - .first() - .get_one::() + .unwrap().first() + .get_one::().unwrap() .unwrap(); assert_eq!( test_user_supplied_template,"{\"times\": [2020-01-01 00:00:00+00, 2020-01-02 00:00:00+00, 2020-01-03 00:00:00+00, 2020-01-04 00:00:00+00, 2020-01-05 00:00:00+00], \"vals\": [30, 45, null, 55.5, 10]}" ); - let test_user_supplied_json_template = client.select( + let test_user_supplied_json_template = client.update( "SELECT toolkit_experimental.to_text(timevector(time, value),'{\"times\": {{ TIMES | json_encode() | safe }}, \"vals\": {{ VALUES | json_encode() | safe }}}') FROM data", None, None, - ).first() - .get_one::() + ).unwrap().first() + .get_one::().unwrap() .unwrap(); assert_eq!( @@ -651,33 +663,39 @@ mod tests { #[should_panic = "All values in the series must be finite"] #[pg_test] pub fn test_format_timevector_panics_on_infinities() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); - client.select( - "CREATE TABLE data(time TIMESTAMPTZ, value DOUBLE PRECISION)", - None, - None, - ); - client.select( - r#"INSERT INTO data VALUES + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); + client + .update( + "CREATE TABLE data(time TIMESTAMPTZ, value DOUBLE PRECISION)", + None, + None, + ) + .unwrap(); + client + .update( + r#"INSERT INTO data VALUES ('2020-1-1', 30.0), ('2020-1-2', 45.0), ('2020-1-3', NULL), ('2020-1-4', 55.5), ('2020-1-6', 'Infinity'), ('2020-1-5', 10.0)"#, - None, - None, - ); + None, + None, + ) + .unwrap(); let test_plotly_template = client - .select( + .update( "SELECT toolkit_experimental.to_plotly(timevector(time, value)) FROM data", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() .unwrap(); assert_eq!(test_plotly_template,"{\"times\": [\n \"2020-01-01 00:00:00+00\",\n \"2020-01-02 00:00:00+00\",\n \"2020-01-03 00:00:00+00\",\n \"2020-01-04 00:00:00+00\",\n \"2020-01-05 00:00:00+00\"\n], \"vals\": [\n \"30\",\n \"45\",\n \"null\",\n \"55.5\",\n \"10\"\n]}" @@ -687,57 +705,65 @@ mod tests { #[pg_test] pub fn timevector_io() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); - client.select( - "CREATE TABLE data(time TIMESTAMPTZ, value DOUBLE PRECISION)", - None, - None, - ); - client.select( - r#"INSERT INTO data VALUES + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); + client + .update( + "CREATE TABLE data(time TIMESTAMPTZ, value DOUBLE PRECISION)", + None, + None, + ) + .unwrap(); + client + .update( + r#"INSERT INTO data VALUES ('2020-1-1', 30.0), ('2020-1-2', 45.0), ('2020-1-3', NULL), ('2020-1-4', 55.5), ('2020-1-5', 10.0)"#, - None, - None, - ); + None, + None, + ) + .unwrap(); let tvec = client - .select("SELECT timevector(time,value)::TEXT FROM data", None, None) + .update("SELECT timevector(time,value)::TEXT FROM data", None, None) + .unwrap() .first() .get_one::() + .unwrap() .unwrap(); let expected = r#"(version:1,num_points:5,flags:3,internal_padding:(0,0,0),points:[(ts:"2020-01-01 00:00:00+00",val:30),(ts:"2020-01-02 00:00:00+00",val:45),(ts:"2020-01-03 00:00:00+00",val:NaN),(ts:"2020-01-04 00:00:00+00",val:55.5),(ts:"2020-01-05 00:00:00+00",val:10)],null_val:[4])"#; assert_eq!(tvec, expected); - let mut unnest = client.select( - &format!("SELECT unnest('{}'::timevector_tstz_f64)::TEXT", expected), - None, - None, - ); + let mut unnest = client + .update( + &format!("SELECT unnest('{}'::timevector_tstz_f64)::TEXT", expected), + None, + None, + ) + .unwrap(); assert_eq!( - unnest.next().unwrap()[1].value(), + unnest.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-01 00:00:00+00\",30)") ); assert_eq!( - unnest.next().unwrap()[1].value(), + unnest.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-02 00:00:00+00\",45)") ); assert_eq!( - unnest.next().unwrap()[1].value(), + unnest.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-03 00:00:00+00\",NaN)") ); assert_eq!( - unnest.next().unwrap()[1].value(), + unnest.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-04 00:00:00+00\",55.5)") ); assert_eq!( - unnest.next().unwrap()[1].value(), + unnest.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-05 00:00:00+00\",10)") ); assert!(unnest.next().is_none()); @@ -746,34 +772,42 @@ mod tests { #[pg_test] pub fn test_arrow_equivalence() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); - client.select( - "CREATE TABLE data(time TIMESTAMPTZ, value DOUBLE PRECISION)", - None, - None, - ); - client.select( - r#"INSERT INTO data VALUES + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); + client + .update( + "CREATE TABLE data(time TIMESTAMPTZ, value DOUBLE PRECISION)", + None, + None, + ) + .unwrap(); + client + .update( + r#"INSERT INTO data VALUES ('1-1-2020', 30.0), ('1-2-2020', 45.0), ('1-3-2020', NULL), ('1-4-2020', 55.5), ('1-5-2020', 10.0)"#, - None, - None, - ); + None, + None, + ) + .unwrap(); - let mut func = client.select( - "SELECT unnest(timevector(time, value))::TEXT FROM data", - None, - None, - ); - let mut op = client.select( - "SELECT (timevector(time, value) -> unnest())::TEXT FROM data", - None, - None, - ); + let mut func = client + .update( + "SELECT unnest(timevector(time, value))::TEXT FROM data", + None, + None, + ) + .unwrap(); + let mut op = client + .update( + "SELECT (timevector(time, value) -> unnest())::TEXT FROM data", + None, + None, + ) + .unwrap(); let mut test = true; while test { @@ -789,15 +823,18 @@ mod tests { #[pg_test] pub fn test_rollup() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); - client.select( - "CREATE TABLE data(time TIMESTAMPTZ, value DOUBLE PRECISION, bucket INTEGER)", - None, - None, - ); - client.select( - r#"INSERT INTO data VALUES + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); + client + .update( + "CREATE TABLE data(time TIMESTAMPTZ, value DOUBLE PRECISION, bucket INTEGER)", + None, + None, + ) + .unwrap(); + client + .update( + r#"INSERT INTO data VALUES ('2020-1-1', 30.0, 1), ('2020-1-2', 45.0, 1), ('2020-1-3', NULL, 2), @@ -806,52 +843,55 @@ mod tests { ('2020-1-6', 13.0, 3), ('2020-1-7', 71.0, 4), ('2020-1-8', 0.0, 4)"#, - None, - None, - ); + None, + None, + ) + .unwrap(); - let mut unnest = client.select( - "SELECT unnest(rollup(tvec))::TEXT + let mut unnest = client + .update( + "SELECT unnest(rollup(tvec))::TEXT FROM ( SELECT timevector(time, value) AS tvec FROM data GROUP BY bucket ORDER BY bucket ) s", - None, - None, - ); + None, + None, + ) + .unwrap(); assert_eq!( - unnest.next().unwrap()[1].value(), + unnest.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-01 00:00:00+00\",30)") ); assert_eq!( - unnest.next().unwrap()[1].value(), + unnest.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-02 00:00:00+00\",45)") ); assert_eq!( - unnest.next().unwrap()[1].value(), + unnest.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-03 00:00:00+00\",NaN)") ); assert_eq!( - unnest.next().unwrap()[1].value(), + unnest.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-04 00:00:00+00\",55.5)") ); assert_eq!( - unnest.next().unwrap()[1].value(), + unnest.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-05 00:00:00+00\",10)") ); assert_eq!( - unnest.next().unwrap()[1].value(), + unnest.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-06 00:00:00+00\",13)") ); assert_eq!( - unnest.next().unwrap()[1].value(), + unnest.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-07 00:00:00+00\",71)") ); assert_eq!( - unnest.next().unwrap()[1].value(), + unnest.next().unwrap()[1].value().unwrap(), Some("(\"2020-01-08 00:00:00+00\",0)") ); assert!(unnest.next().is_none()); @@ -860,46 +900,60 @@ mod tests { #[pg_test] fn test_rollup_preserves_nulls_flag() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); - client.select( - "CREATE TABLE tvecs (vector Timevector_TSTZ_F64)", - None, - None, - ); - client.select( - "INSERT INTO tvecs SELECT timevector('2020-1-1', 20)", - None, - None, - ); - client.select( - "INSERT INTO tvecs SELECT timevector('2020-1-2', 30)", - None, - None, - ); - client.select( - "INSERT INTO tvecs SELECT timevector('2020-1-3', 15)", - None, - None, - ); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); + client + .update( + "CREATE TABLE tvecs (vector Timevector_TSTZ_F64)", + None, + None, + ) + .unwrap(); + client + .update( + "INSERT INTO tvecs SELECT timevector('2020-1-1', 20)", + None, + None, + ) + .unwrap(); + client + .update( + "INSERT INTO tvecs SELECT timevector('2020-1-2', 30)", + None, + None, + ) + .unwrap(); + client + .update( + "INSERT INTO tvecs SELECT timevector('2020-1-3', 15)", + None, + None, + ) + .unwrap(); let tvec = client - .select("SELECT rollup(vector)::TEXT FROM tvecs", None, None) + .update("SELECT rollup(vector)::TEXT FROM tvecs", None, None) + .unwrap() .first() .get_one::() + .unwrap() .unwrap(); let expected = r#"(version:1,num_points:3,flags:1,internal_padding:(0,0,0),points:[(ts:"2020-01-01 00:00:00+00",val:20),(ts:"2020-01-02 00:00:00+00",val:30),(ts:"2020-01-03 00:00:00+00",val:15)],null_val:[0])"#; assert_eq!(tvec, expected); - client.select( - "INSERT INTO tvecs SELECT timevector('2019-1-4', NULL)", - None, - None, - ); + client + .update( + "INSERT INTO tvecs SELECT timevector('2019-1-4', NULL)", + None, + None, + ) + .unwrap(); let tvec = client - .select("SELECT rollup(vector)::TEXT FROM tvecs", None, None) + .update("SELECT rollup(vector)::TEXT FROM tvecs", None, None) + .unwrap() .first() .get_one::() + .unwrap() .unwrap(); let expected = r#"(version:1,num_points:4,flags:2,internal_padding:(0,0,0),points:[(ts:"2020-01-01 00:00:00+00",val:20),(ts:"2020-01-02 00:00:00+00",val:30),(ts:"2020-01-03 00:00:00+00",val:15),(ts:"2019-01-04 00:00:00+00",val:NaN)],null_val:[8])"#; assert_eq!(tvec, expected); @@ -908,11 +962,12 @@ mod tests { #[pg_test] fn test_asof_join() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); - let mut result = client.select( - "WITH s as ( + let mut result = client + .update( + "WITH s as ( SELECT timevector(time, value) AS v1 FROM (VALUES ('2022-10-1 1:00 UTC'::TIMESTAMPTZ, 20.0), @@ -928,20 +983,21 @@ mod tests { ) as v(time, value)) SELECT (v1 -> toolkit_experimental.asof(v2))::TEXT FROM s, t;", - None, - None, - ); + None, + None, + ) + .unwrap(); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(,15,\"2022-10-01 00:30:00+00\")") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(30,45,\"2022-10-01 02:00:00+00\")") ); assert_eq!( - result.next().unwrap()[1].value(), + result.next().unwrap()[1].value().unwrap(), Some("(40,60,\"2022-10-01 03:30:00+00\")") ); assert!(result.next().is_none()); @@ -950,10 +1006,10 @@ mod tests { #[pg_test(error = "both timevectors must be populated for an asof join")] fn test_asof_none() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); - client.select( + client.update( "WITH s as ( SELECT timevector(now(), 0) -> toolkit_experimental.filter($$ $value != 0 $$) AS empty), t as ( @@ -964,16 +1020,16 @@ mod tests { ('2022-10-1 3:30 UTC'::TIMESTAMPTZ, 60.0) ) as v(time, value)) SELECT (valid -> toolkit_experimental.asof(empty)) - FROM s, t;", None, None); + FROM s, t;", None, None).unwrap(); }) } #[pg_test(error = "both timevectors must be populated for an asof join")] fn test_none_asof() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); - client.select( + client.update( "WITH s as ( SELECT timevector(now(), 0) -> toolkit_experimental.filter($$ $value != 0 $$) AS empty), t as ( @@ -984,7 +1040,7 @@ mod tests { ('2022-10-1 3:30 UTC'::TIMESTAMPTZ, 60.0) ) as v(time, value)) SELECT (empty -> toolkit_experimental.asof(valid)) - FROM s, t;", None, None); + FROM s, t;", None, None).unwrap(); }) } } diff --git a/extension/src/time_vector/pipeline.rs b/extension/src/time_vector/pipeline.rs index a1d0eb19..63fe07e0 100644 --- a/extension/src/time_vector/pipeline.rs +++ b/extension/src/time_vector/pipeline.rs @@ -166,8 +166,12 @@ pub fn arrow_add_unstable_element<'p>( )] pub unsafe fn pipeline_support(input: Internal) -> Internal { pipeline_support_helper(input, |old_pipeline, new_element| { - let new_element = - UnstableTimevectorPipeline::from_polymorphic_datum(new_element, false, 0).unwrap(); + let new_element = UnstableTimevectorPipeline::from_polymorphic_datum( + new_element, + false, + pg_sys::Oid::INVALID, + ) + .unwrap(); arrow_add_unstable_element(old_pipeline, new_element) .into_datum() .unwrap() @@ -252,9 +256,12 @@ pub(crate) unsafe fn pipeline_support_helper( let new_element_const: *mut pg_sys::Const = arg2.cast(); - let old_pipeline = - UnstableTimevectorPipeline::from_polymorphic_datum((*old_const).constvalue, false, 0) - .unwrap(); + let old_pipeline = UnstableTimevectorPipeline::from_polymorphic_datum( + (*old_const).constvalue, + false, + pg_sys::Oid::INVALID, + ) + .unwrap(); let new_pipeline = make_new_pipeline(old_pipeline, (*new_element_const).constvalue); let new_const = pg_sys::palloc(size_of::()).cast(); @@ -315,27 +322,33 @@ mod tests { #[pg_test] fn test_pipeline_lttb() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); - client.select( - "CREATE TABLE lttb_pipe (series timevector_tstz_f64)", - None, - None, - ); - client.select( + client + .update( + "CREATE TABLE lttb_pipe (series timevector_tstz_f64)", + None, + None, + ) + .unwrap(); + client.update( "INSERT INTO lttb_pipe \ SELECT timevector(time, val) FROM ( \ SELECT \ @@ -345,16 +358,18 @@ mod tests { ) bar", None, None - ); + ).unwrap(); let val = client - .select( + .update( "SELECT (series -> lttb(17))::TEXT FROM lttb_pipe", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:17,flags:1,internal_padding:(0,0,0),points:[\ @@ -379,13 +394,15 @@ mod tests { ); let val = client - .select( + .update( "SELECT (series -> lttb(8))::TEXT FROM lttb_pipe", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:8,flags:1,internal_padding:(0,0,0),points:[\ @@ -401,13 +418,15 @@ mod tests { ); let val = client - .select( + .update( "SELECT (series -> lttb(8) -> lttb(8))::TEXT FROM lttb_pipe", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:8,flags:1,internal_padding:(0,0,0),points:[\ @@ -423,13 +442,15 @@ mod tests { ); let val = client - .select( + .update( "SELECT (series -> (lttb(8) -> lttb(8) -> lttb(8)))::TEXT FROM lttb_pipe", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:8,flags:1,internal_padding:(0,0,0),points:[\ @@ -448,29 +469,33 @@ mod tests { #[pg_test] fn test_pipeline_folding() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); - let output = client.select( + let output = client.update( "EXPLAIN (verbose) SELECT timevector('2021-01-01'::timestamptz, 0.1) -> round() -> abs() -> round();", None, None - ).nth(1) + ).unwrap().nth(1) .unwrap() - .by_ordinal(1).unwrap() - .value::().unwrap(); + .get_datum_by_ordinal(1).unwrap() + .value::().unwrap().unwrap(); // check that it's executing as if we had input `timevector -> (round() -> abs())` assert_eq!(output.trim(), "Output: \ arrow_run_pipeline(\ diff --git a/extension/src/time_vector/pipeline/aggregation.rs b/extension/src/time_vector/pipeline/aggregation.rs index bf2a46fa..bfc7f551 100644 --- a/extension/src/time_vector/pipeline/aggregation.rs +++ b/extension/src/time_vector/pipeline/aggregation.rs @@ -164,7 +164,8 @@ pub fn pipeline_stats_agg() -> toolkit_experimental::PipelineThenStatsAgg<'stati pub unsafe fn pipeline_stats_agg_support(input: Internal) -> Internal { pipeline_support_helper(input, |old_pipeline, new_element| { let new_element = - PipelineThenStatsAgg::from_polymorphic_datum(new_element, false, 0).unwrap(); + PipelineThenStatsAgg::from_polymorphic_datum(new_element, false, pg_sys::Oid::INVALID) + .unwrap(); finalize_with_stats_agg(old_pipeline, new_element) .into_datum() .unwrap() @@ -258,7 +259,9 @@ pub fn finalize_with_sum<'e>( #[pg_extern(immutable, parallel_safe, schema = "toolkit_experimental")] pub unsafe fn pipeline_sum_support(input: Internal) -> Internal { pipeline_support_helper(input, |old_pipeline, new_element| { - let new_element = PipelineThenSum::from_polymorphic_datum(new_element, false, 0).unwrap(); + let new_element = + PipelineThenSum::from_polymorphic_datum(new_element, false, pg_sys::Oid::INVALID) + .unwrap(); finalize_with_sum(old_pipeline, new_element) .into_datum() .unwrap() @@ -349,7 +352,8 @@ pub fn finalize_with_average<'e>( pub unsafe fn pipeline_average_support(input: Internal) -> Internal { pipeline_support_helper(input, |old_pipeline, new_element| { let new_element = - PipelineThenAverage::from_polymorphic_datum(new_element, false, 0).unwrap(); + PipelineThenAverage::from_polymorphic_datum(new_element, false, pg_sys::Oid::INVALID) + .unwrap(); finalize_with_average(old_pipeline, new_element) .into_datum() .unwrap() @@ -437,7 +441,8 @@ pub fn finalize_with_num_vals<'e>( pub unsafe fn pipeline_num_vals_support(input: Internal) -> Internal { pipeline_support_helper(input, |old_pipeline, new_element| { let new_element = - PipelineThenNumVals::from_polymorphic_datum(new_element, false, 0).unwrap(); + PipelineThenNumVals::from_polymorphic_datum(new_element, false, pg_sys::Oid::INVALID) + .unwrap(); finalize_with_num_vals(old_pipeline, new_element) .into_datum() .unwrap() @@ -520,8 +525,12 @@ pub fn pipeline_counter_agg() -> toolkit_experimental::PipelineThenCounterAgg<'s #[pg_extern(immutable, parallel_safe, schema = "toolkit_experimental")] pub unsafe fn pipeline_counter_agg_support(input: Internal) -> Internal { pipeline_support_helper(input, |old_pipeline, new_element| { - let new_element = - PipelineThenCounterAgg::from_polymorphic_datum(new_element, false, 0).unwrap(); + let new_element = PipelineThenCounterAgg::from_polymorphic_datum( + new_element, + false, + pg_sys::Oid::INVALID, + ) + .unwrap(); finalize_with_counter_agg(old_pipeline, new_element) .into_datum() .unwrap() @@ -548,7 +557,7 @@ pub fn arrow_run_pipeline_then_hyperloglog<'a>( timevector = run_pipeline_elements(timevector, pipeline.elements.iter()); HyperLogLog::build_from( pipeline.hll_size as i32, - PgBuiltInOids::FLOAT8OID as u32, + PgBuiltInOids::FLOAT8OID.into(), None, timevector .iter() @@ -604,8 +613,12 @@ pub fn pipeline_hyperloglog(size: i32) -> toolkit_experimental::PipelineThenHype #[pg_extern(immutable, parallel_safe, schema = "toolkit_experimental")] pub unsafe fn pipeline_hyperloglog_support(input: Internal) -> Internal { pipeline_support_helper(input, |old_pipeline, new_element| { - let new_element = - PipelineThenHyperLogLog::from_polymorphic_datum(new_element, false, 0).unwrap(); + let new_element = PipelineThenHyperLogLog::from_polymorphic_datum( + new_element, + false, + pg_sys::Oid::INVALID, + ) + .unwrap(); finalize_with_hyperloglog(old_pipeline, new_element) .into_datum() .unwrap() @@ -678,8 +691,12 @@ pub fn pipeline_percentile_agg() -> toolkit_experimental::PipelineThenPercentile #[pg_extern(immutable, parallel_safe, schema = "toolkit_experimental")] pub unsafe fn pipeline_percentile_agg_support(input: Internal) -> Internal { pipeline_support_helper(input, |old_pipeline, new_element| { - let new_element = - PipelineThenPercentileAgg::from_polymorphic_datum(new_element, false, 0).unwrap(); + let new_element = PipelineThenPercentileAgg::from_polymorphic_datum( + new_element, + false, + pg_sys::Oid::INVALID, + ) + .unwrap(); finalize_with_percentile_agg(old_pipeline, new_element) .into_datum() .unwrap() @@ -705,20 +722,24 @@ mod tests { #[pg_test] fn test_stats_agg_finalizer() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); // we use a subselect to guarantee order let create_series = "SELECT timevector(time, value) as series FROM \ @@ -729,7 +750,7 @@ mod tests { ('2020-01-05 UTC'::TIMESTAMPTZ, 30.0)) as v(time, value)"; let val = client - .select( + .update( &format!( "SELECT (series -> stats_agg())::TEXT FROM ({}) s", create_series @@ -737,8 +758,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,n:5,sx:100,sx2:250,sx3:0,sx4:21250)" @@ -748,23 +771,27 @@ mod tests { #[pg_test] fn test_stats_agg_pipeline_folding() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); let output = client - .select( + .update( "EXPLAIN (verbose) SELECT \ timevector('1930-04-05'::timestamptz, 123.0) \ -> ceil() -> abs() -> floor() \ @@ -772,11 +799,13 @@ mod tests { None, None, ) + .unwrap() .nth(1) .unwrap() - .by_ordinal(1) + .get_datum_by_ordinal(1) .unwrap() .value::() + .unwrap() .unwrap(); assert_eq!(output.trim(), "Output: (\ arrow_run_pipeline_then_stats_agg(\ @@ -792,20 +821,24 @@ mod tests { #[pg_test] fn test_sum_finalizer() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); // we use a subselect to guarantee order let create_series = "SELECT timevector(time, value) as series FROM \ @@ -816,36 +849,42 @@ mod tests { ('2020-01-05 UTC'::TIMESTAMPTZ, 30.0)) as v(time, value)"; let val = client - .select( + .update( &format!("SELECT (series -> sum())::TEXT FROM ({}) s", create_series), None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(val.unwrap(), "100"); }); } #[pg_test] fn test_sum_pipeline_folding() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); let output = client - .select( + .update( "EXPLAIN (verbose) SELECT \ timevector('1930-04-05'::timestamptz, 123.0) \ -> ceil() -> abs() -> floor() \ @@ -853,11 +892,13 @@ mod tests { None, None, ) + .unwrap() .nth(1) .unwrap() - .by_ordinal(1) + .get_datum_by_ordinal(1) .unwrap() .value::() + .unwrap() .unwrap(); assert_eq!(output.trim(), "Output: \ arrow_pipeline_then_sum(\ @@ -873,20 +914,24 @@ mod tests { #[pg_test] fn test_average_finalizer() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); // we use a subselect to guarantee order let create_series = "SELECT timevector(time, value) as series FROM \ @@ -897,7 +942,7 @@ mod tests { ('2020-01-05 UTC'::TIMESTAMPTZ, 30.0)) as v(time, value)"; let val = client - .select( + .update( &format!( "SELECT (series -> average())::TEXT FROM ({}) s", create_series @@ -905,31 +950,37 @@ mod tests { None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(val.unwrap(), "20"); }); } #[pg_test] fn test_average_pipeline_folding() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); let output = client - .select( + .update( "EXPLAIN (verbose) SELECT \ timevector('1930-04-05'::timestamptz, 123.0) \ -> ceil() -> abs() -> floor() \ @@ -937,11 +988,13 @@ mod tests { None, None, ) + .unwrap() .nth(1) .unwrap() - .by_ordinal(1) + .get_datum_by_ordinal(1) .unwrap() .value::() + .unwrap() .unwrap(); assert_eq!(output.trim(), "Output: \ arrow_pipeline_then_average(\ @@ -957,20 +1010,24 @@ mod tests { #[pg_test] fn test_num_vals_finalizer() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); // we use a subselect to guarantee order let create_series = "SELECT timevector(time, value) as series FROM \ @@ -981,7 +1038,7 @@ mod tests { ('2020-01-05 UTC'::TIMESTAMPTZ, 30.0)) as v(time, value)"; let val = client - .select( + .update( &format!( "SELECT (series -> num_vals())::TEXT FROM ({}) s", create_series @@ -989,31 +1046,37 @@ mod tests { None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(val.unwrap(), "5"); }); } #[pg_test] fn test_num_vals_pipeline_folding() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); let output = client - .select( + .update( "EXPLAIN (verbose) SELECT \ timevector('1930-04-05'::timestamptz, 123.0) \ -> ceil() -> abs() -> floor() \ @@ -1021,11 +1084,13 @@ mod tests { None, None, ) + .unwrap() .nth(1) .unwrap() - .by_ordinal(1) + .get_datum_by_ordinal(1) .unwrap() .value::() + .unwrap() .unwrap(); assert_eq!(output.trim(), "Output: \ arrow_pipeline_then_num_vals(\ @@ -1041,20 +1106,24 @@ mod tests { #[pg_test] fn test_counter_agg_finalizer() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); // we use a subselect to guarantee order let create_series = "SELECT timevector(time, value) as series FROM \ @@ -1065,7 +1134,7 @@ mod tests { ('2020-01-05 UTC'::TIMESTAMPTZ, 30.0)) as v(time, value)"; let val = client - .select( + .update( &format!( "SELECT (series -> sort() -> counter_agg())::TEXT FROM ({}) s", create_series @@ -1073,21 +1142,23 @@ mod tests { None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(val.unwrap(), "(version:1,stats:(n:5,sx:3156624000,sx2:74649600000,sx3:0,sx4:1894671345254400000000,sy:215,sy2:2280,sy3:6720.000000000007,sy4:1788960,sxy:12960000),first:(ts:\"2020-01-01 00:00:00+00\",val:15),second:(ts:\"2020-01-02 00:00:00+00\",val:25),penultimate:(ts:\"2020-01-04 00:00:00+00\",val:10),last:(ts:\"2020-01-05 00:00:00+00\",val:30),reset_sum:45,num_resets:2,num_changes:4,bounds:(is_present:0,has_left:0,has_right:0,padding:(0,0,0,0,0),left:None,right:None))"); - let val = client.select( + let val = client.update( &format!("SELECT series -> sort() -> counter_agg() -> with_bounds('[2020-01-01 UTC, 2020-02-01 UTC)') -> extrapolated_delta('prometheus') FROM ({}) s", create_series), None, None ) - .first() - .get_one::().unwrap(); + .unwrap().first() + .get_one::().unwrap().unwrap(); assert!((val - 67.5).abs() < f64::EPSILON); let output = client - .select( + .update( "EXPLAIN (verbose) SELECT \ timevector('1930-04-05'::timestamptz, 123.0) \ -> ceil() -> abs() -> floor() \ @@ -1095,11 +1166,13 @@ mod tests { None, None, ) + .unwrap() .nth(1) .unwrap() - .by_ordinal(1) + .get_datum_by_ordinal(1) .unwrap() .value::() + .unwrap() .unwrap(); assert_eq!(output.trim(), "Output: \ arrow_run_pipeline_then_counter_agg(\ @@ -1115,20 +1188,24 @@ mod tests { #[pg_test] fn test_hyperloglog_finalizer() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); // we use a subselect to guarantee order let create_series = "SELECT timevector(time, value) as series FROM \ @@ -1144,7 +1221,7 @@ mod tests { ('2020-01-10 UTC'::TIMESTAMPTZ, 5.0)) as v(time, value)"; let val = client - .select( + .update( &format!( "SELECT (series -> hyperloglog(100))::TEXT FROM ({}) s", create_series @@ -1152,12 +1229,14 @@ mod tests { None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(val.unwrap(), "(version:1,log:Sparse(num_compressed:7,element_type:FLOAT8,collation:None,compressed_bytes:28,precision:7,compressed:[136,188,20,7,8,30,244,43,72,69,89,2,72,255,97,27,72,83,248,27,200,110,35,5,8,37,85,12]))"); let val = client - .select( + .update( &format!( "SELECT series -> hyperloglog(100) -> distinct_count() FROM ({}) s", create_series @@ -1165,13 +1244,15 @@ mod tests { None, None, ) + .unwrap() .first() - .get_one::() + .get_one::() + .unwrap() .unwrap(); assert_eq!(val, 7); let output = client - .select( + .update( "EXPLAIN (verbose) SELECT \ timevector('1930-04-05'::timestamptz, 123.0) \ -> ceil() -> abs() -> floor() \ @@ -1179,11 +1260,13 @@ mod tests { None, None, ) + .unwrap() .nth(1) .unwrap() - .by_ordinal(1) + .get_datum_by_ordinal(1) .unwrap() .value::() + .unwrap() .unwrap(); assert_eq!(output.trim(), "Output: \ arrow_run_pipeline_then_hyperloglog(\ @@ -1199,20 +1282,24 @@ mod tests { #[pg_test] fn test_percentile_agg_finalizer() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); // we use a subselect to guarantee order let create_series = "SELECT timevector(time, value) as series FROM \ @@ -1223,7 +1310,7 @@ mod tests { ('2020-01-05 UTC'::TIMESTAMPTZ, 30.0)) as v(time, value)"; let val = client - .select( + .update( &format!( "SELECT (series -> percentile_agg())::TEXT FROM ({}) s", create_series @@ -1231,8 +1318,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,\ @@ -1256,23 +1345,27 @@ mod tests { #[pg_test] fn test_percentile_agg_pipeline_folding() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); let output = client - .select( + .update( "EXPLAIN (verbose) SELECT \ timevector('1930-04-05'::timestamptz, 123.0) \ -> ceil() -> abs() -> floor() \ @@ -1280,11 +1373,13 @@ mod tests { None, None, ) + .unwrap() .nth(1) .unwrap() - .by_ordinal(1) + .get_datum_by_ordinal(1) .unwrap() .value::() + .unwrap() .unwrap(); assert_eq!(output.trim(), "Output: \ arrow_run_pipeline_then_percentile_agg(\ diff --git a/extension/src/time_vector/pipeline/arithmetic.rs b/extension/src/time_vector/pipeline/arithmetic.rs index 3a26d65c..a6659b54 100644 --- a/extension/src/time_vector/pipeline/arithmetic.rs +++ b/extension/src/time_vector/pipeline/arithmetic.rs @@ -293,20 +293,24 @@ mod tests { #[pg_test] fn test_simple_arith_binops() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); // we use a subselect to guarantee order let create_series = "SELECT timevector(time, value) as series FROM \ @@ -317,7 +321,7 @@ mod tests { ('2020-01-05 UTC'::TIMESTAMPTZ, 30.0)) as v(time, value)"; let val = client - .select( + .update( &format!( "SELECT (series -> add(1.0))::TEXT FROM ({}) s", create_series @@ -325,8 +329,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:5,flags:0,internal_padding:(0,0,0),points:[\ @@ -339,7 +345,7 @@ mod tests { ); let val = client - .select( + .update( &format!( "SELECT (series -> sub(3.0))::TEXT FROM ({}) s", create_series @@ -347,8 +353,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:5,flags:0,internal_padding:(0,0,0),points:[\ @@ -361,7 +369,7 @@ mod tests { ); let val = client - .select( + .update( &format!( "SELECT (series -> mul(2.0))::TEXT FROM ({}) s", create_series @@ -369,8 +377,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:5,flags:0,internal_padding:(0,0,0),points:[\ @@ -383,7 +393,7 @@ mod tests { ); let val = client - .select( + .update( &format!( "SELECT (series -> div(5.0))::TEXT FROM ({}) s", create_series @@ -391,8 +401,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:5,flags:0,internal_padding:(0,0,0),points:[\ @@ -405,7 +417,7 @@ mod tests { ); let val = client - .select( + .update( &format!( "SELECT (series -> mod(5.0))::TEXT FROM ({}) s", create_series @@ -413,8 +425,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:5,flags:0,internal_padding:(0,0,0),points:[\ @@ -427,7 +441,7 @@ mod tests { ); let val = client - .select( + .update( &format!( "SELECT (series -> power(2.0))::TEXT FROM ({}) s", create_series @@ -435,8 +449,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:5,flags:0,internal_padding:(0,0,0),points:[\ @@ -449,7 +465,7 @@ mod tests { ); let val = client - .select( + .update( &format!( "SELECT (series -> logn(10.0))::TEXT FROM ({}) s", create_series @@ -457,8 +473,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:5,flags:0,internal_padding:(0,0,0),points:[\ @@ -474,20 +492,24 @@ mod tests { #[pg_test] fn test_simple_arith_unaryops() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); // we use a subselect to guarantee order let create_series = "SELECT timevector(time, value) as series FROM \ @@ -498,13 +520,15 @@ mod tests { ('2020-01-05 UTC'::TIMESTAMPTZ, 30.3)) as v(time, value)"; let val = client - .select( + .update( &format!("SELECT (series -> abs())::TEXT FROM ({}) s", create_series), None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:5,flags:0,internal_padding:(0,0,0),points:[\ @@ -517,13 +541,13 @@ mod tests { ); // TODO re-enable once made stable - // let val = client.select( + // let val = client.update( // &format!("SELECT (series -> cbrt())::TEXT FROM ({}) s", create_series), // None, // None // ) // .first() - // .get_one::(); + // .get_one::().unwrap(); // assert_eq!(val.unwrap(), "[\ // (ts:\"2020-01-04 00:00:00+00\",val:2.943382658441668),\ // (ts:\"2020-01-01 00:00:00+00\",val:-2.161592332945083),\ @@ -533,13 +557,15 @@ mod tests { // ]"); let val = client - .select( + .update( &format!("SELECT (series -> ceil())::TEXT FROM ({}) s", create_series), None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:5,flags:0,internal_padding:(0,0,0),points:[\ @@ -552,7 +578,7 @@ mod tests { ); let val = client - .select( + .update( &format!( "SELECT (series -> floor())::TEXT FROM ({}) s", create_series @@ -560,8 +586,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:5,flags:0,internal_padding:(0,0,0),points:[\ @@ -576,13 +604,13 @@ mod tests { // TODO why are there `null`s here? // Josh - likely JSON can't represent nans correctly... // TODO re-enable once made stable - // let val = client.select( + // let val = client.update( // &format!("SELECT (series -> ln())::TEXT FROM ({}) s", create_series), // None, // None // ) // .first() - // .get_one::(); + // .get_one::().unwrap(); // assert_eq!(val.unwrap(), "[\ // (ts:\"2020-01-04 00:00:00+00\",val:3.2386784521643803),\ // (ts:\"2020-01-01 00:00:00+00\",val:null),\ @@ -592,13 +620,13 @@ mod tests { // ]"); // TODO re-enable once made stable - // let val = client.select( + // let val = client.update( // &format!("SELECT (series -> log10())::TEXT FROM ({}) s", create_series), // None, // None // ) // .first() - // .get_one::(); + // .get_one::().unwrap(); // assert_eq!(val.unwrap(), "[\ // (ts:\"2020-01-04 00:00:00+00\",val:1.4065401804339552),\ // (ts:\"2020-01-01 00:00:00+00\",val:null),\ @@ -608,7 +636,7 @@ mod tests { // ]"); let val = client - .select( + .update( &format!( "SELECT (series -> round())::TEXT FROM ({}) s", create_series @@ -616,8 +644,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:5,flags:0,internal_padding:(0,0,0),points:[\ @@ -630,13 +660,15 @@ mod tests { ); let val = client - .select( + .update( &format!("SELECT (series -> sign())::TEXT FROM ({}) s", create_series), None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:5,flags:0,internal_padding:(0,0,0),points:[\ @@ -649,13 +681,13 @@ mod tests { ); // TODO re-enable once made stable - // let val = client.select( + // let val = client.update( // &format!("SELECT (series -> sqrt())::TEXT FROM ({}) s", create_series), // None, // None // ) // .first() - // .get_one::(); + // .get_one::().unwrap(); // assert_eq!(val.unwrap(), "[\ // (ts:\"2020-01-04 00:00:00+00\",val:5.049752469181039),\ // (ts:\"2020-01-01 00:00:00+00\",val:null),\ @@ -665,7 +697,7 @@ mod tests { // ]"); let val = client - .select( + .update( &format!( "SELECT (series -> trunc())::TEXT FROM ({}) s", create_series @@ -673,8 +705,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:5,flags:0,internal_padding:(0,0,0),points:[\ diff --git a/extension/src/time_vector/pipeline/delta.rs b/extension/src/time_vector/pipeline/delta.rs index e64df043..c434d5a7 100644 --- a/extension/src/time_vector/pipeline/delta.rs +++ b/extension/src/time_vector/pipeline/delta.rs @@ -66,28 +66,35 @@ mod tests { #[pg_test] fn test_pipeline_delta() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); - client.select( - "CREATE TABLE series(time timestamptz, value double precision)", - None, - None, - ); - client.select( - "INSERT INTO series \ + client + .update( + "CREATE TABLE series(time timestamptz, value double precision)", + None, + None, + ) + .unwrap(); + client + .update( + "INSERT INTO series \ VALUES \ ('2020-01-01 UTC'::TIMESTAMPTZ, 10.0), \ ('2020-01-02 UTC'::TIMESTAMPTZ, 25.0), \ @@ -98,18 +105,21 @@ mod tests { ('2020-01-07 UTC'::TIMESTAMPTZ, 30.8), \ ('2020-01-08 UTC'::TIMESTAMPTZ, 30.9), \ ('2020-01-09 UTC'::TIMESTAMPTZ, -427.2)", - None, - None, - ); + None, + None, + ) + .unwrap(); let val = client - .select( + .update( "SELECT (timevector(time, value) -> delta())::TEXT FROM series", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:8,flags:1,internal_padding:(0,0,0),points:[\ diff --git a/extension/src/time_vector/pipeline/expansion.rs b/extension/src/time_vector/pipeline/expansion.rs index aaf7b34e..ab5bdac9 100644 --- a/extension/src/time_vector/pipeline/expansion.rs +++ b/extension/src/time_vector/pipeline/expansion.rs @@ -146,8 +146,12 @@ pub fn arrow_run_pipeline_then_materialize<'a>( #[pg_extern(immutable, parallel_safe, schema = "toolkit_experimental")] pub unsafe fn pipeline_materialize_support(input: pgx::Internal) -> pgx::Internal { pipeline_support_helper(input, |old_pipeline, new_element| { - let new_element = - PipelineForceMaterialize::from_polymorphic_datum(new_element, false, 0).unwrap(); + let new_element = PipelineForceMaterialize::from_polymorphic_datum( + new_element, + false, + pg_sys::Oid::INVALID, + ) + .unwrap(); arrow_force_materialize(old_pipeline, new_element) .into_datum() .unwrap() @@ -170,20 +174,24 @@ mod tests { #[pg_test] fn test_unnest_finalizer() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); // we use a subselect to guarantee order let create_series = "SELECT timevector(time, value) as series FROM \ @@ -194,7 +202,7 @@ mod tests { ('2020-01-05 UTC'::TIMESTAMPTZ, 30.0)) as v(time, value)"; let val = client - .select( + .update( &format!( "SELECT array_agg(val)::TEXT \ FROM (SELECT series -> unnest() as val FROM ({}) s) t", @@ -203,28 +211,34 @@ mod tests { None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(val.unwrap(), "{\"(\\\"2020-01-04 00:00:00+00\\\",25)\",\"(\\\"2020-01-01 00:00:00+00\\\",10)\",\"(\\\"2020-01-03 00:00:00+00\\\",20)\",\"(\\\"2020-01-02 00:00:00+00\\\",15)\",\"(\\\"2020-01-05 00:00:00+00\\\",30)\"}"); }); } #[pg_test] fn test_series_finalizer() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); // we use a subselect to guarantee order let create_series = "SELECT timevector(time, value) as series FROM \ @@ -235,7 +249,7 @@ mod tests { ('2020-01-05 UTC'::TIMESTAMPTZ, 31.0)) as v(time, value)"; let val = client - .select( + .update( &format!( "SELECT (series -> materialize())::TEXT FROM ({}) s", create_series @@ -243,8 +257,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:5,flags:0,internal_padding:(0,0,0),points:[\ @@ -260,25 +276,29 @@ mod tests { #[pg_test] fn test_force_materialize() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); // `-> materialize()` should force materialization, but otherwise the // pipeline-folding optimization should proceed let output = client - .select( + .update( "EXPLAIN (verbose) SELECT \ timevector('2021-01-01'::timestamptz, 0.1) \ -> round() -> abs() \ @@ -287,11 +307,13 @@ mod tests { None, None, ) + .unwrap() .nth(1) .unwrap() - .by_ordinal(1) + .get_datum_by_ordinal(1) .unwrap() .value::() + .unwrap() .unwrap(); assert_eq!(output.trim(), "Output: \ arrow_run_pipeline(\ diff --git a/extension/src/time_vector/pipeline/fill_to.rs b/extension/src/time_vector/pipeline/fill_to.rs index 136aaa1e..8b110123 100644 --- a/extension/src/time_vector/pipeline/fill_to.rs +++ b/extension/src/time_vector/pipeline/fill_to.rs @@ -146,45 +146,53 @@ mod tests { #[pg_test] fn test_pipeline_fill_to() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); - client.select( - "CREATE TABLE series(time timestamptz, value double precision)", - None, - None, - ); - client.select( - "INSERT INTO series \ + client + .update( + "CREATE TABLE series(time timestamptz, value double precision)", + None, + None, + ) + .unwrap(); + client + .update( + "INSERT INTO series \ VALUES \ ('2020-01-01 UTC'::TIMESTAMPTZ, 10.0), \ ('2020-01-03 UTC'::TIMESTAMPTZ, 20.0), \ ('2020-01-04 UTC'::TIMESTAMPTZ, 90.0), \ ('2020-01-06 UTC'::TIMESTAMPTZ, 30), \ ('2020-01-09 UTC'::TIMESTAMPTZ, 40.0)", - None, - None, - ); + None, + None, + ) + .unwrap(); - let val = client.select( + let val = client.update( "SELECT (timevector(time, value) -> fill_to('24 hours', 'locf'))::TEXT FROM series", None, None ) - .first() - .get_one::(); + .unwrap().first() + .get_one::().unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:9,flags:1,internal_padding:(0,0,0),points:[\ @@ -200,13 +208,13 @@ mod tests { ],null_val:[0,0])" ); - let val = client.select( + let val = client.update( "SELECT (timevector(time, value) -> fill_to('24 hours', 'linear'))::TEXT FROM series", None, None ) - .first() - .get_one::(); + .unwrap().first() + .get_one::().unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:9,flags:1,internal_padding:(0,0,0),points:[\ @@ -222,13 +230,13 @@ mod tests { ],null_val:[0,0])" ); - let val = client.select( + let val = client.update( "SELECT (timevector(time, value) -> fill_to('24 hours', 'nearest'))::TEXT FROM series", None, None ) - .first() - .get_one::(); + .unwrap().first() + .get_one::().unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:9,flags:1,internal_padding:(0,0,0),points:[\ @@ -244,13 +252,13 @@ mod tests { ],null_val:[0,0])" ); - let val = client.select( + let val = client.update( "SELECT (timevector(time, value) -> fill_to('10 hours', 'nearest'))::TEXT FROM series", None, None ) - .first() - .get_one::(); + .unwrap().first() + .get_one::().unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:22,flags:1,internal_padding:(0,0,0),points:[\ diff --git a/extension/src/time_vector/pipeline/filter.rs b/extension/src/time_vector/pipeline/filter.rs index 529e9e54..5d3605c2 100644 --- a/extension/src/time_vector/pipeline/filter.rs +++ b/extension/src/time_vector/pipeline/filter.rs @@ -64,46 +64,56 @@ mod tests { #[pg_test] fn test_pipeline_filter_lambda() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); - client.select( - "CREATE TABLE series(time timestamptz, value double precision)", - None, - None, - ); - client.select( - "INSERT INTO series \ + client + .update( + "CREATE TABLE series(time timestamptz, value double precision)", + None, + None, + ) + .unwrap(); + client + .update( + "INSERT INTO series \ VALUES \ ('2020-01-04 UTC'::TIMESTAMPTZ, 25.0), \ ('2020-01-01 UTC'::TIMESTAMPTZ, 10.0), \ ('2020-01-03 UTC'::TIMESTAMPTZ, 20.0), \ ('2020-01-02 UTC'::TIMESTAMPTZ, 15.0), \ ('2020-01-05 UTC'::TIMESTAMPTZ, 30.0)", - None, - None, - ); + None, + None, + ) + .unwrap(); let val = client - .select( + .update( "SELECT (timevector(time, value))::TEXT FROM series", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:5,flags:0,internal_padding:(0,0,0),points:[\ @@ -115,13 +125,13 @@ mod tests { ],null_val:[0])" ); - let val = client.select( + let val = client.update( "SELECT (timevector(time, value) -> filter($$ $time != '2020-01-05't and ($value = 10 or $value = 20) $$))::TEXT FROM series", None, None ) - .first() - .get_one::(); + .unwrap().first() + .get_one::().unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:2,flags:0,internal_padding:(0,0,0),points:[\ diff --git a/extension/src/time_vector/pipeline/lambda.rs b/extension/src/time_vector/pipeline/lambda.rs index 8ed9494f..dd9c7663 100644 --- a/extension/src/time_vector/pipeline/lambda.rs +++ b/extension/src/time_vector/pipeline/lambda.rs @@ -42,7 +42,7 @@ impl<'input> InOutFuncs for Lambda<'input> { } } - fn input(input: &pgx::cstr_core::CStr) -> Self + fn input(input: &std::ffi::CStr) -> Self where Self: Sized, { @@ -444,12 +444,13 @@ mod tests { macro_rules! trace_lambda { ($client: expr, $expr:literal) => { $client - .select( + .update( concat!("SELECT trace_lambda($$ ", $expr, " $$, '2021-01-01', 2.0)"), None, None, ) - .map(|r| r.by_ordinal(1).unwrap().value::().unwrap()) + .unwrap() + .map(|r| r.get::(1).unwrap().unwrap()) .collect() }; } @@ -457,7 +458,7 @@ mod tests { macro_rules! point_lambda { ($client: expr, $expr:literal) => { $client - .select( + .update( concat!( "SELECT point_lambda($$ ", $expr, @@ -466,16 +467,18 @@ mod tests { None, None, ) + .unwrap() .first() .get_one::() .unwrap() + .unwrap() }; } macro_rules! interval_lambda { ($client: expr, $expr:literal) => { $client - .select( + .update( concat!( "SELECT interval_lambda($$ ", $expr, @@ -484,37 +487,43 @@ mod tests { None, None, ) + .unwrap() .first() .get_one::() .unwrap() + .unwrap() }; } macro_rules! f64_lambda { ($client: expr, $expr:literal) => { $client - .select( + .update( concat!("SELECT f64_lambda($$ ", $expr, " $$, now(), 2.0)"), None, None, ) + .unwrap() .first() .get_one::() .unwrap() + .unwrap() }; } macro_rules! bool_lambda { ($client: expr, $expr:literal) => { $client - .select( + .update( concat!("SELECT bool_lambda($$ ", $expr, " $$, now(), 2.0)::text"), None, None, ) + .unwrap() .first() .get_one::() .unwrap() + .unwrap() }; } @@ -544,82 +553,100 @@ mod tests { #[pg_test] fn test_lambda_general() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); - client.select( - "SELECT $$ let $1 = 1.0; 2.0, $1 $$::toolkit_experimental.lambda", - None, - None, - ); - // client.select("SELECT $$ '1 day'i $$::toolkit_experimental.lambda", None, None); - // client.select("SELECT $$ '2020-01-01't $$::toolkit_experimental.lambda", None, None); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) + .unwrap(); + client + .update( + "SELECT $$ let $1 = 1.0; 2.0, $1 $$::toolkit_experimental.lambda", + None, + None, + ) + .unwrap(); + // client.update("SELECT $$ '1 day'i $$::toolkit_experimental.lambda", None, None).unwrap(); + // client.update("SELECT $$ '2020-01-01't $$::toolkit_experimental.lambda", None, None).unwrap(); let res = client - .select("SELECT f64_lambda($$ 1.0 $$, now(), 0.0)::text", None, None) + .update("SELECT f64_lambda($$ 1.0 $$, now(), 0.0)::text", None, None) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(&*res.unwrap(), "1"); let res = client - .select( + .update( "SELECT f64_lambda($$ 1.0 + 1.0 $$, now(), 0.0)::text", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(&*res.unwrap(), "2"); let res = client - .select( + .update( "SELECT f64_lambda($$ 1.0 - 1.0 $$, now(), 0.0)::text", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(&*res.unwrap(), "0"); let res = client - .select( + .update( "SELECT f64_lambda($$ 2.0 * 3.0 $$, now(), 0.0)::text", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(&*res.unwrap(), "6"); let res = client - .select( + .update( "SELECT f64_lambda($$ $value + 3.0 $$, now(), 2.0)::text", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(&*res.unwrap(), "5"); let res = client - .select( + .update( "SELECT f64_lambda($$ 3.0 - 1.0 * 3.0 $$, now(), 2.0)::text", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(&*res.unwrap(), "0"); bool_lambda_eq!(client, "3.0 = 3.0", "true"); @@ -629,43 +656,51 @@ mod tests { bool_lambda_eq!(client, "2.0 != 3.0 and (1 = 1)", "true"); let res = client - .select( + .update( "SELECT ttz_lambda($$ '2020-11-22 13:00:01't $$, now(), 2.0)::text", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(&*res.unwrap(), "2020-11-22 13:00:01+00"); let res = client - .select( + .update( "SELECT ttz_lambda($$ $time $$, '1930-01-12 14:20:21', 2.0)::text", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(&*res.unwrap(), "1930-01-12 14:20:21+00"); let res = client - .select( + .update( "SELECT ttz_lambda($$ '2020-11-22 13:00:01't - '1 day'i $$, now(), 2.0)::text", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(&*res.unwrap(), "2020-11-21 13:00:01+00"); let res = client - .select( + .update( "SELECT ttz_lambda($$ '2020-11-22 13:00:01't + '1 day'i $$, now(), 2.0)::text", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(&*res.unwrap(), "2020-11-23 13:00:01+00"); point_lambda_eq!( @@ -684,20 +719,24 @@ mod tests { #[pg_test] fn test_lambda_comparison() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); bool_lambda_eq!(client, "2.0 < 3.0", "true"); bool_lambda_eq!(client, "2.0 <= 3.0", "true"); @@ -732,20 +771,24 @@ mod tests { #[pg_test] fn test_lambda_function() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); f64_lambda_eq!(client, "pi()", std::f64::consts::PI); @@ -779,20 +822,24 @@ mod tests { #[pg_test] fn test_lambda_unary() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); f64_lambda_eq!(client, "-(2.0)", -2.0f64); f64_lambda_eq!(client, "-(-2.0)", 2.0f64); @@ -808,20 +855,24 @@ mod tests { #[pg_test] fn test_lambda_interval_ops() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); interval_lambda_eq!(client, "'1 day'i + '1 day'i", "2 days"); interval_lambda_eq!(client, "'1 day'i + '1 week'i", "8 days"); @@ -835,20 +886,24 @@ mod tests { #[pg_test] fn test_lambda_variable() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); f64_lambda_eq!(client, "let $foo = 2.0; $foo", 2.0); f64_lambda_eq!(client, "let $foo = -2.0; $foo", -2.0); diff --git a/extension/src/time_vector/pipeline/map.rs b/extension/src/time_vector/pipeline/map.rs index ae3cfa10..6b3d34d8 100644 --- a/extension/src/time_vector/pipeline/map.rs +++ b/extension/src/time_vector/pipeline/map.rs @@ -93,7 +93,10 @@ pub fn map_series_pipeline_element<'e>( } pub fn map_series_element<'a>(function: crate::raw::regproc) -> Element<'a> { - let function: pg_sys::regproc = function.0.value().try_into().unwrap(); + let function: pg_sys::regproc = + unsafe { pg_sys::Oid::from_u32_unchecked(function.0.value() as u32) } + .try_into() + .unwrap(); check_user_function_type(function); Element::MapSeries { function: PgProcId(function), @@ -156,7 +159,7 @@ pub fn map_data_pipeline_element<'e>( let mut nargs: ::std::os::raw::c_int = 0; let rettype = unsafe { pg_sys::get_func_signature( - function.0.value().try_into().unwrap(), + pg_sys::Oid::from_u32_unchecked(function.0.value() as u32), &mut argtypes, &mut nargs, ) @@ -175,7 +178,7 @@ pub fn map_data_pipeline_element<'e>( } Element::MapData { - function: PgProcId(function.0.value().try_into().unwrap()), + function: PgProcId(unsafe { pg_sys::Oid::from_u32_unchecked(function.0.value() as u32) }), } .flatten() } @@ -250,46 +253,56 @@ mod tests { #[pg_test] fn test_pipeline_map_lambda() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); - client.select( - "CREATE TABLE series(time timestamptz, value double precision)", - None, - None, - ); - client.select( - "INSERT INTO series \ + client + .update( + "CREATE TABLE series(time timestamptz, value double precision)", + None, + None, + ) + .unwrap(); + client + .update( + "INSERT INTO series \ VALUES \ ('2020-01-04 UTC'::TIMESTAMPTZ, 25.0), \ ('2020-01-01 UTC'::TIMESTAMPTZ, 10.0), \ ('2020-01-03 UTC'::TIMESTAMPTZ, 20.0), \ ('2020-01-02 UTC'::TIMESTAMPTZ, 15.0), \ ('2020-01-05 UTC'::TIMESTAMPTZ, 30.0)", - None, - None, - ); + None, + None, + ) + .unwrap(); let val = client - .select( + .update( "SELECT (timevector(time, value))::TEXT FROM series", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:5,flags:0,internal_padding:(0,0,0),points:[\ @@ -301,13 +314,13 @@ mod tests { ],null_val:[0])" ); - let val = client.select( + let val = client.update( "SELECT (timevector(time, value) -> map($$ ($time + '1 day'i, $value * 2) $$))::TEXT FROM series", None, None ) - .first() - .get_one::(); + .unwrap().first() + .get_one::().unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:5,flags:0,internal_padding:(0,0,0),points:[\ @@ -323,46 +336,56 @@ mod tests { #[pg_test] fn test_pipeline_map_lambda2() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); - client.select( - "CREATE TABLE series(time timestamptz, value double precision)", - None, - None, - ); - client.select( - "INSERT INTO series \ + client + .update( + "CREATE TABLE series(time timestamptz, value double precision)", + None, + None, + ) + .unwrap(); + client + .update( + "INSERT INTO series \ VALUES \ ('2020-01-04 UTC'::TIMESTAMPTZ, 25.0), \ ('2020-01-01 UTC'::TIMESTAMPTZ, 10.0), \ ('2020-01-03 UTC'::TIMESTAMPTZ, 20.0), \ ('2020-01-02 UTC'::TIMESTAMPTZ, 15.0), \ ('2020-01-05 UTC'::TIMESTAMPTZ, 30.0)", - None, - None, - ); + None, + None, + ) + .unwrap(); let val = client - .select( + .update( "SELECT (timevector(time, value))::TEXT FROM series", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:5,flags:0,internal_padding:(0,0,0),points:[\ @@ -382,73 +405,87 @@ mod tests { (ts:\"2020-01-05 00:00:00+00\",val:1012.2)\ ],null_val:[0])"; let val = client - .select( + .update( "SELECT (timevector(time, value) \ -> map($$ ($time, $value^2 + $value * 2.3 + 43.2) $$))::TEXT \ FROM series", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(val.unwrap(), expected); let val = client - .select( + .update( "SELECT (timevector(time, value) \ -> map($$ ($value^2 + $value * 2.3 + 43.2) $$))::TEXT \ FROM series", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(val.unwrap(), expected); }); } #[pg_test] fn test_pipeline_map_data() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); - client.select( - "CREATE TABLE series(time timestamptz, value double precision)", - None, - None, - ); - client.select( - "INSERT INTO series \ + client + .update( + "CREATE TABLE series(time timestamptz, value double precision)", + None, + None, + ) + .unwrap(); + client + .update( + "INSERT INTO series \ VALUES \ ('2020-01-04 UTC'::TIMESTAMPTZ, 25.0), \ ('2020-01-01 UTC'::TIMESTAMPTZ, 10.0), \ ('2020-01-03 UTC'::TIMESTAMPTZ, 20.0), \ ('2020-01-02 UTC'::TIMESTAMPTZ, 15.0), \ ('2020-01-05 UTC'::TIMESTAMPTZ, 30.0)", - None, - None, - ); + None, + None, + ) + .unwrap(); let val = client - .select( + .update( "SELECT (timevector(time, value))::TEXT FROM series", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:5,flags:0,internal_padding:(0,0,0),points:[\ @@ -460,20 +497,22 @@ mod tests { ],null_val:[0])" ); - client.select( + client.update( "CREATE FUNCTION x2(double precision) RETURNS DOUBLE PRECISION AS 'SELECT $1 * 2;' LANGUAGE SQL", None, None, - ); + ).unwrap(); let val = client - .select( + .update( "SELECT (timevector(time, value) -> map_data('x2'))::TEXT FROM series", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:5,flags:0,internal_padding:(0,0,0),points:[\ @@ -489,46 +528,56 @@ mod tests { #[pg_test] fn test_pipeline_map_series() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); - client.select( - "CREATE TABLE series(time timestamptz, value double precision)", - None, - None, - ); - client.select( - "INSERT INTO series \ + client + .update( + "CREATE TABLE series(time timestamptz, value double precision)", + None, + None, + ) + .unwrap(); + client + .update( + "INSERT INTO series \ VALUES \ ('2020-01-04 UTC'::TIMESTAMPTZ, 25.0), \ ('2020-01-01 UTC'::TIMESTAMPTZ, 10.0), \ ('2020-01-03 UTC'::TIMESTAMPTZ, 20.0), \ ('2020-01-02 UTC'::TIMESTAMPTZ, 15.0), \ ('2020-01-05 UTC'::TIMESTAMPTZ, 30.0)", - None, - None, - ); + None, + None, + ) + .unwrap(); let val = client - .select( + .update( "SELECT (timevector(time, value))::TEXT FROM series", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:5,flags:0,internal_padding:(0,0,0),points:[\ @@ -540,7 +589,7 @@ mod tests { ],null_val:[0])" ); - client.select( + client.update( "CREATE FUNCTION jan_3_x3(timevector_tstz_f64) RETURNS timevector_tstz_f64 AS $$\ SELECT timevector(time, value * 3) \ FROM (SELECT (unnest($1)).*) a \ @@ -548,16 +597,18 @@ mod tests { $$ LANGUAGE SQL", None, None, - ); + ).unwrap(); let val = client - .select( + .update( "SELECT (timevector(time, value) -> map_series('jan_3_x3'))::TEXT FROM series", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:1,flags:1,internal_padding:(0,0,0),points:[\ @@ -570,33 +621,41 @@ mod tests { #[pg_test] #[should_panic = "division by zero"] fn test_pipeline_map_series_failure() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); - client.select( - "CREATE TABLE series(time timestamptz, value double precision)", - None, - None, - ); - client.select( - "INSERT INTO series \ + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) + .unwrap(); + client + .update( + "CREATE TABLE series(time timestamptz, value double precision)", + None, + None, + ) + .unwrap(); + client + .update( + "INSERT INTO series \ VALUES \ ('2020-01-05 UTC'::TIMESTAMPTZ, 30.0)", - None, - None, - ); - client.select( + None, + None, + ) + .unwrap(); + client.update( "CREATE FUNCTION always_fail(timevector_tstz_f64) RETURNS timevector_tstz_f64 AS $$ SELECT 0/0; @@ -604,110 +663,128 @@ mod tests { $$ LANGUAGE SQL", None, None, - ); + ).unwrap(); client - .select( + .update( "SELECT (timevector(time, value) -> map_series('always_fail'))::TEXT FROM series", None, None, ) - .first() - .get_one::(); + .unwrap().first() + .get_one::().unwrap(); }); } #[pg_test] #[should_panic = " returned NULL"] fn test_pipeline_map_series_null() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); - client.select( - "CREATE TABLE series(time timestamptz, value double precision)", - None, - None, - ); - client.select( - "INSERT INTO series \ + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) + .unwrap(); + client + .update( + "CREATE TABLE series(time timestamptz, value double precision)", + None, + None, + ) + .unwrap(); + client + .update( + "INSERT INTO series \ VALUES \ ('2020-01-05 UTC'::TIMESTAMPTZ, 30.0)", - None, - None, - ); - client.select( + None, + None, + ) + .unwrap(); + client.update( "CREATE FUNCTION always_null(timevector_tstz_f64) RETURNS timevector_tstz_f64 AS $$ SELECT NULL::timevector_tstz_f64; $$ LANGUAGE SQL", None, None, - ); + ).unwrap(); client - .select( + .update( "SELECT (timevector(time, value) -> map_series('always_null'))::TEXT FROM series", None, None, ) - .first() - .get_one::(); + .unwrap().first() + .get_one::().unwrap(); }); } #[pg_test] fn test_map_io() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); - client.select( - "CREATE TABLE series(time timestamptz, value double precision)", - None, - None, - ); - client.select( - "INSERT INTO series \ + client + .update( + "CREATE TABLE series(time timestamptz, value double precision)", + None, + None, + ) + .unwrap(); + client + .update( + "INSERT INTO series \ VALUES \ ('2020-01-04 UTC'::TIMESTAMPTZ, 25.0), \ ('2020-01-01 UTC'::TIMESTAMPTZ, 10.0), \ ('2020-01-03 UTC'::TIMESTAMPTZ, 20.0), \ ('2020-01-02 UTC'::TIMESTAMPTZ, 15.0), \ ('2020-01-05 UTC'::TIMESTAMPTZ, 30.0)", - None, - None, - ); + None, + None, + ) + .unwrap(); let val = client - .select( + .update( "SELECT (timevector(time, value))::TEXT FROM series", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:5,flags:0,internal_padding:(0,0,0),points:[\ @@ -719,30 +796,36 @@ mod tests { ],null_val:[0])" ); - client.select( - "CREATE FUNCTION serier(timevector_tstz_f64) RETURNS timevector_tstz_f64 AS $$\ + client + .update( + "CREATE FUNCTION serier(timevector_tstz_f64) RETURNS timevector_tstz_f64 AS $$\ SELECT $1;\ $$ LANGUAGE SQL", - None, - None, - ); + None, + None, + ) + .unwrap(); - client.select( - "CREATE FUNCTION dater(double precision) RETURNS double precision AS $$\ + client + .update( + "CREATE FUNCTION dater(double precision) RETURNS double precision AS $$\ SELECT $1 * 3;\ $$ LANGUAGE SQL", - None, - None, - ); + None, + None, + ) + .unwrap(); let (a, b) = client - .select( + .update( "SELECT map_series('serier')::TEXT, map_data('dater')::TEXT FROM series", None, None, ) + .unwrap() .first() - .get_two::(); + .get_two::() + .unwrap(); let one = "\ (\ version:1,\ @@ -767,7 +850,7 @@ mod tests { // FIXME this doesn't work yet let (a, b) = client - .select( + .update( &format!( "SELECT \ '{}'::UnstableTimevectorPipeline::Text, \ @@ -777,8 +860,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_two::(); + .get_two::() + .unwrap(); assert_eq!((&*a.unwrap(), &*b.unwrap()), (one, two)); }); } diff --git a/extension/src/time_vector/pipeline/sort.rs b/extension/src/time_vector/pipeline/sort.rs index 0e09e5af..f4af26d6 100644 --- a/extension/src/time_vector/pipeline/sort.rs +++ b/extension/src/time_vector/pipeline/sort.rs @@ -65,28 +65,35 @@ mod tests { #[pg_test] fn test_pipeline_sort() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); // using the search path trick for this test b/c the operator is // difficult to spot otherwise. let sp = client - .select( + .update( "SELECT format(' %s, toolkit_experimental',current_setting('search_path'))", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() + .unwrap(); + client + .update(&format!("SET LOCAL search_path TO {}", sp), None, None) .unwrap(); - client.select(&format!("SET LOCAL search_path TO {}", sp), None, None); - client.select( - "CREATE TABLE series(time timestamptz, value double precision)", - None, - None, - ); - client.select( - "INSERT INTO series \ + client + .update( + "CREATE TABLE series(time timestamptz, value double precision)", + None, + None, + ) + .unwrap(); + client + .update( + "INSERT INTO series \ VALUES \ ('2020-01-04 UTC'::TIMESTAMPTZ, 25), \ ('2020-01-01 UTC'::TIMESTAMPTZ, 10), \ @@ -94,18 +101,21 @@ mod tests { ('2020-01-02 UTC'::TIMESTAMPTZ, 15), \ ('2020-01-05 UTC'::TIMESTAMPTZ, 30), \ ('2020-01-02 12:00:00 UTC'::TIMESTAMPTZ, NULL)", - None, - None, - ); + None, + None, + ) + .unwrap(); let val = client - .select( + .update( "SELECT (timevector(time, value))::TEXT FROM series", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:6,flags:2,internal_padding:(0,0,0),points:[\ @@ -119,13 +129,15 @@ mod tests { ); let val = client - .select( + .update( "SELECT (timevector(time, value) -> sort())::TEXT FROM series", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!( val.unwrap(), "(version:1,num_points:6,flags:3,internal_padding:(0,0,0),points:[\ diff --git a/extension/src/time_weighted_average.rs b/extension/src/time_weighted_average.rs index 761ff213..81f7ece0 100644 --- a/extension/src/time_weighted_average.rs +++ b/extension/src/time_weighted_average.rs @@ -5,8 +5,8 @@ use serde::{Deserialize, Serialize}; use crate::{ accessors::{ - toolkit_experimental, AccessorAverage, AccessorFirstTime, AccessorFirstVal, - AccessorLastTime, AccessorLastVal, + AccessorAverage, AccessorFirstTime, AccessorFirstVal, AccessorIntegral, AccessorLastTime, + AccessorLastVal, }, aggregate_utils::in_aggregate_context, duration::DurationUnit, @@ -23,6 +23,10 @@ use time_weighted_average::{ use crate::raw::bytea; +mod accessors; + +use accessors::{TimeWeightInterpolatedAverageAccessor, TimeWeightInterpolatedIntegralAccessor}; + pg_type! { #[derive(Debug)] struct TimeWeightSummary { @@ -425,7 +429,7 @@ pub fn arrow_time_weighted_average_average<'a>( #[opname(->)] pub fn arrow_time_weighted_average_integral<'a>( tws: Option>, - accessor: toolkit_experimental::AccessorIntegral<'a>, + accessor: AccessorIntegral<'a>, ) -> Option { time_weighted_average_integral( tws, @@ -451,12 +455,7 @@ pub fn time_weighted_average_average<'a>(tws: Option>) -> } } -#[pg_extern( - immutable, - parallel_safe, - name = "integral", - schema = "toolkit_experimental" -)] +#[pg_extern(immutable, parallel_safe, name = "integral")] pub fn time_weighted_average_integral<'a>( tws: Option>, unit: default!(String, "'second'"), @@ -488,20 +487,6 @@ fn interpolate<'a>( } } -// Public facing interpolated_average -extension_sql!( - "\n\ - CREATE FUNCTION toolkit_experimental.interpolated_average(tws timeweightsummary,\n\ - start timestamptz,\n\ - duration interval,\n\ - prev timeweightsummary,\n\ - next timeweightsummary) RETURNS DOUBLE PRECISION\n\ - AS $$\n\ - SELECT interpolated_average(tws,start,duration,prev,next) $$ LANGUAGE SQL IMMUTABLE PARALLEL SAFE;\n\ -", - name = "experimental_interpolated_average", requires = [time_weighted_average_interpolated_average] -); - #[pg_extern(immutable, parallel_safe, name = "interpolated_average")] pub fn time_weighted_average_interpolated_average<'a>( tws: Option>, @@ -514,12 +499,33 @@ pub fn time_weighted_average_interpolated_average<'a>( time_weighted_average_average(target) } -#[pg_extern( - immutable, - parallel_safe, - name = "interpolated_integral", - schema = "toolkit_experimental" -)] +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_time_weighted_average_interpolated_average<'a>( + tws: Option>, + accessor: TimeWeightInterpolatedAverageAccessor<'a>, +) -> Option { + let prev = if accessor.flags & 1 == 1 { + Some(accessor.prev.clone().into()) + } else { + None + }; + let next = if accessor.flags & 2 == 2 { + Some(accessor.next.clone().into()) + } else { + None + }; + + time_weighted_average_interpolated_average( + tws, + accessor.timestamp.into(), + accessor.interval.into(), + prev, + next, + ) +} + +#[pg_extern(immutable, parallel_safe, name = "interpolated_integral")] pub fn time_weighted_average_interpolated_integral<'a>( tws: Option>, start: crate::raw::TimestampTz, @@ -532,6 +538,44 @@ pub fn time_weighted_average_interpolated_integral<'a>( time_weighted_average_integral(target, unit) } +#[pg_operator(immutable, parallel_safe)] +#[opname(->)] +pub fn arrow_time_weighted_average_interpolated_integral<'a>( + tws: Option>, + accessor: TimeWeightInterpolatedIntegralAccessor<'a>, +) -> Option { + let prev = if accessor.flags & 1 == 1 { + Some(accessor.prev.clone().into()) + } else { + None + }; + let next = if accessor.flags & 2 == 2 { + Some(accessor.next.clone().into()) + } else { + None + }; + + // Convert from num of milliseconds to DurationUnit and then to string + let unit = match accessor.unit { + 1 => DurationUnit::Microsec, + 1000 => DurationUnit::Millisec, + 1_000_000 => DurationUnit::Second, + 60_000_000 => DurationUnit::Minute, + 3_600_000_000 => DurationUnit::Hour, + _ => todo!(), // This should never be reached, the accessor gets these numbers from microseconds() in duration.rs, which only matches on valid enum values + } + .to_string(); + + time_weighted_average_interpolated_integral( + tws, + accessor.start.into(), + accessor.interval.into(), + prev, + next, + unit, + ) +} + #[cfg(any(test, feature = "pg_test"))] #[pg_schema] mod tests { @@ -541,31 +585,33 @@ mod tests { macro_rules! select_one { ($client:expr, $stmt:expr, $type:ty) => { $client - .select($stmt, None, None) + .update($stmt, None, None) + .unwrap() .first() .get_one::<$type>() .unwrap() + .unwrap() }; } #[pg_test] fn test_time_weight_aggregate() { - Spi::execute(|client| { + Spi::connect(|mut client| { let stmt = "CREATE TABLE test(ts timestamptz, val DOUBLE PRECISION); SET TIME ZONE 'UTC'"; - client.select(stmt, None, None); + client.update(stmt, None, None).unwrap(); // add a point let stmt = "INSERT INTO test VALUES('2020-01-01 00:00:00+00', 10.0)"; - client.select(stmt, None, None); + client.update(stmt, None, None).unwrap(); - let stmt = "SELECT toolkit_experimental.integral(time_weight('Trapezoidal', ts, val), 'hrs') FROM test"; + let stmt = "SELECT integral(time_weight('Trapezoidal', ts, val), 'hrs') FROM test"; assert_eq!(select_one!(client, stmt, f64), 0.0); - let stmt = "SELECT toolkit_experimental.integral(time_weight('LOCF', ts, val), 'msecond') FROM test"; + let stmt = "SELECT integral(time_weight('LOCF', ts, val), 'msecond') FROM test"; assert_eq!(select_one!(client, stmt, f64), 0.0); // add another point let stmt = "INSERT INTO test VALUES('2020-01-01 00:01:00+00', 20.0)"; - client.select(stmt, None, None); + client.update(stmt, None, None).unwrap(); // test basic with 2 points let stmt = "SELECT average(time_weight('Linear', ts, val)) FROM test"; @@ -597,21 +643,21 @@ mod tests { // more values evenly spaced let stmt = "INSERT INTO test VALUES('2020-01-01 00:02:00+00', 10.0), ('2020-01-01 00:03:00+00', 20.0), ('2020-01-01 00:04:00+00', 10.0)"; - client.select(stmt, None, None); + client.update(stmt, None, None).unwrap(); let stmt = "SELECT average(time_weight('Linear', ts, val)) FROM test"; assert!((select_one!(client, stmt, f64) - 15.0).abs() < f64::EPSILON); let stmt = "SELECT average(time_weight('LOCF', ts, val)) FROM test"; assert!((select_one!(client, stmt, f64) - 15.0).abs() < f64::EPSILON); - let stmt = "SELECT toolkit_experimental.integral(time_weight('Linear', ts, val), 'mins') FROM test"; + let stmt = "SELECT integral(time_weight('Linear', ts, val), 'mins') FROM test"; assert!((select_one!(client, stmt, f64) - 60.0).abs() < f64::EPSILON); - let stmt = "SELECT toolkit_experimental.integral(time_weight('LOCF', ts, val), 'hour') FROM test"; + let stmt = "SELECT integral(time_weight('LOCF', ts, val), 'hour') FROM test"; assert!((select_one!(client, stmt, f64) - 1.0).abs() < f64::EPSILON); //non-evenly spaced values let stmt = "INSERT INTO test VALUES('2020-01-01 00:08:00+00', 30.0), ('2020-01-01 00:10:00+00', 10.0), ('2020-01-01 00:10:30+00', 20.0), ('2020-01-01 00:20:00+00', 30.0)"; - client.select(stmt, None, None); + client.update(stmt, None, None).unwrap(); let stmt = "SELECT average(time_weight('Linear', ts, val)) FROM test"; // expected =(15 +15 +15 +15 + 20*4 + 20*2 +15*.5 + 25*9.5) / 20 = 21.25 just taking the midpoints between each point and multiplying by minutes and dividing by total @@ -622,15 +668,15 @@ mod tests { // arrow syntax should be the same assert!((select_one!(client, stmt, f64) - 21.25).abs() < f64::EPSILON); - let stmt = "SELECT toolkit_experimental.integral(time_weight('Linear', ts, val), 'microseconds') FROM test"; + let stmt = "SELECT integral(time_weight('Linear', ts, val), 'microseconds') FROM test"; assert!((select_one!(client, stmt, f64) - 25500000000.00).abs() < f64::EPSILON); let stmt = "SELECT time_weight('Linear', ts, val) \ - ->toolkit_experimental.integral('microseconds') \ + ->integral('microseconds') \ FROM test"; // arrow syntax should be the same assert!((select_one!(client, stmt, f64) - 25500000000.00).abs() < f64::EPSILON); let stmt = "SELECT time_weight('Linear', ts, val) \ - ->toolkit_experimental.integral() \ + ->integral() \ FROM test"; assert!((select_one!(client, stmt, f64) - 25500.00).abs() < f64::EPSILON); @@ -638,7 +684,7 @@ mod tests { // expected = (10 + 20 + 10 + 20 + 10*4 + 30*2 +10*.5 + 20*9.5) / 20 = 17.75 using last value and carrying for each point assert!((select_one!(client, stmt, f64) - 17.75).abs() < f64::EPSILON); - let stmt = "SELECT toolkit_experimental.integral(time_weight('LOCF', ts, val), 'milliseconds') FROM test"; + let stmt = "SELECT integral(time_weight('LOCF', ts, val), 'milliseconds') FROM test"; assert!((select_one!(client, stmt, f64) - 21300000.0).abs() < f64::EPSILON); //make sure this works with whatever ordering we throw at it @@ -647,9 +693,9 @@ mod tests { let stmt = "SELECT average(time_weight('LOCF', ts, val ORDER BY random())) FROM test"; assert!((select_one!(client, stmt, f64) - 17.75).abs() < f64::EPSILON); - let stmt = "SELECT toolkit_experimental.integral(time_weight('Linear', ts, val ORDER BY random()), 'seconds') FROM test"; + let stmt = "SELECT integral(time_weight('Linear', ts, val ORDER BY random()), 'seconds') FROM test"; assert!((select_one!(client, stmt, f64) - 25500.0).abs() < f64::EPSILON); - let stmt = "SELECT toolkit_experimental.integral(time_weight('LOCF', ts, val ORDER BY random())) FROM test"; + let stmt = "SELECT integral(time_weight('LOCF', ts, val ORDER BY random())) FROM test"; assert!((select_one!(client, stmt, f64) - 21300.0).abs() < f64::EPSILON); // make sure we get the same result if we do multi-level aggregation @@ -662,10 +708,10 @@ mod tests { #[pg_test] fn test_time_weight_io() { - Spi::execute(|client| { - client.select("SET timezone TO 'UTC'", None, None); + Spi::connect(|mut client| { + client.update("SET timezone TO 'UTC'", None, None).unwrap(); let stmt = "CREATE TABLE test(ts timestamptz, val DOUBLE PRECISION)"; - client.select(stmt, None, None); + client.update(stmt, None, None).unwrap(); let linear_time_weight = "SELECT time_weight('Linear', ts, val)::TEXT FROM test"; let locf_time_weight = "SELECT time_weight('LOCF', ts, val)::TEXT FROM test"; @@ -673,7 +719,7 @@ mod tests { // add a couple points let stmt = "INSERT INTO test VALUES('2020-01-01 00:00:00+00', 10.0), ('2020-01-01 00:01:00+00', 20.0)"; - client.select(stmt, None, None); + client.update(stmt, None, None).unwrap(); // test basic with 2 points let expected = "(\ @@ -698,7 +744,7 @@ mod tests { // more values evenly spaced let stmt = "INSERT INTO test VALUES('2020-01-01 00:02:00+00', 10.0), ('2020-01-01 00:03:00+00', 20.0), ('2020-01-01 00:04:00+00', 10.0)"; - client.select(stmt, None, None); + client.update(stmt, None, None).unwrap(); let expected = "(\ version:1,\ @@ -721,7 +767,7 @@ mod tests { //non-evenly spaced values let stmt = "INSERT INTO test VALUES('2020-01-01 00:08:00+00', 30.0), ('2020-01-01 00:10:00+00', 10.0), ('2020-01-01 00:10:30+00', 20.0), ('2020-01-01 00:20:00+00', 30.0)"; - client.select(stmt, None, None); + client.update(stmt, None, None).unwrap(); let expected = "(\ version:1,\ @@ -816,14 +862,15 @@ mod tests { #[pg_test] fn test_time_weight_interpolation() { - Spi::execute(|client| { - client.select( + Spi::connect(|mut client| { + client.update( "CREATE TABLE test(time timestamptz, value double precision, bucket timestamptz)", None, None, - ); - client.select( - r#"INSERT INTO test VALUES + ).unwrap(); + client + .update( + r#"INSERT INTO test VALUES ('2020-1-1 8:00'::timestamptz, 10.0, '2020-1-1'::timestamptz), ('2020-1-1 12:00'::timestamptz, 40.0, '2020-1-1'::timestamptz), ('2020-1-1 16:00'::timestamptz, 20.0, '2020-1-1'::timestamptz), @@ -833,13 +880,15 @@ mod tests { ('2020-1-3 10:00'::timestamptz, 30.0, '2020-1-3'::timestamptz), ('2020-1-3 12:00'::timestamptz, 0.0, '2020-1-3'::timestamptz), ('2020-1-3 16:00'::timestamptz, 35.0, '2020-1-3'::timestamptz)"#, - None, - None, - ); - // test experimental version - let mut experimental_averages = client.select( - r#"SELECT - toolkit_experimental.interpolated_average( + None, + None, + ) + .unwrap(); + + let mut averages = client + .update( + r#"SELECT + interpolated_average( agg, bucket, '1 day'::interval, @@ -851,14 +900,15 @@ mod tests { GROUP BY bucket ) s ORDER BY bucket"#, - None, - None, - ); - // test non_experimental version - let mut averages = client.select( - r#"SELECT - interpolated_average( - agg, + None, + None, + ) + .unwrap(); + // test arrow version + let mut arrow_averages = client + .update( + r#"SELECT + agg -> interpolated_average( bucket, '1 day'::interval, LAG(agg) OVER (ORDER BY bucket), @@ -869,12 +919,14 @@ mod tests { GROUP BY bucket ) s ORDER BY bucket"#, - None, - None, - ); - let mut integrals = client.select( - r#"SELECT - toolkit_experimental.interpolated_integral( + None, + None, + ) + .unwrap(); + let mut integrals = client + .update( + r#"SELECT + interpolated_integral( agg, bucket, '1 day'::interval, @@ -887,13 +939,15 @@ mod tests { GROUP BY bucket ) s ORDER BY bucket"#, - None, - None, - ); + None, + None, + ) + .unwrap(); // verify that default value works - client.select( - r#"SELECT - toolkit_experimental.interpolated_integral( + client + .update( + r#"SELECT + interpolated_integral( agg, bucket, '1 day'::interval, @@ -905,40 +959,41 @@ mod tests { GROUP BY bucket ) s ORDER BY bucket"#, - None, - None, - ); + None, + None, + ) + .unwrap(); // Day 1, 4 hours @ 10, 4 @ 40, 8 @ 20 - let result = experimental_averages.next().unwrap()[1].value(); + let result = averages.next().unwrap()[1].value().unwrap(); assert_eq!(result, Some((4. * 10. + 4. * 40. + 8. * 20.) / 16.)); - assert_eq!(result, averages.next().unwrap()[1].value()); + assert_eq!(result, arrow_averages.next().unwrap()[1].value().unwrap()); assert_eq!( - integrals.next().unwrap()[1].value(), + integrals.next().unwrap()[1].value().unwrap(), Some(4. * 10. + 4. * 40. + 8. * 20.) ); // Day 2, 2 hours @ 20, 10 @ 15, 8 @ 50, 4 @ 25 - let result = experimental_averages.next().unwrap()[1].value(); + let result = averages.next().unwrap()[1].value().unwrap(); assert_eq!( result, Some((2. * 20. + 10. * 15. + 8. * 50. + 4. * 25.) / 24.) ); - assert_eq!(result, averages.next().unwrap()[1].value()); + assert_eq!(result, arrow_averages.next().unwrap()[1].value().unwrap()); assert_eq!( - integrals.next().unwrap()[1].value(), + integrals.next().unwrap()[1].value().unwrap(), Some(2. * 20. + 10. * 15. + 8. * 50. + 4. * 25.) ); // Day 3, 10 hours @ 25, 2 @ 30, 4 @ 0 - let result = experimental_averages.next().unwrap()[1].value(); + let result = averages.next().unwrap()[1].value().unwrap(); assert_eq!(result, Some((10. * 25. + 2. * 30.) / 16.)); - assert_eq!(result, averages.next().unwrap()[1].value()); + assert_eq!(result, arrow_averages.next().unwrap()[1].value().unwrap()); assert_eq!( - integrals.next().unwrap()[1].value(), + integrals.next().unwrap()[1].value().unwrap(), Some(10. * 25. + 2. * 30.) ); - assert!(experimental_averages.next().is_none()); assert!(averages.next().is_none()); + assert!(arrow_averages.next().is_none()); assert!(integrals.next().is_none()); }); } diff --git a/extension/src/time_weighted_average/accessors.rs b/extension/src/time_weighted_average/accessors.rs new file mode 100644 index 00000000..143fbf2b --- /dev/null +++ b/extension/src/time_weighted_average/accessors.rs @@ -0,0 +1,115 @@ +use pgx::*; + +use crate::time_weighted_average::DurationUnit; +use crate::{ + datum_utils::interval_to_ms, + flatten, pg_type, ron_inout_funcs, + time_weighted_average::{TimeWeightMethod, TimeWeightSummary, TimeWeightSummaryData}, +}; + +use tspoint::TSPoint; + +pg_type! { + struct TimeWeightInterpolatedAverageAccessor { + timestamp : i64, + interval : i64, + prev : TimeWeightSummaryData, + pad : [u8;3], + flags : u32, + next : TimeWeightSummaryData, + } +} + +ron_inout_funcs!(TimeWeightInterpolatedAverageAccessor); + +#[pg_extern(immutable, parallel_safe, name = "interpolated_average")] +fn time_weight_interpolated_average_accessor<'a>( + start: crate::raw::TimestampTz, + duration: crate::raw::Interval, + prev: Option>, + next: Option>, +) -> TimeWeightInterpolatedAverageAccessor<'static> { + fn empty_summary<'b>() -> Option> { + Some(unsafe { + flatten!(TimeWeightSummary { + first: TSPoint { ts: 0, val: 0.0 }, + last: TSPoint { ts: 0, val: 0.0 }, + weighted_sum: 0.0, + method: TimeWeightMethod::LOCF, + }) + }) + } + + let flags = u32::from(prev.is_some()) + if next.is_some() { 2 } else { 0 }; + let prev = prev.or_else(empty_summary).unwrap().0; + let next = next.or_else(empty_summary).unwrap().0; + let interval = interval_to_ms(&start, &duration); + crate::build! { + TimeWeightInterpolatedAverageAccessor { + timestamp : start.into(), + interval, + prev, + pad : [0,0,0], + flags, + next, + } + } +} + +pg_type! { + #[derive(Debug)] + struct TimeWeightInterpolatedIntegralAccessor { + start : i64, + interval : i64, + prev : TimeWeightSummaryData, + pad : [u8;3], + unit : u32, + flags: u64, + next : TimeWeightSummaryData, + } +} + +ron_inout_funcs!(TimeWeightInterpolatedIntegralAccessor); + +#[pg_extern(immutable, parallel_safe, name = "interpolated_integral")] +fn time_weight_interpolated_integral_accessor<'a>( + start: crate::raw::TimestampTz, + interval: crate::raw::Interval, + prev: Option>, + next: Option>, + unit: default!(String, "'second'"), +) -> TimeWeightInterpolatedIntegralAccessor<'static> { + fn empty_summary<'b>() -> Option> { + Some(unsafe { + flatten!(TimeWeightSummary { + first: TSPoint { ts: 0, val: 0.0 }, + last: TSPoint { ts: 0, val: 0.0 }, + weighted_sum: 0.0, + method: TimeWeightMethod::LOCF, + }) + }) + } + + let unit = match DurationUnit::from_str(&unit) { + Some(unit) => unit.microseconds(), + None => pgx::error!( + "Unrecognized duration unit: {}. Valid units are: usecond, msecond, second, minute, hour", + unit, + ), + }; + let flags = u64::from(prev.is_some()) + if next.is_some() { 2 } else { 0 }; + let prev = prev.or_else(empty_summary).unwrap().0; + let next = next.or_else(empty_summary).unwrap().0; + let interval = interval_to_ms(&start, &interval); + crate::build! { + TimeWeightInterpolatedIntegralAccessor { + start: start.into(), + interval, + prev, + pad : [0,0,0], + unit, + flags, + next, + } + } +} diff --git a/extension/src/type_builder.rs b/extension/src/type_builder.rs index 533a9e23..26b509a9 100644 --- a/extension/src/type_builder.rs +++ b/extension/src/type_builder.rs @@ -260,7 +260,7 @@ macro_rules! ron_inout_funcs { } } - fn input(input: &pgx::cstr_core::CStr) -> $name<'input> + fn input(input: &std::ffi::CStr) -> $name<'input> where Self: Sized, { diff --git a/extension/src/uddsketch.rs b/extension/src/uddsketch.rs index b60e50bc..f0372734 100644 --- a/extension/src/uddsketch.rs +++ b/extension/src/uddsketch.rs @@ -273,7 +273,7 @@ impl<'input> InOutFuncs for UddSketch<'input> { } } - fn input(input: &pgx::cstr_core::CStr) -> Self + fn input(input: &std::ffi::CStr) -> Self where Self: Sized, { @@ -706,36 +706,46 @@ mod tests { #[pg_test] fn test_aggregate() { - Spi::execute(|client| { - client.select("CREATE TABLE test (data DOUBLE PRECISION)", None, None); - client.select( - "INSERT INTO test SELECT generate_series(0.01, 100, 0.01)", - None, - None, - ); + Spi::connect(|mut client| { + client + .update("CREATE TABLE test (data DOUBLE PRECISION)", None, None) + .unwrap(); + client + .update( + "INSERT INTO test SELECT generate_series(0.01, 100, 0.01)", + None, + None, + ) + .unwrap(); let sanity = client - .select("SELECT COUNT(*) FROM test", None, None) + .update("SELECT COUNT(*) FROM test", None, None) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(Some(10000), sanity); - client.select( - "CREATE VIEW sketch AS \ + client + .update( + "CREATE VIEW sketch AS \ SELECT uddsketch(100, 0.05, data) \ FROM test", - None, - None, - ); + None, + None, + ) + .unwrap(); let sanity = client - .select("SELECT COUNT(*) FROM sketch", None, None) + .update("SELECT COUNT(*) FROM sketch", None, None) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert!(sanity.unwrap_or(0) > 0); let (mean, count) = client - .select( + .update( "SELECT \ mean(uddsketch), \ num_vals(uddsketch) \ @@ -743,14 +753,16 @@ mod tests { None, None, ) + .unwrap() .first() - .get_two::(); + .get_two::() + .unwrap(); apx_eql(mean.unwrap(), 50.005, 0.0001); apx_eql(count.unwrap(), 10000.0, 0.000001); let (mean2, count2) = client - .select( + .update( "SELECT \ uddsketch -> mean(), \ uddsketch -> num_vals() \ @@ -758,13 +770,15 @@ mod tests { None, None, ) + .unwrap() .first() - .get_two::(); + .get_two::() + .unwrap(); assert_eq!(mean, mean2); assert_eq!(count, count2); let (error, error2) = client - .select( + .update( "SELECT \ error(uddsketch), \ uddsketch -> error() \ @@ -772,8 +786,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_two::(); + .get_two::() + .unwrap(); apx_eql(error.unwrap(), 0.05, 0.0001); assert_eq!(error, error2); @@ -783,7 +799,7 @@ mod tests { let approx_percentile = value / 100.0; let (est_val, est_quant) = client - .select( + .update( &format!( "SELECT \ approx_percentile({}, uddsketch), \ @@ -794,8 +810,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_two::(); + .get_two::() + .unwrap(); if i == 0 { pct_eql(est_val.unwrap(), 0.01, 1.0); @@ -806,7 +824,7 @@ mod tests { } let (est_val2, est_quant2) = client - .select( + .update( &format!( "SELECT \ uddsketch->approx_percentile({}), \ @@ -817,8 +835,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_two::(); + .get_two::() + .unwrap(); assert_eq!(est_val, est_val2); assert_eq!(est_quant, est_quant2); } @@ -827,47 +847,57 @@ mod tests { #[pg_test] fn test_compound_agg() { - Spi::execute(|client| { - client.select( - "CREATE TABLE new_test (device INTEGER, value DOUBLE PRECISION)", - None, - None, - ); - client.select("INSERT INTO new_test SELECT dev, dev - v FROM generate_series(1,10) dev, generate_series(0, 1.0, 0.01) v", None, None); + Spi::connect(|mut client| { + client + .update( + "CREATE TABLE new_test (device INTEGER, value DOUBLE PRECISION)", + None, + None, + ) + .unwrap(); + client.update("INSERT INTO new_test SELECT dev, dev - v FROM generate_series(1,10) dev, generate_series(0, 1.0, 0.01) v", None, None).unwrap(); let sanity = client - .select("SELECT COUNT(*) FROM new_test", None, None) + .update("SELECT COUNT(*) FROM new_test", None, None) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(Some(1010), sanity); - client.select( - "CREATE VIEW sketches AS \ + client + .update( + "CREATE VIEW sketches AS \ SELECT device, uddsketch(20, 0.01, value) \ FROM new_test \ GROUP BY device", - None, - None, - ); + None, + None, + ) + .unwrap(); - client.select( - "CREATE VIEW composite AS \ + client + .update( + "CREATE VIEW composite AS \ SELECT rollup(uddsketch) as uddsketch \ FROM sketches", - None, - None, - ); + None, + None, + ) + .unwrap(); - client.select( - "CREATE VIEW base AS \ + client + .update( + "CREATE VIEW base AS \ SELECT uddsketch(20, 0.01, value) \ FROM new_test", - None, - None, - ); + None, + None, + ) + .unwrap(); let (value, error) = client - .select( + .update( "SELECT \ approx_percentile(0.9, uddsketch), \ error(uddsketch) \ @@ -875,11 +905,13 @@ mod tests { None, None, ) + .unwrap() .first() - .get_two::(); + .get_two::() + .unwrap(); let (test_value, test_error) = client - .select( + .update( "SELECT \ approx_percentile(0.9, uddsketch), \ error(uddsketch) \ @@ -887,8 +919,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_two::(); + .get_two::() + .unwrap(); apx_eql(test_value.unwrap(), value.unwrap(), 0.0001); apx_eql(test_error.unwrap(), error.unwrap(), 0.000001); @@ -898,39 +932,47 @@ mod tests { #[pg_test] fn test_percentile_agg() { - Spi::execute(|client| { - client.select( - "CREATE TABLE pa_test (device INTEGER, value DOUBLE PRECISION)", - None, - None, - ); - client.select("INSERT INTO pa_test SELECT dev, dev - v FROM generate_series(1,10) dev, generate_series(0, 1.0, 0.01) v", None, None); + Spi::connect(|mut client| { + client + .update( + "CREATE TABLE pa_test (device INTEGER, value DOUBLE PRECISION)", + None, + None, + ) + .unwrap(); + client.update("INSERT INTO pa_test SELECT dev, dev - v FROM generate_series(1,10) dev, generate_series(0, 1.0, 0.01) v", None, None).unwrap(); let sanity = client - .select("SELECT COUNT(*) FROM pa_test", None, None) + .update("SELECT COUNT(*) FROM pa_test", None, None) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(Some(1010), sanity); // use the default values for percentile_agg - client.select( - "CREATE VIEW uddsketch_test AS \ + client + .update( + "CREATE VIEW uddsketch_test AS \ SELECT uddsketch(200, 0.001, value) as approx \ FROM pa_test ", - None, - None, - ); + None, + None, + ) + .unwrap(); - client.select( - "CREATE VIEW percentile_agg AS \ + client + .update( + "CREATE VIEW percentile_agg AS \ SELECT percentile_agg(value) as approx \ FROM pa_test", - None, - None, - ); + None, + None, + ) + .unwrap(); let (value, error) = client - .select( + .update( "SELECT \ approx_percentile(0.9, approx), \ error(approx) \ @@ -938,11 +980,13 @@ mod tests { None, None, ) + .unwrap() .first() - .get_two::(); + .get_two::() + .unwrap(); let (test_value, test_error) = client - .select( + .update( "SELECT \ approx_percentile(0.9, approx), \ error(approx) \ @@ -950,8 +994,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_two::(); + .get_two::() + .unwrap(); apx_eql(test_value.unwrap(), value.unwrap(), 0.0001); apx_eql(test_error.unwrap(), error.unwrap(), 0.000001); @@ -960,38 +1006,46 @@ mod tests { } #[pg_test] fn test_approx_percentile_array() { - Spi::execute(|client| { - client.select( - "CREATE TABLE paa_test (device INTEGER, value DOUBLE PRECISION)", - None, - None, - ); - client.select("INSERT INTO paa_test SELECT dev, dev - v FROM generate_series(1,10) dev, generate_series(0, 1.0, 0.01) v", None, None); + Spi::connect(|mut client| { + client + .update( + "CREATE TABLE paa_test (device INTEGER, value DOUBLE PRECISION)", + None, + None, + ) + .unwrap(); + client.update("INSERT INTO paa_test SELECT dev, dev - v FROM generate_series(1,10) dev, generate_series(0, 1.0, 0.01) v", None, None).unwrap(); let sanity = client - .select("SELECT COUNT(*) FROM paa_test", None, None) + .update("SELECT COUNT(*) FROM paa_test", None, None) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(Some(1010), sanity); - client.select( - "CREATE VIEW uddsketch_test AS \ + client + .update( + "CREATE VIEW uddsketch_test AS \ SELECT uddsketch(200, 0.001, value) as approx \ FROM paa_test ", - None, - None, - ); + None, + None, + ) + .unwrap(); - client.select( - "CREATE VIEW percentile_agg AS \ + client + .update( + "CREATE VIEW percentile_agg AS \ SELECT percentile_agg(value) as approx \ FROM paa_test", - None, - None, - ); + None, + None, + ) + .unwrap(); let (value, error) = client - .select( + .update( "SELECT \ toolkit_experimental.approx_percentile_array(array[0.9,0.5,0.2], approx), \ error(approx) \ @@ -999,11 +1053,13 @@ mod tests { None, None, ) + .unwrap() .first() - .get_two::, f64>(); + .get_two::, f64>() + .unwrap(); let (test_value, test_error) = client - .select( + .update( "SELECT \ toolkit_experimental.approx_percentile_array(array[0.9,0.5,0.2], approx), \ error(approx) \ @@ -1011,8 +1067,10 @@ mod tests { None, None, ) + .unwrap() .first() - .get_two::, f64>(); + .get_two::, f64>() + .unwrap(); assert!( test_value .as_ref() @@ -1035,38 +1093,46 @@ mod tests { #[pg_test] fn test_approx_percentile_array_arrow() { - Spi::execute(|client| { - client.select( - "CREATE TABLE paa_test (device INTEGER, value DOUBLE PRECISION)", - None, - None, - ); - client.select("INSERT INTO paa_test SELECT dev, dev - v FROM generate_series(1,10) dev, generate_series(0, 1.0, 0.01) v", None, None); + Spi::connect(|mut client| { + client + .update( + "CREATE TABLE paa_test (device INTEGER, value DOUBLE PRECISION)", + None, + None, + ) + .unwrap(); + client.update("INSERT INTO paa_test SELECT dev, dev - v FROM generate_series(1,10) dev, generate_series(0, 1.0, 0.01) v", None, None).unwrap(); let sanity = client - .select("SELECT COUNT(*) FROM paa_test", None, None) + .update("SELECT COUNT(*) FROM paa_test", None, None) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(Some(1010), sanity); - client.select( - "CREATE VIEW uddsketch_test AS \ + client + .update( + "CREATE VIEW uddsketch_test AS \ SELECT uddsketch(200, 0.001, value) as approx \ FROM paa_test ", - None, - None, - ); + None, + None, + ) + .unwrap(); - client.select( - "CREATE VIEW percentile_agg AS \ + client + .update( + "CREATE VIEW percentile_agg AS \ SELECT percentile_agg(value) as approx \ FROM paa_test", - None, - None, - ); + None, + None, + ) + .unwrap(); let (value, error) = client - .select( + .update( "SELECT \ toolkit_experimental.approx_percentile_array(array[0.9,0.5,0.2], approx), \ error(approx) \ @@ -1074,19 +1140,23 @@ mod tests { None, None, ) + .unwrap() .first() - .get_two::, f64>(); + .get_two::, f64>() + .unwrap(); let (test_value_arrow, test_error_arrow) = client - .select( + .update( "SELECT approx->toolkit_experimental.approx_percentiles(array[0.9,0.5,0.2]), \ error(approx) \ FROM uddsketch_test", None, None, ) + .unwrap() .first() - .get_two::, f64>(); + .get_two::, f64>() + .unwrap(); assert!( test_value_arrow @@ -1110,18 +1180,22 @@ mod tests { #[pg_test] fn uddsketch_io_test() { - Spi::execute(|client| { - client.select("CREATE TABLE io_test (value DOUBLE PRECISION)", None, None); - client.select("INSERT INTO io_test VALUES (-1000), (-100), (-10), (-1), (-0.1), (-0.01), (-0.001), (0), (0.001), (0.01), (0.1), (1), (10), (100), (1000)", None, None); + Spi::connect(|mut client| { + client + .update("CREATE TABLE io_test (value DOUBLE PRECISION)", None, None) + .unwrap(); + client.update("INSERT INTO io_test VALUES (-1000), (-100), (-10), (-1), (-0.1), (-0.01), (-0.001), (0), (0.001), (0.01), (0.1), (1), (10), (100), (1000)", None, None).unwrap(); let sketch = client - .select( + .update( "SELECT uddsketch(10, 0.01, value)::text FROM io_test", None, None, ) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); let expected = "(\ version:1,\ @@ -1147,13 +1221,12 @@ mod tests { assert_eq!(sketch, Some(expected.into())); client - .select( + .update( "CREATE VIEW sketch AS SELECT uddsketch(10, 0.01, value) FROM io_test", None, None, ) - .first() - .get_one::(); + .unwrap(); for cmd in [ "mean(", @@ -1169,14 +1242,18 @@ mod tests { let sql2 = format!("SELECT {}'{}'::uddsketch) FROM sketch", cmd, expected); let expected = client - .select(&sql1, None, None) + .update(&sql1, None, None) + .unwrap() .first() .get_one::() + .unwrap() .unwrap(); let test = client - .select(&sql2, None, None) + .update(&sql2, None, None) + .unwrap() .first() .get_one::() + .unwrap() .unwrap(); assert!((expected - test).abs() < f64::EPSILON); @@ -1215,11 +1292,13 @@ mod tests { #[pg_test] fn test_udd_null_input_yields_null_output() { - Spi::execute(|client| { + Spi::connect(|mut client| { let output = client - .select("SELECT uddsketch(20, 0.01, NULL)::TEXT", None, None) + .update("SELECT uddsketch(20, 0.01, NULL)::TEXT", None, None) + .unwrap() .first() - .get_one::(); + .get_one::() + .unwrap(); assert_eq!(output, None) }) } diff --git a/extension/src/utilities.rs b/extension/src/utilities.rs index 597fdbb1..b71520f4 100644 --- a/extension/src/utilities.rs +++ b/extension/src/utilities.rs @@ -121,61 +121,65 @@ mod tests { #[pg_test] fn test_to_epoch() { - Spi::execute(|client| { + Spi::connect(|mut client| { let test_val = client - .select( + .update( "SELECT toolkit_experimental.to_epoch('2021-01-01 00:00:00+03'::timestamptz)", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() .unwrap(); assert!((test_val - 1609448400f64).abs() < f64::EPSILON); let test_val = client - .select( + .update( "SELECT toolkit_experimental.to_epoch('epoch'::timestamptz)", None, None, ) + .unwrap() .first() .get_one::() + .unwrap() .unwrap(); assert!((test_val - 0f64).abs() < f64::EPSILON); let test_val = client - .select("SELECT toolkit_experimental.to_epoch('epoch'::timestamptz - interval '42 seconds')", None, None) - .first() - .get_one::().unwrap(); + .update("SELECT toolkit_experimental.to_epoch('epoch'::timestamptz - interval '42 seconds')", None, None) + .unwrap().first() + .get_one::().unwrap().unwrap(); assert!((test_val - -42f64).abs() < f64::EPSILON); }); } #[pg_test] fn test_days_in_month() { - Spi::execute(|client| { - let test_val = client.select("SELECT toolkit_experimental.days_in_month('2021-01-01 00:00:00+03'::timestamptz)",None,None,).first().get_one::().unwrap(); + Spi::connect(|mut client| { + let test_val = client.update("SELECT toolkit_experimental.days_in_month('2021-01-01 00:00:00+03'::timestamptz)",None,None,).unwrap().first().get_one::().unwrap().unwrap(); assert_eq!(test_val, 31); }); - Spi::execute(|client| { - let test_val = client.select("SELECT toolkit_experimental.days_in_month('2020-02-03 00:00:00+03'::timestamptz)",None,None,).first().get_one::().unwrap(); + Spi::connect(|mut client| { + let test_val = client.update("SELECT toolkit_experimental.days_in_month('2020-02-03 00:00:00+03'::timestamptz)",None,None,).unwrap().first().get_one::().unwrap().unwrap(); assert_eq!(test_val, 29); }); } #[pg_test] fn test_monthly_normalize() { - Spi::execute(|client| { - let test_val = client.select("SELECT toolkit_experimental.month_normalize(1000,'2021-01-01 00:00:00+03'::timestamptz)",None,None,).first().get_one::().unwrap(); + Spi::connect(|mut client| { + let test_val = client.update("SELECT toolkit_experimental.month_normalize(1000,'2021-01-01 00:00:00+03'::timestamptz)",None,None,).unwrap().first().get_one::().unwrap().unwrap(); assert_eq!(test_val, 981.8548387096774f64); }); - Spi::execute(|client| { - let test_val = client.select("SELECT toolkit_experimental.month_normalize(1000,'2021-01-01 00:00:00+03'::timestamptz,30.5)",None,None,).first().get_one::().unwrap(); + Spi::connect(|mut client| { + let test_val = client.update("SELECT toolkit_experimental.month_normalize(1000,'2021-01-01 00:00:00+03'::timestamptz,30.5)",None,None,).unwrap().first().get_one::().unwrap().unwrap(); assert_eq!(test_val, 983.8709677419355f64); }); - Spi::execute(|client| { - let test_val = client.select("SELECT toolkit_experimental.month_normalize(1000,'2021-01-01 00:00:00+03'::timestamptz,30)",None,None,).first().get_one::().unwrap(); + Spi::connect(|mut client| { + let test_val = client.update("SELECT toolkit_experimental.month_normalize(1000,'2021-01-01 00:00:00+03'::timestamptz,30)",None,None,).unwrap().first().get_one::().unwrap().unwrap(); assert_eq!(test_val, 967.741935483871f64); }); } diff --git a/extension/timescaledb_toolkit.control b/extension/timescaledb_toolkit.control index 2a41926d..1577d608 100644 --- a/extension/timescaledb_toolkit.control +++ b/extension/timescaledb_toolkit.control @@ -5,4 +5,4 @@ superuser = false module_pathname = '$libdir/timescaledb_toolkit' # only for testing, will be removed for real installs # comma-separated list of previous versions this version can be upgraded from # directly. This is used to generate upgrade scripts. -# upgradeable_from = '1.4, 1.5, 1.5.1, 1.5.2, 1.6.0, 1.7.0, 1.8.0, 1.10.0-dev, 1.10.1, 1.11.0, 1.12.0, 1.12.1, 1.13.0, 1.13.1' +# upgradeable_from = '1.4, 1.5, 1.5.1, 1.5.2, 1.6.0, 1.7.0, 1.8.0, 1.10.0-dev, 1.10.1, 1.11.0, 1.12.0, 1.12.1, 1.13.0, 1.13.1, 1.14.0' diff --git a/tests/update/candlestick.md b/tests/update/candlestick.md index 9bfe1cde..3e6fd98f 100644 --- a/tests/update/candlestick.md +++ b/tests/update/candlestick.md @@ -38,8 +38,8 @@ ORDER BY symbol; --------+---------+---------+---------+---------+-------- AAPL | 133.445 | 133.445 | 133.44 | 133.445 | 30 AMZN | 95.225 | 95.225 | 95.225 | 95.225 | 1 - INTC | 29.82 | 29.82 | 29.82 | 29.82 | NULL + INTC | 29.82 | 29.82 | 29.82 | 29.82 | MSFT | 235.5 | 235.5 | 235.5 | 235.5 | 100 PFE | 47.38 | 47.38 | 47.38 | 47.38 | 2 - TSLA | 123.085 | 123.085 | 123.085 | 123.085 | NULL + TSLA | 123.085 | 123.085 | 123.085 | 123.085 | ``` diff --git a/tests/update/heartbeat.md b/tests/update/heartbeat.md new file mode 100644 index 00000000..aad0552e --- /dev/null +++ b/tests/update/heartbeat.md @@ -0,0 +1,53 @@ +# Candlestick Tests + +## Get candlestick values from tick data + + +```sql,creation,min-toolkit-version=1.15.0 +CREATE TABLE liveness(heartbeat TIMESTAMPTZ, start TIMESTAMPTZ); +INSERT INTO liveness VALUES + ('01-01-2020 0:2:20 UTC', '01-01-2020 0:0 UTC'), + ('01-01-2020 0:10 UTC', '01-01-2020 0:0 UTC'), + ('01-01-2020 0:17 UTC', '01-01-2020 0:0 UTC'), + ('01-01-2020 0:30 UTC', '01-01-2020 0:30 UTC'), + ('01-01-2020 0:35 UTC', '01-01-2020 0:30 UTC'), + ('01-01-2020 0:40 UTC', '01-01-2020 0:30 UTC'), + ('01-01-2020 0:35 UTC', '01-01-2020 0:30 UTC'), + ('01-01-2020 0:40 UTC', '01-01-2020 0:30 UTC'), + ('01-01-2020 0:40 UTC', '01-01-2020 0:30 UTC'), + ('01-01-2020 0:50:30 UTC', '01-01-2020 0:30 UTC'), + ('01-01-2020 1:00:30 UTC', '01-01-2020 1:00 UTC'), + ('01-01-2020 1:08 UTC', '01-01-2020 1:00 UTC'), + ('01-01-2020 1:18 UTC', '01-01-2020 1:00 UTC'), + ('01-01-2020 1:28 UTC', '01-01-2020 1:00 UTC'), + ('01-01-2020 1:38:01 UTC', '01-01-2020 1:30 UTC'), + ('01-01-2020 1:40 UTC', '01-01-2020 1:30 UTC'), + ('01-01-2020 1:40:01 UTC', '01-01-2020 1:30 UTC'), + ('01-01-2020 1:50:01 UTC', '01-01-2020 1:30 UTC'), + ('01-01-2020 1:57 UTC', '01-01-2020 1:30 UTC'), + ('01-01-2020 1:59:50 UTC', '01-01-2020 1:30 UTC'); + +CREATE MATERIALIZED VIEW hb AS + SELECT start, + heartbeat_agg(heartbeat, start, '30m', '10m') AS agg + FROM liveness + GROUP BY start; +``` + +```sql,validation,min-toolkit-version=1.15.0 +SELECT + start, + uptime(agg), + interpolated_uptime(agg, LAG(agg) OVER (ORDER by start)) +FROM hb +ORDER BY start; +``` + +```output + start | uptime | interpolated_uptime +------------------------+----------+--------------------- + 2020-01-01 00:00:00+00 | 00:24:40 | 00:24:40 + 2020-01-01 00:30:00+00 | 00:29:30 | 00:29:30 + 2020-01-01 01:00:00+00 | 00:29:30 | 00:30:00 + 2020-01-01 01:30:00+00 | 00:21:59 | 00:29:59 + ``` diff --git a/tests/update/state_agg.md b/tests/update/state_agg.md new file mode 100644 index 00000000..f2695f65 --- /dev/null +++ b/tests/update/state_agg.md @@ -0,0 +1,27 @@ +# `state_agg` tests + +```sql,creation,min-toolkit-version=1.15.0 +CREATE TABLE states_test(ts TIMESTAMPTZ, state TEXT); +INSERT INTO states_test VALUES + ('2020-01-01 00:00:00+00', 'START'), + ('2020-01-01 00:00:11+00', 'OK'), + ('2020-01-01 00:01:00+00', 'ERROR'), + ('2020-01-01 00:01:03+00', 'OK'), + ('2020-01-01 00:02:00+00', 'STOP'); + +CREATE TABLE agg(sa StateAgg); +INSERT INTO agg SELECT state_agg(ts, state) FROM states_test; +``` + +```sql,validation,min-toolkit-version=1.15.0 +SELECT (state_timeline(sa)).* FROM agg; +``` +```output + state | start_time | end_time +-------+------------------------+------------------------ + START | 2020-01-01 00:00:00+00 | 2020-01-01 00:00:11+00 + OK | 2020-01-01 00:00:11+00 | 2020-01-01 00:01:00+00 + ERROR | 2020-01-01 00:01:00+00 | 2020-01-01 00:01:03+00 + OK | 2020-01-01 00:01:03+00 | 2020-01-01 00:02:00+00 + STOP | 2020-01-01 00:02:00+00 | 2020-01-01 00:02:00+00 +``` diff --git a/tests/update/time-weighted-average.md b/tests/update/time-weighted-average.md new file mode 100644 index 00000000..ddea106d --- /dev/null +++ b/tests/update/time-weighted-average.md @@ -0,0 +1,55 @@ +# Time Weighted Average Tests + +## Test integral and interpolated integral + +```sql,creation,min-toolkit-version=1.15.0 +CREATE TABLE time_weight_test(time timestamptz, value double precision, bucket timestamptz); +INSERT INTO time_weight_test VALUES + ('2020-1-1 8:00'::timestamptz, 10.0, '2020-1-1'::timestamptz), + ('2020-1-1 12:00'::timestamptz, 40.0, '2020-1-1'::timestamptz), + ('2020-1-1 16:00'::timestamptz, 20.0, '2020-1-1'::timestamptz), + ('2020-1-2 2:00'::timestamptz, 15.0, '2020-1-2'::timestamptz), + ('2020-1-2 12:00'::timestamptz, 50.0, '2020-1-2'::timestamptz), + ('2020-1-2 20:00'::timestamptz, 25.0, '2020-1-2'::timestamptz), + ('2020-1-3 10:00'::timestamptz, 30.0, '2020-1-3'::timestamptz), + ('2020-1-3 12:00'::timestamptz, 0.0, '2020-1-3'::timestamptz), + ('2020-1-3 16:00'::timestamptz, 35.0, '2020-1-3'::timestamptz); +CREATE MATERIALIZED VIEW twa AS ( + SELECT bucket, time_weight('linear', time, value) as agg + FROM time_weight_test + GROUP BY bucket +); +``` + +```sql,validation,min-toolkit-version=1.15.0 +SELECT + interpolated_integral( + agg, + bucket, + '1 day'::interval, + LAG(agg) OVER (ORDER BY bucket), + LEAD(agg) OVER (ORDER BY bucket), + 'hours') +FROM twa +ORDER BY bucket; +``` + +```output + interpolated_integral +----------------------- + 364 + 758.8571428571429 + 382.1428571428571 +``` + +```sql,validation,min-toolkit-version=1.15.0 +SELECT integral(agg, 'hrs') FROM twa; +``` + +```output + integral +---------- + 200 + 60 + 550 +``` diff --git a/tools/build b/tools/build index a57a5d0c..1efdfdbc 100755 --- a/tools/build +++ b/tools/build @@ -116,12 +116,12 @@ while [ $# -gt 0 ]; do test-doc) find_profile require_pg_version - $nop cargo pgx start $pg + $nop cargo pgx start --package timescaledb_toolkit $pg $nop cargo run --profile $profile -p sql-doctester -- \ -h localhost \ -p $pg_port \ docs - $nop cargo pgx stop $pg + $nop cargo pgx stop --package timescaledb_toolkit $pg ;; test-updates) @@ -130,7 +130,7 @@ while [ $# -gt 0 ]; do find_pg_config require_cargo_pgx require_cargo_pgx_old - $nop cargo pgx start $pg + $nop cargo pgx start --package timescaledb_toolkit $pg $nop cargo run --profile $profile --manifest-path tools/update-tester/Cargo.toml -- full-update-test-source \ -h localhost \ -p $pg_port \ @@ -138,7 +138,7 @@ while [ $# -gt 0 ]; do "$pg_config" \ "$cargo_pgx" \ "$cargo_pgx_old" - $nop cargo pgx stop $pg + $nop cargo pgx stop --package timescaledb_toolkit $pg ;; *) diff --git a/tools/dependencies.sh b/tools/dependencies.sh index 365162bf..129ce5bc 100644 --- a/tools/dependencies.sh +++ b/tools/dependencies.sh @@ -14,9 +14,9 @@ TSDB_PG_VERSIONS='12 13 14' CARGO_EDIT=0.11.2 # Keep synchronized with extension/Cargo.toml and `cargo install --version N.N.N cargo-pgx` in Readme.md . -PGX_VERSION=0.6.1 +PGX_VERSION=0.7.1 -RUST_TOOLCHAIN=1.64.0 +RUST_TOOLCHAIN=1.65.0 RUST_PROFILE=minimal RUST_COMPONENTS=clippy,rustfmt diff --git a/tools/release b/tools/release index 4a9cd1bd..c2507372 100755 --- a/tools/release +++ b/tools/release @@ -166,6 +166,8 @@ minpat=${VERSION#*.} MINOR=${minpat%.*} PATCH=${minpat#*.} +POST_REL_BRANCH=post-$VERSION + # 0. Sanity-check the surroundings. # working directory clean? assert_clean || die 'cowardly refusing to operate on dirty working directory' @@ -228,7 +230,7 @@ for pg in $PG_VERSIONS; do done assert_clean || die 'tools/build should not dirty the working directory' -# 4. Push the branch so release-build-scripts repository +# 4. Push the branch if $push; then $nop git push origin $branch fi @@ -259,7 +261,7 @@ fi # 7. Prepare the main branch for the next release cycle. # Github action gives us a shallow checkout which we must deepen before we can push changes. $nop git fetch --deepen=2147483647 origin $MAIN_BRANCH -$nop git checkout $MAIN_BRANCH +$nop git checkout -b $POST_REL_BRANCH $MAIN_BRANCH # 7a. Update upgradeable_form in control file. $nop sed --in-place "/$UPGRADEABLE_FROM_RE/ { s/'\$/, $VERSION'/ }" $CONTROL @@ -321,9 +323,9 @@ set +e # TODO Carefully attempt to report errors but keep going on all steps after we push the tag in step 6. if $push; then - # 7d. Push to post-$VERSION branch. - $nop git push origin HEAD:post-$VERSION - $nop gh pr create -R timescale/timescaledb-toolkit -B $MAIN_BRANCH --fill -H post-$VERSION + # 7d. Push to $POST_REL_BRANCH branch. + $nop git push origin HEAD:$POST_REL_BRANCH + $nop gh pr create -R timescale/timescaledb-toolkit -B $MAIN_BRANCH --fill -H $POST_REL_BRANCH # 8. File issue for release tasks that are not yet automated. $nop gh issue create -R timescale/timescaledb-toolkit -F- -t "Release $VERSION" <