From 7c7132ea65a6b65253bd83da8806bfa617f5949f Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 17 Sep 2026 06:44:01 +0000 Subject: [PATCH] refactor(flow): execute streaming flows with DataFusion (#8976) * test(mito2): cover regex inverted index pruning Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(flow): execute streaming flows with DataFusion Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(flow): remove legacy streaming runtime Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(flow): avoid retrying stateless sink inserts Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(flow): align stateless writes with sink schema Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(flow): reject stale stateless source schemas Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(flow): validate stateless flow routing Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * Revert "test(mito2): cover regex inverted index pruning" This reverts commit 79e96ac745e810114c568be0bad6a4b849da2422. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(flow): preserve source timestamps in stateless flows Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(flow): address stateless review feedback Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(flow): recover stateless flows after schema changes Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(flow): serialize schema rebuilds and validate retained sources Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(flow): verify streaming recovery and schema changes through SQL Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(flow): cool down failed schema rebuilds Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(flow): explain legacy aggregate recreation requirements Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(flow): adapt stateless provider downcast to DataFusion 55 Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- Cargo.lock | 148 +- src/cmd/src/flownode.rs | 24 +- src/cmd/src/standalone.rs | 18 +- src/flow/AGENTS.md | 83 +- src/flow/Cargo.toml | 3 - src/flow/src/adapter.rs | 1811 ++++++++------- src/flow/src/adapter/flownode_impl.rs | 191 +- src/flow/src/adapter/node_context.rs | 458 ---- src/flow/src/adapter/parse_expr.rs | 245 -- src/flow/src/adapter/refill.rs | 440 ---- src/flow/src/adapter/stat.rs | 47 - src/flow/src/adapter/stateless.rs | 730 ++++++ src/flow/src/adapter/tests.rs | 1082 ++++++++- src/flow/src/adapter/util.rs | 61 +- src/flow/src/adapter/worker.rs | 606 ----- src/flow/src/batching_mode/frontend_client.rs | 118 + src/flow/src/batching_mode/time_window.rs | 8 +- src/flow/src/compute.rs | 23 - src/flow/src/compute/render.rs | 527 ----- src/flow/src/compute/render/map.rs | 425 ---- src/flow/src/compute/render/reduce.rs | 1985 ----------------- src/flow/src/compute/render/src_sink.rs | 245 -- src/flow/src/compute/state.rs | 167 -- src/flow/src/compute/types.rs | 208 -- src/flow/src/df_optimizer.rs | 46 +- src/flow/src/error.rs | 38 +- src/flow/src/expr.rs | 349 +-- src/flow/src/expr/df_func.rs | 300 --- src/flow/src/expr/error.rs | 32 +- src/flow/src/expr/func.rs | 1467 ------------ src/flow/src/expr/id.rs | 43 - src/flow/src/expr/linear.rs | 1174 ---------- src/flow/src/expr/relation.rs | 36 - src/flow/src/expr/relation/accum.rs | 1052 --------- src/flow/src/expr/relation/func.rs | 303 --- src/flow/src/expr/scalar.rs | 877 -------- src/flow/src/expr/signature.rs | 70 - src/flow/src/expr/utils.rs | 348 --- src/flow/src/lib.rs | 11 +- src/flow/src/plan.rs | 270 --- src/flow/src/plan/join.rs | 76 - src/flow/src/plan/reduce.rs | 87 - src/flow/src/repr.rs | 2 +- src/flow/src/repr/relation.rs | 97 +- src/flow/src/server.rs | 422 +--- src/flow/src/test_utils.rs | 66 - src/flow/src/transform.rs | 318 --- src/flow/src/transform/aggr.rs | 803 ------- src/flow/src/transform/expr.rs | 839 ------- src/flow/src/transform/literal.rs | 426 ---- src/flow/src/transform/plan.rs | 276 --- src/flow/src/utils.rs | 1002 +-------- src/operator/src/statement/ddl.rs | 199 +- tests-integration/src/standalone.rs | 19 - .../common/flow/flow_advance_ttl.result | 52 +- .../common/flow/flow_advance_ttl.sql | 35 +- .../common/flow/show_create_flow.result | 4 +- 57 files changed, 3385 insertions(+), 17407 deletions(-) delete mode 100644 src/flow/src/adapter/node_context.rs delete mode 100644 src/flow/src/adapter/parse_expr.rs delete mode 100644 src/flow/src/adapter/refill.rs delete mode 100644 src/flow/src/adapter/stat.rs create mode 100644 src/flow/src/adapter/stateless.rs delete mode 100644 src/flow/src/adapter/worker.rs delete mode 100644 src/flow/src/compute.rs delete mode 100644 src/flow/src/compute/render.rs delete mode 100644 src/flow/src/compute/render/map.rs delete mode 100644 src/flow/src/compute/render/reduce.rs delete mode 100644 src/flow/src/compute/render/src_sink.rs delete mode 100644 src/flow/src/compute/state.rs delete mode 100644 src/flow/src/compute/types.rs delete mode 100644 src/flow/src/expr/df_func.rs delete mode 100644 src/flow/src/expr/func.rs delete mode 100644 src/flow/src/expr/id.rs delete mode 100644 src/flow/src/expr/linear.rs delete mode 100644 src/flow/src/expr/relation.rs delete mode 100644 src/flow/src/expr/relation/accum.rs delete mode 100644 src/flow/src/expr/relation/func.rs delete mode 100644 src/flow/src/expr/scalar.rs delete mode 100644 src/flow/src/expr/signature.rs delete mode 100644 src/flow/src/expr/utils.rs delete mode 100644 src/flow/src/plan.rs delete mode 100644 src/flow/src/plan/join.rs delete mode 100644 src/flow/src/plan/reduce.rs delete mode 100644 src/flow/src/transform.rs delete mode 100644 src/flow/src/transform/aggr.rs delete mode 100644 src/flow/src/transform/expr.rs delete mode 100644 src/flow/src/transform/literal.rs delete mode 100644 src/flow/src/transform/plan.rs diff --git a/Cargo.lock b/Cargo.lock index 61005e91848..db47e98e355 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1289,7 +1289,7 @@ dependencies = [ "proc-macro2", "quote", "regex", - "rustc-hash 2.1.1", + "rustc-hash", "shlex 1.3.0", "syn 2.0.117", ] @@ -1741,12 +1741,6 @@ dependencies = [ "shlex 2.0.1", ] -[[package]] -name = "cc-traits" -version = "2.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "060303ef31ef4a522737e1b1ab68c67916f2a787bb2f4f54f383279adba962b5" - [[package]] name = "cedarwood" version = "0.5.0" @@ -4849,36 +4843,6 @@ dependencies = [ "syn 2.0.117", ] -[[package]] -name = "dfir_rs" -version = "0.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fd3b08cdfdba4b482d762906a16ac2cfd45d9949ae62e03ed3c690cfd7dae5dc" -dependencies = [ - "bincode", - "byteorder", - "bytes", - "futures", - "getrandom 0.2.16", - "itertools 0.13.0", - "lattices", - "pusherator", - "ref-cast", - "regex", - "rustc-hash 1.1.0", - "sealed", - "serde", - "serde_json", - "slotmap", - "smallvec", - "tokio", - "tokio-stream", - "tokio-util", - "tracing", - "variadics", - "web-time", -] - [[package]] name = "diff" version = "0.1.13" @@ -5126,18 +5090,6 @@ dependencies = [ "cfg-if", ] -[[package]] -name = "enum-as-inner" -version = "0.6.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a1e6a265c649f3f5979b601d26f1d05ada116434c87741c9493cb56218f76cbc" -dependencies = [ - "heck 0.5.0", - "proc-macro2", - "quote", - "syn 2.0.117", -] - [[package]] name = "enum_dispatch" version = "0.3.13" @@ -5506,7 +5458,6 @@ dependencies = [ "async-recursion", "async-trait", "bytes", - "cache", "catalog", "chrono", "client", @@ -5536,8 +5487,6 @@ dependencies = [ "datafusion-physical-expr", "datafusion-substrait", "datatypes", - "dfir_rs", - "enum-as-inner", "enum_dispatch", "futures", "get-size2", @@ -7340,7 +7289,7 @@ dependencies = [ "jieba-macros", "phf 0.13.1", "regex", - "rustc-hash 2.1.1", + "rustc-hash", ] [[package]] @@ -7753,33 +7702,6 @@ dependencies = [ "rustversion", ] -[[package]] -name = "lattices" -version = "0.6.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "51222529a85db7b6e228b279b2b9bb7bc1e47d9ee7b680404bdf5440b255913c" -dependencies = [ - "cc-traits", - "lattices_macro", - "ref-cast", - "sealed", - "serde", - "variadics", - "variadics_macro", -] - -[[package]] -name = "lattices_macro" -version = "0.5.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eee67863077d060b8a25c754241a6e6bf234b82207a62df7b3d5387b6d523c3e" -dependencies = [ - "proc-macro-crate 1.3.1", - "proc-macro2", - "quote", - "syn 2.0.117", -] - [[package]] name = "lazy-regex" version = "3.4.1" @@ -11869,16 +11791,6 @@ dependencies = [ "pulldown-cmark", ] -[[package]] -name = "pusherator" -version = "0.0.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3378af79ad42241f1075683daa2a0c06c6654404c3d096a493348395e620f668" -dependencies = [ - "either", - "variadics", -] - [[package]] name = "query" version = "1.3.0-alpha.1" @@ -12001,7 +11913,7 @@ dependencies = [ "pin-project-lite", "quinn-proto", "quinn-udp", - "rustc-hash 2.1.1", + "rustc-hash", "rustls", "socket2 0.5.10", "thiserror 2.0.17", @@ -12023,7 +11935,7 @@ dependencies = [ "rand 0.10.1", "rand_pcg", "ring", - "rustc-hash 2.1.1", + "rustc-hash", "rustls", "rustls-pki-types", "slab", @@ -12984,12 +12896,6 @@ version = "0.1.25" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "989e6739f80c4ad5b13e0fd7fe89531180375b18520cc8c82080e4dc4035b84f" -[[package]] -name = "rustc-hash" -version = "1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" - [[package]] name = "rustc-hash" version = "2.1.1" @@ -13334,18 +13240,6 @@ version = "4.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1c107b6f4780854c8b126e228ea8869f4d7b71260f962fefb57b996b8959ba6b" -[[package]] -name = "sealed" -version = "0.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f4a8caec23b7800fb97971a1c6ae365b6239aaeddfb934d6265f8505e795699d" -dependencies = [ - "heck 0.4.1", - "proc-macro2", - "quote", - "syn 2.0.117", -] - [[package]] name = "secrecy" version = "0.8.0" @@ -13947,15 +13841,6 @@ version = "0.4.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" -[[package]] -name = "slotmap" -version = "1.0.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dbff4acf519f630b3a3ddcfaea6c06b42174d9a44bc70c620e9ed1649d58b82a" -dependencies = [ - "version_check", -] - [[package]] name = "smallvec" version = "1.15.1" @@ -14899,7 +14784,7 @@ dependencies = [ "rayon", "regex", "rust-stemmers", - "rustc-hash 2.1.1", + "rustc-hash", "serde", "serde_json", "sketches-ddsketch", @@ -16329,29 +16214,6 @@ dependencies = [ "ryu", ] -[[package]] -name = "variadics" -version = "0.0.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41d49933ddcb1a4ec2ebb9a2931bc526ed269a9276a456c0595ddcc7cb188896" -dependencies = [ - "hashbrown 0.14.5", - "sealed", -] - -[[package]] -name = "variadics_macro" -version = "0.6.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "55476accb4f25086be8a1d18bed73aef9452cefb5ccd133e122e4350a1543ff3" -dependencies = [ - "proc-macro-crate 1.3.1", - "proc-macro2", - "quote", - "syn 2.0.117", - "variadics", -] - [[package]] name = "vcpkg" version = "0.2.15" diff --git a/src/cmd/src/flownode.rs b/src/cmd/src/flownode.rs index eccfaa63b07..0dfcf85e2e7 100644 --- a/src/cmd/src/flownode.rs +++ b/src/cmd/src/flownode.rs @@ -37,12 +37,9 @@ use common_stat::ResourceStatImpl; use common_telemetry::info; use common_telemetry::logging::{DEFAULT_LOGGING_DIR, TracingOptions}; use common_version::{short_version, verbose_version}; -use flow::{ - FlownodeBuilder, FlownodeInstance, FlownodeServiceBuilder, FrontendClient, FrontendInvoker, -}; +use flow::{FlownodeBuilder, FlownodeInstance, FlownodeServiceBuilder, FrontendClient}; use meta_client::{MetaClientOptions, MetaClientType}; use plugins::flownode::context::GrpcConfigureContext; -use servers::addrs; use servers::configurator::GrpcBuilderConfiguratorRef; use snafu::{OptionExt, ResultExt, ensure}; use tracing_appender::non_blocking::WorkerGuard; @@ -410,25 +407,6 @@ impl StartCommand { .build() .context(StartFlownodeSnafu)?; flownode.setup_services(services); - let flownode = flownode; - - let invoker = FrontendInvoker::build_from( - flownode.flow_engine().streaming_engine(), - catalog_manager.clone(), - cached_meta_backend.clone(), - layered_cache_registry.clone(), - meta_client.clone(), - client, - addrs::resolve_addr(&opts.grpc.bind_addr, Some(&opts.grpc.server_addr)), - ) - .await - .context(StartFlownodeSnafu)?; - flownode - .flow_engine() - .streaming_engine() - // TODO(discord9): refactor and avoid circular reference - .set_frontend_invoker(invoker) - .await; Ok(Instance::new(flownode, guard)) } diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index 6c5ac25f641..4573b0f129f 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -57,7 +57,7 @@ use datanode::datanode::{Datanode, DatanodeBuilder}; use datanode::region_server::RegionServer; use flow::{ FlowDualEngineRef, FlownodeBuilder, FlownodeInstance, FlownodeOptions, FrontendClient, - FrontendInvoker, GrpcQueryHandlerWithBoxedError, + GrpcQueryHandlerWithBoxedError, }; use frontend::frontend::Frontend; use frontend::instance::builder::FrontendBuilder; @@ -703,22 +703,6 @@ impl StartCommand { .set_handler(weak_grpc_handler) .await; - // set the frontend invoker for flownode - let flow_streaming_engine = flow_engine.streaming_engine(); - // flow server need to be able to use frontend to write insert requests back - let invoker = FrontendInvoker::build_from( - flow_streaming_engine.clone(), - catalog_manager.clone(), - kv_backend.clone(), - layered_cache_registry.clone(), - procedure_executor, - node_manager.clone(), - fe_instance.frontend_peer_addr().to_string(), - ) - .await - .context(StartFlownodeSnafu)?; - flow_streaming_engine.set_frontend_invoker(invoker).await; - let servers = Services::new(opts, fe_instance.clone(), plugins.clone()) .build() .context(error::StartFrontendSnafu)?; diff --git a/src/flow/AGENTS.md b/src/flow/AGENTS.md index 0ac0aa0e243..0b024e96fbd 100644 --- a/src/flow/AGENTS.md +++ b/src/flow/AGENTS.md @@ -1,20 +1,20 @@ # flow — Agent & Contributor Guide Navigation aid for `src/flow`. Keep it short and point to code. Paths are -relative to the repo root. +relative to the repository root. -Repo-wide rules that apply here: [`.agents/architecture-invariants.md`](../../.agents/architecture-invariants.md). +Repo-wide rules that apply to here: [`.agents/architecture-invariants.md`](../../.agents/architecture-invariants.md). ## What this crate does -Flownode is the stream-processing engine behind continuous aggregation / -materialized views. It has two execution paths: +Flownode provides two flow profiles: -- **Batching mode** (the actively developed path): splits data - into time windows and periodically runs aggregation SQL through the frontend, - writing results back to a sink table. -- **Streaming mode** (the legacy dataflow path): an incremental DFIR/dataflow - compute graph that processes row-level diffs. +- **Streaming mode**: a single stateless DataFusion query profile. Each mirror + insert is materialized as a transient input table, executed against the + retained logical plan, and written directly to the sink through + `batching_mode::frontend_client::FrontendClient`. +- **Batching mode**: splits data into time windows and periodically runs + aggregation SQL through the frontend, writing results back to a sink table. Users cannot select a mode directly: `flow_type` is a reserved internal option. `StatementExecutor::determine_flow_type` in `src/operator/src/statement/ddl.rs` @@ -26,16 +26,11 @@ changing routing rules. | Module | Path | Purpose | | --- | --- | --- | -| `engine` | `src/flow/src/engine.rs` | `FlowEngine` trait: create/remove/flush/insert lifecycle | -| `adapter` | `src/flow/src/adapter.rs`, `src/flow/src/adapter/` | `StreamingEngine`, worker pool, dual-engine dispatch, table sources/sinks | -| `batching_mode` | `src/flow/src/batching_mode.rs`, `src/flow/src/batching_mode/` | `BatchingEngine`, task scheduling, time windows, frontend client, checkpoints | -| `compute` | `src/flow/src/compute/` | Streaming dataflow render/state | -| `expr` | `src/flow/src/expr.rs`, `src/flow/src/expr/` | Scalar/aggregate expressions and Map-Filter-Project | -| `plan` | `src/flow/src/plan.rs` | `TypedPlan` (reduce/join/MFP) | -| `transform` | `src/flow/src/transform.rs` | Substrait → flow plan | -| `df_optimizer` | `src/flow/src/df_optimizer.rs` | SQL → DataFusion logical plan → optimized plan | -| `repr` | `src/flow/src/repr.rs` | `Row`, `DiffRow`, `Batch`, `RelationDesc` | -| `server` | `src/flow/src/server.rs` | gRPC `Flow` service, `FlownodeBuilder`/`FlownodeInstance` | +| `engine` | `src/flow/src/engine.rs` | `FlowEngine` lifecycle contract | +| `adapter` | `src/flow/src/adapter.rs`, `src/flow/src/adapter/` | Stateless streaming execution, dual-engine dispatch, shared source schema/default normalization, and sink creation | +| `batching_mode` | `src/flow/src/batching_mode.rs`, `src/flow/src/batching_mode/` | `BatchingEngine`, scheduling, time windows, frontend client, and checkpoints | +| `repr` | `src/flow/src/repr.rs` | Shared row and relation schema representations | +| `server` | `src/flow/src/server.rs` | gRPC `Flow` service and flownode builders | | `heartbeat` | `src/flow/src/heartbeat.rs` | Reports flownode state/stats to metasrv | Flow metadata lives in `common-meta`, not here: @@ -43,53 +38,37 @@ Flow metadata lives in `common-meta`, not here: ## Data flow -`Frontend → Flownode (gRPC)` → `FlowService` (`server.rs`) → -`FlowDualEngine` (`adapter/flownode_impl.rs`) routes by `FlowType`: +`Frontend → Flownode (gRPC)` → `FlowService` (`server.rs`) → `FlowDualEngine` +(`adapter/flownode_impl.rs`) routes by `FlowType`: +- Streaming: mirror rows are normalized against the source schema, evaluated by + the retained stateless DataFusion plan, and sent to the sink via + `FrontendClient`. - Batching: marks dirty windows; a task later runs aggregation SQL via the - frontend client and writes the sink table (`batching_mode/task.rs`). -- Streaming: worker threads apply incremental diffs and push to the sink - (`adapter/worker.rs`, `compute/render.rs`). + frontend client and writes the sink table. + +`FlowType::Streaming` and mirror routing are shared contracts; preserve both +when changing dispatch. Preserve batching behavior and the shared schema, +default, and sink-creation helpers. ## Public surface -- gRPC `Flow` service in `src/flow/src/server.rs` - (`handle_create_remove`, `handle_mirror_request`, `handle_mark_dirty_time_window`). -- Sink writes for the streaming engine go through `FrontendInvoker` - (`row_inserts`, `row_deletes`). +- gRPC `Flow` service in `src/flow/src/server.rs`. - `FlowEngine` trait in `src/flow/src/engine.rs`. - Started from the `cmd` crate via `FlownodeBuilder` / `FlownodeInstance`. -## When you change X, also touch Y - -- **Flow definition / options**: validation in `common-meta`'s `ddl/create_flow.rs` - and the serialized `FlowInfoValue` in `common-meta`'s `key/flow/`. -- **New scalar/aggregate function** (`expr/`): also wire up evaluation in - `compute/render.rs` (streaming) and ensure batching SQL handles it. -- **Persisted flow metadata**: keep `FlowInfoValue` backward compatible - (`serde(default)` / `serde(alias)`). -- **Shared routing, metadata, or sink contracts**: check both engines. An - engine-specific implementation change does not automatically need a mirror - change in the other path. - ## Testing ```bash cargo nextest run -p flow ``` -Helpers in `src/flow/src/test_utils.rs` (test context, test query engine). +Stateless streaming tests are under `src/flow/src/adapter/stateless.rs`. ## Gotchas -- Batching vs streaming differ in latency, state, and execution. Confirm the - selected mode before reasoning about a flow. -- Streaming workers are `!Send`; cross-thread interaction goes through - `WorkerHandle`, not the worker directly. -- Internal flow timestamps (`repr::Timestamp`, ms) are not necessarily the - table's time column; window functions key off the diff timestamp. - -## Maintenance contract - -Update this file when the dual-engine routing, the gRPC surface, or the flow -metadata contract (shared with `common-meta`) changes. +- Batching and streaming differ in latency and state. Confirm the selected mode + before reasoning about a flow. +- Streaming is finite and stateless per mirror request; it has no worker graph, + background runtime, or replay state. +- Internal flow timestamps and table time columns are separate contracts. diff --git a/src/flow/Cargo.toml b/src/flow/Cargo.toml index e5603fd8d82..4586e612693 100644 --- a/src/flow/Cargo.toml +++ b/src/flow/Cargo.toml @@ -14,7 +14,6 @@ arrow-schema.workspace = true async-recursion = "1.0" async-trait.workspace = true bytes.workspace = true -cache.workspace = true catalog.workspace = true chrono.workspace = true client.workspace = true @@ -43,8 +42,6 @@ datafusion-expr.workspace = true datafusion-physical-expr.workspace = true datafusion-substrait.workspace = true datatypes.workspace = true -dfir_rs = { version = "0.13.0", default-features = false } -enum-as-inner = "0.6.0" enum_dispatch = "0.3" futures.workspace = true get-size2 = "0.1.2" diff --git a/src/flow/src/adapter.rs b/src/flow/src/adapter.rs index c2f95e12588..c98510bd7d5 100644 --- a/src/flow/src/adapter.rs +++ b/src/flow/src/adapter.rs @@ -12,110 +12,385 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! for getting data from source and sending results to sink -//! and communicating with other parts of the database -#![warn(unused_imports)] +//! Flow source schema management and stateless streaming execution. use std::collections::BTreeMap; use std::sync::Arc; -use std::time::{Duration, Instant, SystemTime}; +#[cfg(test)] +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::{AtomicBool, Ordering}; -use api::v1::{RowDeleteRequest, RowDeleteRequests, RowInsertRequest, RowInsertRequests}; use common_base::memory_limit::MemoryLimit; use common_config::Configurable; use common_error::ext::BoxedError; +use common_meta::distributed_time_constants::BASE_HEARTBEAT_INTERVAL; use common_meta::key::TableMetadataManagerRef; use common_options::memory::MemoryOptions; -use common_runtime::JoinHandle; +use common_recordbatch::map_dictionary_to_values_data_type; use common_stat::get_total_cpu_cores; use common_telemetry::logging::{LoggingOptions, TracingOptions}; -use common_telemetry::{debug, info, trace}; -use datatypes::schema::ColumnSchema; -use datatypes::value::Value; -use greptime_proto::v1; -use itertools::{EitherOrBoth, Itertools}; +use common_telemetry::{error, info}; +use datafusion_common::TableReference; +use datafusion_common::tree_node::TreeNode; +use datafusion_expr::logical_plan::Distinct; +use datafusion_expr::{Expr, LogicalPlan}; +use datatypes::schema::{ColumnSchema, SchemaRef}; +use itertools::Itertools; use meta_client::MetaClientOptions; use query::QueryEngine; use query::options::QueryOptions; use serde::{Deserialize, Serialize}; use servers::grpc::GrpcOptions; use servers::http::HttpOptions; -use session::context::QueryContext; -use snafu::{OptionExt, ResultExt, ensure}; +use snafu::{IntoError, OptionExt, ResultExt, ensure}; use store_api::storage::{ConcreteDataType, RegionId}; -use table::metadata::TableId; -use tokio::sync::broadcast::error::TryRecvError; -use tokio::sync::{Mutex, RwLock, broadcast, watch}; +use tokio::sync::RwLock; +use tokio::time::Instant; -pub(crate) use crate::adapter::node_context::FlownodeContext; -use crate::adapter::refill::RefillTask; +use crate::adapter::stateless::StatelessFlow; use crate::adapter::table_source::ManagedTableSource; -use crate::adapter::util::relation_desc_to_column_schemas_with_fallback; -pub(crate) use crate::adapter::worker::{Worker, WorkerHandle, create_worker}; +use crate::adapter::util::{ + relation_desc_to_column_schemas_with_fallback, table_info_value_to_relation_desc, +}; use crate::batching_mode::BatchingModeOptions; -use crate::compute::ErrCollector; -use crate::df_optimizer::sql_to_flow_plan; -use crate::error::{EvalSnafu, ExternalSnafu, InternalSnafu, InvalidQuerySnafu, UnexpectedSnafu}; -use crate::expr::Batch; -use crate::metrics::{METRIC_FLOW_INSERT_ELAPSED, METRIC_FLOW_ROWS, METRIC_FLOW_RUN_INTERVAL_MS}; -use crate::repr::{self, BATCH_SIZE, DiffRow, RelationDesc, Row}; +use crate::batching_mode::frontend_client::FrontendClient; +use crate::batching_mode::utils::sql_to_df_plan; +use crate::error::{ + DatafusionSnafu, Error, ExternalSnafu, FlowNotFoundSnafu, InsertIntoFlowSnafu, InternalSnafu, + InvalidQuerySnafu, UnexpectedSnafu, +}; +use crate::repr::{ColumnType, DiffRow, RelationDesc, Row}; use crate::{CreateFlowArgs, FlowId, TableName}; pub(crate) mod flownode_impl; -mod parse_expr; -pub(crate) mod refill; -mod stat; +pub(crate) mod stateless; +pub(crate) mod table_source; #[cfg(test)] mod tests; pub(crate) mod util; -mod worker; -pub(crate) mod node_context; -pub(crate) mod table_source; - -use crate::FrontendInvoker; -use crate::error::Error; - -fn expire_after_secs_to_millis(expire_after_secs: i64) -> Result { - ensure!( - expire_after_secs >= 0, - InvalidQuerySnafu { - reason: format!("EXPIRE AFTER must be non-negative, got {expire_after_secs} seconds"), - } - ); - - expire_after_secs - .checked_mul(1_000) - .with_context(|| InvalidQuerySnafu { - reason: format!( - "EXPIRE AFTER value {expire_after_secs} seconds cannot be represented in milliseconds" - ), +/// Converts a retained plan's output to logical schemas. Only a direct source +/// column (optionally wrapped in an alias) carries source semantics; computed +/// expressions deliberately use the physical output field metadata instead. +fn output_column_schemas( + plan: &LogicalPlan, + source_schema: &SchemaRef, +) -> Result<(Vec, Vec>), Error> { + // Plain DISTINCT preserves the selected columns and their source semantics. Inspect its + // input for lineage, while leaving computed expressions without source metadata. + let expressions = match plan { + LogicalPlan::Projection(projection) => Some(&projection.expr), + LogicalPlan::Distinct(Distinct::All(input)) => match input.as_ref() { + LogicalPlan::Projection(projection) => Some(&projection.expr), + _ => None, + }, + _ => None, + }; + let is_pass_through = matches!(plan, LogicalPlan::TableScan(_) | LogicalPlan::Filter(_)) + || matches!(plan, LogicalPlan::Distinct(Distinct::All(input)) if matches!(input.as_ref(), LogicalPlan::TableScan(_) | LogicalPlan::Filter(_))); + let mut lineage = Vec::new(); + let columns = plan + .schema() + .fields() + .iter() + .enumerate() + .map(|(idx, field)| { + let mut column = ColumnSchema::try_from(field.as_ref()) + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + let source_index = expressions + .and_then(|exprs| { + let expr = exprs.get(idx)?; + let expr = match expr { + Expr::Column(column) => Some(column), + Expr::Alias(alias) => match alias.expr.as_ref() { + Expr::Column(column) => Some(column), + _ => None, + }, + _ => None, + }?; + source_schema.column_index_by_name(&expr.name) + }) + .or_else(|| { + is_pass_through + .then(|| source_schema.column_index_by_name(field.name())) + .flatten() + }); + if let Some(source_index) = source_index { + let mut source = source_schema.column_schemas()[source_index].clone(); + source.name = field.name().clone(); + source.data_type = map_dictionary_to_values_data_type(&source.data_type); + column = source; + } else { + column.data_type = map_dictionary_to_values_data_type(&column.data_type); + // DataFusion may propagate field metadata through a computed expression. Such + // an expression has no source-column lineage and must not inherit its PK/time + // semantics into the sink relation. + if expressions.is_some() { + column = column.with_time_index(false); + } + } + lineage.push(source_index); + Ok(column) }) + .collect::, Error>>()?; + Ok((columns, lineage)) } -// `GREPTIME_TIMESTAMP` is not used to distinguish when table is created automatically by flow -pub const AUTO_CREATED_PLACEHOLDER_TS_COL: &str = "__ts_placeholder"; +fn relation_desc_from_output( + columns: &[ColumnSchema], + lineage: &[Option], + source_primary_key_indices: &[usize], +) -> RelationDesc { + let keys = source_primary_key_indices + .iter() + .filter_map(|source_index| { + lineage + .iter() + .position(|index| index == &Some(*source_index)) + }) + .collect_vec(); + let time_index = columns.iter().position(ColumnSchema::is_time_index); + RelationDesc { + typ: crate::repr::RelationType { + column_types: columns + .iter() + .map(|column| ColumnType::new(column.data_type.clone(), column.is_nullable())) + .collect(), + keys: if keys.is_empty() { + vec![] + } else { + vec![crate::repr::Key::from(keys)] + }, + time_index, + auto_columns: vec![], + }, + names: columns + .iter() + .map(|column| Some(column.name.clone())) + .collect(), + } +} +fn default_num_workers() -> usize { + get_total_cpu_cores().div_ceil(2) +} + +/// Returns whether an existing sink uses the legacy explicit source-time-index layout. +/// +/// This check is deliberately separate from [`resolve_sink_layout`]: ordinary auto-column +/// resolution remains the first choice, and this compatibility path is only for a sink whose +/// trailing column is a real, user-defined time index. +pub(crate) fn is_explicit_source_timestamp_compatibility( + output_schema: &[ColumnSchema], + output_lineage: &[Option], + sink_schema: &[ColumnSchema], + source_schema: &SchemaRef, +) -> bool { + if sink_schema.len() != output_schema.len() + 1 || output_lineage.len() != output_schema.len() { + return false; + } + if !output_schema + .iter() + .zip(&sink_schema[..output_schema.len()]) + .all(|(output, sink)| output.data_type == sink.data_type) + { + return false; + } + + let Some(source_timestamp_index) = source_schema.timestamp_index() else { + return false; + }; + let sink_timestamp = &sink_schema[output_schema.len()]; + sink_timestamp.data_type == source_schema.column_schemas()[source_timestamp_index].data_type + && sink_timestamp.data_type.is_timestamp() + && sink_timestamp.is_time_index() + && sink_timestamp.default_constraint().is_some() + && sink_timestamp.name != AUTO_CREATED_UPDATE_AT_TS_COL + && sink_timestamp.name != AUTO_CREATED_PLACEHOLDER_TS_COL + && !output_lineage.contains(&Some(source_timestamp_index)) +} + +pub const AUTO_CREATED_PLACEHOLDER_TS_COL: &str = "__ts_placeholder"; pub const AUTO_CREATED_UPDATE_AT_TS_COL: &str = "update_at"; -/// Flow config that exists both in standalone&distributed mode +/// Resolves the columns appended by the flow, validating the complete output/sink layout. +/// +/// A sink with the same arity as the query is an ordinary sink, even when its last +/// column happens to be named `update_at`. Auto columns are only inferred from the +/// arity difference and their exact trailing layout. +pub(crate) fn resolve_sink_layout( + output_schema: &[ColumnSchema], + sink_schema: &[ColumnSchema], +) -> Result, Error> { + ensure!( + sink_schema.len() >= output_schema.len() && sink_schema.len() - output_schema.len() <= 2, + InvalidQuerySnafu { + reason: format!( + "Flow output has {} columns, but sink has {} columns; only zero, one, or two trailing auto columns are supported", + output_schema.len(), + sink_schema.len() + ) + } + ); + let suffix_len = sink_schema.len() - output_schema.len(); + for (idx, (output, sink)) in output_schema.iter().zip(sink_schema.iter()).enumerate() { + ensure!( + output.data_type == sink.data_type, + InvalidQuerySnafu { + reason: format!( + "Flow output column {idx} has type {:?}, but sink column {} has type {:?}", + output.data_type, sink.name, sink.data_type + ) + } + ); + } + let suffix = &sink_schema[output_schema.len()..]; + match suffix_len { + 0 => Ok(vec![]), + 1 => { + let column = &suffix[0]; + ensure!( + column.name == AUTO_CREATED_UPDATE_AT_TS_COL && column.data_type.is_timestamp(), + InvalidQuerySnafu { + reason: format!( + "The trailing sink column must be timestamp {}", + AUTO_CREATED_UPDATE_AT_TS_COL + ) + } + ); + Ok(suffix.to_vec()) + } + 2 => { + let update_at = &suffix[0]; + let placeholder = &suffix[1]; + ensure!( + update_at.name == AUTO_CREATED_UPDATE_AT_TS_COL + && update_at.data_type.is_timestamp() + && placeholder.name == AUTO_CREATED_PLACEHOLDER_TS_COL + && placeholder.data_type.is_timestamp() + && placeholder.is_time_index(), + InvalidQuerySnafu { + reason: "The two trailing sink columns must be timestamp update_at followed by timestamp time-index __ts_placeholder".to_string() + } + ); + Ok(suffix.to_vec()) + } + _ => unreachable!(), + } +} + +/// Legacy helper for callers that only have a sink schema. New stateless flows +/// resolve the suffix against both schemas and store it in `StatelessFlow`. +pub(crate) fn sink_output_column_count(sink_schema: &[ColumnSchema]) -> Result { + let mut count = sink_schema.len(); + if sink_schema + .last() + .is_some_and(|column| column.name == AUTO_CREATED_PLACEHOLDER_TS_COL) + { + let placeholder = &sink_schema[count - 1]; + ensure!( + placeholder.data_type.is_timestamp() && placeholder.is_time_index(), + InvalidQuerySnafu { + reason: format!( + "Auto-created sink column {} must be a timestamp time index", + AUTO_CREATED_PLACEHOLDER_TS_COL + ) + } + ); + count -= 1; + } + if sink_schema + .get(count.saturating_sub(1)) + .is_some_and(|column| column.name == AUTO_CREATED_UPDATE_AT_TS_COL) + { + ensure!( + sink_schema[count - 1].data_type.is_timestamp(), + InvalidQuerySnafu { + reason: format!( + "Auto-created sink column {} must be a timestamp", + AUTO_CREATED_UPDATE_AT_TS_COL + ) + } + ); + count -= 1; + } + Ok(count) +} + +pub(crate) fn validate_sink_layout( + output_schema: &[ColumnSchema], + sink_schema: &[ColumnSchema], +) -> Result<(), Error> { + resolve_sink_layout(output_schema, sink_schema).map(|_| ()) +} + +pub(crate) fn validate_auto_column_names(output_schema: &[ColumnSchema]) -> Result<(), Error> { + for column in output_schema { + ensure!( + column.name != AUTO_CREATED_UPDATE_AT_TS_COL + && column.name != AUTO_CREATED_PLACEHOLDER_TS_COL, + InvalidQuerySnafu { + reason: format!( + "Flow output column {} is reserved for an auto-created sink column", + column.name + ) + } + ); + } + Ok(()) +} + +pub(crate) fn validate_sink_layout_with_suffix( + output_schema: &[ColumnSchema], + sink_schema: &[ColumnSchema], + suffix: &[ColumnSchema], +) -> Result<(), Error> { + ensure!( + sink_schema.len() == output_schema.len() + suffix.len() + && sink_schema[output_schema.len()..] == *suffix, + InvalidQuerySnafu { + reason: "Stored sink auto-column layout no longer matches the sink schema".to_string() + } + ); + for (idx, (output, sink)) in output_schema + .iter() + .zip(&sink_schema[..output_schema.len()]) + .enumerate() + { + ensure!( + output.data_type == sink.data_type, + InvalidQuerySnafu { + reason: format!( + "Flow output column {idx} has type {:?}, but sink column {} has type {:?}", + output.data_type, sink.name, sink.data_type + ) + } + ); + } + Ok(()) +} + #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] #[serde(default)] pub struct FlowConfig { + /// Deprecated and ignored. Flow workers have been removed. + #[deprecated(note = "flow workers have been removed; this field is ignored")] + #[serde(default = "default_num_workers")] pub num_workers: usize, pub batching_mode: BatchingModeOptions, } +#[allow(deprecated)] impl Default for FlowConfig { fn default() -> Self { Self { - num_workers: (get_total_cpu_cores() / 2).max(1), + num_workers: get_total_cpu_cores().div_ceil(2), batching_mode: BatchingModeOptions::default(), } } } -/// Options for flow node #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] #[serde(default)] pub struct FlownodeOptions { @@ -140,8 +415,6 @@ impl Default for FlownodeOptions { meta_client: None, logging: LoggingOptions::default(), tracing: TracingOptions::default(), - // flownode's query option is set to 1 to throttle flow's query so - // that it won't use too much cpu or memory query: QueryOptions { parallelism: 1, allow_query_fallback: false, @@ -156,328 +429,557 @@ impl Default for FlownodeOptions { impl Configurable for FlownodeOptions { fn validate_sanitize(&mut self) -> common_config::error::Result<()> { - if self.flow.num_workers == 0 { - self.flow.num_workers = (get_total_cpu_cores() / 2).max(1); - } Ok(()) } } -/// Arc-ed FlowNodeManager, cheaper to clone pub type FlowStreamingEngineRef = Arc; -/// FlowNodeManager manages the state of all tasks in the flow node, which should be run on the same thread -/// -/// The choice of timestamp is just using current system timestamp for now -/// -pub struct StreamingEngine { - /// The handler to the worker that will run the dataflow - /// which is `!Send` so a handle is used - pub worker_handles: Vec, - /// The selector to select a worker to run the dataflow - worker_selector: Mutex, - /// The query engine that will be used to parse the query and convert it to a dataflow plan - pub query_engine: Arc, - /// Getting table name and table schema from table info manager - table_info_source: ManagedTableSource, - frontend_invoker: RwLock>, - /// contains mapping from table name to global id, and table schema - node_context: RwLock, - /// Contains all refill tasks - refill_tasks: RwLock>, - flow_err_collectors: RwLock>, - src_send_buf_lens: RwLock>>, - tick_manager: FlowTickManager, - /// This node id is only available in distributed mode, on standalone mode this is guaranteed to be `None` - pub node_id: Option, - /// Lock for flushing, will be `read` by `handle_inserts` and `write` by `flush_flow` - /// - /// So that a series of event like `inserts -> flush` can be handled correctly - flush_lock: RwLock<()>, +#[derive(Default)] +struct StatelessFlowRuntime { + flow: Option>, + failed_rebuild: Option<(u32, Instant)>, } -/// Building FlownodeManager -impl StreamingEngine { - /// set frontend invoker - pub async fn set_frontend_invoker(&self, frontend: FrontendInvoker) { - *self.frontend_invoker.write().await = Some(frontend); - } +struct StatelessFlowSlot { + runtime: Arc>, + /// Cleared at registry-detach time, fencing creators that already captured this slot. + active: std::sync::atomic::AtomicBool, + #[cfg(test)] + rebuild_attempts: AtomicUsize, +} - /// Create **without** setting `frontend_invoker` +fn validate_captured_slot( + slot: &StatelessFlowSlot, + current_source_table_id: Option, + expected_table_id: table::metadata::TableId, + flow_id: FlowId, +) -> Result<(), Error> { + ensure!( + slot.active.load(Ordering::Acquire), + FlowNotFoundSnafu { id: flow_id } + ); + let current_source_table_id = + current_source_table_id.context(FlowNotFoundSnafu { id: flow_id })?; + ensure!( + current_source_table_id == expected_table_id, + InvalidQuerySnafu { + reason: format!("Flow {flow_id} source table changed while it was selected") + } + ); + Ok(()) +} + +pub struct StreamingEngine { + pub query_engine: Arc, + pub frontend_client: Arc, + table_info_source: ManagedTableSource, + stateless_flows: RwLock>>, + pub node_id: Option, +} + +impl StreamingEngine { pub fn new( node_id: Option, query_engine: Arc, table_meta: TableMetadataManagerRef, + frontend_client: Arc, ) -> Self { - let srv_map = ManagedTableSource::new( + let table_info_source = ManagedTableSource::new( table_meta.table_info_manager().clone(), table_meta.table_name_manager().clone(), ); - let node_context = FlownodeContext::new(Box::new(srv_map.clone()) as _); - let tick_manager = FlowTickManager::new(); - let worker_handles = Vec::new(); - StreamingEngine { - worker_handles, - worker_selector: Mutex::new(0), + Self { query_engine, - table_info_source: srv_map, - frontend_invoker: RwLock::new(None), - node_context: RwLock::new(node_context), - refill_tasks: Default::default(), - flow_err_collectors: Default::default(), - src_send_buf_lens: Default::default(), - tick_manager, + frontend_client, + table_info_source, + stateless_flows: Default::default(), node_id, - flush_lock: RwLock::new(()), } } - /// Create a flownode manager with one worker - pub fn new_with_workers<'s>( - node_id: Option, - query_engine: Arc, - table_meta: TableMetadataManagerRef, - num_workers: usize, - ) -> (Self, Vec>) { - let mut zelf = Self::new(node_id, query_engine, table_meta); - - let workers: Vec<_> = (0..num_workers) - .map(|_| { - let (handle, worker) = create_worker(); - zelf.add_worker_handle(handle); - worker - }) - .collect(); - (zelf, workers) - } - - /// add a worker handler to manager, meaning this corresponding worker is under it's manage - pub fn add_worker_handle(&mut self, handle: WorkerHandle) { - self.worker_handles.push(handle); - } -} - -#[derive(Debug)] -pub enum DiffRequest { - Insert(Vec<(Row, repr::Timestamp)>), - Delete(Vec<(Row, repr::Timestamp)>), -} - -impl DiffRequest { - pub fn len(&self) -> usize { - match self { - Self::Insert(v) => v.len(), - Self::Delete(v) => v.len(), - } - } - - pub fn is_empty(&self) -> bool { - self.len() == 0 - } -} - -pub fn batches_to_rows_req(batches: Vec) -> Result, Error> { - let mut reqs = Vec::new(); - for batch in batches { - let mut rows = Vec::with_capacity(batch.row_count()); - for i in 0..batch.row_count() { - let row = batch.get_row(i).context(EvalSnafu)?; - rows.push((Row::new(row), 0)); - } - reqs.push(DiffRequest::Insert(rows)); - } - Ok(reqs) -} - -/// This impl block contains methods to send writeback requests to frontend -impl StreamingEngine { - /// Return the number of requests it made - pub async fn send_writeback_requests(&self) -> Result { - let all_reqs = self.generate_writeback_request().await?; - if all_reqs.is_empty() || all_reqs.iter().all(|v| v.1.is_empty()) { - return Ok(0); - } - let mut req_cnt = 0; - for (table_name, reqs) in all_reqs { - if reqs.is_empty() { - continue; - } - let (catalog, schema) = (table_name[0].clone(), table_name[1].clone()); - let ctx = Arc::new(QueryContext::with(&catalog, &schema)); - - let (is_ts_placeholder, proto_schema) = match self - .try_fetch_existing_table(&table_name) - .await? - .context(UnexpectedSnafu { - reason: format!("Table not found: {}", table_name.join(".")), - }) { - Ok(r) => r, - Err(e) => { - if self - .table_info_source - .get_opt_table_id_from_name(&table_name) - .await? - .is_none() - { - // deal with both flow&sink table no longer exists - // but some output is still in output buf - common_telemetry::warn!(e; "Table `{}` no longer exists, skip writeback", table_name.join(".")); - continue; - } else { - return Err(e); - } - } - }; - let schema_len = proto_schema.len(); - - let total_rows = reqs.iter().map(|r| r.len()).sum::(); - trace!( - "Sending {} writeback requests to table {}, reqs total rows={}", - reqs.len(), - table_name.join("."), - reqs.iter().map(|r| r.len()).sum::() - ); - - METRIC_FLOW_ROWS - .with_label_values(&["out-streaming"]) - .inc_by(total_rows as u64); - - let now = self.tick_manager.tick(); - for req in reqs { - match req { - DiffRequest::Insert(insert) => { - let rows_proto: Vec = insert - .into_iter() - .map(|(mut row, _ts)| { - // extend `update_at` col if needed - // if schema include a millisecond timestamp here, and result row doesn't have it, add it - if row.len() < proto_schema.len() - && proto_schema[row.len()].datatype - == greptime_proto::v1::ColumnDataType::TimestampMillisecond - as i32 - { - row.extend([Value::from( - common_time::Timestamp::new_millisecond(now), - )]); - } - // ts col, if auto create - if is_ts_placeholder { - ensure!( - row.len() == schema_len - 1, - InternalSnafu { - reason: format!( - "Row len mismatch, expect {} got {}", - schema_len - 1, - row.len() - ) - } - ); - row.extend([Value::from( - common_time::Timestamp::new_millisecond(0), - )]); - } - if row.len() != proto_schema.len() { - UnexpectedSnafu { - reason: format!( - "Flow output row length mismatch, expect {} got {}, the columns in schema are: {:?}", - proto_schema.len(), - row.len(), - proto_schema.iter().map(|c|&c.column_name).collect_vec() - ), - } - .fail()?; - } - Ok(row.into()) - }) - .collect::, Error>>()?; - let table_name = table_name.last().unwrap().clone(); - let req = RowInsertRequest { - table_name, - rows: Some(v1::Rows { - schema: proto_schema.clone(), - rows: rows_proto, - }), - }; - req_cnt += 1; - self.frontend_invoker - .read() - .await - .as_ref() - .with_context(|| UnexpectedSnafu { - reason: "Expect a frontend invoker for flownode to write back", - })? - .row_inserts(RowInsertRequests { inserts: vec![req] }, ctx.clone()) - .await - .map_err(BoxedError::new) - .with_context(|_| ExternalSnafu {})?; - } - DiffRequest::Delete(remove) => { - info!("original remove rows={:?}", remove); - let rows_proto: Vec = remove - .into_iter() - .map(|(mut row, _ts)| { - row.extend(Some(Value::from( - common_time::Timestamp::new_millisecond(0), - ))); - row.into() - }) - .collect::>(); - let table_name = table_name.last().unwrap().clone(); - let req = RowDeleteRequest { - table_name, - rows: Some(v1::Rows { - schema: proto_schema.clone(), - rows: rows_proto, - }), - }; - - req_cnt += 1; - self.frontend_invoker - .read() - .await - .as_ref() - .with_context(|| UnexpectedSnafu { - reason: "Expect a frontend invoker for flownode to write back", - })? - .row_deletes(RowDeleteRequests { deletes: vec![req] }, ctx.clone()) - .await - .map_err(BoxedError::new) - .with_context(|_| ExternalSnafu {})?; - } - } - } - } - Ok(req_cnt) - } - - /// Generate writeback request for all sink table - pub async fn generate_writeback_request( + pub async fn handle_write_request( &self, - ) -> Result>, Error> { - trace!("Start to generate writeback request"); - let mut output = BTreeMap::new(); - let mut total_row_count = 0; - for (name, sink_recv) in self - .node_context + region_id: RegionId, + rows: Vec, + batch_datatypes: &[ConcreteDataType], + source_schema_version: u32, + ) -> Result<(), Error> { + let table_id = region_id.table_id(); + let flow_ids = self.flow_ids_for_table(table_id).await; + let mut failed_flow_ids = Vec::new(); + let mut first_error = None; + for (flow_id, slot) in flow_ids { + let result = self + .execute_flow( + flow_id, + slot, + table_id, + rows.clone(), + batch_datatypes, + source_schema_version, + ) + .await; + if let Err(err) = result { + error!(err; "Failed to insert into flow={}, region_id={}", flow_id, region_id); + failed_flow_ids.push(flow_id); + if first_error.is_none() { + first_error = Some(BoxedError::new(err)); + } + } + } + match first_error { + Some(source) => Err(InsertIntoFlowSnafu { + region_id: u64::from(region_id), + flow_ids: failed_flow_ids, + } + .into_error(source)), + None => Ok(()), + } + } + + async fn flow_ids_for_table( + &self, + table_id: table::metadata::TableId, + ) -> Vec<(FlowId, Arc)> { + let slots = self + .stateless_flows + .read() + .await + .iter() + .map(|(id, slot)| (*id, Arc::clone(slot))) + .collect::>(); + let mut flows = Vec::new(); + for (id, slot) in slots { + if slot + .runtime + .read() + .await + .flow + .as_ref() + .is_some_and(|flow| flow.source_table_id == table_id) + { + flows.push((id, slot)); + } + } + flows + } + + async fn execute_flow( + &self, + flow_id: FlowId, + slot: Arc, + expected_table_id: table::metadata::TableId, + rows: Vec, + batch_datatypes: &[ConcreteDataType], + source_schema_version: u32, + ) -> Result { + // Keep the captured lifecycle slot rather than resolving the ID again after a drop. + let mut guard = slot.runtime.clone().read_owned().await; + let current = guard + .flow + .as_ref() + .context(FlowNotFoundSnafu { id: flow_id })?; + validate_captured_slot( + &slot, + Some(current.source_table_id), + expected_table_id, + flow_id, + )?; + if current.source_schema_version != source_schema_version { + if guard + .failed_rebuild + .as_ref() + .is_some_and(|(version, until)| { + *version == source_schema_version && *until > Instant::now() + }) + { + return InvalidQuerySnafu { + reason: format!( + "Flow {flow_id} schema rebuild for source version {source_schema_version} is cooling down" + ), + } + .fail(); + } + drop(guard); + // Serialize rebuilds with definition replacement using the existing publication + // lease. Another write may already have rebuilt this schema while we waited. + let mut published = slot.runtime.clone().write_owned().await; + let current = published + .flow + .as_ref() + .context(FlowNotFoundSnafu { id: flow_id })?; + validate_captured_slot( + &slot, + Some(current.source_table_id), + expected_table_id, + flow_id, + )?; + if current.source_schema_version != source_schema_version { + if published + .failed_rebuild + .as_ref() + .is_some_and(|(version, until)| { + *version == source_schema_version && *until > Instant::now() + }) + { + return InvalidQuerySnafu { + reason: format!( + "Flow {flow_id} schema rebuild for source version {source_schema_version} is cooling down" + ), + } + .fail(); + } + #[cfg(test)] + slot.rebuild_attempts.fetch_add(1, Ordering::Relaxed); + let replacement = self + .build_stateless_flow(¤t.create_args, false) + .await + .and_then(|replacement| { + ensure!( + replacement.source_table_id == expected_table_id + && replacement.source_schema_version == source_schema_version, + InvalidQuerySnafu { + reason: format!( + "Source schema changed while rebuilding flow {flow_id}" + ) + } + ); + Ok(replacement) + }); + match replacement { + Ok(replacement) => { + published.flow = Some(Arc::new(replacement)); + published.failed_rebuild = None; + } + Err(error) => { + // Reuse the default heartbeat retry baseline; this is neither a + // negotiated interval nor a cache-freshness guarantee. + published.failed_rebuild = Some(( + source_schema_version, + Instant::now() + BASE_HEARTBEAT_INTERVAL, + )); + return Err(error); + } + } + } + // Do not open a replacement/drop gap between preparation and sink execution. + guard = tokio::sync::OwnedRwLockWriteGuard::downgrade(published); + } + let flow = guard + .flow + .as_ref() + .context(FlowNotFoundSnafu { id: flow_id })?; + // This is the last check before planning. In particular, a request normalized against + // an old source schema is never allowed to reach a newly published plan. + let latest = self + .table_info_source + .get_table_info_value(&flow.source_table_id) + .await? + .context(UnexpectedSnafu { + reason: "Source table metadata is missing", + })? + .table_info + .meta + .schema + .version(); + ensure!( + source_schema_version == latest && flow.source_schema_version == latest, + InvalidQuerySnafu { + reason: format!("Source schema version changed before flow {flow_id} execution") + } + ); + stateless::execute( + flow, + &rows, + batch_datatypes, + &self.query_engine, + &self.frontend_client, + latest, + ) + .await + } + + pub async fn remove_flow_inner(&self, flow_id: FlowId) -> Result<(), Error> { + // Detach first: this is the tombstone/linearization point. A concurrent create must + // install a different slot and can never republish into the removed one. + let slot = self + .stateless_flows .write() .await - .sink_receiver - .iter_mut() - .map(|(n, (_s, r))| (n, r)) + .remove(&flow_id) + .context(FlowNotFoundSnafu { id: flow_id })?; + slot.active.store(false, Ordering::Release); + let mut runtime = slot.runtime.write().await; + runtime.flow.take(); + runtime.failed_rebuild = None; + Ok(()) + } + + async fn publish_initial_flow( + &self, + slot: Arc, + flow: Arc, + create_if_not_exists: bool, + or_replace: bool, + ) -> Result { + let mut runtime = slot.runtime.write().await; + ensure!( + slot.active.load(Ordering::Acquire), + FlowNotFoundSnafu { + id: flow.create_args.flow_id + } + ); + if runtime.flow.is_some() && !or_replace { + if create_if_not_exists { + return Ok(false); + } + return crate::error::FlowAlreadyExistSnafu { + id: flow.create_args.flow_id, + } + .fail(); + } + let latest = self + .table_info_source + .get_table_info_value(&flow.source_table_id) + .await? + .context(UnexpectedSnafu { + reason: "Source table metadata is missing", + })? + .table_info + .meta + .schema + .version(); + ensure!( + latest == flow.source_schema_version, + InvalidQuerySnafu { + reason: format!( + "Source schema changed while building flow: built version {}, current version {latest}", + flow.source_schema_version + ) + } + ); + runtime.flow = Some(flow); + runtime.failed_rebuild = None; + Ok(true) + } + + pub async fn create_flow_inner(&self, args: CreateFlowArgs) -> Result, Error> { + let flow_id = args.flow_id; + let flow = Arc::new(self.build_stateless_flow(&args, true).await?); + let slot = { + let mut slots = self.stateless_flows.write().await; + if let Some(slot) = slots.get(&flow_id) { + Arc::clone(slot) + } else { + let slot = Arc::new(StatelessFlowSlot { + runtime: Arc::new(RwLock::new(StatelessFlowRuntime::default())), + active: AtomicBool::new(true), + #[cfg(test)] + rebuild_attempts: AtomicUsize::new(0), + }); + slots.insert(flow_id, Arc::clone(&slot)); + slot + } + }; + let published = self + .publish_initial_flow( + slot.clone(), + flow, + args.create_if_not_exists, + args.or_replace, + ) + .await?; + ensure!( + self.stateless_flows + .read() + .await + .get(&flow_id) + .is_some_and(|current| Arc::ptr_eq(current, &slot)), + FlowNotFoundSnafu { id: flow_id } + ); + if published { + info!("Successfully create flow with id={flow_id}"); + Ok(Some(flow_id)) + } else { + Ok(None) + } + } + + async fn build_stateless_flow( + &self, + args: &CreateFlowArgs, + create_sink: bool, + ) -> Result { + let CreateFlowArgs { + flow_id, + sink_table_name, + source_table_ids, + sql, + query_ctx, + .. + } = args; + ensure!( + source_table_ids.len() == 1, + InvalidQuerySnafu { + reason: "Stateless streaming flow does not support multiple source tables", + } + ); + + let query_ctx = query_ctx.clone().map(Arc::new).context(UnexpectedSnafu { + reason: "Query context is missing", + })?; + let source_table_id = source_table_ids[0]; + let source_table_name = self + .table_info_source + .get_table_name(&source_table_id) + .await?; + let source_table_info = self + .table_info_source + .get_table_info_value(&source_table_id) + .await? + .context(UnexpectedSnafu { + reason: "Source table metadata is missing", + })?; + let flow_plan = + sql_to_df_plan(query_ctx.clone(), self.query_engine.clone(), sql, true).await?; + stateless::validate_plan(&flow_plan)?; + let source_meta = source_table_info.table_info.meta; + let source_schema = source_meta.schema; + let source_primary_key_indices = source_meta.primary_key_indices; + stateless::validate_source_scan(&flow_plan, source_table_id, &source_schema)?; + let (inferred_schema, lineage) = output_column_schemas(&flow_plan, &source_schema)?; + let inferred_relation = + relation_desc_from_output(&inferred_schema, &lineage, &source_primary_key_indices); + let sink_exists = self.fetch_table_pk_schema(sink_table_name).await?.is_some(); + if !sink_exists { + validate_auto_column_names(&inferred_schema)?; + } + if !sink_exists + && create_sink + && !self + .create_table_from_relation( + &format!("flow-id={flow_id}"), + sink_table_name, + &inferred_relation, + ) + .await? { - let mut batches = Vec::new(); - while let Ok(batch) = sink_recv.try_recv() { - total_row_count += batch.row_count(); - batches.push(batch); + return UnexpectedSnafu { + reason: format!("Failed to create table {sink_table_name:?}"), + } + .fail(); + } + ensure!( + sink_exists || create_sink, + UnexpectedSnafu { + reason: format!("Sink table metadata is missing: {sink_table_name:?}"), + } + ); + + // Fetch the metadata after validation or auto-creation. The sink's layout is the insert + // contract: flow output aliases must not leak into the request schema. + let (sink_primary_keys, _, sink_schema) = self + .fetch_table_pk_schema(sink_table_name) + .await? + .context(UnexpectedSnafu { + reason: format!("Sink table metadata is missing: {sink_table_name:?}"), + })?; + let (auto_columns, plan) = match resolve_sink_layout(&inferred_schema, &sink_schema) { + Ok(auto_columns) => (auto_columns, flow_plan), + Err(normal_error) => { + // Appending a hidden source timestamp changes DISTINCT's key. It is therefore + // never a valid compatibility rewrite for a DISTINCT flow. + let has_distinct = { + let mut found = false; + flow_plan + .apply(|node| { + if matches!(node, LogicalPlan::Distinct(Distinct::All(_))) { + found = true; + } + Ok(datafusion_common::tree_node::TreeNodeRecursion::Continue) + }) + .context(DatafusionSnafu { + context: "Failed to inspect flow plan", + })?; + found + }; + if has_distinct + || !sink_exists + || !is_explicit_source_timestamp_compatibility( + &inferred_schema, + &lineage, + &sink_schema, + &source_schema, + ) + { + return Err(normal_error); + } + let source_timestamp_index = source_schema.timestamp_index().unwrap(); + let source_timestamp_name = + &source_schema.column_schemas()[source_timestamp_index].name; + let plan = stateless::rewrite_source_timestamp( + flow_plan, + &TableReference::full( + source_table_name[0].clone(), + source_table_name[1].clone(), + source_table_name[2].clone(), + ), + source_timestamp_name, + )?; + let (effective_output, _) = output_column_schemas(&plan, &source_schema)?; + ensure!( + effective_output.len() == sink_schema.len() + && effective_output + .iter() + .zip(&sink_schema) + .all(|(output, sink)| output.data_type == sink.data_type), + InvalidQuerySnafu { + reason: "Compatibility plan output does not match the full sink schema" + } + ); + (vec![], plan) + } + }; + + Ok(StatelessFlow { + source_table_id, + source_table_name, + source_schema_version: source_schema.version(), + source_schema, + sink_table_name: sink_table_name.clone(), + sink_schema, + sink_primary_keys, + auto_columns, + plan, + query_ctx, + create_args: args.clone(), + }) + } + + pub async fn flush_flow_inner(&self, _flow_id: FlowId) -> Result { + Ok(0) + } + + pub(crate) async fn stateless_flow_ids(&self) -> Vec { + let slots = self.stateless_flows.read().await; + let mut ids = Vec::new(); + for (id, slot) in slots.iter() { + if slot.runtime.read().await.flow.is_some() { + ids.push(*id); } - let reqs = batches_to_rows_req(batches)?; - output.insert(name.clone(), reqs); } - trace!("Prepare writeback req: total row count={}", total_row_count); - Ok(output) + ids + } + + pub async fn flow_exist_inner(&self, flow_id: FlowId) -> Result { + let slot = self.stateless_flows.read().await.get(&flow_id).cloned(); + Ok(match slot { + Some(slot) => slot.runtime.read().await.flow.is_some(), + None => false, + }) } - /// Fetch table schema and primary key from table info source, if table not exist return None async fn fetch_table_pk_schema( &self, table_name: &TableName, @@ -499,30 +1001,22 @@ impl StreamingEngine { .into_iter() .map(|i| schema[i].name.clone()) .collect_vec(); - let time_index = meta.schema.timestamp_index(); - Ok(Some((primary_keys, time_index, schema))) + Ok(Some((primary_keys, meta.schema.timestamp_index(), schema))) } else { Ok(None) } } - /// return (primary keys, schema and if the table have a placeholder timestamp column) - /// schema of the table comes from flow's output plan - /// - /// adjust to add `update_at` column and ts placeholder if needed async fn adjust_auto_created_table_schema( &self, schema: &RelationDesc, ) -> Result<(Vec, Vec, bool), Error> { - // TODO(discord9): consider remove buggy auto create by schema - - // TODO(discord9): use default key from schema let primary_keys = schema .typ() .keys .first() - .map(|v| { - v.column_indices + .map(|key| { + key.column_indices .iter() .map(|i| { schema @@ -533,431 +1027,210 @@ impl StreamingEngine { .collect_vec() }) .unwrap_or_default(); - let update_at = ColumnSchema::new( + let mut columns = relation_desc_to_column_schemas_with_fallback(schema); + columns.push(ColumnSchema::new( AUTO_CREATED_UPDATE_AT_TS_COL, ConcreteDataType::timestamp_millisecond_datatype(), true, - ); - - let original_schema = relation_desc_to_column_schemas_with_fallback(schema); - - let mut with_auto_added_col = original_schema.clone(); - with_auto_added_col.push(update_at); - - // if no time index, add one as placeholder + )); let no_time_index = schema.typ().time_index.is_none(); if no_time_index { - let ts_col = ColumnSchema::new( - AUTO_CREATED_PLACEHOLDER_TS_COL, - ConcreteDataType::timestamp_millisecond_datatype(), - true, - ) - .with_time_index(true); - with_auto_added_col.push(ts_col); - } - - Ok((primary_keys, with_auto_added_col, no_time_index)) - } -} - -/// Flow Runtime related methods -impl StreamingEngine { - /// run in common_runtime background runtime - pub fn run_background( - self: Arc, - shutdown: Option>, - ) -> JoinHandle<()> { - info!("Starting flownode manager's background task"); - common_runtime::spawn_global(async move { self.run(shutdown).await }) - } - - /// log all flow errors - pub async fn log_all_errors(&self) { - for (f_id, f_err) in self.flow_err_collectors.read().await.iter() { - let all_errors = f_err.get_all().await; - if !all_errors.is_empty() { - let all_errors = all_errors - .into_iter() - .map(|i| format!("{:?}", i)) - .join("\n"); - common_telemetry::error!("Flow {} has following errors: {}", f_id, all_errors); - } - } - } - - /// Trigger dataflow running, and then send writeback request to the source sender - /// - /// note that this method didn't handle input mirror request, as this should be handled by grpc server - pub async fn run(&self, mut shutdown: Option>) { - debug!("Starting to run"); - let default_interval = Duration::from_secs(1); - let mut tick_interval = tokio::time::interval(default_interval); - // burst mode, so that if we miss a tick, we will run immediately to fully utilize the cpu - tick_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Burst); - let mut avg_spd = 0; // rows/sec - let mut since_last_run = tokio::time::Instant::now(); - let run_per_trace = 10; - let mut run_cnt = 0; - loop { - // TODO(discord9): only run when new inputs arrive or scheduled to - let row_cnt = self.run_available(false).await.unwrap_or_else(|err| { - common_telemetry::error!(err;"Run available errors"); - 0 - }); - - if let Err(err) = self.send_writeback_requests().await { - common_telemetry::error!(err;"Send writeback request errors"); - }; - self.log_all_errors().await; - - // determine if need to shutdown - match &shutdown.as_mut().map(|s| s.try_recv()) { - Some(Ok(())) => { - info!("Shutdown flow's main loop"); - break; - } - Some(Err(TryRecvError::Empty)) => (), - Some(Err(TryRecvError::Closed)) => { - common_telemetry::error!("Shutdown channel is closed"); - break; - } - Some(Err(TryRecvError::Lagged(num))) => { - common_telemetry::error!( - "Shutdown channel is lagged by {}, meaning multiple shutdown cmd have been issued", - num - ); - break; - } - None => (), - } - - // for now we want to batch rows until there is around `BATCH_SIZE` rows in send buf - // before trigger a run of flow's worker - let wait_for = since_last_run.elapsed(); - - // last runs insert speed - let cur_spd = row_cnt * 1000 / wait_for.as_millis().max(1) as usize; - // rapid increase, slow decay - avg_spd = if cur_spd > avg_spd { - cur_spd - } else { - (9 * avg_spd + cur_spd) / 10 - }; - let new_wait = BATCH_SIZE * 1000 / avg_spd.max(1); //in ms - let new_wait = Duration::from_millis(new_wait as u64).min(default_interval); - - // print trace every `run_per_trace` times so that we can see if there is something wrong - // but also not get flooded with trace - if run_cnt >= run_per_trace { - trace!("avg_spd={} r/s, cur_spd={} r/s", avg_spd, cur_spd); - trace!("Wait for {} ms, row_cnt={}", new_wait.as_millis(), row_cnt); - run_cnt = 0; - } else { - run_cnt += 1; - } - - METRIC_FLOW_RUN_INTERVAL_MS.set(new_wait.as_millis() as i64); - since_last_run = tokio::time::Instant::now(); - tokio::select! { - _ = tick_interval.tick() => (), - _ = tokio::time::sleep(new_wait) => () - } - } - // flow is now shutdown, drop frontend_invoker early so a ref cycle(in standalone mode) can be prevent: - // FlowWorkerManager.frontend_invoker -> FrontendInvoker.inserter - // -> Inserter.node_manager -> NodeManager.flownode -> Flownode.flow_streaming_engine.frontend_invoker - self.frontend_invoker.write().await.take(); - } - - /// Run all available subgraph in the flow node - /// This will try to run all dataflow in this node - /// - /// set `blocking` to true to wait until worker finish running - /// false to just trigger run and return immediately - /// return numbers of rows send to worker(Inaccuary) - /// TODO(discord9): add flag for subgraph that have input since last run - pub async fn run_available(&self, blocking: bool) -> Result { - let mut row_cnt = 0; - - let now = self.tick_manager.tick(); - for worker in self.worker_handles.iter() { - // TODO(discord9): consider how to handle error in individual worker - worker.run_available(now, blocking).await?; - } - // check row send and rows remain in send buf - let flush_res = if blocking { - let ctx = self.node_context.read().await; - ctx.flush_all_sender().await - } else { - match self.node_context.try_read() { - Ok(ctx) => ctx.flush_all_sender().await, - Err(_) => return Ok(row_cnt), - } - }; - match flush_res { - Ok(r) => { - common_telemetry::trace!("Total flushed {} rows", r); - row_cnt += r; - } - Err(err) => { - common_telemetry::error!("Flush send buf errors: {:?}", err); - } - }; - - Ok(row_cnt) - } - - /// send write request to related source sender - pub async fn handle_write_request( - &self, - region_id: RegionId, - rows: Vec, - batch_datatypes: &[ConcreteDataType], - ) -> Result<(), Error> { - let rows_len = rows.len(); - let table_id = region_id.table_id(); - let _timer = METRIC_FLOW_INSERT_ELAPSED - .with_label_values(&[table_id.to_string().as_str()]) - .start_timer(); - self.node_context - .read() - .await - .send(table_id, rows, batch_datatypes) - .await?; - trace!( - "Handling write request for table_id={} with {} rows", - table_id, rows_len - ); - Ok(()) - } -} - -/// Create&Remove flow -impl StreamingEngine { - /// remove a flow by it's id - pub async fn remove_flow_inner(&self, flow_id: FlowId) -> Result<(), Error> { - for handle in self.worker_handles.iter() { - if handle.contains_flow(flow_id).await? { - handle.remove_flow(flow_id).await?; - break; - } - } - self.node_context.write().await.remove_flow(flow_id); - Ok(()) - } - - /// Return task id if a new task is created, otherwise return None - /// - /// steps to create task: - /// 1. parse query into typed plan(and optional parse expire_after expr) - /// 2. render source/sink with output table id and used input table id - pub async fn create_flow_inner(&self, args: CreateFlowArgs) -> Result, Error> { - let CreateFlowArgs { - flow_id, - sink_table_name, - source_table_ids, - create_if_not_exists, - or_replace, - expire_after: expire_after_secs, - eval_interval: _, - comment, - sql, - flow_options, - query_ctx, - .. - } = args; - let expire_after = expire_after_secs - .map(expire_after_secs_to_millis) - .transpose()?; - - let mut node_ctx = self.node_context.write().await; - // assign global id to source and sink table - for source in &source_table_ids { - node_ctx - .assign_global_id_to_table(&self.table_info_source, None, Some(*source)) - .await?; - } - node_ctx - .assign_global_id_to_table(&self.table_info_source, Some(sink_table_name.clone()), None) - .await?; - - node_ctx.register_task_src_sink(flow_id, &source_table_ids, sink_table_name.clone()); - - node_ctx.query_context = query_ctx.map(Arc::new); - // construct a active dataflow state with it - let flow_plan = sql_to_flow_plan(&mut node_ctx, &self.query_engine, &sql).await?; - - debug!("Flow {:?}'s Plan is {:?}", flow_id, flow_plan); - - // check schema against actual table schema if exists - // if not exist create sink table immediately - if let Some((_, _, real_schema)) = self.fetch_table_pk_schema(&sink_table_name).await? { - let auto_schema = relation_desc_to_column_schemas_with_fallback(&flow_plan.schema); - - // for column schema, only `data_type` need to be check for equality - // since one can omit flow's column name when write flow query - // print a user friendly error message about mismatch and how to correct them - for (idx, zipped) in auto_schema - .iter() - .zip_longest(real_schema.iter()) - .enumerate() - { - match zipped { - EitherOrBoth::Both(auto, real) => { - if auto.data_type != real.data_type { - InvalidQuerySnafu { - reason: format!( - "Column {}(name is '{}', flow inferred name is '{}')'s data type mismatch, expect {:?} got {:?}", - idx, - real.name, - auto.name, - real.data_type, - auto.data_type - ), - } - .fail()?; - } - } - EitherOrBoth::Right(real) if real.data_type.is_timestamp() => { - // if table is auto created, the last one or two column should be timestamp(update at and ts placeholder) - continue; - } - _ => InvalidQuerySnafu { - reason: format!( - "schema length mismatched, expected {} found {}", - real_schema.len(), - auto_schema.len() - ), - } - .fail()?, - } - } - } else { - // assign inferred schema to sink table - // create sink table - let did_create = self - .create_table_from_relation( - &format!("flow-id={flow_id}"), - &sink_table_name, - &flow_plan.schema, + columns.push( + ColumnSchema::new( + AUTO_CREATED_PLACEHOLDER_TS_COL, + ConcreteDataType::timestamp_millisecond_datatype(), + true, ) - .await?; - if !did_create { - UnexpectedSnafu { - reason: format!("Failed to create table {:?}", sink_table_name), + .with_time_index(true), + ); + } + Ok((primary_keys, columns, no_time_index)) + } +} + +impl StreamingEngine { + async fn handle_inserts_inner( + &self, + request: api::v1::region::InsertRequests, + ) -> Result<(), Error> { + // A mirrored envelope can contain several regions of one source. Normalize each request, + // then concatenate by source so DISTINCT is evaluated once over the complete envelope. + let mut grouped: BTreeMap<_, (_, Vec, Vec, u32)> = + BTreeMap::new(); + let mut poisoned_tables = std::collections::HashSet::new(); + let mut first_error = None; + for write_request in request.requests { + let region_id = RegionId::from(write_request.region_id); + let table_id = region_id.table_id(); + if poisoned_tables.contains(&table_id) { + continue; + } + match self.handle_insert_request(write_request).await { + Ok((rows, types, version)) => { + let entry = grouped + .entry(table_id) + .or_insert_with(|| (region_id, vec![], types.clone(), version)); + if entry.2 != types || entry.3 != version { + poisoned_tables.insert(table_id); + grouped.remove(&table_id); + let err = InvalidQuerySnafu { reason: format!("Source table {table_id} metadata changed within one insert envelope") }.build(); + let ids: Vec = self + .flow_ids_for_table(table_id) + .await + .into_iter() + .map(|(id, _)| id) + .collect(); + let err = InsertIntoFlowSnafu { + region_id: u64::from(region_id), + flow_ids: ids, + } + .into_error(BoxedError::new(err)); + error!(err; "Failed to normalize flow insert request for region_id={region_id}"); + if first_error.is_none() { + first_error = Some(err); + } + } else { + entry.1.extend(rows); + } + } + Err(err) => { + let ids: Vec = self + .flow_ids_for_table(table_id) + .await + .into_iter() + .map(|(id, _)| id) + .collect(); + poisoned_tables.insert(table_id); + grouped.remove(&table_id); + let err = InsertIntoFlowSnafu { + region_id: u64::from(region_id), + flow_ids: ids, + } + .into_error(BoxedError::new(err)); + error!(err; "Failed to normalize flow insert request for region_id={region_id}"); + if first_error.is_none() { + first_error = Some(err); + } } - .fail()?; } } + for (_, (region_id, rows, types, version)) in grouped { + if let Err(err) = self + .handle_write_request(region_id, rows, &types, version) + .await + && first_error.is_none() + { + first_error = Some(err); + } + } + match first_error { + Some(err) => Err(err), + None => Ok(()), + } + } - node_ctx.add_flow_plan(flow_id, flow_plan.clone()); - - let _ = comment; - let _ = flow_options; - - // TODO(discord9): add more than one handles - let sink_id = node_ctx.table_repr.get_by_name(&sink_table_name).unwrap().1; - let sink_sender = node_ctx.get_sink_by_global_id(&sink_id)?; - - let source_ids = source_table_ids - .iter() - .map(|id| node_ctx.table_repr.get_by_table_id(id).unwrap().1) - .collect_vec(); - let source_receivers = source_ids - .iter() - .map(|id| { - node_ctx - .get_source_by_global_id(id) - .map(|s| s.get_receiver()) - }) - .collect::, _>>()?; - let err_collector = ErrCollector::default(); - self.flow_err_collectors - .write() - .await - .insert(flow_id, err_collector.clone()); - // TODO(discord9): load balance? - let handle = self.get_worker_handle_for_create_flow().await; - let create_request = worker::Request::Create { - flow_id, - plan: flow_plan, - sink_id, - sink_sender, - source_ids, - src_recvs: source_receivers, - expire_after, - or_replace, - create_if_not_exists, - err_collector, + async fn handle_insert_request( + &self, + write_request: api::v1::region::InsertRequest, + ) -> Result<(Vec, Vec, u32), Error> { + let region_id = write_request.region_id; + let table_id = RegionId::from(region_id).table_id(); + let (insert_schema, rows_proto) = write_request + .rows + .map(|r| (r.schema, r.rows)) + .unwrap_or_default(); + let now = common_time::util::current_time_millis(); + let (table_types, fetch_order, source_schema_version) = { + // Fetch the source metadata once for both current-schema normalization and + // retained-plan validation. In particular, do not execute a retained plan + // against rows normalized with a newer schema. + let table_info = self + .table_info_source + .get_table_info_value(&table_id) + .await? + .context(UnexpectedSnafu { + reason: format!("Table metadata is missing for table id {table_id}"), + })?; + let source_schema_version = table_info.table_info.meta.schema.version(); + let table_schema = table_info_value_to_relation_desc(table_info)?; + let defaults = table_schema + .default_values + .iter() + .zip(table_schema.relation_desc.typ().column_types.iter()) + .map(|(value, ty)| { + value.as_ref().and_then(|value| { + value.create_default(ty.scalar_type(), ty.nullable()).ok() + }) + }) + .collect_vec(); + let types = table_schema + .relation_desc + .typ() + .column_types + .iter() + .map(|ty| ty.scalar_type.clone()) + .collect_vec(); + let names = table_schema + .relation_desc + .names + .iter() + .enumerate() + .map(|(idx, name)| { + name.clone().context(InternalSnafu { + reason: format!("Column {idx} of table {table_id} has no name"), + }) + }) + .collect::, _>>()?; + let input_columns = insert_schema + .iter() + .enumerate() + .map(|(idx, column)| (&column.column_name, idx)) + .collect::>(); + let order = names + .iter() + .zip(defaults) + .map(|(name, default)| { + input_columns + .get(name) + .copied() + .map(FetchFromRow::Idx) + .or_else(|| default.map(FetchFromRow::Default)) + .with_context(|| UnexpectedSnafu { + reason: format!("Column not found: {name}"), + }) + }) + .collect::, _>>()?; + (types, order, source_schema_version) }; - - handle.create_flow(create_request).await?; - info!("Successfully create flow with id={}", flow_id); - Ok(Some(flow_id)) - } - - pub async fn flush_flow_inner(&self, flow_id: FlowId) -> Result { - debug!("Starting to flush flow_id={:?}", flow_id); - // lock to make sure writes before flush are written to flow - // and immediately drop to prevent following writes to be blocked - drop(self.flush_lock.write().await); - let flushed_input_rows = self.node_context.read().await.flush_all_sender().await?; - let rows_send = self.run_available(true).await?; - let row = self.send_writeback_requests().await?; - debug!( - "Done to flush flow_id={:?} with {} input rows flushed, {} rows sent and {} output rows flushed", - flow_id, flushed_input_rows, rows_send, row - ); - Ok(row) - } - - pub async fn flow_exist_inner(&self, flow_id: FlowId) -> Result { - let mut exist = false; - for handle in self.worker_handles.iter() { - if handle.contains_flow(flow_id).await? { - exist = true; - break; - } - } - Ok(exist) + let rows = rows_proto + .into_iter() + .map(|row| { + let row = Row::from(row); + let row = fetch_order + .iter() + .map(|item| item.fetch(&row)) + .collect_vec(); + (Row::new(row), now, 1) + }) + .collect_vec(); + Ok((rows, table_types, source_schema_version)) } } -/// FlowTickManager is a manager for flow tick, which trakc flow execution progress -/// -/// TODO(discord9): better way to do it, and not expose flow tick even to other flow to avoid -/// TSO coord mess -#[derive(Clone, Debug)] -pub struct FlowTickManager { - /// The starting instant of the flow, used with `start_timestamp` to calculate the current timestamp - start: Instant, - /// The timestamp when the flow started - start_timestamp: repr::Timestamp, +#[derive(Debug, Clone)] +enum FetchFromRow { + Idx(usize), + Default(datatypes::value::Value), } -impl Default for FlowTickManager { - fn default() -> Self { - Self::new() - } -} - -impl FlowTickManager { - pub fn new() -> Self { - FlowTickManager { - start: Instant::now(), - start_timestamp: SystemTime::now() - .duration_since(SystemTime::UNIX_EPOCH) - .unwrap() - .as_millis() as repr::Timestamp, +impl FetchFromRow { + fn fetch(&self, row: &Row) -> datatypes::value::Value { + match self { + Self::Idx(idx) => row + .get(*idx) + .cloned() + .unwrap_or(datatypes::value::Value::Null), + Self::Default(value) => value.clone(), } } - - /// Return the current timestamp in milliseconds - /// - /// TODO(discord9): reconsider since `tick()` require a monotonic clock and also need to survive recover later - pub fn tick(&self) -> repr::Timestamp { - let current = Instant::now(); - let since_the_epoch = current - self.start; - since_the_epoch.as_millis() as repr::Timestamp + self.start_timestamp - } } diff --git a/src/flow/src/adapter/flownode_impl.rs b/src/flow/src/adapter/flownode_impl.rs index 18e7119c354..62e8597dbb6 100644 --- a/src/flow/src/adapter/flownode_impl.rs +++ b/src/flow/src/adapter/flownode_impl.rs @@ -33,13 +33,12 @@ use common_meta::key::flow::FlowMetadataManager; use common_meta::key::flow::flow_info::FlowScheduleConfig; use common_meta::key::flow::flow_state::FlowStat; use common_runtime::JoinHandle; -use common_telemetry::{error, info, trace, warn}; -use datatypes::value::Value; +use common_telemetry::{error, info, warn}; use futures::TryStreamExt; use itertools::Itertools; use operator::utils::try_to_session_query_context; use session::context::QueryContextBuilder; -use snafu::{IntoError, OptionExt, ResultExt, ensure}; +use snafu::{OptionExt, ResultExt, ensure}; use store_api::storage::{RegionId, TableId}; use tokio::sync::{Mutex, RwLock}; @@ -48,11 +47,10 @@ use crate::batching_mode::engine::BatchingEngine; use crate::engine::{FlowEngine, FlowStatProvider}; use crate::error::{ CreateFlowSnafu, ExternalSnafu, FlowNotFoundSnafu, FlowNotRecoveredSnafu, - IllegalCheckTaskStateSnafu, InsertIntoFlowSnafu, InternalSnafu, JoinTaskSnafu, ListFlowsSnafu, - NoAvailableFrontendSnafu, SyncCheckTaskSnafu, UnexpectedSnafu, UnsupportedSnafu, + IllegalCheckTaskStateSnafu, InternalSnafu, JoinTaskSnafu, ListFlowsSnafu, + NoAvailableFrontendSnafu, SyncCheckTaskSnafu, UnsupportedSnafu, }; use crate::metrics::{METRIC_FLOW_ROWS, METRIC_FLOW_TASK_COUNT}; -use crate::repr::{self, DiffRow}; use crate::utils::StateReportHandler; use crate::{Error, FlowId}; @@ -165,23 +163,7 @@ impl FlowDualEngine { } pub async fn gen_state_report(&self) -> FlowStat { - let streaming = self.streaming_engine.flow_stat().await; - let batching = self.batching_engine.flow_stat().await; - - let mut state_size = streaming.state_size; - state_size.extend(batching.state_size); - - let mut last_exec_time_map = streaming.last_exec_time_map; - last_exec_time_map.extend(batching.last_exec_time_map); - - let mut start_time_map = streaming.start_time_map; - start_time_map.extend(batching.start_time_map); - - FlowStat { - state_size, - last_exec_time_map, - start_time_map, - } + self.batching_engine.flow_stat().await } /// Start state report task, which receives a sender from heartbeat task and sends report back. @@ -702,6 +684,11 @@ impl FlowEngine for FlowDualEngine { async fn remove_flow(&self, flow_id: FlowId) -> Result<(), Error> { let flow_type = self.src_table2flow.read().await.get_flow_type(flow_id); + let flow_type = match flow_type { + Some(flow_type) => Some(flow_type), + None if self.streaming_engine.flow_exist(flow_id).await? => Some(FlowType::Streaming), + None => None, + }; match flow_type { Some(FlowType::Batching) => self.batching_engine.remove_flow(flow_id).await, @@ -742,11 +729,12 @@ impl FlowEngine for FlowDualEngine { async fn flow_exist(&self, flow_id: FlowId) -> Result { let flow_type = self.src_table2flow.read().await.get_flow_type(flow_id); - // not using `flow_type.is_some()` to make sure the flow is actually exist in the underlying engine + // Check the underlying registry so stateless flows survive a missing + // routing-map entry during recovery. match flow_type { Some(FlowType::Batching) => self.batching_engine.flow_exist(flow_id).await, Some(FlowType::Streaming) => self.streaming_engine.flow_exist(flow_id).await, - None => Ok(false), + None => self.streaming_engine.flow_exist(flow_id).await, } } @@ -990,13 +978,7 @@ impl FlowEngine for StreamingEngine { } async fn list_flows(&self) -> Result, Error> { - Ok(self - .flow_err_collectors - .read() - .await - .keys() - .cloned() - .collect::>()) + Ok(self.stateless_flow_ids().await) } async fn handle_flow_inserts( @@ -1017,151 +999,6 @@ impl FlowEngine for StreamingEngine { } } -/// Simple helper enum for fetching value from row with default value -#[derive(Debug, Clone)] -enum FetchFromRow { - Idx(usize), - Default(Value), -} - -impl FetchFromRow { - /// Panic if idx is out of bound - fn fetch(&self, row: &repr::Row) -> Value { - match self { - FetchFromRow::Idx(idx) => row.get(*idx).unwrap().clone(), - FetchFromRow::Default(v) => v.clone(), - } - } -} - -impl StreamingEngine { - async fn handle_inserts_inner( - &self, - request: InsertRequests, - ) -> std::result::Result<(), Error> { - // using try_read to ensure two things: - // 1. flush wouldn't happen until inserts before it is inserted - // 2. inserts happening concurrently with flush wouldn't be block by flush - let _flush_lock = self.flush_lock.try_read(); - for write_request in request.requests { - let region_id = write_request.region_id; - let table_id = RegionId::from(region_id).table_id(); - - let (insert_schema, rows_proto) = write_request - .rows - .map(|r| (r.schema, r.rows)) - .unwrap_or_default(); - - // TODO(discord9): reconsider time assignment mechanism - let now = self.tick_manager.tick(); - - let (table_types, fetch_order) = { - let ctx = self.node_context.read().await; - - // TODO(discord9): also check schema version so that altered table can be reported - let table_schema = ctx.table_source.table_from_id(&table_id).await?; - let default_vals = table_schema - .default_values - .iter() - .zip(table_schema.relation_desc.typ().column_types.iter()) - .map(|(v, ty)| { - v.as_ref().and_then(|v| { - match v.create_default(ty.scalar_type(), ty.nullable()) { - Ok(v) => Some(v), - Err(err) => { - common_telemetry::error!(err; "Failed to create default value"); - None - } - } - }) - }) - .collect_vec(); - - let table_types = table_schema - .relation_desc - .typ() - .column_types - .clone() - .into_iter() - .map(|t| t.scalar_type) - .collect_vec(); - let table_col_names = table_schema.relation_desc.names; - let table_col_names = table_col_names - .iter().enumerate() - .map(|(idx,name)| match name { - Some(name) => Ok(name.clone()), - None => InternalSnafu { - reason: format!("Expect column {idx} of table id={table_id} to have name in table schema, found None"), - } - .fail(), - }) - .collect::, _>>()?; - let name_to_col = HashMap::<_, _>::from_iter( - insert_schema - .iter() - .enumerate() - .map(|(i, name)| (&name.column_name, i)), - ); - - let fetch_order: Vec = table_col_names - .iter() - .zip(default_vals) - .map(|(col_name, col_default_val)| { - name_to_col - .get(col_name) - .copied() - .map(FetchFromRow::Idx) - .or_else(|| col_default_val.clone().map(FetchFromRow::Default)) - .with_context(|| UnexpectedSnafu { - reason: format!( - "Column not found: {}, default_value: {:?}", - col_name, col_default_val - ), - }) - }) - .try_collect()?; - - trace!("Reordering columns: {:?}", fetch_order); - (table_types, fetch_order) - }; - - // TODO(discord9): use column instead of row - let rows: Vec = rows_proto - .into_iter() - .map(|r| { - let r = repr::Row::from(r); - let reordered = fetch_order.iter().map(|i| i.fetch(&r)).collect_vec(); - repr::Row::new(reordered) - }) - .map(|r| (r, now, 1)) - .collect_vec(); - if let Err(err) = self - .handle_write_request(region_id.into(), rows, &table_types) - .await - { - let err = BoxedError::new(err); - let flow_ids = self - .node_context - .read() - .await - .get_flow_ids(table_id) - .into_iter() - .flatten() - .cloned() - .collect_vec(); - let err = InsertIntoFlowSnafu { - region_id, - flow_ids, - } - .into_error(err); - common_telemetry::error!(err; "Failed to handle write request"); - return Err(err); - } - } - Ok(()) - } -} - #[cfg(test)] mod tests { use std::collections::HashMap; diff --git a/src/flow/src/adapter/node_context.rs b/src/flow/src/adapter/node_context.rs deleted file mode 100644 index bcddcbb891d..00000000000 --- a/src/flow/src/adapter/node_context.rs +++ /dev/null @@ -1,458 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! Node context, prone to change with every incoming requests - -use std::collections::{BTreeMap, BTreeSet, HashMap}; -use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; - -use common_recordbatch::RecordBatch; -use common_telemetry::trace; -use datatypes::prelude::ConcreteDataType; -use session::context::QueryContext; -use snafu::{OptionExt, ResultExt}; -use table::metadata::TableId; -use tokio::sync::{RwLock, broadcast, mpsc}; - -use crate::adapter::table_source::FlowTableSource; -use crate::adapter::{FlowId, ManagedTableSource, TableName}; -use crate::error::{Error, EvalSnafu, TableNotFoundSnafu}; -use crate::expr::error::InternalSnafu; -use crate::expr::{Batch, GlobalId}; -use crate::metrics::METRIC_FLOW_INPUT_BUF_SIZE; -use crate::plan::TypedPlan; -use crate::repr::{BATCH_SIZE, BROADCAST_CAP, DiffRow, RelationDesc, SEND_BUF_CAP}; - -/// A context that holds the information of the dataflow -#[derive(Debug)] -pub struct FlownodeContext { - /// mapping from source table to tasks, useful for schedule which task to run when a source table is updated - pub source_to_tasks: BTreeMap>, - /// mapping from task to sink table, useful for sending data back to the client when a task is done running - pub flow_to_sink: BTreeMap, - pub flow_plans: BTreeMap, - pub sink_to_flow: BTreeMap, - /// broadcast sender for source table, any incoming write request will be sent to the source table's corresponding sender - /// - /// Note that we are getting insert requests with table id, so we should use table id as the key - pub source_sender: BTreeMap, - /// broadcast receiver for sink table, there should only be one receiver, and it will receive all the data from the sink table - /// - /// and send it back to the client, since we are mocking the sink table as a client, we should use table name as the key - /// note that the sink receiver should only have one, and we are using broadcast as mpsc channel here - pub sink_receiver: - BTreeMap, mpsc::UnboundedReceiver)>, - /// can query the schema of the table source, from metasrv with local cache - pub table_source: Box, - /// All the tables that have been registered in the worker - pub table_repr: IdToNameMap, - pub query_context: Option>, -} - -impl FlownodeContext { - pub fn new(table_source: Box) -> Self { - Self { - source_to_tasks: Default::default(), - flow_to_sink: Default::default(), - flow_plans: Default::default(), - sink_to_flow: Default::default(), - source_sender: Default::default(), - sink_receiver: Default::default(), - table_source, - table_repr: Default::default(), - query_context: Default::default(), - } - } - - pub fn get_flow_ids(&self, table_id: TableId) -> Option<&BTreeSet> { - self.source_to_tasks.get(&table_id) - } -} - -/// a simple broadcast sender with backpressure, bounded capacity and blocking on send when send buf is full -/// note that it wouldn't evict old data, so it's possible to block forever if the receiver is slow -/// -/// receiver still use tokio broadcast channel, since only sender side need to know -/// backpressure and adjust dataflow running duration to avoid blocking -#[derive(Debug)] -pub struct SourceSender { - // TODO(discord9): make it all Vec? - sender: broadcast::Sender, - send_buf_tx: mpsc::Sender, - send_buf_rx: RwLock>, - send_buf_row_cnt: AtomicUsize, -} - -impl Default for SourceSender { - fn default() -> Self { - // TODO(discord9): the capacity is arbitrary, we can adjust it later, might also want to limit the max number of rows in send buf - let (send_buf_tx, send_buf_rx) = mpsc::channel(SEND_BUF_CAP); - Self { - // TODO(discord9): found a better way then increase this to prevent lagging and hence missing input data - sender: broadcast::Sender::new(SEND_BUF_CAP), - send_buf_tx, - send_buf_rx: RwLock::new(send_buf_rx), - send_buf_row_cnt: AtomicUsize::new(0), - } - } -} - -impl SourceSender { - /// max number of iterations to try flush send buf - const MAX_ITERATIONS: usize = 16; - pub fn get_receiver(&self) -> broadcast::Receiver { - self.sender.subscribe() - } - - /// send as many as possible rows from send buf - /// until send buf is empty or broadchannel is full - pub async fn try_flush(&self) -> Result { - let mut row_cnt = 0; - loop { - let mut send_buf = self.send_buf_rx.write().await; - // if inner sender channel is empty or send buf is empty, there - // is nothing to do for now, just break - if self.sender.len() >= BROADCAST_CAP || send_buf.is_empty() { - break; - } - // TODO(discord9): send rows instead so it's just moving a point - if let Some(batch) = send_buf.recv().await { - let len = batch.row_count(); - if let Err(prev_row_cnt) = - self.send_buf_row_cnt - .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |x| x.checked_sub(len)) - { - common_telemetry::error!( - "send buf row count underflow, prev = {}, len = {}", - prev_row_cnt, - len - ); - } - row_cnt += len; - self.sender - .send(batch) - .map_err(|err| { - InternalSnafu { - reason: format!("Failed to send row, error = {:?}", err), - } - .build() - }) - .with_context(|_| EvalSnafu)?; - } - } - if row_cnt > 0 { - trace!("Source Flushed {} rows", row_cnt); - METRIC_FLOW_INPUT_BUF_SIZE.sub(row_cnt as _); - trace!( - "Remaining Source Send buf.len() = {}", - METRIC_FLOW_INPUT_BUF_SIZE.get() - ); - } - - Ok(row_cnt) - } - - /// return number of rows it actual send(including what's in the buffer) - pub async fn send_rows( - &self, - rows: Vec, - batch_datatypes: &[ConcreteDataType], - ) -> Result { - METRIC_FLOW_INPUT_BUF_SIZE.add(rows.len() as _); - // important for backpressure. if send buf is full, block until it's not - while self.send_buf_row_cnt.load(Ordering::SeqCst) >= BATCH_SIZE * 4 { - tokio::task::yield_now().await; - } - - // row count metrics is approx so relaxed order is ok - let batch = Batch::try_from_rows_with_types( - rows.into_iter().map(|(row, _, _)| row).collect(), - batch_datatypes, - ) - .context(EvalSnafu)?; - common_telemetry::trace!("Send one batch to worker with {} rows", batch.row_count()); - - self.send_buf_row_cnt - .fetch_add(batch.row_count(), Ordering::SeqCst); - self.send_buf_tx.send(batch).await.map_err(|e| { - crate::error::InternalSnafu { - reason: format!("Failed to send row, error = {:?}", e), - } - .build() - })?; - - Ok(0) - } - - /// send record batch - pub async fn send_record_batch(&self, batch: RecordBatch) -> Result { - let row_cnt = batch.num_rows(); - let batch = Batch::try_from(batch)?; - - self.send_buf_row_cnt.fetch_add(row_cnt, Ordering::SeqCst); - - self.send_buf_tx.send(batch).await.map_err(|e| { - crate::error::InternalSnafu { - reason: format!("Failed to send batch, error = {:?}", e), - } - .build() - })?; - Ok(row_cnt) - } -} - -impl FlownodeContext { - /// return number of rows it actual send(including what's in the buffer) - /// - /// TODO(discord9): make this concurrent - pub async fn send( - &self, - table_id: TableId, - rows: Vec, - batch_datatypes: &[ConcreteDataType], - ) -> Result { - let sender = self - .source_sender - .get(&table_id) - .with_context(|| TableNotFoundSnafu { - name: table_id.to_string(), - })?; - sender.send_rows(rows, batch_datatypes).await - } - - pub async fn send_rb(&self, table_id: TableId, batch: RecordBatch) -> Result { - let sender = self - .source_sender - .get(&table_id) - .with_context(|| TableNotFoundSnafu { - name: table_id.to_string(), - })?; - sender.send_record_batch(batch).await - } - - /// flush all sender's buf - /// - /// return numbers being sent - pub async fn flush_all_sender(&self) -> Result { - let mut sum = 0; - for sender in self.source_sender.values() { - sender.try_flush().await.inspect(|x| sum += x)?; - } - Ok(sum) - } -} - -impl FlownodeContext { - /// mapping source table to task, and sink table to task in worker context - /// - /// also add their corresponding broadcast sender/receiver - pub fn register_task_src_sink( - &mut self, - task_id: FlowId, - source_table_ids: &[TableId], - sink_table_name: TableName, - ) { - for source_table_id in source_table_ids { - self.add_source_sender_if_not_exist(*source_table_id); - self.source_to_tasks - .entry(*source_table_id) - .or_default() - .insert(task_id); - } - - self.add_sink_receiver(sink_table_name.clone()); - self.flow_to_sink.insert(task_id, sink_table_name.clone()); - self.sink_to_flow.insert(sink_table_name, task_id); - } - - /// add flow plan to worker context - pub fn add_flow_plan(&mut self, task_id: FlowId, plan: TypedPlan) { - self.flow_plans.insert(task_id, plan); - } - - pub fn get_flow_plan(&self, task_id: &FlowId) -> Option { - self.flow_plans.get(task_id).cloned() - } - - /// remove flow from worker context - pub fn remove_flow(&mut self, task_id: FlowId) { - if let Some(sink_table_name) = self.flow_to_sink.remove(&task_id) { - self.sink_to_flow.remove(&sink_table_name); - } - for (source_table_id, tasks) in self.source_to_tasks.iter_mut() { - tasks.remove(&task_id); - if tasks.is_empty() { - self.source_sender.remove(source_table_id); - } - } - self.flow_plans.remove(&task_id); - } - - /// try add source sender, if already exist, do nothing - pub fn add_source_sender_if_not_exist(&mut self, table_id: TableId) { - let _sender = self.source_sender.entry(table_id).or_default(); - } - - pub fn add_sink_receiver(&mut self, table_name: TableName) { - self.sink_receiver - .entry(table_name) - .or_insert_with(mpsc::unbounded_channel); - } - - pub fn get_source_by_global_id(&self, id: &GlobalId) -> Result<&SourceSender, Error> { - let table_id = self - .table_repr - .get_by_global_id(id) - .with_context(|| TableNotFoundSnafu { - name: format!("Global Id = {:?}", id), - })? - .1 - .with_context(|| TableNotFoundSnafu { - name: format!("Table Id = {:?}", id), - })?; - self.source_sender - .get(&table_id) - .with_context(|| TableNotFoundSnafu { - name: table_id.to_string(), - }) - } - - pub fn get_sink_by_global_id( - &self, - id: &GlobalId, - ) -> Result, Error> { - let table_name = self - .table_repr - .get_by_global_id(id) - .with_context(|| TableNotFoundSnafu { - name: format!("{:?}", id), - })? - .0 - .with_context(|| TableNotFoundSnafu { - name: format!("Global Id = {:?}", id), - })?; - self.sink_receiver - .get(&table_name) - .map(|(s, _r)| s.clone()) - .with_context(|| TableNotFoundSnafu { - name: table_name.join("."), - }) - } -} - -impl FlownodeContext { - /// Retrieves a GlobalId and table schema representing a table previously registered by calling the [register_table] function. - /// - /// Returns an error if no table has been registered with the provided names - pub async fn table(&self, name: &TableName) -> Result<(GlobalId, RelationDesc), Error> { - let id = self - .table_repr - .get_by_name(name) - .map(|(_tid, gid)| gid) - .with_context(|| TableNotFoundSnafu { - name: name.join("."), - })?; - let schema = self.table_source.table(name).await?; - Ok((id, schema.relation_desc)) - } - - /// Assign a global id to a table, if already assigned, return the existing global id - /// - /// require at least one of `table_name` or `table_id` to be `Some` - /// - /// and will try to fetch the schema from table info manager(if table exist now) - /// - /// NOTE: this will not actually render the table into collection referred as GlobalId - /// merely creating a mapping from table id to global id - pub async fn assign_global_id_to_table( - &mut self, - srv_map: &ManagedTableSource, - mut table_name: Option, - table_id: Option, - ) -> Result { - // if we can find by table name/id. not assign it - if let Some(gid) = table_name - .as_ref() - .and_then(|table_name| self.table_repr.get_by_name(table_name)) - .map(|(_, gid)| gid) - .or_else(|| { - table_id - .and_then(|id| self.table_repr.get_by_table_id(&id)) - .map(|(_, gid)| gid) - }) - { - Ok(gid) - } else { - let global_id = self.new_global_id(); - - // table id is Some meaning db must have created the table - if let Some(table_id) = table_id { - let known_table_name = srv_map.get_table_name(&table_id).await?; - table_name = table_name.or(Some(known_table_name)); - } // if we don't have table id, it means database haven't assign one yet or we don't need it - - // still update the mapping with new global id - self.table_repr.insert(table_name, table_id, global_id); - Ok(global_id) - } - } - - /// Get a new global id - pub fn new_global_id(&self) -> GlobalId { - GlobalId::User(self.table_repr.global_id_to_name_id.len() as u64) - } -} - -/// A tri-directional map that maps table name, table id, and global id -#[derive(Default, Debug)] -pub struct IdToNameMap { - name_to_global_id: HashMap, - id_to_global_id: HashMap, - global_id_to_name_id: BTreeMap, Option)>, -} - -impl IdToNameMap { - pub fn new() -> Self { - Default::default() - } - - pub fn insert(&mut self, name: Option, id: Option, global_id: GlobalId) { - name.clone() - .and_then(|name| self.name_to_global_id.insert(name.clone(), global_id)); - id.and_then(|id| self.id_to_global_id.insert(id, global_id)); - self.global_id_to_name_id.insert(global_id, (name, id)); - } - - pub fn get_by_name(&self, name: &TableName) -> Option<(Option, GlobalId)> { - self.name_to_global_id.get(name).map(|global_id| { - let (_name, id) = self.global_id_to_name_id.get(global_id).unwrap(); - (*id, *global_id) - }) - } - - pub fn get_by_table_id(&self, id: &TableId) -> Option<(Option, GlobalId)> { - self.id_to_global_id.get(id).map(|global_id| { - let (name, _id) = self.global_id_to_name_id.get(global_id).unwrap(); - (name.clone(), *global_id) - }) - } - - pub fn get_by_global_id( - &self, - global_id: &GlobalId, - ) -> Option<(Option, Option)> { - self.global_id_to_name_id.get(global_id).cloned() - } -} diff --git a/src/flow/src/adapter/parse_expr.rs b/src/flow/src/adapter/parse_expr.rs deleted file mode 100644 index 3f84eed810b..00000000000 --- a/src/flow/src/adapter/parse_expr.rs +++ /dev/null @@ -1,245 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! parse expr like "ts <= now() - interval '5 m'" - -use nom::IResult; -use nom::branch::alt; -use nom::bytes::complete::{tag, tag_no_case}; -use nom::character::complete::{alphanumeric1, digit0, multispace0}; -use nom::combinator::peek; -use nom::sequence::tuple; - -use crate::repr; - -#[test] -fn test_parse_duration() { - let input = "1 h 5 m 42 second"; - let (remain, ttl) = parse_duration(input).unwrap(); - assert_eq!(remain, ""); - assert_eq!(ttl, (3600 + 5 * 60 + 42) * 1000); -} - -#[test] -fn test_parse_fixed() { - let input = "timestamp < now() - INTERVAL '5m 42s'"; - let (remain, ttl) = parse_fixed(input).unwrap(); - assert_eq!(remain, ""); - assert_eq!(ttl, (5 * 60 + 42) * 1000); -} - -pub fn parse_fixed(input: &str) -> IResult<&str, i64> { - let (r, _) = tuple(( - multispace0, - tag_no_case("timestamp"), - multispace0, - tag("<"), - multispace0, - tag_no_case("now()"), - multispace0, - tag("-"), - multispace0, - tag_no_case("interval"), - multispace0, - ))(input)?; - tuple((tag("'"), parse_duration, tag("'")))(r).map(|(r, (_, ttl, _))| (r, ttl)) -} - -/// parse duration and return ttl, currently only support time part of psql interval type -pub fn parse_duration(input: &str) -> IResult<&str, i64> { - let mut intervals = vec![]; - let mut remain = input; - while peek(parse_quality)(remain).is_ok() { - let (r, number) = parse_quality(remain)?; - let (r, unit) = parse_time_unit(r)?; - intervals.push((number, unit)); - remain = r; - } - let mut total = 0; - for (number, unit) in intervals { - let number = match unit { - TimeUnit::Second => number, - TimeUnit::Minute => number * 60, - TimeUnit::Hour => number * 60 * 60, - }; - total += number; - } - total *= 1000; - Ok((remain, total)) -} - -enum Expr { - Col(String), - Now, - Duration(repr::Duration), - Binary { - left: Box, - op: String, - right: Box, - }, -} - -fn parse_expr(input: &str) -> IResult<&str, Expr> { - parse_expr_bp(input, 0) -} - -/// a simple pratt parser -fn parse_expr_bp(input: &str, min_bp: u8) -> IResult<&str, Expr> { - let (mut input, mut lhs): (&str, Expr) = parse_item(input)?; - loop { - let (r, op) = parse_op(input)?; - let (_, (l_bp, r_bp)) = infix_binding_power(op)?; - if l_bp < min_bp { - return Ok((input, lhs)); - } - let (r, rhs) = parse_expr_bp(r, r_bp)?; - input = r; - lhs = Expr::Binary { - left: Box::new(lhs), - op: op.to_string(), - right: Box::new(rhs), - }; - } -} - -fn parse_op(input: &str) -> IResult<&str, &str> { - alt((parse_add_sub, parse_cmp))(input) -} - -fn parse_item(input: &str) -> IResult<&str, Expr> { - if let Ok((r, name)) = parse_col_name(input) { - Ok((r, Expr::Col(name.to_string()))) - } else if let Ok((r, _now)) = parse_now(input) { - Ok((r, Expr::Now)) - } else if let Ok((_r, _num)) = parse_quality(input) { - todo!() - } else { - todo!() - } -} - -fn infix_binding_power(op: &str) -> IResult<&str, (u8, u8)> { - let ret = match op { - "<" | ">" | "<=" | ">=" => (1, 2), - "+" | "-" => (3, 4), - _ => { - return Err(nom::Err::Error(nom::error::Error::new( - op, - nom::error::ErrorKind::Fail, - ))); - } - }; - Ok((op, ret)) -} - -fn parse_col_name(input: &str) -> IResult<&str, &str> { - tuple((multispace0, alphanumeric1, multispace0))(input).map(|(r, (_, name, _))| (r, name)) -} - -fn parse_now(input: &str) -> IResult<&str, &str> { - tag_no_case("now()")(input) -} - -fn parse_add_sub(input: &str) -> IResult<&str, &str> { - tuple((multispace0, alt((tag("+"), tag("-"))), multispace0))(input) - .map(|(r, (_, op, _))| (r, op)) -} - -fn parse_cmp(input: &str) -> IResult<&str, &str> { - tuple(( - multispace0, - alt((tag("<="), tag(">="), tag("<"), tag(">"))), - multispace0, - ))(input) - .map(|(r, (_, op, _))| (r, op)) -} - -/// parse a number with optional sign -fn parse_quality(input: &str) -> IResult<&str, repr::Duration> { - tuple(( - multispace0, - alt((tag("+"), tag("-"), tag(""))), - digit0, - multispace0, - ))(input) - .map(|(r, (_, sign, name, _))| (r, sign, name)) - .and_then(|(r, sign, name)| { - let num = name.parse::().map_err(|_| { - nom::Err::Error(nom::error::Error::new(input, nom::error::ErrorKind::Digit)) - })?; - let num = match sign { - "+" => num, - "-" => -num, - _ => num, - }; - Ok((r, num)) - }) -} - -#[derive(Debug, Clone)] -enum TimeUnit { - Second, - Minute, - Hour, -} - -#[derive(Debug, Clone)] -enum DateUnit { - Day, - Month, - Year, -} - -fn parse_time_unit(input: &str) -> IResult<&str, TimeUnit> { - fn to_second(input: &str) -> IResult<&str, TimeUnit> { - alt(( - tag_no_case("second"), - tag_no_case("seconds"), - tag_no_case("S"), - ))(input) - .map(move |(r, _)| (r, TimeUnit::Second)) - } - fn to_minute(input: &str) -> IResult<&str, TimeUnit> { - alt(( - tag_no_case("minute"), - tag_no_case("minutes"), - tag_no_case("m"), - ))(input) - .map(move |(r, _)| (r, TimeUnit::Minute)) - } - fn to_hour(input: &str) -> IResult<&str, TimeUnit> { - alt((tag_no_case("hour"), tag_no_case("hours"), tag_no_case("h")))(input) - .map(move |(r, _)| (r, TimeUnit::Hour)) - } - - tuple(( - multispace0, - alt(( - to_second, to_minute, - to_hour, /* - tag_no_case("day"), - tag_no_case("days"), - tag_no_case("d"), - tag_no_case("month"), - tag_no_case("months"), - tag_no_case("m"), - tag_no_case("year"), - tag_no_case("years"), - tag_no_case("y"), - */ - )), - multispace0, - ))(input) - .map(|(r, (_, unit, _))| (r, unit)) -} diff --git a/src/flow/src/adapter/refill.rs b/src/flow/src/adapter/refill.rs deleted file mode 100644 index 82caf650db3..00000000000 --- a/src/flow/src/adapter/refill.rs +++ /dev/null @@ -1,440 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! This module contains the refill flow task, which is used to refill flow with given table id and a time range. - -use std::collections::BTreeSet; -use std::sync::Arc; - -use catalog::CatalogManagerRef; -use client::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; -use common_error::ext::BoxedError; -use common_meta::key::flow::FlowMetadataManagerRef; -use common_recordbatch::{RecordBatch, RecordBatches, SendableRecordBatchStream}; -use common_runtime::JoinHandle; -use common_telemetry::error; -use datatypes::value::Value; -use futures::StreamExt; -use query::parser::QueryLanguageParser; -use session::context::QueryContextBuilder; -use snafu::{OptionExt, ResultExt, ensure}; -use table::metadata::TableId; - -use crate::adapter::table_source::ManagedTableSource; -use crate::adapter::{FlowId, FlowStreamingEngineRef, StreamingEngine}; -use crate::error::{FlowNotFoundSnafu, JoinTaskSnafu, UnexpectedSnafu}; -use crate::expr::error::ExternalSnafu; -use crate::expr::utils::find_plan_time_window_expr_lower_bound; -use crate::repr::RelationDesc; -use crate::server::get_all_flow_ids; -use crate::{Error, FrontendInvoker}; - -impl StreamingEngine { - /// Create and start refill flow tasks in background - pub async fn create_and_start_refill_flow_tasks( - self: &FlowStreamingEngineRef, - flow_metadata_manager: &FlowMetadataManagerRef, - catalog_manager: &CatalogManagerRef, - ) -> Result<(), Error> { - let tasks = self - .create_refill_flow_tasks(flow_metadata_manager, catalog_manager) - .await?; - self.starting_refill_flows(tasks).await?; - Ok(()) - } - - /// Create a series of tasks to refill flow - pub async fn create_refill_flow_tasks( - &self, - flow_metadata_manager: &FlowMetadataManagerRef, - catalog_manager: &CatalogManagerRef, - ) -> Result, Error> { - let nodeid = self.node_id.map(|c| c as u64); - - let flow_ids = get_all_flow_ids(flow_metadata_manager, catalog_manager, nodeid).await?; - let mut refill_tasks = Vec::new(); - 'flow_id_loop: for flow_id in flow_ids { - let info = flow_metadata_manager - .flow_info_manager() - .get(flow_id) - .await - .map_err(BoxedError::new) - .context(ExternalSnafu)? - .context(FlowNotFoundSnafu { id: flow_id })?; - - // TODO(discord9): also check flow is already running - for src_table in info.source_table_ids() { - // check if source table still exists - if !self.table_info_source.check_table_exist(src_table).await? { - error!( - "Source table id = {:?} not found while refill flow_id={}, consider re-create the flow if necessary", - src_table, flow_id - ); - continue 'flow_id_loop; - } - } - - let expire_after = info - .expire_after() - .map(super::expire_after_secs_to_millis) - .transpose()?; - // TODO(discord9): better way to get last point - let now = self.tick_manager.tick(); - let plan = self - .node_context - .read() - .await - .get_flow_plan(&FlowId::from(flow_id)) - .context(FlowNotFoundSnafu { id: flow_id })?; - let time_range = if let Some(expire_after) = expire_after { - let low_bound = common_time::Timestamp::new_millisecond(now - expire_after); - let real_low_bound = find_plan_time_window_expr_lower_bound(&plan, low_bound)?; - real_low_bound.map(|l| (l, common_time::Timestamp::new_millisecond(now))) - } else { - None - }; - - common_telemetry::debug!( - "Time range for refill flow_id={} is {:?}", - flow_id, - time_range - ); - - for src_table in info.source_table_ids() { - let time_index_col = self - .table_info_source - .get_time_index_column_from_table_id(*src_table) - .await? - .1; - let time_index_name = time_index_col.name; - let task = RefillTask::create( - flow_id as u64, - *src_table, - time_range, - &time_index_name, - &self.table_info_source, - ) - .await?; - refill_tasks.push(task); - } - } - Ok(refill_tasks) - } - - /// Starting to refill flows, if any error occurs, will rebuild the flow and retry - pub(crate) async fn starting_refill_flows( - self: &FlowStreamingEngineRef, - tasks: Vec, - ) -> Result<(), Error> { - // TODO(discord9): add a back pressure mechanism - let frontend_invoker = - self.frontend_invoker - .read() - .await - .clone() - .context(UnexpectedSnafu { - reason: "frontend invoker is not set", - })?; - - for mut task in tasks { - task.start_running(self.clone(), &frontend_invoker).await?; - // TODO(discord9): save refill tasks to a map and check if it's finished when necessary - // i.e. when system table need query it's state - self.refill_tasks - .write() - .await - .insert(task.data.flow_id, task); - } - Ok(()) - } -} - -/// Task to refill flow with given table id and a time range -pub struct RefillTask { - data: TaskData, - state: TaskState<()>, -} - -#[derive(Clone)] -struct TaskData { - flow_id: FlowId, - table_id: TableId, - table_schema: RelationDesc, -} - -impl TaskData { - /// validate that incoming batch's schema is the same as table schema(by comparing types&names) - fn validate_schema(table_schema: &RelationDesc, rb: &RecordBatch) -> Result<(), Error> { - let rb_schema = &rb.schema; - ensure!( - rb_schema.column_schemas().len() == table_schema.len()?, - UnexpectedSnafu { - reason: format!( - "RecordBatch schema length does not match table schema length, {}!={}", - rb_schema.column_schemas().len(), - table_schema.len()? - ) - } - ); - for (i, rb_col) in rb_schema.column_schemas().iter().enumerate() { - let (rb_name, rb_ty) = (rb_col.name.as_str(), &rb_col.data_type); - let (table_name, table_ty) = ( - table_schema.names[i].as_ref(), - &table_schema.typ().column_types[i].scalar_type, - ); - ensure!( - Some(rb_name) == table_name.map(|c| c.as_str()), - UnexpectedSnafu { - reason: format!( - "Mismatch in column names: expected {:?}, found {}", - table_name, rb_name - ) - } - ); - - ensure!( - rb_ty == table_ty, - UnexpectedSnafu { - reason: format!( - "Mismatch in column types for {}: expected {:?}, found {:?}", - rb_name, table_ty, rb_ty - ) - } - ); - } - Ok(()) - } -} - -/// Refill task state -enum TaskState { - /// Task is not started - Prepared { sql: String }, - /// Task is running - Running { - handle: JoinHandle>, - }, - /// Task is finished - Finished { res: Result }, -} - -impl TaskState { - fn new(sql: String) -> Self { - Self::Prepared { sql } - } -} - -mod test_send { - use std::collections::BTreeMap; - - use tokio::sync::RwLock; - - use super::*; - fn is_send() {} - fn foo() { - is_send::>(); - is_send::(); - is_send::>(); - is_send::>>(); - } -} - -impl TaskState<()> { - /// check if task is finished - async fn is_finished(&mut self) -> Result { - match self { - Self::Finished { .. } => Ok(true), - Self::Running { handle } => Ok(if handle.is_finished() { - *self = Self::Finished { - res: handle.await.context(JoinTaskSnafu)?, - }; - true - } else { - false - }), - _ => Ok(false), - } - } - - fn start_running( - &mut self, - task_data: &TaskData, - manager: FlowStreamingEngineRef, - mut output_stream: SendableRecordBatchStream, - ) -> Result<(), Error> { - let data = (*task_data).clone(); - let handle: JoinHandle> = common_runtime::spawn_global(async move { - while let Some(rb) = output_stream.next().await { - let rb = match rb { - Ok(rb) => rb, - Err(err) => Err(BoxedError::new(err)).context(ExternalSnafu)?, - }; - TaskData::validate_schema(&data.table_schema, &rb)?; - - // send rb into flow node - manager - .node_context - .read() - .await - .send_rb(data.table_id, rb) - .await?; - } - common_telemetry::info!( - "Refill successful for source table_id={}, flow_id={}", - data.table_id, - data.flow_id - ); - Ok(()) - }); - *self = Self::Running { handle }; - - Ok(()) - } -} - -/// Query stream of RefillTask, simply wrap RecordBatches and RecordBatchStream and check output is not `AffectedRows` -enum QueryStream { - Batches { batches: RecordBatches }, - Stream { stream: SendableRecordBatchStream }, -} - -impl TryFrom for QueryStream { - type Error = Error; - fn try_from(value: common_query::Output) -> Result { - match value.data { - common_query::OutputData::Stream(stream) => Ok(QueryStream::Stream { stream }), - common_query::OutputData::RecordBatches(batches) => { - Ok(QueryStream::Batches { batches }) - } - _ => UnexpectedSnafu { - reason: format!("Unexpected output data type: {:?}", value.data), - } - .fail(), - } - } -} - -impl QueryStream { - fn try_into_stream(self) -> Result { - match self { - Self::Batches { batches } => Ok(batches.as_stream()), - Self::Stream { stream } => Ok(stream), - } - } -} - -impl RefillTask { - /// Query with "select * from table WHERE time >= range_start and time < range_end" - pub async fn create( - flow_id: FlowId, - table_id: TableId, - time_range: Option<(common_time::Timestamp, common_time::Timestamp)>, - time_col_name: &str, - table_src: &ManagedTableSource, - ) -> Result { - let (table_name, table_schema) = table_src.get_table_name_schema(&table_id).await?; - let all_col_names: BTreeSet<_> = table_schema - .relation_desc - .iter_names() - .flatten() - .map(|s| s.as_str()) - .collect(); - - if !all_col_names.contains(time_col_name) { - UnexpectedSnafu { - reason: format!( - "Can't find column {} in table {} while refill flow", - time_col_name, - table_name.join(".") - ), - } - .fail()?; - } - - let sql = if let Some(time_range) = time_range { - format!( - "select * from {0} where {1} >= {2} and {1} < {3}", - table_name.join("."), - time_col_name, - Value::from(time_range.0), - Value::from(time_range.1), - ) - } else { - format!("select * from {0}", table_name.join(".")) - }; - - Ok(RefillTask { - data: TaskData { - flow_id, - table_id, - table_schema: table_schema.relation_desc, - }, - state: TaskState::new(sql), - }) - } - - /// Start running the task in background, non-blocking - pub async fn start_running( - &mut self, - manager: FlowStreamingEngineRef, - invoker: &FrontendInvoker, - ) -> Result<(), Error> { - let TaskState::Prepared { sql } = &mut self.state else { - UnexpectedSnafu { - reason: "task is not prepared", - } - .fail()? - }; - - // we don't need information from query context in this query so a default query context is enough - let query_ctx = Arc::new( - QueryContextBuilder::default() - .current_catalog(DEFAULT_CATALOG_NAME.to_string()) - .current_schema(DEFAULT_SCHEMA_NAME.to_string()) - .build(), - ); - - let stmt_exec = invoker.statement_executor(); - - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx) - .map_err(BoxedError::new) - .context(ExternalSnafu)?; - let plan = stmt_exec - .plan(&stmt, query_ctx.clone()) - .await - .map_err(BoxedError::new) - .context(ExternalSnafu)?; - - let output_data = stmt_exec - .exec_plan(plan, query_ctx) - .await - .map_err(BoxedError::new) - .context(ExternalSnafu)?; - let output_data = output_data - .map_dictionary_to_values() - .map_err(BoxedError::new) - .context(ExternalSnafu)?; - - let output_stream = QueryStream::try_from(output_data)?; - let output_stream = output_stream.try_into_stream()?; - - self.state - .start_running(&self.data, manager, output_stream)?; - Ok(()) - } - - pub async fn is_finished(&mut self) -> Result { - self.state.is_finished().await - } -} diff --git a/src/flow/src/adapter/stat.rs b/src/flow/src/adapter/stat.rs deleted file mode 100644 index 9521a24c5c9..00000000000 --- a/src/flow/src/adapter/stat.rs +++ /dev/null @@ -1,47 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::collections::BTreeMap; - -use common_meta::key::flow::flow_state::FlowStat; - -use crate::StreamingEngine; -use crate::engine::FlowStatProvider; - -impl FlowStatProvider for StreamingEngine { - async fn flow_stat(&self) -> FlowStat { - let mut state_size_map = BTreeMap::new(); - let mut last_exec_time_map = BTreeMap::new(); - let mut start_time_map = BTreeMap::new(); - - for worker in self.worker_handles.iter() { - match worker.get_full_flow_stat().await { - Ok((sizes, exec_times, start_times)) => { - state_size_map.extend(sizes.into_iter().map(|(k, v)| (k as u32, v))); - last_exec_time_map.extend(exec_times.into_iter().map(|(k, v)| (k as u32, v))); - start_time_map.extend(start_times.into_iter().map(|(k, v)| (k as u32, v))); - } - Err(err) => { - common_telemetry::error!(err; "Get full flow stat error"); - } - } - } - - FlowStat { - state_size: state_size_map, - last_exec_time_map, - start_time_map, - } - } -} diff --git a/src/flow/src/adapter/stateless.rs b/src/flow/src/adapter/stateless.rs new file mode 100644 index 00000000000..1c1090a0ccc --- /dev/null +++ b/src/flow/src/adapter/stateless.rs @@ -0,0 +1,730 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Stateless DataFusion execution for streaming flows. + +use std::collections::HashSet; +use std::sync::Arc; + +use api::helper::{to_grpc_value, vectors_to_rows}; +use api::v1::greptime_request::Request; +use api::v1::{RowInsertRequest, RowInsertRequests, Rows}; +use common_error::ext::BoxedError; +use common_query::OutputData; +use common_recordbatch::{RecordBatch, RecordBatches, map_dictionary_to_values_data_type}; +use common_time::Timestamp; +use datafusion::catalog::MemTable; +use datafusion::datasource::{TableProvider, provider_as_source, source_as_provider}; +use datafusion_common::tree_node::{Transformed, TreeNode}; +use datafusion_common::{Column, DFSchema, TableReference}; +use datafusion_expr::logical_plan::{Distinct, Projection}; +use datafusion_expr::{Expr, LogicalPlan}; +use datatypes::schema::{ColumnSchema, SchemaRef}; +use datatypes::value::Value; +use query::QueryEngine; +use session::context::QueryContextRef; +use snafu::{OptionExt, ResultExt, ensure}; +use table::metadata::TableId; +use table::table::adapter::DfTableProviderAdapter; + +use crate::TableName; +use crate::adapter::util::column_schemas_to_proto; +use crate::batching_mode::frontend_client::FrontendClient; +use crate::error::{DatafusionSnafu, Error, ExternalSnafu, InvalidQuerySnafu, UnexpectedSnafu}; +use crate::repr::DiffRow; + +/// The validated, immutable part of one streaming flow. +#[derive(Clone)] +pub(crate) struct StatelessFlow { + pub(crate) source_table_id: TableId, + pub(crate) source_table_name: TableName, + pub(crate) source_schema: SchemaRef, + pub(crate) source_schema_version: u32, + pub(crate) sink_table_name: TableName, + pub(crate) sink_schema: Vec, + pub(crate) sink_primary_keys: Vec, + /// The exact trailing columns resolved when the flow was created. + pub(crate) auto_columns: Vec, + pub(crate) plan: LogicalPlan, + pub(crate) query_ctx: QueryContextRef, + pub(crate) create_args: crate::CreateFlowArgs, +} + +/// Per-request input provider. It owns no catalog or storage state. +#[cfg(test)] +fn test_source_plan(table_name: TableReference, provider: Arc) -> LogicalPlan { + datafusion_expr::LogicalPlanBuilder::scan(table_name, provider_as_source(provider), None) + .unwrap() + .filter(datafusion_expr::col("number").gt(datafusion_expr::lit(1))) + .unwrap() + .project(vec![datafusion_expr::col("number")]) + .unwrap() + .build() + .unwrap() +} + +fn input_provider(batch: &RecordBatch) -> Result, Error> { + let arrow_batch = batch.df_record_batch().clone(); + let provider = MemTable::try_new(arrow_batch.schema(), vec![vec![arrow_batch]]).context( + DatafusionSnafu { + context: "Failed to create transient flow input provider", + }, + )?; + Ok(Arc::new(provider)) +} + +/// Adds the source timestamp to every supported plan node that has to carry it through a +/// filter or projection. The expression is appended only after the visible expressions, so the +/// sink contract remains positional. +pub(crate) fn rewrite_source_timestamp( + plan: LogicalPlan, + source_name: &TableReference, + source_timestamp_name: &str, +) -> Result { + let mut names = HashSet::new(); + plan.apply(|node| { + names.extend( + node.schema() + .fields() + .iter() + .map(|field| field.name().clone()), + ); + Ok(datafusion_common::tree_node::TreeNodeRecursion::Continue) + }) + .context(DatafusionSnafu { + context: "Failed to inspect streaming flow plan schema", + })?; + let mut hidden_name = "__flow_source_timestamp".to_string(); + let mut suffix = 0; + while !names.insert(hidden_name.clone()) { + suffix += 1; + hidden_name = format!("__flow_source_timestamp_{suffix}"); + } + let visible_count = plan.schema().fields().len(); + let source_timestamp = Column::from_name(source_timestamp_name); + let mut plan = plan + .transform_up_with_subqueries(|node| match node { + LogicalPlan::TableScan(mut scan) => { + if scan.table_name.resolved_eq(source_name) + && let Some(projection) = &mut scan.projection + { + let timestamp_index = scan + .source + .schema() + .index_of(source_timestamp_name) + .map_err(|_| { + datafusion::error::DataFusionError::Plan( + "Source timestamp is absent from source scan".into(), + ) + })?; + if !projection.contains(×tamp_index) { + projection.push(timestamp_index); + let schema = scan.source.schema(); + scan.projected_schema = Arc::new(DFSchema::new_with_metadata( + projection + .iter() + .map(|index| { + ( + Some(scan.table_name.clone()), + Arc::new(schema.field(*index).clone()), + ) + }) + .collect(), + schema.metadata().clone(), + )?); + } + } + Ok(Transformed::yes(LogicalPlan::TableScan(scan))) + } + LogicalPlan::Projection(mut projection) => { + let hidden_expr = if projection + .input + .schema() + .fields() + .iter() + .any(|field| field.name() == &hidden_name) + { + Expr::Column(Column::from_name(hidden_name.clone())) + } else { + Expr::Column(source_timestamp.clone()) + }; + projection.expr.push(hidden_expr.alias(hidden_name.clone())); + let projection = Projection::try_new(projection.expr, projection.input)?; + Ok(Transformed::yes(LogicalPlan::Projection(projection))) + } + _ => Ok(Transformed::no(node)), + }) + .context(DatafusionSnafu { + context: "Failed to add source timestamp to streaming flow plan", + })? + .data; + + // A plan ending at a scan or filter has no projection at which to give the carried value its + // hidden name. Add one only in that case; a projection below a filter already carries it. + if !plan + .schema() + .fields() + .iter() + .any(|field| field.name() == &hidden_name) + { + let expressions = plan + .schema() + .fields() + .iter() + .take(visible_count) + .map(|field| Expr::Column(Column::from_name(field.name()))) + .chain(std::iter::once( + Expr::Column(source_timestamp).alias(hidden_name), + )) + .collect::>(); + plan = Projection::try_new(expressions, Arc::new(plan)) + .map(LogicalPlan::Projection) + .context(DatafusionSnafu { + context: "Failed to finalize source timestamp in streaming flow plan", + })?; + } + Ok(plan) +} + +fn replace_source( + plan: LogicalPlan, + source_name: &TableReference, + provider: Arc, +) -> Result { + let mut scan_count = 0; + let mut replaced_scan_count = 0; + let plan = plan + .transform_up_with_subqueries(|node| match node { + LogicalPlan::TableScan(mut scan) => { + scan_count += 1; + if scan.table_name.resolved_eq(source_name) { + replaced_scan_count += 1; + scan.source = provider_as_source(provider.clone()); + let schema = scan.source.schema(); + scan.projected_schema = if let Some(projection) = &scan.projection { + Arc::new(DFSchema::new_with_metadata( + projection + .iter() + .map(|index| { + ( + Some(scan.table_name.clone()), + Arc::new(schema.field(*index).clone()), + ) + }) + .collect(), + schema.metadata().clone(), + )?) + } else { + Arc::new(DFSchema::try_from_qualified_schema( + scan.table_name.clone(), + &schema, + )?) + }; + Ok(Transformed::yes(LogicalPlan::TableScan(scan))) + } else { + Ok(Transformed::no(LogicalPlan::TableScan(scan))) + } + } + LogicalPlan::Join(_) | LogicalPlan::Aggregate(_) => { + Err(datafusion::error::DataFusionError::Plan( + "Streaming flow supports only a single projection/filter source".into(), + )) + } + _ => Ok(Transformed::no(node)), + }) + .context(DatafusionSnafu { + context: "Failed to substitute transient flow input provider", + })? + .data; + ensure!( + scan_count == 1 && replaced_scan_count == 1, + InvalidQuerySnafu { + reason: format!( + "Expected one source scan matching {:?}, found {scan_count} scans and {replaced_scan_count} matches", + source_name + ) + } + ); + Ok(plan) +} + +/// Validates the deliberately small stateless streaming SQL subset. +pub(crate) fn validate_plan(plan: &LogicalPlan) -> Result<(), Error> { + let mut scans = 0; + plan.apply(|node| { + match node { + LogicalPlan::TableScan(_) => scans += 1, + LogicalPlan::Projection(_) | LogicalPlan::Filter(_) => {} + LogicalPlan::Aggregate(_) => { + return Err(datafusion::error::DataFusionError::Plan( + "Aggregation is unsupported in streaming flows. Recreate the flow to select batching mode. A source table with TTL=instant must use persisted retention first. Aggregation SQL without a time window requires EVAL INTERVAL.".into(), + )); + } + // DISTINCT is evaluated against this request's transient input only. DISTINCT ON + // has ordering/selection semantics beyond the supported stateless subset. + LogicalPlan::Distinct(Distinct::All(_)) => {} + _ => { + return Err(datafusion::error::DataFusionError::Plan( + "Streaming flow supports only projection and filter over one source scan" + .into(), + )); + } + } + Ok(datafusion_common::tree_node::TreeNodeRecursion::Continue) + }) + .context(DatafusionSnafu { + context: "Failed to validate streaming flow plan", + })?; + ensure!( + scans == 1, + InvalidQuerySnafu { + reason: format!("Expected one source scan, found {scans}") + } + ); + Ok(()) +} + +/// Ensures the retained scan was planned against the source metadata captured for this flow. +pub(crate) fn validate_source_scan( + plan: &LogicalPlan, + source_table_id: TableId, + source_schema: &SchemaRef, +) -> Result<(), Error> { + plan.apply(|node| { + if let LogicalPlan::TableScan(scan) = node { + let provider = source_as_provider(&scan.source)?; + let provider = provider + .downcast_ref::() + .ok_or_else(|| { + datafusion::error::DataFusionError::Plan( + "Streaming flow source scan does not use a table provider".into(), + ) + })?; + let table_info = provider.table().table_info(); + if table_info.ident.table_id != source_table_id + || table_info.meta.schema.as_ref() != source_schema.as_ref() + { + return Err(datafusion::error::DataFusionError::Plan(format!( + "Streaming flow source scan does not match source table {source_table_id} schema version {}", + source_schema.version() + ))); + } + } + Ok(datafusion_common::tree_node::TreeNodeRecursion::Continue) + }) + .context(DatafusionSnafu { + context: "Failed to validate streaming flow source scan", + })?; + Ok(()) +} + +/// Rejects execution when the source metadata changed after the flow plan was retained. +/// The outer adapter replans a flow when it observes a new source schema version; this guard +/// rejects a request if the version changes again before execution. +fn validate_source_schema_version( + retained_version: u32, + current_version: u32, +) -> Result<(), Error> { + ensure!( + retained_version == current_version, + InvalidQuerySnafu { + reason: format!( + "Source schema version changed from {retained_version} to {current_version}; recreate or recover the flow before writing" + ) + } + ); + Ok(()) +} + +fn synthesize_auto_values(columns: &[ColumnSchema], now: Timestamp) -> Result, Error> { + columns + .iter() + .map(|column| { + let timestamp_type = column.data_type.as_timestamp().context(InvalidQuerySnafu { + reason: format!("Auto sink column {} is not a timestamp", column.name), + })?; + let value = if column.name == crate::adapter::AUTO_CREATED_UPDATE_AT_TS_COL { + now.convert_to(timestamp_type.unit()) + .context(InvalidQuerySnafu { + reason: "Current timestamp cannot be represented in sink timestamp unit", + })? + } else if column.name == crate::adapter::AUTO_CREATED_PLACEHOLDER_TS_COL { + Timestamp::new(0, timestamp_type.unit()) + } else { + return InvalidQuerySnafu { + reason: format!("Unsupported auto sink column {}", column.name), + } + .fail(); + }; + Ok(Value::Timestamp(value)) + }) + .collect() +} + +/// Executes one mirror write using only the supplied batch and writes its output. +pub(crate) async fn execute( + flow: &StatelessFlow, + rows: &[DiffRow], + batch_datatypes: &[datatypes::data_type::ConcreteDataType], + query_engine: &Arc, + frontend_client: &Arc, + current_source_schema_version: u32, +) -> Result { + validate_source_schema_version(flow.source_schema_version, current_source_schema_version)?; + let values = rows.iter().map(|(row, _, _)| row.clone()).collect(); + let batch = crate::expr::Batch::try_from_rows_with_types(values, batch_datatypes) + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + let batch = RecordBatch::new(flow.source_schema.clone(), batch.batch().to_vec()) + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + let provider = input_provider(&batch)?; + let source_ref = TableReference::full( + flow.source_table_name[0].clone(), + flow.source_table_name[1].clone(), + flow.source_table_name[2].clone(), + ); + let plan = replace_source(flow.plan.clone(), &source_ref, provider)?; + let output = query_engine + .execute(plan, flow.query_ctx.clone()) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + let batches = match output.data { + OutputData::RecordBatches(batches) => batches, + OutputData::Stream(stream) => RecordBatches::try_collect(stream) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)?, + OutputData::AffectedRows(_) => { + return UnexpectedSnafu { + reason: "Streaming flow query returned affected rows", + } + .fail(); + } + }; + + let output_schema = batches + .schema() + .column_schemas() + .iter() + .cloned() + .map(|mut column| { + column.data_type = map_dictionary_to_values_data_type(&column.data_type); + column + }) + .collect::>(); + crate::adapter::validate_sink_layout_with_suffix( + &output_schema, + &flow.sink_schema, + &flow.auto_columns, + )?; + let mut output_rows = Vec::new(); + for batch in batches { + let vectors = datatypes::vectors::Helper::try_into_vectors(batch.columns()) + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + output_rows.extend(vectors_to_rows(vectors.iter(), batch.num_rows())); + } + if output_rows.is_empty() { + return Ok(0); + } + + // Auto columns are deliberately synthesized here rather than in the query plan. This keeps + // one current timestamp for the whole request and preserves the sink's timestamp precision. + let auto_values = synthesize_auto_values(&flow.auto_columns, Timestamp::current_millis())? + .into_iter() + .map(to_grpc_value) + .collect::>(); + for row in &mut output_rows { + row.values.extend(auto_values.iter().cloned()); + } + + let output_row_count = output_rows.len(); + let proto_schema = column_schemas_to_proto(flow.sink_schema.clone(), &flow.sink_primary_keys)?; + let request = Request::RowInserts(RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: flow.sink_table_name[2].clone(), + rows: Some(Rows { + schema: proto_schema, + rows: output_rows, + }), + }], + }); + let mut peer = None; + frontend_client + .handle_insert_once( + request, + &flow.sink_table_name[0], + &flow.sink_table_name[1], + &mut peer, + ) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + Ok(output_row_count) +} + +#[cfg(test)] +mod tests { + use datafusion::catalog::MemTable; + use datafusion::logical_expr::LogicalPlanBuilder; + use datatypes::data_type::ConcreteDataType; + use datatypes::schema::{ColumnSchema, Schema}; + use datatypes::vectors::{Int32Vector, TimestampMillisecondVector}; + use session::context::QueryContext; + + use super::*; + + #[test] + fn validation_accepts_distinct_over_one_source() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "number", + ConcreteDataType::int32_datatype(), + false, + )])); + let plan = LogicalPlanBuilder::scan( + TableReference::bare("source"), + provider_as_source(provider(&schema, 1)), + None, + ) + .unwrap() + .project(vec![datafusion_expr::col("number")]) + .unwrap() + .distinct() + .unwrap() + .build() + .unwrap(); + assert!(validate_plan(&plan).is_ok()); + } + + #[test] + fn validation_rejects_plan_without_source_scan() { + let plan = LogicalPlan::EmptyRelation(datafusion_expr::logical_plan::EmptyRelation { + produce_one_row: false, + schema: Arc::new(DFSchema::empty()), + }); + assert!(validate_plan(&plan).is_err()); + } + + fn provider(schema: &SchemaRef, value: i32) -> Arc { + let batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice([value])) as datatypes::prelude::VectorRef], + ) + .unwrap(); + let arrow = batch.df_record_batch().clone(); + Arc::new(MemTable::try_new(arrow.schema(), vec![vec![arrow]]).unwrap()) + } + + #[test] + fn source_schema_version_must_match_retained_plan() { + assert!(validate_source_schema_version(7, 7).is_ok()); + + let error = validate_source_schema_version(7, 8).unwrap_err(); + assert!(matches!(error, Error::InvalidQuery { reason, .. } if + reason.contains("Source schema version changed from 7 to 8") + && reason.contains("recreate or recover") + )); + } + + #[test] + fn auto_values_use_the_sink_timestamp_units() { + let update_at = ColumnSchema::new( + crate::adapter::AUTO_CREATED_UPDATE_AT_TS_COL, + datatypes::data_type::ConcreteDataType::timestamp_second_datatype(), + true, + ); + let placeholder = ColumnSchema::new( + crate::adapter::AUTO_CREATED_PLACEHOLDER_TS_COL, + datatypes::data_type::ConcreteDataType::timestamp_nanosecond_datatype(), + true, + ); + let values = synthesize_auto_values( + &[update_at], + Timestamp::new(1_234, common_time::timestamp::TimeUnit::Millisecond), + ) + .unwrap(); + assert_eq!(values.len(), 1); + assert_eq!( + values[0].as_timestamp().unwrap().unit(), + common_time::timestamp::TimeUnit::Second + ); + assert!(values[0].as_timestamp().unwrap().value() > 0); + + let values = synthesize_auto_values( + &[ + ColumnSchema::new( + crate::adapter::AUTO_CREATED_UPDATE_AT_TS_COL, + datatypes::data_type::ConcreteDataType::timestamp_microsecond_datatype(), + true, + ), + placeholder, + ], + Timestamp::new(1_234, common_time::timestamp::TimeUnit::Millisecond), + ) + .unwrap(); + assert_eq!(values[1].as_timestamp().unwrap().value(), 0); + assert_eq!( + values[1].as_timestamp().unwrap().unit(), + common_time::timestamp::TimeUnit::Nanosecond + ); + } + + #[test] + fn auto_values_reject_arbitrary_columns() { + let column = ColumnSchema::new( + "other", + datatypes::data_type::ConcreteDataType::timestamp_millisecond_datatype(), + true, + ); + assert!(synthesize_auto_values(&[column], Timestamp::current_millis()).is_err()); + } + + #[test] + fn rewrite_source_timestamp_appends_collision_free_hidden_output() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new("number", ConcreteDataType::int32_datatype(), false), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + "__flow_source_timestamp", + ConcreteDataType::int32_datatype(), + false, + ), + ])); + let batch = RecordBatch::new( + schema.clone(), + vec![ + Arc::new(Int32Vector::from_slice([2])) as datatypes::prelude::VectorRef, + Arc::new(TimestampMillisecondVector::from_slice([42])) + as datatypes::prelude::VectorRef, + Arc::new(Int32Vector::from_slice([7])) as datatypes::prelude::VectorRef, + ], + ) + .unwrap(); + let plan = datafusion_expr::LogicalPlanBuilder::scan( + TableReference::bare("source"), + provider_as_source(input_provider(&batch).unwrap()), + None, + ) + .unwrap() + .filter(datafusion_expr::col("number").gt(datafusion_expr::lit(1))) + .unwrap() + .project(vec![datafusion_expr::col("number")]) + .unwrap() + .build() + .unwrap(); + let rewritten = + rewrite_source_timestamp(plan, &TableReference::bare("source"), "ts").unwrap(); + assert_eq!(rewritten.schema().fields().len(), 2); + assert!( + rewritten + .schema() + .field(1) + .name() + .starts_with("__flow_source_timestamp") + ); + } + + #[tokio::test] + async fn finite_projection_filter_does_not_retain_previous_batch() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "number", + ConcreteDataType::int32_datatype(), + false, + )])); + let initial_provider = provider(&schema, 0); + let plan = test_source_plan(TableReference::bare("source"), initial_provider); + assert!(validate_plan(&plan).is_ok()); + + let engine = crate::test_utils::create_test_query_engine(); + let run = |input: Arc| { + let engine = engine.clone(); + let plan = plan.clone(); + async move { + let plan = replace_source(plan, &TableReference::bare("source"), input).unwrap(); + let output = engine.execute(plan, QueryContext::arc()).await.unwrap(); + match output.data { + OutputData::Stream(stream) => RecordBatches::try_collect(stream) + .await + .unwrap() + .iter() + .map(|batch| batch.num_rows()) + .sum(), + OutputData::RecordBatches(batches) => { + batches.iter().map(|batch| batch.num_rows()).sum() + } + OutputData::AffectedRows(_) => 0, + } + } + }; + + assert_eq!(run(provider(&schema, 1)).await, 0); + assert_eq!(run(provider(&schema, 2)).await, 1); + } + + #[tokio::test] + async fn finite_distinct_collapses_duplicates_per_request() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "number", + ConcreteDataType::int32_datatype(), + false, + )])); + let input = |values: &[i32]| { + let batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice(values)) as datatypes::prelude::VectorRef], + ) + .unwrap(); + input_provider(&batch).unwrap() + }; + let plan = LogicalPlanBuilder::scan( + TableReference::bare("source"), + provider_as_source(input(&[1])), + None, + ) + .unwrap() + .project(vec![datafusion_expr::col("number")]) + .unwrap() + .distinct() + .unwrap() + .build() + .unwrap(); + assert!(validate_plan(&plan).is_ok()); + + let engine = crate::test_utils::create_test_query_engine(); + let run = |provider: Arc| { + let engine = engine.clone(); + let plan = plan.clone(); + async move { + let plan = replace_source(plan, &TableReference::bare("source"), provider).unwrap(); + let output = engine.execute(plan, QueryContext::arc()).await.unwrap(); + let batches = match output.data { + OutputData::Stream(stream) => RecordBatches::try_collect(stream).await.unwrap(), + OutputData::RecordBatches(batches) => batches, + OutputData::AffectedRows(_) => panic!("unexpected affected rows"), + }; + batches.iter().map(|batch| batch.num_rows()).sum::() + } + }; + + assert_eq!(run(input(&[1, 1, 2])).await, 2); + assert_eq!(run(input(&[1, 1])).await, 1); + } +} diff --git a/src/flow/src/adapter/tests.rs b/src/flow/src/adapter/tests.rs index df6a58c46ce..6bc02f60786 100644 --- a/src/flow/src/adapter/tests.rs +++ b/src/flow/src/adapter/tests.rs @@ -15,12 +15,257 @@ //! Mock test for adapter module //! TODO(discord9): write mock test -use datatypes::schema::{ColumnSchema, SchemaBuilder}; -use store_api::storage::ConcreteDataType; +use api::v1::SemanticType; +use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, TimeUnit as ArrowTimeUnit}; +use datafusion::catalog::MemTable; +use datafusion::datasource::provider_as_source; +use datafusion_common::TableReference; +use datafusion_expr::LogicalPlanBuilder; +use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema, Schema, SchemaBuilder}; +use store_api::storage::{ConcreteDataType, TableId}; use table::metadata::{TableInfo, TableInfoBuilder, TableMetaBuilder}; use super::*; +#[test] +fn stateless_output_aliases_are_matched_by_position() { + let output = vec![ColumnSchema::new( + "output_alias", + ConcreteDataType::int32_datatype(), + false, + )]; + let sink = vec![ColumnSchema::new( + "sink_column", + ConcreteDataType::int32_datatype(), + false, + )]; + assert!(validate_sink_layout(&output, &sink).is_ok()); + let proto = crate::adapter::util::column_schemas_to_proto(sink, &[]).unwrap(); + assert_eq!(proto[0].column_name, "sink_column"); +} + +#[test] +fn stateless_sink_schema_has_tag_and_timestamp_semantics() { + let schema = vec![ + ColumnSchema::new("host", ConcreteDataType::string_datatype(), false), + ColumnSchema::new("ts", ConcreteDataType::timestamp_second_datatype(), false) + .with_time_index(true), + ]; + let proto = + crate::adapter::util::column_schemas_to_proto(schema, &["host".to_string()]).unwrap(); + assert_eq!(proto[0].semantic_type, SemanticType::Tag as i32); + assert_eq!(proto[1].semantic_type, SemanticType::Timestamp as i32); +} + +#[test] +fn stateless_resolves_suffix_by_output_arity() { + let ordinary = ColumnSchema::new("value", ConcreteDataType::int32_datatype(), false); + let update_at = ColumnSchema::new( + AUTO_CREATED_UPDATE_AT_TS_COL, + ConcreteDataType::timestamp_second_datatype(), + true, + ); + // Equal arity is an ordinary sink, despite the reserved-looking name. + assert!( + resolve_sink_layout( + &[ordinary.clone(), update_at.clone()], + &[ordinary.clone(), update_at.clone()] + ) + .unwrap() + .is_empty() + ); + assert_eq!( + resolve_sink_layout( + std::slice::from_ref(&ordinary), + &[ordinary.clone(), update_at] + ) + .unwrap() + .len(), + 1 + ); +} + +#[test] +fn stateless_explicit_timestamp_compatibility_requires_default_and_lineage_absence() { + let source = Arc::new(Schema::new(vec![ + ColumnSchema::new("value", ConcreteDataType::int32_datatype(), false), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ])); + let output = vec![ColumnSchema::new( + "value", + ConcreteDataType::int32_datatype(), + false, + )]; + let sink_ts = ColumnSchema::new( + "event_time", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true) + .with_default_constraint(Some(ColumnDefaultConstraint::Function("now()".into()))) + .unwrap(); + assert!(is_explicit_source_timestamp_compatibility( + &output, + &[Some(0)], + &[output[0].clone(), sink_ts.clone()], + &source, + )); + assert!(!is_explicit_source_timestamp_compatibility( + &output, + &[Some(1)], + &[output[0].clone(), sink_ts], + &source, + )); + let sink_ts_without_default = ColumnSchema::new( + "event_time", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true); + assert!(!is_explicit_source_timestamp_compatibility( + &output, + &[Some(0)], + &[output[0].clone(), sink_ts_without_default], + &source, + )); +} + +#[test] +fn stateless_rejects_reserved_auto_names_for_auto_sink() { + assert!( + validate_auto_column_names(&[ColumnSchema::new( + AUTO_CREATED_UPDATE_AT_TS_COL, + ConcreteDataType::int32_datatype(), + true, + )]) + .is_err() + ); + assert!( + validate_auto_column_names(&[ColumnSchema::new( + AUTO_CREATED_PLACEHOLDER_TS_COL, + ConcreteDataType::int32_datatype(), + true, + )]) + .is_err() + ); +} + +#[test] +fn stateless_distinct_preserves_direct_column_lineage() { + let source = Arc::new(Schema::new(vec![ + ColumnSchema::new("number", ConcreteDataType::int32_datatype(), false), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ])); + let provider = MemTable::try_new( + Arc::new(datafusion::arrow::datatypes::Schema::new(vec![ + Field::new("number", ArrowDataType::Int32, false), + Field::new( + "ts", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + ), + ])), + vec![vec![]], + ) + .unwrap(); + let plan = LogicalPlanBuilder::scan( + TableReference::bare("source"), + provider_as_source(Arc::new(provider)), + None, + ) + .unwrap() + .project(vec![datafusion_expr::col("number").alias("dis")]) + .unwrap() + .distinct() + .unwrap() + .build() + .unwrap(); + + let (output, lineage) = super::output_column_schemas(&plan, &source).unwrap(); + assert_eq!(output[0].name, "dis"); + assert_eq!(lineage, vec![Some(0)]); + let relation = super::relation_desc_from_output(&output, &lineage, &[0]); + assert_eq!(relation.typ.keys[0].column_indices, vec![0]); +} + +#[test] +fn stateless_normalizes_dictionary_output_type() { + let field = Field::new_dictionary("host", ArrowDataType::UInt32, ArrowDataType::Utf8, true); + let arrow_schema = Arc::new(datafusion::arrow::datatypes::Schema::new(vec![field])); + let provider = MemTable::try_new(arrow_schema.clone(), vec![vec![]]).unwrap(); + let plan = LogicalPlanBuilder::scan( + TableReference::bare("source"), + provider_as_source(Arc::new(provider)), + None, + ) + .unwrap() + .build() + .unwrap(); + let source = Arc::new(Schema::new(vec![ColumnSchema::new( + "host", + ConcreteDataType::string_datatype(), + true, + )])); + let (output, lineage) = super::output_column_schemas(&plan, &source).unwrap(); + assert_eq!(output[0].data_type, ConcreteDataType::string_datatype()); + assert_eq!(lineage, vec![Some(0)]); + + let relation = super::relation_desc_from_output(&output, &lineage, &[0]); + assert_eq!(relation.typ.keys[0].column_indices, vec![0]); +} + +#[test] +fn stateless_allows_only_trailing_auto_columns() { + let ordinary = ColumnSchema::new("value", ConcreteDataType::int32_datatype(), false); + let update_at = ColumnSchema::new( + AUTO_CREATED_UPDATE_AT_TS_COL, + ConcreteDataType::timestamp_second_datatype(), + true, + ); + let placeholder = ColumnSchema::new( + AUTO_CREATED_PLACEHOLDER_TS_COL, + ConcreteDataType::timestamp_microsecond_datatype(), + true, + ) + .with_time_index(true); + + assert_eq!( + sink_output_column_count(&[ordinary.clone(), update_at.clone()]).unwrap(), + 1 + ); + assert_eq!( + sink_output_column_count(&[ordinary.clone(), update_at, placeholder]).unwrap(), + 1 + ); + assert!( + validate_sink_layout( + std::slice::from_ref(&ordinary), + std::slice::from_ref(&ordinary) + ) + .is_ok() + ); + assert!( + validate_sink_layout( + &[ordinary], + &[ + ColumnSchema::new("value", ConcreteDataType::int32_datatype(), false), + ColumnSchema::new("unexpected", ConcreteDataType::int32_datatype(), false), + ] + ) + .is_err() + ); +} + pub fn new_test_table_info_with_name>( table_id: TableId, table_name: &str, @@ -63,23 +308,830 @@ pub fn new_test_table_info_with_name>( fn mock_harness_flow_node_manager() {} #[test] -fn test_expire_after_secs_to_millis() { - assert_eq!(expire_after_secs_to_millis(300).unwrap(), 300_000); - assert_eq!(expire_after_secs_to_millis(0).unwrap(), 0); +fn stateless_captured_slot_rejects_inactive_or_detached_slot() { + let slot = super::StatelessFlowSlot { + runtime: Arc::new(tokio::sync::RwLock::new( + super::StatelessFlowRuntime::default(), + )), + active: std::sync::atomic::AtomicBool::new(false), + rebuild_attempts: std::sync::atomic::AtomicUsize::new(0), + }; - let max_secs = i64::MAX / 1_000; - assert_eq!( - expire_after_secs_to_millis(max_secs).unwrap(), - max_secs * 1_000 - ); + assert!(super::validate_captured_slot(&slot, None, 1, 42).is_err()); } #[test] -fn test_expire_after_secs_to_millis_invalid() { - let invalid_values = [i64::MAX / 1_000 + 1, -1]; +fn stateless_captured_slot_rejects_source_mismatch() { + let slot = super::StatelessFlowSlot { + runtime: Arc::new(tokio::sync::RwLock::new( + super::StatelessFlowRuntime::default(), + )), + active: std::sync::atomic::AtomicBool::new(true), + rebuild_attempts: std::sync::atomic::AtomicUsize::new(0), + }; - for invalid_secs in invalid_values { - let error = expire_after_secs_to_millis(invalid_secs).unwrap_err(); - assert!(matches!(error, Error::InvalidQuery { .. })); + assert!(super::validate_captured_slot(&slot, Some(2), 1, 42).is_err()); +} + +#[derive(Default)] +struct RecordingSink { + inserts: std::sync::Mutex>, + failed_calls: std::sync::atomic::AtomicUsize, +} + +#[async_trait::async_trait] +impl crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError for RecordingSink { + async fn do_query( + &self, + request: api::v1::greptime_request::Request, + _: session::context::QueryContextRef, + ) -> std::result::Result { + let api::v1::greptime_request::Request::RowInserts(request) = request else { + panic!("unexpected frontend request"); + }; + if request + .inserts + .iter() + .any(|insert| insert.table_name == "failed_sink") + { + self.failed_calls + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + return Err(BoxedError::new( + InvalidQuerySnafu { + reason: "injected sink failure", + } + .build(), + )); + } + let count = request + .inserts + .iter() + .map(|insert| insert.rows.as_ref().unwrap().rows.len()) + .sum(); + self.inserts.lock().unwrap().extend(request.inserts); + Ok(common_query::Output::new_with_affected_rows(count)) } } + +struct StreamingHarness { + engine: StreamingEngine, + metadata: TableMetadataManagerRef, + catalog: Arc, + sink: Arc, +} + +impl StreamingHarness { + async fn new() -> Self { + let metadata = Arc::new(common_meta::key::TableMetadataManager::new(Arc::new( + common_meta::kv_backend::memory::MemoryKvBackend::new(), + ))); + metadata.init().await.unwrap(); + let catalog = catalog::memory::new_memory_catalog_manager().unwrap(); + let query = query::QueryEngineFactory::new( + catalog.clone(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + let sink = Arc::new(RecordingSink::default()); + let handler: Arc< + dyn crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError, + > = sink.clone(); + let frontend = + FrontendClient::from_grpc_handler(Arc::downgrade(&handler), QueryOptions::default()); + Self { + engine: StreamingEngine::new(None, query, metadata.clone(), Arc::new(frontend)), + metadata, + catalog, + sink, + } + } + + async fn table(&self, id: u32, name: &str) -> TableInfo { + let info = new_test_table_info_with_name(id, name, []); + self.metadata + .create_table_metadata( + info.clone(), + common_meta::key::table_route::TableRouteValue::physical(vec![]), + Default::default(), + ) + .await + .unwrap(); + self.register(&info); + info + } + + fn register(&self, info: &TableInfo) { + self.catalog + .register_table_sync(catalog::RegisterTableRequest { + catalog: info.catalog_name.clone(), + schema: info.schema_name.clone(), + table_name: info.name.clone(), + table_id: info.ident.table_id, + table: table::test_util::EmptyTable::from_table_info(info), + }) + .unwrap(); + } + + async fn flow(&self, id: FlowId, source: u32, sink: &str, sql: &str) { + self.engine + .create_flow_inner(CreateFlowArgs { + flow_id: id, + source_table_ids: vec![source], + sink_table_name: ["greptime".into(), "public".into(), sink.into()], + create_if_not_exists: false, + or_replace: false, + expire_after: None, + eval_interval: None, + comment: None, + sql: sql.into(), + flow_options: Default::default(), + query_ctx: Some(session::context::QueryContext::arc().as_ref().clone()), + eval_schedule: None, + }) + .await + .unwrap(); + } + + fn take_numbers(&self) -> Vec<(String, Vec)> { + let inserts = std::mem::take(&mut *self.sink.inserts.lock().unwrap()); + inserts + .into_iter() + .map(|insert| { + let mut values = insert + .rows + .unwrap() + .rows + .into_iter() + .map(|row| { + let Some(api::v1::value::ValueData::I32Value(value)) = + row.values[0].value_data + else { + panic!("expected int32 output"); + }; + value + }) + .collect::>(); + values.sort_unstable(); + (insert.table_name, values) + }) + .collect() + } +} + +fn mirror_request(table: u32, region: u32, values: &[i32]) -> api::v1::region::InsertRequest { + use api::v1::value::ValueData; + api::v1::region::InsertRequest { + region_id: RegionId::new(table, region).as_u64(), + rows: Some(api::v1::Rows { + schema: util::column_schemas_to_proto( + vec![ + ColumnSchema::new("number", ConcreteDataType::int32_datatype(), true), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ], + &["number".into()], + ) + .unwrap(), + rows: values + .iter() + .map(|value| api::v1::Row { + values: vec![ + api::v1::Value { + value_data: Some(ValueData::I32Value(*value)), + }, + api::v1::Value { + value_data: Some(ValueData::TimestampMillisecondValue(1)), + }, + ], + }) + .collect(), + }), + ..Default::default() + } +} + +#[tokio::test] +async fn stateless_rejects_aggregate_flow_with_actionable_diagnostic() { + let h = StreamingHarness::new().await; + h.table(1, "source").await; + + let err = h + .engine + .create_flow_inner(CreateFlowArgs { + flow_id: 1, + source_table_ids: vec![1], + sink_table_name: ["greptime".into(), "public".into(), "sink".into()], + create_if_not_exists: true, + or_replace: true, + expire_after: None, + eval_interval: None, + comment: None, + sql: "SELECT number, count(*) FROM source GROUP BY number".into(), + flow_options: Default::default(), + query_ctx: Some(session::context::QueryContext::arc().as_ref().clone()), + eval_schedule: None, + }) + .await + .unwrap_err(); + + let message = err.to_string(); + assert!(message.contains( + "Aggregation is unsupported in streaming flows. Recreate the flow to select batching mode." + )); + assert!( + message.contains("A source table with TTL=instant must use persisted retention first.") + ); + assert!(message.contains("Aggregation SQL without a time window requires EVAL INTERVAL.")); + assert!(!h.engine.flow_exist_inner(1).await.unwrap()); + assert!(h.take_numbers().is_empty()); +} + +#[tokio::test] +async fn stateless_failed_flow_and_table_do_not_starve_healthy_sinks() { + let h = StreamingHarness::new().await; + h.table(1, "source_a").await; + h.table(2, "source_b").await; + h.table(3, "failed_sink").await; + h.table(4, "healthy_a").await; + h.table(5, "healthy_b").await; + h.flow(1, 1, "failed_sink", "SELECT number, ts FROM source_a") + .await; + h.flow(2, 1, "healthy_a", "SELECT number, ts FROM source_a") + .await; + h.flow(3, 2, "healthy_b", "SELECT number, ts FROM source_b") + .await; + assert!( + h.engine + .handle_inserts_inner(api::v1::region::InsertRequests { + requests: vec![mirror_request(1, 0, &[11]), mirror_request(2, 0, &[22])], + }) + .await + .is_err() + ); + assert_eq!( + h.take_numbers(), + vec![ + ("healthy_a".into(), vec![11]), + ("healthy_b".into(), vec![22]) + ] + ); + let failed_slot = h + .engine + .flow_ids_for_table(1) + .await + .into_iter() + .find(|(id, _)| *id == 1) + .unwrap() + .1; + assert!(failed_slot.runtime.read().await.failed_rebuild.is_none()); + assert!( + h.engine + .handle_inserts_inner(api::v1::region::InsertRequests { + requests: vec![mirror_request(1, 0, &[12])], + }) + .await + .is_err() + ); + assert_eq!(h.take_numbers(), vec![("healthy_a".into(), vec![12])]); + assert_eq!( + h.sink + .failed_calls + .load(std::sync::atomic::Ordering::Relaxed), + 2 + ); + assert!(failed_slot.runtime.read().await.failed_rebuild.is_none()); + + let mut malformed = mirror_request(1, 1, &[99]); + malformed.rows.as_mut().unwrap().schema.pop(); + assert!( + h.engine + .handle_inserts_inner(api::v1::region::InsertRequests { + requests: vec![ + mirror_request(1, 0, &[33]), + malformed, + mirror_request(1, 2, &[44]), + mirror_request(2, 0, &[55]) + ], + }) + .await + .is_err() + ); + assert_eq!(h.take_numbers(), vec![("healthy_b".into(), vec![55])]); + h.engine.remove_flow_inner(1).await.unwrap(); + h.engine + .handle_inserts_inner(api::v1::region::InsertRequests { + requests: vec![mirror_request(1, 0, &[66])], + }) + .await + .unwrap(); + assert_eq!(h.take_numbers(), vec![("healthy_a".into(), vec![66])]); +} + +#[tokio::test] +async fn stateless_distinct_groups_regions_without_retaining_previous_envelope() { + let h = StreamingHarness::new().await; + h.table(1, "source").await; + h.table(2, "sink").await; + h.flow( + 1, + 1, + "sink", + "SELECT DISTINCT number AS value, ts FROM source", + ) + .await; + for _ in 0..2 { + h.engine + .handle_inserts_inner(api::v1::region::InsertRequests { + requests: vec![mirror_request(1, 0, &[1, 2]), mirror_request(1, 1, &[2, 3])], + }) + .await + .unwrap(); + assert_eq!(h.take_numbers(), vec![("sink".into(), vec![1, 2, 3])]); + } +} + +#[tokio::test] +async fn stateless_schema_bump_rebuilds_for_current_and_subsequent_writes() { + let h = StreamingHarness::new().await; + let mut source = h.table(1, "source").await; + h.table(2, "sink").await; + h.flow(1, 1, "sink", "SELECT number, ts FROM source").await; + let current = h + .metadata + .table_info_manager() + .get(1) + .await + .unwrap() + .unwrap(); + let mut columns = source.meta.schema.column_schemas().to_vec(); + columns.push(ColumnSchema::new( + "extra", + ConcreteDataType::int32_datatype(), + true, + )); + source.meta.schema = Arc::new( + SchemaBuilder::try_from(columns) + .unwrap() + .version(124) + .build() + .unwrap(), + ); + h.metadata + .update_table_info(¤t, None, source.clone()) + .await + .unwrap(); + // The metadata is newer than the catalog provider: do not publish a plan + // carrying old column indices under the new version. + assert!( + h.engine + .handle_inserts_inner(api::v1::region::InsertRequests { + requests: vec![mirror_request(1, 0, &[99])], + }) + .await + .is_err() + ); + assert!(h.take_numbers().is_empty()); + h.catalog + .deregister_table_sync(catalog::DeregisterTableRequest { + catalog: source.catalog_name.clone(), + schema: source.schema_name.clone(), + table_name: source.name.clone(), + }) + .unwrap(); + h.register(&source); + let slot = h + .engine + .flow_ids_for_table(1) + .await + .into_iter() + .find(|(id, _)| *id == 1) + .unwrap() + .1; + assert_eq!( + slot.runtime + .read() + .await + .flow + .as_ref() + .unwrap() + .source_schema_version, + 123 + ); + // The stale-provider failure above is cooled down; expire this private test deadline + // after repairing the catalog so the existing rebuild behavior remains deterministic. + slot.runtime.write().await.failed_rebuild = Some((124, tokio::time::Instant::now())); + let (left_rows, types, version) = h + .engine + .handle_insert_request(mirror_request(1, 0, &[5])) + .await + .unwrap(); + let (right_rows, _, _) = h + .engine + .handle_insert_request(mirror_request(1, 1, &[6])) + .await + .unwrap(); + // Queue both readers behind a writer. Releasing it grants both readers, + // so neither rebuild can publish before both writes observe the old runtime. + let lease = slot.runtime.write().await; + let left = h + .engine + .execute_flow(1, slot.clone(), 1, left_rows, &types, version); + let right = h + .engine + .execute_flow(1, slot.clone(), 1, right_rows, &types, version); + tokio::pin!(left, right); + assert!(futures::poll!(&mut left).is_pending()); + assert!(futures::poll!(&mut right).is_pending()); + drop(lease); + let (left, right) = tokio::join!(left, right); + left.unwrap(); + right.unwrap(); + let mut inserted = h.take_numbers(); + inserted.sort(); + assert_eq!( + inserted, + vec![("sink".into(), vec![5]), ("sink".into(), vec![6])] + ); + assert_eq!( + slot.runtime + .read() + .await + .flow + .as_ref() + .unwrap() + .source_schema_version, + 124 + ); + for number in [7, 8] { + h.engine + .handle_inserts_inner(api::v1::region::InsertRequests { + requests: vec![mirror_request(1, 0, &[number])], + }) + .await + .unwrap(); + assert_eq!(h.take_numbers(), vec![("sink".into(), vec![number])]); + } +} + +#[tokio::test] +async fn stateless_failed_schema_rebuild_is_cooled_down_and_retried() { + let h = StreamingHarness::new().await; + let mut source = h.table(1, "source").await; + let mut sink = h.table(2, "sink").await; + h.table(3, "healthy_sink").await; + h.flow(1, 1, "sink", "SELECT * FROM source").await; + h.flow(2, 1, "healthy_sink", "SELECT number, ts FROM source") + .await; + let slot = h + .engine + .flow_ids_for_table(1) + .await + .into_iter() + .find(|(id, _)| *id == 1) + .unwrap() + .1; + let current = h + .metadata + .table_info_manager() + .get(1) + .await + .unwrap() + .unwrap(); + let mut columns = source.meta.schema.column_schemas().to_vec(); + columns.push(ColumnSchema::new( + "extra", + ConcreteDataType::int32_datatype(), + true, + )); + source.meta.schema = Arc::new( + SchemaBuilder::try_from(columns.clone()) + .unwrap() + .version(124) + .build() + .unwrap(), + ); + h.metadata + .update_table_info(¤t, None, source.clone()) + .await + .unwrap(); + h.catalog + .deregister_table_sync(catalog::DeregisterTableRequest { + catalog: source.catalog_name.clone(), + schema: source.schema_name.clone(), + table_name: source.name.clone(), + }) + .unwrap(); + h.register(&source); + + // Queue two stale runtimes before either can take the publication writer. The second + // must recheck the failed attempt under the writer rather than rebuilding again. + let (left_rows, types, version) = h + .engine + .handle_insert_request(mirror_request(1, 0, &[1])) + .await + .unwrap(); + let (right_rows, _, _) = h + .engine + .handle_insert_request(mirror_request(1, 1, &[2])) + .await + .unwrap(); + let lease = slot.runtime.write().await; + let left = h + .engine + .execute_flow(1, slot.clone(), 1, left_rows, &types, version); + let right = h + .engine + .execute_flow(1, slot.clone(), 1, right_rows, &types, version); + tokio::pin!(left, right); + assert!(futures::poll!(&mut left).is_pending()); + assert!(futures::poll!(&mut right).is_pending()); + drop(lease); + let (left, right) = tokio::join!(left, right); + let errors = [ + left.unwrap_err().to_string(), + right.unwrap_err().to_string(), + ]; + assert!( + errors + .iter() + .any(|error| error.contains("Flow output has 3 columns, but sink has 2 columns")) + ); + assert!(errors.iter().any(|error| { + error.contains("Flow 1 schema rebuild for source version 124 is cooling down") + })); + assert_eq!( + slot.rebuild_attempts + .load(std::sync::atomic::Ordering::Relaxed), + 1 + ); + let deadline = slot.runtime.read().await.failed_rebuild.unwrap(); + + // SELECT * is still invalid, while the sibling flow keeps producing output. + assert!( + h.engine + .handle_inserts_inner(api::v1::region::InsertRequests { + requests: vec![mirror_request(1, 0, &[3])], + }) + .await + .is_err() + ); + assert_eq!(h.take_numbers(), vec![("healthy_sink".into(), vec![3])]); + assert_eq!(slot.runtime.read().await.failed_rebuild, Some(deadline)); + assert_eq!( + slot.rebuild_attempts + .load(std::sync::atomic::Ordering::Relaxed), + 1 + ); + + // A new source version bypasses the old version's cooldown immediately. + let current = h + .metadata + .table_info_manager() + .get(1) + .await + .unwrap() + .unwrap(); + source.meta.schema = Arc::new( + SchemaBuilder::try_from(columns.clone()) + .unwrap() + .version(125) + .build() + .unwrap(), + ); + h.metadata + .update_table_info(¤t, None, source.clone()) + .await + .unwrap(); + h.catalog + .deregister_table_sync(catalog::DeregisterTableRequest { + catalog: source.catalog_name.clone(), + schema: source.schema_name.clone(), + table_name: source.name.clone(), + }) + .unwrap(); + h.register(&source); + assert!( + h.engine + .handle_inserts_inner(api::v1::region::InsertRequests { + requests: vec![mirror_request(1, 0, &[5])], + }) + .await + .is_err() + ); + assert_eq!( + slot.rebuild_attempts + .load(std::sync::atomic::Ordering::Relaxed), + 2 + ); + assert_eq!(h.take_numbers(), vec![("healthy_sink".into(), vec![5])]); + + // Make the rebuilt SELECT * layout valid, then expire only this test's private deadline. + let current_sink = h + .metadata + .table_info_manager() + .get(2) + .await + .unwrap() + .unwrap(); + sink.meta.schema = Arc::new( + SchemaBuilder::try_from(columns) + .unwrap() + .version(124) + .build() + .unwrap(), + ); + h.metadata + .update_table_info(¤t_sink, None, sink.clone()) + .await + .unwrap(); + h.catalog + .deregister_table_sync(catalog::DeregisterTableRequest { + catalog: source.catalog_name.clone(), + schema: source.schema_name.clone(), + table_name: source.name.clone(), + }) + .unwrap(); + h.catalog + .deregister_table_sync(catalog::DeregisterTableRequest { + catalog: sink.catalog_name.clone(), + schema: sink.schema_name.clone(), + table_name: sink.name.clone(), + }) + .unwrap(); + h.register(&source); + h.register(&sink); + slot.runtime.write().await.failed_rebuild = Some((125, tokio::time::Instant::now())); + h.engine + .handle_inserts_inner(api::v1::region::InsertRequests { + requests: vec![mirror_request(1, 0, &[6])], + }) + .await + .unwrap(); + assert_eq!( + slot.rebuild_attempts + .load(std::sync::atomic::Ordering::Relaxed), + 3 + ); + assert!(slot.runtime.read().await.failed_rebuild.is_none()); + let mut inserted = h.take_numbers(); + inserted.sort(); + assert_eq!( + inserted, + vec![("healthy_sink".into(), vec![6]), ("sink".into(), vec![6])] + ); + + // IF NOT EXISTS is a no-op and preserves cooldown; a successful replacement clears it. + let current = h + .metadata + .table_info_manager() + .get(1) + .await + .unwrap() + .unwrap(); + let mut columns = source.meta.schema.column_schemas().to_vec(); + columns.push(ColumnSchema::new( + "extra2", + ConcreteDataType::int32_datatype(), + true, + )); + source.meta.schema = Arc::new( + SchemaBuilder::try_from(columns) + .unwrap() + .version(126) + .build() + .unwrap(), + ); + h.metadata + .update_table_info(¤t, None, source.clone()) + .await + .unwrap(); + h.catalog + .deregister_table_sync(catalog::DeregisterTableRequest { + catalog: source.catalog_name.clone(), + schema: source.schema_name.clone(), + table_name: source.name.clone(), + }) + .unwrap(); + h.register(&source); + assert!( + h.engine + .handle_inserts_inner(api::v1::region::InsertRequests { + requests: vec![mirror_request(1, 0, &[7])], + }) + .await + .is_err() + ); + assert_eq!(h.take_numbers(), vec![("healthy_sink".into(), vec![7])]); + let cooldown = slot.runtime.read().await.failed_rebuild.unwrap(); + let mut args = slot + .runtime + .read() + .await + .flow + .as_ref() + .unwrap() + .create_args + .clone(); + args.or_replace = true; + assert!(h.engine.create_flow_inner(args.clone()).await.is_err()); + assert_eq!(slot.runtime.read().await.failed_rebuild, Some(cooldown)); + let current_sink = h + .metadata + .table_info_manager() + .get(2) + .await + .unwrap() + .unwrap(); + sink.meta.schema = source.meta.schema.clone(); + h.metadata + .update_table_info(¤t_sink, None, sink.clone()) + .await + .unwrap(); + h.catalog + .deregister_table_sync(catalog::DeregisterTableRequest { + catalog: source.catalog_name.clone(), + schema: source.schema_name.clone(), + table_name: source.name.clone(), + }) + .unwrap(); + h.catalog + .deregister_table_sync(catalog::DeregisterTableRequest { + catalog: sink.catalog_name.clone(), + schema: sink.schema_name.clone(), + table_name: sink.name.clone(), + }) + .unwrap(); + h.register(&source); + h.register(&sink); + args.or_replace = false; + args.create_if_not_exists = true; + h.engine.create_flow_inner(args.clone()).await.unwrap(); + assert_eq!(slot.runtime.read().await.failed_rebuild, Some(cooldown)); + args.create_if_not_exists = false; + args.or_replace = true; + h.engine.create_flow_inner(args).await.unwrap(); + assert!(slot.runtime.read().await.failed_rebuild.is_none()); +} + +#[tokio::test] +async fn stateless_rejects_wrong_provider_identity_and_detached_lifecycle() { + let h = StreamingHarness::new().await; + let source = h.table(1, "source").await; + h.table(2, "sink").await; + h.flow(1, 1, "sink", "SELECT number, ts FROM source").await; + let old_slot = h.engine.flow_ids_for_table(1).await.pop().unwrap().1; + old_slot.runtime.write().await.failed_rebuild = Some((123, tokio::time::Instant::now())); + let (rows, types, version) = h + .engine + .handle_insert_request(mirror_request(1, 0, &[9])) + .await + .unwrap(); + h.engine.remove_flow_inner(1).await.unwrap(); + h.flow(1, 1, "sink", "SELECT number, ts FROM source").await; + let fresh_slot = h.engine.flow_ids_for_table(1).await.pop().unwrap().1; + assert!(fresh_slot.runtime.read().await.failed_rebuild.is_none()); + assert!( + h.engine + .execute_flow(1, old_slot, 1, rows, &types, version) + .await + .is_err() + ); + assert!(h.take_numbers().is_empty()); + + h.catalog + .deregister_table_sync(catalog::DeregisterTableRequest { + catalog: source.catalog_name.clone(), + schema: source.schema_name.clone(), + table_name: source.name.clone(), + }) + .unwrap(); + let mut wrong = source.clone(); + wrong.ident.table_id = 99; + h.register(&wrong); + let args = h + .engine + .stateless_flows + .read() + .await + .get(&1) + .unwrap() + .runtime + .read() + .await + .flow + .as_ref() + .unwrap() + .create_args + .clone(); + assert!(h.engine.build_stateless_flow(&args, false).await.is_err()); + assert!(h.take_numbers().is_empty()); +} diff --git a/src/flow/src/adapter/util.rs b/src/flow/src/adapter/util.rs index b9e73e296c3..f4b1d0396d6 100644 --- a/src/flow/src/adapter/util.rs +++ b/src/flow/src/adapter/util.rs @@ -14,43 +14,24 @@ //! Util functions for adapter -use std::sync::Arc; - use api::helper::ColumnDataTypeWrapper; use api::v1::column_def::options_from_column_schema; -use api::v1::{ColumnDataType, ColumnDataTypeExtension, CreateTableExpr, SemanticType}; +use api::v1::{ColumnDataType, ColumnDataTypeExtension, SemanticType}; use common_error::ext::BoxedError; use common_meta::key::table_info::TableInfoValue; -use common_meta::rpc::ddl::TriggerReason; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema}; use itertools::Itertools; use operator::expr_helper; -use session::context::QueryContextBuilder; -use snafu::{OptionExt, ResultExt}; +use snafu::ResultExt; use table::table_reference::TableReference; use crate::StreamingEngine; use crate::adapter::table_source::TableDesc; -use crate::adapter::{AUTO_CREATED_PLACEHOLDER_TS_COL, TableName, WorkerHandle}; -use crate::error::{Error, ExternalSnafu, UnexpectedSnafu}; +use crate::adapter::{AUTO_CREATED_PLACEHOLDER_TS_COL, TableName}; +use crate::error::{Error, ExternalSnafu}; use crate::repr::{ColumnType, RelationDesc, RelationType}; impl StreamingEngine { - /// Get a worker handle for creating flow, using round robin to select a worker - pub(crate) async fn get_worker_handle_for_create_flow(&self) -> &WorkerHandle { - let use_idx = { - let mut selector = self.worker_selector.lock().await; - if *selector >= self.worker_handles.len() { - *selector = 0 - }; - let use_idx = *selector; - *selector += 1; - use_idx - }; - // Safety: selector is always in bound - &self.worker_handles[use_idx] - } - /// Create table from given schema(will adjust to add auto column if needed), return true if table is created pub(crate) async fn create_table_from_relation( &self, @@ -81,7 +62,9 @@ impl StreamingEngine { .map_err(BoxedError::new) .context(ExternalSnafu)?; - self.submit_create_sink_table_ddl(create_expr).await?; + self.frontend_client + .create(create_expr, &table_name[0], &table_name[1]) + .await?; Ok(true) } @@ -109,36 +92,6 @@ impl StreamingEngine { Ok(None) } } - - /// submit a create table ddl - pub(crate) async fn submit_create_sink_table_ddl( - &self, - mut create_table: CreateTableExpr, - ) -> Result<(), Error> { - let stmt_exec = { - self.frontend_invoker - .read() - .await - .as_ref() - .map(|f| f.statement_executor()) - } - .context(UnexpectedSnafu { - reason: "Failed to get statement executor", - })?; - let ctx = Arc::new( - QueryContextBuilder::default() - .current_catalog(create_table.catalog_name.clone()) - .current_schema(create_table.schema_name.clone()) - .build(), - ); - stmt_exec - .create_table_inner(&mut create_table, None, ctx, TriggerReason::AutoCreate) - .await - .map_err(BoxedError::new) - .context(ExternalSnafu)?; - - Ok(()) - } } pub fn table_info_value_to_relation_desc( diff --git a/src/flow/src/adapter/worker.rs b/src/flow/src/adapter/worker.rs deleted file mode 100644 index 71fbacb18ff..00000000000 --- a/src/flow/src/adapter/worker.rs +++ /dev/null @@ -1,606 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! For single-thread flow worker - -use std::collections::{BTreeMap, VecDeque}; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; - -use common_telemetry::info; -use dfir_rs::scheduled::graph::Dfir; -use enum_as_inner::EnumAsInner; -use snafu::ensure; -use tokio::sync::{Mutex, broadcast, mpsc, oneshot}; - -use crate::adapter::FlowId; -use crate::compute::{Context, DataflowState, ErrCollector}; -use crate::error::{Error, FlowAlreadyExistSnafu, InternalSnafu, UnexpectedSnafu}; -use crate::expr::{Batch, GlobalId}; -use crate::plan::TypedPlan; -use crate::repr::{self, DiffRow}; - -pub type SharedBuf = Arc>>; - -type ReqId = usize; - -/// Create both worker(`!Send`) and worker handle(`Send + Sync`) -pub fn create_worker<'a>() -> (WorkerHandle, Worker<'a>) { - let (itc_client, itc_server) = create_inter_thread_call(); - let worker_handle = WorkerHandle { - itc_client, - shutdown: AtomicBool::new(false), - }; - let worker = Worker { - task_states: BTreeMap::new(), - itc_server: Arc::new(Mutex::new(itc_server)), - }; - (worker_handle, worker) -} - -/// ActiveDataflowState is a wrapper around `Dfir` and `DataflowState` -pub(crate) struct ActiveDataflowState<'subgraph> { - df: Dfir<'subgraph>, - state: DataflowState, - err_collector: ErrCollector, -} - -impl std::fmt::Debug for ActiveDataflowState<'_> { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("ActiveDataflowState") - .field("df", &"") - .field("state", &self.state) - .field("err_collector", &self.err_collector) - .finish() - } -} - -impl Default for ActiveDataflowState<'_> { - fn default() -> Self { - ActiveDataflowState { - df: Dfir::new(), - state: DataflowState::default(), - err_collector: ErrCollector::default(), - } - } -} - -impl<'subgraph> ActiveDataflowState<'subgraph> { - /// Create a new render context, assigned with given global id - pub fn new_ctx<'ctx>(&'ctx mut self, global_id: GlobalId) -> Context<'ctx, 'subgraph> - where - 'subgraph: 'ctx, - { - Context { - id: global_id, - df: &mut self.df, - compute_state: &mut self.state, - err_collector: self.err_collector.clone(), - input_collection: Default::default(), - local_scope: Default::default(), - input_collection_batch: Default::default(), - local_scope_batch: Default::default(), - } - } - - pub fn set_current_ts(&mut self, ts: repr::Timestamp) { - self.state.set_current_ts(ts); - } - - pub fn set_last_exec_time(&mut self, ts: repr::Timestamp) { - self.state.set_last_exec_time(ts); - } - - /// Run all available subgraph - /// - /// return true if any subgraph actually executed - pub fn run_available(&mut self) -> bool { - self.state.run_available_with_schedule(&mut self.df) - } -} - -#[derive(Debug)] -pub struct WorkerHandle { - itc_client: InterThreadCallClient, - shutdown: AtomicBool, -} - -impl WorkerHandle { - /// create task, return task id - pub async fn create_flow(&self, create_reqs: Request) -> Result, Error> { - ensure!( - matches!(create_reqs, Request::Create { .. }), - InternalSnafu { - reason: format!( - "Flow Node/Worker itc failed, expect Request::Create, found {create_reqs:?}" - ), - } - ); - - let ret = self.itc_client.call_with_resp(create_reqs).await?; - ret.into_create().map_err(|ret| { - InternalSnafu { - reason: format!( - "Flow Node/Worker itc failed, expect Response::Create, found {ret:?}" - ), - } - .build() - })? - } - - /// remove task, return task id - pub async fn remove_flow(&self, flow_id: FlowId) -> Result { - let req = Request::Remove { flow_id }; - - let ret = self.itc_client.call_with_resp(req).await?; - - ret.into_remove().map_err(|ret| { - InternalSnafu { - reason: format!("Flow Node/Worker failed, expect Response::Remove, found {ret:?}"), - } - .build() - }) - } - - /// trigger running the worker, will not block, and will run the worker parallelly - /// - /// will set the current timestamp to `now` for all dataflows before running them - /// - /// `blocking` indicate whether it will wait til all dataflows are finished computing if true or - /// just start computing and return immediately if false - /// - /// the returned error is unrecoverable, and the worker should be shutdown/rebooted - pub async fn run_available(&self, now: repr::Timestamp, blocking: bool) -> Result<(), Error> { - common_telemetry::trace!("Running available with blocking={}", blocking); - if blocking { - let resp = self - .itc_client - .call_with_resp(Request::RunAvail { now, blocking }) - .await?; - common_telemetry::trace!("Running available with response={:?}", resp); - Ok(()) - } else { - self.itc_client - .call_no_resp(Request::RunAvail { now, blocking }) - } - } - - pub async fn contains_flow(&self, flow_id: FlowId) -> Result { - let req = Request::ContainTask { flow_id }; - let ret = self.itc_client.call_with_resp(req).await?; - - ret.into_contain_task().map_err(|ret| { - InternalSnafu { - reason: format!( - "Flow Node/Worker itc failed, expect Response::ContainTask, found {ret:?}" - ), - } - .build() - }) - } - - /// shutdown the worker - pub fn shutdown(&self) -> Result<(), Error> { - if !self.shutdown.fetch_or(true, Ordering::SeqCst) { - self.itc_client.call_no_resp(Request::Shutdown) - } else { - UnexpectedSnafu { - reason: "Worker already shutdown", - } - .fail() - } - } - - pub async fn get_full_flow_stat( - &self, - ) -> Result< - ( - BTreeMap, - BTreeMap, - BTreeMap, - ), - Error, - > { - let ret = self - .itc_client - .call_with_resp(Request::QueryFullFlowStat) - .await?; - ret.into_query_full_flow_stat().map_err(|ret| { - InternalSnafu { - reason: format!( - "Flow Node/Worker get_full_flow_stat failed, expected Response::QueryFullFlowStat, found {ret:?}" - ), - } - .build() - }) - } -} - -impl Drop for WorkerHandle { - fn drop(&mut self) { - if let Err(ret) = self.shutdown() { - common_telemetry::error!( - ret; - "While dropping Worker Handle, failed to shutdown worker, worker might be in inconsistent state." - ); - } else { - info!("Flow Worker shutdown due to Worker Handle dropped.") - } - } -} - -/// The actual worker that does the work and contain active state -#[derive(Debug)] -pub struct Worker<'subgraph> { - /// Task states - pub(crate) task_states: BTreeMap>, - itc_server: Arc>, -} - -impl<'s> Worker<'s> { - #[allow(clippy::too_many_arguments)] - pub fn create_flow( - &mut self, - flow_id: FlowId, - plan: TypedPlan, - sink_id: GlobalId, - sink_sender: mpsc::UnboundedSender, - source_ids: &[GlobalId], - src_recvs: Vec>, - // TODO(discord9): set expire duration for all arrangement and compare to sys timestamp instead - expire_after: Option, - or_replace: bool, - create_if_not_exists: bool, - err_collector: ErrCollector, - ) -> Result, Error> { - let already_exist = self.task_states.contains_key(&flow_id); - match (create_if_not_exists, or_replace, already_exist) { - // if replace, ignore that old flow exists - (_, true, true) => { - info!("Replacing flow with id={}", flow_id); - } - (false, false, true) => FlowAlreadyExistSnafu { id: flow_id }.fail()?, - // already exists, and not replace, return None - (true, false, true) => { - info!("Flow with id={} already exists, do nothing", flow_id); - return Ok(None); - } - // continue as normal - (_, _, false) => (), - } - - let mut cur_task_state = ActiveDataflowState::<'s> { - err_collector, - ..Default::default() - }; - cur_task_state.state.set_expire_after(expire_after); - - { - let mut ctx = cur_task_state.new_ctx(sink_id); - for (source_id, src_recv) in source_ids.iter().zip(src_recvs) { - let bundle = ctx.render_source_batch(src_recv)?; - ctx.insert_global_batch(*source_id, bundle); - } - - let rendered = ctx.render_plan_batch(plan)?; - ctx.render_unbounded_sink_batch(rendered, sink_sender); - } - self.task_states.insert(flow_id, cur_task_state); - Ok(Some(flow_id)) - } - - /// remove task, return true if a task is removed - pub fn remove_flow(&mut self, flow_id: FlowId) -> bool { - self.task_states.remove(&flow_id).is_some() - } - - /// Run the worker, blocking, until shutdown signal is received - pub fn run(&mut self) { - loop { - let (req, ret_tx) = if let Some(ret) = self.itc_server.blocking_lock().blocking_recv() { - ret - } else { - common_telemetry::error!( - "Worker's itc server has been closed unexpectedly, shutting down worker now." - ); - break; - }; - - let ret = self.handle_req(req); - match (ret, ret_tx) { - (Ok(Some(resp)), Some(ret_tx)) => { - if let Err(err) = ret_tx.send(resp) { - common_telemetry::error!( - err; - "Result receiver is dropped, can't send result" - ); - }; - } - (Ok(None), None) => continue, - (Ok(Some(resp)), None) => { - common_telemetry::error!( - "Expect no result for current request, but found {resp:?}" - ) - } - (Ok(None), Some(_)) => { - common_telemetry::error!("Expect result for current request, but found nothing") - } - (Err(()), _) => { - break; - } - } - } - } - - /// run with tick acquired from tick manager(usually means system time) - /// TODO(discord9): better tick management - pub fn run_tick(&mut self, now: repr::Timestamp) { - for (_flow_id, task_state) in self.task_states.iter_mut() { - task_state.set_current_ts(now); - task_state.set_last_exec_time(now); - task_state.run_available(); - } - } - /// handle request, return response if any, Err if receive shutdown signal - /// - /// return `Err(())` if receive shutdown request - fn handle_req(&mut self, req: Request) -> Result, ()> { - let ret = match req { - Request::Create { - flow_id, - plan, - sink_id, - sink_sender, - source_ids, - src_recvs, - expire_after, - or_replace, - create_if_not_exists, - err_collector, - } => { - let task_create_result = self.create_flow( - flow_id, - plan, - sink_id, - sink_sender, - &source_ids, - src_recvs, - expire_after, - or_replace, - create_if_not_exists, - err_collector, - ); - Some(Response::Create { - result: task_create_result, - }) - } - Request::Remove { flow_id } => { - let ret = self.remove_flow(flow_id); - Some(Response::Remove { result: ret }) - } - Request::RunAvail { now, blocking } => { - self.run_tick(now); - if blocking { - Some(Response::RunAvail) - } else { - None - } - } - Request::ContainTask { flow_id } => { - let ret = self.task_states.contains_key(&flow_id); - Some(Response::ContainTask { result: ret }) - } - Request::Shutdown => return Err(()), - Request::QueryFullFlowStat => { - let mut state_size = BTreeMap::new(); - let mut last_exec_time_map = BTreeMap::new(); - let mut start_time_map = BTreeMap::new(); - for (flow_id, task_state) in self.task_states.iter() { - state_size.insert(*flow_id, task_state.state.get_state_size()); - if let Some(t) = task_state.state.last_exec_time() { - last_exec_time_map.insert(*flow_id, t); - } - if let Some(t) = task_state.state.start_time() { - start_time_map.insert(*flow_id, t); - } - } - Some(Response::QueryFullFlowStat { - state_size, - last_exec_time_map, - start_time_map, - }) - } - }; - Ok(ret) - } -} - -#[derive(Debug, EnumAsInner)] -pub enum Request { - Create { - flow_id: FlowId, - plan: TypedPlan, - sink_id: GlobalId, - sink_sender: mpsc::UnboundedSender, - source_ids: Vec, - src_recvs: Vec>, - expire_after: Option, - or_replace: bool, - create_if_not_exists: bool, - err_collector: ErrCollector, - }, - Remove { - flow_id: FlowId, - }, - /// Trigger the worker to run, useful after input buffer is full - RunAvail { - now: repr::Timestamp, - blocking: bool, - }, - ContainTask { - flow_id: FlowId, - }, - Shutdown, - QueryFullFlowStat, -} - -#[derive(Debug, EnumAsInner)] -enum Response { - Create { - result: Result, Error>, - // TODO(discord9): add flow err_collector - }, - Remove { - result: bool, - }, - ContainTask { - result: bool, - }, - RunAvail, - QueryFullFlowStat { - state_size: BTreeMap, - last_exec_time_map: BTreeMap, - start_time_map: BTreeMap, - }, -} - -fn create_inter_thread_call() -> (InterThreadCallClient, InterThreadCallServer) { - let (arg_send, arg_recv) = mpsc::unbounded_channel(); - let client = InterThreadCallClient { - arg_sender: arg_send, - }; - let server = InterThreadCallServer { arg_recv }; - (client, server) -} - -#[derive(Debug)] -struct InterThreadCallClient { - arg_sender: mpsc::UnboundedSender<(Request, Option>)>, -} - -impl InterThreadCallClient { - /// call without response - fn call_no_resp(&self, req: Request) -> Result<(), Error> { - self.arg_sender.send((req, None)).map_err(from_send_error) - } - - /// call with response - async fn call_with_resp(&self, req: Request) -> Result { - let (tx, rx) = oneshot::channel(); - self.arg_sender - .send((req, Some(tx))) - .map_err(from_send_error)?; - rx.await.map_err(|_| { - InternalSnafu { - reason: "Sender is dropped", - } - .build() - }) - } -} - -#[derive(Debug)] -struct InterThreadCallServer { - pub arg_recv: mpsc::UnboundedReceiver<(Request, Option>)>, -} - -impl InterThreadCallServer { - pub async fn recv(&mut self) -> Option<(Request, Option>)> { - self.arg_recv.recv().await - } - - pub fn blocking_recv(&mut self) -> Option<(Request, Option>)> { - self.arg_recv.blocking_recv() - } -} - -fn from_send_error(err: mpsc::error::SendError) -> Error { - InternalSnafu { - // this `err` will simply display `channel closed` - reason: format!( - "Worker's receiver channel have been closed unexpected: {}", - err - ), - } - .build() -} - -#[cfg(test)] -mod test { - use tokio::sync::oneshot; - - use super::*; - use crate::expr::Id; - use crate::plan::Plan; - use crate::repr::RelationType; - - #[test] - fn drop_handle() { - let (tx, rx) = oneshot::channel(); - let worker_thread_handle = std::thread::spawn(move || { - let (handle, mut worker) = create_worker(); - tx.send(handle).unwrap(); - worker.run(); - }); - let handle = rx.blocking_recv().unwrap(); - drop(handle); - worker_thread_handle.join().unwrap(); - } - - #[tokio::test] - pub async fn test_simple_get_with_worker_and_handle() { - let (tx, rx) = oneshot::channel(); - let worker_thread_handle = std::thread::spawn(move || { - let (handle, mut worker) = create_worker(); - tx.send(handle).unwrap(); - worker.run(); - }); - let handle = rx.await.unwrap(); - let src_ids = vec![GlobalId::User(1)]; - let (tx, rx) = broadcast::channel::(1024); - let (sink_tx, mut sink_rx) = mpsc::unbounded_channel::(); - let (flow_id, plan) = ( - 1, - TypedPlan { - plan: Plan::Get { - id: Id::Global(GlobalId::User(1)), - }, - schema: RelationType::new(vec![]).into_unnamed(), - }, - ); - let create_reqs = Request::Create { - flow_id, - plan, - sink_id: GlobalId::User(1), - sink_sender: sink_tx, - source_ids: src_ids, - src_recvs: vec![rx], - expire_after: None, - or_replace: false, - create_if_not_exists: true, - err_collector: ErrCollector::default(), - }; - assert_eq!( - handle.create_flow(create_reqs).await.unwrap(), - Some(flow_id) - ); - tx.send(Batch::empty()).unwrap(); - handle.run_available(0, true).await.unwrap(); - let (state_size, _, _) = handle.get_full_flow_stat().await.unwrap(); - assert_eq!(state_size.len(), 1); - assert_eq!(sink_rx.recv().await.unwrap(), Batch::empty()); - drop(handle); - worker_thread_handle.join().unwrap(); - } -} diff --git a/src/flow/src/batching_mode/frontend_client.rs b/src/flow/src/batching_mode/frontend_client.rs index e932933ff84..0207837bdc7 100644 --- a/src/flow/src/batching_mode/frontend_client.rs +++ b/src/flow/src/batching_mode/frontend_client.rs @@ -555,6 +555,81 @@ impl FrontendClient { } } + /// Handle an insert request with one attempt. + /// + /// Unlike [`Self::handle`], this does not use the batching retry policy. It + /// is intended for stateless streaming sinks, where retrying an insert can + /// duplicate rows. + pub(crate) async fn handle_insert_once( + &self, + req: api::v1::greptime_request::Request, + catalog: &str, + schema: &str, + peer_desc: &mut Option, + ) -> Result { + match self { + FrontendClient::Distributed { .. } => { + let db = self.get_random_active_frontend(catalog, schema).await?; + + *peer_desc = Some(PeerDesc::Dist { + peer: db.peer.clone(), + }); + + db.database + .handle(req.clone()) + .await + .with_context(|_| InvalidRequestSnafu { + context: format!("Failed to handle request at {:?}: {:?}", db.peer, req), + }) + } + FrontendClient::Standalone { + database_client, + query, + } => { + let ctx = QueryContextBuilder::default() + .current_catalog(catalog.to_string()) + .current_schema(schema.to_string()) + .extensions(HashMap::from([( + QUERY_PARALLELISM_HINT.to_string(), + query.parallelism.to_string(), + )])) + .build(); + let ctx = Arc::new(ctx); + let database_client = { + database_client + .handler + .lock() + .unwrap() + .as_ref() + .context(UnexpectedSnafu { + reason: "Standalone's frontend instance is not set", + })? + .upgrade() + .context(UnexpectedSnafu { + reason: "Failed to upgrade database client", + })? + }; + let resp: common_query::Output = database_client + .do_query(req, ctx) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + match resp.data { + common_query::OutputData::AffectedRows(rows) => rows.try_into().map_err(|_| { + UnexpectedSnafu { + reason: format!("Failed to convert rows to u32: {}", rows), + } + .build() + }), + _ => UnexpectedSnafu { + reason: "Unexpected output data", + } + .fail(), + } + } + } + } + /// Handle a request to frontend pub(crate) async fn handle( &self, @@ -700,6 +775,7 @@ impl std::fmt::Display for PeerDesc { #[cfg(test)] mod tests { use std::pin::Pin; + use std::sync::atomic::{AtomicUsize, Ordering}; use std::task::{Context, Poll}; use std::time::Duration; @@ -773,6 +849,11 @@ mod tests { #[derive(Debug)] struct MetricsHandler; + #[derive(Debug)] + struct InsertOnceHandler { + calls: Arc, + } + #[derive(Debug)] struct ExtensionAwareHandler; @@ -808,6 +889,18 @@ mod tests { } } + #[async_trait::async_trait] + impl GrpcQueryHandlerWithBoxedError for InsertOnceHandler { + async fn do_query( + &self, + _query: Request, + _ctx: QueryContextRef, + ) -> std::result::Result { + self.calls.fetch_add(1, Ordering::SeqCst); + Ok(Output::new_with_affected_rows(1)) + } + } + #[async_trait::async_trait] impl GrpcQueryHandlerWithBoxedError for MetricsHandler { async fn do_query( @@ -999,6 +1092,31 @@ mod tests { ); } + #[tokio::test] + async fn test_handle_insert_once_calls_standalone_handler_once() { + let calls = Arc::new(AtomicUsize::new(0)); + let handler: Arc = Arc::new(InsertOnceHandler { + calls: calls.clone(), + }); + let client = + FrontendClient::from_grpc_handler(Arc::downgrade(&handler), QueryOptions::default()); + let mut peer_desc = None; + + let affected_rows = client + .handle_insert_once( + Request::RowInserts(api::v1::RowInsertRequests { inserts: vec![] }), + "greptime", + "public", + &mut peer_desc, + ) + .await + .unwrap(); + + assert_eq!(affected_rows, 1); + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert!(peer_desc.is_none()); + } + #[tokio::test] async fn test_query_with_terminal_metrics_tracks_watermark_in_standalone_mode() { let handler: Arc = Arc::new(MetricsHandler); diff --git a/src/flow/src/batching_mode/time_window.rs b/src/flow/src/batching_mode/time_window.rs index 5dbd08150ac..1fb6b781d17 100644 --- a/src/flow/src/batching_mode/time_window.rs +++ b/src/flow/src/batching_mode/time_window.rs @@ -57,7 +57,6 @@ use crate::error::{ ArrowSnafu, DatafusionSnafu, DatatypesSnafu, ExternalSnafu, PlanSnafu, TimeSnafu, UnexpectedSnafu, }; -use crate::expr::error::DataTypeSnafu; /// Represents a test timestamp in seconds since the Unix epoch. const DEFAULT_TEST_TIMESTAMP: Timestamp = Timestamp::new_second(17_0000_0000); @@ -292,9 +291,10 @@ impl TimeWindowExpr { let mut vector = cdt.create_mutable_vector(rows.rows.len()); for row in rows.rows { let value = pb_value_to_value_ref(&row.values[ts_col_index], None); - vector.try_push_value_ref(&value).context(DataTypeSnafu { - msg: "Failed to convert rows to columns", - })?; + vector + .try_push_value_ref(&value) + .map_err(BoxedError::new) + .context(ExternalSnafu)?; } let vector = vector.to_vector(); diff --git a/src/flow/src/compute.rs b/src/flow/src/compute.rs deleted file mode 100644 index 8463039dcd8..00000000000 --- a/src/flow/src/compute.rs +++ /dev/null @@ -1,23 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! Build and Compute the dataflow - -mod render; -mod state; -mod types; - -pub(crate) use render::Context; -pub(crate) use state::DataflowState; -pub(crate) use types::ErrCollector; diff --git a/src/flow/src/compute/render.rs b/src/flow/src/compute/render.rs deleted file mode 100644 index 5be923aeac1..00000000000 --- a/src/flow/src/compute/render.rs +++ /dev/null @@ -1,527 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! In this file, `render` means convert a static `Plan` into a Executable Dataflow -//! -//! And the [`Context`] is the environment for the render process, it contains all the necessary information for the render process - -use std::collections::BTreeMap; - -use dfir_rs::scheduled::graph::Dfir; -use dfir_rs::scheduled::graph_ext::GraphExt; -use dfir_rs::scheduled::port::{PortCtx, SEND}; -use itertools::Itertools; -use snafu::OptionExt; - -use crate::compute::state::{DataflowState, Scheduler}; -use crate::compute::types::{Collection, CollectionBundle, ErrCollector, Toff}; -use crate::error::{Error, InvalidQuerySnafu, NotImplementedSnafu}; -use crate::expr::{self, Batch, GlobalId, LocalId}; -use crate::plan::{Plan, TypedPlan}; -use crate::repr::{self, DiffRow, RelationType}; - -mod map; -mod reduce; -mod src_sink; - -/// The Context for build a Operator with id of `GlobalId` -pub struct Context<'referred, 'df> { - pub id: GlobalId, - pub df: &'referred mut Dfir<'df>, - pub compute_state: &'referred mut DataflowState, - /// a list of all collections being used in the operator - /// - /// TODO(discord9): remove extra clone by counting usage and remove it on last usage? - pub input_collection: BTreeMap, - /// used by `Get`/`Let` Plan for getting/setting local variables - /// - /// TODO(discord9): consider if use Vec<(LocalId, CollectionBundle)> instead - pub local_scope: Vec>, - /// a list of all collections being used in the operator - /// - /// TODO(discord9): remove extra clone by counting usage and remove it on last usage? - pub input_collection_batch: BTreeMap>, - /// used by `Get`/`Let` Plan for getting/setting local variables - /// - /// TODO(discord9): consider if use Vec<(LocalId, CollectionBundle)> instead - pub local_scope_batch: Vec>>, - // Collect all errors in this operator's evaluation - pub err_collector: ErrCollector, -} - -impl Drop for Context<'_, '_> { - fn drop(&mut self) { - for bundle in std::mem::take(&mut self.input_collection) - .into_values() - .chain( - std::mem::take(&mut self.local_scope) - .into_iter() - .flat_map(|v| v.into_iter()) - .map(|(_k, v)| v), - ) - { - bundle.collection.into_inner().drop(self.df); - drop(bundle.arranged); - } - - for bundle in std::mem::take(&mut self.input_collection_batch) - .into_values() - .chain( - std::mem::take(&mut self.local_scope_batch) - .into_iter() - .flat_map(|v| v.into_iter()) - .map(|(_k, v)| v), - ) - { - bundle.collection.into_inner().drop(self.df); - drop(bundle.arranged); - } - // The automatically generated "drop glue" which recursively calls the destructors of all the fields (including the now empty `input_collection`) - } -} - -impl Context<'_, '_> { - pub fn insert_global(&mut self, id: GlobalId, collection: CollectionBundle) { - self.input_collection.insert(id, collection); - } - - pub fn insert_local(&mut self, id: LocalId, collection: CollectionBundle) { - if let Some(last) = self.local_scope.last_mut() { - last.insert(id, collection); - } else { - let first = BTreeMap::from([(id, collection)]); - self.local_scope.push(first); - } - } - - pub fn insert_global_batch(&mut self, id: GlobalId, collection: CollectionBundle) { - self.input_collection_batch.insert(id, collection); - } - - pub fn insert_local_batch(&mut self, id: LocalId, collection: CollectionBundle) { - if let Some(last) = self.local_scope_batch.last_mut() { - last.insert(id, collection); - } else { - let first = BTreeMap::from([(id, collection)]); - self.local_scope_batch.push(first); - } - } -} - -impl Context<'_, '_> { - /// Like `render_plan` but in Batch Mode - pub fn render_plan_batch(&mut self, plan: TypedPlan) -> Result, Error> { - match plan.plan { - Plan::Constant { rows } => Ok(self.render_constant_batch(rows, &plan.schema.typ)), - Plan::Get { id } => self.get_batch_by_id(id), - Plan::Let { id, value, body } => self.eval_batch_let(id, value, body), - Plan::Mfp { input, mfp } => self.render_mfp_batch(input, mfp, &plan.schema.typ), - Plan::Reduce { - input, - key_val_plan, - reduce_plan, - } => self.render_reduce_batch(input, &key_val_plan, &reduce_plan, &plan.schema.typ), - Plan::Join { .. } => NotImplementedSnafu { - reason: "Join is still WIP", - } - .fail(), - Plan::Union { .. } => NotImplementedSnafu { - reason: "Union is still WIP", - } - .fail(), - } - } - - /// Interpret plan to dataflow and prepare them for execution - /// - /// return the output handler of this plan - pub fn render_plan(&mut self, plan: TypedPlan) -> Result { - match plan.plan { - Plan::Constant { rows } => Ok(self.render_constant(rows)), - Plan::Get { id } => self.get_by_id(id), - Plan::Let { id, value, body } => self.eval_let(id, value, body), - Plan::Mfp { input, mfp } => self.render_mfp(input, mfp), - Plan::Reduce { - input, - key_val_plan, - reduce_plan, - } => self.render_reduce(input, key_val_plan, reduce_plan, plan.schema.typ), - Plan::Join { .. } => NotImplementedSnafu { - reason: "Join is still WIP", - } - .fail(), - Plan::Union { .. } => NotImplementedSnafu { - reason: "Union is still WIP", - } - .fail(), - } - } - - /// render Constant, take all rows that have a timestamp not greater than the current time - /// This function is primarily used for testing - /// Always assume input is sorted by timestamp - pub fn render_constant_batch( - &mut self, - rows: Vec, - output_type: &RelationType, - ) -> CollectionBundle { - let (send_port, recv_port) = self.df.make_edge::<_, Toff>("constant_batch"); - let mut per_time: BTreeMap> = Default::default(); - for (key, group) in &rows.into_iter().chunk_by(|(_row, ts, _diff)| *ts) { - per_time.entry(key).or_default().extend(group); - } - - let now = self.compute_state.current_time_ref(); - // TODO(discord9): better way to schedule future run - let scheduler = self.compute_state.get_scheduler(); - let scheduler_inner = scheduler.clone(); - let err_collector = self.err_collector.clone(); - - let output_type = output_type.clone(); - - let subgraph_id = - self.df - .add_subgraph_source("ConstantBatch", send_port, move |_ctx, send_port| { - // find the first timestamp that is greater than now - // use filter_map - - let mut after = per_time.split_off(&(*now.borrow() + 1)); - // swap - std::mem::swap(&mut per_time, &mut after); - let not_great_than_now = after; - - not_great_than_now.into_iter().for_each(|(_ts, rows)| { - err_collector.run(|| { - let rows = rows.into_iter().map(|(row, _ts, _diff)| row).collect(); - let batch = Batch::try_from_rows_with_types( - rows, - &output_type - .column_types - .iter() - .map(|ty| ty.scalar_type().clone()) - .collect_vec(), - )?; - send_port.give(vec![batch]); - Ok(()) - }); - }); - // schedule the next run - if let Some(next_run_time) = per_time.keys().next().copied() { - scheduler_inner.schedule_at(next_run_time); - } - }); - scheduler.set_cur_subgraph(subgraph_id); - - CollectionBundle::from_collection(Collection::from_port(recv_port)) - } - - /// render Constant, take all rows that have a timestamp not greater than the current time - /// - /// Always assume input is sorted by timestamp - pub fn render_constant(&mut self, rows: Vec) -> CollectionBundle { - let (send_port, recv_port) = self.df.make_edge::<_, Toff>("constant"); - let mut per_time: BTreeMap> = Default::default(); - for (key, group) in &rows.into_iter().chunk_by(|(_row, ts, _diff)| *ts) { - per_time.entry(key).or_default().extend(group); - } - - let now = self.compute_state.current_time_ref(); - // TODO(discord9): better way to schedule future run - let scheduler = self.compute_state.get_scheduler(); - let scheduler_inner = scheduler.clone(); - - let subgraph_id = - self.df - .add_subgraph_source("Constant", send_port, move |_ctx, send_port| { - // find the first timestamp that is greater than now - // use filter_map - - let mut after = per_time.split_off(&(*now.borrow() + 1)); - // swap - std::mem::swap(&mut per_time, &mut after); - let not_great_than_now = after; - - not_great_than_now.into_iter().for_each(|(_ts, rows)| { - send_port.give(rows); - }); - // schedule the next run - if let Some(next_run_time) = per_time.keys().next().copied() { - scheduler_inner.schedule_at(next_run_time); - } - }); - scheduler.set_cur_subgraph(subgraph_id); - - CollectionBundle::from_collection(Collection::from_port(recv_port)) - } - - pub fn get_batch_by_id(&mut self, id: expr::Id) -> Result, Error> { - let ret = match id { - expr::Id::Local(local) => { - let bundle = self - .local_scope_batch - .iter() - .rev() - .find_map(|scope| scope.get(&local)) - .with_context(|| InvalidQuerySnafu { - reason: format!("Local variable {:?} not found", local), - })?; - bundle.clone(self.df) - } - expr::Id::Global(id) => { - let bundle = - self.input_collection_batch - .get(&id) - .with_context(|| InvalidQuerySnafu { - reason: format!("Collection {:?} not found", id), - })?; - bundle.clone(self.df) - } - }; - Ok(ret) - } - - pub fn get_by_id(&mut self, id: expr::Id) -> Result { - let ret = match id { - expr::Id::Local(local) => { - let bundle = self - .local_scope - .iter() - .rev() - .find_map(|scope| scope.get(&local)) - .with_context(|| InvalidQuerySnafu { - reason: format!("Local variable {:?} not found", local), - })?; - bundle.clone(self.df) - } - expr::Id::Global(id) => { - let bundle = self - .input_collection - .get(&id) - .with_context(|| InvalidQuerySnafu { - reason: format!("Collection {:?} not found", id), - })?; - bundle.clone(self.df) - } - }; - Ok(ret) - } - - /// Eval `Let` operator, useful for assigning a value to a local variable - pub fn eval_batch_let( - &mut self, - id: LocalId, - value: Box, - body: Box, - ) -> Result, Error> { - let value = self.render_plan_batch(*value)?; - - self.local_scope_batch.push(Default::default()); - self.insert_local_batch(id, value); - let ret = self.render_plan_batch(*body)?; - Ok(ret) - } - - /// Eval `Let` operator, useful for assigning a value to a local variable - pub fn eval_let( - &mut self, - id: LocalId, - value: Box, - body: Box, - ) -> Result { - let value = self.render_plan(*value)?; - - self.local_scope.push(Default::default()); - self.insert_local(id, value); - let ret = self.render_plan(*body)?; - Ok(ret) - } -} - -/// The Common argument for all `Subgraph` in the render process -struct SubgraphArg<'a, T = Toff> { - now: repr::Timestamp, - err_collector: &'a ErrCollector, - scheduler: &'a Scheduler, - send: &'a PortCtx, -} - -#[cfg(test)] -mod test { - use std::cell::RefCell; - use std::rc::Rc; - - use dfir_rs::scheduled::graph::Dfir; - use dfir_rs::scheduled::graph_ext::GraphExt; - use dfir_rs::scheduled::handoff::VecHandoff; - use pretty_assertions::assert_eq; - - use super::*; - use crate::repr::Row; - pub fn run_and_check( - state: &mut DataflowState, - df: &mut Dfir, - time_range: std::ops::Range, - expected: BTreeMap>, - output: Rc>>, - ) { - for now in time_range { - state.set_current_ts(now); - state.run_available_with_schedule(df); - if !state.get_err_collector().is_empty() { - panic!( - "Errors occur: {:?}", - state.get_err_collector().get_all_blocking() - ) - } - assert!(state.get_err_collector().is_empty()); - if let Some(expected) = expected.get(&now) { - assert_eq!(*output.borrow(), *expected, "at ts={}", now); - } else { - assert_eq!(*output.borrow(), vec![], "at ts={}", now); - }; - output.borrow_mut().clear(); - } - } - - pub fn get_output_handle( - ctx: &mut Context, - mut bundle: CollectionBundle, - ) -> Rc>> { - let collection = bundle.collection; - let _arranged = bundle.arranged.pop_first().unwrap().1; - let output = Rc::new(RefCell::new(vec![])); - let output_inner = output.clone(); - let _subgraph = ctx.df.add_subgraph_sink( - "test_render_constant", - collection.into_inner(), - move |_ctx, recv| { - let data = recv.take_inner(); - let res = data.into_iter().flat_map(|v| v.into_iter()).collect_vec(); - output_inner.borrow_mut().clear(); - output_inner.borrow_mut().extend(res); - }, - ); - output - } - - pub fn harness_test_ctx<'r, 'h>( - df: &'r mut Dfir<'h>, - state: &'r mut DataflowState, - ) -> Context<'r, 'h> { - let err_collector = state.get_err_collector(); - Context { - id: GlobalId::User(0), - df, - compute_state: state, - input_collection: BTreeMap::new(), - local_scope: Default::default(), - input_collection_batch: BTreeMap::new(), - local_scope_batch: Default::default(), - err_collector, - } - } - - /// test if constant operator works properly - /// that is it only emit once, not multiple times - #[test] - fn test_render_constant() { - let mut df = Dfir::new(); - let mut state = DataflowState::default(); - let mut ctx = harness_test_ctx(&mut df, &mut state); - - let rows = vec![ - (Row::empty(), 1, 1), - (Row::empty(), 2, 1), - (Row::empty(), 3, 1), - ]; - let collection = ctx.render_constant(rows); - let collection = collection.collection.clone(ctx.df); - let cnt = Rc::new(RefCell::new(0)); - let cnt_inner = cnt.clone(); - let res_subgraph_id = ctx.df.add_subgraph_sink( - "test_render_constant", - collection.into_inner(), - move |_ctx, recv| { - let data = recv.take_inner(); - *cnt_inner.borrow_mut() += data.iter().map(|v| v.len()).sum::(); - }, - ); - ctx.compute_state.set_current_ts(2); - ctx.compute_state.run_available_with_schedule(ctx.df); - assert_eq!(*cnt.borrow(), 2); - - ctx.compute_state.set_current_ts(3); - ctx.compute_state.run_available_with_schedule(ctx.df); - // to get output - ctx.df.schedule_subgraph(res_subgraph_id); - ctx.df.run_available(); - - assert_eq!(*cnt.borrow(), 3); - } - - /// a simple example to show how to use source and sink - #[test] - fn example_source_sink() { - let mut df = Dfir::new(); - let (send_port, recv_port) = df.make_edge::<_, VecHandoff>("test_handoff"); - df.add_subgraph_source("test_handoff_source", send_port, move |_ctx, send| { - for i in 0..10 { - send.give(vec![i]); - } - }); - - let sum = Rc::new(RefCell::new(0)); - let sum_move = sum.clone(); - let sink = df.add_subgraph_sink("test_handoff_sink", recv_port, move |_ctx, recv| { - let data = recv.take_inner(); - *sum_move.borrow_mut() += data.iter().sum::(); - }); - - df.run_available(); - assert_eq!(sum.borrow().to_owned(), 45); - df.schedule_subgraph(sink); - df.run_available(); - - assert_eq!(sum.borrow().to_owned(), 45); - } - - #[test] - fn test_tee_auto_schedule() { - use dfir_rs::scheduled::handoff::TeeingHandoff as Toff; - let mut df = Dfir::new(); - let (send_port, recv_port) = df.make_edge::<_, Toff>("test_handoff"); - let source = df.add_subgraph_source("test_handoff_source", send_port, move |_ctx, send| { - for i in 0..10 { - send.give(vec![i]); - } - }); - let teed_recv_port = recv_port.tee(&mut df); - - let sum = Rc::new(RefCell::new(0)); - let sum_move = sum.clone(); - let _sink = df.add_subgraph_sink("test_handoff_sink", teed_recv_port, move |_ctx, recv| { - let data = recv.take_inner(); - *sum_move.borrow_mut() += data.iter().flat_map(|i| i.iter()).sum::(); - }); - drop(recv_port); - - df.run_available(); - assert_eq!(sum.borrow().to_owned(), 45); - - df.schedule_subgraph(source); - df.run_available(); - - assert_eq!(sum.borrow().to_owned(), 90); - } -} diff --git a/src/flow/src/compute/render/map.rs b/src/flow/src/compute/render/map.rs deleted file mode 100644 index 77636c5b3aa..00000000000 --- a/src/flow/src/compute/render/map.rs +++ /dev/null @@ -1,425 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::collections::BTreeMap; - -use dfir_rs::scheduled::graph_ext::GraphExt; -use dfir_rs::scheduled::port::{PortCtx, SEND}; -use itertools::Itertools; -use snafu::OptionExt; - -use crate::compute::render::Context; -use crate::compute::state::Scheduler; -use crate::compute::types::{Arranged, Collection, CollectionBundle, ErrCollector, Toff}; -use crate::error::{Error, PlanSnafu}; -use crate::expr::{Batch, EvalError, MapFilterProject, MfpPlan, ScalarExpr}; -use crate::plan::TypedPlan; -use crate::repr::{self, DiffRow, KeyValDiffRow, RelationType, Row}; -use crate::utils::ArrangeHandler; - -impl Context<'_, '_> { - /// Like `render_mfp` but in batch mode - pub fn render_mfp_batch( - &mut self, - input: Box, - mfp: MapFilterProject, - _output_type: &RelationType, - ) -> Result, Error> { - let input = self.render_plan_batch(*input)?; - - let (out_send_port, out_recv_port) = self.df.make_edge::<_, Toff>("mfp_batch"); - - // This closure capture following variables: - let mfp_plan = MfpPlan::create_from(mfp)?; - - let err_collector = self.err_collector.clone(); - - // TODO(discord9): better way to schedule future run - let scheduler = self.compute_state.get_scheduler(); - - let subgraph = self.df.add_subgraph_in_out( - "mfp_batch", - input.collection.into_inner(), - out_send_port, - move |_ctx, recv, send| { - // mfp only need to passively receive updates from recvs - let src_data = recv.take_inner().into_iter().flat_map(|v| v.into_iter()); - - let output_batches = src_data - .filter_map(|mut input_batch| { - err_collector.run(|| { - let res_batch = mfp_plan.mfp.eval_batch_into(&mut input_batch)?; - Ok(res_batch) - }) - }) - .collect_vec(); - - send.give(output_batches); - }, - ); - - // register current subgraph in scheduler for future scheduling - scheduler.set_cur_subgraph(subgraph); - - let bundle = - CollectionBundle::from_collection(Collection::::from_port(out_recv_port)); - Ok(bundle) - } - - /// render MapFilterProject, will only emit the `rows` once. Assume all incoming row's sys time being `now`` and ignore the row's stated sys time - /// TODO(discord9): schedule mfp operator to run when temporal filter need - /// - /// `MapFilterProject`(`mfp` for short) is scheduled to run when there is enough amount of input updates - /// ***or*** when a future update in it's output buffer(a `Arrangement`) is supposed to emit now. - // There is a false positive in using `Vec` as key due to `Value` have `bytes` variant - #[allow(clippy::mutable_key_type)] - pub fn render_mfp( - &mut self, - input: Box, - mfp: MapFilterProject, - ) -> Result { - let input = self.render_plan(*input)?; - // TODO(discord9): consider if check if contain temporal to determine if - // need arrange or not, or does this added complexity worth it - let (out_send_port, out_recv_port) = self.df.make_edge::<_, Toff>("mfp"); - - let output_arity = mfp.output_arity(); - - // default to have a arrange with only future updates, so it can be empty if no temporal filter is applied - // as stream only sends current updates and etc. - let arrange_handler = self.compute_state.new_arrange(None); - let arrange_handler_inner = - arrange_handler - .clone_future_only() - .with_context(|| PlanSnafu { - reason: "No write is expected at this point", - })?; - - // This closure capture following variables: - let mfp_plan = MfpPlan::create_from(mfp)?; - let now = self.compute_state.current_time_ref(); - - let err_collector = self.err_collector.clone(); - - // TODO(discord9): better way to schedule future run - let scheduler = self.compute_state.get_scheduler(); - let scheduler_inner = scheduler.clone(); - - let subgraph = self.df.add_subgraph_in_out( - "mfp", - input.collection.into_inner(), - out_send_port, - move |_ctx, recv, send| { - // mfp only need to passively receive updates from recvs - let data = recv.take_inner().into_iter().flat_map(|v| v.into_iter()); - - mfp_subgraph( - &arrange_handler_inner, - data, - &mfp_plan, - *now.borrow(), - &err_collector, - &scheduler_inner, - send, - ); - }, - ); - - // register current subgraph in scheduler for future scheduling - scheduler.set_cur_subgraph(subgraph); - - let arranged = BTreeMap::from([( - (0..output_arity).map(ScalarExpr::Column).collect_vec(), - Arranged::new(arrange_handler), - )]); - - let bundle = CollectionBundle { - collection: Collection::from_port(out_recv_port), - arranged, - }; - Ok(bundle) - } -} - -fn mfp_subgraph( - arrange: &ArrangeHandler, - input: impl IntoIterator, - mfp_plan: &MfpPlan, - now: repr::Timestamp, - err_collector: &ErrCollector, - scheduler: &Scheduler, - send: &PortCtx, -) { - // all updates that should be send immediately - let mut output_now = vec![]; - let run_mfp = || { - let mut all_updates = eval_mfp_core(input, mfp_plan, now, err_collector); - all_updates.retain(|(kv, ts, d)| { - if *ts > now { - true - } else { - output_now.push((kv.clone(), *ts, *d)); - false - } - }); - let future_updates = all_updates; - - arrange.write().apply_updates(now, future_updates)?; - Ok(()) - }; - err_collector.run(run_mfp); - - // Deal with output: - // 1. Read all updates that were emitted between the last time this arrangement had updates and the current time. - // 2. Output the updates. - // 3. Truncate all updates within that range. - let from = arrange.read().last_compaction_time(); - let from = from.unwrap_or(repr::Timestamp::MIN); - let range = ( - std::ops::Bound::Excluded(from), - std::ops::Bound::Included(now), - ); - - // find all updates that need to be send from arrangement - let output_kv = arrange.read().get_updates_in_range(range); - - err_collector.run(|| { - snafu::ensure!( - mfp_plan.is_temporal() || output_kv.is_empty(), - crate::expr::error::InternalSnafu { - reason: "Output from future should be empty since temporal filter is not applied" - } - ); - Ok(()) - }); - - // the output is expected to be key -> empty val - let output = output_kv - .into_iter() - .chain(output_now) // chain previous immediately send updates - .map(|((key, _v), ts, diff)| (key, ts, diff)) - .collect_vec(); - // send output - send.give(output); - - let run_compaction = || { - arrange.write().compact_to(now)?; - Ok(()) - }; - err_collector.run(run_compaction); - - // schedule next time this subgraph should run - scheduler.schedule_for_arrange(&arrange.read(), now); -} - -/// The core of evaluating MFP operator, given a MFP and a input, evaluate the MFP operator, -/// return the output updates **And** possibly any number of errors that occurred during the evaluation -fn eval_mfp_core( - input: impl IntoIterator, - mfp_plan: &MfpPlan, - now: repr::Timestamp, - err_collector: &ErrCollector, -) -> Vec { - let mut all_updates = Vec::new(); - for (mut row, _sys_time, diff) in input.into_iter() { - // this updates is expected to be only zero, one or two rows - let updates = mfp_plan.evaluate::(&mut row.inner, now, diff); - // TODO(discord9): refactor error handling - // Expect error in a single row to not interrupt the whole evaluation - let updates = updates - .filter_map(|r| match r { - Ok((key, ts, diff)) => Some(((key, Row::empty()), ts, diff)), - Err((err, _ts, _diff)) => { - err_collector.push_err(err); - None - } - }) - .collect_vec(); - - all_updates.extend(updates); - } - all_updates -} - -#[cfg(test)] -mod test { - - use datatypes::data_type::ConcreteDataType; - use dfir_rs::scheduled::graph::Dfir; - - use super::*; - use crate::compute::render::test::{get_output_handle, harness_test_ctx, run_and_check}; - use crate::compute::state::DataflowState; - use crate::expr::{self, BinaryFunc, GlobalId}; - use crate::plan::Plan; - use crate::repr::{ColumnType, RelationType}; - - /// test if temporal filter works properly - /// namely: if mfp operator can schedule a delete at the correct time - #[test] - fn test_render_mfp_with_temporal() { - let mut df = Dfir::new(); - let mut state = DataflowState::default(); - let mut ctx = harness_test_ctx(&mut df, &mut state); - - let rows = vec![ - (Row::new(vec![1i64.into()]), 0, 1), - (Row::new(vec![2i64.into()]), 0, 1), - (Row::new(vec![3i64.into()]), 0, 1), - ]; - let collection = ctx.render_constant(rows.clone()); - ctx.insert_global(GlobalId::User(1), collection); - let input_plan = Plan::Get { - id: expr::Id::Global(GlobalId::User(1)), - }; - let typ = RelationType::new(vec![ColumnType::new_nullable( - ConcreteDataType::int64_datatype(), - )]); - // temporal filter: now <= col(0) < now + 4 - let mfp = MapFilterProject::new(1) - .filter(vec![ - ScalarExpr::Column(0) - .call_unary(expr::UnaryFunc::Cast( - ConcreteDataType::timestamp_microsecond_datatype(), - )) - .call_binary( - ScalarExpr::CallUnmaterializable(expr::UnmaterializableFunc::Now), - BinaryFunc::Gte, - ), - ScalarExpr::Column(0) - .call_binary( - ScalarExpr::literal(4i64.into(), ConcreteDataType::int64_datatype()), - BinaryFunc::SubInt64, - ) - .call_unary(expr::UnaryFunc::Cast( - ConcreteDataType::timestamp_microsecond_datatype(), - )) - .call_binary( - ScalarExpr::CallUnmaterializable(expr::UnmaterializableFunc::Now), - BinaryFunc::Lt, - ), - ]) - .unwrap(); - - let bundle = ctx - .render_mfp(Box::new(input_plan.with_types(typ.into_unnamed())), mfp) - .unwrap(); - let output = get_output_handle(&mut ctx, bundle); - // drop ctx here to simulate actual process of compile first, run later scenario - drop(ctx); - // expected output at given time - let expected_output = BTreeMap::from([ - ( - 0, // time - vec![ - (Row::new(vec![1i64.into()]), 0, 1), - (Row::new(vec![2i64.into()]), 0, 1), - (Row::new(vec![3i64.into()]), 0, 1), - ], - ), - ( - 2, // time - vec![(Row::new(vec![1i64.into()]), 2, -1)], - ), - ( - 3, // time - vec![(Row::new(vec![2i64.into()]), 3, -1)], - ), - ( - 4, // time - vec![(Row::new(vec![3i64.into()]), 4, -1)], - ), - ]); - run_and_check(&mut state, &mut df, 0..5, expected_output, output); - } - - /// test if mfp operator without temporal filter works properly - /// that is it filter the rows correctly - #[test] - fn test_render_mfp() { - let mut df = Dfir::new(); - let mut state = DataflowState::default(); - let mut ctx = harness_test_ctx(&mut df, &mut state); - - let rows = vec![ - (Row::new(vec![1.into()]), 1, 1), - (Row::new(vec![2.into()]), 2, 1), - (Row::new(vec![3.into()]), 3, 1), - ]; - let collection = ctx.render_constant(rows.clone()); - ctx.insert_global(GlobalId::User(1), collection); - let input_plan = Plan::Get { - id: expr::Id::Global(GlobalId::User(1)), - }; - let typ = RelationType::new(vec![ColumnType::new_nullable( - ConcreteDataType::int64_datatype(), - )]); - // filter: col(0)>1 - let mfp = MapFilterProject::new(1) - .filter(vec![ScalarExpr::Column(0).call_binary( - ScalarExpr::literal(1.into(), ConcreteDataType::int32_datatype()), - BinaryFunc::Gt, - )]) - .unwrap(); - let bundle = ctx - .render_mfp(Box::new(input_plan.with_types(typ.into_unnamed())), mfp) - .unwrap(); - - let output = get_output_handle(&mut ctx, bundle); - drop(ctx); - let expected = BTreeMap::from([ - (2, vec![(Row::new(vec![2.into()]), 2, 1)]), - (3, vec![(Row::new(vec![3.into()]), 3, 1)]), - ]); - run_and_check(&mut state, &mut df, 1..5, expected, output); - } - - /// test if mfp operator can run multiple times within same tick - #[test] - fn test_render_mfp_multiple_times() { - let mut df = Dfir::new(); - let mut state = DataflowState::default(); - let mut ctx = harness_test_ctx(&mut df, &mut state); - - let (sender, recv) = tokio::sync::broadcast::channel(1000); - let collection = ctx.render_source(recv).unwrap(); - ctx.insert_global(GlobalId::User(1), collection); - let input_plan = Plan::Get { - id: expr::Id::Global(GlobalId::User(1)), - }; - let typ = RelationType::new(vec![ColumnType::new_nullable( - ConcreteDataType::int64_datatype(), - )]); - // filter: col(0)>1 - let mfp = MapFilterProject::new(1) - .filter(vec![ScalarExpr::Column(0).call_binary( - ScalarExpr::literal(1.into(), ConcreteDataType::int32_datatype()), - BinaryFunc::Gt, - )]) - .unwrap(); - let bundle = ctx - .render_mfp(Box::new(input_plan.with_types(typ.into_unnamed())), mfp) - .unwrap(); - - let output = get_output_handle(&mut ctx, bundle); - drop(ctx); - sender.send((Row::new(vec![2.into()]), 0, 1)).unwrap(); - state.run_available_with_schedule(&mut df); - assert_eq!(output.borrow().len(), 1); - output.borrow_mut().clear(); - sender.send((Row::new(vec![3.into()]), 0, 1)).unwrap(); - state.run_available_with_schedule(&mut df); - assert_eq!(output.borrow().len(), 1); - } -} diff --git a/src/flow/src/compute/render/reduce.rs b/src/flow/src/compute/render/reduce.rs deleted file mode 100644 index e8b69af1283..00000000000 --- a/src/flow/src/compute/render/reduce.rs +++ /dev/null @@ -1,1985 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::collections::{BTreeMap, BTreeSet}; -use std::ops::Range; -use std::sync::Arc; - -use arrow::array::new_null_array; -use common_telemetry::trace; -use datatypes::data_type::ConcreteDataType; -use datatypes::prelude::DataType; -use datatypes::value::{ListValue, Value}; -use datatypes::vectors::{BooleanVector, NullVector}; -use dfir_rs::scheduled::graph_ext::GraphExt; -use itertools::Itertools; -use snafu::{OptionExt, ResultExt, ensure}; - -use crate::compute::render::{Context, SubgraphArg}; -use crate::compute::types::{Arranged, Collection, CollectionBundle, ErrCollector, Toff}; -use crate::error::{Error, NotImplementedSnafu, PlanSnafu}; -use crate::expr::error::{ArrowSnafu, DataAlreadyExpiredSnafu, DataTypeSnafu, InternalSnafu}; -use crate::expr::{Accum, Accumulator, Batch, EvalError, ScalarExpr, VectorDiff}; -use crate::plan::{AccumulablePlan, AggrWithIndex, KeyValPlan, ReducePlan, TypedPlan}; -use crate::repr::{self, DiffRow, KeyValDiffRow, RelationType, Row}; -use crate::utils::{ArrangeHandler, ArrangeReader, ArrangeWriter, KeyExpiryManager}; - -impl Context<'_, '_> { - const REDUCE_BATCH: &'static str = "reduce_batch"; - /// Like `render_reduce`, but for batch mode, and only barebone implementation - /// no support for distinct aggregation for now - // There is a false positive in using `Vec` as key due to `Value` have `bytes` variant - #[allow(clippy::mutable_key_type)] - pub fn render_reduce_batch( - &mut self, - input: Box, - key_val_plan: &KeyValPlan, - reduce_plan: &ReducePlan, - output_type: &RelationType, - ) -> Result, Error> { - let accum_plan = if let ReducePlan::Accumulable(accum_plan) = reduce_plan { - if !accum_plan.distinct_aggrs.is_empty() { - NotImplementedSnafu { - reason: "Distinct aggregation is not supported in batch mode", - } - .fail()? - } - accum_plan.clone() - } else { - NotImplementedSnafu { - reason: "Only accumulable reduce plan is supported in batch mode", - } - .fail()? - }; - - let input = self.render_plan_batch(*input)?; - - // first assembly key&val to separate key and val columns(since this is batch mode) - // Then stream kvs through a reduce operator - - // the output is concat from key and val - let output_key_arity = key_val_plan.key_plan.output_arity(); - - // TODO(discord9): config global expire time from self - let arrange_handler = self.compute_state.new_arrange(None); - - if let (Some(time_index), Some(expire_after)) = - (output_type.time_index, self.compute_state.expire_after()) - { - let expire_man = - KeyExpiryManager::new(Some(expire_after), Some(ScalarExpr::Column(time_index))); - arrange_handler.write().set_expire_state(expire_man); - } - - // reduce need full arrangement to be able to query all keys - let arrange_handler_inner = arrange_handler.clone_full_arrange().context(PlanSnafu { - reason: "No write is expected at this point", - })?; - let key_val_plan = key_val_plan.clone(); - - let output_type = output_type.clone(); - - let now = self.compute_state.current_time_ref(); - - let err_collector = self.err_collector.clone(); - - // TODO(discord9): better way to schedule future run - let scheduler = self.compute_state.get_scheduler(); - - let scheduler_inner = scheduler.clone(); - - let (out_send_port, out_recv_port) = - self.df.make_edge::<_, Toff>(Self::REDUCE_BATCH); - - let subgraph = self.df.add_subgraph_in_out( - Self::REDUCE_BATCH, - input.collection.into_inner(), - out_send_port, - move |_ctx, recv, send| { - let now = *(now.borrow()); - let arrange = arrange_handler_inner.clone(); - // mfp only need to passively receive updates from recvs - let src_data = recv - .take_inner() - .into_iter() - .flat_map(|v| v.into_iter()) - .collect_vec(); - - reduce_batch_subgraph( - &arrange, - src_data, - &key_val_plan, - &accum_plan, - &output_type, - SubgraphArg { - now, - err_collector: &err_collector, - scheduler: &scheduler_inner, - send, - }, - ) - }, - ); - - scheduler.set_cur_subgraph(subgraph); - - // by default the key of output arrange - let arranged = BTreeMap::from([( - (0..output_key_arity).map(ScalarExpr::Column).collect_vec(), - Arranged::new(arrange_handler), - )]); - - let bundle = CollectionBundle { - collection: Collection::from_port(out_recv_port), - arranged, - }; - Ok(bundle) - } - - const REDUCE: &'static str = "reduce"; - /// render `Plan::Reduce` into executable dataflow - // There is a false positive in using `Vec` as key due to `Value` have `bytes` variant - #[allow(clippy::mutable_key_type)] - pub fn render_reduce( - &mut self, - input: Box, - key_val_plan: KeyValPlan, - reduce_plan: ReducePlan, - output_type: RelationType, - ) -> Result { - let input = self.render_plan(*input)?; - // first assembly key&val that's ((Row, Row), tick, diff) - // Then stream kvs through a reduce operator - - // the output is concat from key and val - let output_key_arity = key_val_plan.key_plan.output_arity(); - - // TODO(discord9): config global expire time from self - let arrange_handler = self.compute_state.new_arrange(None); - - if let (Some(time_index), Some(expire_after)) = - (output_type.time_index, self.compute_state.expire_after()) - { - let expire_man = - KeyExpiryManager::new(Some(expire_after), Some(ScalarExpr::Column(time_index))); - arrange_handler.write().set_expire_state(expire_man); - } - - // reduce need full arrangement to be able to query all keys - let arrange_handler_inner = arrange_handler.clone_full_arrange().context(PlanSnafu { - reason: "No write is expected at this point", - })?; - - let distinct_input = self.add_accum_distinct_input_arrange(&reduce_plan); - - let reduce_arrange = ReduceArrange { - output_arrange: arrange_handler_inner, - distinct_input, - }; - - let now = self.compute_state.current_time_ref(); - - let err_collector = self.err_collector.clone(); - - // TODO(discord9): better way to schedule future run - let scheduler = self.compute_state.get_scheduler(); - let scheduler_inner = scheduler.clone(); - - let (out_send_port, out_recv_port) = self.df.make_edge::<_, Toff>(Self::REDUCE); - - let subgraph = self.df.add_subgraph_in_out( - Self::REDUCE, - input.collection.into_inner(), - out_send_port, - move |_ctx, recv, send| { - // mfp only need to passively receive updates from recvs - let data = recv - .take_inner() - .into_iter() - .flat_map(|v| v.into_iter()) - .collect_vec(); - - reduce_subgraph( - &reduce_arrange, - data, - &key_val_plan, - &reduce_plan, - SubgraphArg { - now: *now.borrow(), - err_collector: &err_collector, - scheduler: &scheduler_inner, - send, - }, - ); - }, - ); - - scheduler.set_cur_subgraph(subgraph); - - // by default the key of output arrange - let arranged = BTreeMap::from([( - (0..output_key_arity).map(ScalarExpr::Column).collect_vec(), - Arranged::new(arrange_handler), - )]); - - let bundle = CollectionBundle { - collection: Collection::from_port(out_recv_port), - arranged, - }; - Ok(bundle) - } - - /// Contrast to it name, it's for adding distinct input for - /// accumulable reduce plan with distinct input, - /// like `select COUNT(DISTINCT col) from table` - /// - /// The return value is optional a list of arrangement, which is created for distinct input, and should be the - /// same length as the distinct aggregation in accumulable reduce plan - fn add_accum_distinct_input_arrange( - &mut self, - reduce_plan: &ReducePlan, - ) -> Option> { - match reduce_plan { - ReducePlan::Distinct => None, - ReducePlan::Accumulable(AccumulablePlan { distinct_aggrs, .. }) => { - (!distinct_aggrs.is_empty()).then(|| { - std::iter::repeat_with(|| { - let arr = self.compute_state.new_arrange(None); - arr.set_full_arrangement(true); - arr - }) - .take(distinct_aggrs.len()) - .collect() - }) - } - } - } -} - -fn from_accum_values_to_live_accums( - accums: Vec, - len: usize, -) -> Result>, EvalError> { - let accum_ranges = from_val_to_slice_idx(accums.first().cloned(), len)?; - let mut accum_list = vec![]; - for range in accum_ranges.iter() { - accum_list.push(accums.get(range.clone()).unwrap_or_default().to_vec()); - } - Ok(accum_list) -} - -/// All arrange(aka state) used in reduce operator -pub struct ReduceArrange { - /// The output arrange of reduce operator - output_arrange: ArrangeHandler, - /// The distinct input arrangement for accumulable reduce plan - /// only used when accumulable reduce plan has distinct aggregation - distinct_input: Option>, -} - -fn batch_split_by_key_val( - batch: &Batch, - key_val_plan: &KeyValPlan, - err_collector: &ErrCollector, -) -> (Batch, Batch) { - let row_count = batch.row_count(); - let mut key_batch = Batch::empty(); - let mut val_batch = Batch::empty(); - - err_collector.run(|| { - if key_val_plan.key_plan.output_arity() != 0 { - key_batch = key_val_plan.key_plan.eval_batch_into(&mut batch.clone())?; - } - - if key_val_plan.val_plan.output_arity() != 0 { - val_batch = key_val_plan.val_plan.eval_batch_into(&mut batch.clone())?; - } - Ok(()) - }); - - // deal with empty key or val - if key_batch.row_count() == 0 && key_batch.column_count() == 0 { - key_batch.set_row_count(row_count); - } - - if val_batch.row_count() == 0 && val_batch.column_count() == 0 { - val_batch.set_row_count(row_count); - } - - (key_batch, val_batch) -} - -/// split a row into key and val by evaluate the key and val plan -fn split_rows_to_key_val( - rows: impl IntoIterator, - key_val_plan: KeyValPlan, - err_collector: ErrCollector, -) -> impl IntoIterator { - let mut row_buf = Row::new(vec![]); - rows.into_iter().filter_map( - move |(mut row, sys_time, diff): DiffRow| -> Option { - err_collector.run(|| { - let len = row.len(); - if let Some(key) = key_val_plan - .key_plan - .evaluate_into(&mut row.inner, &mut row_buf)? - { - // reuse the row as buffer - row.inner.resize(len, Value::Null); - // val_plan is not supported to carry any filter predicate, - let val = key_val_plan - .val_plan - .evaluate_into(&mut row.inner, &mut row_buf)? - .context(InternalSnafu { - reason: "val_plan should not contain any filter predicate", - })?; - Ok(Some(((key, val), sys_time, diff))) - } else { - Ok(None) - } - })? - }, - ) -} - -fn reduce_batch_subgraph( - arrange: &ArrangeHandler, - src_data: impl IntoIterator, - key_val_plan: &KeyValPlan, - accum_plan: &AccumulablePlan, - output_type: &RelationType, - SubgraphArg { - now, - err_collector, - scheduler: _, - send, - }: SubgraphArg>, -) { - let mut key_to_many_vals = BTreeMap::>::new(); - let mut input_row_count = 0; - let mut input_batch_count = 0; - - for batch in src_data { - input_batch_count += 1; - input_row_count += batch.row_count(); - err_collector.run(|| { - let (key_batch, val_batch) = - batch_split_by_key_val(&batch, key_val_plan, err_collector); - ensure!( - key_batch.row_count() == val_batch.row_count(), - InternalSnafu { - reason: format!( - "Key and val batch should have the same row count, found {} and {}", - key_batch.row_count(), - val_batch.row_count() - ) - } - ); - - let mut distinct_keys = BTreeSet::new(); - for row_idx in 0..key_batch.row_count() { - let key_row = key_batch.get_row(row_idx)?; - let key_row = Row::new(key_row); - - if distinct_keys.contains(&key_row) { - continue; - } else { - distinct_keys.insert(key_row.clone()); - } - } - - let key_data_types = output_type - .column_types - .iter() - .map(|t| t.scalar_type.clone()) - .collect_vec(); - - // TODO(discord9): here reduce numbers of eq to minimal by keeping slicing key/val batch - for key_row in distinct_keys { - let key_scalar_value = { - let mut key_scalar_value = Vec::with_capacity(key_row.len()); - for (key_idx, key) in key_row.iter().enumerate() { - let v = - key.try_to_scalar_value(&key.data_type()) - .context(DataTypeSnafu { - msg: "can't convert key values to datafusion value", - })?; - - let key_data_type = key_data_types.get(key_idx).context(InternalSnafu { - reason: format!( - "Key index out of bound, expected at most {} but got {}", - output_type.column_types.len(), - key_idx - ), - })?; - - // if incoming value's datatype is null, it need to be handled specially, see below - if key_data_type.as_arrow_type() != v.data_type() - && !v.data_type().is_null() - { - crate::expr::error::InternalSnafu { - reason: format!( - "Key data type mismatch, expected {:?} but got {:?}", - key_data_type.as_arrow_type(), - v.data_type() - ), - } - .fail()? - } - - // handle single null key - let arrow_value = if v.data_type().is_null() { - let ret = new_null_array(&arrow::datatypes::DataType::Null, 1); - arrow::array::Scalar::new(ret) - } else { - v.to_scalar().context(crate::expr::error::DatafusionSnafu { - context: "can't convert key values to arrow value", - })? - }; - key_scalar_value.push(arrow_value); - } - key_scalar_value - }; - - // first compute equal from separate columns - let eq_results = key_scalar_value - .into_iter() - .zip(key_batch.batch().iter()) - .map(|(key, col)| { - // TODO(discord9): this takes half of the cpu! And this is redundant amount of `eq`! - - // note that if lhs is a null, we still need to get all rows that are null! But can't use `eq` since - // it will return null if input have null, so we need to use `is_null` instead - if arrow::array::Datum::get(&key).0.data_type().is_null() { - arrow::compute::kernels::boolean::is_null( - col.to_arrow_array().as_ref() as _ - ) - } else { - arrow::compute::kernels::cmp::eq( - &key, - &col.to_arrow_array().as_ref() as _, - ) - } - }) - .try_collect::<_, Vec<_>, _>() - .context(ArrowSnafu { - context: "Failed to compare key values", - })?; - - // then combine all equal results to finally found equal key rows - let opt_eq_mask = eq_results - .into_iter() - .fold(None, |acc, v| match acc { - Some(Ok(acc)) => Some(arrow::compute::kernels::boolean::and(&acc, &v)), - Some(Err(_)) => acc, - None => Some(Ok(v)), - }) - .transpose() - .context(ArrowSnafu { - context: "Failed to combine key comparison results", - })?; - - let key_eq_mask = if let Some(eq_mask) = opt_eq_mask { - BooleanVector::from(eq_mask) - } else { - // if None, meaning key_batch's column number is zero, which means - // the key is empty, so we just return a mask of all true - // meaning taking all values - BooleanVector::from(vec![true; key_batch.row_count()]) - }; - // TODO: both slice and mutate remaining batch - - let cur_val_batch = val_batch.filter(&key_eq_mask)?; - - key_to_many_vals - .entry(key_row) - .or_default() - .push(cur_val_batch); - } - - Ok(()) - }); - } - - trace!( - "Reduce take {} batches, {} rows", - input_batch_count, input_row_count - ); - - // write lock the arrange for the rest of the function body - // to prevent wired race condition - let mut arrange = arrange.write(); - let mut all_arrange_updates = Vec::with_capacity(key_to_many_vals.len()); - - let mut all_output_dict = BTreeMap::new(); - - for (key, val_batches) in key_to_many_vals { - err_collector.run(|| -> Result<(), _> { - let (accums, _, _) = arrange.get(now, &key).unwrap_or_default(); - let accum_list = - from_accum_values_to_live_accums(accums.unpack(), accum_plan.simple_aggrs.len())?; - - let mut accum_output = AccumOutput::new(); - for AggrWithIndex { - expr, - input_idx, - output_idx, - } in accum_plan.simple_aggrs.iter() - { - let cur_accum_value = accum_list.get(*output_idx).cloned().unwrap_or_default(); - let mut cur_accum = if cur_accum_value.is_empty() { - Accum::new_accum(&expr.func.clone())? - } else { - Accum::try_into_accum(&expr.func, cur_accum_value)? - }; - - for val_batch in val_batches.iter() { - // if batch is empty, input null instead - let cur_input = val_batch - .batch() - .get(*input_idx) - .cloned() - .unwrap_or_else(|| Arc::new(NullVector::new(val_batch.row_count()))); - let len = cur_input.len(); - cur_accum.update_batch(&expr.func, VectorDiff::from(cur_input))?; - - trace!("Reduce accum after take {} rows: {:?}", len, cur_accum); - } - let final_output = cur_accum.eval(&expr.func)?; - trace!("Reduce accum final output: {:?}", final_output); - accum_output.insert_output(*output_idx, final_output); - - let cur_accum_value = cur_accum.into_state(); - accum_output.insert_accum(*output_idx, cur_accum_value); - } - - let (new_accums, res_val_row) = accum_output.into_accum_output()?; - - let arrange_update = ((key.clone(), Row::new(new_accums)), now, 1); - all_arrange_updates.push(arrange_update); - - all_output_dict.insert(key, Row::from(res_val_row)); - - Ok(()) - }); - } - - err_collector.run(|| { - arrange.apply_updates(now, all_arrange_updates)?; - arrange.compact_to(now) - }); - // release the lock - drop(arrange); - - // this output part is not supposed to be resource intensive - // (because for every batch there wouldn't usually be as many output row?), - // so we can do some costly operation here - let output_types = output_type - .column_types - .iter() - .map(|t| t.scalar_type.clone()) - .collect_vec(); - - err_collector.run(|| { - let column_cnt = output_types.len(); - let row_cnt = all_output_dict.len(); - - let mut output_builder = output_types - .into_iter() - .map(|t| t.create_mutable_vector(row_cnt)) - .collect_vec(); - - for (key, val) in all_output_dict { - for (i, v) in key.into_iter().chain(val.into_iter()).enumerate() { - output_builder - .get_mut(i) - .context(InternalSnafu{ - reason: format!( - "Output builder should have the same length as the row, expected at most {} but got {}", - column_cnt - 1, - i - ) - })? - .try_push_value_ref(&v.as_value_ref()) - .context(DataTypeSnafu { - msg: "Failed to push value", - })?; - } - } - - let output_columns = output_builder - .into_iter() - .map(|mut b| b.to_vector()) - .collect_vec(); - - let output_batch = Batch::try_new(output_columns, row_cnt)?; - - trace!("Reduce output batch: {:?}", output_batch); - - send.give(vec![output_batch]); - - Ok(()) - }); -} - -/// reduce subgraph, reduce the input data into a single row -/// output is concat from key and val -fn reduce_subgraph( - ReduceArrange { - output_arrange: arrange, - distinct_input, - }: &ReduceArrange, - data: impl IntoIterator, - key_val_plan: &KeyValPlan, - reduce_plan: &ReducePlan, - SubgraphArg { - now, - err_collector, - scheduler, - send, - }: SubgraphArg, -) { - let key_val = split_rows_to_key_val(data, key_val_plan.clone(), err_collector.clone()); - // from here for distinct reduce and accum reduce, things are drastically different - // for distinct reduce the arrange store the output, - // but for accum reduce the arrange store the accum state, and output is - // evaluated from the accum state if there is need to update - match reduce_plan { - ReducePlan::Distinct => reduce_distinct_subgraph( - arrange, - key_val, - SubgraphArg { - now, - err_collector, - scheduler, - send, - }, - ), - ReducePlan::Accumulable(accum_plan) => reduce_accum_subgraph( - arrange, - distinct_input, - key_val, - accum_plan, - SubgraphArg { - now, - err_collector, - scheduler, - send, - }, - ), - }; -} - -/// return distinct rows(distinct by row's key) from the input, but do not update the arrangement -/// -/// if the same key already exist, we only preserve the oldest value(It make sense for distinct input over key) -fn eval_distinct_core( - arrange: ArrangeReader, - kv: impl IntoIterator, - now: repr::Timestamp, - err_collector: &ErrCollector, -) -> Vec { - let _ = err_collector; - - // note that we also need to keep track of the distinct rows inside the current input - // hence the `inner_map` to keeping track of the distinct rows - let mut inner_map = BTreeMap::new(); - kv.into_iter() - .filter_map(|((key, val), ts, diff)| { - // first check inner_map, then check the arrangement to make sure getting the newest value - let old_val = inner_map - .get(&key) - .cloned() - .or_else(|| arrange.get(now, &key)); - - let new_key_val = match (old_val, diff) { - // a new distinct row - (None, 1) => Some(((key, val), ts, diff)), - // if diff from newest value, also do update - (Some(old_val), diff) if old_val.0 == val && old_val.2 != diff => { - Some(((key, val), ts, diff)) - } - _ => None, - }; - - if let Some(((k, v), t, d)) = new_key_val.clone() { - // update the inner_map, so later updates can be checked against it - inner_map.insert(k, (v, t, d)); - } - new_key_val - }) - .collect_vec() -} - -/// eval distinct reduce plan, output the distinct, and update the arrangement -/// -/// This function is extracted because also want to use it to update distinct input of accumulable reduce plan -fn update_reduce_distinct_arrange( - arrange: &ArrangeHandler, - kv: impl IntoIterator, - now: repr::Timestamp, - err_collector: &ErrCollector, -) -> impl Iterator { - let result_updates = eval_distinct_core(arrange.read(), kv, now, err_collector); - - err_collector.run(|| { - arrange.write().apply_updates(now, result_updates)?; - Ok(()) - }); - - // Deal with output: - - // 1. Read all updates that were emitted between the last time this arrangement had updates and the current time. - let from = arrange.read().last_compaction_time(); - let from = from.unwrap_or(repr::Timestamp::MIN); - let range = ( - std::ops::Bound::Excluded(from), - std::ops::Bound::Included(now), - ); - let output_kv = arrange.read().get_updates_in_range(range); - - // 2. Truncate all updates stored in arrangement within that range. - let run_compaction = || { - arrange.write().compact_to(now)?; - Ok(()) - }; - err_collector.run(run_compaction); - - // 3. Output the updates. - // output is concat from key and val - output_kv.into_iter().map(|((mut key, v), ts, diff)| { - key.extend(v.into_iter()); - (key, ts, diff) - }) -} - -/// eval distinct reduce plan, output the distinct, and update the arrangement -/// -/// invariant: it'is assumed `kv`'s time is always <= now, -/// since it's from a Collection Bundle, where future inserts are stored in arrange -fn reduce_distinct_subgraph( - arrange: &ArrangeHandler, - kv: impl IntoIterator, - SubgraphArg { - now, - err_collector, - scheduler: _, - send, - }: SubgraphArg, -) { - let ret = update_reduce_distinct_arrange(arrange, kv, now, err_collector).collect_vec(); - - // no future updates should exist here - if arrange.read().get_next_update_time(&now).is_some() { - err_collector.push_err( - InternalSnafu { - reason: "No future updates should exist in the reduce distinct arrangement", - } - .build(), - ); - } - - send.give(ret); -} - -/// eval accumulable reduce plan by eval aggregate function and reduce the result -/// -/// TODO(discord9): eval distinct by adding distinct input arrangement -/// -/// invariant: it'is assumed `kv`'s time is always <= now, -/// since it's from a Collection Bundle, where future inserts are stored in arrange -/// -/// the data being send is just new rows that represent the new output after given input is processed -/// -/// i.e: for example before new updates comes in, the output of query `SELECT sum(number), count(number) FROM table` -/// is (10,2(), and after new updates comes in, the output is (15,3), then the new row being send is ((15, 3), now, 1) -/// -/// while it will also update key -> accums's value, for example if it is empty before, it will become something like -/// |offset| accum for sum | accum for count | -/// where offset is a single value holding the end offset of each accumulator -/// and the rest is the actual accumulator values which could be multiple values -fn reduce_accum_subgraph( - arrange: &ArrangeHandler, - distinct_input: &Option>, - kv: impl IntoIterator, - accum_plan: &AccumulablePlan, - SubgraphArg { - now, - err_collector, - scheduler, - send, - }: SubgraphArg, -) { - let AccumulablePlan { - full_aggrs, - simple_aggrs, - distinct_aggrs, - } = accum_plan; - let mut key_to_vals = BTreeMap::>::new(); - - for ((key, val), _tick, diff) in kv { - // it is assumed that value is in order of insertion - let vals = key_to_vals.entry(key).or_default(); - vals.push((val, diff)); - } - - let mut all_updates = Vec::with_capacity(key_to_vals.len()); - let mut all_outputs = Vec::with_capacity(key_to_vals.len()); - // lock the arrange for write for the rest of function body - // so to prevent wired race condition since we are going to update the arrangement by write after read - // TODO(discord9): consider key-based lock - let mut arrange = arrange.write(); - for (key, value_diffs) in key_to_vals { - if let Some(expire_man) = &arrange.get_expire_state() { - let mut is_expired = false; - err_collector.run(|| { - if let Some(expired) = expire_man.get_expire_duration(now, &key)? { - is_expired = true; - // expired data is ignored in computation, and a simple warning is logged - common_telemetry::warn!( - "Data already expired: {}", - DataAlreadyExpiredSnafu { - expired_by: expired, - } - .build() - ); - Ok(()) - } else { - Ok(()) - } - }); - if is_expired { - // errors already collected, we can just continue to next key - continue; - } - } - let col_diffs = { - let row_len = value_diffs[0].0.len(); - let res = err_collector.run(|| get_col_diffs(value_diffs, row_len)); - match res { - Some(res) => res, - // TODO(discord9): consider better error handling other than - // just skip the row and logging error - None => continue, - } - }; - let (accums, _, _) = arrange.get(now, &key).unwrap_or_default(); - - let accums = accums.inner; - - // deser accums from offsets - let accum_ranges = { - let res = err_collector - .run(|| from_val_to_slice_idx(accums.first().cloned(), full_aggrs.len())); - if let Some(res) = res { - res - } else { - // errors is collected, we can just continue and send error back through `err_collector` - continue; - } - }; - - let mut accum_output = AccumOutput::new(); - eval_simple_aggrs( - simple_aggrs, - &accums, - &accum_ranges, - &col_diffs, - &mut accum_output, - err_collector, - ); - - // for distinct input - eval_distinct_aggrs( - distinct_aggrs, - distinct_input, - &accums, - &accum_ranges, - &col_diffs, - &mut accum_output, - SubgraphArg { - now, - err_collector, - scheduler, - send, - }, - ); - - // get and append results - err_collector.run(|| { - let (new_accums, res_val_row) = accum_output.into_accum_output()?; - - // construct the updates and save it - all_updates.push(((key.clone(), Row::new(new_accums)), now, 1)); - let mut key_val = key; - key_val.extend(res_val_row); - all_outputs.push((key_val, now, 1)); - Ok(()) - }); - } - err_collector.run(|| { - arrange.apply_updates(now, all_updates)?; - arrange.compact_to(now) - }); - - // for all arranges involved, schedule next time this subgraph should run - // no future updates should exist here - let all_arrange_used = distinct_input - .iter() - .flatten() - .map(|d| d.write()) - .chain(std::iter::once(arrange)); - check_no_future_updates(all_arrange_used, err_collector, now); - - send.give(all_outputs); -} - -fn get_col_diffs( - value_diffs: Vec<(Row, repr::Diff)>, - row_len: usize, -) -> Result>, EvalError> { - ensure!( - value_diffs.iter().all(|(row, _)| row.len() == row_len), - InternalSnafu { - reason: "value_diffs should have rows with equal length" - } - ); - let ret = (0..row_len) - .map(|i| { - value_diffs - .iter() - .map(|(row, diff)| (row.get(i).cloned().unwrap(), *diff)) - .collect_vec() - }) - .collect_vec(); - Ok(ret) -} - -/// Eval simple aggregate functions with no distinct input -fn eval_simple_aggrs( - simple_aggrs: &Vec, - accums: &[Value], - accum_ranges: &[Range], - col_diffs: &[Vec<(Value, i64)>], - accum_output: &mut AccumOutput, - err_collector: &ErrCollector, -) { - for AggrWithIndex { - expr, - input_idx, - output_idx, - } in simple_aggrs - { - let cur_accum_range = accum_ranges[*output_idx].clone(); // range of current accum - let cur_old_accum = accums - .get(cur_accum_range) - .unwrap_or_default() - .iter() - .cloned(); - let cur_col_diff = col_diffs[*input_idx].iter().cloned(); - - // actual eval aggregation function - if let Some((res, new_accum)) = - err_collector.run(|| expr.func.eval_diff_accumulable(cur_old_accum, cur_col_diff)) - { - accum_output.insert_accum(*output_idx, new_accum); - accum_output.insert_output(*output_idx, res); - } // else just collect error and continue - } -} - -/// Accumulate the output of aggregation functions -/// -/// The accum is a map from index to the accumulator of the aggregation function -/// -/// The output is a map from index to the output of the aggregation function -#[derive(Debug)] -struct AccumOutput { - accum: BTreeMap>, - output: BTreeMap, -} - -impl AccumOutput { - fn new() -> Self { - Self { - accum: BTreeMap::new(), - output: BTreeMap::new(), - } - } - - fn insert_accum(&mut self, idx: usize, v: Vec) { - self.accum.insert(idx, v); - } - - fn insert_output(&mut self, idx: usize, v: Value) { - self.output.insert(idx, v); - } - - /// return (accums, output) - fn into_accum_output(self) -> Result<(Vec, Vec), EvalError> { - if self.accum.is_empty() && self.output.is_empty() { - return Ok((vec![], vec![])); - } - ensure!( - !self.accum.is_empty() && self.accum.len() == self.output.len(), - InternalSnafu { - reason: format!( - "Accum and output should have the non-zero and same length, found accum.len() = {}, output.len() = {}", - self.accum.len(), - self.output.len() - ) - } - ); - // make output vec from output map - if let Some(kv) = self.accum.last_key_value() { - ensure!( - *kv.0 == self.accum.len() - 1, - InternalSnafu { - reason: "Accum should be a continuous range" - } - ); - } - if let Some(kv) = self.output.last_key_value() { - ensure!( - *kv.0 == self.output.len() - 1, - InternalSnafu { - reason: "Output should be a continuous range" - } - ); - } - - let accums = self.accum.into_values().collect_vec(); - let new_accums = from_accums_to_offsetted_accum(accums); - let output = self.output.into_values().collect_vec(); - Ok((new_accums, output)) - } -} - -/// Eval distinct aggregate functions with distinct input arrange -fn eval_distinct_aggrs( - distinct_aggrs: &Vec, - distinct_input: &Option>, - accums: &[Value], - accum_ranges: &[Range], - col_diffs: &[Vec<(Value, i64)>], - accum_output: &mut AccumOutput, - SubgraphArg { - now, - err_collector, - scheduler: _, - send: _, - }: SubgraphArg, -) { - for AggrWithIndex { - expr, - input_idx, - output_idx, - } in distinct_aggrs - { - let cur_accum_range = accum_ranges[*output_idx].clone(); // range of current accum - let cur_old_accum = accums - .get(cur_accum_range) - .unwrap_or_default() - .iter() - .cloned(); - let cur_col_diff = col_diffs[*input_idx].iter().cloned(); - // first filter input with distinct - let input_arrange = distinct_input - .as_ref() - .and_then(|v| v[*input_idx].clone_full_arrange()) - .expect("A full distinct input arrangement should exist"); - let kv = cur_col_diff.map(|(v, d)| ((Row::new(vec![v]), Row::empty()), now, d)); - let col_diff_distinct = - update_reduce_distinct_arrange(&input_arrange, kv, now, err_collector).map( - |(row, _ts, diff)| (row.get(0).expect("Row should not be empty").clone(), diff), - ); - let col_diff_distinct = { - let res = col_diff_distinct.collect_vec(); - res.into_iter() - }; - // actual eval aggregation function - let (res, new_accum) = expr - .func - .eval_diff_accumulable(cur_old_accum, col_diff_distinct) - .unwrap(); - accum_output.insert_accum(*output_idx, new_accum); - accum_output.insert_output(*output_idx, res); - } -} - -fn check_no_future_updates<'a>( - all_arrange_used: impl IntoIterator>, - err_collector: &ErrCollector, - now: repr::Timestamp, -) { - for arrange in all_arrange_used { - if arrange.get_next_update_time(&now).is_some() { - err_collector.push_err( - InternalSnafu { - reason: "No future updates should exist in the reduce distinct arrangement", - } - .build(), - ); - } - } -} - -/// convert a list of accumulators to a vector of values with first value as offset of end of each accumulator -fn from_accums_to_offsetted_accum(new_accums: Vec>) -> Vec { - let offset = new_accums - .iter() - .map(|v| v.len() as u64) - .scan(1, |state, x| { - *state += x; - Some(*state) - }) - .map(Value::from) - .collect::>(); - let first = ListValue::new(offset, Arc::new(ConcreteDataType::uint64_datatype())); - let first = Value::List(first); - // construct new_accums - - std::iter::once(first) - .chain(new_accums.into_iter().flatten()) - .collect::>() -} - -/// Convert a value to a list of slice index -fn from_val_to_slice_idx( - value: Option, - expected_len: usize, -) -> Result>, EvalError> { - let offset_end = if let Some(value) = value { - let list = value - .as_list() - .with_context(|_| DataTypeSnafu { - msg: "Accum's first element should be a list", - })? - .context(InternalSnafu { - reason: "Accum's first element should be a list", - })?; - let ret: Vec = list - .items() - .iter() - .map(|v| { - v.as_u64().map(|j| j as usize).context(InternalSnafu { - reason: "End offset should be a list of u64", - }) - }) - .try_collect()?; - ensure!( - ret.len() == expected_len, - InternalSnafu { - reason: "Offset List should have the same length as full_aggrs" - } - ); - Ok(ret) - } else { - Ok(vec![1usize; expected_len]) - }?; - let accum_ranges = (0..expected_len) - .map(|idx| { - if idx == 0 { - // note that the first element is the offset list - debug_assert!( - offset_end[0] >= 1, - "Offset should be at least 1: {:?}", - &offset_end - ); - 1..offset_end[0] - } else { - offset_end[idx - 1]..offset_end[idx] - } - }) - .collect_vec(); - Ok(accum_ranges) -} - -// mainly for reduce's test -// TODO(discord9): add tests for accum ser/de -#[cfg(test)] -mod test { - - use std::time::Duration; - - use common_time::Timestamp; - use datatypes::data_type::{ConcreteDataType, ConcreteDataType as CDT}; - use dfir_rs::scheduled::graph::Dfir; - - use super::*; - use crate::compute::render::test::{get_output_handle, harness_test_ctx, run_and_check}; - use crate::compute::state::DataflowState; - use crate::expr::{ - self, AggregateExpr, AggregateFunc, BinaryFunc, GlobalId, MapFilterProject, UnaryFunc, - }; - use crate::plan::Plan; - use crate::repr::{ColumnType, RelationType}; - - /// SELECT sum(number) FROM numbers_with_ts GROUP BY tumble(ts, '1 second', '2021-07-01 00:00:00') - /// input table columns: number, ts - /// expected: sum(number), window_start, window_end - #[test] - fn test_tumble_group_by() { - let mut df = Dfir::new(); - let mut state = DataflowState::default(); - let mut ctx = harness_test_ctx(&mut df, &mut state); - const START: i64 = 1625097600000; - let rows = vec![ - (1u32, START + 1000), - (2u32, START + 1500), - (3u32, START + 2000), - (1u32, START + 2500), - (2u32, START + 3000), - (3u32, START + 3500), - ]; - let rows = rows - .into_iter() - .map(|(number, ts)| { - ( - Row::new(vec![number.into(), Timestamp::new_millisecond(ts).into()]), - 1, - 1, - ) - }) - .collect_vec(); - - let collection = ctx.render_constant(rows.clone()); - ctx.insert_global(GlobalId::User(1), collection); - - let aggr_expr = AggregateExpr { - func: AggregateFunc::SumUInt32, - expr: ScalarExpr::Column(0), - distinct: false, - }; - let expected = TypedPlan { - schema: RelationType::new(vec![ - ColumnType::new(CDT::uint64_datatype(), true), // sum(number) - ColumnType::new(CDT::timestamp_millisecond_datatype(), false), // window start - ColumnType::new(CDT::timestamp_millisecond_datatype(), false), // window end - ]) - .into_unnamed(), - // TODO(discord9): mfp indirectly ref to key columns - /* - .with_key(vec![1]) - .with_time_index(Some(0)),*/ - plan: Plan::Mfp { - input: Box::new( - Plan::Reduce { - input: Box::new( - Plan::Get { - id: crate::expr::Id::Global(GlobalId::User(1)), - } - .with_types( - RelationType::new(vec![ - ColumnType::new(ConcreteDataType::uint32_datatype(), false), - ColumnType::new( - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ), - ]) - .into_unnamed(), - ), - ), - key_val_plan: KeyValPlan { - key_plan: MapFilterProject::new(2) - .map(vec![ - ScalarExpr::Column(1).call_unary( - UnaryFunc::TumbleWindowFloor { - window_size: Duration::from_nanos(1_000_000_000), - start_time: Some(Timestamp::new_millisecond( - 1625097600000, - )), - }, - ), - ScalarExpr::Column(1).call_unary( - UnaryFunc::TumbleWindowCeiling { - window_size: Duration::from_nanos(1_000_000_000), - start_time: Some(Timestamp::new_millisecond( - 1625097600000, - )), - }, - ), - ]) - .unwrap() - .project(vec![2, 3]) - .unwrap() - .into_safe(), - val_plan: MapFilterProject::new(2) - .project(vec![0, 1]) - .unwrap() - .into_safe(), - }, - reduce_plan: ReducePlan::Accumulable(AccumulablePlan { - full_aggrs: vec![aggr_expr.clone()], - simple_aggrs: vec![AggrWithIndex::new(aggr_expr.clone(), 0, 0)], - distinct_aggrs: vec![], - }), - } - .with_types( - RelationType::new(vec![ - ColumnType::new(CDT::timestamp_millisecond_datatype(), false), // window start - ColumnType::new(CDT::timestamp_millisecond_datatype(), false), // window end - ColumnType::new(CDT::uint64_datatype(), true), //sum(number) - ]) - .with_key(vec![1]) - .with_time_index(Some(0)) - .into_unnamed(), - ), - ), - mfp: MapFilterProject::new(3) - .map(vec![ - ScalarExpr::Column(2), - ScalarExpr::Column(3), - ScalarExpr::Column(0), - ScalarExpr::Column(1), - ]) - .unwrap() - .project(vec![4, 5, 6]) - .unwrap(), - }, - }; - - let bundle = ctx.render_plan(expected).unwrap(); - - let output = get_output_handle(&mut ctx, bundle); - drop(ctx); - let expected = BTreeMap::from([( - 1, - vec![ - ( - Row::new(vec![ - 3u64.into(), - Timestamp::new_millisecond(START + 1000).into(), - Timestamp::new_millisecond(START + 2000).into(), - ]), - 1, - 1, - ), - ( - Row::new(vec![ - 4u64.into(), - Timestamp::new_millisecond(START + 2000).into(), - Timestamp::new_millisecond(START + 3000).into(), - ]), - 1, - 1, - ), - ( - Row::new(vec![ - 5u64.into(), - Timestamp::new_millisecond(START + 3000).into(), - Timestamp::new_millisecond(START + 4000).into(), - ]), - 1, - 1, - ), - ], - )]); - run_and_check(&mut state, &mut df, 1..2, expected, output); - } - - /// select avg(number) from number; - #[test] - fn test_avg_eval() { - let mut df = Dfir::new(); - let mut state = DataflowState::default(); - let mut ctx = harness_test_ctx(&mut df, &mut state); - - let rows = vec![ - (Row::new(vec![1u32.into()]), 1, 1), - (Row::new(vec![2u32.into()]), 1, 1), - (Row::new(vec![3u32.into()]), 1, 1), - (Row::new(vec![1u32.into()]), 1, 1), - (Row::new(vec![2u32.into()]), 1, 1), - (Row::new(vec![3u32.into()]), 1, 1), - ]; - let collection = ctx.render_constant(rows.clone()); - ctx.insert_global(GlobalId::User(1), collection); - - let aggr_exprs = vec![ - AggregateExpr { - func: AggregateFunc::SumUInt32, - expr: ScalarExpr::Column(0), - distinct: false, - }, - AggregateExpr { - func: AggregateFunc::Count, - expr: ScalarExpr::Column(0), - distinct: false, - }, - ]; - let avg_expr = ScalarExpr::If { - cond: Box::new(ScalarExpr::Column(1).call_binary( - ScalarExpr::Literal(Value::from(0u32), CDT::int64_datatype()), - BinaryFunc::NotEq, - )), - then: Box::new(ScalarExpr::Column(0).call_binary( - ScalarExpr::Column(1).call_unary(UnaryFunc::Cast(CDT::uint64_datatype())), - BinaryFunc::DivUInt64, - )), - els: Box::new(ScalarExpr::Literal(Value::Null, CDT::uint64_datatype())), - }; - let expected = TypedPlan { - schema: RelationType::new(vec![ColumnType::new(CDT::uint64_datatype(), true)]) - .into_unnamed(), - plan: Plan::Mfp { - input: Box::new( - Plan::Reduce { - input: Box::new( - Plan::Get { - id: crate::expr::Id::Global(GlobalId::User(1)), - } - .with_types( - RelationType::new(vec![ColumnType::new( - ConcreteDataType::int64_datatype(), - false, - )]) - .into_unnamed(), - ), - ), - key_val_plan: KeyValPlan { - key_plan: MapFilterProject::new(1) - .project(vec![]) - .unwrap() - .into_safe(), - val_plan: MapFilterProject::new(1) - .project(vec![0]) - .unwrap() - .into_safe(), - }, - reduce_plan: ReducePlan::Accumulable(AccumulablePlan { - full_aggrs: aggr_exprs.clone(), - simple_aggrs: vec![ - AggrWithIndex::new(aggr_exprs[0].clone(), 0, 0), - AggrWithIndex::new(aggr_exprs[1].clone(), 0, 1), - ], - distinct_aggrs: vec![], - }), - } - .with_types( - RelationType::new(vec![ - ColumnType::new(ConcreteDataType::uint32_datatype(), true), - ColumnType::new(ConcreteDataType::int64_datatype(), true), - ]) - .into_unnamed(), - ), - ), - mfp: MapFilterProject::new(2) - .map(vec![ - avg_expr, - // TODO(discord9): optimize mfp so to remove indirect ref - ScalarExpr::Column(2), - ]) - .unwrap() - .project(vec![3]) - .unwrap(), - }, - }; - - let bundle = ctx.render_plan(expected).unwrap(); - - let output = get_output_handle(&mut ctx, bundle); - drop(ctx); - let expected = BTreeMap::from([(1, vec![(Row::new(vec![2u64.into()]), 1, 1)])]); - run_and_check(&mut state, &mut df, 1..2, expected, output); - } - - /// SELECT DISTINCT col FROM table - /// - /// table schema: - /// | name | type | - /// |------|-------| - /// | col | Int64 | - #[test] - fn test_basic_distinct() { - let mut df = Dfir::new(); - let mut state = DataflowState::default(); - let mut ctx = harness_test_ctx(&mut df, &mut state); - - let rows = vec![ - (Row::new(vec![1i64.into()]), 1, 1), - (Row::new(vec![2i64.into()]), 2, 1), - (Row::new(vec![3i64.into()]), 3, 1), - (Row::new(vec![1i64.into()]), 4, 1), - (Row::new(vec![2i64.into()]), 5, 1), - (Row::new(vec![3i64.into()]), 6, 1), - ]; - let collection = ctx.render_constant(rows.clone()); - ctx.insert_global(GlobalId::User(1), collection); - let input_plan = Plan::Get { - id: expr::Id::Global(GlobalId::User(1)), - }; - let typ = RelationType::new(vec![ColumnType::new_nullable( - ConcreteDataType::int64_datatype(), - )]); - let key_val_plan = KeyValPlan { - key_plan: MapFilterProject::new(1).project([0]).unwrap().into_safe(), - val_plan: MapFilterProject::new(1).project([]).unwrap().into_safe(), - }; - let reduce_plan = ReducePlan::Distinct; - let bundle = ctx - .render_reduce( - Box::new(input_plan.with_types(typ.into_unnamed())), - key_val_plan, - reduce_plan, - RelationType::empty(), - ) - .unwrap(); - - let output = get_output_handle(&mut ctx, bundle); - drop(ctx); - let expected = BTreeMap::from([( - 6, - vec![ - (Row::new(vec![1i64.into()]), 1, 1), - (Row::new(vec![2i64.into()]), 2, 1), - (Row::new(vec![3i64.into()]), 3, 1), - ], - )]); - run_and_check(&mut state, &mut df, 6..7, expected, output); - } - - /// Batch Mode Reduce Evaluation - /// SELECT SUM(col) FROM table - /// - /// table schema: - /// | name | type | - /// |------|-------| - /// | col | Int64 | - #[test] - fn test_basic_batch_reduce_accum() { - let mut df = Dfir::new(); - let mut state = DataflowState::default(); - let now = state.current_time_ref(); - let mut ctx = harness_test_ctx(&mut df, &mut state); - - let rows = vec![ - (Row::new(vec![Value::Null]), -1, 1), - (Row::new(vec![1i64.into()]), 0, 1), - (Row::new(vec![Value::Null]), 1, 1), - (Row::new(vec![2i64.into()]), 2, 1), - (Row::new(vec![3i64.into()]), 3, 1), - (Row::new(vec![1i64.into()]), 4, 1), - (Row::new(vec![2i64.into()]), 5, 1), - (Row::new(vec![3i64.into()]), 6, 1), - ]; - let input_plan = Plan::Constant { rows: rows.clone() }; - - let typ = RelationType::new(vec![ColumnType::new_nullable( - ConcreteDataType::int64_datatype(), - )]); - let key_val_plan = KeyValPlan { - key_plan: MapFilterProject::new(1).project([]).unwrap().into_safe(), - val_plan: MapFilterProject::new(1).project([0]).unwrap().into_safe(), - }; - - let simple_aggrs = vec![AggrWithIndex::new( - AggregateExpr { - func: AggregateFunc::SumInt64, - expr: ScalarExpr::Column(0), - distinct: false, - }, - 0, - 0, - )]; - let accum_plan = AccumulablePlan { - full_aggrs: vec![AggregateExpr { - func: AggregateFunc::SumInt64, - expr: ScalarExpr::Column(0), - distinct: false, - }], - simple_aggrs, - distinct_aggrs: vec![], - }; - - let reduce_plan = ReducePlan::Accumulable(accum_plan); - let bundle = ctx - .render_reduce_batch( - Box::new(input_plan.with_types(typ.into_unnamed())), - &key_val_plan, - &reduce_plan, - &RelationType::new(vec![ColumnType::new(CDT::int64_datatype(), true)]), - ) - .unwrap(); - - { - let now_inner = now.clone(); - let expected = BTreeMap::>::from([ - (-1, vec![]), - (0, vec![1i64]), - (1, vec![1i64]), - (2, vec![3i64]), - (3, vec![6i64]), - (4, vec![7i64]), - (5, vec![9i64]), - (6, vec![12i64]), - ]); - let collection = bundle.collection; - ctx.df - .add_subgraph_sink("test_sink", collection.into_inner(), move |_ctx, recv| { - let now = *now_inner.borrow(); - let data = recv.take_inner(); - let res = data.into_iter().flat_map(|v| v.into_iter()).collect_vec(); - - if let Some(expected) = expected.get(&now) { - let batch = expected.iter().map(|v| Value::from(*v)).collect_vec(); - let batch = Batch::try_from_rows_with_types( - vec![batch.into()], - &[CDT::int64_datatype()], - ) - .unwrap(); - assert_eq!(res.first(), Some(&batch)); - } - }); - drop(ctx); - - for now in 1..7 { - state.set_current_ts(now); - state.run_available_with_schedule(&mut df); - if !state.get_err_collector().is_empty() { - panic!( - "Errors occur: {:?}", - state.get_err_collector().get_all_blocking() - ) - } - } - } - } - - /// SELECT SUM(col) FROM table - /// - /// table schema: - /// | name | type | - /// |------|-------| - /// | col | Int64 | - #[test] - fn test_basic_reduce_accum() { - let mut df = Dfir::new(); - let mut state = DataflowState::default(); - let mut ctx = harness_test_ctx(&mut df, &mut state); - - let rows = vec![ - (Row::new(vec![1i64.into()]), 1, 1), - (Row::new(vec![2i64.into()]), 2, 1), - (Row::new(vec![3i64.into()]), 3, 1), - (Row::new(vec![1i64.into()]), 4, 1), - (Row::new(vec![2i64.into()]), 5, 1), - (Row::new(vec![3i64.into()]), 6, 1), - ]; - let collection = ctx.render_constant(rows.clone()); - ctx.insert_global(GlobalId::User(1), collection); - let input_plan = Plan::Get { - id: expr::Id::Global(GlobalId::User(1)), - }; - let typ = RelationType::new(vec![ColumnType::new_nullable( - ConcreteDataType::int64_datatype(), - )]); - let key_val_plan = KeyValPlan { - key_plan: MapFilterProject::new(1).project([]).unwrap().into_safe(), - val_plan: MapFilterProject::new(1).project([0]).unwrap().into_safe(), - }; - - let simple_aggrs = vec![AggrWithIndex::new( - AggregateExpr { - func: AggregateFunc::SumInt64, - expr: ScalarExpr::Column(0), - distinct: false, - }, - 0, - 0, - )]; - let accum_plan = AccumulablePlan { - full_aggrs: vec![AggregateExpr { - func: AggregateFunc::SumInt64, - expr: ScalarExpr::Column(0), - distinct: false, - }], - simple_aggrs, - distinct_aggrs: vec![], - }; - - let reduce_plan = ReducePlan::Accumulable(accum_plan); - let bundle = ctx - .render_reduce( - Box::new(input_plan.with_types(typ.into_unnamed())), - key_val_plan, - reduce_plan, - RelationType::empty(), - ) - .unwrap(); - - let output = get_output_handle(&mut ctx, bundle); - drop(ctx); - let expected = BTreeMap::from([ - (1, vec![(Row::new(vec![1i64.into()]), 1, 1)]), - (2, vec![(Row::new(vec![3i64.into()]), 2, 1)]), - (3, vec![(Row::new(vec![6i64.into()]), 3, 1)]), - (4, vec![(Row::new(vec![7i64.into()]), 4, 1)]), - (5, vec![(Row::new(vec![9i64.into()]), 5, 1)]), - (6, vec![(Row::new(vec![12i64.into()]), 6, 1)]), - ]); - run_and_check(&mut state, &mut df, 1..7, expected, output); - } - - /// SELECT SUM(DISTINCT col) FROM table - /// - /// table schema: - /// | name | type | - /// |------|-------| - /// | col | Int64 | - /// - /// this test include even more insert/delete case to cover all case for eval_distinct_core - #[test] - fn test_delete_reduce_distinct_accum() { - let mut df = Dfir::new(); - let mut state = DataflowState::default(); - let mut ctx = harness_test_ctx(&mut df, &mut state); - - let rows = vec![ - // same tick - (Row::new(vec![1i64.into()]), 1, 1), - (Row::new(vec![1i64.into()]), 1, -1), - // next tick - (Row::new(vec![1i64.into()]), 2, 1), - (Row::new(vec![1i64.into()]), 3, -1), - // repeat in same tick - (Row::new(vec![1i64.into()]), 4, 1), - (Row::new(vec![1i64.into()]), 4, -1), - (Row::new(vec![1i64.into()]), 4, 1), - ]; - let collection = ctx.render_constant(rows.clone()); - ctx.insert_global(GlobalId::User(1), collection); - let input_plan = Plan::Get { - id: expr::Id::Global(GlobalId::User(1)), - }; - let typ = RelationType::new(vec![ColumnType::new_nullable( - ConcreteDataType::int64_datatype(), - )]); - let key_val_plan = KeyValPlan { - key_plan: MapFilterProject::new(1).project([]).unwrap().into_safe(), - val_plan: MapFilterProject::new(1).project([0]).unwrap().into_safe(), - }; - - let distinct_aggrs = vec![AggrWithIndex::new( - AggregateExpr { - func: AggregateFunc::SumInt64, - expr: ScalarExpr::Column(0), - distinct: false, - }, - 0, - 0, - )]; - let accum_plan = AccumulablePlan { - full_aggrs: vec![AggregateExpr { - func: AggregateFunc::SumInt64, - expr: ScalarExpr::Column(0), - distinct: true, - }], - simple_aggrs: vec![], - distinct_aggrs, - }; - - let reduce_plan = ReducePlan::Accumulable(accum_plan); - let bundle = ctx - .render_reduce( - Box::new(input_plan.with_types(typ.into_unnamed())), - key_val_plan, - reduce_plan, - RelationType::empty(), - ) - .unwrap(); - - let output = get_output_handle(&mut ctx, bundle); - drop(ctx); - let expected = BTreeMap::from([ - (1, vec![(Row::new(vec![0i64.into()]), 1, 1)]), - (2, vec![(Row::new(vec![1i64.into()]), 2, 1)]), - (3, vec![(Row::new(vec![0i64.into()]), 3, 1)]), - (4, vec![(Row::new(vec![1i64.into()]), 4, 1)]), - ]); - run_and_check(&mut state, &mut df, 1..7, expected, output); - } - - /// SELECT SUM(DISTINCT col) FROM table - /// - /// table schema: - /// | name | type | - /// |------|-------| - /// | col | Int64 | - /// - /// this test include insert and delete which should cover all case for eval_distinct_core - #[test] - fn test_basic_reduce_distinct_accum() { - let mut df = Dfir::new(); - let mut state = DataflowState::default(); - let mut ctx = harness_test_ctx(&mut df, &mut state); - - let rows = vec![ - (Row::new(vec![1i64.into()]), 1, 1), - (Row::new(vec![1i64.into()]), 1, -1), - (Row::new(vec![2i64.into()]), 2, 1), - (Row::new(vec![3i64.into()]), 3, 1), - (Row::new(vec![1i64.into()]), 4, 1), - (Row::new(vec![2i64.into()]), 5, 1), - (Row::new(vec![3i64.into()]), 6, 1), - (Row::new(vec![1i64.into()]), 7, 1), - ]; - let collection = ctx.render_constant(rows.clone()); - ctx.insert_global(GlobalId::User(1), collection); - let input_plan = Plan::Get { - id: expr::Id::Global(GlobalId::User(1)), - }; - let typ = RelationType::new(vec![ColumnType::new_nullable( - ConcreteDataType::int64_datatype(), - )]); - let key_val_plan = KeyValPlan { - key_plan: MapFilterProject::new(1).project([]).unwrap().into_safe(), - val_plan: MapFilterProject::new(1).project([0]).unwrap().into_safe(), - }; - - let distinct_aggrs = vec![AggrWithIndex::new( - AggregateExpr { - func: AggregateFunc::SumInt64, - expr: ScalarExpr::Column(0), - distinct: false, - }, - 0, - 0, - )]; - let accum_plan = AccumulablePlan { - full_aggrs: vec![AggregateExpr { - func: AggregateFunc::SumInt64, - expr: ScalarExpr::Column(0), - distinct: true, - }], - simple_aggrs: vec![], - distinct_aggrs, - }; - - let reduce_plan = ReducePlan::Accumulable(accum_plan); - let bundle = ctx - .render_reduce( - Box::new(input_plan.with_types(typ.into_unnamed())), - key_val_plan, - reduce_plan, - RelationType::empty(), - ) - .unwrap(); - - let output = get_output_handle(&mut ctx, bundle); - drop(ctx); - let expected = BTreeMap::from([ - (1, vec![(Row::new(vec![0i64.into()]), 1, 1)]), - (2, vec![(Row::new(vec![2i64.into()]), 2, 1)]), - (3, vec![(Row::new(vec![5i64.into()]), 3, 1)]), - (4, vec![(Row::new(vec![6i64.into()]), 4, 1)]), - (5, vec![(Row::new(vec![6i64.into()]), 5, 1)]), - (6, vec![(Row::new(vec![6i64.into()]), 6, 1)]), - (7, vec![(Row::new(vec![6i64.into()]), 7, 1)]), - ]); - run_and_check(&mut state, &mut df, 1..7, expected, output); - } - - /// SELECT SUM(col), SUM(DISTINCT col) FROM table - /// - /// table schema: - /// | name | type | - /// |------|-------| - /// | col | Int64 | - #[test] - fn test_composite_reduce_distinct_accum() { - let mut df = Dfir::new(); - let mut state = DataflowState::default(); - let mut ctx = harness_test_ctx(&mut df, &mut state); - - let rows = vec![ - (Row::new(vec![1i64.into()]), 1, 1), - (Row::new(vec![2i64.into()]), 2, 1), - (Row::new(vec![3i64.into()]), 3, 1), - (Row::new(vec![1i64.into()]), 4, 1), - (Row::new(vec![2i64.into()]), 5, 1), - (Row::new(vec![3i64.into()]), 6, 1), - (Row::new(vec![1i64.into()]), 7, 1), - ]; - let collection = ctx.render_constant(rows.clone()); - ctx.insert_global(GlobalId::User(1), collection); - let input_plan = Plan::Get { - id: expr::Id::Global(GlobalId::User(1)), - }; - let typ = RelationType::new(vec![ColumnType::new_nullable( - ConcreteDataType::int64_datatype(), - )]); - let key_val_plan = KeyValPlan { - key_plan: MapFilterProject::new(1).project([]).unwrap().into_safe(), - val_plan: MapFilterProject::new(1).project([0]).unwrap().into_safe(), - }; - let simple_aggrs = vec![AggrWithIndex::new( - AggregateExpr { - func: AggregateFunc::SumInt64, - expr: ScalarExpr::Column(0), - distinct: false, - }, - 0, - 0, - )]; - let distinct_aggrs = vec![AggrWithIndex::new( - AggregateExpr { - func: AggregateFunc::SumInt64, - expr: ScalarExpr::Column(0), - distinct: true, - }, - 0, - 1, - )]; - let accum_plan = AccumulablePlan { - full_aggrs: vec![ - AggregateExpr { - func: AggregateFunc::SumInt64, - expr: ScalarExpr::Column(0), - distinct: false, - }, - AggregateExpr { - func: AggregateFunc::SumInt64, - expr: ScalarExpr::Column(0), - distinct: true, - }, - ], - simple_aggrs, - distinct_aggrs, - }; - - let reduce_plan = ReducePlan::Accumulable(accum_plan); - let bundle = ctx - .render_reduce( - Box::new(input_plan.with_types(typ.into_unnamed())), - key_val_plan, - reduce_plan, - RelationType::empty(), - ) - .unwrap(); - - let output = get_output_handle(&mut ctx, bundle); - drop(ctx); - let expected = BTreeMap::from([ - (1, vec![(Row::new(vec![1i64.into(), 1i64.into()]), 1, 1)]), - (2, vec![(Row::new(vec![3i64.into(), 3i64.into()]), 2, 1)]), - (3, vec![(Row::new(vec![6i64.into(), 6i64.into()]), 3, 1)]), - (4, vec![(Row::new(vec![7i64.into(), 6i64.into()]), 4, 1)]), - (5, vec![(Row::new(vec![9i64.into(), 6i64.into()]), 5, 1)]), - (6, vec![(Row::new(vec![12i64.into(), 6i64.into()]), 6, 1)]), - (7, vec![(Row::new(vec![13i64.into(), 6i64.into()]), 7, 1)]), - ]); - run_and_check(&mut state, &mut df, 1..7, expected, output); - } -} diff --git a/src/flow/src/compute/render/src_sink.rs b/src/flow/src/compute/render/src_sink.rs deleted file mode 100644 index 1bf3699ceb9..00000000000 --- a/src/flow/src/compute/render/src_sink.rs +++ /dev/null @@ -1,245 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! Source and Sink for the dataflow - -use std::collections::BTreeMap; - -use common_telemetry::{debug, trace}; -use dfir_rs::scheduled::graph_ext::GraphExt; -use itertools::Itertools; -use snafu::OptionExt; -use tokio::sync::broadcast::error::TryRecvError; -use tokio::sync::{broadcast, mpsc}; - -use crate::compute::render::Context; -use crate::compute::types::{Arranged, Collection, CollectionBundle, Toff}; -use crate::error::{Error, PlanSnafu}; -use crate::expr::error::InternalSnafu; -use crate::expr::{Batch, EvalError}; -use crate::repr::{DiffRow, Row}; - -#[allow(clippy::mutable_key_type)] -impl Context<'_, '_> { - /// simply send the batch to downstream, without fancy features like buffering - pub fn render_source_batch( - &mut self, - mut src_recv: broadcast::Receiver, - ) -> Result, Error> { - debug!("Rendering Source Batch"); - let (send_port, recv_port) = self.df.make_edge::<_, Toff>("source_batch"); - - let schd = self.compute_state.get_scheduler(); - let inner_schd = schd.clone(); - let now = self.compute_state.current_time_ref(); - let err_collector = self.err_collector.clone(); - - let sub = self - .df - .add_subgraph_source("source_batch", send_port, move |_ctx, send| { - let mut total_batches = vec![]; - let mut total_row_count = 0; - loop { - match src_recv.try_recv() { - Ok(batch) => { - total_row_count += batch.row_count(); - total_batches.push(batch); - } - Err(TryRecvError::Empty) => { - break; - } - Err(TryRecvError::Lagged(lag_offset)) => { - // use `err_collector` instead of `error!` to locate which operator caused the error - err_collector.run(|| -> Result<(), EvalError> { - InternalSnafu { - reason: format!("Flow missing {} rows behind", lag_offset), - } - .fail() - }); - break; - } - Err(TryRecvError::Closed) => { - err_collector.run(|| -> Result<(), EvalError> { - InternalSnafu { - reason: "Source Batch Channel is closed".to_string(), - } - .fail() - }); - break; - } - } - } - - trace!( - "Send {} rows in {} batches", - total_row_count, - total_batches.len() - ); - send.give(total_batches); - - let now = *now.borrow(); - // always schedule source to run at now so we can - // repeatedly run source if needed - inner_schd.schedule_at(now); - }); - schd.set_cur_subgraph(sub); - let bundle = CollectionBundle::from_collection(Collection::::from_port(recv_port)); - Ok(bundle) - } - - /// Render a source which comes from brocast channel into the dataflow - /// will immediately send updates not greater than `now` and buffer the rest in arrangement - pub fn render_source( - &mut self, - mut src_recv: broadcast::Receiver, - ) -> Result { - debug!("Rendering Source"); - let (send_port, recv_port) = self.df.make_edge::<_, Toff>("source"); - let arrange_handler = self.compute_state.new_arrange(None); - let arrange_handler_inner = - arrange_handler - .clone_future_only() - .with_context(|| PlanSnafu { - reason: "No write is expected at this point", - })?; - - let schd = self.compute_state.get_scheduler(); - let inner_schd = schd.clone(); - let now = self.compute_state.current_time_ref(); - let err_collector = self.err_collector.clone(); - - let sub = self - .df - .add_subgraph_source("source", send_port, move |_ctx, send| { - let now = *now.borrow(); - // write lock to prevent unexpected mutation - let mut arranged = arrange_handler_inner.write(); - let arr = arranged.get_updates_in_range(..=now); - err_collector.run(|| arranged.compact_to(now)); - - let prev_avail = arr.into_iter().map(|((k, _), t, d)| (k, t, d)); - let mut to_send = Vec::new(); - let mut to_arrange = Vec::new(); - // TODO(discord9): handling tokio broadcast error - loop { - match src_recv.try_recv() { - Ok((r, t, d)) => { - if t <= now { - to_send.push((r, t, d)); - } else { - to_arrange.push(((r, Row::empty()), t, d)); - } - } - Err(TryRecvError::Empty) => { - break; - } - Err(TryRecvError::Lagged(lag_offset)) => { - common_telemetry::error!("Flow missing {} rows behind", lag_offset); - break; - } - Err(err) => { - err_collector.run(|| -> Result<(), EvalError> { - InternalSnafu { - reason: format!( - "Error receiving from broadcast channel: {}", - err - ), - } - .fail() - }); - } - } - } - let all = prev_avail.chain(to_send).collect_vec(); - if !to_arrange.is_empty() { - debug!("Source Operator buffered {} rows", to_arrange.len()); - } - err_collector.run(|| arranged.apply_updates(now, to_arrange)); - send.give(all); - // always schedule source to run at now so we can repeatedly run source if needed - inner_schd.schedule_at(now); - }); - schd.set_cur_subgraph(sub); - let arranged = Arranged::new(arrange_handler); - arranged.writer.borrow_mut().replace(sub); - let arranged = BTreeMap::from([(vec![], arranged)]); - Ok(CollectionBundle { - collection: Collection::from_port(recv_port), - arranged, - }) - } - - pub fn render_unbounded_sink_batch( - &mut self, - bundle: CollectionBundle, - sender: mpsc::UnboundedSender, - ) { - let CollectionBundle { - collection, - arranged: _, - } = bundle; - - let _sink = self.df.add_subgraph_sink( - "UnboundedSinkBatch", - collection.into_inner(), - move |_ctx, recv| { - let data = recv.take_inner(); - let mut row_count = 0; - let mut batch_count = 0; - for batch in data.into_iter().flat_map(|i| i.into_iter()) { - row_count += batch.row_count(); - batch_count += 1; - // if the sender is closed unexpectedly, stop sending - if sender.is_closed() || sender.send(batch).is_err() { - common_telemetry::error!("UnboundedSinkBatch is closed"); - break; - } - } - trace!("sink send {} rows in {} batches", row_count, batch_count); - }, - ); - } - - pub fn render_unbounded_sink( - &mut self, - bundle: CollectionBundle, - sender: mpsc::UnboundedSender, - ) { - let CollectionBundle { - collection, - arranged: _, - } = bundle; - - let _sink = self.df.add_subgraph_sink( - "UnboundedSink", - collection.into_inner(), - move |_ctx, recv| { - let data = recv.take_inner(); - debug!( - "render_unbounded_sink: send {} rows", - data.iter().map(|i| i.len()).sum::() - ); - for row in data.into_iter().flat_map(|i| i.into_iter()) { - // if the sender is closed, stop sending - if sender.is_closed() { - common_telemetry::error!("UnboundedSink is closed"); - break; - } - // TODO(discord9): handling tokio error - let _ = sender.send(row); - } - }, - ); - } -} diff --git a/src/flow/src/compute/state.rs b/src/flow/src/compute/state.rs deleted file mode 100644 index b71633f517b..00000000000 --- a/src/flow/src/compute/state.rs +++ /dev/null @@ -1,167 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::cell::RefCell; -use std::collections::{BTreeMap, VecDeque}; -use std::rc::Rc; - -use dfir_rs::scheduled::SubgraphId; -use dfir_rs::scheduled::graph::Dfir; -use get_size2::GetSize; - -use crate::compute::types::ErrCollector; -use crate::repr::{self, Timestamp}; -use crate::utils::{ArrangeHandler, Arrangement}; - -/// input/output of a dataflow -/// One `ComputeState` manage the input/output/schedule of one `Dfir` -#[derive(Debug, Default)] -pub struct DataflowState { - /// it is important to use a deque to maintain the order of subgraph here - /// TODO(discord9): consider dedup? Also not necessary for hydroflow itself also do dedup when schedule - schedule_subgraph: Rc>>>, - /// Frontier (in sys time) before which updates should not be emitted. - /// - /// We *must* apply it to sinks, to ensure correct outputs. - /// We *should* apply it to sources and imported shared state, because it improves performance. - /// Which means it's also the current time in temporal filter to get current correct result - as_of: Rc>, - /// error collector local to this `ComputeState`, - /// useful for distinguishing errors from different `Dfir` - err_collector: ErrCollector, - /// save all used arrange in this dataflow, since usually there is no delete operation - /// we can just keep track of all used arrange and schedule subgraph when they need to be updated - arrange_used: Vec, - /// the time arrangement need to be expired after a certain time in milliseconds - expire_after: Option, - /// the last time each subgraph executed - last_exec_time: Option, - /// the time the flow first executed, in unix timestamp milliseconds - start_time: Option, -} - -impl DataflowState { - pub fn new_arrange(&mut self, name: Option>) -> ArrangeHandler { - let arrange = name.map(Arrangement::new_with_name).unwrap_or_default(); - - let arr = ArrangeHandler::from(arrange); - // mark this arrange as used in this dataflow - self.arrange_used.push( - arr.clone_future_only() - .expect("No write happening at this point"), - ); - arr - } - - /// schedule all subgraph that need to run with time <= `as_of` and run_available() - /// - /// return true if any subgraph actually executed - #[allow(clippy::swap_with_temporary)] - pub fn run_available_with_schedule(&mut self, df: &mut Dfir) -> bool { - // first split keys <= as_of into another map - let mut before = self - .schedule_subgraph - .borrow_mut() - .split_off(&(*self.as_of.borrow() + 1)); - std::mem::swap(&mut before, &mut self.schedule_subgraph.borrow_mut()); - for (_, v) in before { - for subgraph in v { - df.schedule_subgraph(subgraph); - } - } - df.run_available() - } - pub fn get_scheduler(&self) -> Scheduler { - Scheduler { - schedule_subgraph: self.schedule_subgraph.clone(), - cur_subgraph: Rc::new(RefCell::new(None)), - } - } - - /// return a handle to the current time, will update when `as_of` is updated - /// - /// so it can keep track of the current time even in a closure that is called later - pub fn current_time_ref(&self) -> Rc> { - self.as_of.clone() - } - - pub fn current_ts(&self) -> Timestamp { - *self.as_of.borrow() - } - - pub fn set_current_ts(&mut self, ts: Timestamp) { - self.as_of.replace(ts); - } - - pub fn get_err_collector(&self) -> ErrCollector { - self.err_collector.clone() - } - - pub fn set_expire_after(&mut self, after: Option) { - self.expire_after = after; - } - - pub fn expire_after(&self) -> Option { - self.expire_after - } - - pub fn get_state_size(&self) -> usize { - self.arrange_used.iter().map(|x| x.read().get_size()).sum() - } - - pub fn set_last_exec_time(&mut self, time: Timestamp) { - self.last_exec_time = Some(time); - if self.start_time.is_none() { - // start_time is recorded at the completion of the first execution - // (post-execution), consistent with how last_exec_time is recorded. - self.start_time = Some(time); - } - } - - pub fn last_exec_time(&self) -> Option { - self.last_exec_time - } - - /// Returns the time the flow first executed, in unix timestamp milliseconds. - pub fn start_time(&self) -> Option { - self.start_time - } -} - -#[derive(Debug, Clone)] -pub struct Scheduler { - // this scheduler is shared with `DataflowState`, so it can schedule subgraph - schedule_subgraph: Rc>>>, - cur_subgraph: Rc>>, -} - -impl Scheduler { - pub fn schedule_at(&self, next_run_time: Timestamp) { - let mut schedule_subgraph = self.schedule_subgraph.borrow_mut(); - let subgraph = self.cur_subgraph.borrow(); - let subgraph = subgraph.as_ref().expect("Set SubgraphId before schedule"); - let subgraph_queue = schedule_subgraph.entry(next_run_time).or_default(); - subgraph_queue.push_back(*subgraph); - } - - pub fn schedule_for_arrange(&self, arrange: &Arrangement, now: Timestamp) { - if let Some(i) = arrange.get_next_update_time(&now) { - self.schedule_at(i) - } - } - - pub fn set_cur_subgraph(&self, subgraph: SubgraphId) { - self.cur_subgraph.replace(Some(subgraph)); - } -} diff --git a/src/flow/src/compute/types.rs b/src/flow/src/compute/types.rs deleted file mode 100644 index 9e33928130a..00000000000 --- a/src/flow/src/compute/types.rs +++ /dev/null @@ -1,208 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::cell::RefCell; -use std::collections::{BTreeMap, VecDeque}; -use std::rc::Rc; -use std::sync::Arc; - -use common_error::ext::ErrorExt; -use dfir_rs::scheduled::SubgraphId; -use dfir_rs::scheduled::graph::Dfir; -use dfir_rs::scheduled::handoff::TeeingHandoff; -use dfir_rs::scheduled::port::RecvPort; -use itertools::Itertools; -use tokio::sync::Mutex; - -use crate::expr::{Batch, EvalError, ScalarExpr}; -use crate::metrics::METRIC_FLOW_ERRORS; -use crate::repr::DiffRow; -use crate::utils::ArrangeHandler; - -pub type Toff = TeeingHandoff; - -/// A collection, represent a collections of data that is received from a handoff. -pub struct Collection { - /// represent a stream of updates recv from this port - stream: RecvPort>, -} - -impl Collection { - pub fn from_port(port: RecvPort>) -> Self { - Collection { stream: port } - } - - /// clone a collection, require a mutable reference to the hydroflow instance - /// - /// Note: need to be the same hydroflow instance that this collection is created from - pub fn clone(&self, df: &mut Dfir) -> Self { - Collection { - stream: self.stream.tee(df), - } - } - - pub fn into_inner(self) -> RecvPort> { - self.stream - } -} - -/// Arranged is a wrapper around `ArrangeHandler` that maintain a list of readers and a writer -pub struct Arranged { - pub arrangement: ArrangeHandler, - pub writer: Rc>>, - /// maintain a list of readers for the arrangement for the ease of scheduling - pub readers: Rc>>, -} - -impl Arranged { - pub fn new(arr: ArrangeHandler) -> Self { - Self { - arrangement: arr, - writer: Default::default(), - readers: Default::default(), - } - } - - /// Copy it's future only updates, internally `Rc-ed` so it's cheap to copy - pub fn try_copy_future(&self) -> Option { - self.arrangement - .clone_future_only() - .map(|arrangement| Arranged { - arrangement, - readers: self.readers.clone(), - writer: self.writer.clone(), - }) - } -} - -/// A bundle of the various ways a collection can be represented. -/// -/// This type maintains the invariant that it does contain at least one(or both) valid -/// source of data, either a collection or at least one arrangement. This is for convenience -/// of reading the data from the collection. -/// -// TODO(discord9): make T default to Batch and obsolete the Row Mode -pub struct CollectionBundle { - /// This is useful for passively reading the new updates from the collection - /// - /// Invariant: the timestamp of the updates should always not greater than now, since future updates should be stored in the arrangement - pub collection: Collection, - /// the key [`ScalarExpr`] indicate how the keys(also a [`Row`]) used in Arranged is extract from collection's [`Row`] - /// So it is the "index" of the arrangement - /// - /// The `Arranged` is the actual data source, it can be used to read the data from the collection by - /// using the key indicated by the `Vec` - /// There is a false positive in using `Vec` as key due to `ScalarExpr::Literal` - /// contain a `Value` which have `bytes` variant - #[allow(clippy::mutable_key_type)] - pub arranged: BTreeMap, Arranged>, -} - -pub trait GenericBundle { - fn is_batch(&self) -> bool; - - fn try_as_batch(&self) -> Option<&CollectionBundle> { - None - } - - fn try_as_row(&self) -> Option<&CollectionBundle> { - None - } -} - -impl GenericBundle for CollectionBundle { - fn is_batch(&self) -> bool { - true - } - - fn try_as_batch(&self) -> Option<&CollectionBundle> { - Some(self) - } -} - -impl GenericBundle for CollectionBundle { - fn is_batch(&self) -> bool { - false - } - - fn try_as_row(&self) -> Option<&CollectionBundle> { - Some(self) - } -} - -impl CollectionBundle { - pub fn from_collection(collection: Collection) -> Self { - Self { - collection, - arranged: BTreeMap::default(), - } - } -} - -impl CollectionBundle { - pub fn clone(&self, df: &mut Dfir) -> Self { - Self { - collection: self.collection.clone(df), - arranged: self - .arranged - .iter() - .map(|(k, v)| (k.clone(), v.try_copy_future().unwrap())) - .collect(), - } - } -} - -/// A thread local error collector, used to collect errors during the evaluation of the plan -/// -/// usually only the first error matters, but store all of them just in case -/// -/// Using a `VecDeque` to preserve the order of errors -/// when running dataflow continuously and need errors in order -#[derive(Debug, Default, Clone)] -pub struct ErrCollector { - pub inner: Arc>>, -} - -impl ErrCollector { - pub fn get_all_blocking(&self) -> Vec { - self.inner.blocking_lock().drain(..).collect_vec() - } - pub async fn get_all(&self) -> Vec { - self.inner.lock().await.drain(..).collect_vec() - } - - pub fn is_empty(&self) -> bool { - self.inner.blocking_lock().is_empty() - } - - pub fn push_err(&self, err: EvalError) { - METRIC_FLOW_ERRORS - .with_label_values(&[err.status_code().as_ref()]) - .inc(); - self.inner.blocking_lock().push_back(err) - } - - pub fn run(&self, f: F) -> Option - where - F: FnOnce() -> Result, - { - match f() { - Ok(r) => Some(r), - Err(e) => { - self.push_err(e); - None - } - } - } -} diff --git a/src/flow/src/df_optimizer.rs b/src/flow/src/df_optimizer.rs index 1cffe464fd9..106546488a8 100644 --- a/src/flow/src/df_optimizer.rs +++ b/src/flow/src/df_optimizer.rs @@ -29,19 +29,13 @@ use datafusion::optimizer::optimize_projections::OptimizeProjections; use datafusion::optimizer::simplify_expressions::SimplifyExpressions; use datafusion::optimizer::{Analyzer, AnalyzerRule, Optimizer, OptimizerContext}; use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion, TreeNodeVisitor}; -use query::QueryEngine; use query::optimizer::count_wildcard::CountWildcardToTimeIndexRule; -use query::parser::QueryLanguageParser; -use query::query_engine::DefaultSerializer; use session::context::QueryContextRef; use snafu::ResultExt; + /// note here we are using the `substrait_proto_df` crate from the `substrait` module and /// rename it to `substrait_proto` -use substrait::DFLogicalSubstraitConvertor; - -use crate::adapter::FlownodeContext; -use crate::error::{DatafusionSnafu, Error, ExternalSnafu, UnexpectedSnafu}; -use crate::plan::TypedPlan; +use crate::error::{DatafusionSnafu, Error, ExternalSnafu}; // TODO(discord9): use `Analyzer` to manage rules if more `AnalyzerRule` is needed pub async fn apply_df_optimizer( @@ -83,42 +77,6 @@ pub async fn apply_df_optimizer( Ok(plan) } -/// To reuse existing code for parse sql, the sql is first parsed into a datafusion logical plan, -/// then to a substrait plan, and finally to a flow plan. -pub async fn sql_to_flow_plan( - ctx: &mut FlownodeContext, - engine: &Arc, - sql: &str, -) -> Result { - let query_ctx = ctx.query_context.clone().ok_or_else(|| { - UnexpectedSnafu { - reason: "Query context is missing", - } - .build() - })?; - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx) - .map_err(BoxedError::new) - .context(ExternalSnafu)?; - let plan = engine - .planner() - .plan(&stmt, query_ctx.clone()) - .await - .map_err(BoxedError::new) - .context(ExternalSnafu)?; - - let opted_plan = apply_df_optimizer(plan, &query_ctx).await?; - - // TODO(discord9): add df optimization - let sub_plan = DFLogicalSubstraitConvertor {} - .to_sub_plan(&opted_plan, DefaultSerializer) - .map_err(BoxedError::new) - .context(ExternalSnafu)?; - - let flow_plan = TypedPlan::from_substrait_plan(ctx, &sub_plan).await?; - - Ok(flow_plan) -} - /// This rule check all group by exprs, and make sure they are also in select clause in a aggr query #[derive(Debug)] struct CheckGroupByRule {} diff --git a/src/flow/src/error.rs b/src/flow/src/error.rs index 2dcefd95c86..8f3cec0763e 100644 --- a/src/flow/src/error.rs +++ b/src/flow/src/error.rs @@ -25,12 +25,11 @@ use common_error::{ use common_macro::stack_trace_debug; use common_telemetry::common_error::ext::ErrorExt; use common_telemetry::common_error::status_code::StatusCode; -use snafu::{Location, ResultExt, Snafu}; +use snafu::{Location, Snafu}; use tonic::codegen::http::HeaderValue; use tonic::metadata::MetadataMap; use crate::FlowId; -use crate::expr::EvalError; /// This error is used to represent all possible errors that can occur in the flow module. #[derive(Snafu)] @@ -101,14 +100,6 @@ pub enum Error { location: Location, }, - /// TODO(discord9): add detailed location of column - #[snafu(display("Failed to eval stream"))] - Eval { - source: EvalError, - #[snafu(implicit)] - location: Location, - }, - #[snafu(display("Table not found: {name}"))] TableNotFound { name: String, @@ -161,13 +152,6 @@ pub enum Error { location: Location, }, - #[snafu(display("Not implement in flow: {reason}"))] - NotImplemented { - reason: String, - #[snafu(implicit)] - location: Location, - }, - #[snafu(display("Flow plan error: {reason}"))] Plan { reason: String, @@ -182,13 +166,6 @@ pub enum Error { location: Location, }, - #[snafu(display("Unsupported temporal filter: {reason}"))] - UnsupportedTemporalFilter { - reason: String, - #[snafu(implicit)] - location: Location, - }, - #[snafu(display("Datatypes error: {source} with extra message: {extra}"))] Datatypes { source: datatypes::Error, @@ -330,8 +307,7 @@ pub type Result = std::result::Result; impl ErrorExt for Error { fn status_code(&self) -> StatusCode { match self { - Self::Eval { .. } - | Self::JoinTask { .. } + Self::JoinTask { .. } | Self::Datafusion { .. } | Self::InsertIntoFlow { .. } | Self::NoAvailableFrontend { .. } @@ -349,9 +325,7 @@ impl ErrorExt for Error { Self::Unexpected { .. } | Self::SyncCheckTask { .. } | Self::IllegalCheckTaskState { .. } => StatusCode::Unexpected, - Self::NotImplemented { .. } - | Self::UnsupportedTemporalFilter { .. } - | Self::Unsupported { .. } => StatusCode::Unsupported, + Self::Unsupported { .. } => StatusCode::Unsupported, Self::External { source, .. } => source.status_code(), Self::Internal { .. } | Self::CacheRequired { .. } => StatusCode::Internal, Self::StartServer { source, .. } | Self::ShutdownServer { source, .. } => { @@ -405,9 +379,3 @@ impl ErrorExt for Error { } define_into_tonic_status!(Error); - -impl From for Error { - fn from(e: EvalError) -> Self { - Err::<(), _>(e).context(EvalSnafu).unwrap_err() - } -} diff --git a/src/flow/src/expr.rs b/src/flow/src/expr.rs index 5c0359e55fb..6c7f7609f1e 100644 --- a/src/flow/src/expr.rs +++ b/src/flow/src/expr.rs @@ -12,373 +12,78 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! for declare Expression in dataflow, including map, reduce, id and join(TODO!) etc. +//! Small row-batch conversion helper used by stateless streaming. -mod df_func; pub(crate) mod error; -pub(crate) mod func; -mod id; -mod linear; -pub(crate) mod relation; -mod scalar; -mod signature; -pub(crate) mod utils; - -use arrow::compute::FilterBuilder; use common_recordbatch::RecordBatch; -use datatypes::prelude::{ConcreteDataType, DataType}; -use datatypes::value::Value; -use datatypes::vectors::{BooleanVector, Helper, VectorRef}; -pub(crate) use df_func::{DfScalarFunction, RawDfScalarFn}; -pub(crate) use error::{EvalError, InvalidArgumentSnafu}; -pub(crate) use func::{BinaryFunc, UnaryFunc, UnmaterializableFunc, VariadicFunc}; -pub(crate) use id::{GlobalId, Id, LocalId}; +use datatypes::data_type::DataType; +use datatypes::prelude::ConcreteDataType; +use datatypes::vectors::{Helper, VectorRef}; use itertools::Itertools; -pub(crate) use linear::{MapFilterProject, MfpPlan, SafeMfpPlan}; -pub(crate) use relation::{Accum, Accumulator, AggregateExpr, AggregateFunc}; -pub(crate) use scalar::{ScalarExpr, TypedExpr}; use snafu::{ResultExt, ensure}; use crate::Error; use crate::error::DatatypesSnafu; -use crate::expr::error::{ArrowSnafu, DataTypeSnafu}; -use crate::repr::Diff; +use crate::repr::Row; -pub const TUMBLE_START: &str = "tumble_start"; -pub const TUMBLE_END: &str = "tumble_end"; - -/// A batch of vectors with the same length but without schema, only useful in dataflow -/// -/// somewhere cheap to clone since it just contains a list of VectorRef(which is a `Arc`). #[derive(Debug, Clone)] pub struct Batch { batch: Vec, row_count: usize, - /// describe if corresponding rows in batch is insert or delete, None means all rows are insert - diffs: Option, } impl TryFrom for Batch { type Error = Error; - fn try_from(value: RecordBatch) -> Result { - let columns = value.columns(); - let batch = Helper::try_into_vectors(columns).context(DatatypesSnafu { - extra: "failed to convert Arrow array to vector when building Flow batch", - })?; Ok(Self { + batch: Helper::try_into_vectors(value.columns()).context(DatatypesSnafu { + extra: "failed to convert Arrow array to vector", + })?, row_count: value.num_rows(), - batch, - diffs: None, }) } } -impl PartialEq for Batch { - fn eq(&self, other: &Self) -> bool { - let mut batch_eq = true; - if self.batch.len() != other.batch.len() { - return false; - } - for (left, right) in self.batch.iter().zip(other.batch.iter()) { - batch_eq = batch_eq - && ::eq(&left.to_arrow_array(), &right.to_arrow_array()); - } - - let diff_eq = match (&self.diffs, &other.diffs) { - (Some(left), Some(right)) => { - ::eq(&left.to_arrow_array(), &right.to_arrow_array()) - } - (None, None) => true, - _ => false, - }; - batch_eq && diff_eq && self.row_count == other.row_count - } -} - -impl Eq for Batch {} - -impl Default for Batch { - fn default() -> Self { - Self::empty() - } -} - impl Batch { - /// Get batch from rows, will try best to determine data type pub fn try_from_rows_with_types( - rows: Vec, - batch_datatypes: &[ConcreteDataType], - ) -> Result { + rows: Vec, + types: &[ConcreteDataType], + ) -> Result { if rows.is_empty() { - return Ok(Self::empty()); + return Ok(Self { + batch: vec![], + row_count: 0, + }); } let len = rows.len(); - let mut builder = batch_datatypes + let mut builders = types .iter() .map(|ty| ty.create_mutable_vector(len)) .collect_vec(); + ensure!( + rows.iter().all(|row| row.len() == builders.len()), + error::InvalidArgumentSnafu { + reason: "row length does not match schema".to_string() + } + ); for row in rows { - ensure!( - row.len() == builder.len(), - InvalidArgumentSnafu { - reason: format!( - "row length not match, expect {}, found {}", - builder.len(), - row.len() - ) - } - ); for (idx, value) in row.iter().enumerate() { - builder[idx] + builders[idx] .try_push_value_ref(&value.as_value_ref()) - .context(DataTypeSnafu { - msg: "Failed to convert rows to columns", + .context(error::DataTypeSnafu { + msg: "failed to convert rows to columns", })?; } } - - let columns = builder.into_iter().map(|mut b| b.to_vector()).collect_vec(); - let batch = Self::try_new(columns, len)?; - Ok(batch) - } - - pub fn empty() -> Self { - Self { - batch: vec![], - row_count: 0, - diffs: None, - } - } - pub fn try_new(batch: Vec, row_count: usize) -> Result { - ensure!( - batch.iter().map(|v| v.len()).all_equal() - && batch.first().map(|v| v.len() == row_count).unwrap_or(true), - InvalidArgumentSnafu { - reason: "All columns should have same length".to_string() - } - ); Ok(Self { - batch, - row_count, - diffs: None, + batch: builders.into_iter().map(|mut b| b.to_vector()).collect(), + row_count: len, }) } - - pub fn new_unchecked(batch: Vec, row_count: usize) -> Self { - Self { - batch, - row_count, - diffs: None, - } - } - pub fn batch(&self) -> &[VectorRef] { &self.batch } - - pub fn batch_mut(&mut self) -> &mut Vec { - &mut self.batch - } - pub fn row_count(&self) -> usize { self.row_count } - - pub fn set_row_count(&mut self, row_count: usize) { - self.row_count = row_count; - } - - pub fn column_count(&self) -> usize { - self.batch.len() - } - - pub fn get_row(&self, idx: usize) -> Result, EvalError> { - ensure!( - idx < self.row_count, - InvalidArgumentSnafu { - reason: format!( - "Expect row index to be less than {}, found {}", - self.row_count, idx - ) - } - ); - let mut ret = Vec::with_capacity(self.column_count()); - ret.extend(self.batch.iter().map(|v| v.get(idx))); - Ok(ret) - } - - /// Slices the `Batch`, returning a new `Batch`. - pub fn slice(&self, offset: usize, length: usize) -> Result { - let batch = self - .batch() - .iter() - .map(|v| v.slice(offset, length)) - .collect_vec(); - Batch::try_new(batch, length) - } - - /// append another batch to self - /// - /// NOTE: This is expensive since it will create new vectors for each column - pub fn append_batch(&mut self, other: Batch) -> Result<(), EvalError> { - ensure!( - self.batch.len() == other.batch.len() - || self.batch.is_empty() - || other.batch.is_empty(), - InvalidArgumentSnafu { - reason: format!( - "Expect two batch to have same numbers of column, found {} and {} columns", - self.batch.len(), - other.batch.len() - ) - } - ); - - if self.batch.is_empty() { - self.batch = other.batch; - self.row_count = other.row_count; - return Ok(()); - } else if other.batch.is_empty() { - return Ok(()); - } - - let dts = { - let max_len = self.batch.len().max(other.batch.len()); - let mut dts = Vec::with_capacity(max_len); - for i in 0..max_len { - if let Some(v) = self.batch().get(i) - && !v.data_type().is_null() - { - dts.push(v.data_type()) - } else if let Some(v) = other.batch().get(i) - && !v.data_type().is_null() - { - dts.push(v.data_type()) - } else { - // both are null, so we will push null type - dts.push(datatypes::prelude::ConcreteDataType::null_datatype()) - } - } - - dts - }; - - let batch_builders = dts - .iter() - .map(|dt| dt.create_mutable_vector(self.row_count() + other.row_count())) - .collect_vec(); - - let mut result = vec![]; - let self_row_count = self.row_count(); - let other_row_count = other.row_count(); - for (idx, mut builder) in batch_builders.into_iter().enumerate() { - builder - .extend_slice_of(self.batch()[idx].as_ref(), 0, self_row_count) - .context(DataTypeSnafu { - msg: "Failed to extend vector", - })?; - builder - .extend_slice_of(other.batch()[idx].as_ref(), 0, other_row_count) - .context(DataTypeSnafu { - msg: "Failed to extend vector", - })?; - result.push(builder.to_vector()); - } - self.batch = result; - self.row_count = self_row_count + other_row_count; - Ok(()) - } - - /// filter the batch with given predicate - pub fn filter(&self, predicate: &BooleanVector) -> Result { - let len = predicate.as_boolean_array().true_count(); - let filter_builder = FilterBuilder::new(predicate.as_boolean_array()).optimize(); - let filter_pred = filter_builder.build(); - let filtered = self - .batch() - .iter() - .map(|col| filter_pred.filter(col.to_arrow_array().as_ref())) - .try_collect::<_, Vec<_>, _>() - .context(ArrowSnafu { - context: "Failed to filter val batches", - })?; - let res_vector = Helper::try_into_vectors(&filtered).context(DataTypeSnafu { - msg: "can't convert arrow array to vector", - })?; - Self::try_new(res_vector, len) - } -} - -/// Vector with diff to note the insert and delete -pub(crate) struct VectorDiff { - vector: VectorRef, - diff: Option, -} - -impl From for VectorDiff { - fn from(vector: VectorRef) -> Self { - Self { vector, diff: None } - } -} - -impl VectorDiff { - fn len(&self) -> usize { - self.vector.len() - } - - fn try_new(vector: VectorRef, diff: Option) -> Result { - ensure!( - diff.as_ref().is_none_or(|diff| diff.len() == vector.len()), - InvalidArgumentSnafu { - reason: "Length of vector and diff should be the same" - } - ); - Ok(Self { vector, diff }) - } -} - -impl IntoIterator for VectorDiff { - type Item = (Value, Diff); - type IntoIter = VectorDiffIter; - - fn into_iter(self) -> Self::IntoIter { - VectorDiffIter { - vector: self.vector, - diff: self.diff, - idx: 0, - } - } -} - -/// iterator for VectorDiff -pub(crate) struct VectorDiffIter { - vector: VectorRef, - diff: Option, - idx: usize, -} - -impl std::iter::Iterator for VectorDiffIter { - type Item = (Value, Diff); - - fn next(&mut self) -> Option { - if self.idx >= self.vector.len() { - return None; - } - let value = self.vector.get(self.idx); - // +1 means insert, -1 means delete, and default to +1 insert when diff is not provided - let diff = if let Some(diff) = self.diff.as_ref() { - if let Ok(diff_at) = diff.get(self.idx).try_into() { - diff_at - } else { - common_telemetry::warn!("Invalid diff value at index {}", self.idx); - return None; - } - } else { - 1 - }; - - self.idx += 1; - Some((value, diff)) - } } diff --git a/src/flow/src/expr/df_func.rs b/src/flow/src/expr/df_func.rs deleted file mode 100644 index 9093425323e..00000000000 --- a/src/flow/src/expr/df_func.rs +++ /dev/null @@ -1,300 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! Porting Datafusion scalar function to our scalar function to be used in dataflow - -use std::sync::Arc; - -use arrow::array::RecordBatchOptions; -use bytes::BytesMut; -use common_error::ext::BoxedError; -use common_recordbatch::DfRecordBatch; -use common_telemetry::debug; -use datafusion_physical_expr::PhysicalExpr; -use datatypes::data_type::DataType; -use datatypes::value::Value; -use datatypes::vectors::VectorRef; -use prost::Message; -use snafu::{IntoError, ResultExt}; -use substrait::error::{DecodeRelSnafu, EncodeRelSnafu}; -use substrait::substrait_proto_df::proto::expression::ScalarFunction; - -use crate::error::Error; -use crate::expr::error::{ - ArrowSnafu, DatafusionSnafu as EvalDatafusionSnafu, EvalError, ExternalSnafu, - InvalidArgumentSnafu, -}; -use crate::expr::{Batch, ScalarExpr}; -use crate::repr::RelationDesc; -use crate::transform::{FunctionExtensions, from_scalar_fn_to_df_fn_impl}; - -/// A way to represent a scalar function that is implemented in Datafusion -#[derive(Debug, Clone)] -pub struct DfScalarFunction { - /// The raw bytes encoded datafusion scalar function - pub(crate) raw_fn: RawDfScalarFn, - // TODO(discord9): directly from datafusion expr - /// The implementation of the function - pub(crate) fn_impl: Arc, - /// The input schema of the function - pub(crate) df_schema: Arc, -} - -impl DfScalarFunction { - pub fn new(raw_fn: RawDfScalarFn, fn_impl: Arc) -> Result { - Ok(Self { - df_schema: Arc::new(raw_fn.input_schema.to_df_schema()?), - raw_fn, - fn_impl, - }) - } - - pub async fn try_from_raw_fn(raw_fn: RawDfScalarFn) -> Result { - Ok(Self { - fn_impl: raw_fn.get_fn_impl().await?, - df_schema: Arc::new(raw_fn.input_schema.to_df_schema()?), - raw_fn, - }) - } - - /// Evaluate a batch of expressions using input values - pub fn eval_batch(&self, batch: &Batch, exprs: &[ScalarExpr]) -> Result { - let row_count = batch.row_count(); - let batch: Vec<_> = exprs - .iter() - .map(|expr| expr.eval_batch(batch)) - .collect::>()?; - - let schema = self.df_schema.inner().clone(); - - let arrays = batch - .iter() - .map(|array| array.to_arrow_array()) - .collect::>(); - let rb = DfRecordBatch::try_new_with_options(schema, arrays, &RecordBatchOptions::new().with_row_count(Some(row_count))).map_err(|err| { - ArrowSnafu { - context: - "Failed to create RecordBatch from values when eval_batch datafusion scalar function", - } - .into_error(err) - })?; - - let len = rb.num_rows(); - - let res = self.fn_impl.evaluate(&rb).context(EvalDatafusionSnafu { - context: "Failed to evaluate datafusion scalar function", - })?; - let res = common_query::columnar_value::ColumnarValue::try_from(&res) - .map_err(BoxedError::new) - .context(ExternalSnafu)?; - let res_vec = res - .try_into_vector(len) - .map_err(BoxedError::new) - .context(ExternalSnafu)?; - - Ok(res_vec) - } - - /// eval a list of expressions using input values - fn eval_args(values: &[Value], exprs: &[ScalarExpr]) -> Result, EvalError> { - exprs - .iter() - .map(|expr| expr.eval(values)) - .collect::>() - } - - // TODO(discord9): add RecordBatch support - pub fn eval(&self, values: &[Value], exprs: &[ScalarExpr]) -> Result { - // first eval exprs to construct values to feed to datafusion - let values: Vec<_> = Self::eval_args(values, exprs)?; - if values.is_empty() { - return InvalidArgumentSnafu { - reason: "values is empty".to_string(), - } - .fail(); - } - // TODO(discord9): make cols all array length of one - let mut cols = vec![]; - for (idx, typ) in self - .raw_fn - .input_schema - .typ() - .column_types - .iter() - .enumerate() - { - let typ = typ.scalar_type(); - let mut array = typ.create_mutable_vector(1); - array.push_value_ref(&values[idx].as_value_ref()); - cols.push(array.to_vector().to_arrow_array()); - } - let schema = self.df_schema.inner().clone(); - let rb = DfRecordBatch::try_new_with_options( - schema, - cols, - &RecordBatchOptions::new().with_row_count(Some(1)), - ) - .map_err(|err| { - ArrowSnafu { - context: - "Failed to create RecordBatch from values when eval datafusion scalar function", - } - .into_error(err) - })?; - - let res = self.fn_impl.evaluate(&rb).context(EvalDatafusionSnafu { - context: "Failed to evaluate datafusion scalar function", - })?; - let res = common_query::columnar_value::ColumnarValue::try_from(&res) - .map_err(BoxedError::new) - .context(ExternalSnafu)?; - let res_vec = res - .try_into_vector(1) - .map_err(BoxedError::new) - .context(ExternalSnafu)?; - let res_val = res_vec - .try_get(0) - .map_err(BoxedError::new) - .context(ExternalSnafu)?; - Ok(res_val) - } -} - -#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct RawDfScalarFn { - /// The raw bytes encoded datafusion scalar function, - /// due to substrait have too many layers of nested struct and `ScalarFunction` 's derive is different - /// for simplicity's sake - /// so we store bytes instead of `ScalarFunction` here - /// but in unit test we will still compare decoded struct(using `f_decoded` field in Debug impl) - pub(crate) f: bytes::BytesMut, - /// The input schema of the function - pub(crate) input_schema: RelationDesc, - /// Extension contains mapping from function reference to function name - pub(crate) extensions: FunctionExtensions, -} - -impl std::fmt::Debug for RawDfScalarFn { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("RawDfScalarFn") - .field("f", &self.f) - .field("f_decoded", &ScalarFunction::decode(&mut self.f.as_ref())) - .field("df_schema", &self.input_schema) - .field("extensions", &self.extensions) - .finish() - } -} - -impl RawDfScalarFn { - pub fn from_proto( - f: &substrait::substrait_proto_df::proto::expression::ScalarFunction, - input_schema: RelationDesc, - extensions: FunctionExtensions, - ) -> Result { - let mut buf = BytesMut::new(); - f.encode(&mut buf) - .context(EncodeRelSnafu) - .map_err(BoxedError::new) - .context(crate::error::ExternalSnafu)?; - Ok(Self { - f: buf, - input_schema, - extensions, - }) - } - async fn get_fn_impl(&self) -> Result, Error> { - let f = ScalarFunction::decode(&mut self.f.as_ref()) - .context(DecodeRelSnafu) - .map_err(BoxedError::new) - .context(crate::error::ExternalSnafu)?; - debug!("Decoded scalar function: {:?}", f); - - let input_schema = &self.input_schema; - let extensions = &self.extensions; - - from_scalar_fn_to_df_fn_impl(&f, input_schema, extensions).await - } -} - -impl std::cmp::PartialEq for DfScalarFunction { - fn eq(&self, other: &Self) -> bool { - self.raw_fn.eq(&other.raw_fn) - } -} - -// can't derive Eq because of Arc not eq, so implement it manually -impl std::cmp::Eq for DfScalarFunction {} - -impl std::cmp::PartialOrd for DfScalarFunction { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) - } -} -impl std::cmp::Ord for DfScalarFunction { - fn cmp(&self, other: &Self) -> std::cmp::Ordering { - self.raw_fn.cmp(&other.raw_fn) - } -} -impl std::hash::Hash for DfScalarFunction { - fn hash(&self, state: &mut H) { - self.raw_fn.hash(state); - } -} - -#[cfg(test)] -mod test { - - use datatypes::prelude::ConcreteDataType; - use substrait::substrait_proto_df::proto::expression::literal::LiteralType; - use substrait::substrait_proto_df::proto::expression::{Literal, RexType}; - use substrait::substrait_proto_df::proto::function_argument::ArgType; - use substrait::substrait_proto_df::proto::{Expression, FunctionArgument}; - - use super::*; - use crate::repr::{ColumnType, RelationType}; - - #[tokio::test] - async fn test_df_scalar_function() { - let raw_scalar_func = ScalarFunction { - function_reference: 0, - arguments: vec![FunctionArgument { - arg_type: Some(ArgType::Value(Expression { - rex_type: Some(RexType::Literal(Literal { - nullable: false, - type_variation_reference: 0, - literal_type: Some(LiteralType::I64(-1)), - })), - })), - }], - output_type: None, - ..Default::default() - }; - let input_schema = RelationDesc::try_new( - RelationType::new(vec![ColumnType::new_nullable( - ConcreteDataType::null_datatype(), - )]), - vec!["null_column".to_string()], - ) - .unwrap(); - let extensions = FunctionExtensions::from_iter(vec![(0, "abs")]); - let raw_fn = RawDfScalarFn::from_proto(&raw_scalar_func, input_schema, extensions).unwrap(); - let df_func = DfScalarFunction::try_from_raw_fn(raw_fn).await.unwrap(); - assert_eq!( - df_func - .eval(&[Value::Null], &[ScalarExpr::Column(0)]) - .unwrap(), - Value::Int64(1) - ); - } -} diff --git a/src/flow/src/expr/error.rs b/src/flow/src/expr/error.rs index b29098b9d8f..ac5319314aa 100644 --- a/src/flow/src/expr/error.rs +++ b/src/flow/src/expr/error.rs @@ -31,12 +31,6 @@ use snafu::{Location, Snafu}; #[snafu(visibility(pub))] #[stack_trace_debug] pub enum EvalError { - #[snafu(display("Division by zero"))] - DivisionByZero { - #[snafu(implicit)] - location: Location, - }, - #[snafu(display("Type mismatch: expected {expected}, actual {actual}"))] TypeMismatch { expected: ConcreteDataType, @@ -84,26 +78,12 @@ pub enum EvalError { location: Location, }, - #[snafu(display("Optimize error: {reason}"))] - Optimize { - reason: String, - #[snafu(implicit)] - location: Location, - }, - #[snafu(display("Overflowed during evaluation"))] Overflow { #[snafu(implicit)] location: Location, }, - #[snafu(display("Incoming data already expired by {} ms", expired_by))] - DataAlreadyExpired { - expired_by: i64, - #[snafu(implicit)] - location: Location, - }, - #[snafu(display("Arrow error: {error:?}, context: {context}"))] Arrow { #[snafu(source)] @@ -134,20 +114,16 @@ impl ErrorExt for EvalError { fn status_code(&self) -> StatusCode { use EvalError::*; match self { - DivisionByZero { .. } - | TypeMismatch { .. } + TypeMismatch { .. } | TryFromValue { .. } - | DataAlreadyExpired { .. } | InvalidArgument { .. } | Overflow { .. } => StatusCode::InvalidArguments, CastValue { source, .. } | DataType { source, .. } => source.status_code(), - Internal { .. } - | Optimize { .. } - | Arrow { .. } - | Datafusion { .. } - | External { .. } => StatusCode::Internal, + Internal { .. } | Arrow { .. } | Datafusion { .. } | External { .. } => { + StatusCode::Internal + } } } diff --git a/src/flow/src/expr/func.rs b/src/flow/src/expr/func.rs deleted file mode 100644 index cda70773a39..00000000000 --- a/src/flow/src/expr/func.rs +++ /dev/null @@ -1,1467 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! This module contains the definition of functions that can be used in expressions. - -use std::collections::HashMap; -use std::sync::{Arc, OnceLock}; -use std::time::Duration; - -use arrow::array::{ArrayRef, BooleanArray}; -use common_error::ext::BoxedError; -use common_time::Timestamp; -use common_time::timestamp::TimeUnit; -use datafusion_expr::Operator; -use datatypes::data_type::ConcreteDataType; -use datatypes::prelude::DataType; -use datatypes::types::cast; -use datatypes::value::Value; -use datatypes::vectors::{BooleanVector, Helper, TimestampMillisecondVector, VectorRef}; -use serde::{Deserialize, Serialize}; -use smallvec::smallvec; -use snafu::{OptionExt, ResultExt, ensure}; -use strum::{EnumIter, IntoEnumIterator}; -use substrait::df_logical_plan::consumer::name_to_op; - -use crate::error::{Error, ExternalSnafu, InvalidQuerySnafu, PlanSnafu, UnexpectedSnafu}; -use crate::expr::error::{ - ArrowSnafu, CastValueSnafu, DataTypeSnafu, DivisionByZeroSnafu, EvalError, OverflowSnafu, - TryFromValueSnafu, TypeMismatchSnafu, -}; -use crate::expr::signature::{GenericFn, Signature}; -use crate::expr::{Batch, InvalidArgumentSnafu, ScalarExpr, TUMBLE_END, TUMBLE_START, TypedExpr}; -use crate::repr::{self, value_to_internal_ts}; - -/// UnmaterializableFunc is a function that can't be eval independently, -/// and require special handling -#[derive(Ord, PartialOrd, Clone, Debug, Eq, PartialEq, Hash)] -pub enum UnmaterializableFunc { - Now, - CurrentSchema, - TumbleWindow { - ts: Box, - window_size: Duration, - start_time: Option, - }, -} - -impl UnmaterializableFunc { - /// Return the signature of the function - pub fn signature(&self) -> Signature { - match self { - Self::Now => Signature { - input: smallvec![], - // TODO(yingwen): Maybe return timestamp. - output: ConcreteDataType::timestamp_millisecond_datatype(), - generic_fn: GenericFn::Now, - }, - Self::CurrentSchema => Signature { - input: smallvec![], - output: ConcreteDataType::string_datatype(), - generic_fn: GenericFn::CurrentSchema, - }, - Self::TumbleWindow { .. } => Signature { - input: smallvec![ConcreteDataType::timestamp_millisecond_datatype()], - output: ConcreteDataType::timestamp_millisecond_datatype(), - generic_fn: GenericFn::TumbleWindow, - }, - } - } - - pub fn is_valid_func_name(name: &str) -> bool { - matches!( - name.to_lowercase().as_str(), - "now" | "current_schema" | "tumble" - ) - } - - /// Create a UnmaterializableFunc from a string of the function name - pub fn from_str_args(name: &str, _args: Vec) -> Result { - match name.to_lowercase().as_str() { - "now" => Ok(Self::Now), - "current_schema" => Ok(Self::CurrentSchema), - _ => InvalidQuerySnafu { - reason: format!("Unknown unmaterializable function: {}", name), - } - .fail(), - } - } -} - -/// UnaryFunc is a function that takes one argument. Also notice this enum doesn't contain function arguments, -/// because the arguments are stored in the expression. (except `cast` function, which requires a type argument) -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Deserialize, Serialize, Hash)] -pub enum UnaryFunc { - Not, - IsNull, - IsTrue, - IsFalse, - StepTimestamp, - Cast(ConcreteDataType), - TumbleWindowFloor { - window_size: Duration, - start_time: Option, - }, - TumbleWindowCeiling { - window_size: Duration, - start_time: Option, - }, -} - -impl UnaryFunc { - /// Return the signature of the function - pub fn signature(&self) -> Signature { - match self { - Self::IsNull => Signature { - input: smallvec![ConcreteDataType::null_datatype()], - output: ConcreteDataType::boolean_datatype(), - generic_fn: GenericFn::IsNull, - }, - Self::Not | Self::IsTrue | Self::IsFalse => Signature { - input: smallvec![ConcreteDataType::boolean_datatype()], - output: ConcreteDataType::boolean_datatype(), - generic_fn: match self { - Self::Not => GenericFn::Not, - Self::IsTrue => GenericFn::IsTrue, - Self::IsFalse => GenericFn::IsFalse, - _ => unreachable!(), - }, - }, - Self::StepTimestamp => Signature { - input: smallvec![ConcreteDataType::timestamp_millisecond_datatype()], - output: ConcreteDataType::timestamp_millisecond_datatype(), - generic_fn: GenericFn::StepTimestamp, - }, - Self::Cast(to) => Signature { - input: smallvec![ConcreteDataType::null_datatype()], - output: to.clone(), - generic_fn: GenericFn::Cast, - }, - Self::TumbleWindowFloor { .. } => Signature { - input: smallvec![ConcreteDataType::timestamp_millisecond_datatype()], - output: ConcreteDataType::timestamp_millisecond_datatype(), - generic_fn: GenericFn::TumbleWindow, - }, - Self::TumbleWindowCeiling { .. } => Signature { - input: smallvec![ConcreteDataType::timestamp_millisecond_datatype()], - output: ConcreteDataType::timestamp_millisecond_datatype(), - generic_fn: GenericFn::TumbleWindow, - }, - } - } - - pub fn is_valid_func_name(name: &str) -> bool { - matches!( - name.to_lowercase().as_str(), - "not" | "is_null" | "is_true" | "is_false" | "step_timestamp" | "cast" - ) - } - - /// Create a UnaryFunc from a string of the function name and given argument type(optional) - pub fn from_str_and_type( - name: &str, - arg_type: Option, - ) -> Result { - match name { - "not" => Ok(Self::Not), - "is_null" => Ok(Self::IsNull), - "is_true" => Ok(Self::IsTrue), - "is_false" => Ok(Self::IsFalse), - "step_timestamp" => Ok(Self::StepTimestamp), - "cast" => { - let arg_type = arg_type.with_context(|| InvalidQuerySnafu { - reason: "cast function requires a type argument".to_string(), - })?; - Ok(UnaryFunc::Cast(arg_type)) - } - _ => InvalidQuerySnafu { - reason: format!("Unknown unary function: {}", name), - } - .fail(), - } - } - - pub fn eval_batch(&self, batch: &Batch, expr: &ScalarExpr) -> Result { - let arg_col = expr.eval_batch(batch)?; - match self { - Self::Not => { - let arrow_array = arg_col.to_arrow_array(); - let bool_array = arrow_array - .as_any() - .downcast_ref::() - .context({ - TypeMismatchSnafu { - expected: ConcreteDataType::boolean_datatype(), - actual: arg_col.data_type(), - } - })?; - let ret = arrow::compute::not(bool_array).context(ArrowSnafu { context: "not" })?; - let ret = BooleanVector::from(ret); - Ok(Arc::new(ret)) - } - Self::IsNull => { - let arrow_array = arg_col.to_arrow_array(); - let ret = arrow::compute::is_null(&arrow_array) - .context(ArrowSnafu { context: "is_null" })?; - let ret = BooleanVector::from(ret); - Ok(Arc::new(ret)) - } - Self::IsTrue | Self::IsFalse => { - let arrow_array = arg_col.to_arrow_array(); - let bool_array = arrow_array - .as_any() - .downcast_ref::() - .context({ - TypeMismatchSnafu { - expected: ConcreteDataType::boolean_datatype(), - actual: arg_col.data_type(), - } - })?; - - if matches!(self, Self::IsTrue) { - Ok(Arc::new(BooleanVector::from(bool_array.clone()))) - } else { - let ret = - arrow::compute::not(bool_array).context(ArrowSnafu { context: "not" })?; - Ok(Arc::new(BooleanVector::from(ret))) - } - } - Self::StepTimestamp => { - let timestamp_array = get_timestamp_array(&arg_col)?; - let timestamp_array_ref = timestamp_array - .as_any() - .downcast_ref::() - .context({ - TypeMismatchSnafu { - expected: ConcreteDataType::boolean_datatype(), - actual: ConcreteDataType::from_arrow_type(timestamp_array.data_type()), - } - })?; - - let ret = arrow::compute::unary(timestamp_array_ref, |arr| arr + 1); - let ret = TimestampMillisecondVector::from(ret); - Ok(Arc::new(ret)) - } - Self::Cast(to) => { - let arrow_array = arg_col.to_arrow_array(); - let ret = arrow::compute::cast(&arrow_array, &to.as_arrow_type()) - .context(ArrowSnafu { context: "cast" })?; - let vector = Helper::try_into_vector(ret).context(DataTypeSnafu { - msg: "Fail to convert to Vector", - })?; - Ok(vector) - } - Self::TumbleWindowFloor { - window_size, - start_time, - } => { - let timestamp_array = get_timestamp_array(&arg_col)?; - let date_array_ref = timestamp_array - .as_any() - .downcast_ref::() - .context({ - TypeMismatchSnafu { - expected: ConcreteDataType::boolean_datatype(), - actual: ConcreteDataType::from_arrow_type(timestamp_array.data_type()), - } - })?; - - let start_time = start_time.map(|t| t.value()); - let window_size = window_size.as_millis() as repr::Duration; - - let ret = arrow::compute::unary(date_array_ref, |ts| { - get_window_start(ts, window_size, start_time) - }); - - let ret = TimestampMillisecondVector::from(ret); - Ok(Arc::new(ret)) - } - Self::TumbleWindowCeiling { - window_size, - start_time, - } => { - let timestamp_array = get_timestamp_array(&arg_col)?; - let date_array_ref = timestamp_array - .as_any() - .downcast_ref::() - .context({ - TypeMismatchSnafu { - expected: ConcreteDataType::boolean_datatype(), - actual: ConcreteDataType::from_arrow_type(timestamp_array.data_type()), - } - })?; - - let start_time = start_time.map(|t| t.value()); - let window_size = window_size.as_millis() as repr::Duration; - - let ret = arrow::compute::unary(date_array_ref, |ts| { - get_window_start(ts, window_size, start_time) + window_size - }); - - let ret = TimestampMillisecondVector::from(ret); - Ok(Arc::new(ret)) - } - } - } - - pub fn from_tumble_func(name: &str, args: &[TypedExpr]) -> Result<(Self, TypedExpr), Error> { - match name.to_lowercase().as_str() { - TUMBLE_START | TUMBLE_END => { - let ts = args.first().context(InvalidQuerySnafu { - reason: "Tumble window function requires a timestamp argument", - })?; - let window_size = { - let window_size_untyped = args - .get(1) - .and_then(|expr| expr.expr.as_literal()) - .context(InvalidQuerySnafu { - reason: "Tumble window function requires a window size argument", - })?; - if let Some(window_size) = window_size_untyped.as_string() { - // cast as interval - let interval = cast( - Value::from(window_size), - &ConcreteDataType::interval_day_time_datatype(), - ) - .map_err(BoxedError::new) - .context(ExternalSnafu)? - .as_interval_day_time() - .context(UnexpectedSnafu { - reason: "Expect window size arg to be interval after successful cast" - .to_string(), - })?; - Duration::from_millis(interval.as_millis() as u64) - } else if let Some(interval) = window_size_untyped.as_interval_day_time() { - Duration::from_millis(interval.as_millis() as u64) - } else { - InvalidQuerySnafu { - reason: format!( - "Tumble window function requires window size argument to be either a interval or a string describe a interval, found {:?}", - window_size_untyped - ) - }.fail()? - } - }; - - // start time argument is optional - let start_time = match args.get(2) { - Some(start_time) => { - if let Some(value) = start_time.expr.as_literal() { - // cast as timestamp - let ret = cast( - value, - &ConcreteDataType::timestamp_millisecond_datatype(), - ) - .map_err(BoxedError::new) - .context(ExternalSnafu)? - .as_timestamp() - .context(UnexpectedSnafu { - reason: - "Expect start time arg to be timestamp after successful cast" - .to_string(), - })?; - Some(ret) - } else { - UnexpectedSnafu { - reason: "Expect start time arg to be literal", - } - .fail()? - } - } - None => None, - }; - - if name == TUMBLE_START { - Ok(( - Self::TumbleWindowFloor { - window_size, - start_time, - }, - ts.clone(), - )) - } else if name == TUMBLE_END { - Ok(( - Self::TumbleWindowCeiling { - window_size, - start_time, - }, - ts.clone(), - )) - } else { - unreachable!() - } - } - _ => crate::error::InternalSnafu { - reason: format!("Unknown tumble kind function: {}", name), - } - .fail()?, - } - } - - /// Evaluate the function with given values and expression - /// - /// # Arguments - /// - /// - `values`: The values to be used in the evaluation - /// - /// - `expr`: The expression to be evaluated and use as argument, will extract the value from the `values` and evaluate the expression - pub fn eval(&self, values: &[Value], expr: &ScalarExpr) -> Result { - let arg = expr.eval(values)?; - match self { - Self::Not => { - let bool = if let Value::Boolean(bool) = arg { - Ok(bool) - } else { - TypeMismatchSnafu { - expected: ConcreteDataType::boolean_datatype(), - actual: arg.data_type(), - } - .fail()? - }?; - Ok(Value::from(!bool)) - } - Self::IsNull => Ok(Value::from(arg.is_null())), - Self::IsTrue | Self::IsFalse => { - let bool = if let Value::Boolean(bool) = arg { - Ok(bool) - } else { - TypeMismatchSnafu { - expected: ConcreteDataType::boolean_datatype(), - actual: arg.data_type(), - } - .fail()? - }?; - if matches!(self, Self::IsTrue) { - Ok(Value::from(bool)) - } else { - Ok(Value::from(!bool)) - } - } - Self::StepTimestamp => { - let ty = arg.data_type(); - if let Value::Timestamp(timestamp) = arg { - let timestamp = Timestamp::new_millisecond(timestamp.value() + 1); - Ok(Value::from(timestamp)) - } else if let Ok(v) = value_to_internal_ts(arg) { - let timestamp = Timestamp::new_millisecond(v + 1); - Ok(Value::from(timestamp)) - } else { - TypeMismatchSnafu { - expected: ConcreteDataType::timestamp_millisecond_datatype(), - actual: ty, - } - .fail()? - } - } - Self::Cast(to) => { - let arg_ty = arg.data_type(); - cast(arg, to).context({ - CastValueSnafu { - from: arg_ty, - to: to.clone(), - } - }) - } - Self::TumbleWindowFloor { - window_size, - start_time, - } => { - let ts = get_ts_as_millisecond(arg)?; - let start_time = start_time.map(|t| t.value()); - let window_size = window_size.as_millis() as repr::Duration; - let window_start = get_window_start(ts, window_size, start_time); - - let ret = Timestamp::new_millisecond(window_start); - Ok(Value::from(ret)) - } - Self::TumbleWindowCeiling { - window_size, - start_time, - } => { - let ts = get_ts_as_millisecond(arg)?; - let start_time = start_time.map(|t| t.value()); - let window_size = window_size.as_millis() as repr::Duration; - let window_start = get_window_start(ts, window_size, start_time); - - let window_end = window_start + window_size; - let ret = Timestamp::new_millisecond(window_end); - Ok(Value::from(ret)) - } - } - } -} - -fn get_timestamp_array(vector: &VectorRef) -> Result { - let arrow_array = vector.to_arrow_array(); - let timestamp_array = if *arrow_array.data_type() - == ConcreteDataType::timestamp_millisecond_datatype().as_arrow_type() - { - arrow_array - } else { - arrow::compute::cast( - &arrow_array, - &ConcreteDataType::timestamp_millisecond_datatype().as_arrow_type(), - ) - .context(ArrowSnafu { - context: "Trying to cast to timestamp in StepTimestamp", - })? - }; - Ok(timestamp_array) -} - -fn get_window_start( - ts: repr::Timestamp, - window_size: repr::Duration, - start_time: Option, -) -> repr::Timestamp { - let start_time = start_time.unwrap_or(0); - // left close right open - if ts >= start_time { - start_time + (ts - start_time) / window_size * window_size - } else { - start_time + (ts - start_time) / window_size * window_size - - if ((start_time - ts) % window_size) != 0 { - window_size - } else { - 0 - } - } -} - -#[test] -fn test_get_window_start() { - assert_eq!(get_window_start(1, 3, None), 0); - assert_eq!(get_window_start(3, 3, None), 3); - assert_eq!(get_window_start(0, 3, None), 0); - - assert_eq!(get_window_start(-1, 3, None), -3); - assert_eq!(get_window_start(-3, 3, None), -3); -} - -fn get_ts_as_millisecond(arg: Value) -> Result { - let ts = if let Some(ts) = arg.as_timestamp() { - ts.convert_to(TimeUnit::Millisecond) - .context(OverflowSnafu)? - .value() - } else { - InvalidArgumentSnafu { - reason: "Expect input to be timestamp or datetime type", - } - .fail()? - }; - Ok(ts) -} - -/// BinaryFunc is a function that takes two arguments. -/// Also notice this enum doesn't contain function arguments, since the arguments are stored in the expression. -/// -/// TODO(discord9): support more binary functions for more types -#[derive( - Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Deserialize, Serialize, Hash, EnumIter, -)] -pub enum BinaryFunc { - Eq, - NotEq, - Lt, - Lte, - Gt, - Gte, - AddInt16, - AddInt32, - AddInt64, - AddUInt16, - AddUInt32, - AddUInt64, - AddFloat32, - AddFloat64, - SubInt16, - SubInt32, - SubInt64, - SubUInt16, - SubUInt32, - SubUInt64, - SubFloat32, - SubFloat64, - MulInt16, - MulInt32, - MulInt64, - MulUInt16, - MulUInt32, - MulUInt64, - MulFloat32, - MulFloat64, - DivInt16, - DivInt32, - DivInt64, - DivUInt16, - DivUInt32, - DivUInt64, - DivFloat32, - DivFloat64, - ModInt16, - ModInt32, - ModInt64, - ModUInt16, - ModUInt32, - ModUInt64, -} - -/// Generate binary function signature based on the function and the input types -/// The user can provide custom signature for some functions in the form of a regular match arm, -/// and the rest will be generated according to the provided list of functions like this: -/// ```ignore -/// AddInt16=>(int16_datatype,Add), -/// ``` -/// which expand to: -/// ```ignore, rust -/// Self::AddInt16 => Signature { -/// input: smallvec![ -/// ConcreteDataType::int16_datatype(), -/// ConcreteDataType::int16_datatype(), -/// ], -/// output: ConcreteDataType::int16_datatype(), -/// generic_fn: GenericFn::Add, -/// }, -/// ```` -macro_rules! generate_binary_signature { - ($value:ident, { $($user_arm:tt)* }, - [ $( - $auto_arm:ident=>($con_type:ident,$generic:ident) - ),* - ]) => { - match $value { - $($user_arm)*, - $( - Self::$auto_arm => Signature { - input: smallvec![ - ConcreteDataType::$con_type(), - ConcreteDataType::$con_type(), - ], - output: ConcreteDataType::$con_type(), - generic_fn: GenericFn::$generic, - }, - )* - } - }; -} - -static SPECIALIZATION: OnceLock> = - OnceLock::new(); - -impl BinaryFunc { - /// Use null type to ref to any type - pub fn signature(&self) -> Signature { - generate_binary_signature!(self, { - Self::Eq | Self::NotEq | Self::Lt | Self::Lte | Self::Gt | Self::Gte => Signature { - input: smallvec![ - ConcreteDataType::null_datatype(), - ConcreteDataType::null_datatype() - ], - output: ConcreteDataType::boolean_datatype(), - generic_fn: match self { - Self::Eq => GenericFn::Eq, - Self::NotEq => GenericFn::NotEq, - Self::Lt => GenericFn::Lt, - Self::Lte => GenericFn::Lte, - Self::Gt => GenericFn::Gt, - Self::Gte => GenericFn::Gte, - _ => unreachable!(), - }, - } - }, - [ - AddInt16=>(int16_datatype,Add), - AddInt32=>(int32_datatype,Add), - AddInt64=>(int64_datatype,Add), - AddUInt16=>(uint16_datatype,Add), - AddUInt32=>(uint32_datatype,Add), - AddUInt64=>(uint64_datatype,Add), - AddFloat32=>(float32_datatype,Add), - AddFloat64=>(float64_datatype,Add), - SubInt16=>(int16_datatype,Sub), - SubInt32=>(int32_datatype,Sub), - SubInt64=>(int64_datatype,Sub), - SubUInt16=>(uint16_datatype,Sub), - SubUInt32=>(uint32_datatype,Sub), - SubUInt64=>(uint64_datatype,Sub), - SubFloat32=>(float32_datatype,Sub), - SubFloat64=>(float64_datatype,Sub), - MulInt16=>(int16_datatype,Mul), - MulInt32=>(int32_datatype,Mul), - MulInt64=>(int64_datatype,Mul), - MulUInt16=>(uint16_datatype,Mul), - MulUInt32=>(uint32_datatype,Mul), - MulUInt64=>(uint64_datatype,Mul), - MulFloat32=>(float32_datatype,Mul), - MulFloat64=>(float64_datatype,Mul), - DivInt16=>(int16_datatype,Div), - DivInt32=>(int32_datatype,Div), - DivInt64=>(int64_datatype,Div), - DivUInt16=>(uint16_datatype,Div), - DivUInt32=>(uint32_datatype,Div), - DivUInt64=>(uint64_datatype,Div), - DivFloat32=>(float32_datatype,Div), - DivFloat64=>(float64_datatype,Div), - ModInt16=>(int16_datatype,Mod), - ModInt32=>(int32_datatype,Mod), - ModInt64=>(int64_datatype,Mod), - ModUInt16=>(uint16_datatype,Mod), - ModUInt32=>(uint32_datatype,Mod), - ModUInt64=>(uint64_datatype,Mod) - ] - ) - } - - pub fn add(input_type: ConcreteDataType) -> Result { - Self::specialization(GenericFn::Add, input_type) - } - - pub fn sub(input_type: ConcreteDataType) -> Result { - Self::specialization(GenericFn::Sub, input_type) - } - - pub fn mul(input_type: ConcreteDataType) -> Result { - Self::specialization(GenericFn::Mul, input_type) - } - - pub fn div(input_type: ConcreteDataType) -> Result { - Self::specialization(GenericFn::Div, input_type) - } - - /// Get the specialization of the binary function based on the generic function and the input type - pub fn specialization(generic: GenericFn, input_type: ConcreteDataType) -> Result { - let rule = SPECIALIZATION.get_or_init(|| { - let mut spec = HashMap::new(); - for func in BinaryFunc::iter() { - let sig = func.signature(); - spec.insert((sig.generic_fn, sig.input[0].clone()), func); - } - spec - }); - rule.get(&(generic, input_type.clone())) - .cloned() - .with_context(|| InvalidQuerySnafu { - reason: format!( - "No specialization found for binary function {:?} with input type {:?}", - generic, input_type - ), - }) - } - - /// try it's best to infer types from the input types and expressions - /// - /// if it can't found out types, will return None - pub(crate) fn infer_type_from( - generic: GenericFn, - arg_exprs: &[ScalarExpr], - arg_types: &[Option], - ) -> Result { - let ret = match (arg_types[0].as_ref(), arg_types[1].as_ref()) { - (Some(t1), Some(t2)) => { - ensure!( - t1 == t2, - InvalidQuerySnafu { - reason: format!( - "Binary function {:?} requires both arguments to have the same type, left={:?}, right={:?}", - generic, t1, t2 - ), - } - ); - t1.clone() - } - (Some(t), None) | (None, Some(t)) => t.clone(), - _ => arg_exprs[0] - .as_literal() - .map(|lit| lit.data_type()) - .or_else(|| arg_exprs[1].as_literal().map(|lit| lit.data_type())) - .with_context(|| InvalidQuerySnafu { - reason: format!( - "Binary function {:?} requires at least one argument with known type", - generic - ), - })?, - }; - Ok(ret) - } - - pub fn is_valid_func_name(name: &str) -> bool { - matches!( - name.to_lowercase().as_str(), - "eq" | "equal" - | "not_eq" - | "not_equal" - | "lt" - | "lte" - | "gt" - | "gte" - | "add" - | "sub" - | "subtract" - | "mul" - | "multiply" - | "div" - | "divide" - | "mod" - ) - } - - /// choose the appropriate specialization based on the input types - /// return a specialization of the binary function and it's actual input and output type(so no null type present) - /// - /// will try it best to extract from `arg_types` and `arg_exprs` to get the input types - /// if `arg_types` is not enough, it will try to extract from `arg_exprs` if `arg_exprs` is literal with known type - pub fn from_str_expr_and_type( - name: &str, - arg_exprs: &[ScalarExpr], - arg_types: &[Option], - ) -> Result<(Self, Signature), Error> { - // this `name_to_op` if error simply return a similar message of `unsupported function xxx` so - let op = name_to_op(name).with_context(|| InvalidQuerySnafu { - reason: format!("Unsupported binary function: {}", name), - })?; - - // get first arg type and make sure if both is some, they are the same - let generic_fn = { - match op { - Operator::Eq => GenericFn::Eq, - Operator::NotEq => GenericFn::NotEq, - Operator::Lt => GenericFn::Lt, - Operator::LtEq => GenericFn::Lte, - Operator::Gt => GenericFn::Gt, - Operator::GtEq => GenericFn::Gte, - Operator::Plus => GenericFn::Add, - Operator::Minus => GenericFn::Sub, - Operator::Multiply => GenericFn::Mul, - Operator::Divide => GenericFn::Div, - Operator::Modulo => GenericFn::Mod, - _ => { - return InvalidQuerySnafu { - reason: format!("Unsupported binary function: {}", name), - } - .fail(); - } - } - }; - let need_type = matches!( - generic_fn, - GenericFn::Add | GenericFn::Sub | GenericFn::Mul | GenericFn::Div | GenericFn::Mod - ); - - ensure!( - arg_exprs.len() == 2 && arg_types.len() == 2, - PlanSnafu { - reason: "Binary function requires exactly 2 arguments".to_string() - } - ); - - let arg_type = Self::infer_type_from(generic_fn, arg_exprs, arg_types)?; - - // if type is not needed, we can erase input type to null to find correct functions for - // functions that do not need type - let query_input_type = if need_type { - arg_type.clone() - } else { - ConcreteDataType::null_datatype() - }; - - let spec_fn = Self::specialization(generic_fn, query_input_type)?; - - let signature = Signature { - input: smallvec![arg_type.clone(), arg_type], - output: spec_fn.signature().output, - generic_fn, - }; - - Ok((spec_fn, signature)) - } - - pub fn eval_batch( - &self, - batch: &Batch, - expr1: &ScalarExpr, - expr2: &ScalarExpr, - ) -> Result { - let left = expr1.eval_batch(batch)?; - let left = left.to_arrow_array(); - let right = expr2.eval_batch(batch)?; - let right = right.to_arrow_array(); - - let arrow_array: ArrayRef = match self { - Self::Eq => Arc::new( - arrow::compute::kernels::cmp::eq(&left, &right) - .context(ArrowSnafu { context: "eq" })?, - ), - Self::NotEq => Arc::new( - arrow::compute::kernels::cmp::neq(&left, &right) - .context(ArrowSnafu { context: "neq" })?, - ), - Self::Lt => Arc::new( - arrow::compute::kernels::cmp::lt(&left, &right) - .context(ArrowSnafu { context: "lt" })?, - ), - Self::Lte => Arc::new( - arrow::compute::kernels::cmp::lt_eq(&left, &right) - .context(ArrowSnafu { context: "lte" })?, - ), - Self::Gt => Arc::new( - arrow::compute::kernels::cmp::gt(&left, &right) - .context(ArrowSnafu { context: "gt" })?, - ), - Self::Gte => Arc::new( - arrow::compute::kernels::cmp::gt_eq(&left, &right) - .context(ArrowSnafu { context: "gte" })?, - ), - - Self::AddInt16 - | Self::AddInt32 - | Self::AddInt64 - | Self::AddUInt16 - | Self::AddUInt32 - | Self::AddUInt64 - | Self::AddFloat32 - | Self::AddFloat64 => arrow::compute::kernels::numeric::add(&left, &right) - .context(ArrowSnafu { context: "add" })?, - - Self::SubInt16 - | Self::SubInt32 - | Self::SubInt64 - | Self::SubUInt16 - | Self::SubUInt32 - | Self::SubUInt64 - | Self::SubFloat32 - | Self::SubFloat64 => arrow::compute::kernels::numeric::sub(&left, &right) - .context(ArrowSnafu { context: "sub" })?, - - Self::MulInt16 - | Self::MulInt32 - | Self::MulInt64 - | Self::MulUInt16 - | Self::MulUInt32 - | Self::MulUInt64 - | Self::MulFloat32 - | Self::MulFloat64 => arrow::compute::kernels::numeric::mul(&left, &right) - .context(ArrowSnafu { context: "mul" })?, - - Self::DivInt16 - | Self::DivInt32 - | Self::DivInt64 - | Self::DivUInt16 - | Self::DivUInt32 - | Self::DivUInt64 - | Self::DivFloat32 - | Self::DivFloat64 => arrow::compute::kernels::numeric::div(&left, &right) - .context(ArrowSnafu { context: "div" })?, - - Self::ModInt16 - | Self::ModInt32 - | Self::ModInt64 - | Self::ModUInt16 - | Self::ModUInt32 - | Self::ModUInt64 => arrow::compute::kernels::numeric::rem(&left, &right) - .context(ArrowSnafu { context: "rem" })?, - }; - - let vector = Helper::try_into_vector(arrow_array).context(DataTypeSnafu { - msg: "Fail to convert to Vector", - })?; - Ok(vector) - } - - /// Evaluate the function with given values and expression - /// - /// # Arguments - /// - /// - `values`: The values to be used in the evaluation - /// - /// - `expr1`: The first arg to be evaluated, will extract the value from the `values` and evaluate the expression - /// - /// - `expr2`: The second arg to be evaluated - pub fn eval( - &self, - values: &[Value], - expr1: &ScalarExpr, - expr2: &ScalarExpr, - ) -> Result { - let left = expr1.eval(values)?; - let right = expr2.eval(values)?; - match self { - Self::Eq => Ok(Value::from(left == right)), - Self::NotEq => Ok(Value::from(left != right)), - Self::Lt => Ok(Value::from(left < right)), - Self::Lte => Ok(Value::from(left <= right)), - Self::Gt => Ok(Value::from(left > right)), - Self::Gte => Ok(Value::from(left >= right)), - - Self::AddInt16 => Ok(add::(left, right)?), - Self::AddInt32 => Ok(add::(left, right)?), - Self::AddInt64 => Ok(add::(left, right)?), - Self::AddUInt16 => Ok(add::(left, right)?), - Self::AddUInt32 => Ok(add::(left, right)?), - Self::AddUInt64 => Ok(add::(left, right)?), - Self::AddFloat32 => Ok(add::(left, right)?), - Self::AddFloat64 => Ok(add::(left, right)?), - - Self::SubInt16 => Ok(sub::(left, right)?), - Self::SubInt32 => Ok(sub::(left, right)?), - Self::SubInt64 => Ok(sub::(left, right)?), - Self::SubUInt16 => Ok(sub::(left, right)?), - Self::SubUInt32 => Ok(sub::(left, right)?), - Self::SubUInt64 => Ok(sub::(left, right)?), - Self::SubFloat32 => Ok(sub::(left, right)?), - Self::SubFloat64 => Ok(sub::(left, right)?), - - Self::MulInt16 => Ok(mul::(left, right)?), - Self::MulInt32 => Ok(mul::(left, right)?), - Self::MulInt64 => Ok(mul::(left, right)?), - Self::MulUInt16 => Ok(mul::(left, right)?), - Self::MulUInt32 => Ok(mul::(left, right)?), - Self::MulUInt64 => Ok(mul::(left, right)?), - Self::MulFloat32 => Ok(mul::(left, right)?), - Self::MulFloat64 => Ok(mul::(left, right)?), - - Self::DivInt16 => Ok(div::(left, right)?), - Self::DivInt32 => Ok(div::(left, right)?), - Self::DivInt64 => Ok(div::(left, right)?), - Self::DivUInt16 => Ok(div::(left, right)?), - Self::DivUInt32 => Ok(div::(left, right)?), - Self::DivUInt64 => Ok(div::(left, right)?), - Self::DivFloat32 => Ok(div::(left, right)?), - Self::DivFloat64 => Ok(div::(left, right)?), - - Self::ModInt16 => Ok(rem::(left, right)?), - Self::ModInt32 => Ok(rem::(left, right)?), - Self::ModInt64 => Ok(rem::(left, right)?), - Self::ModUInt16 => Ok(rem::(left, right)?), - Self::ModUInt32 => Ok(rem::(left, right)?), - Self::ModUInt64 => Ok(rem::(left, right)?), - } - } - - /// Reverse the comparison operator, i.e. `a < b` becomes `b > a`, - /// equal and not equal are unchanged. - pub fn reverse_compare(&self) -> Result { - let ret = match &self { - BinaryFunc::Eq => BinaryFunc::Eq, - BinaryFunc::NotEq => BinaryFunc::NotEq, - BinaryFunc::Lt => BinaryFunc::Gt, - BinaryFunc::Lte => BinaryFunc::Gte, - BinaryFunc::Gt => BinaryFunc::Lt, - BinaryFunc::Gte => BinaryFunc::Lte, - _ => { - return InvalidQuerySnafu { - reason: format!("Expect a comparison operator, found {:?}", self), - } - .fail(); - } - }; - Ok(ret) - } -} - -/// VariadicFunc is a function that takes a variable number of arguments. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Deserialize, Serialize, Hash)] -pub enum VariadicFunc { - And, - Or, -} - -impl VariadicFunc { - /// Return the signature of the function - pub fn signature(&self) -> Signature { - Signature { - input: smallvec![ConcreteDataType::boolean_datatype()], - output: ConcreteDataType::boolean_datatype(), - generic_fn: match self { - Self::And => GenericFn::And, - Self::Or => GenericFn::Or, - }, - } - } - - pub fn is_valid_func_name(name: &str) -> bool { - matches!(name.to_lowercase().as_str(), "and" | "or") - } - - /// Create a VariadicFunc from a string of the function name and given argument types(optional) - pub fn from_str_and_types( - name: &str, - arg_types: &[Option], - ) -> Result { - // TODO(discord9): future variadic funcs to be added might need to check arg_types - let _ = arg_types; - match name { - "and" => Ok(Self::And), - "or" => Ok(Self::Or), - _ => InvalidQuerySnafu { - reason: format!("Unknown variadic function: {}", name), - } - .fail(), - } - } - - pub fn eval_batch(&self, batch: &Batch, exprs: &[ScalarExpr]) -> Result { - ensure!( - !exprs.is_empty(), - InvalidArgumentSnafu { - reason: format!("Variadic function {:?} requires at least 1 arguments", self) - } - ); - let args = exprs - .iter() - .map(|expr| expr.eval_batch(batch).map(|v| v.to_arrow_array())) - .collect::, _>>()?; - let mut iter = args.into_iter(); - - let first = iter.next().unwrap(); - let mut left = first - .as_any() - .downcast_ref::() - .context({ - TypeMismatchSnafu { - expected: ConcreteDataType::boolean_datatype(), - actual: ConcreteDataType::from_arrow_type(first.data_type()), - } - })? - .clone(); - - for right in iter { - let right = right.as_any().downcast_ref::().context({ - TypeMismatchSnafu { - expected: ConcreteDataType::boolean_datatype(), - actual: ConcreteDataType::from_arrow_type(right.data_type()), - } - })?; - left = match self { - Self::And => { - arrow::compute::and(&left, right).context(ArrowSnafu { context: "and" })? - } - Self::Or => { - arrow::compute::or(&left, right).context(ArrowSnafu { context: "or" })? - } - } - } - - Ok(Arc::new(BooleanVector::from(left))) - } - - /// Evaluate the function with given values and expressions - pub fn eval(&self, values: &[Value], exprs: &[ScalarExpr]) -> Result { - match self { - VariadicFunc::And => and(values, exprs), - VariadicFunc::Or => or(values, exprs), - } - } -} - -fn and(values: &[Value], exprs: &[ScalarExpr]) -> Result { - // If any is false, then return false. Else, if any is null, then return null. Else, return true. - let mut null = false; - for expr in exprs { - match expr.eval(values) { - Ok(Value::Boolean(true)) => {} - Ok(Value::Boolean(false)) => return Ok(Value::Boolean(false)), // short-circuit - Ok(Value::Null) => null = true, - Err(this_err) => { - return Err(this_err); - } // retain first error encountered - Ok(x) => InvalidArgumentSnafu { - reason: format!( - "`and()` only support boolean type, found value {:?} of type {:?}", - x, - x.data_type() - ), - } - .fail()?, - } - } - match null { - true => Ok(Value::Null), - false => Ok(Value::Boolean(true)), - } -} - -fn or(values: &[Value], exprs: &[ScalarExpr]) -> Result { - // If any is false, then return false. Else, if any is null, then return null. Else, return true. - let mut null = false; - for expr in exprs { - match expr.eval(values) { - Ok(Value::Boolean(true)) => return Ok(Value::Boolean(true)), // short-circuit - Ok(Value::Boolean(false)) => {} - Ok(Value::Null) => null = true, - Err(this_err) => { - return Err(this_err); - } // retain first error encountered - Ok(x) => InvalidArgumentSnafu { - reason: format!( - "`or()` only support boolean type, found value {:?} of type {:?}", - x, - x.data_type() - ), - } - .fail()?, - } - } - match null { - true => Ok(Value::Null), - false => Ok(Value::Boolean(false)), - } -} - -fn add(left: Value, right: Value) -> Result -where - T: TryFrom + num_traits::Num, - Value: From, -{ - let left = T::try_from(left).map_err(|e| TryFromValueSnafu { msg: e.to_string() }.build())?; - let right = T::try_from(right).map_err(|e| TryFromValueSnafu { msg: e.to_string() }.build())?; - Ok(Value::from(left + right)) -} - -fn sub(left: Value, right: Value) -> Result -where - T: TryFrom + num_traits::Num, - Value: From, -{ - let left = T::try_from(left).map_err(|e| TryFromValueSnafu { msg: e.to_string() }.build())?; - let right = T::try_from(right).map_err(|e| TryFromValueSnafu { msg: e.to_string() }.build())?; - Ok(Value::from(left - right)) -} - -fn mul(left: Value, right: Value) -> Result -where - T: TryFrom + num_traits::Num, - Value: From, -{ - let left = T::try_from(left).map_err(|e| TryFromValueSnafu { msg: e.to_string() }.build())?; - let right = T::try_from(right).map_err(|e| TryFromValueSnafu { msg: e.to_string() }.build())?; - Ok(Value::from(left * right)) -} - -fn div(left: Value, right: Value) -> Result -where - T: TryFrom + num_traits::Num, - >::Error: std::fmt::Debug, - Value: From, -{ - let left = T::try_from(left).map_err(|e| TryFromValueSnafu { msg: e.to_string() }.build())?; - let right = T::try_from(right).map_err(|e| TryFromValueSnafu { msg: e.to_string() }.build())?; - if right.is_zero() { - return Err(DivisionByZeroSnafu {}.build()); - } - Ok(Value::from(left / right)) -} - -fn rem(left: Value, right: Value) -> Result -where - T: TryFrom + num_traits::Num, - >::Error: std::fmt::Debug, - Value: From, -{ - let left = T::try_from(left).map_err(|e| TryFromValueSnafu { msg: e.to_string() }.build())?; - let right = T::try_from(right).map_err(|e| TryFromValueSnafu { msg: e.to_string() }.build())?; - Ok(Value::from(left % right)) -} - -#[cfg(test)] -mod test { - use std::sync::Arc; - - use datatypes::vectors::Vector; - use pretty_assertions::assert_eq; - - use super::*; - - #[test] - fn test_tumble_batch() { - let timestamp_vector = TimestampMillisecondVector::from_vec(vec![1, 2, 10, 13, 14, 20, 25]); - let tumble_start = UnaryFunc::TumbleWindowFloor { - window_size: Duration::from_millis(10), - start_time: None, - }; - let tumble_end = UnaryFunc::TumbleWindowCeiling { - window_size: Duration::from_millis(10), - start_time: None, - }; - - let len = timestamp_vector.len(); - let batch = Batch::try_new(vec![Arc::new(timestamp_vector)], len).unwrap(); - let arg = ScalarExpr::Column(0); - - let start = tumble_start.eval_batch(&batch, &arg).unwrap(); - let end = tumble_end.eval_batch(&batch, &arg).unwrap(); - assert_eq!( - start.to_arrow_array().as_ref(), - TimestampMillisecondVector::from_vec(vec![0, 0, 10, 10, 10, 20, 20]) - .to_arrow_array() - .as_ref() - ); - - assert_eq!( - end.to_arrow_array().as_ref(), - TimestampMillisecondVector::from_vec(vec![10, 10, 20, 20, 20, 30, 30]) - .to_arrow_array() - .as_ref() - ); - - let ts_ms_vector = TimestampMillisecondVector::from_vec(vec![1, 2, 10, 13, 14, 20, 25]); - let batch = Batch::try_new(vec![Arc::new(ts_ms_vector)], len).unwrap(); - - let start = tumble_start.eval_batch(&batch, &arg).unwrap(); - let end = tumble_end.eval_batch(&batch, &arg).unwrap(); - - assert_eq!( - start.to_arrow_array().as_ref(), - TimestampMillisecondVector::from_vec(vec![0, 0, 10, 10, 10, 20, 20]) - .to_arrow_array() - .as_ref() - ); - - assert_eq!( - end.to_arrow_array().as_ref(), - TimestampMillisecondVector::from_vec(vec![10, 10, 20, 20, 20, 30, 30]) - .to_arrow_array() - .as_ref() - ); - } - - #[test] - fn test_num_ops() { - let left = Value::from(10); - let right = Value::from(3); - let res = add::(left.clone(), right.clone()).unwrap(); - assert_eq!(res, Value::from(13)); - let res = sub::(left.clone(), right.clone()).unwrap(); - assert_eq!(res, Value::from(7)); - let res = mul::(left.clone(), right.clone()).unwrap(); - assert_eq!(res, Value::from(30)); - let res = div::(left.clone(), right.clone()).unwrap(); - assert_eq!(res, Value::from(3)); - let res = rem::(left, right).unwrap(); - assert_eq!(res, Value::from(1)); - - let values = vec![Value::from(true), Value::from(false)]; - let exprs = vec![ScalarExpr::Column(0), ScalarExpr::Column(1)]; - let res = and(&values, &exprs).unwrap(); - assert_eq!(res, Value::from(false)); - let res = or(&values, &exprs).unwrap(); - assert_eq!(res, Value::from(true)); - } - - /// test if the binary function specialization works - /// whether from direct type or from the expression that is literal - #[test] - fn test_binary_func_spec() { - assert_eq!( - BinaryFunc::from_str_expr_and_type( - "add", - &[ScalarExpr::Column(0), ScalarExpr::Column(0)], - &[ - Some(ConcreteDataType::int32_datatype()), - Some(ConcreteDataType::int32_datatype()) - ] - ) - .unwrap(), - (BinaryFunc::AddInt32, BinaryFunc::AddInt32.signature()) - ); - - assert_eq!( - BinaryFunc::from_str_expr_and_type( - "add", - &[ScalarExpr::Column(0), ScalarExpr::Column(0)], - &[Some(ConcreteDataType::int32_datatype()), None] - ) - .unwrap(), - (BinaryFunc::AddInt32, BinaryFunc::AddInt32.signature()) - ); - - assert_eq!( - BinaryFunc::from_str_expr_and_type( - "add", - &[ScalarExpr::Column(0), ScalarExpr::Column(0)], - &[Some(ConcreteDataType::int32_datatype()), None] - ) - .unwrap(), - (BinaryFunc::AddInt32, BinaryFunc::AddInt32.signature()) - ); - - assert_eq!( - BinaryFunc::from_str_expr_and_type( - "add", - &[ScalarExpr::Column(0), ScalarExpr::Column(0)], - &[Some(ConcreteDataType::int32_datatype()), None] - ) - .unwrap(), - (BinaryFunc::AddInt32, BinaryFunc::AddInt32.signature()) - ); - - assert_eq!( - BinaryFunc::from_str_expr_and_type( - "add", - &[ - ScalarExpr::Literal(Value::from(1i32), ConcreteDataType::int32_datatype()), - ScalarExpr::Column(0) - ], - &[None, None] - ) - .unwrap(), - (BinaryFunc::AddInt32, BinaryFunc::AddInt32.signature()) - ); - - // this testcase make sure the specialization can find actual type from expression and fill in signature - assert_eq!( - BinaryFunc::from_str_expr_and_type( - "equal", - &[ - ScalarExpr::Literal(Value::from(1i32), ConcreteDataType::int32_datatype()), - ScalarExpr::Column(0) - ], - &[None, None] - ) - .unwrap(), - ( - BinaryFunc::Eq, - Signature { - input: smallvec![ - ConcreteDataType::int32_datatype(), - ConcreteDataType::int32_datatype() - ], - output: ConcreteDataType::boolean_datatype(), - generic_fn: GenericFn::Eq - } - ) - ); - - matches!( - BinaryFunc::from_str_expr_and_type( - "add", - &[ScalarExpr::Column(0), ScalarExpr::Column(0)], - &[None, None] - ), - Err(Error::InvalidQuery { .. }) - ); - } - - #[test] - fn test_cast_int() { - let interval = cast( - Value::from("1 second"), - &ConcreteDataType::interval_day_time_datatype(), - ) - .unwrap(); - assert_eq!( - interval, - Value::from(common_time::IntervalDayTime::new(0, 1000)) - ); - } -} diff --git a/src/flow/src/expr/id.rs b/src/flow/src/expr/id.rs deleted file mode 100644 index f88baa70ded..00000000000 --- a/src/flow/src/expr/id.rs +++ /dev/null @@ -1,43 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! `Id` is used to identify a dataflow component in plan like `Plan::Get{id: Id}`, this could be a source of data for an arrangement. - -use serde::{Deserialize, Serialize}; - -/// Global id's scope is in Current Flow node, and is cross-dataflow -#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Serialize, Deserialize)] -pub enum GlobalId { - /// System namespace. - System(u64), - /// User namespace. - User(u64), - /// Transient namespace. - Transient(u64), - /// Dummy id for query being explained - Explain, -} - -/// Local id is used in local scope created by `Plan::Let{id: LocalId, value, body}` -#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Serialize, Deserialize)] -pub struct LocalId(pub(crate) u64); - -/// Id is used to identify a dataflow component in plan like `Plan::Get{id: Id}` -#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Serialize, Deserialize)] -pub enum Id { - /// An identifier that refers to a local component of a dataflow. - Local(LocalId), - /// An identifier that refers to a global dataflow. - Global(GlobalId), -} diff --git a/src/flow/src/expr/linear.rs b/src/flow/src/expr/linear.rs deleted file mode 100644 index 9729433518a..00000000000 --- a/src/flow/src/expr/linear.rs +++ /dev/null @@ -1,1174 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! define MapFilterProject which is a compound operator that can be applied row-by-row. - -use std::collections::{BTreeMap, BTreeSet}; - -use arrow::array::BooleanArray; -use arrow::buffer::BooleanBuffer; -use arrow::compute::FilterBuilder; -use common_telemetry::trace; -use datatypes::prelude::ConcreteDataType; -use datatypes::value::Value; -use datatypes::vectors::{BooleanVector, Helper}; -use itertools::Itertools; -use snafu::{OptionExt, ResultExt, ensure}; - -use crate::error::{Error, InvalidQuerySnafu}; -use crate::expr::error::{ArrowSnafu, DataTypeSnafu, EvalError, InternalSnafu, TypeMismatchSnafu}; -use crate::expr::{Batch, InvalidArgumentSnafu, ScalarExpr}; -use crate::repr::{self, Diff, Row, value_to_internal_ts}; - -/// A compound operator that can be applied row-by-row. -/// -/// In practice, this operator is a sequence of map, filter, and project in arbitrary order, -/// which can and is stored by reordering the sequence's -/// apply order to a `map` first, `filter` second and `project` third order. -/// -/// input is a row(a sequence of values), which is also being used for store intermediate results, -/// like `map` operator can append new columns to the row according to it's expressions, -/// `filter` operator decide whether this entire row can even be output by decide whether the row satisfy the predicates, -/// `project` operator decide which columns of the row should be output. -/// -/// This operator integrates the map, filter, and project operators. -/// It applies a sequences of map expressions, which are allowed to -/// refer to previous expressions, interleaved with predicates which -/// must be satisfied for an output to be produced. If all predicates -/// evaluate to `Value::Boolean(True)` the data at the identified columns are -/// collected and produced as output in a packed `Row`. -/// -/// This operator is a "builder" and its contents may contain expressions -/// that are not yet executable. For example, it may contain temporal -/// expressions in `self.expressions`, even though this is not something -/// we can directly evaluate. The plan creation methods will defensively -/// ensure that the right thing happens. -#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)] -pub struct MapFilterProject { - /// A sequence of expressions that should be appended to the row. - /// - /// Many of these expressions may not be produced in the output, - /// and may only be present as common subexpressions. - pub expressions: Vec, - /// Expressions that must evaluate to `Datum::True` for the output - /// row to be produced. - /// - /// Each entry is prepended with a column identifier indicating - /// the column *before* which the predicate should first be applied. - /// Most commonly this would be one plus the largest column identifier - /// in the predicate's referred columns, but it could be larger to implement - /// guarded evaluation of predicates. - /// Put it in another word, the first element of the tuple means - /// the predicates can't be evaluated until that number of columns is formed. - /// - /// This list should be sorted by the first field. - pub predicates: Vec<(usize, ScalarExpr)>, - /// A sequence of column identifiers whose data form the output row. - pub projection: Vec, - /// The expected number of input columns. - /// - /// This is needed to ensure correct identification of newly formed - /// columns in the output. - pub input_arity: usize, -} - -impl MapFilterProject { - /// Create a no-op operator for an input of a supplied arity. - pub fn new(input_arity: usize) -> Self { - Self { - expressions: Vec::new(), - predicates: Vec::new(), - projection: (0..input_arity).collect(), - input_arity, - } - } - - pub fn get_nth_expr(&self, n: usize) -> Option { - let idx = *self.projection.get(n)?; - if idx < self.input_arity { - Some(ScalarExpr::Column(idx)) - } else { - // find direct ref to input's expr - - let mut expr = self.expressions.get(idx - self.input_arity)?; - loop { - match expr { - ScalarExpr::Column(prev) => { - if *prev < self.input_arity { - return Some(ScalarExpr::Column(*prev)); - } else { - expr = self.expressions.get(*prev - self.input_arity)?; - continue; - } - } - _ => return Some(expr.clone()), - } - } - } - } - - /// The number of columns expected in the output row. - pub fn output_arity(&self) -> usize { - self.projection.len() - } - - /// Given two mfps, return an mfp that applies one - /// followed by the other. - /// Note that the arguments are in the opposite order - /// from how function composition is usually written in mathematics. - pub fn compose(before: Self, after: Self) -> Result { - let (m, f, p) = after.into_map_filter_project(); - before.map(m)?.filter(f)?.project(p) - } - - /// True if the operator describes the identity transformation. - pub fn is_identity(&self) -> bool { - self.expressions.is_empty() - && self.predicates.is_empty() - // identity if projection is the identity permutation - && self.projection.len() == self.input_arity - && self.projection.iter().enumerate().all(|(i, p)| i == *p) - } - - /// Retain only the indicated columns in the presented order. - /// - /// i.e. before: `self.projection = [1, 2, 0], columns = [1, 0]` - /// ```mermaid - /// flowchart TD - /// col-0 - /// col-1 - /// col-2 - /// projection --> |0|col-1 - /// projection --> |1|col-2 - /// projection --> |2|col-0 - /// ``` - /// - /// after: `self.projection = [2, 1]` - /// ```mermaid - /// flowchart TD - /// col-0 - /// col-1 - /// col-2 - /// project("project:[1,2,0]") - /// project - /// project -->|0| col-1 - /// project -->|1| col-2 - /// project -->|2| col-0 - /// new_project("apply new project:[1,0]") - /// new_project -->|0| col-2 - /// new_project -->|1| col-1 - /// ``` - pub fn project(mut self, columns: I) -> Result - where - I: IntoIterator + std::fmt::Debug, - { - self.projection = columns - .into_iter() - .map(|c| self.projection.get(c).cloned().ok_or(c)) - .collect::, _>>() - .map_err(|c| { - InvalidQuerySnafu { - reason: format!( - "column index {} out of range, expected at most {} columns", - c, - self.projection.len() - ), - } - .build() - })?; - Ok(self) - } - - /// Retain only rows satisfying these predicates. - /// - /// This method introduces predicates as eagerly as they can be evaluated, - /// which may not be desired for predicates that may cause exceptions. - /// If fine manipulation is required, the predicates can be added manually. - /// - /// simply added to the end of the predicates list - /// - /// while paying attention to column references maintained by `self.projection` - /// - /// so `self.projection = [1, 2, 0], filter = [0]+[1]>0`: - /// becomes: - /// ```mermaid - /// flowchart TD - /// col-0 - /// col-1 - /// col-2 - /// project("first project:[1,2,0]") - /// project - /// project -->|0| col-1 - /// project -->|1| col-2 - /// project -->|2| col-0 - /// filter("then filter:[0]+[1]>0") - /// filter -->|0| col-1 - /// filter --> |1| col-2 - /// ``` - pub fn filter(mut self, predicates: I) -> Result - where - I: IntoIterator, - { - for mut predicate in predicates { - // Correct column references. - predicate.permute(&self.projection[..])?; - - // Validate column references. - let referred_columns = predicate.get_all_ref_columns(); - for c in referred_columns.iter() { - // current row len include input columns and previous number of expressions - let cur_row_len = self.input_arity + self.expressions.len(); - ensure!( - *c < cur_row_len, - InvalidQuerySnafu { - reason: format!( - "column index {} out of range, expected at most {} columns", - c, cur_row_len - ) - } - ); - } - - // Insert predicate as eagerly as it can be evaluated: - // just after the largest column in its support is formed. - let max_support = referred_columns - .into_iter() - .max() - .map(|c| c + 1) - .unwrap_or(0); - self.predicates.push((max_support, predicate)) - } - // Stable sort predicates by position at which they take effect. - self.predicates - .sort_by_key(|(position, _predicate)| *position); - Ok(self) - } - - /// Append the result of evaluating expressions to each row. - /// - /// simply append `expressions` to `self.expressions` - /// - /// while paying attention to column references maintained by `self.projection` - /// - /// hence, before apply map with a previously non-trivial projection would be like: - /// before: - /// ```mermaid - /// flowchart TD - /// col-0 - /// col-1 - /// col-2 - /// projection --> |0|col-1 - /// projection --> |1|col-2 - /// projection --> |2|col-0 - /// ``` - /// after apply map: - /// ```mermaid - /// flowchart TD - /// col-0 - /// col-1 - /// col-2 - /// project("project:[1,2,0]") - /// project - /// project -->|0| col-1 - /// project -->|1| col-2 - /// project -->|2| col-0 - /// map("map:[0]/[1]/[2]") - /// map -->|0|col-1 - /// map -->|1|col-2 - /// map -->|2|col-0 - /// ``` - pub fn map(mut self, expressions: I) -> Result - where - I: IntoIterator, - { - for mut expression in expressions { - // Correct column references. - expression.permute(&self.projection[..])?; - - // Validate column references. - for c in expression.get_all_ref_columns().into_iter() { - // current row len include input columns and previous number of expressions - let current_row_len = self.input_arity + self.expressions.len(); - ensure!( - c < current_row_len, - InvalidQuerySnafu { - reason: format!( - "column index {} out of range, expected at most {} columns", - c, current_row_len - ) - } - ); - } - - // Introduce expression and produce as output. - self.expressions.push(expression); - // Expression by default is projected to output. - let cur_expr_col_num = self.input_arity + self.expressions.len() - 1; - self.projection.push(cur_expr_col_num); - } - - Ok(self) - } - - /// Like [`MapFilterProject::as_map_filter_project`], but consumes `self` rather than cloning. - pub fn into_map_filter_project(self) -> (Vec, Vec, Vec) { - let predicates = self - .predicates - .into_iter() - .map(|(_pos, predicate)| predicate) - .collect(); - (self.expressions, predicates, self.projection) - } - - /// As the arguments to `Map`, `Filter`, and `Project` operators. - /// - /// In principle, this operator can be implemented as a sequence of - /// more elemental operators, likely less efficiently. - pub fn as_map_filter_project(&self) -> (Vec, Vec, Vec) { - self.clone().into_map_filter_project() - } -} - -impl MapFilterProject { - /// Convert the `MapFilterProject` into a safe evaluation plan. Marking it safe to evaluate. - pub fn into_safe(self) -> SafeMfpPlan { - SafeMfpPlan { mfp: self } - } - - /// Optimize the `MapFilterProject` in place. - pub fn optimize(&mut self) { - // TODO(discord9): optimize - } - /// get the mapping of old columns to new columns after the mfp - pub fn get_old_to_new_mapping(&self) -> BTreeMap { - BTreeMap::from_iter( - self.projection - .clone() - .into_iter() - .enumerate() - .map(|(new, old)| { - // `projection` give the new -> old mapping - let mut old = old; - // trace back to the original column - // since there maybe indirect ref to old columns like - // col 2 <- expr=col(2) at pos col 4 <- expr=col(4) at pos col 6 - // ideally such indirect ref should be optimize away - // TODO(discord9): refactor this after impl `optimize()` - while let Some(ScalarExpr::Column(prev)) = if old >= self.input_arity { - // get the correspond expr if not a original column - self.expressions.get(old - self.input_arity) - } else { - // we don't care about non column ref case since only need old to new column mapping - // in which case, the old->new mapping remain the same - None - } { - old = *prev; - if old < self.input_arity { - break; - } - } - (old, new) - }), - ) - } - - /// Lists input columns whose values are used in outputs. - /// - /// It is entirely appropriate to determine the demand of an instance - /// and then both apply a projection to the subject of the instance and - /// `self.permute` this instance. - pub fn demand(&self) -> BTreeSet { - let mut demanded = BTreeSet::new(); - // first, get all columns referenced by predicates - for (_index, pred) in self.predicates.iter() { - demanded.extend(pred.get_all_ref_columns()); - } - // then, get columns referenced by projection which is direct output - demanded.extend(self.projection.iter().cloned()); - - // check every expressions, if a expression is contained in demanded, then all columns it referenced should be added to demanded - for index in (0..self.expressions.len()).rev() { - if demanded.contains(&(self.input_arity + index)) { - demanded.extend(self.expressions[index].get_all_ref_columns()); - } - } - - // only keep demanded columns that are in input - demanded.retain(|col| col < &self.input_arity); - demanded - } - - /// Update input column references, due to an input projection or permutation. - /// - /// The `shuffle` argument remaps expected column identifiers to new locations, - /// with the expectation that `shuffle` describes all input columns, and so the - /// intermediate results will be able to start at position `shuffle.len()`. - /// - /// The supplied `shuffle` may not list columns that are not "demanded" by the - /// instance, and so we should ensure that `self` is optimized to not reference - /// columns that are not demanded. - pub fn permute( - &mut self, - mut shuffle: BTreeMap, - new_input_arity: usize, - ) -> Result<(), Error> { - // check shuffle is valid - let demand = self.demand(); - for d in demand { - ensure!( - shuffle.contains_key(&d), - InvalidQuerySnafu { - reason: format!( - "Demanded column {} is not in shuffle's keys: {:?}", - d, - shuffle.keys() - ) - } - ); - } - ensure!( - shuffle.len() <= new_input_arity, - InvalidQuerySnafu { - reason: format!( - "shuffle's length {} is greater than new_input_arity {}", - shuffle.len(), - self.input_arity - ) - } - ); - - // decompose self into map, filter, project for ease of manipulation - let (mut map, mut filter, mut project) = self.as_map_filter_project(); - for index in 0..map.len() { - // Intermediate columns are just shifted. - shuffle.insert(self.input_arity + index, new_input_arity + index); - } - - for expr in map.iter_mut() { - expr.permute_map(&shuffle)?; - } - for pred in filter.iter_mut() { - pred.permute_map(&shuffle)?; - } - let new_row_len = new_input_arity + map.len(); - for proj in project.iter_mut() { - ensure!( - shuffle[proj] < new_row_len, - InvalidQuerySnafu { - reason: format!( - "shuffled column index {} out of range, expected at most {} columns", - shuffle[proj], new_row_len - ) - } - ); - *proj = shuffle[proj]; - } - *self = Self::new(new_input_arity) - .map(map)? - .filter(filter)? - .project(project)?; - Ok(()) - } -} - -/// A wrapper type which indicates it is safe to simply evaluate all expressions. -#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)] -pub struct SafeMfpPlan { - /// the inner `MapFilterProject` that is safe to evaluate. - pub(crate) mfp: MapFilterProject, -} - -impl SafeMfpPlan { - /// See [`MapFilterProject::permute`]. - pub fn permute(&mut self, map: BTreeMap, new_arity: usize) -> Result<(), Error> { - self.mfp.permute(map, new_arity) - } - - /// similar to [`MapFilterProject::evaluate_into`], just in batch, and rows that don't pass the predicates are not included in the output. - /// - /// so it's not guaranteed that the output will have the same number of rows as the input. - pub fn eval_batch_into(&self, batch: &mut Batch) -> Result { - ensure!( - batch.column_count() == self.mfp.input_arity, - InvalidArgumentSnafu { - reason: format!( - "batch column length {} is not equal to input_arity {}", - batch.column_count(), - self.mfp.input_arity - ), - } - ); - - let passed_predicates = self.eval_batch_inner(batch)?; - let filter = FilterBuilder::new(passed_predicates.as_boolean_array()); - let pred = filter.build(); - let mut result = vec![]; - for col in batch.batch() { - let filtered = pred - .filter(col.to_arrow_array().as_ref()) - .with_context(|_| ArrowSnafu { - context: format!("failed to filter column for mfp operator {:?}", self), - })?; - result.push(Helper::try_into_vector(filtered).context(DataTypeSnafu { - msg: "Failed to convert arrow array to vector", - })?); - } - let projected = self - .mfp - .projection - .iter() - .map(|c| result[*c].clone()) - .collect_vec(); - let row_count = pred.count(); - - Batch::try_new(projected, row_count) - } - - /// similar to [`MapFilterProject::evaluate_into`], just in batch. - pub fn eval_batch_inner(&self, batch: &mut Batch) -> Result { - // mark the columns that have been evaluated and appended to the `batch` - let mut expression = 0; - // preds default to true and will be updated as we evaluate each predicate - let buf = BooleanBuffer::new_set(batch.row_count()); - let arr = BooleanArray::new(buf, None); - let mut all_preds = BooleanVector::from(arr); - - // to compute predicate, need to first compute all expressions used in predicates - for (support, predicate) in self.mfp.predicates.iter() { - while self.mfp.input_arity + expression < *support { - let expr_eval = self.mfp.expressions[expression].eval_batch(batch)?; - batch.batch_mut().push(expr_eval); - expression += 1; - } - let pred_vec = predicate.eval_batch(batch)?; - let pred_arr = pred_vec.to_arrow_array(); - let pred_arr = pred_arr.as_any().downcast_ref::().context({ - TypeMismatchSnafu { - expected: ConcreteDataType::boolean_datatype(), - actual: pred_vec.data_type(), - } - })?; - let all_arr = all_preds.as_boolean_array(); - let res_arr = arrow::compute::and(all_arr, pred_arr).context(ArrowSnafu { - context: format!("failed to compute predicate for mfp operator {:?}", self), - })?; - all_preds = BooleanVector::from(res_arr); - } - - // while evaluated expressions are less than total expressions, keep evaluating - while expression < self.mfp.expressions.len() { - let expr_eval = self.mfp.expressions[expression].eval_batch(batch)?; - batch.batch_mut().push(expr_eval); - expression += 1; - } - - Ok(all_preds) - } - - /// Evaluates the linear operator on a supplied list of datums. - /// - /// The arguments are the initial datums associated with the row, - /// and an appropriately lifetimed arena for temporary allocations - /// needed by scalar evaluation. - /// - /// An `Ok` result will either be `None` if any predicate did not - /// evaluate to `Value::Boolean(true)`, or the values of the columns listed - /// by `self.projection` if all predicates passed. If an error - /// occurs in the evaluation it is returned as an `Err` variant. - /// As the evaluation exits early with failed predicates, it may - /// miss some errors that would occur later in evaluation. - /// - /// The `row` is not cleared first, but emptied if the function - /// returns `Ok(Some(row)). - #[inline(always)] - pub fn evaluate_into( - &self, - values: &mut Vec, - row_buf: &mut Row, - ) -> Result, EvalError> { - ensure!( - values.len() == self.mfp.input_arity, - InvalidArgumentSnafu { - reason: format!( - "values length {} is not equal to input_arity {}", - values.len(), - self.mfp.input_arity - ), - } - ); - let passed_predicates = self.evaluate_inner(values)?; - - if !passed_predicates { - Ok(None) - } else { - row_buf.clear(); - row_buf.extend(self.mfp.projection.iter().map(|c| values[*c].clone())); - Ok(Some(row_buf.clone())) - } - } - - /// Populates `values` with `self.expressions` and tests `self.predicates`. - /// - /// This does not apply `self.projection`, which is up to the calling method. - pub fn evaluate_inner(&self, values: &mut Vec) -> Result { - let mut expression = 0; - for (support, predicate) in self.mfp.predicates.iter() { - while self.mfp.input_arity + expression < *support { - values.push(self.mfp.expressions[expression].eval(&values[..])?); - expression += 1; - } - if predicate.eval(&values[..])? != Value::Boolean(true) { - return Ok(false); - } - } - // while evaluated expressions are less than total expressions, keep evaluating - while expression < self.mfp.expressions.len() { - values.push(self.mfp.expressions[expression].eval(&values[..])?); - expression += 1; - } - Ok(true) - } -} - -impl std::ops::Deref for SafeMfpPlan { - type Target = MapFilterProject; - fn deref(&self) -> &Self::Target { - &self.mfp - } -} - -/// Predicates partitioned into temporal and non-temporal. -/// -/// Temporal predicates require some recognition to determine their -/// structure, and it is best to do that once and re-use the results. -/// -/// There are restrictions on the temporal predicates we currently support. -/// They must directly constrain `MzNow` from below or above, -/// by expressions that do not themselves contain `MzNow`. -/// Conjunctions of such constraints are also ok. -#[derive(Clone, Debug, PartialEq)] -pub struct MfpPlan { - /// Normal predicates to evaluate on `&[Datum]` and expect `Ok(Datum::True)`. - pub(crate) mfp: SafeMfpPlan, - /// TODO(discord9): impl temporal filter later - /// Expressions that when evaluated lower-bound `MzNow`. - pub(crate) lower_bounds: Vec, - /// Expressions that when evaluated upper-bound `MzNow`. - pub(crate) upper_bounds: Vec, -} - -impl MfpPlan { - /// Indicates if the `MfpPlan` contains temporal predicates. That is have outputs that may occur in future. - pub fn is_temporal(&self) -> bool { - !self.lower_bounds.is_empty() || !self.upper_bounds.is_empty() - } - /// find `now` in `predicates` and put them into lower/upper temporal bounds for temporal filter to use - pub fn create_from(mut mfp: MapFilterProject) -> Result { - let mut lower_bounds = Vec::new(); - let mut upper_bounds = Vec::new(); - - let mut temporal = Vec::new(); - - // Optimize, to ensure that temporal predicates are move in to `mfp.predicates`. - mfp.optimize(); - - mfp.predicates.retain(|(_position, predicate)| { - if predicate.contains_temporal() { - temporal.push(predicate.clone()); - false - } else { - true - } - }); - - for predicate in temporal { - let (lower, upper) = predicate.extract_bound()?; - lower_bounds.extend(lower); - upper_bounds.extend(upper); - } - Ok(Self { - mfp: SafeMfpPlan { mfp }, - lower_bounds, - upper_bounds, - }) - } - - /// Indicates if the planned `MapFilterProject` emits exactly its inputs as outputs. - pub fn is_identity(&self) -> bool { - self.mfp.mfp.is_identity() && self.lower_bounds.is_empty() && self.upper_bounds.is_empty() - } - - /// if `lower_bound <= sys_time < upper_bound`, return `[(data, sys_time, +1), (data, min_upper_bound, -1)]` - /// - /// else if `sys_time < lower_bound`, return `[(data, lower_bound, +1), (data, min_upper_bound, -1)]` - /// - /// else if `sys_time >= upper_bound`, return `[None, None]` - /// - /// if eval error appeal in any of those process, corresponding result will be `Err` - pub fn evaluate>( - &self, - values: &mut Vec, - sys_time: repr::Timestamp, - diff: Diff, - ) -> impl Iterator> - { - match self.mfp.evaluate_inner(values) { - Err(e) => { - return Some(Err((e.into(), sys_time, diff))) - .into_iter() - .chain(None); - } - Ok(true) => {} - Ok(false) => { - return None.into_iter().chain(None); - } - } - - let mut lower_bound = sys_time; - let mut upper_bound = None; - - // Track whether we have seen a null in either bound, as this should - // prevent the record from being produced at any time. - let mut null_eval = false; - let ret_err = |e: EvalError| { - Some(Err((e.into(), sys_time, diff))) - .into_iter() - .chain(None) - }; - for l in self.lower_bounds.iter() { - match l.eval(values) { - Ok(v) => { - if v.is_null() { - null_eval = true; - continue; - } - match value_to_internal_ts(v) { - Ok(ts) => lower_bound = lower_bound.max(ts), - Err(e) => return ret_err(e), - } - } - Err(e) => return ret_err(e), - }; - } - - for u in self.upper_bounds.iter() { - if upper_bound != Some(lower_bound) { - match u.eval(values) { - Err(e) => return ret_err(e), - Ok(val) => { - if val.is_null() { - null_eval = true; - continue; - } - let ts = match value_to_internal_ts(val) { - Ok(ts) => ts, - Err(e) => return ret_err(e), - }; - if let Some(upper) = upper_bound { - upper_bound = Some(upper.min(ts)); - } else { - upper_bound = Some(ts); - } - // Force the upper bound to be at least the lower - // bound. - if upper_bound.is_some() && upper_bound < Some(lower_bound) { - upper_bound = Some(lower_bound); - } - } - } - } - } - - if Some(lower_bound) != upper_bound && !null_eval { - if self.mfp.mfp.projection.iter().any(|c| values.len() <= *c) { - trace!("values={:?}, mfp={:?}", &values, &self.mfp.mfp); - let err = InternalSnafu { - reason: format!( - "Index out of bound for mfp={:?} and values={:?}", - &self.mfp.mfp, &values - ), - } - .build(); - return ret_err(err); - } - // safety: already checked that `projection` is not out of bound - let res_row = Row::pack(self.mfp.mfp.projection.iter().map(|c| values[*c].clone())); - let upper_opt = - upper_bound.map(|upper_bound| Ok((res_row.clone(), upper_bound, -diff))); - // if diff==-1, the `upper_opt` will cancel the future `-1` inserted before by previous diff==1 row - let lower = Some(Ok((res_row, lower_bound, diff))); - - lower.into_iter().chain(upper_opt) - } else { - None.into_iter().chain(None) - } - } -} - -#[cfg(test)] -mod test { - use std::sync::Arc; - - use datatypes::data_type::ConcreteDataType; - use datatypes::vectors::{Int32Vector, Int64Vector}; - use pretty_assertions::assert_eq; - - use super::*; - use crate::expr::{BinaryFunc, UnaryFunc, UnmaterializableFunc}; - - #[test] - fn test_mfp_with_time() { - use crate::expr::func::BinaryFunc; - let lte_now = ScalarExpr::Column(0).call_binary( - ScalarExpr::CallUnmaterializable(UnmaterializableFunc::Now), - BinaryFunc::Lte, - ); - assert!(lte_now.contains_temporal()); - - let gt_now_minus_two = ScalarExpr::Column(0) - .call_binary( - ScalarExpr::Literal(Value::from(2i64), ConcreteDataType::int64_datatype()), - BinaryFunc::AddInt64, - ) - .call_binary( - ScalarExpr::CallUnmaterializable(UnmaterializableFunc::Now), - BinaryFunc::Gt, - ); - assert!(gt_now_minus_two.contains_temporal()); - - let mfp = MapFilterProject::new(3) - .filter(vec![ - // col(0) <= now() - lte_now, - // col(0) + 2 > now() - gt_now_minus_two, - ]) - .unwrap() - .project(vec![0]) - .unwrap(); - - let mfp = MfpPlan::create_from(mfp).unwrap(); - let expected = vec![ - ( - 0, - vec![ - (Row::new(vec![Value::from(4i64)]), 4, 1), - (Row::new(vec![Value::from(4i64)]), 6, -1), - ], - ), - ( - 5, - vec![ - (Row::new(vec![Value::from(4i64)]), 5, 1), - (Row::new(vec![Value::from(4i64)]), 6, -1), - ], - ), - (10, vec![]), - ]; - for (sys_time, expected) in expected { - let mut values = vec![Value::from(4i64), Value::from(2i64), Value::from(3i64)]; - let ret = mfp - .evaluate::(&mut values, sys_time, 1) - .collect::, _>>() - .unwrap(); - assert_eq!(ret, expected); - } - } - - #[test] - fn test_mfp() { - use crate::expr::func::BinaryFunc; - let mfp = MapFilterProject::new(3) - .map(vec![ - ScalarExpr::Column(0).call_binary(ScalarExpr::Column(1), BinaryFunc::Lt), - ScalarExpr::Column(1).call_binary(ScalarExpr::Column(2), BinaryFunc::Lt), - ]) - .unwrap() - .project(vec![3, 4]) - .unwrap(); - assert!(!mfp.is_identity()); - let mfp = MapFilterProject::compose(mfp, MapFilterProject::new(2)).unwrap(); - { - let mfp_0 = mfp.as_map_filter_project(); - let same = MapFilterProject::new(3) - .map(mfp_0.0) - .unwrap() - .filter(mfp_0.1) - .unwrap() - .project(mfp_0.2) - .unwrap(); - assert_eq!(mfp, same); - } - assert_eq!(mfp.demand().len(), 3); - let mut mfp = mfp; - mfp.permute(BTreeMap::from([(0, 2), (2, 0), (1, 1)]), 3) - .unwrap(); - assert_eq!( - mfp, - MapFilterProject::new(3) - .map(vec![ - ScalarExpr::Column(2).call_binary(ScalarExpr::Column(1), BinaryFunc::Lt), - ScalarExpr::Column(1).call_binary(ScalarExpr::Column(0), BinaryFunc::Lt), - ]) - .unwrap() - .project(vec![3, 4]) - .unwrap() - ); - let safe_mfp = SafeMfpPlan { mfp }; - let mut values = vec![Value::from(4), Value::from(2), Value::from(3)]; - let ret = safe_mfp - .evaluate_into(&mut values, &mut Row::empty()) - .unwrap() - .unwrap(); - assert_eq!(ret, Row::pack(vec![Value::from(false), Value::from(true)])); - let ty = [ - ConcreteDataType::int32_datatype(), - ConcreteDataType::int32_datatype(), - ConcreteDataType::int32_datatype(), - ]; - // batch mode - let mut batch = Batch::try_from_rows_with_types( - vec![Row::from(vec![ - Value::from(4), - Value::from(2), - Value::from(3), - ])], - &ty, - ) - .unwrap(); - let ret = safe_mfp.eval_batch_into(&mut batch).unwrap(); - - assert_eq!( - ret, - Batch::try_from_rows_with_types( - vec![Row::from(vec![Value::from(false), Value::from(true)])], - &[ - ConcreteDataType::boolean_datatype(), - ConcreteDataType::boolean_datatype(), - ], - ) - .unwrap() - ); - } - - #[test] - fn manipulation_mfp() { - // give a input of 4 columns - let mfp = MapFilterProject::new(4); - // append a expression to the mfp'input row that get the sum of the first 3 columns - let mfp = mfp - .map(vec![ - ScalarExpr::Column(0) - .call_binary(ScalarExpr::Column(1), BinaryFunc::AddInt32) - .call_binary(ScalarExpr::Column(2), BinaryFunc::AddInt32), - ]) - .unwrap(); - // only retain sum result - let mfp = mfp.project(vec![4]).unwrap(); - // accept only if the sum is greater than 10 - let mfp = mfp - .filter(vec![ScalarExpr::Column(0).call_binary( - ScalarExpr::Literal(Value::from(10i32), ConcreteDataType::int32_datatype()), - BinaryFunc::Gt, - )]) - .unwrap(); - let input1 = vec![ - Value::from(4), - Value::from(2), - Value::from(3), - Value::from("abc"), - ]; - let safe_mfp = SafeMfpPlan { mfp }; - let ret = safe_mfp - .evaluate_into(&mut input1.clone(), &mut Row::empty()) - .unwrap(); - assert_eq!(ret, None); - - let input_type = [ - ConcreteDataType::int32_datatype(), - ConcreteDataType::int32_datatype(), - ConcreteDataType::int32_datatype(), - ConcreteDataType::string_datatype(), - ]; - - let mut input1_batch = - Batch::try_from_rows_with_types(vec![Row::new(input1)], &input_type).unwrap(); - let ret_batch = safe_mfp.eval_batch_into(&mut input1_batch).unwrap(); - assert_eq!( - ret_batch, - Batch::try_new(vec![Arc::new(Int32Vector::from_vec(vec![]))], 0).unwrap() - ); - - let input2 = vec![ - Value::from(5), - Value::from(2), - Value::from(4), - Value::from("abc"), - ]; - let ret = safe_mfp - .evaluate_into(&mut input2.clone(), &mut Row::empty()) - .unwrap(); - assert_eq!(ret, Some(Row::pack(vec![Value::from(11)]))); - - let mut input2_batch = - Batch::try_from_rows_with_types(vec![Row::new(input2)], &input_type).unwrap(); - let ret_batch = safe_mfp.eval_batch_into(&mut input2_batch).unwrap(); - assert_eq!( - ret_batch, - Batch::try_new(vec![Arc::new(Int32Vector::from_vec(vec![11]))], 1).unwrap() - ); - } - - #[test] - fn test_permute() { - let mfp = MapFilterProject::new(3) - .map(vec![ - ScalarExpr::Column(0).call_binary(ScalarExpr::Column(1), BinaryFunc::Lt), - ]) - .unwrap() - .filter(vec![ - ScalarExpr::Column(0).call_binary(ScalarExpr::Column(1), BinaryFunc::Gt), - ]) - .unwrap() - .project(vec![0, 1]) - .unwrap(); - assert_eq!(mfp.demand(), BTreeSet::from([0, 1])); - let mut less = mfp.clone(); - less.permute(BTreeMap::from([(1, 0), (0, 1)]), 2).unwrap(); - - let mut more = mfp.clone(); - more.permute(BTreeMap::from([(0, 1), (1, 2), (2, 0)]), 4) - .unwrap(); - } - - #[test] - fn mfp_test_cast_and_filter() { - let mfp = MapFilterProject::new(3) - .map(vec![ScalarExpr::Column(0).call_unary(UnaryFunc::Cast( - ConcreteDataType::int32_datatype(), - ))]) - .unwrap() - .filter(vec![ - ScalarExpr::Column(3).call_binary(ScalarExpr::Column(1), BinaryFunc::Gt), - ]) - .unwrap() - .project([0, 1, 2]) - .unwrap(); - let input1 = vec![ - Value::from(4i64), - Value::from(2), - Value::from(3), - Value::from(53), - ]; - let safe_mfp = SafeMfpPlan { mfp }; - let ret = safe_mfp.evaluate_into(&mut input1.clone(), &mut Row::empty()); - assert!(matches!(ret, Err(EvalError::InvalidArgument { .. }))); - - let input_type = [ - ConcreteDataType::int64_datatype(), - ConcreteDataType::int32_datatype(), - ConcreteDataType::int32_datatype(), - ConcreteDataType::int32_datatype(), - ]; - let mut input1_batch = - Batch::try_from_rows_with_types(vec![Row::new(input1)], &input_type).unwrap(); - let ret_batch = safe_mfp.eval_batch_into(&mut input1_batch); - assert!(matches!(ret_batch, Err(EvalError::InvalidArgument { .. }))); - - let input2 = vec![Value::from(4i64), Value::from(2), Value::from(3)]; - let ret = safe_mfp - .evaluate_into(&mut input2.clone(), &mut Row::empty()) - .unwrap(); - assert_eq!(ret, Some(Row::new(input2.clone()))); - - let input_type = [ - ConcreteDataType::int64_datatype(), - ConcreteDataType::int32_datatype(), - ConcreteDataType::int32_datatype(), - ]; - let input2_batch = - Batch::try_from_rows_with_types(vec![Row::new(input2)], &input_type).unwrap(); - let ret_batch = safe_mfp.eval_batch_into(&mut input2_batch.clone()).unwrap(); - assert_eq!(ret_batch, input2_batch); - - let input3 = vec![Value::from(4i64), Value::from(5), Value::from(2)]; - let ret = safe_mfp - .evaluate_into(&mut input3.clone(), &mut Row::empty()) - .unwrap(); - assert_eq!(ret, None); - - let input3_batch = - Batch::try_from_rows_with_types(vec![Row::new(input3)], &input_type).unwrap(); - let ret_batch = safe_mfp.eval_batch_into(&mut input3_batch.clone()).unwrap(); - assert_eq!( - ret_batch, - Batch::try_new( - vec![ - Arc::new(Int64Vector::from_vec(Default::default())), - Arc::new(Int32Vector::from_vec(Default::default())), - Arc::new(Int32Vector::from_vec(Default::default())) - ], - 0 - ) - .unwrap() - ); - } - - #[test] - fn test_mfp_out_of_order() { - let mfp = MapFilterProject::new(3) - .project(vec![2, 1, 0]) - .unwrap() - .filter(vec![ - ScalarExpr::Column(0).call_binary(ScalarExpr::Column(1), BinaryFunc::Gt), - ]) - .unwrap() - .map(vec![ - ScalarExpr::Column(0).call_binary(ScalarExpr::Column(1), BinaryFunc::Lt), - ]) - .unwrap() - .project(vec![3]) - .unwrap(); - let input1 = vec![Value::from(2), Value::from(3), Value::from(4)]; - let safe_mfp = SafeMfpPlan { mfp }; - let ret = safe_mfp.evaluate_into(&mut input1.clone(), &mut Row::empty()); - assert_eq!(ret.unwrap(), Some(Row::new(vec![Value::from(false)]))); - - let input_type = [ - ConcreteDataType::int32_datatype(), - ConcreteDataType::int32_datatype(), - ConcreteDataType::int32_datatype(), - ]; - let mut input1_batch = - Batch::try_from_rows_with_types(vec![Row::new(input1)], &input_type).unwrap(); - let ret_batch = safe_mfp.eval_batch_into(&mut input1_batch).unwrap(); - - assert_eq!( - ret_batch, - Batch::try_new(vec![Arc::new(BooleanVector::from(vec![false]))], 1).unwrap() - ); - } - #[test] - fn test_mfp_chore() { - // project keeps permute columns until it becomes the identity permutation - let mfp = MapFilterProject::new(3) - .project([1, 2, 0]) - .unwrap() - .project([1, 2, 0]) - .unwrap() - .project([1, 2, 0]) - .unwrap(); - assert_eq!(mfp, MapFilterProject::new(3)); - } -} diff --git a/src/flow/src/expr/relation.rs b/src/flow/src/expr/relation.rs deleted file mode 100644 index b5d7e4ef207..00000000000 --- a/src/flow/src/expr/relation.rs +++ /dev/null @@ -1,36 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! Describes an aggregation function and it's input expression. - -pub(crate) use accum::{Accum, Accumulator}; -pub(crate) use func::AggregateFunc; - -use crate::expr::ScalarExpr; - -mod accum; -mod func; - -/// Describes an aggregation expression. -#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)] -pub struct AggregateExpr { - /// Names the aggregation function. - pub func: AggregateFunc, - /// An expression which extracts from each row the input to `func`. - /// TODO(discord9): currently unused in render phase(because AccumulablePlan remember each Aggr Expr's input/output column), - /// so it only used in generate KeyValPlan from AggregateExpr - pub expr: ScalarExpr, - /// Should the aggregation be applied only to distinct results in each group. - pub distinct: bool, -} diff --git a/src/flow/src/expr/relation/accum.rs b/src/flow/src/expr/relation/accum.rs deleted file mode 100644 index d0d25e97687..00000000000 --- a/src/flow/src/expr/relation/accum.rs +++ /dev/null @@ -1,1052 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! Accumulators for aggregate functions that's is accumulatable. i.e. sum/count -//! -//! Accumulator will only be restore from row and being updated every time dataflow need process a new batch of rows. -//! So the overhead is acceptable. -//! -//! Currently support sum, count, any, all and min/max(with one caveat that min/max can't support delete with aggregate). -//! TODO: think of better ways to not ser/de every time a accum needed to be updated, since it's in a tight loop - -use std::any::type_name; -use std::fmt::Display; - -use common_decimal::Decimal128; -use datatypes::data_type::ConcreteDataType; -use datatypes::value::{OrderedF32, OrderedF64, OrderedFloat, Value}; -use enum_dispatch::enum_dispatch; -use serde::{Deserialize, Serialize}; -use snafu::ensure; - -use crate::expr::error::{InternalSnafu, OverflowSnafu, TryFromValueSnafu, TypeMismatchSnafu}; -use crate::expr::signature::GenericFn; -use crate::expr::{AggregateFunc, EvalError}; -use crate::repr::Diff; - -/// Accumulates values for the various types of accumulable aggregations. -#[enum_dispatch] -pub trait Accumulator: Sized { - fn into_state(self) -> Vec; - - fn update( - &mut self, - aggr_fn: &AggregateFunc, - value: Value, - diff: Diff, - ) -> Result<(), EvalError>; - - fn update_batch(&mut self, aggr_fn: &AggregateFunc, value_diffs: I) -> Result<(), EvalError> - where - I: IntoIterator, - { - for (v, d) in value_diffs { - self.update(aggr_fn, v, d)?; - } - Ok(()) - } - - fn eval(&self, aggr_fn: &AggregateFunc) -> Result; -} - -/// Bool accumulator, used for `Any` `All` `Max/MinBool` -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -pub struct Bool { - /// The number of `true` values observed. - trues: Diff, - /// The number of `false` values observed. - falses: Diff, -} - -impl Bool { - /// Expect two `Diff` type values, one for `true` and one for `false`. - pub fn try_from_iter(iter: &mut I) -> Result - where - I: Iterator, - { - Ok(Self { - trues: Diff::try_from(iter.next().ok_or_else(fail_accum::)?) - .map_err(err_try_from_val)?, - falses: Diff::try_from(iter.next().ok_or_else(fail_accum::)?) - .map_err(err_try_from_val)?, - }) - } -} - -impl TryFrom> for Bool { - type Error = EvalError; - - fn try_from(state: Vec) -> Result { - ensure!( - state.len() == 2, - InternalSnafu { - reason: "Bool Accumulator state should have 2 values", - } - ); - let mut iter = state.into_iter(); - - Self::try_from_iter(&mut iter) - } -} - -impl Accumulator for Bool { - fn into_state(self) -> Vec { - vec![self.trues.into(), self.falses.into()] - } - - /// Null values are ignored - fn update( - &mut self, - aggr_fn: &AggregateFunc, - value: Value, - diff: Diff, - ) -> Result<(), EvalError> { - ensure!( - matches!( - aggr_fn, - AggregateFunc::Any - | AggregateFunc::All - | AggregateFunc::MaxBool - | AggregateFunc::MinBool - ), - InternalSnafu { - reason: format!( - "Bool Accumulator does not support this aggregation function: {:?}", - aggr_fn - ), - } - ); - - match value { - Value::Boolean(true) => self.trues += diff, - Value::Boolean(false) => self.falses += diff, - Value::Null => (), // ignore nulls - x => { - return Err(TypeMismatchSnafu { - expected: ConcreteDataType::boolean_datatype(), - actual: x.data_type(), - } - .build()); - } - }; - Ok(()) - } - - fn eval(&self, aggr_fn: &AggregateFunc) -> Result { - match aggr_fn { - AggregateFunc::Any => Ok(Value::from(self.trues > 0)), - AggregateFunc::All => Ok(Value::from(self.falses == 0)), - AggregateFunc::MaxBool => Ok(Value::from(self.trues > 0)), - AggregateFunc::MinBool => Ok(Value::from(self.falses == 0)), - _ => Err(InternalSnafu { - reason: format!( - "Bool Accumulator does not support this aggregation function: {:?}", - aggr_fn - ), - } - .build()), - } - } -} - -/// Accumulates simple numeric values for sum over integer. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -pub struct SimpleNumber { - /// The accumulation of all non-NULL values observed. - accum: i128, - /// The number of non-NULL values observed. - non_nulls: Diff, -} - -impl SimpleNumber { - /// Expect one `Decimal128` and one `Diff` type values. - /// The `Decimal128` type is used to store the sum of all non-NULL values. - /// The `Diff` type is used to count the number of non-NULL values. - pub fn try_from_iter(iter: &mut I) -> Result - where - I: Iterator, - { - Ok(Self { - accum: Decimal128::try_from(iter.next().ok_or_else(fail_accum::)?) - .map_err(err_try_from_val)? - .val(), - non_nulls: Diff::try_from(iter.next().ok_or_else(fail_accum::)?) - .map_err(err_try_from_val)?, - }) - } -} - -impl TryFrom> for SimpleNumber { - type Error = EvalError; - - fn try_from(state: Vec) -> Result { - ensure!( - state.len() == 2, - InternalSnafu { - reason: "Number Accumulator state should have 2 values", - } - ); - let mut iter = state.into_iter(); - Self::try_from_iter(&mut iter) - } -} - -impl Accumulator for SimpleNumber { - fn into_state(self) -> Vec { - vec![ - Value::Decimal128(Decimal128::new(self.accum, 38, 0)), - self.non_nulls.into(), - ] - } - - fn update( - &mut self, - aggr_fn: &AggregateFunc, - value: Value, - diff: Diff, - ) -> Result<(), EvalError> { - ensure!( - matches!( - aggr_fn, - AggregateFunc::SumInt16 - | AggregateFunc::SumInt32 - | AggregateFunc::SumInt64 - | AggregateFunc::SumUInt16 - | AggregateFunc::SumUInt32 - | AggregateFunc::SumUInt64 - ), - InternalSnafu { - reason: format!( - "SimpleNumber Accumulator does not support this aggregation function: {:?}", - aggr_fn - ), - } - ); - - let v = match (aggr_fn, value) { - (AggregateFunc::SumInt16, Value::Int16(x)) => i128::from(x), - (AggregateFunc::SumInt32, Value::Int32(x)) => i128::from(x), - (AggregateFunc::SumInt64, Value::Int64(x)) => i128::from(x), - (AggregateFunc::SumUInt16, Value::UInt16(x)) => i128::from(x), - (AggregateFunc::SumUInt32, Value::UInt32(x)) => i128::from(x), - (AggregateFunc::SumUInt64, Value::UInt64(x)) => i128::from(x), - (_f, Value::Null) => return Ok(()), // ignore null - (f, v) => { - let expected_datatype = f.signature().input; - return Err(TypeMismatchSnafu { - expected: expected_datatype[0].clone(), - actual: v.data_type(), - } - .build())?; - } - }; - - self.accum += v * i128::from(diff); - - self.non_nulls += diff; - Ok(()) - } - - fn eval(&self, aggr_fn: &AggregateFunc) -> Result { - match aggr_fn { - AggregateFunc::SumInt16 | AggregateFunc::SumInt32 | AggregateFunc::SumInt64 => { - i64::try_from(self.accum) - .map_err(|_e| OverflowSnafu {}.build()) - .map(Value::from) - } - AggregateFunc::SumUInt16 | AggregateFunc::SumUInt32 | AggregateFunc::SumUInt64 => { - u64::try_from(self.accum) - .map_err(|_e| OverflowSnafu {}.build()) - .map(Value::from) - } - _ => Err(InternalSnafu { - reason: format!( - "SimpleNumber Accumulator does not support this aggregation function: {:?}", - aggr_fn - ), - } - .build()), - } - } -} -/// Accumulates float values for sum over floating numbers. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -pub struct Float { - /// Accumulates non-special float values, i.e. not NaN, +inf, -inf. - /// accum will be set to zero if `non_nulls` is zero. - accum: OrderedF64, - /// Counts +inf - pos_infs: Diff, - /// Counts -inf - neg_infs: Diff, - /// Counts NaNs - nans: Diff, - /// Counts non-NULL values - non_nulls: Diff, -} - -impl Float { - /// Expect first value to be `OrderedF64` and the rest four values to be `Diff` type values. - pub fn try_from_iter(iter: &mut I) -> Result - where - I: Iterator, - { - let mut ret = Self { - accum: OrderedF64::try_from(iter.next().ok_or_else(fail_accum::)?) - .map_err(err_try_from_val)?, - pos_infs: Diff::try_from(iter.next().ok_or_else(fail_accum::)?) - .map_err(err_try_from_val)?, - neg_infs: Diff::try_from(iter.next().ok_or_else(fail_accum::)?) - .map_err(err_try_from_val)?, - nans: Diff::try_from(iter.next().ok_or_else(fail_accum::)?) - .map_err(err_try_from_val)?, - non_nulls: Diff::try_from(iter.next().ok_or_else(fail_accum::)?) - .map_err(err_try_from_val)?, - }; - - // This prevent counter-intuitive behavior of summing over no values having non-zero results - if ret.non_nulls == 0 { - ret.accum = OrderedFloat::from(0.0); - } - - Ok(ret) - } -} - -impl TryFrom> for Float { - type Error = EvalError; - - fn try_from(state: Vec) -> Result { - ensure!( - state.len() == 5, - InternalSnafu { - reason: "Float Accumulator state should have 5 values", - } - ); - - let mut iter = state.into_iter(); - - let mut ret = Self { - accum: OrderedF64::try_from(iter.next().unwrap()).map_err(err_try_from_val)?, - pos_infs: Diff::try_from(iter.next().unwrap()).map_err(err_try_from_val)?, - neg_infs: Diff::try_from(iter.next().unwrap()).map_err(err_try_from_val)?, - nans: Diff::try_from(iter.next().unwrap()).map_err(err_try_from_val)?, - non_nulls: Diff::try_from(iter.next().unwrap()).map_err(err_try_from_val)?, - }; - - // This prevent counter-intuitive behavior of summing over no values - if ret.non_nulls == 0 { - ret.accum = OrderedFloat::from(0.0); - } - - Ok(ret) - } -} - -impl Accumulator for Float { - fn into_state(self) -> Vec { - vec![ - self.accum.into(), - self.pos_infs.into(), - self.neg_infs.into(), - self.nans.into(), - self.non_nulls.into(), - ] - } - - /// sum ignore null - fn update( - &mut self, - aggr_fn: &AggregateFunc, - value: Value, - diff: Diff, - ) -> Result<(), EvalError> { - ensure!( - matches!( - aggr_fn, - AggregateFunc::SumFloat32 | AggregateFunc::SumFloat64 - ), - InternalSnafu { - reason: format!( - "Float Accumulator does not support this aggregation function: {:?}", - aggr_fn - ), - } - ); - - let x = match (aggr_fn, value) { - (AggregateFunc::SumFloat32, Value::Float32(x)) => OrderedF64::from(*x as f64), - (AggregateFunc::SumFloat64, Value::Float64(x)) => OrderedF64::from(x), - (_f, Value::Null) => return Ok(()), // ignore null - (f, v) => { - let expected_datatype = f.signature().input; - return Err(TypeMismatchSnafu { - expected: expected_datatype[0].clone(), - actual: v.data_type(), - } - .build())?; - } - }; - - if x.is_nan() { - self.nans += diff; - } else if x.is_infinite() { - if x.is_sign_positive() { - self.pos_infs += diff; - } else { - self.neg_infs += diff; - } - } else { - self.accum += *(x * OrderedF64::from(diff as f64)); - } - - self.non_nulls += diff; - Ok(()) - } - - fn eval(&self, aggr_fn: &AggregateFunc) -> Result { - match aggr_fn { - AggregateFunc::SumFloat32 => Ok(Value::Float32(OrderedF32::from(self.accum.0 as f32))), - AggregateFunc::SumFloat64 => Ok(Value::Float64(self.accum)), - _ => Err(InternalSnafu { - reason: format!( - "Float Accumulator does not support this aggregation function: {:?}", - aggr_fn - ), - } - .build()), - } - } -} - -/// Accumulates a single `Ord`ed `Value`, useful for min/max aggregations. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -pub struct OrdValue { - val: Option, - non_nulls: Diff, -} - -impl OrdValue { - pub fn try_from_iter(iter: &mut I) -> Result - where - I: Iterator, - { - Ok(Self { - val: { - let v = iter.next().ok_or_else(fail_accum::)?; - if v == Value::Null { None } else { Some(v) } - }, - non_nulls: Diff::try_from(iter.next().ok_or_else(fail_accum::)?) - .map_err(err_try_from_val)?, - }) - } -} - -impl TryFrom> for OrdValue { - type Error = EvalError; - - fn try_from(state: Vec) -> Result { - ensure!( - state.len() == 2, - InternalSnafu { - reason: "OrdValue Accumulator state should have 2 values", - } - ); - - let mut iter = state.into_iter(); - - Ok(Self { - val: { - let v = iter.next().unwrap(); - if v == Value::Null { None } else { Some(v) } - }, - non_nulls: Diff::try_from(iter.next().unwrap()).map_err(err_try_from_val)?, - }) - } -} - -impl Accumulator for OrdValue { - fn into_state(self) -> Vec { - vec![self.val.unwrap_or(Value::Null), self.non_nulls.into()] - } - - /// min/max try to find results in all non-null values, if all values are null, the result is null. - /// count(col_name) gives the number of non-null values, count(*) gives the number of rows including nulls. - /// TODO(discord9): add count(*) as a aggr function - fn update( - &mut self, - aggr_fn: &AggregateFunc, - value: Value, - diff: Diff, - ) -> Result<(), EvalError> { - ensure!( - aggr_fn.is_max() || aggr_fn.is_min() || matches!(aggr_fn, AggregateFunc::Count), - InternalSnafu { - reason: format!( - "OrdValue Accumulator does not support this aggregation function: {:?}", - aggr_fn - ), - } - ); - if diff <= 0 && (aggr_fn.is_max() || aggr_fn.is_min()) { - return Err(InternalSnafu { - reason: "OrdValue Accumulator does not support non-monotonic input for min/max aggregation".to_string(), - }.build()); - } - - // if aggr_fn is count, the incoming value type doesn't matter in type checking - // otherwise, type need to be the same or value can be null - let check_type_aggr_fn_and_arg_value = - ty_eq_without_precision(value.data_type(), aggr_fn.signature().input[0].clone()) - || matches!(aggr_fn, AggregateFunc::Count) - || value.is_null(); - let check_type_aggr_fn_and_self_val = self - .val - .as_ref() - .map(|zelf| { - ty_eq_without_precision(zelf.data_type(), aggr_fn.signature().input[0].clone()) - }) - .unwrap_or(true) - || matches!(aggr_fn, AggregateFunc::Count); - - if !check_type_aggr_fn_and_arg_value { - return Err(TypeMismatchSnafu { - expected: aggr_fn.signature().input[0].clone(), - actual: value.data_type(), - } - .build()); - } else if !check_type_aggr_fn_and_self_val { - return Err(TypeMismatchSnafu { - expected: aggr_fn.signature().input[0].clone(), - actual: self - .val - .as_ref() - .map(|v| v.data_type()) - .unwrap_or(ConcreteDataType::null_datatype()), - } - .build()); - } - - let is_null = value.is_null(); - if is_null { - return Ok(()); - } - - if !is_null { - // compile count(*) to count(true) to include null/non-nulls - // And the counts of non-null values are updated here - self.non_nulls += diff; - - match aggr_fn.signature().generic_fn { - GenericFn::Max => { - self.val = self - .val - .clone() - .map(|v| v.max(value.clone())) - .or_else(|| Some(value)) - } - GenericFn::Min => { - self.val = self - .val - .clone() - .map(|v| v.min(value.clone())) - .or_else(|| Some(value)) - } - - GenericFn::Count => (), - _ => unreachable!("already checked by ensure!"), - } - }; - // min/max ignore nulls - - Ok(()) - } - - fn eval(&self, aggr_fn: &AggregateFunc) -> Result { - if aggr_fn.is_max() || aggr_fn.is_min() { - Ok(self.val.clone().unwrap_or(Value::Null)) - } else if matches!(aggr_fn, AggregateFunc::Count) { - Ok(self.non_nulls.into()) - } else { - Err(InternalSnafu { - reason: format!( - "OrdValue Accumulator does not support this aggregation function: {:?}", - aggr_fn - ), - } - .build()) - } - } -} - -/// Accumulates values for the various types of accumulable aggregations. -/// -/// We assume that there are not more than 2^32 elements for the aggregation. -/// Thus we can perform a summation over i32 in an i64 accumulator -/// and not worry about exceeding its bounds. -/// -/// The float accumulator performs accumulation with tolerance for floating point error. -/// -/// TODO(discord9): check for overflowing -#[enum_dispatch(Accumulator)] -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -pub enum Accum { - /// Accumulates boolean values. - Bool(Bool), - /// Accumulates simple numeric values. - SimpleNumber(SimpleNumber), - /// Accumulates float values. - Float(Float), - /// Accumulate Values that impl `Ord` - OrdValue(OrdValue), -} - -impl Accum { - /// create a new accumulator from given aggregate function - pub fn new_accum(aggr_fn: &AggregateFunc) -> Result { - Ok(match aggr_fn { - AggregateFunc::Any - | AggregateFunc::All - | AggregateFunc::MaxBool - | AggregateFunc::MinBool => Self::from(Bool { - trues: 0, - falses: 0, - }), - AggregateFunc::SumInt16 - | AggregateFunc::SumInt32 - | AggregateFunc::SumInt64 - | AggregateFunc::SumUInt16 - | AggregateFunc::SumUInt32 - | AggregateFunc::SumUInt64 => Self::from(SimpleNumber { - accum: 0, - non_nulls: 0, - }), - AggregateFunc::SumFloat32 | AggregateFunc::SumFloat64 => Self::from(Float { - accum: OrderedF64::from(0.0), - pos_infs: 0, - neg_infs: 0, - nans: 0, - non_nulls: 0, - }), - f if f.is_max() || f.is_min() || matches!(f, AggregateFunc::Count) => { - Self::from(OrdValue { - val: None, - non_nulls: 0, - }) - } - f => { - return Err(InternalSnafu { - reason: format!( - "Accumulator does not support this aggregation function: {:?}", - f - ), - } - .build()); - } - }) - } - - pub fn try_from_iter( - aggr_fn: &AggregateFunc, - iter: &mut impl Iterator, - ) -> Result { - match aggr_fn { - AggregateFunc::Any - | AggregateFunc::All - | AggregateFunc::MaxBool - | AggregateFunc::MinBool => Ok(Self::from(Bool::try_from_iter(iter)?)), - AggregateFunc::SumInt16 - | AggregateFunc::SumInt32 - | AggregateFunc::SumInt64 - | AggregateFunc::SumUInt16 - | AggregateFunc::SumUInt32 - | AggregateFunc::SumUInt64 => Ok(Self::from(SimpleNumber::try_from_iter(iter)?)), - AggregateFunc::SumFloat32 | AggregateFunc::SumFloat64 => { - Ok(Self::from(Float::try_from_iter(iter)?)) - } - f if f.is_max() || f.is_min() || matches!(f, AggregateFunc::Count) => { - Ok(Self::from(OrdValue::try_from_iter(iter)?)) - } - f => Err(InternalSnafu { - reason: format!( - "Accumulator does not support this aggregation function: {:?}", - f - ), - } - .build()), - } - } - - /// try to convert a vector of value into given aggregate function's accumulator - pub fn try_into_accum(aggr_fn: &AggregateFunc, state: Vec) -> Result { - match aggr_fn { - AggregateFunc::Any - | AggregateFunc::All - | AggregateFunc::MaxBool - | AggregateFunc::MinBool => Ok(Self::from(Bool::try_from(state)?)), - AggregateFunc::SumInt16 - | AggregateFunc::SumInt32 - | AggregateFunc::SumInt64 - | AggregateFunc::SumUInt16 - | AggregateFunc::SumUInt32 - | AggregateFunc::SumUInt64 => Ok(Self::from(SimpleNumber::try_from(state)?)), - AggregateFunc::SumFloat32 | AggregateFunc::SumFloat64 => { - Ok(Self::from(Float::try_from(state)?)) - } - f if f.is_max() || f.is_min() || matches!(f, AggregateFunc::Count) => { - Ok(Self::from(OrdValue::try_from(state)?)) - } - f => Err(InternalSnafu { - reason: format!( - "Accumulator does not support this aggregation function: {:?}", - f - ), - } - .build()), - } - } -} - -fn fail_accum() -> EvalError { - InternalSnafu { - reason: format!( - "list of values exhausted before a accum of type {} can be build from it", - type_name::() - ), - } - .build() -} - -fn err_try_from_val(reason: T) -> EvalError { - TryFromValueSnafu { - msg: reason.to_string(), - } - .build() -} - -/// compare type while ignore their precision, including `TimeStamp`, `Time`, -/// `Duration`, `Interval` -fn ty_eq_without_precision(left: ConcreteDataType, right: ConcreteDataType) -> bool { - left == right - || matches!(left, ConcreteDataType::Timestamp(..)) - && matches!(right, ConcreteDataType::Timestamp(..)) - || matches!(left, ConcreteDataType::Time(..)) && matches!(right, ConcreteDataType::Time(..)) - || matches!(left, ConcreteDataType::Duration(..)) - && matches!(right, ConcreteDataType::Duration(..)) - || matches!(left, ConcreteDataType::Interval(..)) - && matches!(right, ConcreteDataType::Interval(..)) -} - -#[allow(clippy::too_many_lines)] -#[cfg(test)] -mod test { - use common_time::Timestamp; - - use super::*; - - #[test] - fn test_accum() { - let testcases = vec![ - ( - AggregateFunc::SumInt32, - vec![(Value::Int32(1), 1), (Value::Null, 1)], - ( - Value::Int64(1), - vec![Value::Decimal128(Decimal128::new(1, 38, 0)), 1i64.into()], - ), - ), - ( - AggregateFunc::SumFloat32, - vec![(Value::Float32(OrderedF32::from(1.0)), 1), (Value::Null, 1)], - ( - Value::Float32(OrderedF32::from(1.0)), - vec![ - Value::Float64(OrderedF64::from(1.0)), - 0i64.into(), - 0i64.into(), - 0i64.into(), - 1i64.into(), - ], - ), - ), - ( - AggregateFunc::MaxInt32, - vec![(Value::Int32(1), 1), (Value::Int32(2), 1), (Value::Null, 1)], - (Value::Int32(2), vec![Value::Int32(2), 2i64.into()]), - ), - ( - AggregateFunc::MinInt32, - vec![(Value::Int32(2), 1), (Value::Int32(1), 1), (Value::Null, 1)], - (Value::Int32(1), vec![Value::Int32(1), 2i64.into()]), - ), - ( - AggregateFunc::MaxFloat32, - vec![ - (Value::Float32(OrderedF32::from(1.0)), 1), - (Value::Float32(OrderedF32::from(2.0)), 1), - (Value::Null, 1), - ], - ( - Value::Float32(OrderedF32::from(2.0)), - vec![Value::Float32(OrderedF32::from(2.0)), 2i64.into()], - ), - ), - ( - AggregateFunc::MaxDateTime, - vec![ - (Value::Timestamp(Timestamp::from(0)), 1), - (Value::Timestamp(Timestamp::from(1)), 1), - (Value::Null, 1), - ], - ( - Value::Timestamp(Timestamp::from(1)), - vec![Value::Timestamp(Timestamp::from(1)), 2i64.into()], - ), - ), - ( - AggregateFunc::Count, - vec![ - (Value::Int32(1), 1), - (Value::Int32(2), 1), - (Value::Null, 1), - (Value::Null, 1), - ], - (2i64.into(), vec![Value::Null, 2i64.into()]), - ), - ( - AggregateFunc::Any, - vec![ - (Value::Boolean(false), 1), - (Value::Boolean(false), 1), - (Value::Boolean(true), 1), - (Value::Null, 1), - ], - ( - Value::Boolean(true), - vec![Value::from(1i64), Value::from(2i64)], - ), - ), - ( - AggregateFunc::All, - vec![ - (Value::Boolean(false), 1), - (Value::Boolean(false), 1), - (Value::Boolean(true), 1), - (Value::Null, 1), - ], - ( - Value::Boolean(false), - vec![Value::from(1i64), Value::from(2i64)], - ), - ), - ( - AggregateFunc::MaxBool, - vec![ - (Value::Boolean(false), 1), - (Value::Boolean(false), 1), - (Value::Boolean(true), 1), - (Value::Null, 1), - ], - ( - Value::Boolean(true), - vec![Value::from(1i64), Value::from(2i64)], - ), - ), - ( - AggregateFunc::MinBool, - vec![ - (Value::Boolean(false), 1), - (Value::Boolean(false), 1), - (Value::Boolean(true), 1), - (Value::Null, 1), - ], - ( - Value::Boolean(false), - vec![Value::from(1i64), Value::from(2i64)], - ), - ), - ]; - - for (aggr_fn, input, (eval_res, state)) in testcases { - let create_and_insert = || -> Result { - let mut acc = Accum::new_accum(&aggr_fn)?; - acc.update_batch(&aggr_fn, input.clone())?; - let row = acc.into_state(); - let acc = Accum::try_into_accum(&aggr_fn, row.clone())?; - let alter_acc = Accum::try_from_iter(&aggr_fn, &mut row.into_iter())?; - assert_eq!(acc, alter_acc); - Ok(acc) - }; - let acc = match create_and_insert() { - Ok(acc) => acc, - Err(err) => panic!( - "Failed to create accum for {:?} with input {:?} with error: {:?}", - aggr_fn, input, err - ), - }; - - if acc.eval(&aggr_fn).unwrap() != eval_res { - panic!( - "Failed to eval accum for {:?} with input {:?}, expect {:?}, got {:?}", - aggr_fn, - input, - eval_res, - acc.eval(&aggr_fn).unwrap() - ); - } - let actual_state = acc.into_state(); - if actual_state != state { - panic!( - "Failed to cast into state from accum for {:?} with input {:?}, expect state {:?}, got state {:?}", - aggr_fn, input, state, actual_state - ); - } - } - } - #[test] - fn test_fail_path_accum() { - { - let bool_accum = Bool::try_from(vec![Value::Null]); - assert!(matches!(bool_accum, Err(EvalError::Internal { .. }))); - } - - { - let mut bool_accum = Bool::try_from(vec![1i64.into(), 1i64.into()]).unwrap(); - // serde - let bool_accum_serde = serde_json::to_string(&bool_accum).unwrap(); - let bool_accum_de = serde_json::from_str::(&bool_accum_serde).unwrap(); - assert_eq!(bool_accum, bool_accum_de); - assert!(matches!( - bool_accum.update(&AggregateFunc::MaxDate, 1.into(), 1), - Err(EvalError::Internal { .. }) - )); - assert!(matches!( - bool_accum.update(&AggregateFunc::Any, 1.into(), 1), - Err(EvalError::TypeMismatch { .. }) - )); - assert!(matches!( - bool_accum.eval(&AggregateFunc::MaxDate), - Err(EvalError::Internal { .. }) - )); - } - - { - let ret = SimpleNumber::try_from(vec![Value::Null]); - assert!(matches!(ret, Err(EvalError::Internal { .. }))); - let mut accum = - SimpleNumber::try_from(vec![Decimal128::new(0, 38, 0).into(), 0i64.into()]) - .unwrap(); - - assert!(matches!( - accum.update(&AggregateFunc::All, 0.into(), 1), - Err(EvalError::Internal { .. }) - )); - assert!(matches!( - accum.update(&AggregateFunc::SumInt64, 0i32.into(), 1), - Err(EvalError::TypeMismatch { .. }) - )); - assert!(matches!( - accum.eval(&AggregateFunc::All), - Err(EvalError::Internal { .. }) - )); - accum - .update(&AggregateFunc::SumInt64, 1i64.into(), 1) - .unwrap(); - accum - .update(&AggregateFunc::SumInt64, i64::MAX.into(), 1) - .unwrap(); - assert!(matches!( - accum.eval(&AggregateFunc::SumInt64), - Err(EvalError::Overflow { .. }) - )); - } - - { - let ret = Float::try_from(vec![2f64.into(), 0i64.into(), 0i64.into(), 0i64.into()]); - assert!(matches!(ret, Err(EvalError::Internal { .. }))); - let mut accum = Float::try_from(vec![ - 2f64.into(), - 0i64.into(), - 0i64.into(), - 0i64.into(), - 1i64.into(), - ]) - .unwrap(); - accum - .update(&AggregateFunc::SumFloat64, 2f64.into(), -1) - .unwrap(); - assert!(matches!( - accum.update(&AggregateFunc::All, 0.into(), 1), - Err(EvalError::Internal { .. }) - )); - assert!(matches!( - accum.update(&AggregateFunc::SumFloat64, 0.0f32.into(), 1), - Err(EvalError::TypeMismatch { .. }) - )); - // no record, no accum - assert_eq!( - accum.eval(&AggregateFunc::SumFloat64).unwrap(), - 0.0f64.into() - ); - - assert!(matches!( - accum.eval(&AggregateFunc::All), - Err(EvalError::Internal { .. }) - )); - - accum - .update(&AggregateFunc::SumFloat64, f64::INFINITY.into(), 1) - .unwrap(); - accum - .update(&AggregateFunc::SumFloat64, (-f64::INFINITY).into(), 1) - .unwrap(); - accum - .update(&AggregateFunc::SumFloat64, f64::NAN.into(), 1) - .unwrap(); - } - - { - let ret = OrdValue::try_from(vec![Value::Null]); - assert!(matches!(ret, Err(EvalError::Internal { .. }))); - let mut accum = OrdValue::try_from(vec![Value::Null, 0i64.into()]).unwrap(); - assert!(matches!( - accum.update(&AggregateFunc::All, 0.into(), 1), - Err(EvalError::Internal { .. }) - )); - accum - .update(&AggregateFunc::MaxInt16, 1i16.into(), 1) - .unwrap(); - assert!(matches!( - accum.update(&AggregateFunc::MaxInt16, 0i32.into(), 1), - Err(EvalError::TypeMismatch { .. }) - )); - assert!(matches!( - accum.update(&AggregateFunc::MaxInt16, 0i16.into(), -1), - Err(EvalError::Internal { .. }) - )); - accum - .update(&AggregateFunc::MaxInt16, Value::Null, 1) - .unwrap(); - } - - // insert uint64 into max_int64 should fail - { - let mut accum = OrdValue::try_from(vec![Value::Null, 0i64.into()]).unwrap(); - assert!(matches!( - accum.update(&AggregateFunc::MaxInt64, 0u64.into(), 1), - Err(EvalError::TypeMismatch { .. }) - )); - } - } -} diff --git a/src/flow/src/expr/relation/func.rs b/src/flow/src/expr/relation/func.rs deleted file mode 100644 index 35a43957d8c..00000000000 --- a/src/flow/src/expr/relation/func.rs +++ /dev/null @@ -1,303 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::collections::HashMap; -use std::sync::OnceLock; - -use datatypes::prelude::ConcreteDataType; -use datatypes::value::Value; -use datatypes::vectors::VectorRef; -use serde::{Deserialize, Serialize}; -use smallvec::smallvec; -use snafu::OptionExt; -use strum::{EnumIter, IntoEnumIterator}; - -use crate::error::{Error, InvalidQuerySnafu}; -use crate::expr::VectorDiff; -use crate::expr::error::EvalError; -use crate::expr::relation::accum::{Accum, Accumulator}; -use crate::expr::signature::{GenericFn, Signature}; -use crate::repr::Diff; - -/// Aggregate functions that can be applied to a group of rows. -/// -/// `Mean` function is deliberately not included as it can be computed from `Sum` and `Count`, whose state can be better managed. -/// -/// type of the input and output of the aggregate function: -/// -/// `sum(i*)->i64, sum(u*)->u64` -/// -/// `count()->i64` -/// -/// `min/max(T)->T` -#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize, Hash, EnumIter)] -pub enum AggregateFunc { - MaxInt16, - MaxInt32, - MaxInt64, - MaxUInt16, - MaxUInt32, - MaxUInt64, - MaxFloat32, - MaxFloat64, - MaxBool, - MaxString, - MaxDate, - MaxDateTime, - MaxTimestamp, - MaxTime, - MaxDuration, - MaxInterval, - - MinInt16, - MinInt32, - MinInt64, - MinUInt16, - MinUInt32, - MinUInt64, - MinFloat32, - MinFloat64, - MinBool, - MinString, - MinDate, - MinDateTime, - MinTimestamp, - MinTime, - MinDuration, - MinInterval, - - SumInt16, - SumInt32, - SumInt64, - SumUInt16, - SumUInt32, - SumUInt64, - SumFloat32, - SumFloat64, - - Count, - Any, - All, -} - -impl AggregateFunc { - /// if this function is a `max` - pub fn is_max(&self) -> bool { - self.signature().generic_fn == GenericFn::Max - } - - /// if this function is a `min` - pub fn is_min(&self) -> bool { - self.signature().generic_fn == GenericFn::Min - } - - /// Eval value, diff with accumulator - /// - /// Expect self to be accumulable aggregate function, i.e. sum/count - /// - /// TODO(discord9): deal with overflow&better accumulator - pub fn eval_diff_accumulable( - &self, - accum: A, - value_diffs: I, - ) -> Result<(Value, Vec), EvalError> - where - A: IntoIterator, - I: IntoIterator, - { - let mut accum = accum.into_iter().peekable(); - - let mut accum = if accum.peek().is_none() { - Accum::new_accum(self)? - } else { - Accum::try_from_iter(self, &mut accum)? - }; - accum.update_batch(self, value_diffs)?; - let res = accum.eval(self)?; - Ok((res, accum.into_state())) - } - - /// return output value and new accumulator state - pub fn eval_batch( - &self, - accum: A, - vector: VectorRef, - diff: Option, - ) -> Result<(Value, Vec), EvalError> - where - A: IntoIterator, - { - let mut accum = accum.into_iter().peekable(); - - let mut accum = if accum.peek().is_none() { - Accum::new_accum(self)? - } else { - Accum::try_from_iter(self, &mut accum)? - }; - - let vector_diff = VectorDiff::try_new(vector, diff)?; - - accum.update_batch(self, vector_diff)?; - - let res = accum.eval(self)?; - Ok((res, accum.into_state())) - } -} - -/// Generate signature for each aggregate function -macro_rules! generate_signature { - ($value:ident, - { $($user_arm:tt)* }, - [ $( - $auto_arm:ident=>($($arg:ident),*) - ),* - ] - ) => { - match $value { - $($user_arm)*, - $( - Self::$auto_arm => gen_one_siginature!($($arg),*), - )* - } - }; -} - -/// Generate one match arm with optional arguments -macro_rules! gen_one_siginature { - ( - $con_type:ident, $generic:ident - ) => { - Signature { - input: smallvec![ConcreteDataType::$con_type(), ConcreteDataType::$con_type(),], - output: ConcreteDataType::$con_type(), - generic_fn: GenericFn::$generic, - } - }; - ( - $in_type:ident, $out_type:ident, $generic:ident - ) => { - Signature { - input: smallvec![ConcreteDataType::$in_type()], - output: ConcreteDataType::$out_type(), - generic_fn: GenericFn::$generic, - } - }; -} - -static SPECIALIZATION: OnceLock> = - OnceLock::new(); - -impl AggregateFunc { - /// Create a `AggregateFunc` from a string of the function name and given argument type(optional) - /// given an None type will be treated as null type, - /// which in turn for AggregateFunc like `Count` will be treated as any type - pub fn from_str_and_type( - name: &str, - arg_type: Option, - ) -> Result { - let rule = SPECIALIZATION.get_or_init(|| { - let mut spec = HashMap::new(); - for func in Self::iter() { - let sig = func.signature(); - spec.insert((sig.generic_fn, sig.input[0].clone()), func); - } - spec - }); - - let generic_fn = match name { - "max" => GenericFn::Max, - "min" => GenericFn::Min, - "sum" => GenericFn::Sum, - "count" => GenericFn::Count, - "bool_or" => GenericFn::Any, - "bool_and" => GenericFn::All, - _ => { - return InvalidQuerySnafu { - reason: format!("Unknown aggregate function: {}", name), - } - .fail(); - } - }; - let input_type = if matches!(generic_fn, GenericFn::Count) { - ConcreteDataType::null_datatype() - } else { - arg_type.unwrap_or_else(ConcreteDataType::null_datatype) - }; - rule.get(&(generic_fn, input_type.clone())) - .cloned() - .with_context(|| InvalidQuerySnafu { - reason: format!( - "No specialization found for aggregate function {:?} with input type {:?}", - generic_fn, input_type - ), - }) - } - - /// all concrete datatypes with precision types will be returned with largest possible variant - /// as a exception, count have a signature of `null -> i64`, but it's actually `anytype -> i64` - /// - /// TODO(discorcd9): fix signature for sum unsign -> u64 sum signed -> i64 - pub fn signature(&self) -> Signature { - generate_signature!(self, { - AggregateFunc::Count => Signature { - input: smallvec![ConcreteDataType::null_datatype()], - output: ConcreteDataType::int64_datatype(), - generic_fn: GenericFn::Count, - } - },[ - MaxInt16 => (int16_datatype, Max), - MaxInt32 => (int32_datatype, Max), - MaxInt64 => (int64_datatype, Max), - MaxUInt16 => (uint16_datatype, Max), - MaxUInt32 => (uint32_datatype, Max), - MaxUInt64 => (uint64_datatype, Max), - MaxFloat32 => (float32_datatype, Max), - MaxFloat64 => (float64_datatype, Max), - MaxBool => (boolean_datatype, Max), - MaxString => (string_datatype, Max), - MaxDate => (date_datatype, Max), - MaxDateTime => (timestamp_microsecond_datatype, Max), - MaxTimestamp => (timestamp_second_datatype, Max), - MaxTime => (time_second_datatype, Max), - MaxDuration => (duration_second_datatype, Max), - MaxInterval => (interval_year_month_datatype, Max), - MinInt16 => (int16_datatype, Min), - MinInt32 => (int32_datatype, Min), - MinInt64 => (int64_datatype, Min), - MinUInt16 => (uint16_datatype, Min), - MinUInt32 => (uint32_datatype, Min), - MinUInt64 => (uint64_datatype, Min), - MinFloat32 => (float32_datatype, Min), - MinFloat64 => (float64_datatype, Min), - MinBool => (boolean_datatype, Min), - MinString => (string_datatype, Min), - MinDate => (date_datatype, Min), - MinDateTime => (timestamp_microsecond_datatype, Min), - MinTimestamp => (timestamp_second_datatype, Min), - MinTime => (time_second_datatype, Min), - MinDuration => (duration_second_datatype, Min), - MinInterval => (interval_year_month_datatype, Min), - SumInt16 => (int16_datatype, int64_datatype, Sum), - SumInt32 => (int32_datatype, int64_datatype, Sum), - SumInt64 => (int64_datatype, int64_datatype, Sum), - SumUInt16 => (uint16_datatype, uint64_datatype, Sum), - SumUInt32 => (uint32_datatype, uint64_datatype, Sum), - SumUInt64 => (uint64_datatype, uint64_datatype, Sum), - SumFloat32 => (float32_datatype, Sum), - SumFloat64 => (float64_datatype, Sum), - Any => (boolean_datatype, Any), - All => (boolean_datatype, All) - ]) - } -} diff --git a/src/flow/src/expr/scalar.rs b/src/flow/src/expr/scalar.rs deleted file mode 100644 index af16381e4cf..00000000000 --- a/src/flow/src/expr/scalar.rs +++ /dev/null @@ -1,877 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! Scalar expressions. - -use std::collections::{BTreeMap, BTreeSet}; -use std::sync::Arc; - -use arrow::array::{ArrayData, ArrayRef, BooleanArray, make_array}; -use arrow::buffer::BooleanBuffer; -use arrow::compute::or_kleene; -use common_error::ext::BoxedError; -use datafusion::physical_expr_common::datum::compare_with_eq; -use datafusion_common::DataFusionError; -use datatypes::prelude::{ConcreteDataType, DataType}; -use datatypes::value::Value; -use datatypes::vectors::{BooleanVector, Helper, VectorRef}; -use dfir_rs::lattices::cc_traits::Iter; -use itertools::Itertools; -use snafu::{OptionExt, ResultExt, ensure}; - -use crate::error::{ - DatafusionSnafu, Error, InvalidQuerySnafu, UnexpectedSnafu, UnsupportedTemporalFilterSnafu, -}; -use crate::expr::error::{ - ArrowSnafu, DataTypeSnafu, EvalError, InvalidArgumentSnafu, OptimizeSnafu, TypeMismatchSnafu, -}; -use crate::expr::func::{BinaryFunc, UnaryFunc, UnmaterializableFunc, VariadicFunc}; -use crate::expr::{Batch, DfScalarFunction}; -use crate::repr::ColumnType; -/// A scalar expression with a known type. -#[derive(Ord, PartialOrd, Clone, Debug, Eq, PartialEq, Hash)] -pub struct TypedExpr { - /// The expression. - pub expr: ScalarExpr, - /// The type of the expression. - pub typ: ColumnType, -} - -impl TypedExpr { - pub fn new(expr: ScalarExpr, typ: ColumnType) -> Self { - Self { expr, typ } - } -} - -/// A scalar expression, which can be evaluated to a value. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub enum ScalarExpr { - /// A column of the input row - Column(usize), - /// A literal value. - /// Extra type info to know original type even when it is null - Literal(Value, ConcreteDataType), - /// A call to an unmaterializable function. - /// - /// These functions cannot be evaluated by `ScalarExpr::eval`. They must - /// be transformed away by a higher layer. - CallUnmaterializable(UnmaterializableFunc), - CallUnary { - func: UnaryFunc, - expr: Box, - }, - CallBinary { - func: BinaryFunc, - expr1: Box, - expr2: Box, - }, - CallVariadic { - func: VariadicFunc, - exprs: Vec, - }, - CallDf { - /// invariant: the input args set inside this [`DfScalarFunction`] is - /// always col(0) to col(n-1) where n is the length of `expr` - df_scalar_fn: DfScalarFunction, - exprs: Vec, - }, - /// Conditionally evaluated expressions. - /// - /// It is important that `then` and `els` only be evaluated if - /// `cond` is true or not, respectively. This is the only way - /// users can guard execution (other logical operator do not - /// short-circuit) and we need to preserve that. - If { - cond: Box, - then: Box, - els: Box, - }, - InList { - expr: Box, - list: Vec, - }, -} - -impl ScalarExpr { - pub fn with_type(self, typ: ColumnType) -> TypedExpr { - TypedExpr::new(self, typ) - } - - /// try to determine the type of the expression - pub fn typ(&self, context: &[ColumnType]) -> Result { - match self { - ScalarExpr::Column(i) => context.get(*i).cloned().ok_or_else(|| { - UnexpectedSnafu { - reason: format!("column index {} out of range of len={}", i, context.len()), - } - .build() - }), - ScalarExpr::Literal(_, typ) => Ok(ColumnType::new_nullable(typ.clone())), - ScalarExpr::CallUnmaterializable(func) => { - Ok(ColumnType::new_nullable(func.signature().output)) - } - ScalarExpr::CallUnary { func, .. } => { - Ok(ColumnType::new_nullable(func.signature().output)) - } - ScalarExpr::CallBinary { func, .. } => { - Ok(ColumnType::new_nullable(func.signature().output)) - } - ScalarExpr::CallVariadic { func, .. } => { - Ok(ColumnType::new_nullable(func.signature().output)) - } - ScalarExpr::If { then, .. } => then.typ(context), - ScalarExpr::CallDf { df_scalar_fn, .. } => { - let arrow_typ = df_scalar_fn - .fn_impl - // TODO(discord9): get scheme from args instead? - .data_type(df_scalar_fn.df_schema.as_arrow()) - .context({ - DatafusionSnafu { - context: "Failed to get data type from datafusion scalar function", - } - })?; - let typ = ConcreteDataType::try_from(&arrow_typ) - .map_err(BoxedError::new) - .context(crate::error::ExternalSnafu)?; - Ok(ColumnType::new_nullable(typ)) - } - ScalarExpr::InList { expr, .. } => expr.typ(context), - } - } -} - -impl ScalarExpr { - pub fn cast(self, typ: ConcreteDataType) -> Self { - ScalarExpr::CallUnary { - func: UnaryFunc::Cast(typ), - expr: Box::new(self), - } - } - - /// apply optimization to the expression, like flatten variadic function - pub fn optimize(&mut self) { - self.flatten_variadic_fn(); - } - - /// Because Substrait's `And`/`Or` function is binary, but FlowPlan's - /// `And`/`Or` function is variadic, we need to flatten the `And` function if multiple `And`/`Or` functions are nested. - fn flatten_variadic_fn(&mut self) { - if let ScalarExpr::CallVariadic { func, exprs } = self { - let mut new_exprs = vec![]; - for expr in std::mem::take(exprs) { - if let ScalarExpr::CallVariadic { - func: inner_func, - exprs: mut inner_exprs, - } = expr - { - if *func == inner_func { - for inner_expr in inner_exprs.iter_mut() { - inner_expr.flatten_variadic_fn(); - } - new_exprs.extend(inner_exprs); - } - } else { - new_exprs.push(expr); - } - } - *exprs = new_exprs; - } - } -} - -impl ScalarExpr { - /// Call a unary function on this expression. - pub fn call_unary(self, func: UnaryFunc) -> Self { - ScalarExpr::CallUnary { - func, - expr: Box::new(self), - } - } - - /// Call a binary function on this expression and another. - pub fn call_binary(self, other: Self, func: BinaryFunc) -> Self { - ScalarExpr::CallBinary { - func, - expr1: Box::new(self), - expr2: Box::new(other), - } - } - - pub fn eval_batch(&self, batch: &Batch) -> Result { - match self { - ScalarExpr::Column(i) => Ok(batch.batch()[*i].clone()), - ScalarExpr::Literal(val, dt) => Ok(Helper::try_from_scalar_value( - val.try_to_scalar_value(dt).context(DataTypeSnafu { - msg: "Failed to convert literal to scalar value", - })?, - batch.row_count(), - None, - ) - .context(DataTypeSnafu { - msg: "Failed to convert scalar value to vector ref when parsing literal", - })?), - ScalarExpr::CallUnmaterializable(_) => OptimizeSnafu { - reason: "Can't eval unmaterializable function", - } - .fail()?, - ScalarExpr::CallUnary { func, expr } => func.eval_batch(batch, expr), - ScalarExpr::CallBinary { func, expr1, expr2 } => func.eval_batch(batch, expr1, expr2), - ScalarExpr::CallVariadic { func, exprs } => func.eval_batch(batch, exprs), - ScalarExpr::CallDf { - df_scalar_fn, - exprs, - } => df_scalar_fn.eval_batch(batch, exprs), - ScalarExpr::If { cond, then, els } => Self::eval_if_then(batch, cond, then, els), - ScalarExpr::InList { expr, list } => Self::eval_in_list(batch, expr, list), - } - } - - fn eval_in_list( - batch: &Batch, - expr: &ScalarExpr, - list: &[ScalarExpr], - ) -> Result { - let eval_list = list - .iter() - .map(|e| e.eval_batch(batch)) - .collect::, _>>()?; - let eval_expr = expr.eval_batch(batch)?; - - ensure!( - eval_list - .iter() - .all(|v| v.data_type() == eval_expr.data_type()), - TypeMismatchSnafu { - expected: eval_expr.data_type(), - actual: eval_list - .iter() - .find(|v| v.data_type() != eval_expr.data_type()) - .map(|v| v.data_type()) - .unwrap(), - } - ); - - let lhs = eval_expr.to_arrow_array(); - - let found = eval_list - .iter() - .map(|v| v.to_arrow_array()) - .try_fold( - BooleanArray::new(BooleanBuffer::new_unset(batch.row_count()), None), - |result, in_list_elem| -> Result { - let rhs = compare_with_eq(&lhs, &in_list_elem, false)?; - - Ok(or_kleene(&result, &rhs)?) - }, - ) - .with_context(|_| crate::expr::error::DatafusionSnafu { - context: "Failed to compare eval_expr with eval_list", - })?; - - let res = BooleanVector::from(found); - - Ok(Arc::new(res)) - } - - /// NOTE: this if then eval impl assume all given expr are pure, and will not change the state of the world - /// since it will evaluate both then and else branch and filter the result - fn eval_if_then( - batch: &Batch, - cond: &ScalarExpr, - then: &ScalarExpr, - els: &ScalarExpr, - ) -> Result { - let conds = cond.eval_batch(batch)?; - let bool_conds = conds - .as_any() - .downcast_ref::() - .context({ - TypeMismatchSnafu { - expected: ConcreteDataType::boolean_datatype(), - actual: conds.data_type(), - } - })? - .as_boolean_array(); - - let indices = bool_conds - .into_iter() - .enumerate() - .map(|(idx, b)| { - ( - match b { - Some(true) => 0, // then branch vector - Some(false) => 1, // else branch vector - None => 2, // null vector - }, - idx, - ) - }) - .collect_vec(); - - let then_input_vec = then.eval_batch(batch)?; - let else_input_vec = els.eval_batch(batch)?; - - ensure!( - then_input_vec.data_type() == else_input_vec.data_type(), - TypeMismatchSnafu { - expected: then_input_vec.data_type(), - actual: else_input_vec.data_type(), - } - ); - - ensure!( - then_input_vec.len() == else_input_vec.len() - && then_input_vec.len() == batch.row_count(), - InvalidArgumentSnafu { - reason: format!( - "then and else branch must have the same length(found {} and {}) which equals input batch's row count(which is {})", - then_input_vec.len(), - else_input_vec.len(), - batch.row_count() - ) - } - ); - - fn new_nulls(dt: &arrow_schema::DataType, len: usize) -> ArrayRef { - let data = ArrayData::new_null(dt, len); - make_array(data) - } - - let null_input_vec = new_nulls( - &then_input_vec.data_type().as_arrow_type(), - batch.row_count(), - ); - - let interleave_values = vec![ - then_input_vec.to_arrow_array(), - else_input_vec.to_arrow_array(), - null_input_vec, - ]; - let int_ref: Vec<_> = interleave_values.iter().map(|x| x.as_ref()).collect(); - - let interleave_res_arr = - arrow::compute::interleave(&int_ref, &indices).context(ArrowSnafu { - context: "Failed to interleave output arrays", - })?; - let res_vec = Helper::try_into_vector(interleave_res_arr).context(DataTypeSnafu { - msg: "Failed to convert arrow array to vector", - })?; - Ok(res_vec) - } - - /// Eval this expression with the given values. - /// - /// TODO(discord9): add tests to make sure `eval_batch` is the same as `eval` in - /// most cases - pub fn eval(&self, values: &[Value]) -> Result { - match self { - ScalarExpr::Column(index) => Ok(values[*index].clone()), - ScalarExpr::Literal(row_res, _ty) => Ok(row_res.clone()), - ScalarExpr::CallUnmaterializable(_) => OptimizeSnafu { - reason: "Can't eval unmaterializable function".to_string(), - } - .fail(), - ScalarExpr::CallUnary { func, expr } => func.eval(values, expr), - ScalarExpr::CallBinary { func, expr1, expr2 } => func.eval(values, expr1, expr2), - ScalarExpr::CallVariadic { func, exprs } => func.eval(values, exprs), - ScalarExpr::If { cond, then, els } => match cond.eval(values) { - Ok(Value::Boolean(true)) => then.eval(values), - Ok(Value::Boolean(false)) => els.eval(values), - _ => InvalidArgumentSnafu { - reason: "if condition must be boolean".to_string(), - } - .fail(), - }, - ScalarExpr::CallDf { - df_scalar_fn, - exprs, - } => df_scalar_fn.eval(values, exprs), - ScalarExpr::InList { expr, list } => { - let eval_expr = expr.eval(values)?; - let eval_list = list - .iter() - .map(|v| v.eval(values)) - .collect::, _>>()?; - let found = eval_list.iter().any(|item| *item == eval_expr); - Ok(Value::Boolean(found)) - } - } - } - - /// Rewrites column indices with their value in `permutation`. - /// - /// This method is applicable even when `permutation` is not a - /// strict permutation, and it only needs to have entries for - /// each column referenced in `self`. - pub fn permute(&mut self, permutation: &[usize]) -> Result<(), Error> { - // check first so that we don't end up with a partially permuted expression - ensure!( - self.get_all_ref_columns() - .into_iter() - .all(|i| i < permutation.len()), - InvalidQuerySnafu { - reason: format!( - "permutation {:?} is not a valid permutation for expression {:?}", - permutation, self - ), - } - ); - - self.visit_mut_post_nolimit(&mut |e| { - if let ScalarExpr::Column(old_i) = e { - *old_i = permutation[*old_i]; - } - Ok(()) - })?; - Ok(()) - } - - /// Rewrites column indices with their value in `permutation`. - /// - /// This method is applicable even when `permutation` is not a - /// strict permutation, and it only needs to have entries for - /// each column referenced in `self`. - pub fn permute_map(&mut self, permutation: &BTreeMap) -> Result<(), Error> { - // check first so that we don't end up with a partially permuted expression - ensure!( - self.get_all_ref_columns() - .is_subset(&permutation.keys().cloned().collect()), - InvalidQuerySnafu { - reason: format!( - "permutation {:?} is not a valid permutation for expression {:?}", - permutation, self - ), - } - ); - - self.visit_mut_post_nolimit(&mut |e| { - if let ScalarExpr::Column(old_i) = e { - *old_i = permutation[old_i]; - } - Ok(()) - }) - } - - /// Returns the set of columns that are referenced by `self`. - pub fn get_all_ref_columns(&self) -> BTreeSet { - let mut support = BTreeSet::new(); - self.visit_post_nolimit(&mut |e| { - if let ScalarExpr::Column(i) = e { - support.insert(*i); - } - Ok(()) - }) - .unwrap(); - support - } - - /// Return true if the expression is a column reference. - pub fn is_column(&self) -> bool { - matches!(self, ScalarExpr::Column(_)) - } - - /// Cast the expression to a column reference if it is one. - pub fn as_column(&self) -> Option { - if let ScalarExpr::Column(i) = self { - Some(*i) - } else { - None - } - } - - /// Cast the expression to a literal if it is one. - pub fn as_literal(&self) -> Option { - if let ScalarExpr::Literal(lit, _column_type) = self { - Some(lit.clone()) - } else { - None - } - } - - /// Return true if the expression is a literal. - pub fn is_literal(&self) -> bool { - matches!(self, ScalarExpr::Literal(..)) - } - - /// Return true if the expression is a literal true. - pub fn is_literal_true(&self) -> bool { - Some(Value::Boolean(true)) == self.as_literal() - } - - /// Return true if the expression is a literal false. - pub fn is_literal_false(&self) -> bool { - Some(Value::Boolean(false)) == self.as_literal() - } - - /// Return true if the expression is a literal null. - pub fn is_literal_null(&self) -> bool { - Some(Value::Null) == self.as_literal() - } - - /// Build a literal null - pub fn literal_null() -> Self { - ScalarExpr::Literal(Value::Null, ConcreteDataType::null_datatype()) - } - - /// Build a literal from value and type - pub fn literal(res: Value, typ: ConcreteDataType) -> Self { - ScalarExpr::Literal(res, typ) - } - - /// Build a literal false - pub fn literal_false() -> Self { - ScalarExpr::Literal(Value::Boolean(false), ConcreteDataType::boolean_datatype()) - } - - /// Build a literal true - pub fn literal_true() -> Self { - ScalarExpr::Literal(Value::Boolean(true), ConcreteDataType::boolean_datatype()) - } -} - -impl ScalarExpr { - /// visit post-order without stack call limit, but may cause stack overflow - fn visit_post_nolimit(&self, f: &mut F) -> Result<(), EvalError> - where - F: FnMut(&Self) -> Result<(), EvalError>, - { - self.visit_children(|e| e.visit_post_nolimit(f))?; - f(self) - } - - fn visit_children(&self, mut f: F) -> Result<(), EvalError> - where - F: FnMut(&Self) -> Result<(), EvalError>, - { - match self { - ScalarExpr::Column(_) - | ScalarExpr::Literal(_, _) - | ScalarExpr::CallUnmaterializable(_) => Ok(()), - ScalarExpr::CallUnary { expr, .. } => f(expr), - ScalarExpr::CallBinary { expr1, expr2, .. } => { - f(expr1)?; - f(expr2) - } - ScalarExpr::CallVariadic { exprs, .. } => { - for expr in exprs { - f(expr)?; - } - Ok(()) - } - ScalarExpr::If { cond, then, els } => { - f(cond)?; - f(then)?; - f(els) - } - ScalarExpr::CallDf { - df_scalar_fn: _, - exprs, - } => { - for expr in exprs { - f(expr)?; - } - Ok(()) - } - ScalarExpr::InList { expr, list } => { - f(expr)?; - for item in list { - f(item)?; - } - Ok(()) - } - } - } - - fn visit_mut_post_nolimit(&mut self, f: &mut F) -> Result<(), Error> - where - F: FnMut(&mut Self) -> Result<(), Error>, - { - self.visit_mut_children(|e: &mut Self| e.visit_mut_post_nolimit(f))?; - f(self) - } - - fn visit_mut_children(&mut self, mut f: F) -> Result<(), Error> - where - F: FnMut(&mut Self) -> Result<(), Error>, - { - match self { - ScalarExpr::Column(_) - | ScalarExpr::Literal(_, _) - | ScalarExpr::CallUnmaterializable(_) => Ok(()), - ScalarExpr::CallUnary { expr, .. } => f(expr), - ScalarExpr::CallBinary { expr1, expr2, .. } => { - f(expr1)?; - f(expr2) - } - ScalarExpr::CallVariadic { exprs, .. } => { - for expr in exprs { - f(expr)?; - } - Ok(()) - } - ScalarExpr::If { cond, then, els } => { - f(cond)?; - f(then)?; - f(els) - } - ScalarExpr::CallDf { - df_scalar_fn: _, - exprs, - } => { - for expr in exprs { - f(expr)?; - } - Ok(()) - } - ScalarExpr::InList { expr, list } => { - f(expr)?; - for item in list { - f(item)?; - } - Ok(()) - } - } - } -} - -impl ScalarExpr { - /// if expr contains function `Now` - pub fn contains_temporal(&self) -> bool { - let mut contains = false; - self.visit_post_nolimit(&mut |e| { - if let ScalarExpr::CallUnmaterializable(UnmaterializableFunc::Now) = e { - contains = true; - } - Ok(()) - }) - .unwrap(); - contains - } - - /// extract lower or upper bound of `Now` for expr, where `lower bound <= expr < upper bound` - /// - /// returned bool indicates whether the bound is upper bound: - /// - /// false for lower bound, true for upper bound - /// TODO(discord9): allow simple transform like `now() + a < b` to `now() < b - a` - pub fn extract_bound(&self) -> Result<(Option, Option), Error> { - let unsupported_err = |msg: &str| { - UnsupportedTemporalFilterSnafu { - reason: msg.to_string(), - } - .fail() - }; - - let Self::CallBinary { - mut func, - mut expr1, - mut expr2, - } = self.clone() - else { - return unsupported_err("Not a binary expression"); - }; - - // TODO(discord9): support simple transform like `now() + a < b` to `now() < b - a` - - let expr1_is_now = *expr1 == ScalarExpr::CallUnmaterializable(UnmaterializableFunc::Now); - let expr2_is_now = *expr2 == ScalarExpr::CallUnmaterializable(UnmaterializableFunc::Now); - - if !(expr1_is_now ^ expr2_is_now) { - return unsupported_err("None of the sides of the comparison is `now()`"); - } - - if expr2_is_now { - std::mem::swap(&mut expr1, &mut expr2); - func = BinaryFunc::reverse_compare(&func)?; - } - - let step = |expr: ScalarExpr| expr.call_unary(UnaryFunc::StepTimestamp); - match func { - // now == expr2 -> now <= expr2 && now < expr2 + 1 - BinaryFunc::Eq => Ok((Some(*expr2.clone()), Some(step(*expr2)))), - // now < expr2 -> now < expr2 - BinaryFunc::Lt => Ok((None, Some(*expr2))), - // now <= expr2 -> now < expr2 + 1 - BinaryFunc::Lte => Ok((None, Some(step(*expr2)))), - // now > expr2 -> now >= expr2 + 1 - BinaryFunc::Gt => Ok((Some(step(*expr2)), None)), - // now >= expr2 -> now >= expr2 - BinaryFunc::Gte => Ok((Some(*expr2), None)), - _ => unreachable!("Already checked"), - } - } -} - -#[cfg(test)] -mod test { - use datatypes::vectors::{Int32Vector, Vector}; - use pretty_assertions::assert_eq; - - use super::*; - - #[test] - fn test_extract_bound() { - let test_list: [(ScalarExpr, Result<_, EvalError>); 5] = [ - // col(0) == now - ( - ScalarExpr::CallBinary { - func: BinaryFunc::Eq, - expr1: Box::new(ScalarExpr::CallUnmaterializable(UnmaterializableFunc::Now)), - expr2: Box::new(ScalarExpr::Column(0)), - }, - Ok(( - Some(ScalarExpr::Column(0)), - Some(ScalarExpr::CallUnary { - func: UnaryFunc::StepTimestamp, - expr: Box::new(ScalarExpr::Column(0)), - }), - )), - ), - // now < col(0) - ( - ScalarExpr::CallBinary { - func: BinaryFunc::Lt, - expr1: Box::new(ScalarExpr::CallUnmaterializable(UnmaterializableFunc::Now)), - expr2: Box::new(ScalarExpr::Column(0)), - }, - Ok((None, Some(ScalarExpr::Column(0)))), - ), - // now <= col(0) - ( - ScalarExpr::CallBinary { - func: BinaryFunc::Lte, - expr1: Box::new(ScalarExpr::CallUnmaterializable(UnmaterializableFunc::Now)), - expr2: Box::new(ScalarExpr::Column(0)), - }, - Ok(( - None, - Some(ScalarExpr::CallUnary { - func: UnaryFunc::StepTimestamp, - expr: Box::new(ScalarExpr::Column(0)), - }), - )), - ), - // now > col(0) -> now >= col(0) + 1 - ( - ScalarExpr::CallBinary { - func: BinaryFunc::Gt, - expr1: Box::new(ScalarExpr::CallUnmaterializable(UnmaterializableFunc::Now)), - expr2: Box::new(ScalarExpr::Column(0)), - }, - Ok(( - Some(ScalarExpr::CallUnary { - func: UnaryFunc::StepTimestamp, - expr: Box::new(ScalarExpr::Column(0)), - }), - None, - )), - ), - // now >= col(0) - ( - ScalarExpr::CallBinary { - func: BinaryFunc::Gte, - expr1: Box::new(ScalarExpr::CallUnmaterializable(UnmaterializableFunc::Now)), - expr2: Box::new(ScalarExpr::Column(0)), - }, - Ok((Some(ScalarExpr::Column(0)), None)), - ), - ]; - for (expr, expected) in test_list.into_iter() { - let actual = expr.extract_bound(); - // EvalError is not Eq, so we need to compare the error message - match (actual, expected) { - (Ok(l), Ok(r)) => assert_eq!(l, r), - (l, r) => panic!("expected: {:?}, actual: {:?}", r, l), - } - } - } - - #[test] - fn test_bad_permute() { - let mut expr = ScalarExpr::Column(4); - let permutation = vec![1, 2, 3]; - let res = expr.permute(&permutation); - assert!(matches!(res, Err(Error::InvalidQuery { .. }))); - - let mut expr = ScalarExpr::Column(0); - let permute_map = BTreeMap::from([(1, 2), (3, 4)]); - let res = expr.permute_map(&permute_map); - assert!(matches!(res, Err(Error::InvalidQuery { .. }))); - } - - #[test] - fn test_eval_batch_if_then() { - // TODO(discord9): add more tests - { - let expr = ScalarExpr::If { - cond: Box::new(ScalarExpr::Column(0).call_binary( - ScalarExpr::literal(Value::from(0), ConcreteDataType::int32_datatype()), - BinaryFunc::Eq, - )), - then: Box::new(ScalarExpr::literal( - Value::from(42), - ConcreteDataType::int32_datatype(), - )), - els: Box::new(ScalarExpr::literal( - Value::from(37), - ConcreteDataType::int32_datatype(), - )), - }; - let raw = vec![ - None, - Some(0), - Some(1), - None, - None, - Some(0), - Some(0), - Some(1), - Some(1), - ]; - let raw_len = raw.len(); - let vectors = vec![Int32Vector::from(raw).slice(0, raw_len)]; - - let batch = Batch::try_new(vectors, raw_len).unwrap(); - let expected = Int32Vector::from(vec![ - None, - Some(42), - Some(37), - None, - None, - Some(42), - Some(42), - Some(37), - Some(37), - ]) - .slice(0, raw_len); - assert_eq!(expr.eval_batch(&batch).unwrap(), expected); - - let raw = vec![Some(0)]; - let raw_len = raw.len(); - let vectors = vec![Int32Vector::from(raw).slice(0, raw_len)]; - - let batch = Batch::try_new(vectors, raw_len).unwrap(); - let expected = Int32Vector::from(vec![Some(42)]).slice(0, raw_len); - assert_eq!(expr.eval_batch(&batch).unwrap(), expected); - - let raw: Vec> = vec![]; - let raw_len = raw.len(); - let vectors = vec![Int32Vector::from(raw).slice(0, raw_len)]; - - let batch = Batch::try_new(vectors, raw_len).unwrap(); - let expected = Int32Vector::from(vec![]).slice(0, raw_len); - assert_eq!(expr.eval_batch(&batch).unwrap(), expected); - } - } -} diff --git a/src/flow/src/expr/signature.rs b/src/flow/src/expr/signature.rs deleted file mode 100644 index 526d00b96ed..00000000000 --- a/src/flow/src/expr/signature.rs +++ /dev/null @@ -1,70 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! Function signature, useful for type checking and function resolution. - -use datatypes::data_type::ConcreteDataType; -use serde::{Deserialize, Serialize}; -use smallvec::SmallVec; - -/// Function signature -/// -/// TODO(discord9): use `common_query::signature::Signature` crate -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Deserialize, Serialize, Hash)] -pub struct Signature { - /// the input types, usually not great than two input arg - pub input: SmallVec<[ConcreteDataType; 2]>, - /// Output type - pub output: ConcreteDataType, - /// Generic function - pub generic_fn: GenericFn, -} - -/// Generic function category -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Deserialize, Serialize, Hash)] -pub enum GenericFn { - // aggregate func - Max, - Min, - Sum, - Count, - Any, - All, - // unary func - Not, - IsNull, - IsTrue, - IsFalse, - StepTimestamp, - Cast, - // binary func - Eq, - NotEq, - Lt, - Lte, - Gt, - Gte, - Add, - Sub, - Mul, - Div, - Mod, - // variadic func - And, - Or, - // unmaterized func - Now, - CurrentSchema, - TumbleWindow, -} diff --git a/src/flow/src/expr/utils.rs b/src/flow/src/expr/utils.rs deleted file mode 100644 index 6b9fda5337b..00000000000 --- a/src/flow/src/expr/utils.rs +++ /dev/null @@ -1,348 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! This module contains utility functions for expressions. - -use std::cmp::Ordering; -use std::collections::BTreeMap; - -use datatypes::value::Value; -use snafu::{OptionExt, ensure}; - -use crate::Result; -use crate::error::UnexpectedSnafu; -use crate::expr::ScalarExpr; -use crate::plan::TypedPlan; - -/// Find lower bound for time `current` in given `plan` for the time window expr. -/// -/// i.e. for time window expr being `date_bin(INTERVAL '5 minutes', ts) as time_window` and `current="2021-07-01 00:01:01.000"`, -/// return `Some("2021-07-01 00:00:00.000")` -/// -/// if `plan` doesn't contain a `TIME INDEX` column, return `None` -pub fn find_plan_time_window_expr_lower_bound( - plan: &TypedPlan, - current: common_time::Timestamp, -) -> Result> { - let typ = plan.schema.typ(); - let Some(mut time_index) = typ.time_index else { - return Ok(None); - }; - - let mut cur_plan = plan; - let mut expr_time_index; - - loop { - // follow upward and find deepest time index expr that is not a column ref - expr_time_index = Some(cur_plan.plan.get_nth_expr(time_index).context( - UnexpectedSnafu { - reason: "Failed to find time index expr", - }, - )?); - - if let Some(ScalarExpr::Column(i)) = expr_time_index { - time_index = i; - } else { - break; - } - if let Some(input) = cur_plan.plan.get_first_input_plan() { - cur_plan = input; - } else { - break; - } - } - - let expr_time_index = expr_time_index.context(UnexpectedSnafu { - reason: "Failed to find time index expr", - })?; - - let ts_col = expr_time_index - .get_all_ref_columns() - .first() - .cloned() - .context(UnexpectedSnafu { - reason: "Failed to find time index column", - })?; - - find_time_window_lower_bound(&expr_time_index, ts_col, current) -} - -/// Find the lower bound of time window in given `expr` and `current` timestamp. -/// -/// i.e. for `current="2021-07-01 00:01:01.000"` and `expr=date_bin(INTERVAL '5 minutes', ts) as time_window` and `ts_col=ts`, -/// return `Some("2021-07-01 00:00:00.000")` since it's the lower bound -/// of current time window given the current timestamp -/// -/// if return None, meaning this time window have no lower bound -pub fn find_time_window_lower_bound( - expr: &ScalarExpr, - ts_col_idx: usize, - current: common_time::Timestamp, -) -> Result> { - let all_ref_columns = expr.get_all_ref_columns(); - - ensure!( - all_ref_columns.contains(&ts_col_idx), - UnexpectedSnafu { - reason: format!( - "Expected column {} to be referenced in expression {expr:?}", - ts_col_idx - ), - } - ); - - ensure!( - all_ref_columns.len() == 1, - UnexpectedSnafu { - reason: format!( - "Expect only one column to be referenced in expression {expr:?}, found {all_ref_columns:?}" - ), - } - ); - - let permute_map = BTreeMap::from([(ts_col_idx, 0usize)]); - - let mut rewrote_expr = expr.clone(); - - rewrote_expr.permute_map(&permute_map)?; - - fn eval_to_timestamp(expr: &ScalarExpr, values: &[Value]) -> Result { - let val = expr.eval(values)?; - if let Value::Timestamp(ts) = val { - Ok(ts) - } else { - UnexpectedSnafu { - reason: format!("Expected timestamp in expression {expr:?} but got {val:?}"), - } - .fail()? - } - } - - let cur_time_window = eval_to_timestamp(&rewrote_expr, &[current.into()])?; - - // search to find the lower bound - let mut offset: i64 = 1; - let lower_bound; - let mut upper_bound = Some(current); - // first expontial probe to found a range for binary search - loop { - let Some(next_val) = current.value().checked_sub(offset) else { - // no lower bound - return Ok(None); - }; - - let prev_time_probe = common_time::Timestamp::new(next_val, current.unit()); - - let prev_time_window = eval_to_timestamp(&rewrote_expr, &[prev_time_probe.into()])?; - - match prev_time_window.cmp(&cur_time_window) { - Ordering::Less => { - lower_bound = Some(prev_time_probe); - break; - } - Ordering::Equal => { - upper_bound = Some(prev_time_probe); - } - Ordering::Greater => { - UnexpectedSnafu { - reason: format!( - "Unsupported time window expression {rewrote_expr:?}, expect monotonic increasing for time window expression {expr:?}" - ), - } - .fail()? - } - } - - let Some(new_offset) = offset.checked_mul(2) else { - // no lower bound - return Ok(None); - }; - offset = new_offset; - } - - // binary search for the lower bound - - ensure!( - lower_bound.map(|v| v.unit()) == upper_bound.map(|v| v.unit()), - UnexpectedSnafu { - reason: format!( - " unit mismatch for time window expression {expr:?}, found {lower_bound:?} and {upper_bound:?}" - ), - } - ); - - let output_unit = lower_bound.expect("should have lower bound").unit(); - - let mut low = lower_bound.expect("should have lower bound").value(); - let mut high = upper_bound.expect("should have upper bound").value(); - while low < high { - let mid = (low + high) / 2; - let mid_probe = common_time::Timestamp::new(mid, output_unit); - let mid_time_window = eval_to_timestamp(&rewrote_expr, &[mid_probe.into()])?; - - match mid_time_window.cmp(&cur_time_window) { - Ordering::Less => low = mid + 1, - Ordering::Equal => high = mid, - Ordering::Greater => UnexpectedSnafu { - reason: format!("Binary search failed for time window expression {expr:?}"), - } - .fail()?, - } - } - - let final_lower_bound_for_time_window = common_time::Timestamp::new(low, output_unit); - - Ok(Some(final_lower_bound_for_time_window)) -} - -#[cfg(test)] -mod test { - use pretty_assertions::assert_eq; - - use super::*; - use crate::plan::{Plan, TypedPlan}; - use crate::test_utils::{create_test_ctx, create_test_query_engine, sql_to_substrait}; - - #[tokio::test] - async fn test_plan_time_window_lower_bound() { - let testcases = [ - // no time index - ( - "SELECT date_bin('5 minutes', ts) FROM numbers_with_ts;", - "2021-07-01 00:01:01.000", - None, - ), - // time index - ( - "SELECT date_bin('5 minutes', ts) as time_window FROM numbers_with_ts GROUP BY time_window;", - "2021-07-01 00:01:01.000", - Some("2021-07-01 00:00:00.000"), - ), - // time index with other fields - ( - "SELECT sum(number) as sum_up, date_bin('5 minutes', ts) as time_window FROM numbers_with_ts GROUP BY time_window;", - "2021-07-01 00:01:01.000", - Some("2021-07-01 00:00:00.000"), - ), - // time index with other pks - ( - "SELECT number, date_bin('5 minutes', ts) as time_window FROM numbers_with_ts GROUP BY time_window, number;", - "2021-07-01 00:01:01.000", - Some("2021-07-01 00:00:00.000"), - ), - ]; - let engine = create_test_query_engine(); - - for (sql, current, expected) in &testcases { - let plan = sql_to_substrait(engine.clone(), sql).await; - let mut ctx = create_test_ctx(); - let flow_plan = TypedPlan::from_substrait_plan(&mut ctx, &plan) - .await - .unwrap(); - - let current = common_time::Timestamp::from_str(current, None).unwrap(); - - let expected = - expected.map(|expected| common_time::Timestamp::from_str(expected, None).unwrap()); - - assert_eq!( - find_plan_time_window_expr_lower_bound(&flow_plan, current).unwrap(), - expected - ); - } - } - - #[tokio::test] - async fn test_timewindow_lower_bound() { - let testcases = [ - ( - ("'5 minutes'", "ts", Some("2021-07-01 00:00:00.000")), - "2021-07-01 00:01:01.000", - "2021-07-01 00:00:00.000", - ), - ( - ("'5 minutes'", "ts", None), - "2021-07-01 00:01:01.000", - "2021-07-01 00:00:00.000", - ), - ( - ("'5 minutes'", "ts", None), - "2021-07-01 00:00:00.000", - "2021-07-01 00:00:00.000", - ), - // test edge cases - ( - ("'5 minutes'", "ts", None), - "2021-07-01 00:05:00.000", - "2021-07-01 00:05:00.000", - ), - ( - ("'5 minutes'", "ts", None), - "2021-07-01 00:04:59.999", - "2021-07-01 00:00:00.000", - ), - ( - ("'5 minutes'", "ts", None), - "2021-07-01 00:04:59.999999999", - "2021-07-01 00:00:00.000", - ), - ( - ("'5 minutes'", "ts", None), - "2021-07-01 00:04:59.999999999999", - "2021-07-01 00:00:00.000", - ), - ( - ("'5 minutes'", "ts", None), - "2021-07-01 00:04:59.999999999999999", - "2021-07-01 00:00:00.000", - ), - ]; - let engine = create_test_query_engine(); - - for (args, current, expected) in testcases { - let sql = if let Some(origin) = args.2 { - format!( - "SELECT date_bin({}, {}, '{origin}') FROM numbers_with_ts;", - args.0, args.1 - ) - } else { - format!( - "SELECT date_bin({}, {}) FROM numbers_with_ts;", - args.0, args.1 - ) - }; - let plan = sql_to_substrait(engine.clone(), &sql).await; - let mut ctx = create_test_ctx(); - let flow_plan = TypedPlan::from_substrait_plan(&mut ctx, &plan) - .await - .unwrap(); - - let expr = { - let mfp = flow_plan.plan; - let Plan::Mfp { mfp, .. } = mfp else { - unreachable!() - }; - mfp.expressions[0].clone() - }; - - let current = common_time::Timestamp::from_str(current, None).unwrap(); - - let res = find_time_window_lower_bound(&expr, 1, current).unwrap(); - - let expected = Some(common_time::Timestamp::from_str(expected, None).unwrap()); - - assert_eq!(res, expected); - } - } -} diff --git a/src/flow/src/lib.rs b/src/flow/src/lib.rs index e55e7bab763..acba63f51d3 100644 --- a/src/flow/src/lib.rs +++ b/src/flow/src/lib.rs @@ -12,9 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! This crate manage dataflow in Greptime, including adapter, expr, plan, repr and utils. -//! It can transform substrait plan into it's own plan and execute it. -//! It also contains definition of expression, adapter and plan, and internal state management. +//! Flow execution and flownode services. #![allow(dead_code)] #![warn(clippy::too_many_lines)] @@ -25,17 +23,14 @@ // allow unused for now because it should be use later mod adapter; pub(crate) mod batching_mode; -mod compute; mod df_optimizer; pub(crate) mod engine; pub mod error; mod expr; pub mod heartbeat; mod metrics; -mod plan; mod repr; mod server; -mod transform; mod utils; #[cfg(test)] @@ -46,8 +41,6 @@ pub use adapter::{FlowConfig, FlowStreamingEngineRef, StreamingEngine}; pub use batching_mode::frontend_client::{FrontendClient, GrpcQueryHandlerWithBoxedError}; pub(crate) use engine::{CreateFlowArgs, FlowId, TableName}; pub use error::{Error, Result}; -pub use server::{ - FlownodeBuilder, FlownodeInstance, FlownodeServer, FlownodeServiceBuilder, FrontendInvoker, -}; +pub use server::{FlownodeBuilder, FlownodeInstance, FlownodeServer, FlownodeServiceBuilder}; pub use crate::adapter::FlownodeOptions; diff --git a/src/flow/src/plan.rs b/src/flow/src/plan.rs deleted file mode 100644 index b2c91015e06..00000000000 --- a/src/flow/src/plan.rs +++ /dev/null @@ -1,270 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! This module contain basic definition for dataflow's plan -//! that can be translate to hydro dataflow - -mod join; -mod reduce; - -use std::collections::BTreeSet; - -use crate::error::Error; -use crate::expr::{GlobalId, Id, LocalId, MapFilterProject, SafeMfpPlan, ScalarExpr, TypedExpr}; -use crate::plan::join::JoinPlan; -pub(crate) use crate::plan::reduce::{AccumulablePlan, AggrWithIndex, KeyValPlan, ReducePlan}; -use crate::repr::{DiffRow, RelationDesc}; - -/// A plan for a dataflow component. But with type to indicate the output type of the relation. -#[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd)] -pub struct TypedPlan { - /// output type of the relation - pub schema: RelationDesc, - /// The untyped plan. - pub plan: Plan, -} - -impl TypedPlan { - /// directly apply a mfp to the plan - pub fn mfp(self, mfp: SafeMfpPlan) -> Result { - let new_type = self.schema.apply_mfp(&mfp)?; - let mfp = mfp.mfp; - let plan = match self.plan { - Plan::Mfp { - input, - mfp: old_mfp, - } => Plan::Mfp { - input, - mfp: MapFilterProject::compose(old_mfp, mfp)?, - }, - _ => Plan::Mfp { - input: Box::new(self), - mfp, - }, - }; - Ok(TypedPlan { - schema: new_type, - plan, - }) - } - - /// project the plan to the given expressions - pub fn projection(self, exprs: Vec) -> Result { - let input_arity = self.schema.typ.column_types.len(); - let output_arity = exprs.len(); - - let (exprs, _expr_typs): (Vec<_>, Vec<_>) = exprs - .into_iter() - .map(|TypedExpr { expr, typ }| (expr, typ)) - .unzip(); - let mfp = MapFilterProject::new(input_arity) - .map(exprs)? - .project(input_arity..input_arity + output_arity)? - .into_safe(); - let out_typ = self.schema.apply_mfp(&mfp)?; - - let mfp = mfp.mfp; - // special case for mfp to compose when the plan is already mfp - let plan = match self.plan { - Plan::Mfp { - input, - mfp: old_mfp, - } => Plan::Mfp { - input, - mfp: MapFilterProject::compose(old_mfp, mfp)?, - }, - _ => Plan::Mfp { - input: Box::new(self), - mfp, - }, - }; - Ok(TypedPlan { - schema: out_typ, - plan, - }) - } - - /// Add a new filter to the plan, will filter out the records that do not satisfy the filter - pub fn filter(self, filter: TypedExpr) -> Result { - let typ = self.schema.clone(); - let plan = match self.plan { - Plan::Mfp { - input, - mfp: old_mfp, - } => Plan::Mfp { - input, - mfp: old_mfp.filter(vec![filter.expr])?, - }, - _ => Plan::Mfp { - input: Box::new(self), - mfp: MapFilterProject::new(typ.typ.column_types.len()).filter(vec![filter.expr])?, - }, - }; - Ok(TypedPlan { schema: typ, plan }) - } -} - -/// TODO(discord9): support `TableFunc`（by define FlatMap that map 1 to n) -/// Plan describe how to transform data in dataflow -/// -/// This can be considered as a physical plan in dataflow, which describe how to transform data in a streaming manner. -#[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd)] -pub enum Plan { - /// A constant collection of rows. - Constant { rows: Vec }, - /// Get CDC data from an source, be it external reference to an existing source or an internal - /// reference to a `Let` identifier - Get { id: Id }, - /// Create a temporary collection from given `value`, and make this bind only available - /// in scope of `body` - /// - /// Similar to this rust code snippet: - /// ```rust, ignore - /// { - /// let id = value; - /// body - /// } - Let { - id: LocalId, - value: Box, - body: Box, - }, - /// Map, Filter, and Project operators. Chained together. - Mfp { - /// The input collection. - input: Box, - /// Linear operator to apply to each record. - mfp: MapFilterProject, - }, - /// Reduce operator, aggregation by key assembled from KeyValPlan - Reduce { - /// The input collection. - input: Box, - /// A plan for changing input records into key, value pairs. - key_val_plan: KeyValPlan, - /// A plan for performing the reduce. - /// - /// The implementation of reduction has several different strategies based - /// on the properties of the reduction, and the input itself. - reduce_plan: ReducePlan, - }, - /// A multiway relational equijoin, with fused map, filter, and projection. - /// - /// This stage performs a multiway join among `inputs`, using the equality - /// constraints expressed in `plan`. The plan also describes the implementation - /// strategy we will use, and any pushed down per-record work. - Join { - /// An ordered list of inputs that will be joined. - inputs: Vec, - /// Detailed information about the implementation of the join. - /// - /// This includes information about the implementation strategy, but also - /// any map, filter, project work that we might follow the join with, but - /// potentially pushed down into the implementation of the join. - plan: JoinPlan, - }, - /// Adds the contents of the input collections. - /// - /// Importantly, this is *multiset* union, so the multiplicities of records will - /// add. This is in contrast to *set* union, where the multiplicities would be - /// capped at one. A set union can be formed with `Union` followed by `Reduce` - /// implementing the "distinct" operator. - Union { - /// The input collections - inputs: Vec, - /// Whether to consolidate the output, e.g., cancel negated records. - consolidate_output: bool, - }, -} - -impl Plan { - pub fn with_types(self, schema: RelationDesc) -> TypedPlan { - TypedPlan { schema, plan: self } - } -} - -impl Plan { - /// Get nth expr using column ref - pub fn get_nth_expr(&self, n: usize) -> Option { - match self { - Self::Mfp { mfp, .. } => mfp.get_nth_expr(n), - Self::Reduce { key_val_plan, .. } => key_val_plan.get_nth_expr(n), - _ => None, - } - } - - /// Get the first input plan if exists - pub fn get_first_input_plan(&self) -> Option<&TypedPlan> { - match self { - Plan::Let { value, .. } => Some(value), - Plan::Mfp { input, .. } => Some(input), - Plan::Reduce { input, .. } => Some(input), - Plan::Join { inputs, .. } => inputs.first(), - Plan::Union { inputs, .. } => inputs.first(), - _ => None, - } - } - - /// Get mutable ref to the first input plan if exists - pub fn get_mut_first_input_plan(&mut self) -> Option<&mut TypedPlan> { - match self { - Plan::Let { value, .. } => Some(value), - Plan::Mfp { input, .. } => Some(input), - Plan::Reduce { input, .. } => Some(input), - Plan::Join { inputs, .. } => inputs.first_mut(), - Plan::Union { inputs, .. } => inputs.first_mut(), - _ => None, - } - } - - /// Find all the used collection in the plan - pub fn find_used_collection(&self) -> BTreeSet { - fn recur_find_use(plan: &Plan, used: &mut BTreeSet) { - match plan { - Plan::Get { id } => { - match id { - Id::Local(_) => (), - Id::Global(g) => { - used.insert(*g); - } - }; - } - Plan::Let { value, body, .. } => { - recur_find_use(&value.plan, used); - recur_find_use(&body.plan, used); - } - Plan::Mfp { input, .. } => { - recur_find_use(&input.plan, used); - } - Plan::Reduce { input, .. } => { - recur_find_use(&input.plan, used); - } - Plan::Join { inputs, .. } => { - for input in inputs { - recur_find_use(&input.plan, used); - } - } - Plan::Union { inputs, .. } => { - for input in inputs { - recur_find_use(&input.plan, used); - } - } - _ => {} - } - } - let mut ret = Default::default(); - recur_find_use(self, &mut ret); - ret - } -} diff --git a/src/flow/src/plan/join.rs b/src/flow/src/plan/join.rs deleted file mode 100644 index 411b00ee024..00000000000 --- a/src/flow/src/plan/join.rs +++ /dev/null @@ -1,76 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use crate::expr::ScalarExpr; -use crate::plan::SafeMfpPlan; - -/// TODO(discord9): consider impl more join strategies -#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)] -pub enum JoinPlan { - Linear(LinearJoinPlan), -} - -/// Determine if a given row should stay in the output. And apply a map filter project before output the row -#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)] -pub struct JoinFilter { - /// each element in the outer vector will check if each expr in itself can be eval to same value - /// if not, the row will be filtered out. Useful for equi-join(join based on equality of some columns) - pub ready_equivalences: Vec>, - /// Apply a map filter project before output the row - pub before: SafeMfpPlan, -} - -/// A plan for the execution of a linear join. -/// -/// A linear join is a sequence of stages, each of which introduces -/// a new collection. Each stage is represented by a [LinearStagePlan]. -#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)] -pub struct LinearJoinPlan { - /// The source relation from which we start the join. - pub source_relation: usize, - /// The arrangement to use for the source relation, if any - pub source_key: Option>, - /// An initial closure to apply before any stages. - /// - /// Values of `None` indicate the identity closure. - pub initial_closure: Option, - /// A *sequence* of stages to apply one after the other. - pub stage_plans: Vec, - /// A concluding filter to apply after the last stage. - /// - /// Values of `None` indicate the identity closure. - pub final_closure: Option, -} - -/// A plan for the execution of one stage of a linear join. -/// -/// Each stage is a binary join between the current accumulated -/// join results, and a new collection. The former is referred to -/// as the "stream" and the latter the "lookup". -#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)] -pub struct LinearStagePlan { - /// The index of the relation into which we will look up. - pub lookup_relation: usize, - /// The key expressions to use for the stream relation. - pub stream_key: Vec, - /// Columns to retain from the stream relation. - /// These columns are those that are not redundant with `stream_key`, - /// and cannot be read out of the key component of an arrangement. - pub stream_thinning: Vec, - /// The key expressions to use for the lookup relation. - pub lookup_key: Vec, - /// The closure to apply to the concatenation of the key columns, - /// the stream value columns, and the lookup value columns. - pub closure: JoinFilter, -} diff --git a/src/flow/src/plan/reduce.rs b/src/flow/src/plan/reduce.rs deleted file mode 100644 index 65a83756b57..00000000000 --- a/src/flow/src/plan/reduce.rs +++ /dev/null @@ -1,87 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use crate::expr::{AggregateExpr, SafeMfpPlan, ScalarExpr}; - -/// Describe how to extract key-value pair from a `Row` -#[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd)] -pub struct KeyValPlan { - /// Extract key from row - pub key_plan: SafeMfpPlan, - /// Extract value from row - pub val_plan: SafeMfpPlan, -} - -impl KeyValPlan { - /// Get nth expr using column ref - pub fn get_nth_expr(&self, n: usize) -> Option { - self.key_plan.get_nth_expr(n).or_else(|| { - self.val_plan - .get_nth_expr(n - self.key_plan.projection.len()) - }) - } -} - -/// TODO(discord9): def&impl of Hierarchical aggregates(for min/max with support to deletion) and -/// basic aggregates(for other aggregate functions) and mixed aggregate -#[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd)] -pub enum ReducePlan { - /// Plan for not computing any aggregations, just determining the set of - /// distinct keys. - Distinct, - /// Plan for computing only accumulable aggregations. - /// Including simple functions like `sum`, `count`, `min/max`(without deletion) - Accumulable(AccumulablePlan), -} - -/// Accumulable plan for the execution of a reduction. -#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)] -pub struct AccumulablePlan { - /// All of the aggregations we were asked to compute, stored - /// in order. - pub full_aggrs: Vec, - /// All of the non-distinct accumulable aggregates. - /// Each element represents: - /// (index of aggr output, index of value among inputs, aggr expr) - /// These will all be rendered together in one dataflow fragment. - /// - /// Invariant: the output index is the index of the aggregation in `full_aggrs` - /// which means output index is always smaller than the length of `full_aggrs` - pub simple_aggrs: Vec, - /// Same as `simple_aggrs` but for all of the `DISTINCT` accumulable aggregations. - pub distinct_aggrs: Vec, -} - -/// Invariant: the output index is the index of the aggregation in `full_aggrs` -/// which means output index is always smaller than the length of `full_aggrs` -#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)] -pub struct AggrWithIndex { - /// aggregation expression - pub expr: AggregateExpr, - /// index of aggr input among input row - pub input_idx: usize, - /// index of aggr output among output row - pub output_idx: usize, -} - -impl AggrWithIndex { - /// Create a new `AggrWithIndex` - pub fn new(expr: AggregateExpr, input_idx: usize, output_idx: usize) -> Self { - Self { - expr, - input_idx, - output_idx, - } - } -} diff --git a/src/flow/src/repr.rs b/src/flow/src/repr.rs index 715f60594b7..42c1eafd50b 100644 --- a/src/flow/src/repr.rs +++ b/src/flow/src/repr.rs @@ -57,7 +57,7 @@ pub const BROADCAST_CAP: usize = 1024; /// The maximum capacity of the send buffer, to prevent the buffer from growing too large pub const SEND_BUF_CAP: usize = BROADCAST_CAP * 2; -/// Flow worker will try to at least accumulate this many rows before processing them(if one second haven't passed) +/// Batch size used by the batching execution path. pub const BATCH_SIZE: usize = 32 * 16384; /// Convert a value that is or can be converted to Datetime to internal timestamp diff --git a/src/flow/src/repr/relation.rs b/src/flow/src/repr/relation.rs index 817d85f3104..296f1bb4178 100644 --- a/src/flow/src/repr/relation.rs +++ b/src/flow/src/repr/relation.rs @@ -15,12 +15,10 @@ use datafusion_common::DFSchema; use datatypes::data_type::DataType; use datatypes::prelude::ConcreteDataType; -use itertools::Itertools; use serde::{Deserialize, Serialize}; -use snafu::{OptionExt, ResultExt, ensure}; +use snafu::{ResultExt, ensure}; -use crate::error::{DatafusionSnafu, InternalSnafu, InvalidQuerySnafu, Result, UnexpectedSnafu}; -use crate::expr::{SafeMfpPlan, ScalarExpr}; +use crate::error::{DatafusionSnafu, InternalSnafu, InvalidQuerySnafu, Result}; /// a set of column indices that are "keys" for the collection. #[derive(Default, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize, Hash)] @@ -96,73 +94,6 @@ impl RelationType { self } - /// Trying to apply a mpf on current types, will return a new RelationType - /// with the new types, will also try to preserve keys&time index information - /// if the old key&time index columns are preserve in given mfp - /// - /// i.e. old column of size 3, with a mfp's - /// - /// project = `[2, 1]`, - /// - /// the old key = `[1]`, old time index = `[2]`, - /// - /// then new key=`[1]`, new time index=`[0]` - /// - /// note that this function will remove empty keys like key=`[]` will be removed - pub fn apply_mfp(&self, mfp: &SafeMfpPlan) -> Result { - let mfp = &mfp.mfp; - let mut all_types = self.column_types.clone(); - for expr in &mfp.expressions { - let expr_typ = expr.typ(&self.column_types)?; - all_types.push(expr_typ); - } - let all_types = all_types; - let mfp_out_types = mfp - .projection - .iter() - .map(|i| { - all_types.get(*i).cloned().with_context(|| UnexpectedSnafu { - reason: format!( - "MFP index out of bound, len is {}, but the index is {}", - all_types.len(), - *i - ), - }) - }) - .try_collect()?; - - let old_to_new_col = mfp.get_old_to_new_mapping(); - - // since it's just a mfp, we also try to preserve keys&time index information, if they survive mfp transform - let keys = self - .keys - .iter() - .filter_map(|key| { - key.column_indices - .iter() - .map(|old| old_to_new_col.get(old).cloned()) - .collect::>>() - // remove empty keys - .and_then(|v| if v.is_empty() { None } else { Some(v) }) - .map(Key::from) - }) - .collect_vec(); - - let time_index = self - .time_index - .and_then(|old| old_to_new_col.get(&old).cloned()); - let auto_columns = self - .auto_columns - .iter() - .filter_map(|old| old_to_new_col.get(old).cloned()) - .collect_vec(); - Ok(Self { - column_types: mfp_out_types, - keys, - time_index, - auto_columns, - }) - } /// Constructs a `RelationType` representing the relation with no columns and /// no keys. pub fn empty() -> Self { @@ -372,30 +303,6 @@ impl RelationDesc { context: format!("Error when converting to DFSchema: {:?}", arrow_schema), }) } - - /// apply mfp, and also project col names for the projected columns - pub fn apply_mfp(&self, mfp: &SafeMfpPlan) -> Result { - // TODO(discord9): find a way to deduce name at best effect - let names = { - let mfp = &mfp.mfp; - let mut names = self.names.clone(); - for expr in &mfp.expressions { - if let ScalarExpr::Column(i) = expr { - names.push(self.names.get(*i).cloned().flatten()); - } else { - names.push(None); - } - } - mfp.projection - .iter() - .map(|i| names.get(*i).cloned().flatten()) - .collect_vec() - }; - Ok(Self { - typ: self.typ.apply_mfp(mfp)?, - names, - }) - } } impl RelationDesc { diff --git a/src/flow/src/server.rs b/src/flow/src/server.rs index 1b26fbbf15e..bf23d868e65 100644 --- a/src/flow/src/server.rs +++ b/src/flow/src/server.rs @@ -18,52 +18,34 @@ use std::net::SocketAddr; use std::sync::Arc; use api::v1::flow::DirtyWindowRequests; -use api::v1::{RowDeleteRequests, RowInsertRequests}; -use cache::{PARTITION_INFO_CACHE_NAME, TABLE_FLOWNODE_SET_CACHE_NAME, TABLE_ROUTE_CACHE_NAME}; use catalog::CatalogManagerRef; use common_base::Plugins; -use common_datasource::object_store::LocalFileAccess; use common_error::ext::BoxedError; -use common_meta::cache::{LayeredCacheRegistryRef, TableFlownodeSetCacheRef, TableRouteCacheRef}; use common_meta::key::TableMetadataManagerRef; use common_meta::key::flow::FlowMetadataManagerRef; -use common_meta::kv_backend::KvBackendRef; -use common_meta::node_manager::{Flownode, NodeManagerRef}; -use common_meta::procedure_executor::ProcedureExecutorRef; -use common_query::Output; -use common_runtime::JoinHandle; -use common_telemetry::tracing::info; +use common_meta::node_manager::Flownode; use futures::TryStreamExt; use greptime_proto::v1::flow::{FlowRequest, FlowResponse, InsertRequests, flow_server}; use itertools::Itertools; -use operator::delete::Deleter; -use operator::insert::Inserter; -use operator::statement::StatementExecutor; -use partition::cache::PartitionInfoCacheRef; -use partition::manager::PartitionRuleManager; -use query::{QueryEngine, QueryEngineFactory}; +use query::QueryEngineFactory; use servers::add_service; use servers::grpc::builder::GrpcServerBuilder; use servers::grpc::{GrpcServer, GrpcServerConfig}; use servers::http::HttpServerBuilder; use servers::metrics_handler::MetricsHandler; use servers::server::{ServerHandler, ServerHandlers}; -use session::context::QueryContextRef; -use snafu::{OptionExt, ResultExt}; -use tokio::sync::{Mutex, broadcast, oneshot}; +use snafu::ResultExt; use tonic::codec::CompressionEncoding; use tonic::{Request, Response, Status}; use crate::adapter::flownode_impl::{FlowDualEngine, FlowDualEngineRef}; -use crate::adapter::{FlowStreamingEngineRef, create_worker}; use crate::batching_mode::engine::BatchingEngine; use crate::error::{ - CacheRequiredSnafu, DatafusionSnafu, ExternalSnafu, ListFlowsSnafu, ParseAddrSnafu, - ShutdownServerSnafu, StartServerSnafu, UnexpectedSnafu, to_status_with_last_err, + DatafusionSnafu, ExternalSnafu, ListFlowsSnafu, ParseAddrSnafu, ShutdownServerSnafu, + StartServerSnafu, to_status_with_last_err, }; use crate::heartbeat::HeartbeatTask; use crate::metrics::{METRIC_FLOW_PROCESSING_TIME, METRIC_FLOW_ROWS}; -use crate::transform::register_function_to_query_engine; use crate::utils::{SizeReportSender, StateReportHandler}; use crate::{Error, FlownodeOptions, FrontendClient, StreamingEngine}; @@ -163,100 +145,15 @@ pub struct FlownodeServer { /// this struct mostly useful for construct/start and stop the /// flow node server struct FlownodeServerInner { - /// worker shutdown signal, not to be confused with server_shutdown_tx - worker_shutdown_tx: Mutex>, - /// server shutdown signal for shutdown grpc server - server_shutdown_tx: Mutex>, - /// streaming task handler - streaming_task_handler: Mutex>>, - /// state report task handler - state_report_task_handler: Mutex>>, flow_service: FlowService, } impl FlownodeServer { pub fn new(flow_service: FlowService) -> Self { - let (tx, _rx) = broadcast::channel::<()>(1); - let (server_tx, _server_rx) = broadcast::channel::<()>(1); Self { - inner: Arc::new(FlownodeServerInner { - flow_service, - worker_shutdown_tx: Mutex::new(tx), - server_shutdown_tx: Mutex::new(server_tx), - streaming_task_handler: Mutex::new(None), - state_report_task_handler: Mutex::new(None), - }), + inner: Arc::new(FlownodeServerInner { flow_service }), } } - - /// Start the background task for streaming computation. - /// - /// Should be called only after heartbeat is establish, hence can get cluster info - async fn start_workers(&self) -> Result<(), Error> { - let manager_ref = self.inner.flow_service.dual_engine.clone(); - let mut state_report_task_handler = self.inner.state_report_task_handler.lock().await; - let started_state_report_task = state_report_task_handler.is_none(); - if state_report_task_handler.is_none() { - *state_report_task_handler = manager_ref.clone().start_state_report_task().await; - } - drop(state_report_task_handler); - let handle = manager_ref - .streaming_engine() - .run_background(Some(self.inner.worker_shutdown_tx.lock().await.subscribe())); - self.inner - .streaming_task_handler - .lock() - .await - .replace(handle); - - if let Err(err) = self - .inner - .flow_service - .dual_engine - .start_flow_consistent_check_task() - .await - { - self.rollback_started_workers(started_state_report_task) - .await; - return Err(err); - } - - Ok(()) - } - - async fn rollback_started_workers(&self, abort_state_report_task: bool) { - let tx = self.inner.worker_shutdown_tx.lock().await; - if tx.send(()).is_err() { - info!("Receiver dropped, the flow node server has already shutdown"); - } - drop(tx); - - if let Some(handle) = self.inner.streaming_task_handler.lock().await.take() { - handle.abort(); - } - - if abort_state_report_task - && let Some(handle) = self.inner.state_report_task_handler.lock().await.take() - { - handle.abort(); - } - } - - /// Stop the background task for streaming computation. - async fn stop_workers(&self) -> Result<(), Error> { - let tx = self.inner.worker_shutdown_tx.lock().await; - if tx.send(()).is_err() { - info!("Receiver dropped, the flow node server has already shutdown"); - } - // Keep state_report_task_handler alive across worker restarts. - // Dropping it here would permanently lose the report channel receiver. - self.inner - .flow_service - .dual_engine - .stop_flow_consistent_check_task() - .await?; - Ok(()) - } } impl FlownodeServer { @@ -274,6 +171,8 @@ pub struct FlownodeInstance { flownode_server: FlownodeServer, services: ServerHandlers, heartbeat_task: Option, + state_report_task: Option>, + consistent_check_task_started: bool, } impl FlownodeInstance { @@ -282,25 +181,57 @@ impl FlownodeInstance { task.start().await?; } - self.flownode_server.start_workers().await?; + let engine = self.flow_engine(); + // The state-report task owns the only report receiver, so keep it alive + // across an ordinary stop/start cycle of this instance. + if self.state_report_task.is_none() { + self.state_report_task = engine.clone().start_state_report_task().await; + } + if let Err(err) = engine.start_flow_consistent_check_task().await { + self.rollback_background_tasks().await; + return Err(err); + } + self.consistent_check_task_started = true; - self.services.start_all().await.context(StartServerSnafu)?; + if let Err(err) = self.services.start_all().await.context(StartServerSnafu) { + self.rollback_background_tasks().await; + return Err(err); + } Ok(()) } pub async fn shutdown(&mut self) -> Result<(), Error> { - self.services + let services_result = self + .services .shutdown_all() .await - .context(ShutdownServerSnafu)?; + .context(ShutdownServerSnafu); + let tasks_result = self.stop_background_tasks().await; - self.flownode_server.stop_workers().await?; + services_result?; + tasks_result?; + Ok(()) + } + + async fn stop_background_tasks(&mut self) -> Result<(), Error> { + let check_result = if self.consistent_check_task_started { + self.consistent_check_task_started = false; + self.flow_engine().stop_flow_consistent_check_task().await + } else { + Ok(()) + }; if let Some(task) = &self.heartbeat_task { task.shutdown(); } - Ok(()) + check_result + } + + async fn rollback_background_tasks(&mut self) { + if let Err(err) = self.stop_background_tasks().await { + common_telemetry::error!(err; "Failed to roll back flownode background tasks"); + } } pub fn flownode_server(&self) -> &FlownodeServer { @@ -401,10 +332,7 @@ impl FlownodeBuilder { .context(DatafusionSnafu { context: "Failed to build query engine", })?; - let manager = Arc::new( - self.build_manager(query_engine_factory.query_engine()) - .await?, - ); + let manager = Arc::new(self.build_manager(query_engine_factory.query_engine())); let batching = Arc::new(BatchingEngine::new( self.frontend_client.clone(), query_engine_factory.query_engine(), @@ -432,46 +360,19 @@ impl FlownodeBuilder { flownode_server: server, services: ServerHandlers::default(), heartbeat_task, + state_report_task: None, + consistent_check_task_started: false, }; Ok(instance) } - /// build [`FlowWorkerManager`], note this doesn't take ownership of `self`, - /// nor does it actually start running the worker. - async fn build_manager( - &mut self, - query_engine: Arc, - ) -> Result { - let table_meta = self.table_meta.clone(); - - register_function_to_query_engine(&query_engine); - - let num_workers = self.opts.flow.num_workers; - - let node_id = self.opts.node_id.map(|id| id as u32); - - let mut man = StreamingEngine::new(node_id, query_engine, table_meta); - for worker_id in 0..num_workers { - let (tx, rx) = oneshot::channel(); - - let _handle = std::thread::Builder::new() - .name(format!("flow-worker-{}", worker_id)) - .spawn(move || { - let (handle, mut worker) = create_worker(); - let _ = tx.send(handle); - info!("Flow Worker started in new thread"); - worker.run(); - }); - let worker_handle = rx.await.map_err(|e| { - UnexpectedSnafu { - reason: format!("Failed to receive worker handle: {}", e), - } - .build() - })?; - man.add_worker_handle(worker_handle); - } - info!("Flow Node Manager started"); - Ok(man) + fn build_manager(&self, query_engine: Arc) -> StreamingEngine { + StreamingEngine::new( + self.opts.node_id.map(|id| id as u32), + query_engine, + self.table_meta.clone(), + self.frontend_client.clone(), + ) } } @@ -552,134 +453,6 @@ impl<'a> FlownodeServiceBuilder<'a> { } } -/// Basically a tiny frontend that communicates with datanode, different from [`FrontendClient`] which -/// connect to a real frontend instead, this is used for flow's streaming engine. And is for simple query. -/// -/// For heavy query use [`FrontendClient`] which offload computation to frontend, lifting the load from flownode -#[derive(Clone)] -pub struct FrontendInvoker { - inserter: Arc, - deleter: Arc, - statement_executor: Arc, -} - -impl FrontendInvoker { - pub fn new( - inserter: Arc, - deleter: Arc, - statement_executor: Arc, - ) -> Self { - Self { - inserter, - deleter, - statement_executor, - } - } - - pub async fn build_from( - flow_streaming_engine: FlowStreamingEngineRef, - catalog_manager: CatalogManagerRef, - kv_backend: KvBackendRef, - layered_cache_registry: LayeredCacheRegistryRef, - procedure_executor: ProcedureExecutorRef, - node_manager: NodeManagerRef, - origin_frontend_addr: String, - ) -> Result { - let table_route_cache: TableRouteCacheRef = - layered_cache_registry.get().context(CacheRequiredSnafu { - name: TABLE_ROUTE_CACHE_NAME, - })?; - let partition_info_cache: PartitionInfoCacheRef = - layered_cache_registry.get().context(CacheRequiredSnafu { - name: PARTITION_INFO_CACHE_NAME, - })?; - - let partition_manager = Arc::new(PartitionRuleManager::new( - kv_backend.clone(), - table_route_cache.clone(), - partition_info_cache.clone(), - )); - - let table_flownode_cache: TableFlownodeSetCacheRef = - layered_cache_registry.get().context(CacheRequiredSnafu { - name: TABLE_FLOWNODE_SET_CACHE_NAME, - })?; - - // TODO(auto_create_table): flow sink tables are created through a controlled - // `CREATE FLOW` path, not client writes, so they are intentionally exempt from - // the frontend's global auto-create switch. Revisit if flow should honor it. - let inserter = Arc::new(Inserter::new( - catalog_manager.clone(), - partition_manager.clone(), - node_manager.clone(), - table_flownode_cache, - true, - )); - - let deleter = Arc::new(Deleter::new( - catalog_manager.clone(), - partition_manager.clone(), - node_manager.clone(), - )); - - let query_engine = flow_streaming_engine.query_engine.clone(); - - let statement_executor = Arc::new(StatementExecutor::new( - catalog_manager.clone(), - query_engine.clone(), - procedure_executor.clone(), - kv_backend.clone(), - layered_cache_registry.clone(), - inserter.clone(), - partition_manager, - None, - origin_frontend_addr, - LocalFileAccess::Disabled, - )); - - let invoker = FrontendInvoker::new(inserter, deleter, statement_executor); - Ok(invoker) - } -} - -impl FrontendInvoker { - pub async fn row_inserts( - &self, - requests: RowInsertRequests, - ctx: QueryContextRef, - ) -> common_frontend::error::Result { - let _timer = METRIC_FLOW_PROCESSING_TIME - .with_label_values(&["output_insert"]) - .start_timer(); - - self.inserter - .handle_row_inserts(requests, ctx, &self.statement_executor, false, false) - .await - .map_err(BoxedError::new) - .context(common_frontend::error::ExternalSnafu) - } - - pub async fn row_deletes( - &self, - requests: RowDeleteRequests, - ctx: QueryContextRef, - ) -> common_frontend::error::Result { - let _timer = METRIC_FLOW_PROCESSING_TIME - .with_label_values(&["output_delete"]) - .start_timer(); - - self.deleter - .handle_row_deletes(requests, ctx) - .await - .map_err(BoxedError::new) - .context(common_frontend::error::ExternalSnafu) - } - - pub fn statement_executor(&self) -> Arc { - self.statement_executor.clone() - } -} - /// get all flow ids in this flownode pub(crate) async fn get_all_flow_ids( flow_metadata_manager: &FlowMetadataManagerRef, @@ -722,25 +495,16 @@ pub(crate) async fn get_all_flow_ids( #[cfg(test)] mod tests { use std::sync::Arc; - use std::time::Duration; - use api::v1::HealthCheckRequest; - use api::v1::health_check_client::HealthCheckClient; - use api::v1::meta::Role; use catalog::memory::new_memory_catalog_manager; use common_base::Plugins; use common_meta::key::TableMetadataManager; use common_meta::key::flow::FlowMetadataManager; use common_meta::kv_backend::memory::MemoryKvBackend; - use meta_client::client::MetaClient; use query::options::QueryOptions; - use servers::grpc::GRPC_SERVER; use super::*; - use crate::adapter::flownode_impl::FlowDualEngine; use crate::batching_mode::BatchingModeOptions; - use crate::batching_mode::engine::BatchingEngine; - use crate::utils::SizeReportSender; async fn new_test_flownode_server() -> (FlownodeServer, SizeReportSender) { let (frontend_client, _handler) = @@ -766,13 +530,15 @@ mod tests { let catalog_manager = new_memory_catalog_manager().unwrap(); let query_engine = crate::test_utils::create_test_query_engine(); + let frontend_client = Arc::new(frontend_client); let streaming_engine = Arc::new(StreamingEngine::new( node_id, query_engine.clone(), table_meta.clone(), + frontend_client.clone(), )); let batching_engine = Arc::new(BatchingEngine::new( - Arc::new(frontend_client), + frontend_client, query_engine, flow_meta.clone(), table_meta, @@ -793,74 +559,4 @@ mod tests { let server = FlownodeServer::new(FlowService::new(dual_engine)); (server, report_sender) } - - #[tokio::test] - async fn test_state_report_handler_survives_worker_restart() { - let (server, report_sender) = new_test_flownode_server().await; - - server.start_workers().await.unwrap(); - report_sender.query(Duration::from_secs(3)).await.unwrap(); - - server.stop_workers().await.unwrap(); - report_sender.query(Duration::from_secs(3)).await.unwrap(); - - server.start_workers().await.unwrap(); - report_sender.query(Duration::from_secs(3)).await.unwrap(); - - server.stop_workers().await.unwrap(); - } - - #[tokio::test] - async fn test_start_workers_rolls_back_on_check_task_start_failure() { - let batching_opts = BatchingModeOptions { - experimental_frontend_scan_timeout: Duration::from_millis(1), - ..Default::default() - }; - let frontend_client = FrontendClient::from_meta_client( - Arc::new(MetaClient::new(0, Role::Frontend)), - QueryOptions::default(), - batching_opts.clone(), - ) - .unwrap(); - let (server, _report_sender) = - new_test_flownode_server_with_frontend_client(frontend_client, batching_opts, Some(1)) - .await; - - server.start_workers().await.unwrap_err(); - - assert!(server.inner.streaming_task_handler.lock().await.is_none()); - assert!( - server - .inner - .state_report_task_handler - .lock() - .await - .is_none() - ); - } - - #[tokio::test] - async fn test_service_builder_registers_reachable_health_check() { - // Arrange: compose the production gRPC service with an ephemeral local listener. - let (flownode_server, _report_sender) = new_test_flownode_server().await; - let mut opts = FlownodeOptions::default(); - opts.grpc.bind_addr = "127.0.0.1:0".to_string(); - let mut services = FlownodeServiceBuilder::new(&opts) - .with_default_grpc_server(&flownode_server) - .build() - .unwrap(); - services.start_all().await.unwrap(); - let addr = services.addr(GRPC_SERVER).unwrap(); - - // Act: call the shared health handler through the registered production server. - let mut client = HealthCheckClient::connect(format!("http://{addr}")) - .await - .unwrap(); - let result = client.health_check(HealthCheckRequest {}).await; - - services.shutdown_all().await.unwrap(); - - // Assert: the service composition exposes a healthy endpoint. - assert!(result.is_ok()); - } } diff --git a/src/flow/src/test_utils.rs b/src/flow/src/test_utils.rs index a2b568b679e..22ba9cc8f5a 100644 --- a/src/flow/src/test_utils.rs +++ b/src/flow/src/test_utils.rs @@ -22,58 +22,13 @@ use datatypes::schema::Schema; use datatypes::timestamp::TimestampMillisecond; use datatypes::vectors::{TimestampMillisecondVectorBuilder, VectorRef}; use itertools::Itertools; -use prost::Message; use query::QueryEngine; use query::options::QueryOptions; -use query::parser::QueryLanguageParser; -use query::query_engine::DefaultSerializer; -use session::context::QueryContext; /// note here we are using the `substrait_proto_df` crate from the `substrait` module and /// rename it to `substrait_proto` -use substrait::substrait_proto_df as substrait_proto; -use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; -use substrait_proto::proto; use table::table::numbers::{NUMBERS_TABLE_NAME, NumbersTable}; use table::test_util::MemTable; -use crate::adapter::FlownodeContext; -use crate::adapter::node_context::IdToNameMap; -use crate::adapter::table_source::test::FlowDummyTableSource; -use crate::df_optimizer::apply_df_optimizer; -use crate::expr::GlobalId; -use crate::transform::register_function_to_query_engine; - -pub fn create_test_ctx() -> FlownodeContext { - let mut tri_map = IdToNameMap::new(); - { - let gid = GlobalId::User(0); - let name = [ - "greptime".to_string(), - "public".to_string(), - "numbers".to_string(), - ]; - tri_map.insert(Some(name.clone()), Some(1024), gid); - } - - { - let gid = GlobalId::User(1); - let name = [ - "greptime".to_string(), - "public".to_string(), - "numbers_with_ts".to_string(), - ]; - tri_map.insert(Some(name.clone()), Some(1025), gid); - } - - let dummy_source = FlowDummyTableSource::default(); - - let mut ctx = FlownodeContext::new(Box::new(dummy_source)); - ctx.table_repr = tri_map; - ctx.query_context = Some(Arc::new(QueryContext::with("greptime", "public"))); - - ctx -} - pub fn create_test_query_engine() -> Arc { let catalog_list = catalog::memory::new_memory_catalog_manager().unwrap(); let req = RegisterTableRequest { @@ -158,28 +113,7 @@ pub fn create_test_query_engine() -> Arc { ); let engine = factory.query_engine(); - register_function_to_query_engine(&engine); assert_eq!("datafusion", engine.name()); engine } - -pub async fn sql_to_substrait(engine: Arc, sql: &str) -> proto::Plan { - // let engine = create_test_query_engine(); - let stmt = QueryLanguageParser::parse_sql(sql, &QueryContext::arc()).unwrap(); - let plan = engine - .planner() - .plan(&stmt, QueryContext::arc()) - .await - .unwrap(); - let plan = apply_df_optimizer(plan, &QueryContext::arc()) - .await - .unwrap(); - - // encode then decode so to rely on the impl of conversion from logical plan to substrait plan - let bytes = DFLogicalSubstraitConvertor {} - .encode(&plan, DefaultSerializer) - .unwrap(); - - proto::Plan::decode(bytes).unwrap() -} diff --git a/src/flow/src/transform.rs b/src/flow/src/transform.rs deleted file mode 100644 index 9806fa5eb3c..00000000000 --- a/src/flow/src/transform.rs +++ /dev/null @@ -1,318 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! Transform Substrait into execution plan -use std::collections::BTreeMap; -use std::sync::Arc; - -use common_function::function::FunctionRef; -use datafusion::arrow::datatypes::{DataType, TimeUnit}; -use datafusion::logical_expr::ColumnarValue; -use datafusion_expr::{ScalarFunctionArgs, Signature, Volatility}; -use datafusion_substrait::extensions::Extensions; -use query::QueryEngine; -use serde::{Deserialize, Serialize}; -/// note here we are using the `substrait_proto_df` crate from the `substrait` module and -/// rename it to `substrait_proto` -use substrait::substrait_proto_df as substrait_proto; -use substrait_proto::proto::extensions::SimpleExtensionDeclaration; -use substrait_proto::proto::extensions::simple_extension_declaration::MappingType; - -use crate::adapter::FlownodeContext; -use crate::error::{Error, NotImplementedSnafu}; -use crate::expr::{TUMBLE_END, TUMBLE_START}; -/// a simple macro to generate a not implemented error -macro_rules! not_impl_err { - ($($arg:tt)*) => { - NotImplementedSnafu { - reason: format!($($arg)*), - }.fail() - }; -} - -/// generate a plan error -macro_rules! plan_err { - ($($arg:tt)*) => { - PlanSnafu { - reason: format!($($arg)*), - }.fail() - }; -} - -mod aggr; -mod expr; -mod literal; -mod plan; - -pub(crate) use expr::from_scalar_fn_to_df_fn_impl; - -/// In Substrait, a function can be define by an u32 anchor, and the anchor can be mapped to a name -/// -/// So in substrait plan, a ref to a function can be a single u32 anchor instead of a full name in string -#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct FunctionExtensions { - anchor_to_name: BTreeMap, -} - -impl FunctionExtensions { - pub fn from_iter(inner: impl IntoIterator) -> Self { - Self { - anchor_to_name: inner.into_iter().map(|(k, s)| (k, s.to_string())).collect(), - } - } - - /// Create a new FunctionExtensions from a list of SimpleExtensionDeclaration - pub fn try_from_proto(extensions: &[SimpleExtensionDeclaration]) -> Result { - let mut anchor_to_name = BTreeMap::new(); - for e in extensions { - match &e.mapping_type { - Some(ext) => match ext { - MappingType::ExtensionFunction(ext_f) => { - anchor_to_name.insert(ext_f.function_anchor, ext_f.name.clone()); - } - _ => not_impl_err!("Extension type not supported: {ext:?}")?, - }, - None => not_impl_err!("Cannot parse empty extension")?, - } - } - Ok(Self { anchor_to_name }) - } - - /// Get the name of a function by it's anchor - pub fn get(&self, anchor: &u32) -> Option<&String> { - self.anchor_to_name.get(anchor) - } - - pub fn to_extensions(&self) -> Extensions { - Extensions { - functions: self - .anchor_to_name - .iter() - .map(|(k, v)| (*k, v.clone())) - .collect(), - ..Default::default() - } - } -} - -/// register flow-specific functions to the query engine -pub fn register_function_to_query_engine(engine: &Arc) { - let tumble_fn = Arc::new(TumbleFunction::new("tumble")) as FunctionRef; - let tumble_start_fn = Arc::new(TumbleFunction::new(TUMBLE_START)) as FunctionRef; - let tumble_end_fn = Arc::new(TumbleFunction::new(TUMBLE_END)) as FunctionRef; - - engine.register_scalar_function(tumble_fn.into()); - engine.register_scalar_function(tumble_start_fn.into()); - engine.register_scalar_function(tumble_end_fn.into()); -} - -#[derive(Debug)] -pub struct TumbleFunction { - name: String, - signature: Signature, -} - -impl TumbleFunction { - fn new(name: &str) -> Self { - Self { - name: name.to_string(), - signature: Signature::variadic_any(Volatility::Immutable), - } - } -} - -impl std::fmt::Display for TumbleFunction { - fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - write!(f, "{}", self.name.to_ascii_uppercase()) - } -} - -impl common_function::function::Function for TumbleFunction { - fn name(&self) -> &str { - &self.name - } - - fn return_type(&self, _: &[DataType]) -> datafusion_common::Result { - Ok(DataType::Timestamp(TimeUnit::Millisecond, None)) - } - - fn signature(&self) -> &Signature { - &self.signature - } - - fn invoke_with_args(&self, _: ScalarFunctionArgs) -> datafusion_common::Result { - datafusion_common::not_impl_err!("{}", self.name()) - } -} - -#[cfg(test)] -mod test { - use std::sync::Arc; - - use catalog::RegisterTableRequest; - use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, NUMBERS_TABLE_ID}; - use datatypes::data_type::ConcreteDataType as CDT; - use datatypes::prelude::*; - use datatypes::schema::Schema; - use datatypes::timestamp::TimestampMillisecond; - use datatypes::vectors::{TimestampMillisecondVectorBuilder, VectorRef}; - use itertools::Itertools; - use prost::Message; - use query::QueryEngine; - use query::options::QueryOptions; - use query::parser::QueryLanguageParser; - use query::query_engine::DefaultSerializer; - use session::context::QueryContext; - use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; - use substrait_proto::proto; - use table::table::numbers::{NUMBERS_TABLE_NAME, NumbersTable}; - use table::test_util::MemTable; - - use super::*; - use crate::adapter::node_context::IdToNameMap; - use crate::adapter::table_source::test::FlowDummyTableSource; - use crate::df_optimizer::apply_df_optimizer; - use crate::expr::GlobalId; - - pub fn create_test_ctx() -> FlownodeContext { - let mut tri_map = IdToNameMap::new(); - // FIXME(discord9): deprecated, use `numbers_with_ts` instead since this table has no timestamp column - { - let gid = GlobalId::User(0); - let name = [ - "greptime".to_string(), - "public".to_string(), - "numbers".to_string(), - ]; - tri_map.insert(Some(name.clone()), Some(1024), gid); - } - - { - let gid = GlobalId::User(1); - let name = [ - "greptime".to_string(), - "public".to_string(), - "numbers_with_ts".to_string(), - ]; - tri_map.insert(Some(name.clone()), Some(1025), gid); - } - - let dummy_source = FlowDummyTableSource::default(); - - let mut ctx = FlownodeContext::new(Box::new(dummy_source)); - ctx.table_repr = tri_map; - ctx.query_context = Some(Arc::new(QueryContext::with("greptime", "public"))); - - ctx - } - - pub fn create_test_query_engine() -> Arc { - let catalog_list = catalog::memory::new_memory_catalog_manager().unwrap(); - let req = RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: NUMBERS_TABLE_NAME.to_string(), - table_id: NUMBERS_TABLE_ID, - table: NumbersTable::table(NUMBERS_TABLE_ID), - }; - catalog_list.register_table_sync(req).unwrap(); - - let schema = vec![ - datatypes::schema::ColumnSchema::new("number", CDT::uint32_datatype(), false), - datatypes::schema::ColumnSchema::new( - "ts", - CDT::timestamp_millisecond_datatype(), - false, - ), - ]; - let mut columns = vec![]; - let numbers = (1..=10).collect_vec(); - let column: VectorRef = Arc::new(::VectorType::from_vec(numbers)); - columns.push(column); - - let ts = (1..=10).collect_vec(); - let mut builder = TimestampMillisecondVectorBuilder::with_capacity(10); - ts.into_iter() - .map(|v| builder.push(Some(TimestampMillisecond::new(v)))) - .count(); - let column: VectorRef = builder.to_vector_cloned(); - columns.push(column); - - let schema = Arc::new(Schema::new(schema)); - let recordbatch = common_recordbatch::RecordBatch::new(schema, columns).unwrap(); - let table = MemTable::table("numbers_with_ts", recordbatch); - - let req_with_ts = RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: "numbers_with_ts".to_string(), - table_id: 1024, - table, - }; - catalog_list.register_table_sync(req_with_ts).unwrap(); - - let factory = query::QueryEngineFactory::new( - catalog_list, - None, - None, - None, - None, - false, - QueryOptions::default(), - ); - - let engine = factory.query_engine(); - register_function_to_query_engine(&engine); - - assert_eq!("datafusion", engine.name()); - engine - } - - pub async fn sql_to_substrait(engine: Arc, sql: &str) -> proto::Plan { - // let engine = create_test_query_engine(); - let stmt = QueryLanguageParser::parse_sql(sql, &QueryContext::arc()).unwrap(); - let plan = engine - .planner() - .plan(&stmt, QueryContext::arc()) - .await - .unwrap(); - let plan = apply_df_optimizer(plan, &QueryContext::arc()) - .await - .unwrap(); - - // encode then decode so to rely on the impl of conversion from logical plan to substrait plan - let bytes = DFLogicalSubstraitConvertor {} - .encode(&plan, DefaultSerializer) - .unwrap(); - - proto::Plan::decode(bytes).unwrap() - } - - /// TODO(discord9): add more illegal sql tests - #[tokio::test] - async fn test_missing_key_check() { - let engine = create_test_query_engine(); - let sql = "SELECT avg(number) FROM numbers_with_ts GROUP BY tumble(ts, '1 hour'), number"; - - let stmt = QueryLanguageParser::parse_sql(sql, &QueryContext::arc()).unwrap(); - let plan = engine - .planner() - .plan(&stmt, QueryContext::arc()) - .await - .unwrap(); - let plan = apply_df_optimizer(plan, &QueryContext::arc()).await; - - assert!(plan.is_err()); - } -} diff --git a/src/flow/src/transform/aggr.rs b/src/flow/src/transform/aggr.rs deleted file mode 100644 index 7b938524a4e..00000000000 --- a/src/flow/src/transform/aggr.rs +++ /dev/null @@ -1,803 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use itertools::Itertools; -use snafu::OptionExt; -use substrait_proto::proto; -use substrait_proto::proto::aggregate_function::AggregationInvocation; -use substrait_proto::proto::aggregate_rel::{Grouping, Measure}; -use substrait_proto::proto::function_argument::ArgType; - -use crate::error::{Error, NotImplementedSnafu, PlanSnafu}; -use crate::expr::{ - AggregateExpr, AggregateFunc, MapFilterProject, ScalarExpr, TypedExpr, UnaryFunc, -}; -use crate::plan::{AccumulablePlan, AggrWithIndex, KeyValPlan, Plan, ReducePlan, TypedPlan}; -use crate::repr::{ColumnType, RelationDesc, RelationType}; -use crate::transform::{FlownodeContext, FunctionExtensions, substrait_proto}; - -impl TypedExpr { - /// Allow `deprecated` due to the usage of deprecated grouping_expressions on datafusion to substrait side - #[allow(deprecated)] - async fn from_substrait_agg_grouping( - ctx: &mut FlownodeContext, - grouping_expressions: &[proto::Expression], - groupings: &[Grouping], - typ: &RelationDesc, - extensions: &FunctionExtensions, - ) -> Result, Error> { - let _ = ctx; - let mut group_expr = vec![]; - match groupings.len() { - 1 => { - // handle case when deprecated grouping_expressions is referenced by index is empty - let expressions: Box + Send> = if groupings - [0] - .expression_references - .is_empty() - { - Box::new(groupings[0].grouping_expressions.iter()) - } else { - if groupings[0] - .expression_references - .iter() - .any(|idx| *idx as usize >= grouping_expressions.len()) - { - return PlanSnafu { - reason: format!("Invalid grouping expression reference: {:?} for grouping expr: {:?}", - groupings[0].expression_references, - grouping_expressions - ), - }.fail()?; - } - Box::new( - groupings[0] - .expression_references - .iter() - .map(|idx| &grouping_expressions[*idx as usize]), - ) - }; - for e in expressions { - let x = TypedExpr::from_substrait_rex(e, typ, extensions).await?; - group_expr.push(x); - } - } - _ => { - return not_impl_err!( - "Grouping sets not support yet, use union all with group by instead." - ); - } - }; - Ok(group_expr) - } -} - -impl AggregateExpr { - /// Convert list of `Measure` into Flow's AggregateExpr - /// - /// Return both the AggregateExpr and a MapFilterProject that is the final output of the aggregate function - async fn from_substrait_agg_measures( - ctx: &mut FlownodeContext, - measures: &[Measure], - typ: &RelationDesc, - extensions: &FunctionExtensions, - ) -> Result, Error> { - let _ = ctx; - let mut all_aggr_exprs = vec![]; - - for m in measures { - let filter = match m - .filter - .as_ref() - .map(|fil| TypedExpr::from_substrait_rex(fil, typ, extensions)) - { - Some(fut) => Some(fut.await), - None => None, - } - .transpose()?; - - let aggr_expr = match &m.measure { - Some(f) => { - let distinct = match f.invocation { - _ if f.invocation == AggregationInvocation::Distinct as i32 => true, - _ if f.invocation == AggregationInvocation::All as i32 => false, - _ => false, - }; - AggregateExpr::from_substrait_agg_func( - f, typ, extensions, &filter, // TODO(discord9): impl order_by - &None, distinct, - ) - .await? - } - None => { - return not_impl_err!("Aggregate without aggregate function is not supported"); - } - }; - - all_aggr_exprs.extend(aggr_expr); - } - - Ok(all_aggr_exprs) - } - - /// Convert AggregateFunction into Flow's AggregateExpr - /// - /// the returned value is a tuple of AggregateExpr and a optional ScalarExpr that if exist is the final output of the aggregate function - /// since aggr functions like `avg` need to be transform to `sum(x)/cast(count(x) as x_type)` - pub async fn from_substrait_agg_func( - f: &proto::AggregateFunction, - input_schema: &RelationDesc, - extensions: &FunctionExtensions, - filter: &Option, - order_by: &Option>, - distinct: bool, - ) -> Result, Error> { - // TODO(discord9): impl filter - let _ = filter; - let _ = order_by; - let mut args = vec![]; - for arg in &f.arguments { - let arg_expr = match &arg.arg_type { - Some(ArgType::Value(e)) => { - TypedExpr::from_substrait_rex(e, input_schema, extensions).await - } - _ => not_impl_err!("Aggregated function argument non-Value type not supported"), - }?; - args.push(arg_expr); - } - - if args.len() != 1 { - let fn_name = extensions.get(&f.function_reference).cloned(); - return not_impl_err!( - "Aggregated function (name={:?}) with multiple arguments is not supported", - fn_name - ); - } - - let arg = if let Some(first) = args.first() { - first - } else { - return not_impl_err!("Aggregated function without arguments is not supported"); - }; - - let fn_name = extensions - .get(&f.function_reference) - .cloned() - .map(|s| s.to_lowercase()); - - match fn_name.as_ref().map(|s| s.as_ref()) { - Some(function_name) => { - let func = AggregateFunc::from_str_and_type( - function_name, - Some(arg.typ.scalar_type.clone()), - )?; - let exprs = vec![AggregateExpr { - func, - expr: arg.expr.clone(), - distinct, - }]; - Ok(exprs) - } - None => not_impl_err!( - "Aggregated function not found: function anchor = {:?}", - f.function_reference - ), - } - } -} - -impl KeyValPlan { - /// Generate KeyValPlan from AggregateExpr and group_exprs - /// - /// will also change aggregate expr to use column ref if necessary - fn from_substrait_gen_key_val_plan( - aggr_exprs: &mut [AggregateExpr], - group_exprs: &[TypedExpr], - input_arity: usize, - ) -> Result { - let group_expr_val = group_exprs - .iter() - .map(|expr| expr.expr.clone()) - .collect_vec(); - let output_arity = group_expr_val.len(); - let key_plan = MapFilterProject::new(input_arity) - .map(group_expr_val)? - .project(input_arity..input_arity + output_arity)?; - - // val_plan is extracted from aggr_exprs to give aggr function it's necessary input - // and since aggr func need inputs that is column ref, we just add a prefix mfp to transform any expr that is not into a column ref - let val_plan = { - let need_mfp = aggr_exprs.iter().any(|agg| agg.expr.as_column().is_none()); - if need_mfp { - // create mfp from aggr_expr, and modify aggr_expr to use the output column of mfp - let input_exprs = aggr_exprs - .iter_mut() - .enumerate() - .map(|(idx, aggr)| { - let ret = aggr.expr.clone(); - aggr.expr = ScalarExpr::Column(idx); - ret - }) - .collect_vec(); - let aggr_arity = aggr_exprs.len(); - - MapFilterProject::new(input_arity) - .map(input_exprs)? - .project(input_arity..input_arity + aggr_arity)? - } else { - // simply take all inputs as value - MapFilterProject::new(input_arity) - } - }; - Ok(KeyValPlan { - key_plan: key_plan.into_safe(), - val_plan: val_plan.into_safe(), - }) - } -} - -/// find out the column that should be time index in group exprs(which is all columns that should be keys) -/// TODO(discord9): better ways to assign time index -/// for now, it will found the first column that is timestamp or has a tumble window floor function -fn find_time_index_in_group_exprs(group_exprs: &[TypedExpr]) -> Option { - group_exprs.iter().position(|expr| { - matches!( - &expr.expr, - ScalarExpr::CallUnary { - func: UnaryFunc::TumbleWindowFloor { .. }, - expr: _ - } - ) || expr.typ.scalar_type.is_timestamp() - }) -} - -impl TypedPlan { - /// Convert AggregateRel into Flow's TypedPlan - /// - /// The output of aggr plan is: - /// - /// .. - #[async_recursion::async_recursion] - pub async fn from_substrait_agg_rel( - ctx: &mut FlownodeContext, - agg: &proto::AggregateRel, - extensions: &FunctionExtensions, - ) -> Result { - let input = if let Some(input) = agg.input.as_ref() { - TypedPlan::from_substrait_rel(ctx, input, extensions).await? - } else { - return not_impl_err!("Aggregate without an input is not supported"); - }; - - let group_exprs = TypedExpr::from_substrait_agg_grouping( - ctx, - &agg.grouping_expressions, - &agg.groupings, - &input.schema, - extensions, - ) - .await?; - - let time_index = find_time_index_in_group_exprs(&group_exprs); - - let mut aggr_exprs = AggregateExpr::from_substrait_agg_measures( - ctx, - &agg.measures, - &input.schema, - extensions, - ) - .await?; - - let key_val_plan = KeyValPlan::from_substrait_gen_key_val_plan( - &mut aggr_exprs, - &group_exprs, - input.schema.typ.column_types.len(), - )?; - - // output type is group_exprs + aggr_exprs - let output_type = { - let mut output_types = Vec::new(); - // give best effort to get column name - let mut output_names = Vec::new(); - - // first append group_expr as key, then aggr_expr as value - for expr in group_exprs.iter() { - output_types.push(expr.typ.clone()); - let col_name = match &expr.expr { - ScalarExpr::Column(col) => input.schema.get_name(*col).clone(), - // TODO(discord9): impl& use ScalarExpr.display_name, which recursively build expr's name - _ => None, - }; - output_names.push(col_name) - } - - for aggr in &aggr_exprs { - output_types.push(ColumnType::new_nullable( - aggr.func.signature().output.clone(), - )); - // TODO(discord9): find a clever way to name them? - output_names.push(None); - } - // TODO(discord9): try best to get time - if group_exprs.is_empty() { - RelationType::new(output_types) - } else { - RelationType::new(output_types).with_key((0..group_exprs.len()).collect_vec()) - } - .with_time_index(time_index) - .into_named(output_names) - }; - - // copy aggr_exprs to full_aggrs, and split them into simple_aggrs and distinct_aggrs - // also set them input/output column - let full_aggrs = aggr_exprs; - let mut simple_aggrs = Vec::new(); - let mut distinct_aggrs = Vec::new(); - for (output_column, aggr_expr) in full_aggrs.iter().enumerate() { - let input_column = aggr_expr.expr.as_column().with_context(|| PlanSnafu { - reason: "Expect aggregate argument to be transformed into a column at this point", - })?; - if aggr_expr.distinct { - distinct_aggrs.push(AggrWithIndex::new( - aggr_expr.clone(), - input_column, - output_column, - )); - } else { - simple_aggrs.push(AggrWithIndex::new( - aggr_expr.clone(), - input_column, - output_column, - )); - } - } - let accum_plan = AccumulablePlan { - full_aggrs, - simple_aggrs, - distinct_aggrs, - }; - let plan = Plan::Reduce { - input: Box::new(input), - key_val_plan, - reduce_plan: ReducePlan::Accumulable(accum_plan), - }; - // FIX(discord9): deal with key first - return Ok(TypedPlan { - schema: output_type, - plan, - }); - } -} - -#[cfg(test)] -mod test { - - use bytes::BytesMut; - use common_time::IntervalMonthDayNano; - use datatypes::data_type::ConcreteDataType as CDT; - use datatypes::prelude::ConcreteDataType; - use datatypes::value::Value; - use pretty_assertions::assert_eq; - - use super::*; - use crate::expr::{BinaryFunc, DfScalarFunction, GlobalId, RawDfScalarFn}; - use crate::plan::{Plan, TypedPlan}; - use crate::repr::{ColumnType, RelationType}; - use crate::transform::test::{create_test_ctx, create_test_query_engine, sql_to_substrait}; - - #[tokio::test] - async fn test_sum() { - let engine = create_test_query_engine(); - let sql = "SELECT sum(number) FROM numbers"; - let plan = sql_to_substrait(engine.clone(), sql).await; - - let mut ctx = create_test_ctx(); - let flow_plan = TypedPlan::from_substrait_plan(&mut ctx, &plan).await; - - let aggr_expr = AggregateExpr { - func: AggregateFunc::SumUInt64, - expr: ScalarExpr::Column(0), - distinct: false, - }; - let expected = TypedPlan { - schema: RelationType::new(vec![ColumnType::new(CDT::uint64_datatype(), true)]) - .into_named(vec![Some("sum(numbers.number)".to_string())]), - plan: Plan::Reduce { - input: Box::new( - Plan::Get { - id: crate::expr::Id::Global(GlobalId::User(0)), - } - .with_types( - RelationType::new(vec![ColumnType::new( - ConcreteDataType::uint32_datatype(), - false, - )]) - .into_named(vec![Some("number".to_string())]), - ) - .mfp(MapFilterProject::new(1).into_safe()) - .unwrap(), - ), - key_val_plan: KeyValPlan { - key_plan: MapFilterProject::new(1) - .project(vec![]) - .unwrap() - .into_safe(), - val_plan: MapFilterProject::new(1) - .map(vec![ - ScalarExpr::Column(0) - .call_unary(UnaryFunc::Cast(CDT::uint64_datatype())), - ]) - .unwrap() - .project(vec![1]) - .unwrap() - .into_safe(), - }, - reduce_plan: ReducePlan::Accumulable(AccumulablePlan { - full_aggrs: vec![aggr_expr.clone()], - simple_aggrs: vec![AggrWithIndex::new(aggr_expr.clone(), 0, 0)], - distinct_aggrs: vec![], - }), - }, - }; - assert_eq!(flow_plan.unwrap(), expected); - } - - #[tokio::test] - async fn test_distinct_number() { - let engine = create_test_query_engine(); - let sql = "SELECT DISTINCT number FROM numbers"; - let plan = sql_to_substrait(engine.clone(), sql).await; - - let mut ctx = create_test_ctx(); - let flow_plan = TypedPlan::from_substrait_plan(&mut ctx, &plan) - .await - .unwrap(); - - let expected = TypedPlan { - schema: RelationType::new(vec![ - ColumnType::new(CDT::uint32_datatype(), false), // col number - ]) - .with_key(vec![0]) - .into_named(vec![Some("number".to_string())]), - plan: Plan::Reduce { - input: Box::new( - Plan::Get { - id: crate::expr::Id::Global(GlobalId::User(0)), - } - .with_types( - RelationType::new(vec![ColumnType::new( - ConcreteDataType::uint32_datatype(), - false, - )]) - .into_named(vec![Some("number".to_string())]), - ) - .mfp(MapFilterProject::new(1).into_safe()) - .unwrap(), - ), - key_val_plan: KeyValPlan { - key_plan: MapFilterProject::new(1) - .map(vec![ScalarExpr::Column(0)]) - .unwrap() - .project(vec![1]) - .unwrap() - .into_safe(), - val_plan: MapFilterProject::new(1) - .project(vec![0]) - .unwrap() - .into_safe(), - }, - reduce_plan: ReducePlan::Accumulable(AccumulablePlan { - full_aggrs: vec![], - simple_aggrs: vec![], - distinct_aggrs: vec![], - }), - }, - }; - - assert_eq!(flow_plan, expected); - } - - #[tokio::test] - async fn test_sum_group_by() { - let engine = create_test_query_engine(); - let sql = "SELECT sum(number), number FROM numbers GROUP BY number"; - let plan = sql_to_substrait(engine.clone(), sql).await; - - let mut ctx = create_test_ctx(); - let flow_plan = TypedPlan::from_substrait_plan(&mut ctx, &plan) - .await - .unwrap(); - - let aggr_expr = AggregateExpr { - func: AggregateFunc::SumUInt64, - expr: ScalarExpr::Column(0), - distinct: false, - }; - let expected = TypedPlan { - schema: RelationType::new(vec![ - ColumnType::new(CDT::uint64_datatype(), true), // col sum(number) - ColumnType::new(CDT::uint32_datatype(), false), // col number - ]) - .with_key(vec![1]) - .into_named(vec![ - Some("sum(numbers.number)".to_string()), - Some("number".to_string()), - ]), - plan: Plan::Mfp { - input: Box::new( - Plan::Reduce { - input: Box::new( - Plan::Get { - id: crate::expr::Id::Global(GlobalId::User(0)), - } - .with_types( - RelationType::new(vec![ColumnType::new( - ConcreteDataType::uint32_datatype(), - false, - )]) - .into_named(vec![Some("number".to_string())]), - ) - .mfp(MapFilterProject::new(1).into_safe()) - .unwrap(), - ), - key_val_plan: KeyValPlan { - key_plan: MapFilterProject::new(1) - .map(vec![ScalarExpr::Column(0)]) - .unwrap() - .project(vec![1]) - .unwrap() - .into_safe(), - val_plan: MapFilterProject::new(1) - .map(vec![ - ScalarExpr::Column(0) - .call_unary(UnaryFunc::Cast(CDT::uint64_datatype())), - ]) - .unwrap() - .project(vec![1]) - .unwrap() - .into_safe(), - }, - reduce_plan: ReducePlan::Accumulable(AccumulablePlan { - full_aggrs: vec![aggr_expr.clone()], - simple_aggrs: vec![AggrWithIndex::new(aggr_expr.clone(), 0, 0)], - distinct_aggrs: vec![], - }), - } - .with_types( - RelationType::new(vec![ - ColumnType::new(CDT::uint32_datatype(), false), // col number - ColumnType::new(CDT::uint64_datatype(), true), // col sum(number) - ]) - .with_key(vec![0]) - .into_named(vec![Some("number".to_string()), None]), - ), - ), - mfp: MapFilterProject::new(2) - .map(vec![ScalarExpr::Column(1), ScalarExpr::Column(0)]) - .unwrap() - .project(vec![2, 3]) - .unwrap(), - }, - }; - - assert_eq!(flow_plan, expected); - } - - #[tokio::test] - async fn test_sum_add() { - let engine = create_test_query_engine(); - let sql = "SELECT sum(number+number) FROM numbers"; - let plan = sql_to_substrait(engine.clone(), sql).await; - - let mut ctx = create_test_ctx(); - let flow_plan = TypedPlan::from_substrait_plan(&mut ctx, &plan).await; - - let aggr_expr = AggregateExpr { - func: AggregateFunc::SumUInt64, - expr: ScalarExpr::Column(0), - distinct: false, - }; - let expected = TypedPlan { - schema: RelationType::new(vec![ColumnType::new(CDT::uint64_datatype(), true)]) - .into_named(vec![Some( - "sum(numbers.number + numbers.number)".to_string(), - )]), - plan: Plan::Reduce { - input: Box::new( - Plan::Mfp { - input: Box::new( - Plan::Get { - id: crate::expr::Id::Global(GlobalId::User(0)), - } - .with_types( - RelationType::new(vec![ColumnType::new( - ConcreteDataType::uint32_datatype(), - false, - )]) - .into_named(vec![Some("number".to_string())]), - ), - ), - mfp: MapFilterProject::new(1), - } - .with_types( - RelationType::new(vec![ColumnType::new( - ConcreteDataType::uint32_datatype(), - false, - )]) - .into_named(vec![Some("number".to_string())]), - ), - ), - key_val_plan: KeyValPlan { - key_plan: MapFilterProject::new(1) - .project(vec![]) - .unwrap() - .into_safe(), - val_plan: MapFilterProject::new(1) - .map(vec![ - ScalarExpr::Column(0) - .call_binary(ScalarExpr::Column(0), BinaryFunc::AddUInt32) - .call_unary(UnaryFunc::Cast(CDT::uint64_datatype())), - ]) - .unwrap() - .project(vec![1]) - .unwrap() - .into_safe(), - }, - reduce_plan: ReducePlan::Accumulable(AccumulablePlan { - full_aggrs: vec![aggr_expr.clone()], - simple_aggrs: vec![AggrWithIndex::new(aggr_expr.clone(), 0, 0)], - distinct_aggrs: vec![], - }), - }, - }; - assert_eq!(flow_plan.unwrap(), expected); - } - - #[tokio::test] - async fn test_cast_max_min() { - let engine = create_test_query_engine(); - let sql = "SELECT (max(number) - min(number))/30.0, date_bin(INTERVAL '30 second', CAST(ts AS TimestampMillisecond)) as time_window from numbers_with_ts GROUP BY time_window"; - let plan = sql_to_substrait(engine.clone(), sql).await; - - let mut ctx = create_test_ctx(); - let flow_plan = TypedPlan::from_substrait_plan(&mut ctx, &plan).await; - - let aggr_exprs = vec![ - AggregateExpr { - func: AggregateFunc::MaxUInt32, - expr: ScalarExpr::Column(0), - distinct: false, - }, - AggregateExpr { - func: AggregateFunc::MinUInt32, - expr: ScalarExpr::Column(0), - distinct: false, - }, - ]; - let expected = TypedPlan { - schema: RelationType::new(vec![ - ColumnType::new(CDT::float64_datatype(), true), - ColumnType::new(CDT::timestamp_millisecond_datatype(), true), - ]) - .with_time_index(Some(1)) - .into_named(vec![ - Some( - "max(numbers_with_ts.number) - min(numbers_with_ts.number) / Float64(30)" - .to_string(), - ), - Some("time_window".to_string()), - ]), - plan: Plan::Mfp { - input: Box::new( - Plan::Reduce { - input: Box::new( - Plan::Get { - id: crate::expr::Id::Global(GlobalId::User(1)), - } - .with_types( - RelationType::new(vec![ - ColumnType::new(ConcreteDataType::uint32_datatype(), false), - ColumnType::new(ConcreteDataType::timestamp_millisecond_datatype(), false), - ]) - .into_named(vec![ - Some("number".to_string()), - Some("ts".to_string()), - ]), - ) - .mfp(MapFilterProject::new(2).into_safe()) - .unwrap(), - ), - - key_val_plan: KeyValPlan { - key_plan: MapFilterProject::new(2) - .map(vec![ScalarExpr::CallDf { - df_scalar_fn: DfScalarFunction::try_from_raw_fn( - RawDfScalarFn { - f: BytesMut::from( - b"\x08\x02\x1a\x07\x8a\x02\x04\x08\x03\x18\x01\"\x0f\x1a\r\n\x0b\xa2\x02\x08\n\0\x12\x04\x10\x1e \t\"\n\x1a\x08\x12\x06\n\x04\x12\x02\x08\x01".as_ref(), - ), - input_schema: RelationType::new(vec![ColumnType::new( - ConcreteDataType::interval_month_day_nano_datatype(), - true, - ),ColumnType::new( - ConcreteDataType::timestamp_millisecond_datatype(), - false, - )]) - .into_unnamed(), - extensions: FunctionExtensions::from_iter([ - (0, "subtract".to_string()), - (1, "divide".to_string()), - (2, "date_bin".to_string()), - (3, "max".to_string()), - (4, "min".to_string()), - ]), - }, - ) - .await - .unwrap(), - exprs: vec![ - ScalarExpr::Literal( - Value::IntervalMonthDayNano(IntervalMonthDayNano::new(0, 0, 30000000000)), - CDT::interval_month_day_nano_datatype() - ), - ScalarExpr::Column(1) - ], - }]) - .unwrap() - .project(vec![2]) - .unwrap() - .into_safe(), - val_plan: MapFilterProject::new(2) - .into_safe(), - }, - reduce_plan: ReducePlan::Accumulable(AccumulablePlan { - full_aggrs: aggr_exprs.clone(), - simple_aggrs: vec![AggrWithIndex::new(aggr_exprs[0].clone(), 0, 0), - AggrWithIndex::new(aggr_exprs[1].clone(), 0, 1)], - distinct_aggrs: vec![], - }), - } - .with_types( - RelationType::new(vec![ - ColumnType::new( - ConcreteDataType::timestamp_millisecond_datatype(), - true, - ), // time_window - ColumnType::new(ConcreteDataType::uint32_datatype(), true), // max - ColumnType::new(ConcreteDataType::uint32_datatype(), true), // min - ]) - .with_time_index(Some(0)) - .into_unnamed(), - ), - ), - mfp: MapFilterProject::new(3) - .map(vec![ - ScalarExpr::Column(1) - .call_binary(ScalarExpr::Column(2), BinaryFunc::SubUInt32) - .cast(CDT::float64_datatype()) - .call_binary( - ScalarExpr::Literal(Value::from(30.0f64), CDT::float64_datatype()), - BinaryFunc::DivFloat64, - ), - ScalarExpr::Column(0), - ]) - .unwrap() - .project(vec![3, 4]) - .unwrap(), - }, - }; - - assert_eq!(flow_plan.unwrap(), expected); - } -} diff --git a/src/flow/src/transform/expr.rs b/src/flow/src/transform/expr.rs deleted file mode 100644 index a2150e2290b..00000000000 --- a/src/flow/src/transform/expr.rs +++ /dev/null @@ -1,839 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -#![warn(unused_imports)] - -use std::sync::Arc; - -use common_error::ext::BoxedError; -use common_telemetry::debug; -use datafusion::execution::SessionStateBuilder; -use datafusion::functions::all_default_functions; -use datafusion_expr::physical_planning_context::PhysicalPlanningContext; -use datafusion_physical_expr::PhysicalExpr; -use datafusion_substrait::logical_plan::consumer::DefaultSubstraitConsumer; -use datatypes::data_type::ConcreteDataType as CDT; -use snafu::{OptionExt, ResultExt, ensure}; -use substrait_proto::proto::Expression; -use substrait_proto::proto::expression::field_reference::ReferenceType::DirectReference; -use substrait_proto::proto::expression::reference_segment::ReferenceType::StructField; -use substrait_proto::proto::expression::{IfThen, RexType, ScalarFunction}; -use substrait_proto::proto::function_argument::ArgType; - -use crate::error::{ - DatafusionSnafu, DatatypesSnafu, Error, EvalSnafu, ExternalSnafu, InvalidQuerySnafu, - NotImplementedSnafu, PlanSnafu, UnexpectedSnafu, -}; -use crate::expr::{ - BinaryFunc, DfScalarFunction, RawDfScalarFn, ScalarExpr, TUMBLE_END, TUMBLE_START, TypedExpr, - UnaryFunc, UnmaterializableFunc, VariadicFunc, -}; -use crate::repr::{ColumnType, RelationDesc, RelationType}; -use crate::transform::literal::{ - from_substrait_literal, from_substrait_type, to_substrait_literal, -}; -use crate::transform::{FunctionExtensions, substrait_proto}; - -// TODO(discord9): refactor plan to substrait convert of `arrow_cast` function thus remove this function -/// ref to `arrow_schema::datatype` for type name -fn typename_to_cdt(name: &str) -> Result { - let ret = match name { - "Int8" => CDT::int8_datatype(), - "Int16" => CDT::int16_datatype(), - "Int32" => CDT::int32_datatype(), - "Int64" => CDT::int64_datatype(), - "UInt8" => CDT::uint8_datatype(), - "UInt16" => CDT::uint16_datatype(), - "UInt32" => CDT::uint32_datatype(), - "UInt64" => CDT::uint64_datatype(), - "Float32" => CDT::float32_datatype(), - "Float64" => CDT::float64_datatype(), - "Boolean" => CDT::boolean_datatype(), - "String" => CDT::string_datatype(), - "Date" | "Date32" | "Date64" => CDT::date_datatype(), - "Timestamp" => CDT::timestamp_second_datatype(), - "Timestamp(Second, None)" => CDT::timestamp_second_datatype(), - "Timestamp(Millisecond, None)" => CDT::timestamp_millisecond_datatype(), - "Timestamp(Microsecond, None)" => CDT::timestamp_microsecond_datatype(), - "Timestamp(Nanosecond, None)" => CDT::timestamp_nanosecond_datatype(), - "Time32(Second)" | "Time64(Second)" => CDT::time_second_datatype(), - "Time32(Millisecond)" | "Time64(Millisecond)" => CDT::time_millisecond_datatype(), - "Time32(Microsecond)" | "Time64(Microsecond)" => CDT::time_microsecond_datatype(), - "Time32(Nanosecond)" | "Time64(Nanosecond)" => CDT::time_nanosecond_datatype(), - _ => NotImplementedSnafu { - reason: format!("Unrecognized typename: {}", name), - } - .fail()?, - }; - Ok(ret) -} - -/// Convert [`ScalarFunction`] to corresponding Datafusion's [`PhysicalExpr`] -pub(crate) async fn from_scalar_fn_to_df_fn_impl( - f: &ScalarFunction, - input_schema: &RelationDesc, - extensions: &FunctionExtensions, -) -> Result, Error> { - let e = Expression { - rex_type: Some(RexType::ScalarFunction(f.clone())), - }; - let schema = input_schema.to_df_schema()?; - - let extensions = extensions.to_extensions(); - let session_state = SessionStateBuilder::new() - .with_scalar_functions(all_default_functions()) - .build(); - let consumer = DefaultSubstraitConsumer::new(&extensions, &session_state); - let df_expr = - substrait::df_logical_plan::consumer::from_substrait_rex(&consumer, &e, &schema).await; - let expr = df_expr.context({ - DatafusionSnafu { - context: "Failed to convert substrait scalar function to datafusion scalar function", - } - })?; - let phy_expr = datafusion::physical_expr::create_physical_expr( - &expr, - &schema, - &Default::default(), - &PhysicalPlanningContext::default(), - ) - .context(DatafusionSnafu { - context: "Failed to create physical expression from logical expression", - })?; - Ok(phy_expr) -} - -/// Return an [`Expression`](wrapped in a [`FunctionArgument`]) that references the i-th column of the input relation -pub(crate) fn proto_col(i: usize) -> substrait_proto::proto::FunctionArgument { - use substrait_proto::proto::expression; - let expr = Expression { - rex_type: Some(expression::RexType::Selection(Box::new( - expression::FieldReference { - reference_type: Some(expression::field_reference::ReferenceType::DirectReference( - expression::ReferenceSegment { - reference_type: Some( - expression::reference_segment::ReferenceType::StructField(Box::new( - expression::reference_segment::StructField { - field: i as i32, - child: None, - }, - )), - ), - }, - )), - root_type: None, - }, - ))), - }; - substrait_proto::proto::FunctionArgument { - arg_type: Some(substrait_proto::proto::function_argument::ArgType::Value( - expr, - )), - } -} - -fn is_proto_literal(arg: &substrait_proto::proto::FunctionArgument) -> bool { - use substrait_proto::proto::expression; - matches!( - arg.arg_type.as_ref().unwrap(), - ArgType::Value(Expression { - rex_type: Some(expression::RexType::Literal(_)), - }) - ) -} - -fn build_proto_lit( - lit: substrait_proto::proto::expression::Literal, -) -> substrait_proto::proto::FunctionArgument { - use substrait_proto::proto; - proto::FunctionArgument { - arg_type: Some(ArgType::Value(Expression { - rex_type: Some(proto::expression::RexType::Literal(lit)), - })), - } -} - -/// rewrite ScalarFunction's arguments to Columns 0..n so nested exprs are still handled by us instead of datafusion -/// -/// specially, if a argument is a literal, the replacement will not happen -fn rewrite_scalar_function( - f: &ScalarFunction, - arg_typed_exprs: &[TypedExpr], -) -> Result { - let mut f_rewrite = f.clone(); - ensure!( - f_rewrite.arguments.len() == arg_typed_exprs.len(), - crate::error::InternalSnafu { - reason: format!( - "Expect `f_rewrite` and `arg_typed_expr` to be same length, found {} and {}", - f_rewrite.arguments.len(), - arg_typed_exprs.len() - ) - } - ); - for (idx, raw_expr) in f_rewrite.arguments.iter_mut().enumerate() { - // only replace it with col(idx) if it is not literal - // will try best to determine if it is literal, i.e. for function like `cast()` will try - // in both world to understand if it results in a literal - match ( - is_proto_literal(raw_expr), - arg_typed_exprs[idx].expr.is_literal(), - ) { - (false, false) => *raw_expr = proto_col(idx), - (true, _) => (), - (false, true) => { - if let ScalarExpr::Literal(val, ty) = &arg_typed_exprs[idx].expr { - let df_val = val - .try_to_scalar_value(ty) - .map_err(BoxedError::new) - .context(ExternalSnafu)?; - let lit_sub = to_substrait_literal(&df_val)?; - // put const-folded literal back to df to simplify stuff - *raw_expr = build_proto_lit(lit_sub); - } else { - UnexpectedSnafu { - reason: format!( - "Expect value to be literal, but found {:?}", - arg_typed_exprs[idx].expr - ), - } - .fail()? - } - } - } - } - Ok(f_rewrite) -} - -impl TypedExpr { - pub async fn from_substrait_to_datafusion_scalar_func( - f: &ScalarFunction, - arg_typed_exprs: Vec, - extensions: &FunctionExtensions, - ) -> Result { - let (arg_exprs, arg_types): (Vec<_>, Vec<_>) = arg_typed_exprs - .clone() - .into_iter() - .map(|e| (e.expr, e.typ)) - .unzip(); - debug!("Before rewrite: {:?}", f); - let f_rewrite = rewrite_scalar_function(f, &arg_typed_exprs)?; - debug!("After rewrite: {:?}", f_rewrite); - let input_schema = RelationType::new(arg_types).into_unnamed(); - let raw_fn = - RawDfScalarFn::from_proto(&f_rewrite, input_schema.clone(), extensions.clone())?; - - let df_func = DfScalarFunction::try_from_raw_fn(raw_fn).await?; - let expr = ScalarExpr::CallDf { - df_scalar_fn: df_func, - exprs: arg_exprs, - }; - // df already know it's own schema, so not providing here - let ret_type = expr.typ(&[])?; - Ok(TypedExpr::new(expr, ret_type)) - } - - /// Convert ScalarFunction into Flow's ScalarExpr - pub async fn from_substrait_scalar_func( - f: &ScalarFunction, - input_schema: &RelationDesc, - extensions: &FunctionExtensions, - ) -> Result { - let fn_name = - extensions - .get(&f.function_reference) - .with_context(|| NotImplementedSnafu { - reason: format!( - "Aggregated function not found: function reference = {:?}", - f.function_reference - ), - })?; - let arg_len = f.arguments.len(); - let arg_typed_exprs: Vec = { - let mut rets = Vec::new(); - for arg in f.arguments.iter() { - let ret = match &arg.arg_type { - Some(ArgType::Value(e)) => { - TypedExpr::from_substrait_rex(e, input_schema, extensions).await - } - _ => not_impl_err!("Aggregated function argument non-Value type not supported"), - }?; - rets.push(ret); - } - rets - }; - - // literal's type is determined by the function and type of other args - let (arg_exprs, arg_types): (Vec<_>, Vec<_>) = arg_typed_exprs - .clone() - .into_iter() - .map( - |TypedExpr { - expr: arg_val, - typ: arg_type, - }| { - if arg_val.is_literal() { - (arg_val, None) - } else { - (arg_val, Some(arg_type.scalar_type)) - } - }, - ) - .unzip(); - - match arg_len { - 1 if UnaryFunc::is_valid_func_name(fn_name) => { - let func = UnaryFunc::from_str_and_type(fn_name, None)?; - let arg = arg_exprs[0].clone(); - let ret_type = ColumnType::new_nullable(func.signature().output.clone()); - - Ok(TypedExpr::new(arg.call_unary(func), ret_type)) - } - 2 if fn_name == "arrow_cast" => { - let cast_to = arg_exprs[1] - .clone() - .as_literal() - .and_then(|lit| lit.as_string()) - .with_context(|| InvalidQuerySnafu { - reason: "array_cast's second argument must be a literal string", - })?; - let cast_to = typename_to_cdt(&cast_to)?; - let func = UnaryFunc::Cast(cast_to.clone()); - let arg = arg_exprs[0].clone(); - // constant folding here since some datafusion function require it for constant arg(i.e. `DATE_BIN`) - if arg.is_literal() { - let res = func.eval(&[], &arg).context(EvalSnafu)?; - Ok(TypedExpr::new( - ScalarExpr::Literal(res, cast_to.clone()), - ColumnType::new_nullable(cast_to), - )) - } else { - let ret_type = ColumnType::new_nullable(func.signature().output.clone()); - - Ok(TypedExpr::new(arg.call_unary(func), ret_type)) - } - } - 2 if BinaryFunc::is_valid_func_name(fn_name) => { - let (func, signature) = - BinaryFunc::from_str_expr_and_type(fn_name, &arg_exprs, &arg_types[0..2])?; - - // constant folding here - let is_all_literal = arg_exprs.iter().all(|arg| arg.is_literal()); - if is_all_literal { - let res = func - .eval(&[], &arg_exprs[0], &arg_exprs[1]) - .context(EvalSnafu)?; - - // if output type is null, it should be inferred from the input types - let con_typ = signature.output.clone(); - let typ = ColumnType::new_nullable(con_typ.clone()); - return Ok(TypedExpr::new(ScalarExpr::Literal(res, con_typ), typ)); - } - - let mut arg_exprs = arg_exprs; - for (idx, arg_expr) in arg_exprs.iter_mut().enumerate() { - if let ScalarExpr::Literal(val, typ) = arg_expr { - let dest_type = signature.input[idx].clone(); - - // cast val to target_type - let dest_val = if !dest_type.is_null() { - datatypes::types::cast(val.clone(), &dest_type) - .with_context(|_| - DatatypesSnafu{ - extra: format!("Failed to implicitly cast literal {val:?} to type {dest_type:?}") - })? - } else { - val.clone() - }; - *val = dest_val; - *typ = dest_type; - } - } - - let ret_type = ColumnType::new_nullable(func.signature().output.clone()); - let ret_expr = arg_exprs[0].clone().call_binary(arg_exprs[1].clone(), func); - Ok(TypedExpr::new(ret_expr, ret_type)) - } - _var => { - if fn_name == TUMBLE_START || fn_name == TUMBLE_END { - let (func, arg) = UnaryFunc::from_tumble_func(fn_name, &arg_typed_exprs)?; - - let ret_type = ColumnType::new_nullable(func.signature().output.clone()); - - Ok(TypedExpr::new(arg.expr.call_unary(func), ret_type)) - } else if VariadicFunc::is_valid_func_name(fn_name) { - let func = VariadicFunc::from_str_and_types(fn_name, &arg_types)?; - let ret_type = ColumnType::new_nullable(func.signature().output.clone()); - let mut expr = ScalarExpr::CallVariadic { - func, - exprs: arg_exprs, - }; - expr.optimize(); - Ok(TypedExpr::new(expr, ret_type)) - } else if UnmaterializableFunc::is_valid_func_name(fn_name) { - let func = UnmaterializableFunc::from_str_args(fn_name, arg_typed_exprs)?; - let ret_type = ColumnType::new_nullable(func.signature().output.clone()); - Ok(TypedExpr::new( - ScalarExpr::CallUnmaterializable(func), - ret_type, - )) - } else { - let try_as_df = Self::from_substrait_to_datafusion_scalar_func( - f, - arg_typed_exprs, - extensions, - ) - .await?; - Ok(try_as_df) - } - } - } - } - - /// Convert IfThen into Flow's ScalarExpr - pub async fn from_substrait_ifthen_rex( - if_then: &IfThen, - input_schema: &RelationDesc, - extensions: &FunctionExtensions, - ) -> Result { - let ifs: Vec<_> = { - let mut ifs = Vec::new(); - for if_clause in if_then.ifs.iter() { - let proto_if = if_clause.r#if.as_ref().with_context(|| InvalidQuerySnafu { - reason: "IfThen clause without if", - })?; - let proto_then = if_clause.then.as_ref().with_context(|| InvalidQuerySnafu { - reason: "IfThen clause without then", - })?; - let cond = - TypedExpr::from_substrait_rex(proto_if, input_schema, extensions).await?; - let then = - TypedExpr::from_substrait_rex(proto_then, input_schema, extensions).await?; - ifs.push((cond, then)); - } - ifs - }; - // if no else is presented - let els = match if_then - .r#else - .as_ref() - .map(|e| TypedExpr::from_substrait_rex(e, input_schema, extensions)) - { - Some(fut) => Some(fut.await), - None => None, - } - .transpose()? - .unwrap_or_else(|| { - TypedExpr::new( - ScalarExpr::literal_null(), - ColumnType::new_nullable(CDT::null_datatype()), - ) - }); - - fn build_if_then_recur( - mut next_if_then: impl Iterator, - els: TypedExpr, - ) -> TypedExpr { - if let Some((cond, then)) = next_if_then.next() { - // always assume the type of `if`` expr is the same with the `then`` expr - TypedExpr::new( - ScalarExpr::If { - cond: Box::new(cond.expr), - then: Box::new(then.expr), - els: Box::new(build_if_then_recur(next_if_then, els).expr), - }, - then.typ, - ) - } else { - els - } - } - let expr_if = build_if_then_recur(ifs.into_iter(), els); - Ok(expr_if) - } - /// Convert Substrait Rex into Flow's ScalarExpr - #[async_recursion::async_recursion] - pub async fn from_substrait_rex( - e: &Expression, - input_schema: &RelationDesc, - extensions: &FunctionExtensions, - ) -> Result { - match &e.rex_type { - Some(RexType::Literal(lit)) => { - let lit = from_substrait_literal(lit)?; - Ok(TypedExpr::new( - ScalarExpr::Literal(lit.0, lit.1.clone()), - ColumnType::new_nullable(lit.1), - )) - } - Some(RexType::SingularOrList(s)) => { - let substrait_expr = s.value.as_ref().with_context(|| InvalidQuerySnafu { - reason: "SingularOrList expression without value", - })?; - let typed_expr = - TypedExpr::from_substrait_rex(substrait_expr, input_schema, extensions).await?; - // Note that we didn't impl support to in list expr - if !s.options.is_empty() { - let mut list = Vec::with_capacity(s.options.len()); - for opt in s.options.iter() { - let opt_expr = - TypedExpr::from_substrait_rex(opt, input_schema, extensions).await?; - list.push(opt_expr.expr); - } - let in_list_expr = ScalarExpr::InList { - expr: Box::new(typed_expr.expr), - list, - }; - Ok(TypedExpr::new( - in_list_expr, - ColumnType::new_nullable(CDT::boolean_datatype()), - )) - } else { - Ok(typed_expr) - } - } - Some(RexType::Selection(field_ref)) => match &field_ref.reference_type { - Some(DirectReference(direct)) => match &direct.reference_type.as_ref() { - Some(StructField(x)) => match &x.child.as_ref() { - Some(_) => { - not_impl_err!( - "Direct reference StructField with child is not supported" - ) - } - None => { - let column = x.field as usize; - let column_type = input_schema.typ().column_types[column].clone(); - Ok(TypedExpr::new(ScalarExpr::Column(column), column_type)) - } - }, - _ => not_impl_err!( - "Direct reference with types other than StructField is not supported" - ), - }, - _ => not_impl_err!("unsupported field ref type"), - }, - Some(RexType::ScalarFunction(f)) => { - TypedExpr::from_substrait_scalar_func(f, input_schema, extensions).await - } - Some(RexType::IfThen(if_then)) => { - TypedExpr::from_substrait_ifthen_rex(if_then, input_schema, extensions).await - } - Some(RexType::Cast(cast)) => { - let input = cast.input.as_ref().with_context(|| InvalidQuerySnafu { - reason: "Cast expression without input", - })?; - let input = TypedExpr::from_substrait_rex(input, input_schema, extensions).await?; - let cast_type = from_substrait_type(cast.r#type.as_ref().with_context(|| { - InvalidQuerySnafu { - reason: "Cast expression without type", - } - })?)?; - let func = UnaryFunc::from_str_and_type("cast", Some(cast_type.clone()))?; - Ok(TypedExpr::new( - input.expr.call_unary(func), - ColumnType::new_nullable(cast_type), - )) - } - Some(RexType::WindowFunction(_)) => PlanSnafu { - reason: - "Window function is not supported yet. Please use aggregation function instead." - .to_string(), - } - .fail(), - _ => not_impl_err!("unsupported rex_type"), - } - } -} - -#[cfg(test)] -mod test { - use datatypes::prelude::ConcreteDataType; - use datatypes::value::Value; - use pretty_assertions::assert_eq; - - use super::*; - use crate::expr::{GlobalId, MapFilterProject}; - use crate::plan::{Plan, TypedPlan}; - use crate::repr::{self, ColumnType, RelationType}; - use crate::transform::test::{create_test_ctx, create_test_query_engine, sql_to_substrait}; - - /// test if `WHERE` condition can be converted to Flow's ScalarExpr in mfp's filter - #[tokio::test] - async fn test_where_and() { - let engine = create_test_query_engine(); - let sql = - "SELECT number FROM numbers_with_ts WHERE number >= 1 AND number <= 3 AND number!=2"; - let plan = sql_to_substrait(engine.clone(), sql).await; - - let mut ctx = create_test_ctx(); - let flow_plan = TypedPlan::from_substrait_plan(&mut ctx, &plan).await; - - // optimize binary and to variadic and - let filter = ScalarExpr::CallVariadic { - func: VariadicFunc::And, - exprs: vec![ - ScalarExpr::Column(2).call_binary( - ScalarExpr::Literal(Value::from(1i64), CDT::int64_datatype()), - BinaryFunc::Gte, - ), - ScalarExpr::Column(2).call_binary( - ScalarExpr::Literal(Value::from(3i64), CDT::int64_datatype()), - BinaryFunc::Lte, - ), - ScalarExpr::Column(2).call_binary( - ScalarExpr::Literal(Value::from(2i64), CDT::int64_datatype()), - BinaryFunc::NotEq, - ), - ], - }; - let expected = TypedPlan { - schema: RelationType::new(vec![ColumnType::new(CDT::uint32_datatype(), false)]) - .into_named(vec![Some("number".to_string())]), - plan: Plan::Mfp { - input: Box::new( - Plan::Get { - id: crate::expr::Id::Global(GlobalId::User(1)), - } - .with_types( - RelationType::new(vec![ - ColumnType::new(ConcreteDataType::uint32_datatype(), false), - ColumnType::new( - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ), - ]) - .into_named(vec![Some("number".to_string()), Some("ts".to_string())]), - ), - ), - mfp: MapFilterProject::new(2) - .map(vec![ - ScalarExpr::CallUnary { - func: UnaryFunc::Cast(CDT::int64_datatype()), - expr: Box::new(ScalarExpr::Column(0)), - }, - ScalarExpr::Column(0), - ScalarExpr::Column(3), - ]) - .unwrap() - .filter(vec![filter]) - .unwrap() - .project(vec![4]) - .unwrap(), - }, - }; - assert_eq!(flow_plan.unwrap(), expected); - } - - /// case: binary functions&constant folding can happen in converting substrait plan - #[tokio::test] - async fn test_binary_func_and_constant_folding() { - let engine = create_test_query_engine(); - let sql = "SELECT 1+1*2-1/1+1%2==3 FROM numbers"; - let plan = sql_to_substrait(engine.clone(), sql).await; - - let mut ctx = create_test_ctx(); - let flow_plan = TypedPlan::from_substrait_plan(&mut ctx, &plan).await; - - let expected = TypedPlan { - schema: RelationType::new(vec![ColumnType::new(CDT::boolean_datatype(), true)]) - .into_named(vec![Some("Int64(1) + Int64(1) * Int64(2) - Int64(1) / Int64(1) + Int64(1) % Int64(2) = Int64(3)".to_string())]), - plan: Plan::Constant { - rows: vec![( - repr::Row::new(vec![Value::from(true)]), - repr::Timestamp::MIN, - 1, - )], - }, - }; - - assert_eq!(flow_plan.unwrap(), expected); - } - - /// test if the type of the literal is correctly inferred, i.e. in here literal is decoded to be int64, but need to be uint32, - #[tokio::test] - async fn test_implicitly_cast() { - let engine = create_test_query_engine(); - let sql = "SELECT number+1 FROM numbers"; - let plan = sql_to_substrait(engine.clone(), sql).await; - - let mut ctx = create_test_ctx(); - let flow_plan = TypedPlan::from_substrait_plan(&mut ctx, &plan).await; - - let expected = TypedPlan { - schema: RelationType::new(vec![ColumnType::new(CDT::int64_datatype(), true)]) - .into_named(vec![Some("numbers.number + Int64(1)".to_string())]), - plan: Plan::Mfp { - input: Box::new( - Plan::Get { - id: crate::expr::Id::Global(GlobalId::User(0)), - } - .with_types( - RelationType::new(vec![ColumnType::new( - ConcreteDataType::uint32_datatype(), - false, - )]) - .into_named(vec![Some("number".to_string())]), - ), - ), - mfp: MapFilterProject::new(1) - .map(vec![ - ScalarExpr::Column(0) - .call_unary(UnaryFunc::Cast(CDT::int64_datatype())) - .call_binary( - ScalarExpr::Literal(Value::from(1i64), CDT::int64_datatype()), - BinaryFunc::AddInt64, - ), - ]) - .unwrap() - .project(vec![1]) - .unwrap(), - }, - }; - assert_eq!(flow_plan.unwrap(), expected); - } - - #[tokio::test] - async fn test_cast() { - let engine = create_test_query_engine(); - let sql = "SELECT CAST(1 AS INT16) FROM numbers"; - let plan = sql_to_substrait(engine.clone(), sql).await; - - let mut ctx = create_test_ctx(); - let flow_plan = TypedPlan::from_substrait_plan(&mut ctx, &plan).await; - - let expected = TypedPlan { - schema: RelationType::new(vec![ColumnType::new(CDT::int16_datatype(), true)]) - .into_named(vec![Some( - "arrow_cast(Int64(1),Utf8(\"Int16\"))".to_string(), - )]), - plan: Plan::Constant { - // cast of literal is constant folded - rows: vec![(repr::Row::new(vec![Value::from(1i16)]), i64::MIN, 1)], - }, - }; - assert_eq!(flow_plan.unwrap(), expected); - } - - #[tokio::test] - async fn test_select_add() { - let engine = create_test_query_engine(); - let sql = "SELECT number+number FROM numbers"; - let plan = sql_to_substrait(engine.clone(), sql).await; - - let mut ctx = create_test_ctx(); - let flow_plan = TypedPlan::from_substrait_plan(&mut ctx, &plan).await; - - let expected = TypedPlan { - schema: RelationType::new(vec![ColumnType::new(CDT::uint32_datatype(), true)]) - .into_named(vec![Some("numbers.number + numbers.number".to_string())]), - plan: Plan::Mfp { - input: Box::new( - Plan::Get { - id: crate::expr::Id::Global(GlobalId::User(0)), - } - .with_types( - RelationType::new(vec![ColumnType::new( - ConcreteDataType::uint32_datatype(), - false, - )]) - .into_named(vec![Some("number".to_string())]), - ), - ), - mfp: MapFilterProject::new(1) - .map(vec![ - ScalarExpr::Column(0) - .call_binary(ScalarExpr::Column(0), BinaryFunc::AddUInt32), - ]) - .unwrap() - .project(vec![1]) - .unwrap(), - }, - }; - - assert_eq!(flow_plan.unwrap(), expected); - } - - #[tokio::test] - async fn test_func_sig() { - fn lit(v: impl ToString) -> substrait_proto::proto::FunctionArgument { - use substrait_proto::proto::expression; - let expr = Expression { - rex_type: Some(expression::RexType::Literal(expression::Literal { - nullable: false, - type_variation_reference: 0, - literal_type: Some(expression::literal::LiteralType::String(v.to_string())), - })), - }; - substrait_proto::proto::FunctionArgument { - arg_type: Some(substrait_proto::proto::function_argument::ArgType::Value( - expr, - )), - } - } - - let f = substrait_proto::proto::expression::ScalarFunction { - function_reference: 0, - arguments: vec![proto_col(0)], - options: vec![], - output_type: None, - ..Default::default() - }; - let input_schema = - RelationType::new(vec![ColumnType::new(CDT::uint32_datatype(), false)]).into_unnamed(); - let extensions = FunctionExtensions::from_iter([(0, "is_null".to_string())]); - let res = TypedExpr::from_substrait_scalar_func(&f, &input_schema, &extensions) - .await - .unwrap(); - - assert_eq!( - res, - TypedExpr { - expr: ScalarExpr::Column(0).call_unary(UnaryFunc::IsNull), - typ: ColumnType { - scalar_type: CDT::boolean_datatype(), - nullable: true, - }, - } - ); - - let f = substrait_proto::proto::expression::ScalarFunction { - function_reference: 0, - arguments: vec![proto_col(0), proto_col(1)], - options: vec![], - output_type: None, - ..Default::default() - }; - let input_schema = RelationType::new(vec![ - ColumnType::new(CDT::uint32_datatype(), false), - ColumnType::new(CDT::uint32_datatype(), false), - ]) - .into_unnamed(); - let extensions = FunctionExtensions::from_iter([(0, "add".to_string())]); - let res = TypedExpr::from_substrait_scalar_func(&f, &input_schema, &extensions) - .await - .unwrap(); - - assert_eq!( - res, - TypedExpr { - expr: ScalarExpr::Column(0) - .call_binary(ScalarExpr::Column(1), BinaryFunc::AddUInt32,), - typ: ColumnType { - scalar_type: CDT::uint32_datatype(), - nullable: true, - }, - } - ); - } -} diff --git a/src/flow/src/transform/literal.rs b/src/flow/src/transform/literal.rs deleted file mode 100644 index eaebd8ff05c..00000000000 --- a/src/flow/src/transform/literal.rs +++ /dev/null @@ -1,426 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::array::TryFromSliceError; - -use bytes::Bytes; -use common_decimal::Decimal128; -use common_time::timestamp::TimeUnit; -use common_time::{Date, IntervalMonthDayNano, Timestamp}; -use datafusion_common::ScalarValue; -use datatypes::data_type::ConcreteDataType as CDT; -use datatypes::value::Value; -use num_traits::FromBytes; -use snafu::OptionExt; -use substrait::variation_const::{ - DATE_32_TYPE_VARIATION_REF, DATE_64_TYPE_VARIATION_REF, DEFAULT_TYPE_VARIATION_REF, - UNSIGNED_INTEGER_TYPE_VARIATION_REF, -}; -use substrait_proto::proto; -use substrait_proto::proto::expression::Literal; -use substrait_proto::proto::expression::literal::{LiteralType, PrecisionTimestamp}; -use substrait_proto::proto::r#type::Kind; - -use crate::error::{Error, NotImplementedSnafu, PlanSnafu, UnexpectedSnafu}; -use crate::transform::substrait_proto; - -#[derive(Debug)] -enum TimestampPrecision { - Second = 0, - Millisecond = 3, - Microsecond = 6, - Nanosecond = 9, -} - -impl TryFrom for TimestampPrecision { - type Error = Error; - - fn try_from(prec: i32) -> Result { - match prec { - 0 => Ok(Self::Second), - 3 => Ok(Self::Millisecond), - 6 => Ok(Self::Microsecond), - 9 => Ok(Self::Nanosecond), - _ => not_impl_err!("Unsupported precision: {prec}"), - } - } -} - -impl TimestampPrecision { - fn to_time_unit(&self) -> TimeUnit { - match self { - Self::Second => TimeUnit::Second, - Self::Millisecond => TimeUnit::Millisecond, - Self::Microsecond => TimeUnit::Microsecond, - Self::Nanosecond => TimeUnit::Nanosecond, - } - } - - fn to_cdt(&self) -> CDT { - match self { - Self::Second => CDT::timestamp_second_datatype(), - Self::Millisecond => CDT::timestamp_millisecond_datatype(), - Self::Microsecond => CDT::timestamp_microsecond_datatype(), - Self::Nanosecond => CDT::timestamp_nanosecond_datatype(), - } - } -} - -/// TODO(discord9): this is copy from datafusion-substrait since the original function is not public, will be replace once is exported -pub(crate) fn to_substrait_literal(value: &ScalarValue) -> Result { - if value.is_null() { - return not_impl_err!("Unsupported literal: {value:?}"); - } - let (literal_type, type_variation_reference) = match value { - ScalarValue::Boolean(Some(b)) => (LiteralType::Boolean(*b), DEFAULT_TYPE_VARIATION_REF), - ScalarValue::Int8(Some(n)) => (LiteralType::I8(*n as i32), DEFAULT_TYPE_VARIATION_REF), - ScalarValue::UInt8(Some(n)) => ( - LiteralType::I8(*n as i32), - UNSIGNED_INTEGER_TYPE_VARIATION_REF, - ), - ScalarValue::Int16(Some(n)) => (LiteralType::I16(*n as i32), DEFAULT_TYPE_VARIATION_REF), - ScalarValue::UInt16(Some(n)) => ( - LiteralType::I16(*n as i32), - UNSIGNED_INTEGER_TYPE_VARIATION_REF, - ), - ScalarValue::Int32(Some(n)) => (LiteralType::I32(*n), DEFAULT_TYPE_VARIATION_REF), - ScalarValue::UInt32(Some(n)) => ( - LiteralType::I32(*n as i32), - UNSIGNED_INTEGER_TYPE_VARIATION_REF, - ), - ScalarValue::Int64(Some(n)) => (LiteralType::I64(*n), DEFAULT_TYPE_VARIATION_REF), - ScalarValue::UInt64(Some(n)) => ( - LiteralType::I64(*n as i64), - UNSIGNED_INTEGER_TYPE_VARIATION_REF, - ), - ScalarValue::Float32(Some(f)) => (LiteralType::Fp32(*f), DEFAULT_TYPE_VARIATION_REF), - ScalarValue::Float64(Some(f)) => (LiteralType::Fp64(*f), DEFAULT_TYPE_VARIATION_REF), - // TODO(discord9): deal with timezone - ScalarValue::TimestampSecond(Some(t), _) => ( - LiteralType::PrecisionTimestamp(PrecisionTimestamp { - value: *t, - precision: TimestampPrecision::Second as i32, - }), - DEFAULT_TYPE_VARIATION_REF, - ), - ScalarValue::TimestampMillisecond(Some(t), _) => ( - LiteralType::PrecisionTimestamp(PrecisionTimestamp { - value: *t, - precision: TimestampPrecision::Millisecond as i32, - }), - DEFAULT_TYPE_VARIATION_REF, - ), - ScalarValue::TimestampMicrosecond(Some(t), _) => ( - LiteralType::PrecisionTimestamp(PrecisionTimestamp { - value: *t, - precision: TimestampPrecision::Microsecond as i32, - }), - DEFAULT_TYPE_VARIATION_REF, - ), - ScalarValue::TimestampNanosecond(Some(t), _) => ( - LiteralType::PrecisionTimestamp(PrecisionTimestamp { - value: *t, - precision: TimestampPrecision::Nanosecond as i32, - }), - DEFAULT_TYPE_VARIATION_REF, - ), - ScalarValue::Date32(Some(d)) => (LiteralType::Date(*d), DATE_32_TYPE_VARIATION_REF), - _ => ( - not_impl_err!("Unsupported literal: {value:?}")?, - DEFAULT_TYPE_VARIATION_REF, - ), - }; - - Ok(Literal { - nullable: false, - type_variation_reference, - literal_type: Some(literal_type), - }) -} - -/// Convert a Substrait literal into a Value and its ConcreteDataType (So that we can know type even if the value is null) -pub(crate) fn from_substrait_literal(lit: &Literal) -> Result<(Value, CDT), Error> { - let scalar_value = match &lit.literal_type { - Some(LiteralType::Boolean(b)) => (Value::from(*b), CDT::boolean_datatype()), - Some(LiteralType::I8(n)) => match lit.type_variation_reference { - DEFAULT_TYPE_VARIATION_REF => (Value::from(*n as i8), CDT::int8_datatype()), - UNSIGNED_INTEGER_TYPE_VARIATION_REF => (Value::from(*n as u8), CDT::uint8_datatype()), - others => not_impl_err!("Unknown type variation reference {others}",)?, - }, - Some(LiteralType::I16(n)) => match lit.type_variation_reference { - DEFAULT_TYPE_VARIATION_REF => (Value::from(*n as i16), CDT::int16_datatype()), - UNSIGNED_INTEGER_TYPE_VARIATION_REF => (Value::from(*n as u16), CDT::uint16_datatype()), - others => not_impl_err!("Unknown type variation reference {others}",)?, - }, - Some(LiteralType::I32(n)) => match lit.type_variation_reference { - DEFAULT_TYPE_VARIATION_REF => (Value::from(*n), CDT::int32_datatype()), - UNSIGNED_INTEGER_TYPE_VARIATION_REF => (Value::from(*n as u32), CDT::uint32_datatype()), - others => not_impl_err!("Unknown type variation reference {others}",)?, - }, - Some(LiteralType::I64(n)) => match lit.type_variation_reference { - DEFAULT_TYPE_VARIATION_REF => (Value::from(*n), CDT::int64_datatype()), - UNSIGNED_INTEGER_TYPE_VARIATION_REF => (Value::from(*n as u64), CDT::uint64_datatype()), - others => not_impl_err!("Unknown type variation reference {others}",)?, - }, - Some(LiteralType::Fp32(f)) => (Value::from(*f), CDT::float32_datatype()), - Some(LiteralType::Fp64(f)) => (Value::from(*f), CDT::float64_datatype()), - Some(LiteralType::Timestamp(t)) => ( - Value::from(Timestamp::new_microsecond(*t)), - CDT::timestamp_microsecond_datatype(), - ), - Some(LiteralType::PrecisionTimestamp(prec_ts)) => { - let (prec, val) = (prec_ts.precision, prec_ts.value); - let prec = TimestampPrecision::try_from(prec)?; - let unit = prec.to_time_unit(); - let typ = prec.to_cdt(); - (Value::from(Timestamp::new(val, unit)), typ) - } - Some(LiteralType::Date(d)) => (Value::from(Date::new(*d)), CDT::date_datatype()), - Some(LiteralType::String(s)) => (Value::from(s.clone()), CDT::string_datatype()), - Some(LiteralType::Binary(b)) | Some(LiteralType::FixedBinary(b)) => { - (Value::from(b.clone()), CDT::binary_datatype()) - } - Some(LiteralType::Decimal(d)) => { - let value: [u8; 16] = d.value.clone().try_into().map_err(|e| { - PlanSnafu { - reason: format!("Failed to parse decimal value from {e:?}"), - } - .build() - })?; - let p: u8 = d.precision.try_into().map_err(|e| { - PlanSnafu { - reason: format!("Failed to parse decimal precision: {e}"), - } - .build() - })?; - let s: i8 = d.scale.try_into().map_err(|e| { - PlanSnafu { - reason: format!("Failed to parse decimal scale: {e}"), - } - .build() - })?; - let value = i128::from_le_bytes(value); - ( - Value::from(Decimal128::new(value, p, s)), - CDT::decimal128_datatype(p, s), - ) - } - Some(LiteralType::Null(ntype)) => (Value::Null, from_substrait_type(ntype)?), - Some(LiteralType::IntervalDayToSecond(interval)) => from_interval_day_sec(interval)?, - Some(LiteralType::IntervalYearToMonth(interval)) => from_interval_year_month(interval)?, - Some(LiteralType::IntervalCompound(interval_compound)) => { - let interval_day_time = &interval_compound - .interval_day_to_second - .map(|i| from_interval_day_sec(&i)) - .transpose()?; - let interval_year_month = &interval_compound - .interval_year_to_month - .map(|i| from_interval_year_month(&i)) - .transpose()?; - let mut compound = IntervalMonthDayNano::new(0, 0, 0); - if let Some(day_sec) = interval_day_time { - let Value::IntervalDayTime(day_time) = day_sec.0 else { - UnexpectedSnafu { - reason: format!("Expect IntervalDayTime, found {:?}", day_sec), - } - .fail()? - }; - // 1 day in milliseconds = 24 * 60 * 60 * 1000 = 8.64e7 ms = 8.64e13 ns << 2^63 - // so overflow is unexpected - compound.nanoseconds = compound - .nanoseconds - .checked_add(day_time.milliseconds as i64 * 1_000_000) - .with_context(|| UnexpectedSnafu { - reason: format!( - "Overflow when converting interval: {:?}", - interval_compound - ), - })?; - compound.days += day_time.days; - } - - if let Some(year_month) = interval_year_month { - let Value::IntervalYearMonth(year_month) = year_month.0 else { - UnexpectedSnafu { - reason: format!("Expect IntervalYearMonth, found {:?}", year_month), - } - .fail()? - }; - compound.months += year_month.months; - } - - ( - Value::IntervalMonthDayNano(compound), - CDT::interval_month_day_nano_datatype(), - ) - } - _ => not_impl_err!("unsupported literal_type: {:?}", &lit.literal_type)?, - }; - Ok(scalar_value) -} - -fn from_interval_day_sec( - interval: &proto::expression::literal::IntervalDayToSecond, -) -> Result<(Value, CDT), Error> { - let (days, seconds, subseconds) = (interval.days, interval.seconds, interval.subseconds); - let millis = if let Some(prec) = interval.precision_mode { - use substrait_proto::proto::expression::literal::interval_day_to_second::PrecisionMode; - match prec { - PrecisionMode::Precision(e) => { - if e >= 3 { - subseconds - / 10_i64 - .checked_pow((e - 3) as _) - .with_context(|| UnexpectedSnafu { - reason: format!( - "Overflow when converting interval: {:?}", - interval - ), - })? - } else { - subseconds - * 10_i64 - .checked_pow((3 - e) as _) - .with_context(|| UnexpectedSnafu { - reason: format!( - "Overflow when converting interval: {:?}", - interval - ), - })? - } - } - PrecisionMode::Microseconds(_) => subseconds / 1000, - } - } else if subseconds == 0 { - 0 - } else { - not_impl_err!("unsupported subseconds without precision_mode: {subseconds}")? - }; - - let value_interval = common_time::IntervalDayTime::new(days, seconds * 1000 + millis as i32); - - Ok(( - Value::IntervalDayTime(value_interval), - CDT::interval_day_time_datatype(), - )) -} - -fn from_interval_year_month( - interval: &proto::expression::literal::IntervalYearToMonth, -) -> Result<(Value, CDT), Error> { - let value_interval = common_time::IntervalYearMonth::new(interval.years * 12 + interval.months); - - Ok(( - Value::IntervalYearMonth(value_interval), - CDT::interval_year_month_datatype(), - )) -} - -fn from_bytes(i: &Bytes) -> Result -where - for<'a> &'a ::Bytes: - std::convert::TryFrom<&'a [u8], Error = TryFromSliceError>, -{ - let (int_bytes, _rest) = i.split_at(std::mem::size_of::()); - let i = T::from_le_bytes(int_bytes.try_into().map_err(|e| { - UnexpectedSnafu { - reason: format!( - "Expect slice to be {} bytes, found {} bytes, error={:?}", - std::mem::size_of::(), - int_bytes.len(), - e - ), - } - .build() - })?); - Ok(i) -} - -/// convert a Substrait type into a ConcreteDataType -pub fn from_substrait_type(null_type: &substrait_proto::proto::Type) -> Result { - if let Some(kind) = &null_type.kind { - match kind { - Kind::Bool(_) => Ok(CDT::boolean_datatype()), - Kind::I8(integer) => match integer.type_variation_reference { - DEFAULT_TYPE_VARIATION_REF => Ok(CDT::int8_datatype()), - UNSIGNED_INTEGER_TYPE_VARIATION_REF => Ok(CDT::uint8_datatype()), - v => not_impl_err!("Unsupported Substrait type variation {v} of type {kind:?}"), - }, - Kind::I16(integer) => match integer.type_variation_reference { - DEFAULT_TYPE_VARIATION_REF => Ok(CDT::int16_datatype()), - UNSIGNED_INTEGER_TYPE_VARIATION_REF => Ok(CDT::uint16_datatype()), - v => not_impl_err!("Unsupported Substrait type variation {v} of type {kind:?}"), - }, - Kind::I32(integer) => match integer.type_variation_reference { - DEFAULT_TYPE_VARIATION_REF => Ok(CDT::int32_datatype()), - UNSIGNED_INTEGER_TYPE_VARIATION_REF => Ok(CDT::uint32_datatype()), - v => not_impl_err!("Unsupported Substrait type variation {v} of type {kind:?}"), - }, - Kind::I64(integer) => match integer.type_variation_reference { - DEFAULT_TYPE_VARIATION_REF => Ok(CDT::int64_datatype()), - UNSIGNED_INTEGER_TYPE_VARIATION_REF => Ok(CDT::uint64_datatype()), - v => not_impl_err!("Unsupported Substrait type variation {v} of type {kind:?}"), - }, - Kind::Fp32(_) => Ok(CDT::float32_datatype()), - Kind::Fp64(_) => Ok(CDT::float64_datatype()), - Kind::PrecisionTimestamp(ts) => { - Ok(TimestampPrecision::try_from(ts.precision)?.to_cdt()) - } - Kind::Date(date) => match date.type_variation_reference { - DATE_32_TYPE_VARIATION_REF | DATE_64_TYPE_VARIATION_REF => Ok(CDT::date_datatype()), - v => not_impl_err!("Unsupported Substrait type variation {v} of type {kind:?}"), - }, - Kind::Binary(_) => Ok(CDT::binary_datatype()), - Kind::String(_) => Ok(CDT::string_datatype()), - Kind::Decimal(d) => Ok(CDT::decimal128_datatype(d.precision as u8, d.scale as i8)), - _ => not_impl_err!("Unsupported Substrait type: {kind:?}"), - } - } else { - not_impl_err!("Null type without kind is not supported") - } -} - -#[cfg(test)] -mod test { - use pretty_assertions::assert_eq; - - use super::*; - use crate::plan::{Plan, TypedPlan}; - use crate::repr::{self, ColumnType, RelationType}; - use crate::transform::test::{create_test_ctx, create_test_query_engine, sql_to_substrait}; - /// test if literal in substrait plan can be correctly converted to flow plan - #[tokio::test] - async fn test_literal() { - let engine = create_test_query_engine(); - let sql = "SELECT 1 FROM numbers"; - let plan = sql_to_substrait(engine.clone(), sql).await; - - let mut ctx = create_test_ctx(); - let flow_plan = TypedPlan::from_substrait_plan(&mut ctx, &plan).await; - - let expected = TypedPlan { - schema: RelationType::new(vec![ColumnType::new(CDT::int64_datatype(), true)]) - .into_named(vec![Some("Int64(1)".to_string())]), - plan: Plan::Constant { - rows: vec![( - repr::Row::new(vec![Value::Int64(1)]), - repr::Timestamp::MIN, - 1, - )], - }, - }; - - assert_eq!(flow_plan.unwrap(), expected); - } -} diff --git a/src/flow/src/transform/plan.rs b/src/flow/src/transform/plan.rs deleted file mode 100644 index cdd558d84d9..00000000000 --- a/src/flow/src/transform/plan.rs +++ /dev/null @@ -1,276 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::collections::HashSet; - -use itertools::Itertools; -use snafu::OptionExt; -use substrait::substrait_proto_df::proto::{FilterRel, ReadRel}; -use substrait_proto::proto::expression::MaskExpression; -use substrait_proto::proto::read_rel::ReadType; -use substrait_proto::proto::rel::RelType; -use substrait_proto::proto::{Plan as SubPlan, ProjectRel, Rel, plan_rel}; - -use crate::error::{Error, InvalidQuerySnafu, NotImplementedSnafu, PlanSnafu, UnexpectedSnafu}; -use crate::expr::{MapFilterProject, TypedExpr}; -use crate::plan::{Plan, TypedPlan}; -use crate::repr::{self, RelationType}; -use crate::transform::{FlownodeContext, FunctionExtensions, substrait_proto}; - -impl TypedPlan { - /// Convert Substrait Plan into Flow's TypedPlan - pub async fn from_substrait_plan( - ctx: &mut FlownodeContext, - plan: &SubPlan, - ) -> Result { - // Register function extension - let function_extension = FunctionExtensions::try_from_proto(&plan.extensions)?; - - // Parse relations - match plan.relations.len() { - 1 => match plan.relations[0].rel_type.as_ref() { - Some(rt) => match rt { - plan_rel::RelType::Rel(rel) => { - Ok(TypedPlan::from_substrait_rel(ctx, rel, &function_extension).await?) - } - plan_rel::RelType::Root(root) => { - let input = root.input.as_ref().with_context(|| InvalidQuerySnafu { - reason: "Root relation without input", - })?; - - let mut ret = - TypedPlan::from_substrait_rel(ctx, input, &function_extension).await?; - - if !root.names.is_empty() { - ret.schema = ret.schema.clone().try_with_names(root.names.clone())?; - } - - Ok(ret) - } - }, - None => plan_err!("Cannot parse plan relation: None"), - }, - _ => not_impl_err!( - "Substrait plan with more than 1 relation trees not supported. Number of relation trees: {:?}", - plan.relations.len() - ), - } - } - - #[async_recursion::async_recursion] - pub async fn from_substrait_project( - ctx: &mut FlownodeContext, - p: &ProjectRel, - extensions: &FunctionExtensions, - ) -> Result { - let input = if let Some(input) = p.input.as_ref() { - TypedPlan::from_substrait_rel(ctx, input, extensions).await? - } else { - return not_impl_err!("Projection without an input is not supported"); - }; - - // because this `input.schema` is incorrect for pre-expand substrait plan, so we have to use schema before expand multi-value - // function to correctly transform it, and late rewrite it - // TODO(discord9): this logic is obsoleted since now expand happens in datafusion optimizer - let schema_before_expand = { - let input_schema = input.schema.clone(); - let auto_columns: HashSet = - HashSet::from_iter(input_schema.typ().auto_columns.clone()); - let not_auto_added_columns = (0..input_schema.len()?) - .filter(|i| !auto_columns.contains(i)) - .collect_vec(); - let mfp = MapFilterProject::new(input_schema.len()?) - .project(not_auto_added_columns)? - .into_safe(); - - input_schema.apply_mfp(&mfp)? - }; - - let mut exprs: Vec = Vec::with_capacity(p.expressions.len()); - for e in &p.expressions { - let expr = TypedExpr::from_substrait_rex(e, &schema_before_expand, extensions).await?; - exprs.push(expr); - } - let is_literal = exprs.iter().all(|expr| expr.expr.is_literal()); - if is_literal { - let (literals, lit_types): (Vec<_>, Vec<_>) = exprs - .into_iter() - .map(|TypedExpr { expr, typ }| (expr, typ)) - .unzip(); - let typ = RelationType::new(lit_types); - let row = literals - .into_iter() - .map(|lit| lit.as_literal().expect("A literal")) - .collect_vec(); - let row = repr::Row::new(row); - let plan = Plan::Constant { - rows: vec![(row, repr::Timestamp::MIN, 1)], - }; - Ok(TypedPlan { - schema: typ.into_unnamed(), - plan, - }) - } else { - input.projection(exprs) - } - } - - #[async_recursion::async_recursion] - pub async fn from_substrait_filter( - ctx: &mut FlownodeContext, - filter: &FilterRel, - extensions: &FunctionExtensions, - ) -> Result { - let input = if let Some(input) = filter.input.as_ref() { - TypedPlan::from_substrait_rel(ctx, input, extensions).await? - } else { - return not_impl_err!("Filter without an input is not supported"); - }; - - let expr = if let Some(condition) = filter.condition.as_ref() { - TypedExpr::from_substrait_rex(condition, &input.schema, extensions).await? - } else { - return not_impl_err!("Filter without an condition is not valid"); - }; - input.filter(expr) - } - - pub async fn from_substrait_read( - ctx: &mut FlownodeContext, - read: &ReadRel, - _extensions: &FunctionExtensions, - ) -> Result { - if let Some(ReadType::NamedTable(nt)) = &read.read_type { - let query_ctx = ctx.query_context.clone().context(UnexpectedSnafu { - reason: "Query context not found", - })?; - let table_reference = match nt.names.len() { - 1 => [ - query_ctx.current_catalog().to_string(), - query_ctx.current_schema().clone(), - nt.names[0].clone(), - ], - 2 => [ - query_ctx.current_catalog().to_string(), - nt.names[0].clone(), - nt.names[1].clone(), - ], - 3 => [ - nt.names[0].clone(), - nt.names[1].clone(), - nt.names[2].clone(), - ], - _ => InvalidQuerySnafu { - reason: "Expect table to have name", - } - .fail()?, - }; - - let table = ctx.table(&table_reference).await?; - let get_table = Plan::Get { - id: crate::expr::Id::Global(table.0), - }; - let get_table = TypedPlan { - schema: table.1, - plan: get_table, - }; - - if let Some(MaskExpression { - select: Some(projection), - .. - }) = &read.projection - { - let column_indices: Vec = projection - .struct_items - .iter() - .map(|item| item.field as usize) - .collect(); - let input_arity = get_table.schema.typ().column_types.len(); - let mfp = MapFilterProject::new(input_arity).project(column_indices.clone())?; - get_table.mfp(mfp.into_safe()) - } else { - Ok(get_table) - } - } else { - not_impl_err!("Only NamedTable reads are supported") - } - } - - /// Convert Substrait Rel into Flow's TypedPlan - /// TODO(discord9): SELECT DISTINCT(does it get compile with something else?) - pub async fn from_substrait_rel( - ctx: &mut FlownodeContext, - rel: &Rel, - extensions: &FunctionExtensions, - ) -> Result { - match &rel.rel_type { - Some(RelType::Project(p)) => { - Self::from_substrait_project(ctx, p.as_ref(), extensions).await - } - Some(RelType::Filter(filter)) => { - Self::from_substrait_filter(ctx, filter, extensions).await - } - Some(RelType::Read(read)) => Self::from_substrait_read(ctx, read, extensions).await, - Some(RelType::Aggregate(agg)) => { - Self::from_substrait_agg_rel(ctx, agg, extensions).await - } - _ => not_impl_err!("Unsupported relation type: {:?}", rel.rel_type), - } - } -} - -#[cfg(test)] -mod test { - use datatypes::data_type::ConcreteDataType as CDT; - use datatypes::prelude::ConcreteDataType; - use pretty_assertions::assert_eq; - - use super::*; - use crate::expr::GlobalId; - use crate::plan::{Plan, TypedPlan}; - use crate::repr::{ColumnType, RelationType}; - use crate::transform::test::{create_test_ctx, create_test_query_engine, sql_to_substrait}; - - #[tokio::test] - async fn test_select() { - let engine = create_test_query_engine(); - let sql = "SELECT number FROM numbers"; - let plan = sql_to_substrait(engine.clone(), sql).await; - - let mut ctx = create_test_ctx(); - let flow_plan = TypedPlan::from_substrait_plan(&mut ctx, &plan).await; - - let expected = TypedPlan { - schema: RelationType::new(vec![ColumnType::new(CDT::uint32_datatype(), false)]) - .into_named(vec![Some("number".to_string())]), - plan: Plan::Mfp { - input: Box::new( - Plan::Get { - id: crate::expr::Id::Global(GlobalId::User(0)), - } - .with_types( - RelationType::new(vec![ColumnType::new( - ConcreteDataType::uint32_datatype(), - false, - )]) - .into_named(vec![Some("number".to_string())]), - ), - ), - mfp: MapFilterProject::new(1), - }, - }; - - assert_eq!(flow_plan.unwrap(), expected); - } -} diff --git a/src/flow/src/utils.rs b/src/flow/src/utils.rs index 1a9879b9964..575be9138f3 100644 --- a/src/flow/src/utils.rs +++ b/src/flow/src/utils.rs @@ -12,33 +12,17 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! utilities for managing state of dataflow execution - -use std::collections::{BTreeMap, BTreeSet}; -use std::ops::Bound; -use std::sync::Arc; - use common_meta::key::flow::flow_state::FlowStat; -use common_telemetry::trace; use datatypes::value::Value; -use get_size2::GetSize; -use smallvec::{SmallVec, smallvec}; -use tokio::sync::{RwLock, mpsc, oneshot}; -use tokio::time::Instant; +use tokio::sync::{mpsc, oneshot}; use crate::error::InternalSnafu; -use crate::expr::{EvalError, ScalarExpr}; -use crate::repr::{DiffRow, Duration, KeyValDiffRow, Row, Timestamp, value_to_internal_ts}; -/// A batch of updates, arranged by key -pub type Batch = BTreeMap>; - -/// Get a estimate of heap size of a value -pub fn get_value_heap_size(v: &Value) -> usize { - match v { - Value::Binary(bin) => bin.len(), - Value::String(s) => s.len(), - Value::List(list) => list.items().iter().map(get_value_heap_size).sum(), +pub fn get_value_heap_size(value: &Value) -> usize { + match value { + Value::Binary(v) => v.len(), + Value::String(v) => v.len(), + Value::List(v) => v.items().iter().map(get_value_heap_size).sum(), _ => 0, } } @@ -47,995 +31,33 @@ pub fn get_value_heap_size(v: &Value) -> usize { pub struct SizeReportSender { inner: mpsc::Sender>, } - impl SizeReportSender { pub fn new() -> (Self, StateReportHandler) { let (tx, rx) = mpsc::channel(1); - let zelf = Self { inner: tx }; - (zelf, rx) + (Self { inner: tx }, rx) } - - /// Query the size report, will timeout after one second if no response pub async fn query(&self, timeout: std::time::Duration) -> crate::Result { let (tx, rx) = oneshot::channel(); self.inner.send(tx).await.map_err(|_| { InternalSnafu { - reason: "failed to send size report request due to receiver dropped", + reason: "failed to send size report request", } .build() })?; - let timeout = tokio::time::timeout(timeout, rx); - timeout + tokio::time::timeout(timeout, rx) .await - .map_err(|_elapsed| { + .map_err(|_| { InternalSnafu { - reason: "failed to receive size report after one second timeout", + reason: "failed to receive size report after timeout", } .build() })? .map_err(|_| { InternalSnafu { - reason: "failed to receive size report due to sender dropped", + reason: "size report sender dropped", } .build() }) } } - -/// Handle the size report request, and send the report back pub type StateReportHandler = mpsc::Receiver>; - -/// A spine of batches, arranged by timestamp -/// TODO(discord9): consider internally index by key, value, and timestamp for faster lookup -pub type Spine = BTreeMap; - -/// Determine when should a key expire according to it's event timestamp in key. -/// -/// If a key is expired, any future updates to it should be ignored. -/// -/// Note that key is expired by it's event timestamp (contained in the key), not by the time it's inserted (system timestamp). -#[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd)] -pub struct KeyExpiryManager { - /// A map from event timestamp to key, used for expire keys. - event_ts_to_key: BTreeMap>, - - /// Duration after which a key is considered expired, and will be removed from state - key_expiration_duration: Option, - - /// Expression to get timestamp from key row - event_timestamp_from_row: Option, -} - -impl GetSize for KeyExpiryManager { - fn get_heap_size(&self) -> usize { - let row_size = if let Some(row_size) = &self - .event_ts_to_key - .first_key_value() - .map(|(_, v)| v.first().get_heap_size()) - { - *row_size - } else { - 0 - }; - self.event_ts_to_key - .values() - .map(|v| v.len() * row_size + std::mem::size_of::()) - .sum::() - } -} - -impl KeyExpiryManager { - pub fn new( - key_expiration_duration: Option, - event_timestamp_from_row: Option, - ) -> Self { - Self { - event_ts_to_key: Default::default(), - key_expiration_duration, - event_timestamp_from_row, - } - } - - /// Extract event timestamp from key row. - /// - /// If no expire state is set, return None. - pub fn extract_event_ts(&self, row: &Row) -> Result, EvalError> { - let ts = self - .event_timestamp_from_row - .as_ref() - .map(|e| e.eval(&row.inner)) - .transpose()? - .map(value_to_internal_ts) - .transpose()?; - Ok(ts) - } - - /// Return timestamp that should be expired by the time `now` by compute `now - expiration_duration` - pub fn compute_expiration_timestamp(&self, now: Timestamp) -> Option { - self.key_expiration_duration.map(|d| now - d) - } - - /// Update the event timestamp to key mapping. - /// - /// - If given key is expired by now (that is less than `now - expiry_duration`), return the amount of time it's expired. - /// - If it's not expired, return None - pub fn get_expire_duration_and_update_event_ts( - &mut self, - now: Timestamp, - row: &Row, - ) -> Result, EvalError> { - let Some(event_ts) = self.extract_event_ts(row)? else { - return Ok(None); - }; - - self.event_ts_to_key - .entry(event_ts) - .or_default() - .insert(row.clone()); - - if let Some(expire_time) = self.compute_expiration_timestamp(now) - && expire_time > event_ts - { - // return how much time it's expired - return Ok(Some(expire_time - event_ts)); - } - - Ok(None) - } - - /// Get the expire duration of a key, if it's expired by now. - /// - /// Return None if the key is not expired - pub fn get_expire_duration( - &self, - now: Timestamp, - row: &Row, - ) -> Result, EvalError> { - let Some(event_ts) = self.extract_event_ts(row)? else { - return Ok(None); - }; - - if let Some(expire_time) = self.compute_expiration_timestamp(now) - && expire_time > event_ts - { - // return how much time it's expired - return Ok(Some(expire_time - event_ts)); - } - - Ok(None) - } - - /// Remove expired keys from the state, and return an iterator of removed keys with - /// event_ts less than expire time (i.e. now - key_expiration_duration). - pub fn remove_expired_keys(&mut self, now: Timestamp) -> Option> { - let expire_time = self.compute_expiration_timestamp(now)?; - - let mut before = self.event_ts_to_key.split_off(&expire_time); - std::mem::swap(&mut before, &mut self.event_ts_to_key); - - Some(before.into_values().flat_map(|keys| keys.into_iter())) - } -} - -/// A shared state of key-value pair for various state in dataflow execution. -/// -/// i.e: Mfp operator with temporal filter need to store it's future output so that it can add now, and delete later. -/// To get all needed updates in a time span, use [`get_updates_in_range`]. -/// -/// And reduce operator need full state of it's output, so that it can query (and modify by calling [`apply_updates`]) -/// existing state, also need a way to expire keys. To get a key's current value, use [`get`] with time being `now` -/// so it's like: -/// `mfp operator -> arrange(store futures only, no expire) -> reduce operator <-> arrange(full, with key expiring time) -> output` -/// -/// Note the two way arrow between reduce operator and arrange, it's because reduce operator need to query existing state -/// and also need to update existing state. -#[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd)] -pub struct Arrangement { - /// A name or identifier for the arrangement which can be used for debugging or logging purposes. - /// This field is not critical to the functionality but aids in monitoring and management of arrangements. - name: Vec, - - /// Manages a collection of pending updates in a `BTreeMap` where each key is a timestamp and each value is a `Batch` of updates. - /// Updates are grouped into batched based on their timestamps. - /// Each batch covers a range of time from the last key (exclusive) to the current key (inclusive). - /// - /// - Updates with a timestamp (`update_ts`) that falls between two keys are placed in the batch of the higher key. - /// For example, if the keys are `1, 5, 7, 9` and `update_ts` is `6`, the update goes into the batch with key `7`. - /// - Updates with a timestamp before the first key are categorized under the first key. - /// - Updates with a timestamp greater than the highest key result in a new batch being created with that timestamp as the key. - /// - /// The first key represents the current state and includes consolidated updates from the past. It is always set to `now`. - /// Each key should have only one update per batch with a `diff=1` for the batch representing the current time (`now`). - /// - /// Since updates typically occur as a delete followed by an insert, a small vector of size 2 is used to store updates for efficiency. - /// - /// TODO(discord9): Consider balancing the batch size? - spine: Spine, - - /// Indicates whether the arrangement maintains a complete history of updates. - /// - `true`: Maintains all past and future updates, necessary for full state reconstruction at any point in time. - /// - `false`: Only future updates are retained, optimizing for scenarios where past state is irrelevant and conserving resources. - /// Useful for case like `map -> arrange -> reduce`. - full_arrangement: bool, - - /// Indicates whether the arrangement has been modified since its creation. - /// - `true`: The arrangement has been written to, meaning it has received updates. - /// Cloning this arrangement is generally unsafe as it may lead to inconsistencies if the clone is modified independently. - /// However, cloning is safe when both the original and the clone require a full arrangement, as this ensures consistency. - /// - `false`: The arrangement is in its initial state and has not been modified. It can be safely cloned and shared - /// without concerns of carrying over unintended state changes. - is_written: bool, - - /// Manage the expire state of the arrangement. - expire_state: Option, - - /// The time that the last compaction happened, also known as the current time. - last_compaction_time: Option, - - /// Estimated size of the arrangement in heap size. - estimated_size: usize, - last_size_update: Instant, - size_update_interval: tokio::time::Duration, -} - -impl Arrangement { - fn compute_size(&self) -> usize { - self.spine - .values() - .map(|v| { - let per_entry_size = v - .first_key_value() - .map(|(k, v)| { - k.get_heap_size() - + v.len() * v.first().map(|r| r.get_heap_size()).unwrap_or(0) - }) - .unwrap_or(0); - std::mem::size_of::() + v.len() * per_entry_size - }) - .sum::() - + self.expire_state.get_heap_size() - + self.name.get_heap_size() - } - - fn update_and_fetch_size(&mut self) -> usize { - if self.last_size_update.elapsed() > self.size_update_interval { - self.estimated_size = self.compute_size(); - self.last_size_update = Instant::now(); - } - self.estimated_size - } -} - -impl GetSize for Arrangement { - fn get_heap_size(&self) -> usize { - self.estimated_size - } -} - -impl Default for Arrangement { - fn default() -> Self { - Self { - spine: Default::default(), - full_arrangement: false, - is_written: false, - expire_state: None, - last_compaction_time: None, - name: Vec::new(), - estimated_size: 0, - last_size_update: Instant::now(), - size_update_interval: tokio::time::Duration::from_secs(3), - } - } -} - -impl Arrangement { - pub fn new_with_name(name: Vec) -> Self { - Self { - spine: Default::default(), - full_arrangement: false, - is_written: false, - expire_state: None, - last_compaction_time: None, - name, - estimated_size: 0, - last_size_update: Instant::now(), - size_update_interval: tokio::time::Duration::from_secs(3), - } - } - - pub fn get_expire_state(&self) -> Option<&KeyExpiryManager> { - self.expire_state.as_ref() - } - - pub fn set_expire_state(&mut self, expire_state: KeyExpiryManager) { - self.expire_state = Some(expire_state); - } - - /// Apply updates into spine, with no respect of whether the updates are in futures, past, or now. - /// - /// Return the maximum expire time (already expire by how much time) of all updates if any keys is already expired. - pub fn apply_updates( - &mut self, - now: Timestamp, - updates: Vec, - ) -> Result, EvalError> { - self.is_written = true; - - let mut max_expired_by: Option = None; - - for ((key, val), update_ts, diff) in updates { - // check if the key is expired - if let Some(s) = &mut self.expire_state - && let Some(expired_by) = s.get_expire_duration_and_update_event_ts(now, &key)? - { - max_expired_by = max_expired_by.max(Some(expired_by)); - trace!( - "Expired key: {:?}, expired by: {:?} with time being now={}", - key, expired_by, now - ); - continue; - } - - // If the `highest_ts` is less than `update_ts`, we need to create a new batch with key being `update_ts`. - if self - .spine - .last_key_value() - .map(|(highest_ts, _)| *highest_ts < update_ts) - .unwrap_or(true) - { - self.spine.insert(update_ts, Default::default()); - } - - // Get the first batch with key that's greater or equal to `update_ts`. - let (_, batch) = self - .spine - .range_mut(update_ts..) - .next() - .expect("Previous insert should have created the batch"); - - let key_updates = batch.entry(key).or_default(); - key_updates.push((val, update_ts, diff)); - - // a stable sort make updates sort in order of insertion - // without changing the order of updates within same tick - key_updates.sort_by_key(|(_val, ts, _diff)| *ts); - } - self.update_and_fetch_size(); - Ok(max_expired_by) - } - - /// Find out the time of next update in the future that is the next update with `timestamp > now`. - pub fn get_next_update_time(&self, now: &Timestamp) -> Option { - // iter over batches that only have updates of `timestamp>now` and find the first non empty batch, then get the minimum timestamp in that batch - for (_ts, batch) in self.spine.range((Bound::Excluded(now), Bound::Unbounded)) { - let min_ts = batch - .values() - .flat_map(|v| v.iter().map(|(_, ts, _)| *ts).min()) - .min(); - - if min_ts.is_some() { - return min_ts; - } - } - - None - } - - /// Get the last compaction time. - pub fn last_compaction_time(&self) -> Option { - self.last_compaction_time - } - - /// Split spine off at `split_ts`, and return the spine that's before `split_ts` (including `split_ts`). - fn split_spine_le(&mut self, split_ts: &Timestamp) -> Spine { - self.split_batch_at(split_ts); - let mut before = self.spine.split_off(&(split_ts + 1)); - std::mem::swap(&mut before, &mut self.spine); - before - } - - /// Split the batch at `split_ts` into two parts. - fn split_batch_at(&mut self, split_ts: &Timestamp) { - // FAST PATH: - // - // The `split_ts` hit the boundary of a batch, nothing to do. - if self.spine.contains_key(split_ts) { - return; - } - - let Some((_, batch_to_split)) = self.spine.range_mut(split_ts..).next() else { - return; // No batch to split, nothing to do. - }; - - // SLOW PATH: - // - // The `split_ts` is in the middle of a batch, we need to split the batch into two parts. - let mut new_batch = Batch::default(); - - batch_to_split.retain(|key, updates| { - let mut new_updates = SmallVec::default(); - - updates.retain(|(val, ts, diff)| { - if *ts <= *split_ts { - // Move the updates that are less than or equal to `split_ts` to the new batch. - new_updates.push((val.clone(), *ts, *diff)); - } - // Keep the updates that are greater than `split_ts` in the current batch. - *ts > *split_ts - }); - - if !new_updates.is_empty() { - new_batch.insert(key.clone(), new_updates); - } - - // Keep the key in the current batch if it still has updates. - !updates.is_empty() - }); - - if !new_batch.is_empty() { - self.spine.insert(*split_ts, new_batch); - } - } - - /// Advance time to `now` and consolidate all older (`now` included) updates to the first key. - /// - /// Return the maximum expire time(already expire by how much time) of all updates if any keys is already expired. - pub fn compact_to(&mut self, now: Timestamp) -> Result, EvalError> { - let mut max_expired_by: Option = None; - - let batches_to_compact = self.split_spine_le(&now); - self.last_compaction_time = Some(now); - - // If a full arrangement is not needed, we can just discard everything before and including now, - if !self.full_arrangement { - return Ok(None); - } - - // else we update them into current state. - let mut compacting_batch = Batch::default(); - - for (_, batch) in batches_to_compact { - for (key, updates) in batch { - // check if the key is expired - if let Some(s) = &mut self.expire_state - && let Some(expired_by) = - s.get_expire_duration_and_update_event_ts(now, &key)? - { - max_expired_by = max_expired_by.max(Some(expired_by)); - continue; - } - - let mut row = compacting_batch - .remove(&key) - // only one row in the updates during compaction - .and_then(|mut updates| updates.pop()); - - for update in updates { - row = compact_diff_row(row, &update); - } - if let Some(compacted_update) = row { - compacting_batch.insert(key, smallvec![compacted_update]); - } - } - } - - // insert the compacted batch into spine with key being `now` - self.spine.insert(now, compacting_batch); - self.update_and_fetch_size(); - Ok(max_expired_by) - } - - /// Get the updates of the arrangement from the given range of time. - pub fn get_updates_in_range + Clone>( - &self, - range: R, - ) -> Vec { - // Include the next batch in case the range is not aligned with the boundary of a batch. - let batches = match range.end_bound() { - Bound::Included(t) => self.spine.range(range.clone()).chain( - self.spine - .range((Bound::Excluded(t), Bound::Unbounded)) - .next(), - ), - Bound::Excluded(t) => self.spine.range(range.clone()).chain( - self.spine - .range((Bound::Included(t), Bound::Unbounded)) - .next(), - ), - _ => self.spine.range(range.clone()).chain(None), - }; - - let mut res = vec![]; - for (_, batch) in batches { - for (key, updates) in batch { - for (val, ts, diff) in updates { - if range.contains(ts) { - res.push(((key.clone(), val.clone()), *ts, *diff)); - } - } - } - } - res - } - - /// Expire keys in now that are older than expire_time, intended for reducing memory usage and limit late data arrive - pub fn truncate_expired_keys(&mut self, now: Timestamp) { - if let Some(s) = &mut self.expire_state - && let Some(expired_keys) = s.remove_expired_keys(now) - { - for key in expired_keys { - for (_, batch) in self.spine.iter_mut() { - batch.remove(&key); - } - } - } - } - - /// Get current state of things. - /// - /// Useful for query existing keys (i.e. reduce and join operator need to query existing state) - pub fn get(&self, now: Timestamp, key: &Row) -> Option { - // FAST PATH: - // - // If `now <= last_compaction_time`, and it's full arrangement, we can directly return the value - // from the current state (which should be the first batch in the spine if it exist). - if let Some(last_compaction_time) = self.last_compaction_time() - && now <= last_compaction_time - && self.full_arrangement - { - // if the last compaction time's batch is not exist, it means the spine doesn't have it's first batch as current value - return self - .spine - .get(&last_compaction_time) - .and_then(|batch| batch.get(key)) - .and_then(|updates| updates.first().cloned()); - } - - // SLOW PATH: - // - // Accumulate updates from the oldest batch to the batch containing `now`. - - let batches = if self.spine.contains_key(&now) { - // hit the boundary of a batch - self.spine.range(..=now).chain(None) - } else { - // not hit the boundary of a batch, should include the next batch - self.spine.range(..=now).chain( - self.spine - .range((Bound::Excluded(now), Bound::Unbounded)) - .next(), - ) - }; - - let mut final_val = None; - for (ts, batch) in batches { - if let Some(updates) = batch.get(key) { - if *ts <= now { - for update in updates { - final_val = compact_diff_row(final_val, update); - } - } else { - for update in updates.iter().filter(|(_, ts, _)| *ts <= now) { - final_val = compact_diff_row(final_val, update); - } - } - } - } - final_val - } -} - -fn compact_diff_row(old_row: Option, new_row: &DiffRow) -> Option { - let (val, ts, diff) = new_row; - match (old_row, diff) { - (Some((row, _old_ts, old_diff)), diff) if row == *val && old_diff + diff == 0 => { - // the key is deleted now - None - } - (Some((row, _old_ts, old_diff)), diff) if row == *val && old_diff + diff != 0 => { - Some((row, *ts, old_diff + *diff)) - } - // if old val not equal new val, simple consider it as being overwritten, for each key can only have one value - // so it make sense to just replace the old value with new value - _ => Some((val.clone(), *ts, *diff)), - } -} - -/// Simply a type alias for ReadGuard of Arrangement -pub type ArrangeReader<'a> = tokio::sync::RwLockReadGuard<'a, Arrangement>; -/// Simply a type alias for WriteGuard of Arrangement -pub type ArrangeWriter<'a> = tokio::sync::RwLockWriteGuard<'a, Arrangement>; - -/// A handler to the inner Arrangement, can be cloned and shared, useful for query it's inner state -#[derive(Debug, Clone)] -pub struct ArrangeHandler { - inner: Arc>, -} - -impl ArrangeHandler { - /// create a new handler from arrangement - pub fn from(arr: Arrangement) -> Self { - Self { - inner: Arc::new(RwLock::new(arr)), - } - } - - /// write lock the arrangement - pub fn write(&self) -> ArrangeWriter<'_> { - self.inner.blocking_write() - } - - /// read lock the arrangement - pub fn read(&self) -> ArrangeReader<'_> { - self.inner.blocking_read() - } - - /// Clone the handler, but only keep the future updates. - /// - /// It's a cheap operation, since it's `Arc-ed` and only clone the `Arc`. - pub fn clone_future_only(&self) -> Option { - if self.read().is_written { - return None; - } - Some(Self { - inner: self.inner.clone(), - }) - } - - /// Clone the handler, but keep all updates. - /// - /// Prevent illegal clone after the arrange have been written, - /// because that will cause loss of data before clone. - /// - /// It's a cheap operation, since it's `Arc-ed` and only clone the `Arc`. - pub fn clone_full_arrange(&self) -> Option { - { - let zelf = self.read(); - if !zelf.full_arrangement && zelf.is_written { - return None; - } - } - - self.write().full_arrangement = true; - Some(Self { - inner: self.inner.clone(), - }) - } - - pub fn set_full_arrangement(&self, full: bool) { - self.write().full_arrangement = full; - } - - pub fn is_full_arrangement(&self) -> bool { - self.read().full_arrangement - } -} - -#[cfg(test)] -mod test { - use std::borrow::Borrow; - - use datatypes::value::Value; - use itertools::Itertools; - - use super::*; - - fn lit(v: impl Into) -> Row { - Row::new(vec![v.into()]) - } - - fn kv(key: impl Borrow, value: impl Borrow) -> (Row, Row) { - (key.borrow().clone(), value.borrow().clone()) - } - - #[test] - fn test_future_get() { - // test if apply only future updates, whether get(future_time) can operate correctly - let arr = ArrangeHandler::from(Arrangement::default()); - - let mut arr = arr.write(); - - let key = lit("a"); - let updates: Vec = vec![ - (kv(&key, lit("b")), 1 /* ts */, 1 /* diff */), - (kv(&key, lit("c")), 2 /* ts */, 1 /* diff */), - (kv(&key, lit("d")), 3 /* ts */, 1 /* diff */), - ]; - - // all updates above are future updates - arr.apply_updates(0, updates).unwrap(); - - assert_eq!(arr.get(1, &key), Some((lit("b"), 1 /* ts */, 1 /* diff */))); - assert_eq!(arr.get(2, &key), Some((lit("c"), 2 /* ts */, 1 /* diff */))); - assert_eq!(arr.get(3, &key), Some((lit("d"), 3 /* ts */, 1 /* diff */))); - } - - #[test] - fn only_save_future_updates() { - // mfp operator's temporal filter need to record future updates so that it can delete on time - // i.e. insert a record now, delete this record 5 minutes later - // they will only need to keep future updates(if downstream don't need full arrangement that is) - let arr = ArrangeHandler::from(Arrangement::default()); - - { - let arr1 = arr.clone_full_arrange(); - assert!(arr1.is_some()); - let arr2 = arr.clone_future_only(); - assert!(arr2.is_some()); - } - - { - let mut arr = arr.write(); - let updates: Vec = vec![ - (kv(lit("a"), lit("x")), 1 /* ts */, 1 /* diff */), - (kv(lit("b"), lit("y")), 2 /* ts */, 1 /* diff */), - (kv(lit("c"), lit("z")), 3 /* ts */, 1 /* diff */), - ]; - // all updates above are future updates - arr.apply_updates(0, updates).unwrap(); - - assert_eq!( - arr.get_updates_in_range(1..=1), - vec![(kv(lit("a"), lit("x")), 1 /* ts */, 1 /* diff */)] - ); - assert_eq!(arr.spine.len(), 3); - - arr.compact_to(1).unwrap(); - assert_eq!(arr.spine.len(), 3); - - let key = &lit("a"); - assert_eq!(arr.get(3, key), Some((lit("x"), 1 /* ts */, 1 /* diff */))); - let key = &lit("b"); - assert_eq!(arr.get(3, key), Some((lit("y"), 2 /* ts */, 1 /* diff */))); - let key = &lit("c"); - assert_eq!(arr.get(3, key), Some((lit("z"), 3 /* ts */, 1 /* diff */))); - } - - assert!(arr.clone_future_only().is_none()); - { - let arr2 = arr.clone_full_arrange().unwrap(); - let mut arr = arr2.write(); - assert_eq!(arr.spine.len(), 3); - - arr.compact_to(2).unwrap(); - assert_eq!(arr.spine.len(), 2); - let key = &lit("a"); - assert_eq!(arr.get(3, key), Some((lit("x"), 1 /* ts */, 1 /* diff */))); - let key = &lit("b"); - assert_eq!(arr.get(3, key), Some((lit("y"), 2 /* ts */, 1 /* diff */))); - let key = &lit("c"); - assert_eq!(arr.get(3, key), Some((lit("z"), 3 /* ts */, 1 /* diff */))); - } - } - - #[test] - fn test_reduce_expire_keys() { - let mut arr = Arrangement::default(); - let expire_state = KeyExpiryManager { - event_ts_to_key: Default::default(), - key_expiration_duration: Some(10), - event_timestamp_from_row: Some(ScalarExpr::Column(0)), - }; - arr.expire_state = Some(expire_state); - arr.full_arrangement = true; - - let arr = ArrangeHandler::from(arr); - - let updates: Vec = vec![ - (kv(lit(1i64), lit("x")), 1 /* ts */, 1 /* diff */), - (kv(lit(2i64), lit("y")), 2 /* ts */, 1 /* diff */), - (kv(lit(3i64), lit("z")), 3 /* ts */, 1 /* diff */), - ]; - { - let mut arr = arr.write(); - arr.apply_updates(0, updates.clone()).unwrap(); - // repeat the same updates means having multiple updates for the same key - arr.apply_updates(0, updates).unwrap(); - - assert_eq!( - arr.get_updates_in_range(1..=1), - vec![ - (kv(lit(1i64), lit("x")), 1 /* ts */, 1 /* diff */), - (kv(lit(1i64), lit("x")), 1 /* ts */, 1 /* diff */) - ] - ); - assert_eq!(arr.spine.len(), 3); - arr.compact_to(1).unwrap(); - assert_eq!(arr.spine.len(), 3); - } - - { - let mut arr = arr.write(); - assert_eq!(arr.spine.len(), 3); - - arr.truncate_expired_keys(11); - assert_eq!(arr.spine.len(), 3); - let key = &lit(1i64); - assert_eq!(arr.get(11, key), Some((lit("x"), 1 /* ts */, 2 /* diff */))); - let key = &lit(2i64); - assert_eq!(arr.get(11, key), Some((lit("y"), 2 /* ts */, 2 /* diff */))); - let key = &lit(3i64); - assert_eq!(arr.get(11, key), Some((lit("z"), 3 /* ts */, 2 /* diff */))); - - arr.truncate_expired_keys(12); - assert_eq!(arr.spine.len(), 3); - let key = &lit(1i64); - assert_eq!(arr.get(12, key), None); - let key = &lit(2i64); - assert_eq!(arr.get(12, key), Some((lit("y"), 2 /* ts */, 2 /* diff */))); - let key = &lit(3i64); - assert_eq!(arr.get(12, key), Some((lit("z"), 3 /* ts */, 2 /* diff */))); - assert_eq!(arr.expire_state.as_ref().unwrap().event_ts_to_key.len(), 2); - - arr.truncate_expired_keys(13); - assert_eq!(arr.spine.len(), 3); - let key = &lit(1i64); - assert_eq!(arr.get(13, key), None); - let key = &lit(2i64); - assert_eq!(arr.get(13, key), None); - let key = &lit(3i64); - assert_eq!(arr.get(13, key), Some((lit("z"), 3 /* ts */, 2 /* diff */))); - assert_eq!(arr.expire_state.as_ref().unwrap().event_ts_to_key.len(), 1); - } - } - - #[test] - fn test_apply_expired_keys() { - // apply updates with a expired key - let mut arr = Arrangement::default(); - let expire_state = KeyExpiryManager { - event_ts_to_key: Default::default(), - key_expiration_duration: Some(10), - event_timestamp_from_row: Some(ScalarExpr::Column(0)), - }; - arr.expire_state = Some(expire_state); - - let arr = ArrangeHandler::from(arr); - - let updates: Vec = vec![ - (kv(lit(1i64), lit("x")), 1 /* ts */, 1 /* diff */), - (kv(lit(2i64), lit("y")), 2 /* ts */, 1 /* diff */), - ]; - { - let mut arr = arr.write(); - let expired_by = arr.apply_updates(12, updates).unwrap(); - assert_eq!(expired_by, Some(1)); - - let key = &lit(1i64); - assert_eq!(arr.get(12, key), None); - let key = &lit(2i64); - assert_eq!(arr.get(12, key), Some((lit("y"), 2 /* ts */, 1 /* diff */))); - } - } - - /// test if split_spine_le get ranges that are not aligned with batch boundaries - /// this split_spine_le can correctly retrieve all updates in the range, including updates that are in the batches - /// near the boundary of input range - #[test] - fn test_split_off() { - let mut arr = Arrangement::default(); - // manually create batch ..=1 and 2..=3 - arr.spine.insert(1, Batch::default()); - arr.spine.insert(3, Batch::default()); - - let updates = vec![(kv(lit("a"), lit("x")), 2 /* ts */, 1 /* diff */)]; - // updates falls into the range of 2..=3 - arr.apply_updates(2, updates).unwrap(); - - let mut arr1 = arr.clone(); - { - assert_eq!(arr.get_next_update_time(&1), Some(2)); - // split expect to take batch ..=1 and create a new batch 2..=2 (which contains update) - let split = &arr.split_spine_le(&2); - assert_eq!(split.len(), 2); - assert_eq!(split[&2].len(), 1); - - assert_eq!(arr.get_next_update_time(&1), None); - } - - { - // take all updates with timestamp <=1, will get no updates - let split = &arr1.split_spine_le(&1); - assert_eq!(split.len(), 1); - assert_eq!(split[&1].len(), 0); - } - } - - /// test if get ranges is not aligned with boundary of batch, - /// whether can get correct result - #[test] - fn test_get_by_range() { - let mut arr = Arrangement::default(); - - // will form {2: [2, 1], 4: [4,3], 6: [6,5]} three batch - // TODO(discord9): manually set batch - let updates: Vec = vec![ - (kv(lit("a"), lit("")), 2 /* ts */, 1 /* diff */), - (kv(lit("a"), lit("")), 1 /* ts */, 1 /* diff */), - (kv(lit("b"), lit("")), 4 /* ts */, 1 /* diff */), - (kv(lit("c"), lit("")), 3 /* ts */, 1 /* diff */), - (kv(lit("c"), lit("")), 6 /* ts */, 1 /* diff */), - (kv(lit("a"), lit("")), 5 /* ts */, 1 /* diff */), - ]; - arr.apply_updates(0, updates).unwrap(); - assert_eq!( - arr.get_updates_in_range(2..=5), - vec![ - (kv(lit("a"), lit("")), 2 /* ts */, 1 /* diff */), - (kv(lit("b"), lit("")), 4 /* ts */, 1 /* diff */), - (kv(lit("c"), lit("")), 3 /* ts */, 1 /* diff */), - (kv(lit("a"), lit("")), 5 /* ts */, 1 /* diff */), - ] - ); - } - - /// test if get with range unaligned with batch boundary - /// can get correct result - #[test] - fn test_get_unaligned() { - let mut arr = Arrangement::default(); - - // will form {2: [2, 1], 4: [4,3], 6: [6,5]} three batch - // TODO(discord9): manually set batch - let key = &lit("a"); - let updates: Vec = vec![ - (kv(key, lit(1)), 2 /* ts */, 1 /* diff */), - (kv(key, lit(2)), 1 /* ts */, 1 /* diff */), - (kv(key, lit(3)), 4 /* ts */, 1 /* diff */), - (kv(key, lit(4)), 3 /* ts */, 1 /* diff */), - (kv(key, lit(5)), 6 /* ts */, 1 /* diff */), - (kv(key, lit(6)), 5 /* ts */, 1 /* diff */), - ]; - arr.apply_updates(0, updates).unwrap(); - // aligned with batch boundary - assert_eq!(arr.get(2, key), Some((lit(1), 2 /* ts */, 1 /* diff */))); - // unaligned with batch boundary - assert_eq!(arr.get(3, key), Some((lit(4), 3 /* ts */, 1 /* diff */))); - } - - /// test if out of order updates can be sorted correctly - #[test] - fn test_out_of_order_apply_updates() { - let mut arr = Arrangement::default(); - - let key = &lit("a"); - let updates: Vec = vec![ - (kv(key, lit(5)), 6 /* ts */, 1 /* diff */), - (kv(key, lit(2)), 2 /* ts */, -1 /* diff */), - (kv(key, lit(1)), 2 /* ts */, 1 /* diff */), - (kv(key, lit(2)), 1 /* ts */, 1 /* diff */), - (kv(key, lit(3)), 4 /* ts */, 1 /* diff */), - (kv(key, lit(4)), 3 /* ts */, 1 /* diff */), - (kv(key, lit(6)), 5 /* ts */, 1 /* diff */), - ]; - arr.apply_updates(0, updates.clone()).unwrap(); - let sorted = updates - .iter() - .sorted_by_key(|(_, ts, _)| *ts) - .cloned() - .collect_vec(); - assert_eq!(arr.get_updates_in_range(1..7), sorted); - } - - #[test] - fn test_full_arrangement_get_from_first_entry() { - let mut arr = Arrangement::default(); - // will form {3: [1, 2, 3]} - let updates = vec![ - (kv(lit("a"), lit("x")), 3 /* ts */, 1 /* diff */), - (kv(lit("b"), lit("y")), 1 /* ts */, 1 /* diff */), - (kv(lit("b"), lit("y")), 2 /* ts */, -1 /* diff */), - ]; - arr.apply_updates(0, updates).unwrap(); - assert_eq!(arr.get(2, &lit("b")), None /* deleted */); - arr.full_arrangement = true; - assert_eq!(arr.get(2, &lit("b")), None /* still deleted */); - - arr.compact_to(1).unwrap(); - - assert_eq!( - arr.get(1, &lit("b")), - Some((lit("y"), 1, 1)) /* fast path */ - ); - } -} diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index 0e54ba42440..9121d1ef382 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -65,6 +65,7 @@ use common_telemetry::{debug, info, tracing, warn}; use common_time::{Timestamp, Timezone}; use datafusion_common::tree_node::TreeNodeVisitor; use datafusion_expr::LogicalPlan; +use datafusion_expr::logical_plan::Distinct; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::value::Value; @@ -301,13 +302,33 @@ fn determine_flow_type_for_source_state( return Ok(Some(FlowType::Batching)); } - if has_instant_ttl_source_table { - return Ok(Some(FlowType::Streaming)); - } - Ok(None) } +/// The stateless streaming runtime accepts only a single source scan wrapped by +/// projections and filters. Keep this check local to the operator: the flow +/// validator is intentionally private to the flow crate. +fn is_stateless_flow_plan(plan: &LogicalPlan) -> bool { + let mut scans = 0; + let mut supported_nodes = true; + let result = plan.apply_with_subqueries(|node| { + match node { + LogicalPlan::TableScan(_) => scans += 1, + LogicalPlan::Projection(_) | LogicalPlan::Filter(_) => {} + LogicalPlan::Distinct(Distinct::All(_)) => {} + _ => supported_nodes = false, + } + Ok(datafusion_common::tree_node::TreeNodeRecursion::Continue) + }); + + result.is_ok() && supported_nodes && scans == 1 +} + +const INSTANT_TTL_FLOW_QUERY_ERROR: &str = "instant-TTL flow sources support only stateless projection/filter queries over one source; non-stateless queries require a persisted source"; + +const STATELESS_FLOW_QUERY_ERROR: &str = + "flow streaming supports only stateless projection/filter queries over one source"; + impl StatementExecutor { pub fn catalog_manager(&self) -> CatalogManagerRef { self.catalog_manager.clone() @@ -858,6 +879,9 @@ impl StatementExecutor { /// Determines the flow type from source-table state, schedule requirements, /// and SQL shape. + /// + /// Aggregates and persisted-source DISTINCT use batching. Plain DISTINCT is request-local only + /// for instant-TTL sources, where it retains the legacy streaming compatibility behavior. async fn determine_flow_type( &self, expr: &CreateFlowExpr, @@ -941,17 +965,31 @@ impl StatementExecutor { ); let stmt = &stmts[0]; - if is_tql(query_ctx.sql_dialect(), &expr.sql) + let is_tql_query = is_tql(query_ctx.sql_dialect(), &expr.sql) .map_err(BoxedError::new) - .context(ExternalSnafu)? - { + .context(ExternalSnafu)?; + if is_tql_query { + if has_instant_ttl_source_table { + return InvalidSqlSnafu { + err_msg: INSTANT_TTL_FLOW_QUERY_ERROR.to_string(), + } + .fail(); + } return Ok(FlowType::Batching); } - // support tql parse too + // Plan before selecting a mode. In particular, instant-TTL sources + // must not bypass validation of the stateless streaming subset. let plan = match stmt { - // prom ql is only supported in batching mode - Statement::Tql(_) => return Ok(FlowType::Batching), + Statement::Tql(_) => { + if has_instant_ttl_source_table { + return InvalidSqlSnafu { + err_msg: INSTANT_TTL_FLOW_QUERY_ERROR.to_string(), + } + .fail(); + } + return Ok(FlowType::Batching); + } _ => engine .planner() .plan(&QueryStatement::Sql(stmt.clone()), query_ctx) @@ -960,12 +998,15 @@ impl StatementExecutor { .context(ExternalSnafu)?, }; - /// Visitor to find aggregation or distinct - struct FindAggr { + /// Visitor to classify aggregation and DISTINCT separately. Plain DISTINCT is request + /// local only for instant-TTL sources; persisted-source DISTINCT keeps batching semantics. + struct FindQueryShape { is_aggr: bool, + is_distinct: bool, + has_distinct_on: bool, } - impl TreeNodeVisitor<'_> for FindAggr { + impl TreeNodeVisitor<'_> for FindQueryShape { type Node = LogicalPlan; fn f_down( &mut self, @@ -973,25 +1014,58 @@ impl StatementExecutor { ) -> datafusion_common::Result { match node { - LogicalPlan::Aggregate(_) | LogicalPlan::Distinct(_) => { - self.is_aggr = true; - return Ok(datafusion_common::tree_node::TreeNodeRecursion::Stop); - } + LogicalPlan::Aggregate(_) => self.is_aggr = true, + LogicalPlan::Distinct(Distinct::All(_)) => self.is_distinct = true, + LogicalPlan::Distinct(Distinct::On(_)) => self.has_distinct_on = true, _ => (), } Ok(datafusion_common::tree_node::TreeNodeRecursion::Continue) } } - let mut find_aggr = FindAggr { is_aggr: false }; + let mut query_shape = FindQueryShape { + is_aggr: false, + is_distinct: false, + has_distinct_on: false, + }; - plan.visit_with_subqueries(&mut find_aggr) + plan.visit_with_subqueries(&mut query_shape) .context(BuildDfLogicalPlanSnafu)?; - if find_aggr.is_aggr { - Ok(FlowType::Batching) - } else { - Ok(FlowType::Streaming) + if query_shape.has_distinct_on { + return InvalidSqlSnafu { + err_msg: if has_instant_ttl_source_table { + INSTANT_TTL_FLOW_QUERY_ERROR.to_string() + } else { + STATELESS_FLOW_QUERY_ERROR.to_string() + }, + } + .fail(); } + if query_shape.is_aggr { + return if has_instant_ttl_source_table { + InvalidSqlSnafu { + err_msg: INSTANT_TTL_FLOW_QUERY_ERROR.to_string(), + } + .fail() + } else { + Ok(FlowType::Batching) + }; + } + if query_shape.is_distinct && !has_instant_ttl_source_table { + return Ok(FlowType::Batching); + } + + ensure!( + is_stateless_flow_plan(&plan), + InvalidSqlSnafu { + err_msg: if has_instant_ttl_source_table { + INSTANT_TTL_FLOW_QUERY_ERROR.to_string() + } else { + STATELESS_FLOW_QUERY_ERROR.to_string() + }, + } + ); + Ok(FlowType::Streaming) } #[tracing::instrument(skip_all)] @@ -2868,6 +2942,7 @@ mod test { #[cfg(feature = "enterprise")] use api::v1::meta::{ProcedureDetailResponse, ReconcileRequest, ReconcileResponse}; + use arrow::datatypes::{DataType as ArrowDataType, Field as ArrowField}; #[cfg(feature = "enterprise")] use common_meta::cache_invalidator::{CacheInvalidator, CacheInvalidatorRef}; #[cfg(feature = "enterprise")] @@ -2890,6 +2965,9 @@ mod test { use common_meta::rpc::procedure::{ MigrateRegionRequest, MigrateRegionResponse, ProcedureStateResponse, }; + use datafusion::functions_aggregate::expr_fn::count; + use datafusion::logical_expr::builder::LogicalTableSource; + use datafusion::logical_expr::{LogicalPlanBuilder, col}; use session::context::{QueryContext, QueryContextBuilder}; use sql::dialect::GreptimeDbDialect; use sql::parser::{ParseOptions, ParserContext}; @@ -3281,6 +3359,77 @@ SELECT max(c1), min(c2) FROM schema_2.table_2;"; ); } + fn stateless_test_scan(name: &str) -> LogicalPlan { + let schema = arrow::datatypes::Schema::new(vec![ArrowField::new( + "value", + ArrowDataType::Int32, + true, + )]); + LogicalPlanBuilder::scan( + name, + Arc::new(LogicalTableSource::new(Arc::new(schema))), + None, + ) + .unwrap() + .build() + .unwrap() + } + + #[test] + fn test_is_stateless_flow_plan_accepts_scan_projection_and_filter() { + let scan = stateless_test_scan("source"); + assert!(is_stateless_flow_plan(&scan)); + + let projection = LogicalPlanBuilder::from(scan.clone()) + .project(vec![col("value")]) + .unwrap() + .build() + .unwrap(); + assert!(is_stateless_flow_plan(&projection)); + + let filter = LogicalPlanBuilder::from(projection) + .filter(col("value").gt(datafusion_expr::lit(0))) + .unwrap() + .build() + .unwrap(); + assert!(is_stateless_flow_plan(&filter)); + } + + #[test] + fn test_is_stateless_flow_plan_rejects_aggregate_and_multiple_scans() { + let scan = stateless_test_scan("source"); + let aggregate = LogicalPlanBuilder::from(scan.clone()) + .aggregate( + Vec::::new(), + vec![count(col("value"))], + ) + .unwrap() + .build() + .unwrap(); + assert!(!is_stateless_flow_plan(&aggregate)); + + let distinct = LogicalPlanBuilder::from(scan.clone()) + .distinct() + .unwrap() + .build() + .unwrap(); + assert!(is_stateless_flow_plan(&distinct)); + + let distinct_on_aggregate = LogicalPlanBuilder::from(aggregate) + .distinct() + .unwrap() + .build() + .unwrap(); + assert!(!is_stateless_flow_plan(&distinct_on_aggregate)); + + let multiple_scans = LogicalPlanBuilder::from(stateless_test_scan("left")) + .cross_join(stateless_test_scan("right")) + .unwrap() + .build() + .unwrap(); + assert!(!is_stateless_flow_plan(&multiple_scans)); + } + // --- Schedule option tests --- #[test] @@ -3370,11 +3519,11 @@ SELECT max(c1), min(c2) FROM schema_2.table_2;"; } #[test] - fn test_determine_flow_type_for_source_state_instant_ttl_without_missing_sources() { + fn test_determine_flow_type_for_source_state_existing_sources_require_plan() { assert_eq!( determine_flow_type_for_source_state("my_flow", &HashMap::new(), false, true, false) .unwrap(), - Some(FlowType::Streaming) + None ); } diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index 18e1d144024..e875f2d7a01 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -22,7 +22,6 @@ use cache::{ use catalog::information_schema::NoopInformationExtension; use catalog::kvbackend::KvBackendCatalogManagerBuilder; use catalog::process_manager::ProcessManager; -use cmd::error::StartFlownodeSnafu; use common_base::Plugins; use common_catalog::consts::{MIN_USER_FLOW_ID, MIN_USER_TABLE_ID}; use common_config::KvBackendConfig; @@ -55,7 +54,6 @@ use frontend::instance::builder::FrontendBuilder; use frontend::server::Services; use meta_srv::metasrv::{FLOW_ID_SEQ, TABLE_ID_SEQ}; use servers::grpc::GrpcOptions; -use snafu::ResultExt; use standalone::options::StandaloneOptions; use standalone::{StandaloneDatanodeManager, StandaloneRepartitionProcedureFactory}; @@ -312,23 +310,6 @@ impl GreptimeDbStandaloneBuilder { frontend_instance_handler .set_handler(weak_grpc_handler) .await; - - let flow_streaming_engine = flownode.flow_engine().streaming_engine(); - let invoker = flow::FrontendInvoker::build_from( - flow_streaming_engine.clone(), - catalog_manager.clone(), - kv_backend.clone(), - cache_registry.clone(), - procedure_executor.clone(), - node_manager.clone(), - instance.frontend_peer_addr().to_string(), - ) - .await - .context(StartFlownodeSnafu) - .unwrap(); - - flow_streaming_engine.set_frontend_invoker(invoker).await; - procedure_manager.start().await.unwrap(); wal_provider.start().await.unwrap(); diff --git a/tests/cases/standalone/common/flow/flow_advance_ttl.result b/tests/cases/standalone/common/flow/flow_advance_ttl.result index d482d9cbac5..6d87854b961 100644 --- a/tests/cases/standalone/common/flow/flow_advance_ttl.result +++ b/tests/cases/standalone/common/flow/flow_advance_ttl.result @@ -8,7 +8,7 @@ CREATE TABLE distinct_basic ( Affected Rows: 0 --- should fallback to streaming mode when there is no EVAL INTERVAL +-- request-local DISTINCT is supported in streaming mode -- SQLNESS REPLACE id=\d+ id=REDACTED CREATE FLOW test_distinct_basic SINK TO out_distinct_basic AS SELECT @@ -18,8 +18,18 @@ FROM Affected Rows: 0 +-- instant-TTL sources reject non-stateless LIMIT plans +CREATE FLOW test_limit_instant_rejected SINK TO out_limit_instant_rejected AS +SELECT + number +FROM + distinct_basic +LIMIT 1; + +Error: 1004(InvalidArguments), Invalid SQL, error: instant-TTL flow sources support only stateless projection/filter queries over one source; non-stateless queries require a persisted source + -- flow_options should have a flow_type:streaming --- since source table's ttl=instant +-- since source table's ttl=instant and DISTINCT is request-local SELECT flow_name, options FROM INFORMATION_SCHEMA.FLOWS; +---------------------+---------------------------+ @@ -113,9 +123,25 @@ ADMIN FLUSH_TABLE('distinct_basic'); | 0 | +-------------------------------------+ +-- Recover the persisted streaming DISTINCT flow, then replan its first write +-- against an extended source schema without recreating the flow. +-- SQLNESS ARG restart=true +SELECT 1; + ++----------+ +| Int64(1) | ++----------+ +| 1 | ++----------+ + +ALTER TABLE distinct_basic ADD COLUMN extra INT NULL; + +Affected Rows: 0 + INSERT INTO - distinct_basic + distinct_basic (number, ts) VALUES + (23, "2021-07-01 00:00:01.600"), (23, "2021-07-01 00:00:01.600"); Affected Rows: 0 @@ -194,7 +220,7 @@ DROP TABLE distinct_basic; Affected Rows: 0 --- test ttl = 5s +-- test ttl = 5s (DISTINCT remains batching for persisted sources) CREATE TABLE distinct_basic ( "number" INT, ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP, @@ -204,6 +230,22 @@ CREATE TABLE distinct_basic ( Affected Rows: 0 +-- Without a schedule, persisted DISTINCT must reach batching validation, +-- not silently become request-local streaming. +CREATE FLOW test_distinct_persisted_unscheduled +SINK TO out_distinct_persisted_unscheduled AS +SELECT DISTINCT number AS dis FROM distinct_basic; + +Error: 3001(EngineExecuteQuery), Invalid query: SQL batching flow without a time-window expression must specify EVAL INTERVAL to run as an explicit full-query flow + +DROP FLOW IF EXISTS test_distinct_persisted_unscheduled; + +Affected Rows: 0 + +DROP TABLE IF EXISTS out_distinct_persisted_unscheduled; + +Affected Rows: 0 + CREATE FLOW test_distinct_basic SINK TO out_distinct_basic EVAL INTERVAL '1m' AS SELECT DISTINCT number as dis @@ -213,7 +255,7 @@ FROM Affected Rows: 0 -- flow_options should have a flow_type:batching --- since source table's ttl=instant +-- persisted-source DISTINCT retains batching semantics SELECT flow_name, options FROM INFORMATION_SCHEMA.FLOWS; +---------------------+--------------------------+ diff --git a/tests/cases/standalone/common/flow/flow_advance_ttl.sql b/tests/cases/standalone/common/flow/flow_advance_ttl.sql index 4ebb8873a06..9d7e12bf760 100644 --- a/tests/cases/standalone/common/flow/flow_advance_ttl.sql +++ b/tests/cases/standalone/common/flow/flow_advance_ttl.sql @@ -6,7 +6,7 @@ CREATE TABLE distinct_basic ( TIME INDEX(ts) )WITH ('ttl' = 'instant'); --- should fallback to streaming mode when there is no EVAL INTERVAL +-- request-local DISTINCT is supported in streaming mode -- SQLNESS REPLACE id=\d+ id=REDACTED CREATE FLOW test_distinct_basic SINK TO out_distinct_basic AS SELECT @@ -14,8 +14,16 @@ SELECT FROM distinct_basic; +-- instant-TTL sources reject non-stateless LIMIT plans +CREATE FLOW test_limit_instant_rejected SINK TO out_limit_instant_rejected AS +SELECT + number +FROM + distinct_basic +LIMIT 1; + -- flow_options should have a flow_type:streaming --- since source table's ttl=instant +-- since source table's ttl=instant and DISTINCT is request-local SELECT flow_name, options FROM INFORMATION_SCHEMA.FLOWS; SHOW CREATE TABLE distinct_basic; @@ -44,9 +52,17 @@ SELECT number FROM distinct_basic; -- SQLNESS SLEEP 6s ADMIN FLUSH_TABLE('distinct_basic'); +-- Recover the persisted streaming DISTINCT flow, then replan its first write +-- against an extended source schema without recreating the flow. +-- SQLNESS ARG restart=true +SELECT 1; + +ALTER TABLE distinct_basic ADD COLUMN extra INT NULL; + INSERT INTO - distinct_basic + distinct_basic (number, ts) VALUES + (23, "2021-07-01 00:00:01.600"), (23, "2021-07-01 00:00:01.600"); -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | @@ -84,7 +100,7 @@ SELECT count(*) FROM INFORMATION_SCHEMA.FLOWS WHERE flow_name = 'test_distinct_b DROP TABLE distinct_basic; --- test ttl = 5s +-- test ttl = 5s (DISTINCT remains batching for persisted sources) CREATE TABLE distinct_basic ( "number" INT, ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP, @@ -92,6 +108,15 @@ CREATE TABLE distinct_basic ( TIME INDEX(ts) )WITH ('ttl' = '5s'); +-- Without a schedule, persisted DISTINCT must reach batching validation, +-- not silently become request-local streaming. +CREATE FLOW test_distinct_persisted_unscheduled +SINK TO out_distinct_persisted_unscheduled AS +SELECT DISTINCT number AS dis FROM distinct_basic; + +DROP FLOW IF EXISTS test_distinct_persisted_unscheduled; +DROP TABLE IF EXISTS out_distinct_persisted_unscheduled; + CREATE FLOW test_distinct_basic SINK TO out_distinct_basic EVAL INTERVAL '1m' AS SELECT DISTINCT number as dis @@ -99,7 +124,7 @@ FROM distinct_basic; -- flow_options should have a flow_type:batching --- since source table's ttl=instant +-- persisted-source DISTINCT retains batching semantics SELECT flow_name, options FROM INFORMATION_SCHEMA.FLOWS; -- SQLNESS ARG restart=true diff --git a/tests/cases/standalone/common/flow/show_create_flow.result b/tests/cases/standalone/common/flow/show_create_flow.result index 9499618ec60..f529cef1b9f 100644 --- a/tests/cases/standalone/common/flow/show_create_flow.result +++ b/tests/cases/standalone/common/flow/show_create_flow.result @@ -405,12 +405,12 @@ SELECT number FROM out_num_cnt_show; -- should mismatch, hence the old flow remains CREATE OR REPLACE FLOW filter_numbers_show SINK TO out_num_cnt_show AS SELECT number AS n1, number AS n2 FROM numbers_input_show where number > 15; -Error: 3001(EngineExecuteQuery), Invalid query: Column 1(name is 'ts', flow inferred name is 'n2')'s data type mismatch, expect Timestamp(Millisecond(TimestampMillisecondType)) got Int32(Int32Type) +Error: 3001(EngineExecuteQuery), Invalid query: Flow output column 1 has type Int32(Int32Type), but sink column ts has type Timestamp(Millisecond(TimestampMillisecondType)) -- should mismatch, hence the old flow remains CREATE OR REPLACE FLOW filter_numbers_show SINK TO out_num_cnt_show AS SELECT number AS n1, number AS n2, number AS n3 FROM numbers_input_show where number > 15; -Error: 3001(EngineExecuteQuery), Invalid query: Column 1(name is 'ts', flow inferred name is 'n2')'s data type mismatch, expect Timestamp(Millisecond(TimestampMillisecondType)) got Int32(Int32Type) +Error: 3001(EngineExecuteQuery), Invalid query: Flow output has 3 columns, but sink has 2 columns; only zero, one, or two trailing auto columns are supported SELECT flow_definition, source_table_names FROM INFORMATION_SCHEMA.FLOWS WHERE flow_name='filter_numbers_show';