rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-09-25 14:45:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(flow): execute streaming flows with DataFusion (#8976)

Browse Source 
* 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 79e96ac745.

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>
This commit is contained in:
discord9
2026-09-17 06:44:01 +00:00
committed by GitHub
parent 604c88e7e2
commit 7c7132ea65
57 changed files with 3385 additions and 17407 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+5 -143
View File
@@ -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"
src/cmd/src/flownode.rs
+1 -23
View File
@@ -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))
}
src/cmd/src/standalone.rs
+1 -17
View File
@@ -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)?;
src/flow/AGENTS.md
+31 -52
View File
@@ -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.
src/flow/Cargo.toml
-3
View File
@@ -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"
src/flow/src/adapter.rs
+1042 -769
View File
File diff suppressed because it is too large Load Diff
src/flow/src/adapter/flownode_impl.rs
+14 -177
View File
@@ -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<bool, Error> {
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<impl IntoIterator<Item = FlowId>, Error> {
Ok(self
.flow_err_collectors
.read()
.await
.keys()
.cloned()
.collect::<Vec<_>>())
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::<Result<Vec<_>, _>>()?;
let name_to_col = HashMap::<_, _>::from_iter(
insert_schema
.iter()
.enumerate()
.map(|(i, name)| (&name.column_name, i)),
);
let fetch_order: Vec<FetchFromRow> = 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<DiffRow> = 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;
src/flow/src/adapter/node_context.rs
-458
View File
@@ -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<TableId, BTreeSet<FlowId>>,
/// 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<FlowId, TableName>,
pub flow_plans: BTreeMap<FlowId, TypedPlan>,
pub sink_to_flow: BTreeMap<TableName, FlowId>,
/// 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<TableId, SourceSender>,
/// 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<TableName, (mpsc::UnboundedSender<Batch>, mpsc::UnboundedReceiver<Batch>)>,
/// can query the schema of the table source, from metasrv with local cache
pub table_source: Box<dyn FlowTableSource>,
/// All the tables that have been registered in the worker
pub table_repr: IdToNameMap,
pub query_context: Option<Arc<QueryContext>>,
}
impl FlownodeContext {
pub fn new(table_source: Box<dyn FlowTableSource>) -> 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<FlowId>> {
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<DiffRow>?
sender: broadcast::Sender<Batch>,
send_buf_tx: mpsc::Sender<Batch>,
send_buf_rx: RwLock<mpsc::Receiver<Batch>>,
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<Batch> {
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<usize, Error> {
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<DiffRow>,
batch_datatypes: &[ConcreteDataType],
) -> Result<usize, Error> {
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<usize, Error> {
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<DiffRow>,
batch_datatypes: &[ConcreteDataType],
) -> Result<usize, Error> {
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<usize, Error> {
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<usize, Error> {
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<TypedPlan> {
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<mpsc::UnboundedSender<Batch>, 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<TableName>,
table_id: Option<TableId>,
) -> Result<GlobalId, Error> {
// 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<TableName, GlobalId>,
id_to_global_id: HashMap<TableId, GlobalId>,
global_id_to_name_id: BTreeMap<GlobalId, (Option<TableName>, Option<TableId>)>,
}
impl IdToNameMap {
pub fn new() -> Self {
Default::default()
}
pub fn insert(&mut self, name: Option<TableName>, id: Option<TableId>, 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<TableId>, 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<TableName>, 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<TableName>, Option<TableId>)> {
self.global_id_to_name_id.get(global_id).cloned()
}
}
src/flow/src/adapter/parse_expr.rs
-245
View File
@@ -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<Expr>,
op: String,
right: Box<Expr>,
},
}
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::<repr::Duration>().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))
}
src/flow/src/adapter/refill.rs
-440
View File
@@ -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<Vec<RefillTask>, 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<RefillTask>,
) -> 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<T> {
/// Task is not started
Prepared { sql: String },
/// Task is running
Running {
handle: JoinHandle<Result<T, Error>>,
},
/// Task is finished
Finished { res: Result<T, Error> },
}
impl<T> TaskState<T> {
fn new(sql: String) -> Self {
Self::Prepared { sql }
}
}
mod test_send {
use std::collections::BTreeMap;
use tokio::sync::RwLock;
use super::*;
fn is_send<T: Send + Sync>() {}
fn foo() {
is_send::<TaskState<()>>();
is_send::<RefillTask>();
is_send::<BTreeMap<FlowId, RefillTask>>();
is_send::<RwLock<BTreeMap<FlowId, RefillTask>>>();
}
}
impl TaskState<()> {
/// check if task is finished
async fn is_finished(&mut self) -> Result<bool, Error> {
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<Result<(), Error>> = 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<common_query::Output> for QueryStream {
type Error = Error;
fn try_from(value: common_query::Output) -> Result<Self, Self::Error> {
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<SendableRecordBatchStream, Error> {
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<RefillTask, Error> {
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<bool, Error> {
self.state.is_finished().await
}
}
src/flow/src/adapter/stat.rs
-47
View File
@@ -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,
}
}
}
src/flow/src/adapter/stateless.rs
+730
View File
@@ -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<ColumnSchema>,
pub(crate) sink_primary_keys: Vec<String>,
/// The exact trailing columns resolved when the flow was created.
pub(crate) auto_columns: Vec<ColumnSchema>,
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<dyn TableProvider>) -> 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<Arc<dyn TableProvider>, 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<LogicalPlan, Error> {
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(&timestamp_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::<Vec<_>>();
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<dyn TableProvider>,
) -> Result<LogicalPlan, Error> {
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::<DfTableProviderAdapter>()
.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<Vec<Value>, 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<dyn QueryEngine>,
frontend_client: &Arc<FrontendClient>,
current_source_schema_version: u32,
) -> Result<usize, Error> {
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::<Vec<_>>();
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::<Vec<_>>();
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<dyn TableProvider> {
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<dyn TableProvider>| {
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<dyn TableProvider>| {
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::<usize>()
}
};
assert_eq!(run(input(&[1, 1, 2])).await, 2);
assert_eq!(run(input(&[1, 1])).await, 1);
}
}
src/flow/src/adapter/tests.rs
+1067 -15
View File
File diff suppressed because it is too large Load Diff
src/flow/src/adapter/util.rs
+7 -54
View File
@@ -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(
src/flow/src/adapter/worker.rs
-606
View File
@@ -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<Mutex<VecDeque<DiffRow>>>;
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", &"<Dfir>")
.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<Option<FlowId>, 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<bool, Error> {
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<bool, Error> {
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<FlowId, usize>,
BTreeMap<FlowId, i64>,
BTreeMap<FlowId, i64>,
),
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<FlowId, ActiveDataflowState<'subgraph>>,
itc_server: Arc<Mutex<InterThreadCallServer>>,
}
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<Batch>,
source_ids: &[GlobalId],
src_recvs: Vec<broadcast::Receiver<Batch>>,
// TODO(discord9): set expire duration for all arrangement and compare to sys timestamp instead
expire_after: Option<repr::Duration>,
or_replace: bool,
create_if_not_exists: bool,
err_collector: ErrCollector,
) -> Result<Option<FlowId>, 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<Option<Response>, ()> {
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<Batch>,
source_ids: Vec<GlobalId>,
src_recvs: Vec<broadcast::Receiver<Batch>>,
expire_after: Option<repr::Duration>,
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<Option<FlowId>, Error>,
// TODO(discord9): add flow err_collector
},
Remove {
result: bool,
},
ContainTask {
result: bool,
},
RunAvail,
QueryFullFlowStat {
state_size: BTreeMap<FlowId, usize>,
last_exec_time_map: BTreeMap<FlowId, i64>,
start_time_map: BTreeMap<FlowId, i64>,
},
}
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<oneshot::Sender<Response>>)>,
}
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<Response, Error> {
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<oneshot::Sender<Response>>)>,
}
impl InterThreadCallServer {
pub async fn recv(&mut self) -> Option<(Request, Option<oneshot::Sender<Response>>)> {
self.arg_recv.recv().await
}
pub fn blocking_recv(&mut self) -> Option<(Request, Option<oneshot::Sender<Response>>)> {
self.arg_recv.blocking_recv()
}
}
fn from_send_error<T>(err: mpsc::error::SendError<T>) -> 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::<Batch>(1024);
let (sink_tx, mut sink_rx) = mpsc::unbounded_channel::<Batch>();
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();
}
}
src/flow/src/batching_mode/frontend_client.rs
+118
View File
@@ -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<PeerDesc>,
) -> Result<u32, Error> {
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<AtomicUsize>,
}
#[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<Output, BoxedError> {
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<dyn GrpcQueryHandlerWithBoxedError> = 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<dyn GrpcQueryHandlerWithBoxedError> = Arc::new(MetricsHandler);
src/flow/src/batching_mode/time_window.rs
+4 -4
View File
@@ -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();
src/flow/src/compute.rs
-23
View File
@@ -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;
src/flow/src/compute/render.rs
-527
View File
@@ -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<GlobalId, CollectionBundle>,
/// used by `Get`/`Let` Plan for getting/setting local variables
///
/// TODO(discord9): consider if use Vec<(LocalId, CollectionBundle)> instead
pub local_scope: Vec<BTreeMap<LocalId, CollectionBundle>>,
/// 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<GlobalId, CollectionBundle<Batch>>,
/// used by `Get`/`Let` Plan for getting/setting local variables
///
/// TODO(discord9): consider if use Vec<(LocalId, CollectionBundle)> instead
pub local_scope_batch: Vec<BTreeMap<LocalId, CollectionBundle<Batch>>>,
// 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<Batch>) {
self.input_collection_batch.insert(id, collection);
}
pub fn insert_local_batch(&mut self, id: LocalId, collection: CollectionBundle<Batch>) {
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<CollectionBundle<Batch>, 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<CollectionBundle, Error> {
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<DiffRow>,
output_type: &RelationType,
) -> CollectionBundle<Batch> {
let (send_port, recv_port) = self.df.make_edge::<_, Toff<Batch>>("constant_batch");
let mut per_time: BTreeMap<repr::Timestamp, Vec<DiffRow>> = 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<DiffRow>) -> CollectionBundle {
let (send_port, recv_port) = self.df.make_edge::<_, Toff>("constant");
let mut per_time: BTreeMap<repr::Timestamp, Vec<DiffRow>> = 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<CollectionBundle<Batch>, 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<CollectionBundle, Error> {
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<TypedPlan>,
body: Box<TypedPlan>,
) -> Result<CollectionBundle<Batch>, 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<TypedPlan>,
body: Box<TypedPlan>,
) -> Result<CollectionBundle, Error> {
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<SEND, T>,
}
#[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<i64>,
expected: BTreeMap<i64, Vec<DiffRow>>,
output: Rc<RefCell<Vec<DiffRow>>>,
) {
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<RefCell<Vec<DiffRow>>> {
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::<usize>();
},
);
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<i32>>("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::<i32>();
});
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<i32>>("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::<i32>();
});
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);
}
}
src/flow/src/compute/render/map.rs
-425
View File
@@ -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<TypedPlan>,
mfp: MapFilterProject,
_output_type: &RelationType,
) -> Result<CollectionBundle<Batch>, Error> {
let input = self.render_plan_batch(*input)?;
let (out_send_port, out_recv_port) = self.df.make_edge::<_, Toff<Batch>>("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::<Batch>::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<ScalarExpr>` as key due to `Value` have `bytes` variant
#[allow(clippy::mutable_key_type)]
pub fn render_mfp(
&mut self,
input: Box<TypedPlan>,
mfp: MapFilterProject,
) -> Result<CollectionBundle, Error> {
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<Item = DiffRow>,
mfp_plan: &MfpPlan,
now: repr::Timestamp,
err_collector: &ErrCollector,
scheduler: &Scheduler,
send: &PortCtx<SEND, Toff>,
) {
// 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<Item = DiffRow>,
mfp_plan: &MfpPlan,
now: repr::Timestamp,
err_collector: &ErrCollector,
) -> Vec<KeyValDiffRow> {
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::<EvalError>(&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);
}
}
src/flow/src/compute/render/reduce.rs
-1985
View File
File diff suppressed because it is too large Load Diff
src/flow/src/compute/render/src_sink.rs
-245
View File
@@ -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<Batch>,
) -> Result<CollectionBundle<Batch>, Error> {
debug!("Rendering Source Batch");
let (send_port, recv_port) = self.df.make_edge::<_, Toff<Batch>>("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::<Batch>::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<DiffRow>,
) -> Result<CollectionBundle, Error> {
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<Batch>,
sender: mpsc::UnboundedSender<Batch>,
) {
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<DiffRow>,
) {
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::<usize>()
);
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);
}
},
);
}
}
src/flow/src/compute/state.rs
-167
View File
@@ -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<RefCell<BTreeMap<Timestamp, VecDeque<SubgraphId>>>>,
/// 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<RefCell<Timestamp>>,
/// 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<ArrangeHandler>,
/// the time arrangement need to be expired after a certain time in milliseconds
expire_after: Option<Timestamp>,
/// the last time each subgraph executed
last_exec_time: Option<Timestamp>,
/// the time the flow first executed, in unix timestamp milliseconds
start_time: Option<Timestamp>,
}
impl DataflowState {
pub fn new_arrange(&mut self, name: Option<Vec<String>>) -> 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<RefCell<Timestamp>> {
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<repr::Duration>) {
self.expire_after = after;
}
pub fn expire_after(&self) -> Option<Timestamp> {
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<Timestamp> {
self.last_exec_time
}
/// Returns the time the flow first executed, in unix timestamp milliseconds.
pub fn start_time(&self) -> Option<Timestamp> {
self.start_time
}
}
#[derive(Debug, Clone)]
pub struct Scheduler {
// this scheduler is shared with `DataflowState`, so it can schedule subgraph
schedule_subgraph: Rc<RefCell<BTreeMap<Timestamp, VecDeque<SubgraphId>>>>,
cur_subgraph: Rc<RefCell<Option<SubgraphId>>>,
}
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));
}
}
src/flow/src/compute/types.rs
-208
View File
@@ -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<T = DiffRow> = TeeingHandoff<T>;
/// A collection, represent a collections of data that is received from a handoff.
pub struct Collection<T: 'static> {
/// represent a stream of updates recv from this port
stream: RecvPort<TeeingHandoff<T>>,
}
impl<T: 'static + Clone> Collection<T> {
pub fn from_port(port: RecvPort<TeeingHandoff<T>>) -> 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<TeeingHandoff<T>> {
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<RefCell<Option<SubgraphId>>>,
/// maintain a list of readers for the arrangement for the ease of scheduling
pub readers: Rc<RefCell<Vec<SubgraphId>>>,
}
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> {
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<T: 'static = DiffRow> {
/// 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<T>,
/// 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<ScalarExpr>`
/// There is a false positive in using `Vec<ScalarExpr>` as key due to `ScalarExpr::Literal`
/// contain a `Value` which have `bytes` variant
#[allow(clippy::mutable_key_type)]
pub arranged: BTreeMap<Vec<ScalarExpr>, Arranged>,
}
pub trait GenericBundle {
fn is_batch(&self) -> bool;
fn try_as_batch(&self) -> Option<&CollectionBundle<Batch>> {
None
}
fn try_as_row(&self) -> Option<&CollectionBundle<DiffRow>> {
None
}
}
impl GenericBundle for CollectionBundle<Batch> {
fn is_batch(&self) -> bool {
true
}
fn try_as_batch(&self) -> Option<&CollectionBundle<Batch>> {
Some(self)
}
}
impl GenericBundle for CollectionBundle<DiffRow> {
fn is_batch(&self) -> bool {
false
}
fn try_as_row(&self) -> Option<&CollectionBundle<DiffRow>> {
Some(self)
}
}
impl<T: 'static> CollectionBundle<T> {
pub fn from_collection(collection: Collection<T>) -> Self {
Self {
collection,
arranged: BTreeMap::default(),
}
}
}
impl<T: 'static + Clone> CollectionBundle<T> {
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<Mutex<VecDeque<EvalError>>>,
}
impl ErrCollector {
pub fn get_all_blocking(&self) -> Vec<EvalError> {
self.inner.blocking_lock().drain(..).collect_vec()
}
pub async fn get_all(&self) -> Vec<EvalError> {
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<F, R>(&self, f: F) -> Option<R>
where
F: FnOnce() -> Result<R, EvalError>,
{
match f() {
Ok(r) => Some(r),
Err(e) => {
self.push_err(e);
None
}
}
}
}
src/flow/src/df_optimizer.rs
+2 -44
View File
@@ -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<dyn QueryEngine>,
sql: &str,
) -> Result<TypedPlan, Error> {
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 {}
src/flow/src/error.rs
+3 -35
View File
@@ -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<T> = std::result::Result<T, Error>;
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<EvalError> for Error {
fn from(e: EvalError) -> Self {
Err::<(), _>(e).context(EvalSnafu).unwrap_err()
}
}
src/flow/src/expr.rs
+27 -322
View File
@@ -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<VectorRef>,
row_count: usize,
/// describe if corresponding rows in batch is insert or delete, None means all rows are insert
diffs: Option<VectorRef>,
}
impl TryFrom<RecordBatch> for Batch {
type Error = Error;
fn try_from(value: RecordBatch) -> Result<Self, Self::Error> {
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
&& <dyn arrow::array::Array>::eq(&left.to_arrow_array(), &right.to_arrow_array());
}
let diff_eq = match (&self.diffs, &other.diffs) {
(Some(left), Some(right)) => {
<dyn arrow::array::Array>::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<crate::repr::Row>,
batch_datatypes: &[ConcreteDataType],
) -> Result<Self, EvalError> {
rows: Vec<Row>,
types: &[ConcreteDataType],
) -> Result<Self, error::EvalError> {
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<VectorRef>, row_count: usize) -> Result<Self, EvalError> {
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<VectorRef>, 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<VectorRef> {
&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<Vec<Value>, 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<Batch, EvalError> {
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<Self, EvalError> {
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<VectorRef>,
}
impl From<VectorRef> 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<VectorRef>) -> Result<Self, EvalError> {
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<VectorRef>,
idx: usize,
}
impl std::iter::Iterator for VectorDiffIter {
type Item = (Value, Diff);
fn next(&mut self) -> Option<Self::Item> {
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))
}
}
src/flow/src/expr/df_func.rs
-300
View File
@@ -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<dyn PhysicalExpr>,
/// The input schema of the function
pub(crate) df_schema: Arc<datafusion_common::DFSchema>,
}
impl DfScalarFunction {
pub fn new(raw_fn: RawDfScalarFn, fn_impl: Arc<dyn PhysicalExpr>) -> Result<Self, Error> {
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<Self, Error> {
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<VectorRef, EvalError> {
let row_count = batch.row_count();
let batch: Vec<_> = exprs
.iter()
.map(|expr| expr.eval_batch(batch))
.collect::<Result<_, _>>()?;
let schema = self.df_schema.inner().clone();
let arrays = batch
.iter()
.map(|array| array.to_arrow_array())
.collect::<Vec<_>>();
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<Vec<Value>, EvalError> {
exprs
.iter()
.map(|expr| expr.eval(values))
.collect::<Result<_, _>>()
}
// TODO(discord9): add RecordBatch support
pub fn eval(&self, values: &[Value], exprs: &[ScalarExpr]) -> Result<Value, EvalError> {
// 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<Self, Error> {
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<Arc<dyn PhysicalExpr>, 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<dyn PhysicalExpr> 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<std::cmp::Ordering> {
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<H: std::hash::Hasher>(&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)
);
}
}
src/flow/src/expr/error.rs
+4 -28
View File
@@ -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
}
}
}
src/flow/src/expr/func.rs
-1467
View File
File diff suppressed because it is too large Load Diff
src/flow/src/expr/id.rs
-43
View File
@@ -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),
}
src/flow/src/expr/linear.rs
-1174
View File
File diff suppressed because it is too large Load Diff
src/flow/src/expr/relation.rs
-36
View File
@@ -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,
}
src/flow/src/expr/relation/accum.rs
-1052
View File
File diff suppressed because it is too large Load Diff
src/flow/src/expr/relation/func.rs
-303
View File
@@ -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<A, I>(
&self,
accum: A,
value_diffs: I,
) -> Result<(Value, Vec<Value>), EvalError>
where
A: IntoIterator<Item = Value>,
I: IntoIterator<Item = (Value, Diff)>,
{
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<A>(
&self,
accum: A,
vector: VectorRef,
diff: Option<VectorRef>,
) -> Result<(Value, Vec<Value>), EvalError>
where
A: IntoIterator<Item = Value>,
{
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<HashMap<(GenericFn, ConcreteDataType), AggregateFunc>> =
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<ConcreteDataType>,
) -> Result<Self, Error> {
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)
])
}
}
src/flow/src/expr/scalar.rs
-877
View File
@@ -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<ScalarExpr>,
},
CallBinary {
func: BinaryFunc,
expr1: Box<ScalarExpr>,
expr2: Box<ScalarExpr>,
},
CallVariadic {
func: VariadicFunc,
exprs: Vec<ScalarExpr>,
},
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<ScalarExpr>,
},
/// 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<ScalarExpr>,
then: Box<ScalarExpr>,
els: Box<ScalarExpr>,
},
InList {
expr: Box<ScalarExpr>,
list: Vec<ScalarExpr>,
},
}
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<ColumnType, Error> {
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<VectorRef, EvalError> {
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<VectorRef, EvalError> {
let eval_list = list
.iter()
.map(|e| e.eval_batch(batch))
.collect::<Result<Vec<_>, _>>()?;
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<BooleanArray, DataFusionError> {
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<VectorRef, EvalError> {
let conds = cond.eval_batch(batch)?;
let bool_conds = conds
.as_any()
.downcast_ref::<BooleanVector>()
.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<Value, EvalError> {
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::<Result<Vec<_>, _>>()?;
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<usize, usize>) -> 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<usize> {
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<usize> {
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<Value> {
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<F>(&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<F>(&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<F>(&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<F>(&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<Self>, Option<Self>), 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<Option<i32>> = 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);
}
}
}
src/flow/src/expr/signature.rs
-70
View File
@@ -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,
}
src/flow/src/expr/utils.rs
-348
View File
@@ -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<Option<common_time::Timestamp>> {
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<Option<common_time::Timestamp>> {
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<common_time::Timestamp> {
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);
}
}
}
src/flow/src/lib.rs
+2 -9
View File
@@ -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;
src/flow/src/plan.rs
-270
View File
@@ -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<Self, Error> {
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<TypedExpr>) -> Result<Self, Error> {
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<Self, Error> {
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<DiffRow> },
/// 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<TypedPlan>,
body: Box<TypedPlan>,
},
/// Map, Filter, and Project operators. Chained together.
Mfp {
/// The input collection.
input: Box<TypedPlan>,
/// Linear operator to apply to each record.
mfp: MapFilterProject,
},
/// Reduce operator, aggregation by key assembled from KeyValPlan
Reduce {
/// The input collection.
input: Box<TypedPlan>,
/// 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<TypedPlan>,
/// 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<TypedPlan>,
/// 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<ScalarExpr> {
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<GlobalId> {
fn recur_find_use(plan: &Plan, used: &mut BTreeSet<GlobalId>) {
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
}
}
src/flow/src/plan/join.rs
-76
View File
@@ -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<Vec<ScalarExpr>>,
/// 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<Vec<ScalarExpr>>,
/// An initial closure to apply before any stages.
///
/// Values of `None` indicate the identity closure.
pub initial_closure: Option<JoinFilter>,
/// A *sequence* of stages to apply one after the other.
pub stage_plans: Vec<LinearStagePlan>,
/// A concluding filter to apply after the last stage.
///
/// Values of `None` indicate the identity closure.
pub final_closure: Option<JoinFilter>,
}
/// 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<ScalarExpr>,
/// 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<usize>,
/// The key expressions to use for the lookup relation.
pub lookup_key: Vec<ScalarExpr>,
/// The closure to apply to the concatenation of the key columns,
/// the stream value columns, and the lookup value columns.
pub closure: JoinFilter,
}
src/flow/src/plan/reduce.rs
-87
View File
@@ -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<ScalarExpr> {
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<AggregateExpr>,
/// 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<AggrWithIndex>,
/// Same as `simple_aggrs` but for all of the `DISTINCT` accumulable aggregations.
pub distinct_aggrs: Vec<AggrWithIndex>,
}
/// 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,
}
}
}
src/flow/src/repr.rs
+1 -1
View File
@@ -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
src/flow/src/repr/relation.rs
+2 -95
View File
@@ -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<Self> {
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::<Option<Vec<_>>>()
// 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<Self> {
// 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 {
src/flow/src/server.rs
+59 -363
View File
@@ -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<broadcast::Sender<()>>,
/// server shutdown signal for shutdown grpc server
server_shutdown_tx: Mutex<broadcast::Sender<()>>,
/// streaming task handler
streaming_task_handler: Mutex<Option<JoinHandle<()>>>,
/// state report task handler
state_report_task_handler: Mutex<Option<JoinHandle<()>>>,
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<HeartbeatTask>,
state_report_task: Option<common_runtime::JoinHandle<()>>,
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<dyn QueryEngine>,
) -> Result<StreamingEngine, Error> {
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<dyn query::QueryEngine>) -> 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<Inserter>,
deleter: Arc<Deleter>,
statement_executor: Arc<StatementExecutor>,
}
impl FrontendInvoker {
pub fn new(
inserter: Arc<Inserter>,
deleter: Arc<Deleter>,
statement_executor: Arc<StatementExecutor>,
) -> 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<FrontendInvoker, Error> {
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<Output> {
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<Output> {
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<StatementExecutor> {
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());
}
}
src/flow/src/test_utils.rs
-66
View File
@@ -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<dyn QueryEngine> {
let catalog_list = catalog::memory::new_memory_catalog_manager().unwrap();
let req = RegisterTableRequest {
@@ -158,28 +113,7 @@ pub fn create_test_query_engine() -> Arc<dyn QueryEngine> {
);
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<dyn QueryEngine>, 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()
}
src/flow/src/transform.rs
-318
View File
@@ -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<u32, String>,
}
impl FunctionExtensions {
pub fn from_iter(inner: impl IntoIterator<Item = (u32, impl ToString)>) -> 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<Self, Error> {
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<dyn QueryEngine>) {
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<DataType> {
Ok(DataType::Timestamp(TimeUnit::Millisecond, None))
}
fn signature(&self) -> &Signature {
&self.signature
}
fn invoke_with_args(&self, _: ScalarFunctionArgs) -> datafusion_common::Result<ColumnarValue> {
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<dyn QueryEngine> {
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(<u32 as Scalar>::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<dyn QueryEngine>, 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());
}
}
src/flow/src/transform/aggr.rs
-803
View File
@@ -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<Vec<TypedExpr>, 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<dyn Iterator<Item = &proto::Expression> + 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<Vec<AggregateExpr>, 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<TypedExpr>,
order_by: &Option<Vec<TypedExpr>>,
distinct: bool,
) -> Result<Vec<AggregateExpr>, 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<KeyValPlan, Error> {
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<usize> {
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:
///
/// <group_exprs>..<aggr_exprs>
#[async_recursion::async_recursion]
pub async fn from_substrait_agg_rel(
ctx: &mut FlownodeContext,
agg: &proto::AggregateRel,
extensions: &FunctionExtensions,
) -> Result<TypedPlan, Error> {
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);
}
}
src/flow/src/transform/expr.rs
-839
View File
@@ -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<CDT, Error> {
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<Arc<dyn PhysicalExpr>, 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<ScalarFunction, Error> {
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(<literal>)` 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<TypedExpr>,
extensions: &FunctionExtensions,
) -> Result<TypedExpr, Error> {
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<TypedExpr, Error> {
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<TypedExpr> = {
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<TypedExpr, Error> {
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<Item = (TypedExpr, TypedExpr)>,
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<TypedExpr, Error> {
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,
},
}
);
}
}
src/flow/src/transform/literal.rs
-426
View File
@@ -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<i32> for TimestampPrecision {
type Error = Error;
fn try_from(prec: i32) -> Result<Self, Self::Error> {
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<Literal, Error> {
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<T: FromBytes>(i: &Bytes) -> Result<T, Error>
where
for<'a> &'a <T as num_traits::FromBytes>::Bytes:
std::convert::TryFrom<&'a [u8], Error = TryFromSliceError>,
{
let (int_bytes, _rest) = i.split_at(std::mem::size_of::<T>());
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::<T>(),
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<CDT, Error> {
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);
}
}
src/flow/src/transform/plan.rs
-276
View File
@@ -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<TypedPlan, Error> {
// 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<TypedPlan, Error> {
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<usize> =
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<TypedExpr> = 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<TypedPlan, Error> {
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<TypedPlan, Error> {
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<usize> = 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<TypedPlan, Error> {
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);
}
}
src/flow/src/utils.rs
+12 -990
View File
File diff suppressed because it is too large Load Diff
src/operator/src/statement/ddl.rs
+174 -25
View File
@@ -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<datafusion_common::tree_node::TreeNodeRecursion>
{
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::<datafusion_expr::Expr>::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
);
}
tests-integration/src/standalone.rs
-19
View File
@@ -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();
tests/cases/standalone/common/flow/flow_advance_ttl.result
+47 -5
View File
@@ -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;
+---------------------+--------------------------+
tests/cases/standalone/common/flow/flow_advance_ttl.sql
+30 -5
View File
@@ -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
tests/cases/standalone/common/flow/show_create_flow.result
+2 -2
View File
@@ -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';