From b8933e88d0a9b27b42b1de8fc1e425dfec4b2c74 Mon Sep 17 00:00:00 2001 From: luoluoyuyu Date: Tue, 12 May 2026 00:31:47 +0800 Subject: [PATCH 1/7] update --- Cargo.lock | 67 +++++++++++++++ Cargo.toml | 13 +++ src/catalog/Cargo.toml | 11 +++ src/catalog/src/error.rs | 4 + src/catalog/src/lib.rs | 9 +++ src/catalog/src/meta_store.rs | 23 ++++++ src/catalog/src/stream_catalog.rs | 69 ++++++++++++++++ src/catalog_storage/Cargo.toml | 14 ++++ src/catalog_storage/src/lib.rs | 7 ++ src/catalog_storage/src/memory.rs | 63 +++++++++++++++ src/catalog_storage/src/rocksdb.rs | 81 +++++++++++++++++++ src/common/Cargo.toml | 8 ++ src/common/{ => src}/fs_schema.rs | 0 src/common/{mod.rs => src/legacy_mod.rs} | 0 src/common/src/lib.rs | 3 + src/config/Cargo.toml | 17 ++++ src/config/{ => src}/global_config.rs | 0 src/config/{ => src/legacy}/mod.rs | 0 src/config/src/lib.rs | 33 ++++++++ src/config/{ => src}/loader.rs | 0 src/config/{ => src}/log_config.rs | 0 src/config/{ => src}/paths.rs | 0 src/config/{ => src}/python_config.rs | 0 src/config/{ => src}/service_config.rs | 0 src/config/{ => src}/storage.rs | 0 src/config/{ => src}/streaming_job.rs | 0 src/config/{ => src}/system.rs | 0 src/config/{ => src}/wasm_config.rs | 0 src/coordinator/Cargo.toml | 8 ++ .../{ => src/legacy}/analyze/analysis.rs | 0 .../{ => src/legacy}/analyze/analyzer.rs | 0 .../{ => src/legacy}/analyze/mod.rs | 0 .../{ => src/legacy}/coordinator.rs | 0 .../{ => src/legacy}/dataset/data_set.rs | 0 .../legacy}/dataset/execute_result.rs | 0 .../{ => src/legacy}/dataset/mod.rs | 0 .../dataset/show_catalog_tables_result.rs | 0 .../show_create_streaming_table_result.rs | 0 .../dataset/show_create_table_result.rs | 0 .../legacy}/dataset/show_functions_result.rs | 0 .../dataset/show_streaming_tables_result.rs | 0 .../{ => src/legacy}/execution/executor.rs | 0 .../{ => src/legacy}/execution/mod.rs | 0 .../{ => src/legacy}/execution_context.rs | 0 src/coordinator/{ => src/legacy}/mod.rs | 0 .../{ => src/legacy}/plan/ast_utils.rs | 0 .../legacy}/plan/compile_error_plan.rs | 0 .../legacy}/plan/create_function_plan.rs | 0 .../plan/create_python_function_plan.rs | 0 .../legacy}/plan/create_table_plan.rs | 0 .../{ => src/legacy}/plan/ddl_compiler.rs | 0 .../legacy}/plan/drop_function_plan.rs | 0 .../legacy}/plan/drop_streaming_table_plan.rs | 0 .../{ => src/legacy}/plan/drop_table_plan.rs | 0 .../legacy}/plan/logical_plan_visitor.rs | 0 .../legacy}/plan/lookup_table_plan.rs | 0 src/coordinator/{ => src/legacy}/plan/mod.rs | 0 .../{ => src/legacy}/plan/optimizer.rs | 0 .../legacy}/plan/show_catalog_tables_plan.rs | 0 .../plan/show_create_streaming_table_plan.rs | 0 .../legacy}/plan/show_create_table_plan.rs | 0 .../legacy}/plan/show_functions_plan.rs | 0 .../plan/show_streaming_tables_plan.rs | 0 .../legacy}/plan/start_function_plan.rs | 0 .../legacy}/plan/stop_function_plan.rs | 0 .../legacy}/plan/streaming_compiler.rs | 0 .../plan/streaming_table_connector_plan.rs | 0 .../legacy}/plan/streaming_table_plan.rs | 0 .../{ => src/legacy}/plan/visitor.rs | 0 .../{ => src/legacy}/runtime_context.rs | 0 .../legacy}/statement/create_function.rs | 0 .../statement/create_python_function.rs | 0 .../legacy}/statement/create_table.rs | 0 .../legacy}/statement/drop_function.rs | 0 .../legacy}/statement/drop_streaming_table.rs | 0 .../{ => src/legacy}/statement/drop_table.rs | 0 .../{ => src/legacy}/statement/mod.rs | 0 .../legacy}/statement/show_catalog_tables.rs | 0 .../statement/show_create_streaming_table.rs | 0 .../legacy}/statement/show_create_table.rs | 0 .../legacy}/statement/show_functions.rs | 0 .../statement/show_streaming_tables.rs | 0 .../legacy}/statement/start_function.rs | 0 .../legacy}/statement/stop_function.rs | 0 .../legacy}/statement/streaming_table.rs | 0 .../{ => src/legacy}/statement/visitor.rs | 0 .../legacy}/streaming_table_options.rs | 0 src/coordinator/{ => src/legacy}/tool/mod.rs | 0 src/coordinator/src/lib.rs | 3 + src/job_manager/Cargo.toml | 8 ++ src/job_manager/src/lib.rs | 3 + src/lib.rs | 6 +- src/logger/Cargo.toml | 15 ++++ src/logger/src/lib.rs | 9 +++ src/{logging/mod.rs => logger/src/logging.rs} | 0 src/main.rs | 8 +- src/servicer/Cargo.toml | 8 ++ .../src/legacy}/handler.rs | 0 .../src/legacy}/initializer.rs | 0 .../src/legacy}/memory_service.rs | 0 src/{server => servicer/src/legacy}/mod.rs | 0 .../src/legacy}/service.rs | 0 src/servicer/src/lib.rs | 3 + src/sqlparser/Cargo.toml | 8 ++ src/sqlparser/src/lib.rs | 3 + src/streaming_runtime/Cargo.toml | 8 ++ src/streaming_runtime/src/lib.rs | 3 + src/wasm_runtime/Cargo.toml | 8 ++ src/wasm_runtime/src/lib.rs | 3 + 109 files changed, 521 insertions(+), 5 deletions(-) create mode 100644 src/catalog/Cargo.toml create mode 100644 src/catalog/src/error.rs create mode 100644 src/catalog/src/lib.rs create mode 100644 src/catalog/src/meta_store.rs create mode 100644 src/catalog/src/stream_catalog.rs create mode 100644 src/catalog_storage/Cargo.toml create mode 100644 src/catalog_storage/src/lib.rs create mode 100644 src/catalog_storage/src/memory.rs create mode 100644 src/catalog_storage/src/rocksdb.rs create mode 100644 src/common/Cargo.toml rename src/common/{ => src}/fs_schema.rs (100%) rename src/common/{mod.rs => src/legacy_mod.rs} (100%) create mode 100644 src/common/src/lib.rs create mode 100644 src/config/Cargo.toml rename src/config/{ => src}/global_config.rs (100%) rename src/config/{ => src/legacy}/mod.rs (100%) create mode 100644 src/config/src/lib.rs rename src/config/{ => src}/loader.rs (100%) rename src/config/{ => src}/log_config.rs (100%) rename src/config/{ => src}/paths.rs (100%) rename src/config/{ => src}/python_config.rs (100%) rename src/config/{ => src}/service_config.rs (100%) rename src/config/{ => src}/storage.rs (100%) rename src/config/{ => src}/streaming_job.rs (100%) rename src/config/{ => src}/system.rs (100%) rename src/config/{ => src}/wasm_config.rs (100%) create mode 100644 src/coordinator/Cargo.toml rename src/coordinator/{ => src/legacy}/analyze/analysis.rs (100%) rename src/coordinator/{ => src/legacy}/analyze/analyzer.rs (100%) rename src/coordinator/{ => src/legacy}/analyze/mod.rs (100%) rename src/coordinator/{ => src/legacy}/coordinator.rs (100%) rename src/coordinator/{ => src/legacy}/dataset/data_set.rs (100%) rename src/coordinator/{ => src/legacy}/dataset/execute_result.rs (100%) rename src/coordinator/{ => src/legacy}/dataset/mod.rs (100%) rename src/coordinator/{ => src/legacy}/dataset/show_catalog_tables_result.rs (100%) rename src/coordinator/{ => src/legacy}/dataset/show_create_streaming_table_result.rs (100%) rename src/coordinator/{ => src/legacy}/dataset/show_create_table_result.rs (100%) rename src/coordinator/{ => src/legacy}/dataset/show_functions_result.rs (100%) rename src/coordinator/{ => src/legacy}/dataset/show_streaming_tables_result.rs (100%) rename src/coordinator/{ => src/legacy}/execution/executor.rs (100%) rename src/coordinator/{ => src/legacy}/execution/mod.rs (100%) rename src/coordinator/{ => src/legacy}/execution_context.rs (100%) rename src/coordinator/{ => src/legacy}/mod.rs (100%) rename src/coordinator/{ => src/legacy}/plan/ast_utils.rs (100%) rename src/coordinator/{ => src/legacy}/plan/compile_error_plan.rs (100%) rename src/coordinator/{ => src/legacy}/plan/create_function_plan.rs (100%) rename src/coordinator/{ => src/legacy}/plan/create_python_function_plan.rs (100%) rename src/coordinator/{ => src/legacy}/plan/create_table_plan.rs (100%) rename src/coordinator/{ => src/legacy}/plan/ddl_compiler.rs (100%) rename src/coordinator/{ => src/legacy}/plan/drop_function_plan.rs (100%) rename src/coordinator/{ => src/legacy}/plan/drop_streaming_table_plan.rs (100%) rename src/coordinator/{ => src/legacy}/plan/drop_table_plan.rs (100%) rename src/coordinator/{ => src/legacy}/plan/logical_plan_visitor.rs (100%) rename src/coordinator/{ => src/legacy}/plan/lookup_table_plan.rs (100%) rename src/coordinator/{ => src/legacy}/plan/mod.rs (100%) rename src/coordinator/{ => src/legacy}/plan/optimizer.rs (100%) rename src/coordinator/{ => src/legacy}/plan/show_catalog_tables_plan.rs (100%) rename src/coordinator/{ => src/legacy}/plan/show_create_streaming_table_plan.rs (100%) rename src/coordinator/{ => src/legacy}/plan/show_create_table_plan.rs (100%) rename src/coordinator/{ => src/legacy}/plan/show_functions_plan.rs (100%) rename src/coordinator/{ => src/legacy}/plan/show_streaming_tables_plan.rs (100%) rename src/coordinator/{ => src/legacy}/plan/start_function_plan.rs (100%) rename src/coordinator/{ => src/legacy}/plan/stop_function_plan.rs (100%) rename src/coordinator/{ => src/legacy}/plan/streaming_compiler.rs (100%) rename src/coordinator/{ => src/legacy}/plan/streaming_table_connector_plan.rs (100%) rename src/coordinator/{ => src/legacy}/plan/streaming_table_plan.rs (100%) rename src/coordinator/{ => src/legacy}/plan/visitor.rs (100%) rename src/coordinator/{ => src/legacy}/runtime_context.rs (100%) rename src/coordinator/{ => src/legacy}/statement/create_function.rs (100%) rename src/coordinator/{ => src/legacy}/statement/create_python_function.rs (100%) rename src/coordinator/{ => src/legacy}/statement/create_table.rs (100%) rename src/coordinator/{ => src/legacy}/statement/drop_function.rs (100%) rename src/coordinator/{ => src/legacy}/statement/drop_streaming_table.rs (100%) rename src/coordinator/{ => src/legacy}/statement/drop_table.rs (100%) rename src/coordinator/{ => src/legacy}/statement/mod.rs (100%) rename src/coordinator/{ => src/legacy}/statement/show_catalog_tables.rs (100%) rename src/coordinator/{ => src/legacy}/statement/show_create_streaming_table.rs (100%) rename src/coordinator/{ => src/legacy}/statement/show_create_table.rs (100%) rename src/coordinator/{ => src/legacy}/statement/show_functions.rs (100%) rename src/coordinator/{ => src/legacy}/statement/show_streaming_tables.rs (100%) rename src/coordinator/{ => src/legacy}/statement/start_function.rs (100%) rename src/coordinator/{ => src/legacy}/statement/stop_function.rs (100%) rename src/coordinator/{ => src/legacy}/statement/streaming_table.rs (100%) rename src/coordinator/{ => src/legacy}/statement/visitor.rs (100%) rename src/coordinator/{ => src/legacy}/streaming_table_options.rs (100%) rename src/coordinator/{ => src/legacy}/tool/mod.rs (100%) create mode 100644 src/coordinator/src/lib.rs create mode 100644 src/job_manager/Cargo.toml create mode 100644 src/job_manager/src/lib.rs create mode 100644 src/logger/Cargo.toml create mode 100644 src/logger/src/lib.rs rename src/{logging/mod.rs => logger/src/logging.rs} (100%) create mode 100644 src/servicer/Cargo.toml rename src/{server => servicer/src/legacy}/handler.rs (100%) rename src/{server => servicer/src/legacy}/initializer.rs (100%) rename src/{server => servicer/src/legacy}/memory_service.rs (100%) rename src/{server => servicer/src/legacy}/mod.rs (100%) rename src/{server => servicer/src/legacy}/service.rs (100%) create mode 100644 src/servicer/src/lib.rs create mode 100644 src/sqlparser/Cargo.toml create mode 100644 src/sqlparser/src/lib.rs create mode 100644 src/streaming_runtime/Cargo.toml create mode 100644 src/streaming_runtime/src/lib.rs create mode 100644 src/wasm_runtime/Cargo.toml create mode 100644 src/wasm_runtime/src/lib.rs diff --git a/Cargo.lock b/Cargo.lock index 18b01c0d..6bb2fb4d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3800,6 +3800,8 @@ dependencies = [ "datafusion-expr 48.0.1", "datafusion-physical-expr 48.0.1", "datafusion-proto", + "function-stream-config", + "function-stream-logger", "futures", "governor", "itertools 0.14.0", @@ -3838,6 +3840,23 @@ dependencies = [ "xxhash-rust", ] +[[package]] +name = "function-stream-catalog" +version = "0.6.0" +dependencies = [ + "protocol", +] + +[[package]] +name = "function-stream-catalog-storage" +version = "0.6.0" +dependencies = [ + "anyhow", + "function-stream-catalog", + "parking_lot", + "rocksdb", +] + [[package]] name = "function-stream-cli" version = "0.1.0" @@ -3853,6 +3872,54 @@ dependencies = [ "tonic", ] +[[package]] +name = "function-stream-common" +version = "0.6.0" + +[[package]] +name = "function-stream-config" +version = "0.6.0" +dependencies = [ + "serde", + "serde_yaml", + "uuid", +] + +[[package]] +name = "function-stream-coordinator" +version = "0.6.0" + +[[package]] +name = "function-stream-job-manager" +version = "0.6.0" + +[[package]] +name = "function-stream-logger" +version = "0.6.0" +dependencies = [ + "anyhow", + "function-stream-config", + "tracing", + "tracing-appender", + "tracing-subscriber", +] + +[[package]] +name = "function-stream-servicer" +version = "0.6.0" + +[[package]] +name = "function-stream-sqlparser" +version = "0.6.0" + +[[package]] +name = "function-stream-streaming-runtime" +version = "0.6.0" + +[[package]] +name = "function-stream-wasm-runtime" +version = "0.6.0" + [[package]] name = "funty" version = "2.0.0" diff --git a/Cargo.toml b/Cargo.toml index eebf2a6c..95fbe566 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,6 +3,17 @@ members = [ ".", "protocol", "cli/cli", + "src/catalog", + "src/catalog_storage", + "src/common", + "src/config", + "src/coordinator", + "src/job_manager", + "src/logger", + "src/servicer", + "src/sqlparser", + "src/streaming_runtime", + "src/wasm_runtime", ] [package] @@ -35,6 +46,8 @@ tonic = { version = "0.12", features = ["default"] } async-trait = "0.1" num_cpus = "1.0" protocol = { path = "./protocol" } +function-stream-config = { path = "src/config" } +function-stream-logger = { path = "src/logger" } prost = "0.13" rdkafka = { version = "0.38", features = ["cmake-build", "ssl", "gssapi", "curl"] } crossbeam-channel = "0.5" diff --git a/src/catalog/Cargo.toml b/src/catalog/Cargo.toml new file mode 100644 index 00000000..07437967 --- /dev/null +++ b/src/catalog/Cargo.toml @@ -0,0 +1,11 @@ +[package] +name = "function-stream-catalog" +version = "0.6.0" +edition = "2024" + +[lib] +name = "function_stream_catalog" +path = "src/lib.rs" + +[dependencies] +protocol = { path = "../../protocol" } diff --git a/src/catalog/src/error.rs b/src/catalog/src/error.rs new file mode 100644 index 00000000..f5a9378e --- /dev/null +++ b/src/catalog/src/error.rs @@ -0,0 +1,4 @@ +use std::error::Error; + +pub type CatalogError = Box; +pub type CatalogResult = Result; diff --git a/src/catalog/src/lib.rs b/src/catalog/src/lib.rs new file mode 100644 index 00000000..33fb2743 --- /dev/null +++ b/src/catalog/src/lib.rs @@ -0,0 +1,9 @@ +//! Catalog domain types and APIs. + +pub mod error; +pub mod meta_store; +pub mod stream_catalog; + +pub use error::{CatalogError, CatalogResult}; +pub use meta_store::MetaStore; +pub use stream_catalog::{GlobalStreamCatalog, StoredStreamingJob, StreamCatalog}; diff --git a/src/catalog/src/meta_store.rs b/src/catalog/src/meta_store.rs new file mode 100644 index 00000000..b361e03c --- /dev/null +++ b/src/catalog/src/meta_store.rs @@ -0,0 +1,23 @@ +use crate::CatalogResult; + +/// Synchronous metadata key-value backend for catalog records. +pub trait MetaStore: Send + Sync { + fn put(&self, key: &str, value: Vec) -> CatalogResult<()>; + fn get(&self, key: &str) -> CatalogResult>>; + fn delete(&self, key: &str) -> CatalogResult<()>; + fn scan_prefix(&self, prefix: &str) -> CatalogResult)>>; + + /// Atomic apply of many puts (`Some(value)`) and deletes (`None`). + /// + /// Backends should override this with a single transaction or write batch + /// when the storage engine supports it. + fn write_batch(&self, batch: Vec<(String, Option>)>) -> CatalogResult<()> { + for (key, value) in batch { + match value { + Some(value) => self.put(&key, value)?, + None => self.delete(&key)?, + } + } + Ok(()) + } +} diff --git a/src/catalog/src/stream_catalog.rs b/src/catalog/src/stream_catalog.rs new file mode 100644 index 00000000..a5641a34 --- /dev/null +++ b/src/catalog/src/stream_catalog.rs @@ -0,0 +1,69 @@ +use std::sync::Arc; + +use protocol::function_stream_graph::FsProgram; +use protocol::storage as pb; + +use crate::CatalogResult; + +/// One persisted streaming job row from catalog storage. +/// +/// This is intentionally storage-agnostic: the catalog keeps source checkpoint +/// payloads as protocol oneof envelopes and does not inspect source-specific +/// checkpoint data. +#[derive(Debug, Clone)] +pub struct StoredStreamingJob { + pub table_name: String, + pub program: FsProgram, + pub checkpoint_interval_ms: u64, + pub latest_checkpoint_epoch: u64, + pub source_checkpoints: Vec, +} + +/// Interface exposed by the stream catalog manager. +/// +/// The concrete table and planning types are generic so this crate can define +/// the catalog boundary without depending on the monolithic SQL/runtime crates. +pub trait StreamCatalog: Send + Sync { + fn persist_streaming_job( + &self, + table_name: &str, + fs_program: &FsProgram, + comment: &str, + checkpoint_interval_ms: u64, + ) -> CatalogResult<()>; + + fn remove_streaming_job(&self, table_name: &str) -> CatalogResult<()>; + + fn commit_job_checkpoint( + &self, + table_name: &str, + epoch: u64, + source_checkpoints: Vec, + ) -> CatalogResult<()>; + + fn load_streaming_job_definitions(&self) -> CatalogResult>; + + fn add_catalog_table(&self, table: Table) -> CatalogResult<()>; + fn has_catalog_table(&self, name: &str) -> bool; + fn drop_catalog_table(&self, table_name: &str, if_exists: bool) -> CatalogResult<()>; + fn restore_from_store(&self) -> CatalogResult<()>; + fn acquire_planning_context(&self) -> PlanningContext; + fn list_catalog_tables(&self) -> CatalogResult>>; + fn get_catalog_table(&self, name: &str) -> CatalogResult>>; + + fn add_table(&self, table: StreamTable) -> CatalogResult<()>; + fn has_stream_table(&self, name: &str) -> bool; + fn drop_table(&self, table_name: &str, if_exists: bool) -> CatalogResult<()>; + fn list_stream_tables(&self) -> Vec>; + fn get_stream_table(&self, name: &str) -> Option>; +} + +/// Process-global catalog access boundary. +pub trait GlobalStreamCatalog: Send + Sync { + fn init_global(manager: Arc) -> CatalogResult<()>; + fn try_global() -> Option>; + + fn global() -> CatalogResult> { + Self::try_global().ok_or_else(|| "CatalogManager not initialized".into()) + } +} diff --git a/src/catalog_storage/Cargo.toml b/src/catalog_storage/Cargo.toml new file mode 100644 index 00000000..f76fb9d3 --- /dev/null +++ b/src/catalog_storage/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "function-stream-catalog-storage" +version = "0.6.0" +edition = "2024" + +[lib] +name = "function_stream_catalog_storage" +path = "src/lib.rs" + +[dependencies] +anyhow = "1.0" +function-stream-catalog = { path = "../catalog" } +parking_lot = "0.12" +rocksdb = { version = "0.21", features = ["multi-threaded-cf", "lz4"] } diff --git a/src/catalog_storage/src/lib.rs b/src/catalog_storage/src/lib.rs new file mode 100644 index 00000000..6ab4891a --- /dev/null +++ b/src/catalog_storage/src/lib.rs @@ -0,0 +1,7 @@ +//! Persistent catalog storage implementations. + +pub mod memory; +pub mod rocksdb; + +pub use memory::InMemoryMetaStore; +pub use rocksdb::RocksDbMetaStore; diff --git a/src/catalog_storage/src/memory.rs b/src/catalog_storage/src/memory.rs new file mode 100644 index 00000000..df6fd0c0 --- /dev/null +++ b/src/catalog_storage/src/memory.rs @@ -0,0 +1,63 @@ +use std::collections::HashMap; + +use function_stream_catalog::{CatalogResult, MetaStore}; +use parking_lot::RwLock; + +/// In-process KV store for single-node deployments and tests. +pub struct InMemoryMetaStore { + db: RwLock>>, +} + +impl InMemoryMetaStore { + pub fn new() -> Self { + Self { + db: RwLock::new(HashMap::new()), + } + } +} + +impl Default for InMemoryMetaStore { + fn default() -> Self { + Self::new() + } +} + +impl MetaStore for InMemoryMetaStore { + fn put(&self, key: &str, value: Vec) -> CatalogResult<()> { + self.db.write().insert(key.to_string(), value); + Ok(()) + } + + fn get(&self, key: &str) -> CatalogResult>> { + Ok(self.db.read().get(key).cloned()) + } + + fn delete(&self, key: &str) -> CatalogResult<()> { + self.db.write().remove(key); + Ok(()) + } + + fn scan_prefix(&self, prefix: &str) -> CatalogResult)>> { + let db = self.db.read(); + Ok(db + .iter() + .filter(|(key, _)| key.starts_with(prefix)) + .map(|(key, value)| (key.clone(), value.clone())) + .collect()) + } + + fn write_batch(&self, batch: Vec<(String, Option>)>) -> CatalogResult<()> { + let mut db = self.db.write(); + for (key, value) in batch { + match value { + Some(value) => { + db.insert(key, value); + } + None => { + db.remove(&key); + } + } + } + Ok(()) + } +} diff --git a/src/catalog_storage/src/rocksdb.rs b/src/catalog_storage/src/rocksdb.rs new file mode 100644 index 00000000..192b9d80 --- /dev/null +++ b/src/catalog_storage/src/rocksdb.rs @@ -0,0 +1,81 @@ +use std::path::Path; +use std::sync::Arc; + +use anyhow::Context; +use function_stream_catalog::{CatalogResult, MetaStore}; +use rocksdb::{DB, Direction, IteratorMode, Options, WriteBatch}; + +/// RocksDB-backed catalog metadata store. +pub struct RocksDbMetaStore { + db: Arc, +} + +impl RocksDbMetaStore { + pub fn open>(path: P) -> anyhow::Result { + let path = path.as_ref(); + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent) + .with_context(|| format!("stream catalog: create parent directory {parent:?}"))?; + } + let mut opts = Options::default(); + opts.create_if_missing(true); + let db = DB::open(&opts, path) + .with_context(|| format!("stream catalog: open RocksDB at {}", path.display()))?; + Ok(Self { db: Arc::new(db) }) + } +} + +impl MetaStore for RocksDbMetaStore { + fn put(&self, key: &str, value: Vec) -> CatalogResult<()> { + self.db + .put(key.as_bytes(), value.as_slice()) + .map_err(|e| format!("stream catalog store put: {e}").into()) + } + + fn get(&self, key: &str) -> CatalogResult>> { + self.db + .get(key.as_bytes()) + .map_err(|e| format!("stream catalog store get: {e}").into()) + } + + fn delete(&self, key: &str) -> CatalogResult<()> { + self.db + .delete(key.as_bytes()) + .map_err(|e| format!("stream catalog store delete: {e}").into()) + } + + fn scan_prefix(&self, prefix: &str) -> CatalogResult)>> { + let mut out = Vec::new(); + let iter = self + .db + .iterator(IteratorMode::From(prefix.as_bytes(), Direction::Forward)); + for item in iter { + let (key, value) = item.map_err(|e| format!("stream catalog store scan: {e}"))?; + let key = String::from_utf8(key.to_vec()) + .map_err(|e| format!("stream catalog store: invalid utf8 key: {e}"))?; + if !key.starts_with(prefix) { + break; + } + out.push((key, value.to_vec())); + } + Ok(out) + } + + fn write_batch(&self, batch: Vec<(String, Option>)>) -> CatalogResult<()> { + if batch.is_empty() { + return Ok(()); + } + + let mut write_batch = WriteBatch::default(); + for (key, value) in batch { + match value { + Some(value) => write_batch.put(key.as_bytes(), value.as_slice()), + None => write_batch.delete(key.as_bytes()), + } + } + + self.db + .write(write_batch) + .map_err(|e| format!("stream catalog store write_batch: {e}").into()) + } +} diff --git a/src/common/Cargo.toml b/src/common/Cargo.toml new file mode 100644 index 00000000..fb68a526 --- /dev/null +++ b/src/common/Cargo.toml @@ -0,0 +1,8 @@ +[package] +name = "function-stream-common" +version = "0.6.0" +edition = "2024" + +[lib] +name = "function_stream_common" +path = "src/lib.rs" diff --git a/src/common/fs_schema.rs b/src/common/src/fs_schema.rs similarity index 100% rename from src/common/fs_schema.rs rename to src/common/src/fs_schema.rs diff --git a/src/common/mod.rs b/src/common/src/legacy_mod.rs similarity index 100% rename from src/common/mod.rs rename to src/common/src/legacy_mod.rs diff --git a/src/common/src/lib.rs b/src/common/src/lib.rs new file mode 100644 index 00000000..0356f6de --- /dev/null +++ b/src/common/src/lib.rs @@ -0,0 +1,3 @@ +//! Shared types and utilities for FunctionStream crates. + +pub const CRATE_NAME: &str = "function-stream-common"; diff --git a/src/config/Cargo.toml b/src/config/Cargo.toml new file mode 100644 index 00000000..edd2a8db --- /dev/null +++ b/src/config/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "function-stream-config" +version = "0.6.0" +edition = "2024" + +[lib] +name = "function_stream_config" +path = "src/lib.rs" + +[features] +default = ["python"] +python = [] + +[dependencies] +serde = { version = "1.0", features = ["derive"] } +serde_yaml = "0.9" +uuid = { version = "1.0", features = ["v4", "v7"] } diff --git a/src/config/global_config.rs b/src/config/src/global_config.rs similarity index 100% rename from src/config/global_config.rs rename to src/config/src/global_config.rs diff --git a/src/config/mod.rs b/src/config/src/legacy/mod.rs similarity index 100% rename from src/config/mod.rs rename to src/config/src/legacy/mod.rs diff --git a/src/config/src/lib.rs b/src/config/src/lib.rs new file mode 100644 index 00000000..c07c676c --- /dev/null +++ b/src/config/src/lib.rs @@ -0,0 +1,33 @@ +//! Configuration loading and validation. + +pub mod global_config; +pub mod loader; +pub mod log_config; +pub mod paths; +pub mod python_config; +pub mod service_config; +pub mod storage; +pub mod streaming_job; +pub mod system; +pub mod wasm_config; + +// Compatibility shim for files that still reference `crate::config::*`. +pub mod config { + pub use crate::*; +} + +pub use global_config::{ + DEFAULT_OPERATOR_STATE_STORE_MEMORY_BYTES, DEFAULT_STREAMING_RUNTIME_MEMORY_BYTES, GlobalConfig, +}; +pub use loader::load_global_config; +pub use log_config::LogConfig; +#[allow(unused_imports)] +pub use paths::{ + ENV_CONF, ENV_HOME, find_config_file, get_app_log_path, get_conf_dir, get_data_dir, + get_log_path, get_logs_dir, get_project_root, get_python_cache_dir, get_python_cwasm_path, + get_python_wasm_path, get_state_dir, get_state_dir_for_base, get_task_dir, get_wasm_cache_dir, + resolve_path, +}; +#[cfg(feature = "python")] +pub use python_config::PythonConfig; +pub use streaming_job::{DEFAULT_CHECKPOINT_INTERVAL_MS, DEFAULT_PIPELINE_PARALLELISM}; diff --git a/src/config/loader.rs b/src/config/src/loader.rs similarity index 100% rename from src/config/loader.rs rename to src/config/src/loader.rs diff --git a/src/config/log_config.rs b/src/config/src/log_config.rs similarity index 100% rename from src/config/log_config.rs rename to src/config/src/log_config.rs diff --git a/src/config/paths.rs b/src/config/src/paths.rs similarity index 100% rename from src/config/paths.rs rename to src/config/src/paths.rs diff --git a/src/config/python_config.rs b/src/config/src/python_config.rs similarity index 100% rename from src/config/python_config.rs rename to src/config/src/python_config.rs diff --git a/src/config/service_config.rs b/src/config/src/service_config.rs similarity index 100% rename from src/config/service_config.rs rename to src/config/src/service_config.rs diff --git a/src/config/storage.rs b/src/config/src/storage.rs similarity index 100% rename from src/config/storage.rs rename to src/config/src/storage.rs diff --git a/src/config/streaming_job.rs b/src/config/src/streaming_job.rs similarity index 100% rename from src/config/streaming_job.rs rename to src/config/src/streaming_job.rs diff --git a/src/config/system.rs b/src/config/src/system.rs similarity index 100% rename from src/config/system.rs rename to src/config/src/system.rs diff --git a/src/config/wasm_config.rs b/src/config/src/wasm_config.rs similarity index 100% rename from src/config/wasm_config.rs rename to src/config/src/wasm_config.rs diff --git a/src/coordinator/Cargo.toml b/src/coordinator/Cargo.toml new file mode 100644 index 00000000..5322f8a2 --- /dev/null +++ b/src/coordinator/Cargo.toml @@ -0,0 +1,8 @@ +[package] +name = "function-stream-coordinator" +version = "0.6.0" +edition = "2024" + +[lib] +name = "function_stream_coordinator" +path = "src/lib.rs" diff --git a/src/coordinator/analyze/analysis.rs b/src/coordinator/src/legacy/analyze/analysis.rs similarity index 100% rename from src/coordinator/analyze/analysis.rs rename to src/coordinator/src/legacy/analyze/analysis.rs diff --git a/src/coordinator/analyze/analyzer.rs b/src/coordinator/src/legacy/analyze/analyzer.rs similarity index 100% rename from src/coordinator/analyze/analyzer.rs rename to src/coordinator/src/legacy/analyze/analyzer.rs diff --git a/src/coordinator/analyze/mod.rs b/src/coordinator/src/legacy/analyze/mod.rs similarity index 100% rename from src/coordinator/analyze/mod.rs rename to src/coordinator/src/legacy/analyze/mod.rs diff --git a/src/coordinator/coordinator.rs b/src/coordinator/src/legacy/coordinator.rs similarity index 100% rename from src/coordinator/coordinator.rs rename to src/coordinator/src/legacy/coordinator.rs diff --git a/src/coordinator/dataset/data_set.rs b/src/coordinator/src/legacy/dataset/data_set.rs similarity index 100% rename from src/coordinator/dataset/data_set.rs rename to src/coordinator/src/legacy/dataset/data_set.rs diff --git a/src/coordinator/dataset/execute_result.rs b/src/coordinator/src/legacy/dataset/execute_result.rs similarity index 100% rename from src/coordinator/dataset/execute_result.rs rename to src/coordinator/src/legacy/dataset/execute_result.rs diff --git a/src/coordinator/dataset/mod.rs b/src/coordinator/src/legacy/dataset/mod.rs similarity index 100% rename from src/coordinator/dataset/mod.rs rename to src/coordinator/src/legacy/dataset/mod.rs diff --git a/src/coordinator/dataset/show_catalog_tables_result.rs b/src/coordinator/src/legacy/dataset/show_catalog_tables_result.rs similarity index 100% rename from src/coordinator/dataset/show_catalog_tables_result.rs rename to src/coordinator/src/legacy/dataset/show_catalog_tables_result.rs diff --git a/src/coordinator/dataset/show_create_streaming_table_result.rs b/src/coordinator/src/legacy/dataset/show_create_streaming_table_result.rs similarity index 100% rename from src/coordinator/dataset/show_create_streaming_table_result.rs rename to src/coordinator/src/legacy/dataset/show_create_streaming_table_result.rs diff --git a/src/coordinator/dataset/show_create_table_result.rs b/src/coordinator/src/legacy/dataset/show_create_table_result.rs similarity index 100% rename from src/coordinator/dataset/show_create_table_result.rs rename to src/coordinator/src/legacy/dataset/show_create_table_result.rs diff --git a/src/coordinator/dataset/show_functions_result.rs b/src/coordinator/src/legacy/dataset/show_functions_result.rs similarity index 100% rename from src/coordinator/dataset/show_functions_result.rs rename to src/coordinator/src/legacy/dataset/show_functions_result.rs diff --git a/src/coordinator/dataset/show_streaming_tables_result.rs b/src/coordinator/src/legacy/dataset/show_streaming_tables_result.rs similarity index 100% rename from src/coordinator/dataset/show_streaming_tables_result.rs rename to src/coordinator/src/legacy/dataset/show_streaming_tables_result.rs diff --git a/src/coordinator/execution/executor.rs b/src/coordinator/src/legacy/execution/executor.rs similarity index 100% rename from src/coordinator/execution/executor.rs rename to src/coordinator/src/legacy/execution/executor.rs diff --git a/src/coordinator/execution/mod.rs b/src/coordinator/src/legacy/execution/mod.rs similarity index 100% rename from src/coordinator/execution/mod.rs rename to src/coordinator/src/legacy/execution/mod.rs diff --git a/src/coordinator/execution_context.rs b/src/coordinator/src/legacy/execution_context.rs similarity index 100% rename from src/coordinator/execution_context.rs rename to src/coordinator/src/legacy/execution_context.rs diff --git a/src/coordinator/mod.rs b/src/coordinator/src/legacy/mod.rs similarity index 100% rename from src/coordinator/mod.rs rename to src/coordinator/src/legacy/mod.rs diff --git a/src/coordinator/plan/ast_utils.rs b/src/coordinator/src/legacy/plan/ast_utils.rs similarity index 100% rename from src/coordinator/plan/ast_utils.rs rename to src/coordinator/src/legacy/plan/ast_utils.rs diff --git a/src/coordinator/plan/compile_error_plan.rs b/src/coordinator/src/legacy/plan/compile_error_plan.rs similarity index 100% rename from src/coordinator/plan/compile_error_plan.rs rename to src/coordinator/src/legacy/plan/compile_error_plan.rs diff --git a/src/coordinator/plan/create_function_plan.rs b/src/coordinator/src/legacy/plan/create_function_plan.rs similarity index 100% rename from src/coordinator/plan/create_function_plan.rs rename to src/coordinator/src/legacy/plan/create_function_plan.rs diff --git a/src/coordinator/plan/create_python_function_plan.rs b/src/coordinator/src/legacy/plan/create_python_function_plan.rs similarity index 100% rename from src/coordinator/plan/create_python_function_plan.rs rename to src/coordinator/src/legacy/plan/create_python_function_plan.rs diff --git a/src/coordinator/plan/create_table_plan.rs b/src/coordinator/src/legacy/plan/create_table_plan.rs similarity index 100% rename from src/coordinator/plan/create_table_plan.rs rename to src/coordinator/src/legacy/plan/create_table_plan.rs diff --git a/src/coordinator/plan/ddl_compiler.rs b/src/coordinator/src/legacy/plan/ddl_compiler.rs similarity index 100% rename from src/coordinator/plan/ddl_compiler.rs rename to src/coordinator/src/legacy/plan/ddl_compiler.rs diff --git a/src/coordinator/plan/drop_function_plan.rs b/src/coordinator/src/legacy/plan/drop_function_plan.rs similarity index 100% rename from src/coordinator/plan/drop_function_plan.rs rename to src/coordinator/src/legacy/plan/drop_function_plan.rs diff --git a/src/coordinator/plan/drop_streaming_table_plan.rs b/src/coordinator/src/legacy/plan/drop_streaming_table_plan.rs similarity index 100% rename from src/coordinator/plan/drop_streaming_table_plan.rs rename to src/coordinator/src/legacy/plan/drop_streaming_table_plan.rs diff --git a/src/coordinator/plan/drop_table_plan.rs b/src/coordinator/src/legacy/plan/drop_table_plan.rs similarity index 100% rename from src/coordinator/plan/drop_table_plan.rs rename to src/coordinator/src/legacy/plan/drop_table_plan.rs diff --git a/src/coordinator/plan/logical_plan_visitor.rs b/src/coordinator/src/legacy/plan/logical_plan_visitor.rs similarity index 100% rename from src/coordinator/plan/logical_plan_visitor.rs rename to src/coordinator/src/legacy/plan/logical_plan_visitor.rs diff --git a/src/coordinator/plan/lookup_table_plan.rs b/src/coordinator/src/legacy/plan/lookup_table_plan.rs similarity index 100% rename from src/coordinator/plan/lookup_table_plan.rs rename to src/coordinator/src/legacy/plan/lookup_table_plan.rs diff --git a/src/coordinator/plan/mod.rs b/src/coordinator/src/legacy/plan/mod.rs similarity index 100% rename from src/coordinator/plan/mod.rs rename to src/coordinator/src/legacy/plan/mod.rs diff --git a/src/coordinator/plan/optimizer.rs b/src/coordinator/src/legacy/plan/optimizer.rs similarity index 100% rename from src/coordinator/plan/optimizer.rs rename to src/coordinator/src/legacy/plan/optimizer.rs diff --git a/src/coordinator/plan/show_catalog_tables_plan.rs b/src/coordinator/src/legacy/plan/show_catalog_tables_plan.rs similarity index 100% rename from src/coordinator/plan/show_catalog_tables_plan.rs rename to src/coordinator/src/legacy/plan/show_catalog_tables_plan.rs diff --git a/src/coordinator/plan/show_create_streaming_table_plan.rs b/src/coordinator/src/legacy/plan/show_create_streaming_table_plan.rs similarity index 100% rename from src/coordinator/plan/show_create_streaming_table_plan.rs rename to src/coordinator/src/legacy/plan/show_create_streaming_table_plan.rs diff --git a/src/coordinator/plan/show_create_table_plan.rs b/src/coordinator/src/legacy/plan/show_create_table_plan.rs similarity index 100% rename from src/coordinator/plan/show_create_table_plan.rs rename to src/coordinator/src/legacy/plan/show_create_table_plan.rs diff --git a/src/coordinator/plan/show_functions_plan.rs b/src/coordinator/src/legacy/plan/show_functions_plan.rs similarity index 100% rename from src/coordinator/plan/show_functions_plan.rs rename to src/coordinator/src/legacy/plan/show_functions_plan.rs diff --git a/src/coordinator/plan/show_streaming_tables_plan.rs b/src/coordinator/src/legacy/plan/show_streaming_tables_plan.rs similarity index 100% rename from src/coordinator/plan/show_streaming_tables_plan.rs rename to src/coordinator/src/legacy/plan/show_streaming_tables_plan.rs diff --git a/src/coordinator/plan/start_function_plan.rs b/src/coordinator/src/legacy/plan/start_function_plan.rs similarity index 100% rename from src/coordinator/plan/start_function_plan.rs rename to src/coordinator/src/legacy/plan/start_function_plan.rs diff --git a/src/coordinator/plan/stop_function_plan.rs b/src/coordinator/src/legacy/plan/stop_function_plan.rs similarity index 100% rename from src/coordinator/plan/stop_function_plan.rs rename to src/coordinator/src/legacy/plan/stop_function_plan.rs diff --git a/src/coordinator/plan/streaming_compiler.rs b/src/coordinator/src/legacy/plan/streaming_compiler.rs similarity index 100% rename from src/coordinator/plan/streaming_compiler.rs rename to src/coordinator/src/legacy/plan/streaming_compiler.rs diff --git a/src/coordinator/plan/streaming_table_connector_plan.rs b/src/coordinator/src/legacy/plan/streaming_table_connector_plan.rs similarity index 100% rename from src/coordinator/plan/streaming_table_connector_plan.rs rename to src/coordinator/src/legacy/plan/streaming_table_connector_plan.rs diff --git a/src/coordinator/plan/streaming_table_plan.rs b/src/coordinator/src/legacy/plan/streaming_table_plan.rs similarity index 100% rename from src/coordinator/plan/streaming_table_plan.rs rename to src/coordinator/src/legacy/plan/streaming_table_plan.rs diff --git a/src/coordinator/plan/visitor.rs b/src/coordinator/src/legacy/plan/visitor.rs similarity index 100% rename from src/coordinator/plan/visitor.rs rename to src/coordinator/src/legacy/plan/visitor.rs diff --git a/src/coordinator/runtime_context.rs b/src/coordinator/src/legacy/runtime_context.rs similarity index 100% rename from src/coordinator/runtime_context.rs rename to src/coordinator/src/legacy/runtime_context.rs diff --git a/src/coordinator/statement/create_function.rs b/src/coordinator/src/legacy/statement/create_function.rs similarity index 100% rename from src/coordinator/statement/create_function.rs rename to src/coordinator/src/legacy/statement/create_function.rs diff --git a/src/coordinator/statement/create_python_function.rs b/src/coordinator/src/legacy/statement/create_python_function.rs similarity index 100% rename from src/coordinator/statement/create_python_function.rs rename to src/coordinator/src/legacy/statement/create_python_function.rs diff --git a/src/coordinator/statement/create_table.rs b/src/coordinator/src/legacy/statement/create_table.rs similarity index 100% rename from src/coordinator/statement/create_table.rs rename to src/coordinator/src/legacy/statement/create_table.rs diff --git a/src/coordinator/statement/drop_function.rs b/src/coordinator/src/legacy/statement/drop_function.rs similarity index 100% rename from src/coordinator/statement/drop_function.rs rename to src/coordinator/src/legacy/statement/drop_function.rs diff --git a/src/coordinator/statement/drop_streaming_table.rs b/src/coordinator/src/legacy/statement/drop_streaming_table.rs similarity index 100% rename from src/coordinator/statement/drop_streaming_table.rs rename to src/coordinator/src/legacy/statement/drop_streaming_table.rs diff --git a/src/coordinator/statement/drop_table.rs b/src/coordinator/src/legacy/statement/drop_table.rs similarity index 100% rename from src/coordinator/statement/drop_table.rs rename to src/coordinator/src/legacy/statement/drop_table.rs diff --git a/src/coordinator/statement/mod.rs b/src/coordinator/src/legacy/statement/mod.rs similarity index 100% rename from src/coordinator/statement/mod.rs rename to src/coordinator/src/legacy/statement/mod.rs diff --git a/src/coordinator/statement/show_catalog_tables.rs b/src/coordinator/src/legacy/statement/show_catalog_tables.rs similarity index 100% rename from src/coordinator/statement/show_catalog_tables.rs rename to src/coordinator/src/legacy/statement/show_catalog_tables.rs diff --git a/src/coordinator/statement/show_create_streaming_table.rs b/src/coordinator/src/legacy/statement/show_create_streaming_table.rs similarity index 100% rename from src/coordinator/statement/show_create_streaming_table.rs rename to src/coordinator/src/legacy/statement/show_create_streaming_table.rs diff --git a/src/coordinator/statement/show_create_table.rs b/src/coordinator/src/legacy/statement/show_create_table.rs similarity index 100% rename from src/coordinator/statement/show_create_table.rs rename to src/coordinator/src/legacy/statement/show_create_table.rs diff --git a/src/coordinator/statement/show_functions.rs b/src/coordinator/src/legacy/statement/show_functions.rs similarity index 100% rename from src/coordinator/statement/show_functions.rs rename to src/coordinator/src/legacy/statement/show_functions.rs diff --git a/src/coordinator/statement/show_streaming_tables.rs b/src/coordinator/src/legacy/statement/show_streaming_tables.rs similarity index 100% rename from src/coordinator/statement/show_streaming_tables.rs rename to src/coordinator/src/legacy/statement/show_streaming_tables.rs diff --git a/src/coordinator/statement/start_function.rs b/src/coordinator/src/legacy/statement/start_function.rs similarity index 100% rename from src/coordinator/statement/start_function.rs rename to src/coordinator/src/legacy/statement/start_function.rs diff --git a/src/coordinator/statement/stop_function.rs b/src/coordinator/src/legacy/statement/stop_function.rs similarity index 100% rename from src/coordinator/statement/stop_function.rs rename to src/coordinator/src/legacy/statement/stop_function.rs diff --git a/src/coordinator/statement/streaming_table.rs b/src/coordinator/src/legacy/statement/streaming_table.rs similarity index 100% rename from src/coordinator/statement/streaming_table.rs rename to src/coordinator/src/legacy/statement/streaming_table.rs diff --git a/src/coordinator/statement/visitor.rs b/src/coordinator/src/legacy/statement/visitor.rs similarity index 100% rename from src/coordinator/statement/visitor.rs rename to src/coordinator/src/legacy/statement/visitor.rs diff --git a/src/coordinator/streaming_table_options.rs b/src/coordinator/src/legacy/streaming_table_options.rs similarity index 100% rename from src/coordinator/streaming_table_options.rs rename to src/coordinator/src/legacy/streaming_table_options.rs diff --git a/src/coordinator/tool/mod.rs b/src/coordinator/src/legacy/tool/mod.rs similarity index 100% rename from src/coordinator/tool/mod.rs rename to src/coordinator/src/legacy/tool/mod.rs diff --git a/src/coordinator/src/lib.rs b/src/coordinator/src/lib.rs new file mode 100644 index 00000000..6e5ad596 --- /dev/null +++ b/src/coordinator/src/lib.rs @@ -0,0 +1,3 @@ +//! Query planning and job coordination. + +pub const CRATE_NAME: &str = "function-stream-coordinator"; diff --git a/src/job_manager/Cargo.toml b/src/job_manager/Cargo.toml new file mode 100644 index 00000000..1315a9e1 --- /dev/null +++ b/src/job_manager/Cargo.toml @@ -0,0 +1,8 @@ +[package] +name = "function-stream-job-manager" +version = "0.6.0" +edition = "2024" + +[lib] +name = "function_stream_job_manager" +path = "src/lib.rs" diff --git a/src/job_manager/src/lib.rs b/src/job_manager/src/lib.rs new file mode 100644 index 00000000..6877187d --- /dev/null +++ b/src/job_manager/src/lib.rs @@ -0,0 +1,3 @@ +//! Job lifecycle management and scheduling. + +pub const CRATE_NAME: &str = "function-stream-job-manager"; diff --git a/src/lib.rs b/src/lib.rs index a6bb4d28..4c1af5d9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -14,10 +14,12 @@ #![allow(dead_code)] -pub mod config; +pub use function_stream_config as config; +#[path = "coordinator/src/legacy/mod.rs"] pub mod coordinator; -pub mod logging; +pub use function_stream_logger as logging; pub mod runtime; +#[path = "servicer/src/legacy/mod.rs"] pub mod server; pub mod sql; pub mod storage; diff --git a/src/logger/Cargo.toml b/src/logger/Cargo.toml new file mode 100644 index 00000000..8d584f74 --- /dev/null +++ b/src/logger/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "function-stream-logger" +version = "0.6.0" +edition = "2024" + +[lib] +name = "function_stream_logger" +path = "src/lib.rs" + +[dependencies] +anyhow = "1.0" +function-stream-config = { path = "../config" } +tracing = "0.1" +tracing-appender = "0.2" +tracing-subscriber = { version = "0.3", features = ["env-filter"] } diff --git a/src/logger/src/lib.rs b/src/logger/src/lib.rs new file mode 100644 index 00000000..53fb0b31 --- /dev/null +++ b/src/logger/src/lib.rs @@ -0,0 +1,9 @@ +//! Logging setup and helpers. + +pub mod config { + pub use function_stream_config::*; +} + +mod logging; + +pub use logging::init_logging; diff --git a/src/logging/mod.rs b/src/logger/src/logging.rs similarity index 100% rename from src/logging/mod.rs rename to src/logger/src/logging.rs diff --git a/src/main.rs b/src/main.rs index 46da3c7a..ee51cb3e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -12,10 +12,12 @@ #![allow(dead_code)] -mod config; +pub use function_stream_config as config; +#[path = "coordinator/src/legacy/mod.rs"] mod coordinator; -mod logging; +pub use function_stream_logger as logging; mod runtime; +#[path = "servicer/src/legacy/mod.rs"] mod server; mod sql; mod storage; @@ -144,7 +146,7 @@ fn setup_environment() -> Result { config::GlobalConfig::default() }; - logging::init_logging(&config.logging).context("Logging initialization failed")?; + function_stream_logger::init_logging(&config.logging).context("Logging initialization failed")?; log::debug!( "Environment initialized. Data: {}, Conf: {}", diff --git a/src/servicer/Cargo.toml b/src/servicer/Cargo.toml new file mode 100644 index 00000000..08860c22 --- /dev/null +++ b/src/servicer/Cargo.toml @@ -0,0 +1,8 @@ +[package] +name = "function-stream-servicer" +version = "0.6.0" +edition = "2024" + +[lib] +name = "function_stream_servicer" +path = "src/lib.rs" diff --git a/src/server/handler.rs b/src/servicer/src/legacy/handler.rs similarity index 100% rename from src/server/handler.rs rename to src/servicer/src/legacy/handler.rs diff --git a/src/server/initializer.rs b/src/servicer/src/legacy/initializer.rs similarity index 100% rename from src/server/initializer.rs rename to src/servicer/src/legacy/initializer.rs diff --git a/src/server/memory_service.rs b/src/servicer/src/legacy/memory_service.rs similarity index 100% rename from src/server/memory_service.rs rename to src/servicer/src/legacy/memory_service.rs diff --git a/src/server/mod.rs b/src/servicer/src/legacy/mod.rs similarity index 100% rename from src/server/mod.rs rename to src/servicer/src/legacy/mod.rs diff --git a/src/server/service.rs b/src/servicer/src/legacy/service.rs similarity index 100% rename from src/server/service.rs rename to src/servicer/src/legacy/service.rs diff --git a/src/servicer/src/lib.rs b/src/servicer/src/lib.rs new file mode 100644 index 00000000..8a442937 --- /dev/null +++ b/src/servicer/src/lib.rs @@ -0,0 +1,3 @@ +//! Service layer implementations and request handling. + +pub const CRATE_NAME: &str = "function-stream-servicer"; diff --git a/src/sqlparser/Cargo.toml b/src/sqlparser/Cargo.toml new file mode 100644 index 00000000..6489639a --- /dev/null +++ b/src/sqlparser/Cargo.toml @@ -0,0 +1,8 @@ +[package] +name = "function-stream-sqlparser" +version = "0.6.0" +edition = "2024" + +[lib] +name = "function_stream_sqlparser" +path = "src/lib.rs" diff --git a/src/sqlparser/src/lib.rs b/src/sqlparser/src/lib.rs new file mode 100644 index 00000000..b436f895 --- /dev/null +++ b/src/sqlparser/src/lib.rs @@ -0,0 +1,3 @@ +//! SQL parsing facade for FunctionStream. + +pub const CRATE_NAME: &str = "function-stream-sqlparser"; diff --git a/src/streaming_runtime/Cargo.toml b/src/streaming_runtime/Cargo.toml new file mode 100644 index 00000000..732fd51a --- /dev/null +++ b/src/streaming_runtime/Cargo.toml @@ -0,0 +1,8 @@ +[package] +name = "function-stream-streaming-runtime" +version = "0.6.0" +edition = "2024" + +[lib] +name = "function_stream_streaming_runtime" +path = "src/lib.rs" diff --git a/src/streaming_runtime/src/lib.rs b/src/streaming_runtime/src/lib.rs new file mode 100644 index 00000000..a4643ba9 --- /dev/null +++ b/src/streaming_runtime/src/lib.rs @@ -0,0 +1,3 @@ +//! Streaming execution runtime. + +pub const CRATE_NAME: &str = "function-stream-streaming-runtime"; diff --git a/src/wasm_runtime/Cargo.toml b/src/wasm_runtime/Cargo.toml new file mode 100644 index 00000000..4cb1eae2 --- /dev/null +++ b/src/wasm_runtime/Cargo.toml @@ -0,0 +1,8 @@ +[package] +name = "function-stream-wasm-runtime" +version = "0.6.0" +edition = "2024" + +[lib] +name = "function_stream_wasm_runtime" +path = "src/lib.rs" diff --git a/src/wasm_runtime/src/lib.rs b/src/wasm_runtime/src/lib.rs new file mode 100644 index 00000000..3eeefce9 --- /dev/null +++ b/src/wasm_runtime/src/lib.rs @@ -0,0 +1,3 @@ +//! WebAssembly execution runtime. + +pub const CRATE_NAME: &str = "function-stream-wasm-runtime"; From d5e7ddd45411f3d7f74edd5e109e28dd95bdc54b Mon Sep 17 00:00:00 2001 From: luoluoyuyu Date: Tue, 12 May 2026 00:53:22 +0800 Subject: [PATCH 2/7] update --- Cargo.lock | 12 + Cargo.toml | 2 + src/catalog_storage/src/lib.rs | 3 + .../src}/stream_catalog/codec.rs | 0 .../src}/stream_catalog/manager.rs | 0 .../src}/stream_catalog/meta_store.rs | 0 .../src}/stream_catalog/mod.rs | 0 .../src}/stream_catalog/rocksdb_meta_store.rs | 0 .../src}/task/factory.rs | 0 .../src}/task/function_info.rs | 0 .../src}/task/mod.rs | 0 .../src}/task/proto_codec.rs | 0 .../src}/task/rocksdb_storage.rs | 0 .../src}/task/storage.rs | 0 src/coordinator/src/legacy/mod.rs | 2 + src/coordinator/src/legacy/sql_classify.rs | 326 ++++++++++++++++ src/main.rs | 3 +- src/runtime/mod.rs | 12 +- src/runtime_common/Cargo.toml | 15 + .../src}/common/component_state.rs | 0 .../src}/common/mod.rs | 0 .../src}/common/task_completion.rs | 0 src/runtime_common/src/lib.rs | 16 + .../src}/memory/block.rs | 0 .../src}/memory/error.rs | 0 .../src}/memory/global.rs | 0 .../src}/memory/mod.rs | 0 .../src}/memory/pool.rs | 0 .../src}/memory/ticket.rs | 0 src/servicer/src/legacy/handler.rs | 12 +- src/sql/parse.rs | 369 ++---------------- src/storage/mod.rs | 8 + src/streaming_runtime/src/lib.rs | 18 + .../src}/streaming/api/context.rs | 0 .../src}/streaming/api/mod.rs | 0 .../src}/streaming/api/operator.rs | 0 .../src}/streaming/api/source.rs | 0 .../src}/streaming/error.rs | 0 .../src}/streaming/execution/mod.rs | 0 .../streaming/execution/operator_chain.rs | 0 .../src}/streaming/execution/pipeline.rs | 0 .../src}/streaming/execution/source_driver.rs | 0 .../execution/tracker/barrier_aligner.rs | 0 .../src}/streaming/execution/tracker/mod.rs | 0 .../execution/tracker/watermark_tracker.rs | 0 .../src}/streaming/factory/connector/delta.rs | 0 .../factory/connector/dispatchers.rs | 0 .../streaming/factory/connector/filesystem.rs | 0 .../streaming/factory/connector/iceberg.rs | 0 .../src}/streaming/factory/connector/kafka.rs | 0 .../streaming/factory/connector/lancedb.rs | 0 .../src}/streaming/factory/connector/mod.rs | 0 .../src}/streaming/factory/connector/s3.rs | 0 .../factory/connector/sink_props_codec.rs | 0 .../src}/streaming/factory/global/mod.rs | 0 .../factory/global/session_registry.rs | 0 .../src}/streaming/factory/mod.rs | 0 .../streaming/factory/operator_constructor.rs | 0 .../streaming/factory/operator_factory.rs | 0 .../src}/streaming/format/config.rs | 0 .../src}/streaming/format/deserializer.rs | 0 .../src}/streaming/format/encoder.rs | 0 .../src}/streaming/format/json_encoder.rs | 0 .../src}/streaming/format/mod.rs | 0 .../src}/streaming/format/serializer.rs | 0 .../src}/streaming/job/edge_manager.rs | 0 .../src}/streaming/job/job_manager.rs | 0 .../src}/streaming/job/mod.rs | 0 .../src}/streaming/job/models.rs | 0 .../src}/streaming/mod.rs | 0 .../src}/streaming/network/endpoint.rs | 0 .../src}/streaming/network/environment.rs | 0 .../src}/streaming/network/mod.rs | 0 .../grouping/incremental_aggregate.rs | 0 .../src}/streaming/operators/grouping/mod.rs | 0 .../operators/grouping/updating_cache.rs | 0 .../operators/joins/join_instance.rs | 0 .../operators/joins/join_with_expiration.rs | 0 .../src}/streaming/operators/joins/mod.rs | 0 .../src}/streaming/operators/key_by.rs | 0 .../src}/streaming/operators/key_operator.rs | 0 .../src}/streaming/operators/mod.rs | 0 .../src}/streaming/operators/projection.rs | 0 .../streaming/operators/sink/delta/mod.rs | 0 .../operators/sink/filesystem/mod.rs | 0 .../streaming/operators/sink/iceberg/mod.rs | 0 .../streaming/operators/sink/kafka/mod.rs | 0 .../streaming/operators/sink/lancedb/mod.rs | 0 .../src}/streaming/operators/sink/mod.rs | 0 .../src}/streaming/operators/sink/s3/mod.rs | 0 .../streaming/operators/source/kafka/mod.rs | 0 .../src}/streaming/operators/source/mod.rs | 0 .../operators/stateless_physical_executor.rs | 0 .../streaming/operators/value_execution.rs | 0 .../src}/streaming/operators/watermark/mod.rs | 0 .../watermark/watermark_generator.rs | 0 .../src}/streaming/operators/windows/mod.rs | 0 .../windows/session_aggregating_window.rs | 0 .../windows/sliding_aggregating_window.rs | 0 .../windows/tumbling_aggregating_window.rs | 0 .../operators/windows/window_function.rs | 0 .../src}/streaming/protocol/control.rs | 0 .../src}/streaming/protocol/event.rs | 0 .../src}/streaming/protocol/mod.rs | 0 .../src}/streaming/state/error.rs | 0 .../src}/streaming/state/io_manager.rs | 0 .../src}/streaming/state/metrics.rs | 0 .../src}/streaming/state/mod.rs | 0 .../src}/streaming/state/operator_state.rs | 0 .../src}/util/mod.rs | 0 .../src}/util/physical_aggregate.rs | 0 src/wasm_runtime/src/lib.rs | 19 + .../src}/state_backend/error.rs | 0 .../src}/state_backend/factory.rs | 0 .../src}/state_backend/key_builder.rs | 0 .../src}/state_backend/memory/factory.rs | 0 .../src}/state_backend/memory/mod.rs | 0 .../src}/state_backend/memory/store.rs | 0 .../src}/state_backend/mod.rs | 0 .../src}/state_backend/rocksdb/factory.rs | 0 .../src}/state_backend/rocksdb/mod.rs | 0 .../src}/state_backend/rocksdb/store.rs | 0 .../src}/state_backend/server.rs | 0 .../src}/state_backend/store.rs | 0 .../wasm/buffer_and_event/buffer_or_event.rs | 0 .../src}/wasm/buffer_and_event/mod.rs | 0 .../buffer_and_event/stream_element/mod.rs | 0 .../stream_element/stream_element.rs | 0 .../src}/wasm/input/input_protocol.rs | 0 .../src}/wasm/input/input_provider.rs | 0 .../src}/wasm/input/input_runner.rs | 0 .../src}/wasm/input/interface.rs | 0 .../src}/wasm/input/mod.rs | 0 .../src}/wasm/input/protocol/kafka/config.rs | 0 .../input/protocol/kafka/kafka_protocol.rs | 0 .../src}/wasm/input/protocol/kafka/mod.rs | 0 .../src}/wasm/input/protocol/mod.rs | 0 src/{runtime => wasm_runtime/src}/wasm/mod.rs | 0 .../src}/wasm/output/interface.rs | 0 .../src}/wasm/output/mod.rs | 0 .../src}/wasm/output/output_protocol.rs | 0 .../src}/wasm/output/output_provider.rs | 0 .../src}/wasm/output/output_runner.rs | 0 .../output/protocol/kafka/kafka_protocol.rs | 0 .../src}/wasm/output/protocol/kafka/mod.rs | 0 .../output/protocol/kafka/producer_config.rs | 0 .../src}/wasm/output/protocol/mod.rs | 0 .../src}/wasm/processor/function_error.rs | 0 .../src}/wasm/processor/mod.rs | 0 .../src}/wasm/processor/python/mod.rs | 0 .../src}/wasm/processor/python/python_host.rs | 0 .../wasm/processor/python/python_service.rs | 0 .../wasm/processor/wasm/input_strategy.rs | 0 .../src}/wasm/processor/wasm/mod.rs | 0 .../src}/wasm/processor/wasm/thread_pool.rs | 0 .../src}/wasm/processor/wasm/wasm_cache.rs | 0 .../src}/wasm/processor/wasm/wasm_host.rs | 0 .../wasm/processor/wasm/wasm_processor.rs | 0 .../processor/wasm/wasm_processor_trait.rs | 0 .../src}/wasm/processor/wasm/wasm_task.rs | 0 .../src}/wasm/task/builder/mod.rs | 0 .../src}/wasm/task/builder/processor/mod.rs | 0 .../src}/wasm/task/builder/python/mod.rs | 0 .../src}/wasm/task/builder/sink/mod.rs | 0 .../src}/wasm/task/builder/source/mod.rs | 0 .../src}/wasm/task/builder/task_builder.rs | 0 .../src}/wasm/task/control_mailbox.rs | 0 .../src}/wasm/task/lifecycle.rs | 0 .../src}/wasm/task/mod.rs | 0 .../src}/wasm/task/processor_config.rs | 0 .../src}/wasm/task/yaml_keys.rs | 0 .../src}/wasm/taskexecutor/init_context.rs | 0 .../src}/wasm/taskexecutor/mod.rs | 0 .../src}/wasm/taskexecutor/task_manager.rs | 0 174 files changed, 468 insertions(+), 349 deletions(-) rename src/{storage => catalog_storage/src}/stream_catalog/codec.rs (100%) rename src/{storage => catalog_storage/src}/stream_catalog/manager.rs (100%) rename src/{storage => catalog_storage/src}/stream_catalog/meta_store.rs (100%) rename src/{storage => catalog_storage/src}/stream_catalog/mod.rs (100%) rename src/{storage => catalog_storage/src}/stream_catalog/rocksdb_meta_store.rs (100%) rename src/{storage => catalog_storage/src}/task/factory.rs (100%) rename src/{storage => catalog_storage/src}/task/function_info.rs (100%) rename src/{storage => catalog_storage/src}/task/mod.rs (100%) rename src/{storage => catalog_storage/src}/task/proto_codec.rs (100%) rename src/{storage => catalog_storage/src}/task/rocksdb_storage.rs (100%) rename src/{storage => catalog_storage/src}/task/storage.rs (100%) create mode 100644 src/coordinator/src/legacy/sql_classify.rs create mode 100644 src/runtime_common/Cargo.toml rename src/{runtime => runtime_common/src}/common/component_state.rs (100%) rename src/{runtime => runtime_common/src}/common/mod.rs (100%) rename src/{runtime => runtime_common/src}/common/task_completion.rs (100%) create mode 100644 src/runtime_common/src/lib.rs rename src/{runtime => runtime_common/src}/memory/block.rs (100%) rename src/{runtime => runtime_common/src}/memory/error.rs (100%) rename src/{runtime => runtime_common/src}/memory/global.rs (100%) rename src/{runtime => runtime_common/src}/memory/mod.rs (100%) rename src/{runtime => runtime_common/src}/memory/pool.rs (100%) rename src/{runtime => runtime_common/src}/memory/ticket.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/api/context.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/api/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/api/operator.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/api/source.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/error.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/execution/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/execution/operator_chain.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/execution/pipeline.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/execution/source_driver.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/execution/tracker/barrier_aligner.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/execution/tracker/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/execution/tracker/watermark_tracker.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/factory/connector/delta.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/factory/connector/dispatchers.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/factory/connector/filesystem.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/factory/connector/iceberg.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/factory/connector/kafka.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/factory/connector/lancedb.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/factory/connector/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/factory/connector/s3.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/factory/connector/sink_props_codec.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/factory/global/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/factory/global/session_registry.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/factory/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/factory/operator_constructor.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/factory/operator_factory.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/format/config.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/format/deserializer.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/format/encoder.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/format/json_encoder.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/format/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/format/serializer.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/job/edge_manager.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/job/job_manager.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/job/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/job/models.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/network/endpoint.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/network/environment.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/network/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/grouping/incremental_aggregate.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/grouping/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/grouping/updating_cache.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/joins/join_instance.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/joins/join_with_expiration.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/joins/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/key_by.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/key_operator.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/projection.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/sink/delta/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/sink/filesystem/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/sink/iceberg/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/sink/kafka/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/sink/lancedb/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/sink/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/sink/s3/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/source/kafka/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/source/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/stateless_physical_executor.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/value_execution.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/watermark/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/watermark/watermark_generator.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/windows/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/windows/session_aggregating_window.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/windows/sliding_aggregating_window.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/windows/tumbling_aggregating_window.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/operators/windows/window_function.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/protocol/control.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/protocol/event.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/protocol/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/state/error.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/state/io_manager.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/state/metrics.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/state/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/streaming/state/operator_state.rs (100%) rename src/{runtime => streaming_runtime/src}/util/mod.rs (100%) rename src/{runtime => streaming_runtime/src}/util/physical_aggregate.rs (100%) rename src/{storage => wasm_runtime/src}/state_backend/error.rs (100%) rename src/{storage => wasm_runtime/src}/state_backend/factory.rs (100%) rename src/{storage => wasm_runtime/src}/state_backend/key_builder.rs (100%) rename src/{storage => wasm_runtime/src}/state_backend/memory/factory.rs (100%) rename src/{storage => wasm_runtime/src}/state_backend/memory/mod.rs (100%) rename src/{storage => wasm_runtime/src}/state_backend/memory/store.rs (100%) rename src/{storage => wasm_runtime/src}/state_backend/mod.rs (100%) rename src/{storage => wasm_runtime/src}/state_backend/rocksdb/factory.rs (100%) rename src/{storage => wasm_runtime/src}/state_backend/rocksdb/mod.rs (100%) rename src/{storage => wasm_runtime/src}/state_backend/rocksdb/store.rs (100%) rename src/{storage => wasm_runtime/src}/state_backend/server.rs (100%) rename src/{storage => wasm_runtime/src}/state_backend/store.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/buffer_and_event/buffer_or_event.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/buffer_and_event/mod.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/buffer_and_event/stream_element/mod.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/buffer_and_event/stream_element/stream_element.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/input/input_protocol.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/input/input_provider.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/input/input_runner.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/input/interface.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/input/mod.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/input/protocol/kafka/config.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/input/protocol/kafka/kafka_protocol.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/input/protocol/kafka/mod.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/input/protocol/mod.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/mod.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/output/interface.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/output/mod.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/output/output_protocol.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/output/output_provider.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/output/output_runner.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/output/protocol/kafka/kafka_protocol.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/output/protocol/kafka/mod.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/output/protocol/kafka/producer_config.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/output/protocol/mod.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/processor/function_error.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/processor/mod.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/processor/python/mod.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/processor/python/python_host.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/processor/python/python_service.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/processor/wasm/input_strategy.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/processor/wasm/mod.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/processor/wasm/thread_pool.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/processor/wasm/wasm_cache.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/processor/wasm/wasm_host.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/processor/wasm/wasm_processor.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/processor/wasm/wasm_processor_trait.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/processor/wasm/wasm_task.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/task/builder/mod.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/task/builder/processor/mod.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/task/builder/python/mod.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/task/builder/sink/mod.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/task/builder/source/mod.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/task/builder/task_builder.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/task/control_mailbox.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/task/lifecycle.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/task/mod.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/task/processor_config.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/task/yaml_keys.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/taskexecutor/init_context.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/taskexecutor/mod.rs (100%) rename src/{runtime => wasm_runtime/src}/wasm/taskexecutor/task_manager.rs (100%) diff --git a/Cargo.lock b/Cargo.lock index 6bb2fb4d..05587163 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3802,6 +3802,7 @@ dependencies = [ "datafusion-proto", "function-stream-config", "function-stream-logger", + "function-stream-runtime-common", "futures", "governor", "itertools 0.14.0", @@ -3904,6 +3905,17 @@ dependencies = [ "tracing-subscriber", ] +[[package]] +name = "function-stream-runtime-common" +version = "0.6.0" +dependencies = [ + "arrow-array 55.2.0", + "parking_lot", + "serde", + "tokio", + "tracing", +] + [[package]] name = "function-stream-servicer" version = "0.6.0" diff --git a/Cargo.toml b/Cargo.toml index 95fbe566..3dbd8f5c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,6 +10,7 @@ members = [ "src/coordinator", "src/job_manager", "src/logger", + "src/runtime_common", "src/servicer", "src/sqlparser", "src/streaming_runtime", @@ -48,6 +49,7 @@ num_cpus = "1.0" protocol = { path = "./protocol" } function-stream-config = { path = "src/config" } function-stream-logger = { path = "src/logger" } +function-stream-runtime-common = { path = "src/runtime_common" } prost = "0.13" rdkafka = { version = "0.38", features = ["cmake-build", "ssl", "gssapi", "curl"] } crossbeam-channel = "0.5" diff --git a/src/catalog_storage/src/lib.rs b/src/catalog_storage/src/lib.rs index 6ab4891a..da6e8fc2 100644 --- a/src/catalog_storage/src/lib.rs +++ b/src/catalog_storage/src/lib.rs @@ -1,4 +1,7 @@ //! Persistent catalog storage implementations. +//! +//! The stream catalog manager and task persistence (`stream_catalog/`, `task/`) live in this +//! package and are compiled as part of `function-stream` via `#[path]` in `src/storage/mod.rs`. pub mod memory; pub mod rocksdb; diff --git a/src/storage/stream_catalog/codec.rs b/src/catalog_storage/src/stream_catalog/codec.rs similarity index 100% rename from src/storage/stream_catalog/codec.rs rename to src/catalog_storage/src/stream_catalog/codec.rs diff --git a/src/storage/stream_catalog/manager.rs b/src/catalog_storage/src/stream_catalog/manager.rs similarity index 100% rename from src/storage/stream_catalog/manager.rs rename to src/catalog_storage/src/stream_catalog/manager.rs diff --git a/src/storage/stream_catalog/meta_store.rs b/src/catalog_storage/src/stream_catalog/meta_store.rs similarity index 100% rename from src/storage/stream_catalog/meta_store.rs rename to src/catalog_storage/src/stream_catalog/meta_store.rs diff --git a/src/storage/stream_catalog/mod.rs b/src/catalog_storage/src/stream_catalog/mod.rs similarity index 100% rename from src/storage/stream_catalog/mod.rs rename to src/catalog_storage/src/stream_catalog/mod.rs diff --git a/src/storage/stream_catalog/rocksdb_meta_store.rs b/src/catalog_storage/src/stream_catalog/rocksdb_meta_store.rs similarity index 100% rename from src/storage/stream_catalog/rocksdb_meta_store.rs rename to src/catalog_storage/src/stream_catalog/rocksdb_meta_store.rs diff --git a/src/storage/task/factory.rs b/src/catalog_storage/src/task/factory.rs similarity index 100% rename from src/storage/task/factory.rs rename to src/catalog_storage/src/task/factory.rs diff --git a/src/storage/task/function_info.rs b/src/catalog_storage/src/task/function_info.rs similarity index 100% rename from src/storage/task/function_info.rs rename to src/catalog_storage/src/task/function_info.rs diff --git a/src/storage/task/mod.rs b/src/catalog_storage/src/task/mod.rs similarity index 100% rename from src/storage/task/mod.rs rename to src/catalog_storage/src/task/mod.rs diff --git a/src/storage/task/proto_codec.rs b/src/catalog_storage/src/task/proto_codec.rs similarity index 100% rename from src/storage/task/proto_codec.rs rename to src/catalog_storage/src/task/proto_codec.rs diff --git a/src/storage/task/rocksdb_storage.rs b/src/catalog_storage/src/task/rocksdb_storage.rs similarity index 100% rename from src/storage/task/rocksdb_storage.rs rename to src/catalog_storage/src/task/rocksdb_storage.rs diff --git a/src/storage/task/storage.rs b/src/catalog_storage/src/task/storage.rs similarity index 100% rename from src/storage/task/storage.rs rename to src/catalog_storage/src/task/storage.rs diff --git a/src/coordinator/src/legacy/mod.rs b/src/coordinator/src/legacy/mod.rs index 86598bc5..445ec105 100644 --- a/src/coordinator/src/legacy/mod.rs +++ b/src/coordinator/src/legacy/mod.rs @@ -13,6 +13,7 @@ mod analyze; #[allow(clippy::module_inception)] mod coordinator; +mod sql_classify; mod dataset; mod execution; mod execution_context; @@ -24,6 +25,7 @@ mod tool; pub use coordinator::Coordinator; pub use dataset::{DataSet, ShowFunctionsResult}; +pub use sql_classify::classify_statement; pub use statement::{ CreateFunction, CreatePythonFunction, CreateTable, DropFunction, DropStreamingTableStatement, DropTableStatement, PythonModule, ShowCatalogTables, ShowCreateStreamingTable, ShowCreateTable, diff --git a/src/coordinator/src/legacy/sql_classify.rs b/src/coordinator/src/legacy/sql_classify.rs new file mode 100644 index 00000000..af96a66c --- /dev/null +++ b/src/coordinator/src/legacy/sql_classify.rs @@ -0,0 +1,326 @@ +// 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. + +//! Map sqlparser [`Statement`](datafusion::sql::sqlparser::ast::Statement) values into +//! coordinator [`Statement`](super::statement::Statement) trait objects. + +use std::collections::HashMap; + +use datafusion::common::{Result, plan_err}; +use datafusion::error::DataFusionError; +use datafusion::sql::sqlparser::ast::{ + ObjectType, ShowCreateObject, SqlOption, Statement as DFStatement, +}; + +use super::statement::{ + CreateFunction, CreateTable, DropFunction, DropStreamingTableStatement, DropTableStatement, + ShowCatalogTables, ShowCreateStreamingTable, ShowCreateTable, ShowFunctions, + ShowStreamingTables, StartFunction, Statement, StopFunction, StreamingTableStatement, +}; + +/// Convert [`DFStatement`] from the FunctionStream SQL dialect into a coordinator statement. +pub fn classify_statement(stmt: DFStatement) -> Result> { + match stmt { + DFStatement::CreateFunctionWith { options } => { + let properties = sql_options_to_map(&options); + let create_fn = CreateFunction::from_properties(properties) + .map_err(|e| DataFusionError::Plan(format!("CREATE FUNCTION: {e}")))?; + Ok(Box::new(create_fn)) + } + DFStatement::StartFunction { name } => Ok(Box::new(StartFunction::new(name.to_string()))), + DFStatement::StopFunction { name } => Ok(Box::new(StopFunction::new(name.to_string()))), + DFStatement::DropFunction { func_desc, .. } => { + let name = func_desc + .first() + .map(|d| d.name.to_string()) + .unwrap_or_default(); + Ok(Box::new(DropFunction::new(name))) + } + DFStatement::ShowFunctions { .. } => Ok(Box::new(ShowFunctions::new())), + DFStatement::ShowTables { .. } => Ok(Box::new(ShowCatalogTables::new())), + DFStatement::ShowStreamingTable => Ok(Box::new(ShowStreamingTables::new())), + DFStatement::ShowCreate { obj_type, obj_name } => match obj_type { + ShowCreateObject::Table => Ok(Box::new(ShowCreateTable::new(obj_name.to_string()))), + ShowCreateObject::StreamingTable => { + Ok(Box::new(ShowCreateStreamingTable::new(obj_name.to_string()))) + } + _ => plan_err!( + "SHOW CREATE {obj_type} is not supported; use SHOW CREATE TABLE or SHOW CREATE STREAMING TABLE " + ), + }, + s @ DFStatement::CreateTable(_) => Ok(Box::new(CreateTable::new(s))), + s @ DFStatement::CreateStreamingTable { .. } => { + Ok(Box::new(StreamingTableStatement::new(s))) + } + stmt @ DFStatement::Drop { .. } => { + let DFStatement::Drop { + object_type, + names, + if_exists, + .. + } = &stmt + else { + unreachable!() + }; + match object_type { + ObjectType::Table => { + if names.len() != 1 { + return plan_err!( + "DROP TABLE supports exactly one table name per statement" + ); + } + Ok(Box::new(DropTableStatement::new(stmt))) + } + ObjectType::StreamingTable => { + if names.len() != 1 { + return plan_err!( + "DROP STREAMING TABLE supports exactly one table name per statement" + ); + } + let table_name = names[0].to_string(); + Ok(Box::new(DropStreamingTableStatement::new( + table_name, + *if_exists, + ))) + } + _ => plan_err!("Only DROP TABLE and DROP STREAMING TABLE are supported in this SQL frontend"), + } + } + DFStatement::Insert { .. } => plan_err!( + "INSERT is not supported; only CREATE TABLE and CREATE STREAMING TABLE (with AS SELECT) \ + are supported for defining table/query pipelines in this SQL frontend" + ), + other => plan_err!("Unsupported SQL statement: {other}"), + } +} + +fn sql_options_to_map(options: &[SqlOption]) -> HashMap { + options + .iter() + .filter_map(|opt| match opt { + SqlOption::KeyValue { key, value } => Some(( + key.value.clone(), + value.to_string().trim_matches('\'').to_string(), + )), + _ => None, + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::sql::parse::parse_sql; + + fn first_classified(sql: &str) -> Box { + let mut stmts = parse_sql(sql).unwrap(); + assert!(!stmts.is_empty()); + classify_statement(stmts.remove(0)).unwrap() + } + + fn is_type(stmt: &dyn Statement, prefix: &str) -> bool { + format!("{stmt:?}").starts_with(prefix) + } + + #[test] + fn test_parse_create_function() { + let sql = + "CREATE FUNCTION WITH ('function_path'='./test.wasm', 'config_path'='./config.yml')"; + let stmt = first_classified(sql); + assert!(is_type(stmt.as_ref(), "CreateFunction")); + } + + #[test] + fn test_parse_create_function_minimal() { + let sql = "CREATE FUNCTION WITH ('function_path'='./processor.wasm')"; + let stmt = first_classified(sql); + assert!(is_type(stmt.as_ref(), "CreateFunction")); + } + + #[test] + fn test_parse_drop_function() { + let stmt = first_classified("DROP FUNCTION my_task"); + assert!(is_type(stmt.as_ref(), "DropFunction")); + } + + #[test] + fn test_parse_start_function() { + let stmt = first_classified("START FUNCTION my_task"); + assert!(is_type(stmt.as_ref(), "StartFunction")); + } + + #[test] + fn test_parse_stop_function() { + let stmt = first_classified("STOP FUNCTION my_task"); + assert!(is_type(stmt.as_ref(), "StopFunction")); + } + + #[test] + fn test_parse_show_functions() { + let stmt = first_classified("SHOW FUNCTIONS"); + assert!(is_type(stmt.as_ref(), "ShowFunctions")); + } + + #[test] + fn test_parse_show_tables() { + let stmt = first_classified("SHOW TABLES"); + assert!(is_type(stmt.as_ref(), "ShowCatalogTables")); + } + + #[test] + fn test_parse_show_create_table() { + let stmt = first_classified("SHOW CREATE TABLE my_src"); + assert!(is_type(stmt.as_ref(), "ShowCreateTable")); + } + + #[test] + fn test_parse_create_table() { + let stmt = first_classified("CREATE TABLE foo (id INT, name VARCHAR)"); + assert!(is_type(stmt.as_ref(), "CreateTable")); + } + + #[test] + fn test_parse_create_table_connector_source_ddl() { + let sql = concat!( + "CREATE TABLE kafka_src (id BIGINT, ts TIMESTAMP NOT NULL, WATERMARK FOR ts) ", + "WITH ('connector' = 'kafka', 'format' = 'json', 'topic' = 'events')", + ); + let stmt = first_classified(sql); + assert!(is_type(stmt.as_ref(), "CreateTable")); + } + + #[test] + fn test_parse_drop_table() { + let stmt = first_classified("DROP TABLE foo"); + assert!(is_type(stmt.as_ref(), "DropTableStatement")); + } + + #[test] + fn test_parse_drop_table_if_exists() { + let stmt = first_classified("DROP TABLE IF EXISTS foo"); + assert!(is_type(stmt.as_ref(), "DropTableStatement")); + } + + #[test] + fn test_parse_drop_streaming_table() { + let stmt = first_classified("DROP STREAMING TABLE my_sink"); + assert!(is_type(stmt.as_ref(), "DropStreamingTableStatement")); + } + + #[test] + fn test_parse_drop_streaming_table_if_exists() { + let stmt = first_classified("DROP STREAMING TABLE IF EXISTS my_sink"); + assert!(is_type(stmt.as_ref(), "DropStreamingTableStatement")); + } + + #[test] + fn test_parse_show_streaming_tables() { + let stmt = first_classified("SHOW STREAMING TABLES"); + assert!(is_type(stmt.as_ref(), "ShowStreamingTables")); + } + + #[test] + fn test_parse_show_create_streaming_table() { + let stmt = first_classified("SHOW CREATE STREAMING TABLE my_sink"); + assert!(is_type(stmt.as_ref(), "ShowCreateStreamingTable")); + } + + #[test] + fn test_parse_create_streaming_table() { + let sql = concat!( + "CREATE STREAMING TABLE my_sink ", + "WITH ('connector' = 'kafka') ", + "AS SELECT id FROM src", + ); + let stmt = first_classified(sql); + assert!( + is_type(stmt.as_ref(), "StreamingTableStatement"), + "expected StreamingTableStatement, got {:?}", + stmt + ); + } + + #[test] + fn test_parse_create_streaming_table_case_insensitive() { + let sql = concat!( + "create streaming table out_q ", + "with ('connector' = 'memory') ", + "as select 1 as x", + ); + let stmt = first_classified(sql); + assert!(is_type(stmt.as_ref(), "StreamingTableStatement")); + } + + #[test] + fn test_parse_case_insensitive() { + assert!(is_type( + first_classified("create function with ('function_path'='./test.wasm')").as_ref(), + "CreateFunction" + )); + assert!(is_type( + first_classified("show functions").as_ref(), + "ShowFunctions" + )); + assert!(is_type( + first_classified("start function my_task").as_ref(), + "StartFunction" + )); + } + + #[test] + fn test_parse_multiple_statements() { + let sql = concat!( + "CREATE TABLE t1 (id INT); ", + "CREATE STREAMING TABLE sk WITH ('connector' = 'kafka') AS SELECT id FROM t1", + ); + let mut ast = parse_sql(sql).unwrap(); + assert_eq!(ast.len(), 2); + let s0 = classify_statement(ast.remove(0)).unwrap(); + let s1 = classify_statement(ast.remove(0)).unwrap(); + assert!(is_type(s0.as_ref(), "CreateTable")); + assert!(is_type(s1.as_ref(), "StreamingTableStatement")); + } + + #[test] + fn test_classify_unsupported_statement() { + let mut stmts = parse_sql("SELECT 1").unwrap(); + assert!(classify_statement(stmts.remove(0)).is_err()); + } + + #[test] + fn test_insert_not_supported() { + let mut stmts = + parse_sql("INSERT INTO sink SELECT * FROM src").unwrap(); + let err = classify_statement(stmts.remove(0)).unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("INSERT") && msg.contains("not supported"), + "expected explicit INSERT rejection, got: {msg}" + ); + assert!( + msg.contains("CREATE TABLE") || msg.contains("CREATE STREAMING TABLE"), + "error should mention supported alternatives, got: {msg}" + ); + } + + #[test] + fn test_parse_with_extra_properties() { + let sql = r#"CREATE FUNCTION WITH ( + 'function_path'='./test.wasm', + 'config_path'='./config.yml', + 'parallelism'='4', + 'memory-limit'='256mb' + )"#; + let stmt = first_classified(sql); + assert!(is_type(stmt.as_ref(), "CreateFunction")); + } +} diff --git a/src/main.rs b/src/main.rs index ee51cb3e..49978af3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -146,7 +146,8 @@ fn setup_environment() -> Result { config::GlobalConfig::default() }; - function_stream_logger::init_logging(&config.logging).context("Logging initialization failed")?; + function_stream_logger::init_logging(&config.logging) + .context("Logging initialization failed")?; log::debug!( "Environment initialized. Data: {}, Conf: {}", diff --git a/src/runtime/mod.rs b/src/runtime/mod.rs index 8c72b507..28f62306 100644 --- a/src/runtime/mod.rs +++ b/src/runtime/mod.rs @@ -12,10 +12,18 @@ // Runtime module -pub mod common; -pub mod memory; +pub use function_stream_runtime_common::{common, memory}; + +// Implementation sources live under `src/streaming_runtime/src/{streaming,util}/` and are +// compiled here so `crate::sql` / `crate::runtime::memory` paths keep resolving. +#[path = "../streaming_runtime/src/streaming/mod.rs"] pub mod streaming; + +#[path = "../streaming_runtime/src/util/mod.rs"] pub mod util; + +// WASM runtime sources live under `src/wasm_runtime/src/wasm/`; compiled here for `crate::` paths. +#[path = "../wasm_runtime/src/wasm/mod.rs"] pub mod wasm; pub use wasm::input; diff --git a/src/runtime_common/Cargo.toml b/src/runtime_common/Cargo.toml new file mode 100644 index 00000000..0a910f18 --- /dev/null +++ b/src/runtime_common/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "function-stream-runtime-common" +version = "0.6.0" +edition = "2024" + +[lib] +name = "function_stream_runtime_common" +path = "src/lib.rs" + +[dependencies] +arrow-array = "55" +parking_lot = "0.12" +serde = { version = "1.0", features = ["derive"] } +tokio = { version = "1.0", features = ["sync"] } +tracing = "0.1" diff --git a/src/runtime/common/component_state.rs b/src/runtime_common/src/common/component_state.rs similarity index 100% rename from src/runtime/common/component_state.rs rename to src/runtime_common/src/common/component_state.rs diff --git a/src/runtime/common/mod.rs b/src/runtime_common/src/common/mod.rs similarity index 100% rename from src/runtime/common/mod.rs rename to src/runtime_common/src/common/mod.rs diff --git a/src/runtime/common/task_completion.rs b/src/runtime_common/src/common/task_completion.rs similarity index 100% rename from src/runtime/common/task_completion.rs rename to src/runtime_common/src/common/task_completion.rs diff --git a/src/runtime_common/src/lib.rs b/src/runtime_common/src/lib.rs new file mode 100644 index 00000000..3fad765f --- /dev/null +++ b/src/runtime_common/src/lib.rs @@ -0,0 +1,16 @@ +// 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. + +//! Shared runtime building blocks (task lifecycle, memory pool) used by streaming and WASM. + +pub mod common; +pub mod memory; diff --git a/src/runtime/memory/block.rs b/src/runtime_common/src/memory/block.rs similarity index 100% rename from src/runtime/memory/block.rs rename to src/runtime_common/src/memory/block.rs diff --git a/src/runtime/memory/error.rs b/src/runtime_common/src/memory/error.rs similarity index 100% rename from src/runtime/memory/error.rs rename to src/runtime_common/src/memory/error.rs diff --git a/src/runtime/memory/global.rs b/src/runtime_common/src/memory/global.rs similarity index 100% rename from src/runtime/memory/global.rs rename to src/runtime_common/src/memory/global.rs diff --git a/src/runtime/memory/mod.rs b/src/runtime_common/src/memory/mod.rs similarity index 100% rename from src/runtime/memory/mod.rs rename to src/runtime_common/src/memory/mod.rs diff --git a/src/runtime/memory/pool.rs b/src/runtime_common/src/memory/pool.rs similarity index 100% rename from src/runtime/memory/pool.rs rename to src/runtime_common/src/memory/pool.rs diff --git a/src/runtime/memory/ticket.rs b/src/runtime_common/src/memory/ticket.rs similarity index 100% rename from src/runtime/memory/ticket.rs rename to src/runtime_common/src/memory/ticket.rs diff --git a/src/servicer/src/legacy/handler.rs b/src/servicer/src/legacy/handler.rs index 0319e352..c884eb6f 100644 --- a/src/servicer/src/legacy/handler.rs +++ b/src/servicer/src/legacy/handler.rs @@ -26,7 +26,7 @@ use protocol::service::{ use crate::coordinator::{ Coordinator, CreateFunction, CreatePythonFunction, DataSet, DropFunction, PythonModule, - ShowFunctions, ShowFunctionsResult, StartFunction, Statement, StopFunction, + ShowFunctions, ShowFunctionsResult, StartFunction, Statement, StopFunction, classify_statement, }; use crate::sql::parse::parse_sql; @@ -135,12 +135,20 @@ impl FunctionStreamService for FunctionStreamServiceImpl { let timer = Instant::now(); let req = request.into_inner(); - let statements = parse_sql(&req.sql).map_err(|e| { + let ast = parse_sql(&req.sql).map_err(|e| { let detail = e.to_string(); warn!("SQL parse rejection: {}", detail); Status::invalid_argument(detail) })?; + let statements: Result>, _> = + ast.into_iter().map(classify_statement).collect(); + let statements = statements.map_err(|e| { + let detail = e.to_string(); + warn!("SQL classification rejection: {}", detail); + Status::invalid_argument(detail) + })?; + if statements.is_empty() { return Ok(TonicResponse::new(Self::build_success_response( StatusCode::Ok, diff --git a/src/sql/parse.rs b/src/sql/parse.rs index 0c6b9541..8c9d4bb0 100644 --- a/src/sql/parse.rs +++ b/src/sql/parse.rs @@ -10,105 +10,37 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Coordinator-facing SQL parsing (`parse_sql`). +//! FunctionStream SQL parsing (`parse_sql`). //! -//! **Data-definition / pipeline shape (this entry point)** -//! Only these table-related forms are supported: -//! - **`CREATE TABLE ... (cols [, WATERMARK FOR ...]) WITH ('connector' = '...', 'format' = '...', ...)`** -//! connector-backed **source** DDL (no `AS SELECT`; `connector` in `WITH` selects this path) -//! - **`CREATE TABLE ...`** other forms (including `CREATE TABLE ... AS SELECT` where DataFusion accepts it) -//! - **`CREATE STREAMING TABLE ... WITH (...) AS SELECT ...`** (streaming sink DDL) -//! - **`DROP TABLE`** / **`DROP TABLE IF EXISTS`** / **`DROP STREAMING TABLE`** (alias for `DROP TABLE` on the stream catalog) -//! - **`SHOW TABLES`** — list stream catalog tables (connector sources and streaming sinks) -//! - **`SHOW CREATE TABLE `** — best-effort DDL text (full `WITH` / `AS SELECT` may not be stored) +//! This module only performs lexical/syntactic parsing into sqlparser +//! [`Statement`](datafusion::sql::sqlparser::ast::Statement) values using +//! [`FunctionStreamDialect`]. Mapping those AST nodes to coordinator +//! [`Statement`](crate::coordinator::Statement) implementations is done by +//! [`crate::coordinator::classify_statement`]. //! -//! **`INSERT` is not supported** here — use `CREATE TABLE ... AS SELECT` or -//! `CREATE STREAMING TABLE ... AS SELECT` to define the query shape instead. +//! **Data-definition / pipeline shape (supported forms in the dialect)** +//! - **`CREATE TABLE ... (cols [, WATERMARK FOR ...]) WITH (...)`** — connector-backed source DDL +//! - **`CREATE TABLE ...`** other forms (including `AS SELECT` where the dialect accepts it) +//! - **`CREATE STREAMING TABLE ... WITH (...) AS SELECT ...`** +//! - **`DROP TABLE`** / **`DROP STREAMING TABLE`** +//! - **`SHOW TABLES`**, **`SHOW STREAMING TABLE(S)`**, **`SHOW CREATE TABLE`**, **`SHOW CREATE STREAMING TABLE`** //! -//! Other supported statements include function lifecycle (`CREATE FUNCTION WITH`, `START FUNCTION`, …). - -use std::collections::HashMap; +//! **`INSERT` is not supported** at the coordinator layer — use `CREATE TABLE ... AS SELECT` or +//! `CREATE STREAMING TABLE ... AS SELECT` instead (see coordinator classification). use datafusion::common::{Result, plan_err}; use datafusion::error::DataFusionError; -use datafusion::sql::sqlparser::ast::{ - ObjectType, ShowCreateObject, SqlOption, Statement as DFStatement, -}; +use datafusion::sql::sqlparser::ast::Statement as DFStatement; use datafusion::sql::sqlparser::dialect::FunctionStreamDialect; use datafusion::sql::sqlparser::parser::Parser; -use crate::coordinator::{ - CreateFunction, CreateTable, DropFunction, DropStreamingTableStatement, DropTableStatement, - ShowCatalogTables, ShowCreateStreamingTable, ShowCreateTable, ShowFunctions, - ShowStreamingTables, StartFunction, Statement as CoordinatorStatement, StopFunction, - StreamingTableStatement, -}; - -/// Streaming-specific SQL that the sqlparser dialect does not natively handle. -/// -/// Returns `Some(statement)` if the SQL was intercepted, `None` otherwise so -/// the caller falls through to the normal sqlparser pipeline. -fn try_parse_streaming_statement(sql: &str) -> Option> { - let tokens: Vec<&str> = sql.split_whitespace().collect(); - if tokens.is_empty() { - return None; - } - - // SHOW STREAMING TABLES - if tokens.len() == 3 - && tokens[0].eq_ignore_ascii_case("show") - && tokens[1].eq_ignore_ascii_case("streaming") - && tokens[2].eq_ignore_ascii_case("tables") - { - return Some(Box::new(ShowStreamingTables::new())); - } - - // SHOW CREATE STREAMING TABLE - if tokens.len() == 5 - && tokens[0].eq_ignore_ascii_case("show") - && tokens[1].eq_ignore_ascii_case("create") - && tokens[2].eq_ignore_ascii_case("streaming") - && tokens[3].eq_ignore_ascii_case("table") - { - let name = tokens[4].trim_end_matches(';').to_string(); - return Some(Box::new(ShowCreateStreamingTable::new(name))); - } - - // DROP STREAMING TABLE [IF EXISTS] - if tokens.len() >= 4 - && tokens[0].eq_ignore_ascii_case("drop") - && tokens[1].eq_ignore_ascii_case("streaming") - && tokens[2].eq_ignore_ascii_case("table") - { - let (if_exists, name_idx) = if tokens.len() >= 6 - && tokens[3].eq_ignore_ascii_case("if") - && tokens[4].eq_ignore_ascii_case("exists") - { - (true, 5) - } else { - (false, 3) - }; - - if name_idx >= tokens.len() { - return None; - } - let name = tokens[name_idx].trim_end_matches(';').to_string(); - return Some(Box::new(DropStreamingTableStatement::new(name, if_exists))); - } - - None -} - -pub fn parse_sql(query: &str) -> Result>> { +/// Parse SQL text into zero or more dialect [`Statement`](DFStatement) nodes. +pub fn parse_sql(query: &str) -> Result> { let trimmed = query.trim(); if trimmed.is_empty() { return plan_err!("Query is empty"); } - if let Some(stmt) = try_parse_streaming_statement(trimmed) { - return Ok(vec![stmt]); - } - let dialect = FunctionStreamDialect {}; let statements = Parser::parse_sql(&dialect, trimmed) .map_err(|e| DataFusionError::Plan(format!("SQL parse error: {e}")))?; @@ -117,249 +49,23 @@ pub fn parse_sql(query: &str) -> Result>> { return plan_err!("No SQL statements found"); } - statements.into_iter().map(classify_statement).collect() -} - -fn classify_statement(stmt: DFStatement) -> Result> { - match stmt { - DFStatement::CreateFunctionWith { options } => { - let properties = sql_options_to_map(&options); - let create_fn = CreateFunction::from_properties(properties) - .map_err(|e| DataFusionError::Plan(format!("CREATE FUNCTION: {e}")))?; - Ok(Box::new(create_fn)) - } - DFStatement::StartFunction { name } => Ok(Box::new(StartFunction::new(name.to_string()))), - DFStatement::StopFunction { name } => Ok(Box::new(StopFunction::new(name.to_string()))), - DFStatement::DropFunction { func_desc, .. } => { - let name = func_desc - .first() - .map(|d| d.name.to_string()) - .unwrap_or_default(); - Ok(Box::new(DropFunction::new(name))) - } - DFStatement::ShowFunctions { .. } => Ok(Box::new(ShowFunctions::new())), - DFStatement::ShowTables { .. } => Ok(Box::new(ShowCatalogTables::new())), - DFStatement::ShowCreate { obj_type, obj_name } => { - if obj_type != ShowCreateObject::Table { - return plan_err!( - "SHOW CREATE {obj_type} is not supported; use SHOW CREATE TABLE " - ); - } - Ok(Box::new(ShowCreateTable::new(obj_name.to_string()))) - } - s @ DFStatement::CreateTable(_) => Ok(Box::new(CreateTable::new(s))), - s @ DFStatement::CreateStreamingTable { .. } => { - Ok(Box::new(StreamingTableStatement::new(s))) - } - stmt @ DFStatement::Drop { .. } => { - { - let DFStatement::Drop { - object_type, names, .. - } = &stmt - else { - unreachable!() - }; - if *object_type != ObjectType::Table { - return plan_err!("Only DROP TABLE is supported in this SQL frontend"); - } - if names.len() != 1 { - return plan_err!("DROP TABLE supports exactly one table name per statement"); - } - } - Ok(Box::new(DropTableStatement::new(stmt))) - } - DFStatement::Insert { .. } => plan_err!( - "INSERT is not supported; only CREATE TABLE and CREATE STREAMING TABLE (with AS SELECT) \ - are supported for defining table/query pipelines in this SQL frontend" - ), - other => plan_err!("Unsupported SQL statement: {other}"), - } -} - -/// Convert Vec (KeyValue pairs) into HashMap. -fn sql_options_to_map(options: &[SqlOption]) -> HashMap { - options - .iter() - .filter_map(|opt| match opt { - SqlOption::KeyValue { key, value } => Some(( - key.value.clone(), - value.to_string().trim_matches('\'').to_string(), - )), - _ => None, - }) - .collect() + Ok(statements) } #[cfg(test)] mod tests { use super::*; - fn first_stmt(sql: &str) -> Box { - let mut stmts = parse_sql(sql).unwrap(); - assert!(!stmts.is_empty()); - stmts.remove(0) - } - - fn is_type(stmt: &dyn CoordinatorStatement, prefix: &str) -> bool { - format!("{:?}", stmt).starts_with(prefix) - } - - #[test] - fn test_parse_create_function() { - let sql = - "CREATE FUNCTION WITH ('function_path'='./test.wasm', 'config_path'='./config.yml')"; - let stmt = first_stmt(sql); - assert!(is_type(stmt.as_ref(), "CreateFunction")); - } - - #[test] - fn test_parse_create_function_minimal() { - let sql = "CREATE FUNCTION WITH ('function_path'='./processor.wasm')"; - let stmt = first_stmt(sql); - assert!(is_type(stmt.as_ref(), "CreateFunction")); - } - - #[test] - fn test_parse_drop_function() { - let stmt = first_stmt("DROP FUNCTION my_task"); - assert!(is_type(stmt.as_ref(), "DropFunction")); - } - - #[test] - fn test_parse_start_function() { - let stmt = first_stmt("START FUNCTION my_task"); - assert!(is_type(stmt.as_ref(), "StartFunction")); - } - - #[test] - fn test_parse_stop_function() { - let stmt = first_stmt("STOP FUNCTION my_task"); - assert!(is_type(stmt.as_ref(), "StopFunction")); - } - - #[test] - fn test_parse_show_functions() { - let stmt = first_stmt("SHOW FUNCTIONS"); - assert!(is_type(stmt.as_ref(), "ShowFunctions")); - } - - #[test] - fn test_parse_show_tables() { - let stmt = first_stmt("SHOW TABLES"); - assert!(is_type(stmt.as_ref(), "ShowCatalogTables")); - } - - #[test] - fn test_parse_show_create_table() { - let stmt = first_stmt("SHOW CREATE TABLE my_src"); - assert!(is_type(stmt.as_ref(), "ShowCreateTable")); - } - #[test] - fn test_parse_create_table() { - let stmt = first_stmt("CREATE TABLE foo (id INT, name VARCHAR)"); - assert!(is_type(stmt.as_ref(), "CreateTable")); - } - - #[test] - fn test_parse_create_table_connector_source_ddl() { - let sql = concat!( - "CREATE TABLE kafka_src (id BIGINT, ts TIMESTAMP NOT NULL, WATERMARK FOR ts) ", - "WITH ('connector' = 'kafka', 'format' = 'json', 'topic' = 'events')", - ); - let stmt = first_stmt(sql); - assert!(is_type(stmt.as_ref(), "CreateTable")); - } - - #[test] - fn test_parse_drop_table() { - let stmt = first_stmt("DROP TABLE foo"); - assert!(is_type(stmt.as_ref(), "DropTableStatement")); - } - - #[test] - fn test_parse_drop_table_if_exists() { - let stmt = first_stmt("DROP TABLE IF EXISTS foo"); - assert!(is_type(stmt.as_ref(), "DropTableStatement")); - } - - #[test] - fn test_parse_drop_streaming_table() { - let stmt = first_stmt("DROP STREAMING TABLE my_sink"); - assert!(is_type(stmt.as_ref(), "DropStreamingTableStatement")); - } - - #[test] - fn test_parse_drop_streaming_table_if_exists() { - let stmt = first_stmt("DROP STREAMING TABLE IF EXISTS my_sink"); - assert!(is_type(stmt.as_ref(), "DropStreamingTableStatement")); - } - - #[test] - fn test_parse_show_streaming_tables() { - let stmt = first_stmt("SHOW STREAMING TABLES"); - assert!(is_type(stmt.as_ref(), "ShowStreamingTables")); - } - - #[test] - fn test_parse_show_create_streaming_table() { - let stmt = first_stmt("SHOW CREATE STREAMING TABLE my_sink"); - assert!(is_type(stmt.as_ref(), "ShowCreateStreamingTable")); - } - - /// `CREATE STREAMING TABLE` is the sink DDL supported by FunctionStream (not `CREATE STREAM TABLE`). - #[test] - fn test_parse_create_streaming_table() { - let sql = concat!( - "CREATE STREAMING TABLE my_sink ", - "WITH ('connector' = 'kafka') ", - "AS SELECT id FROM src", - ); - let stmt = first_stmt(sql); - assert!( - is_type(stmt.as_ref(), "StreamingTableStatement"), - "expected StreamingTableStatement, got {:?}", - stmt - ); - } - - #[test] - fn test_parse_create_streaming_table_case_insensitive() { - let sql = concat!( - "create streaming table out_q ", - "with ('connector' = 'memory') ", - "as select 1 as x", - ); - let stmt = first_stmt(sql); - assert!(is_type(stmt.as_ref(), "StreamingTableStatement")); - } - - #[test] - fn test_parse_case_insensitive() { - assert!(is_type( - first_stmt("create function with ('function_path'='./test.wasm')").as_ref(), - "CreateFunction" - )); - assert!(is_type( - first_stmt("show functions").as_ref(), - "ShowFunctions" - )); - assert!(is_type( - first_stmt("start function my_task").as_ref(), - "StartFunction" - )); - } - - #[test] - fn test_parse_multiple_statements() { + fn test_parse_multiple_statements_ast() { let sql = concat!( "CREATE TABLE t1 (id INT); ", "CREATE STREAMING TABLE sk WITH ('connector' = 'kafka') AS SELECT id FROM t1", ); let stmts = parse_sql(sql).unwrap(); assert_eq!(stmts.len(), 2); - assert!(is_type(stmts[0].as_ref(), "CreateTable")); - assert!(is_type(stmts[1].as_ref(), "StreamingTableStatement")); + assert!(matches!(stmts[0], DFStatement::CreateTable(_))); + assert!(matches!(stmts[1], DFStatement::CreateStreamingTable { .. })); } #[test] @@ -369,34 +75,9 @@ mod tests { } #[test] - fn test_parse_unsupported_statement() { - let result = parse_sql("SELECT 1"); - assert!(result.is_err()); - } - - #[test] - fn test_insert_not_supported() { - let err = parse_sql("INSERT INTO sink SELECT * FROM src").unwrap_err(); - let msg = err.to_string(); - assert!( - msg.contains("INSERT") && msg.contains("not supported"), - "expected explicit INSERT rejection, got: {msg}" - ); - assert!( - msg.contains("CREATE TABLE") || msg.contains("CREATE STREAMING TABLE"), - "error should mention supported alternatives, got: {msg}" - ); - } - - #[test] - fn test_parse_with_extra_properties() { - let sql = r#"CREATE FUNCTION WITH ( - 'function_path'='./test.wasm', - 'config_path'='./config.yml', - 'parallelism'='4', - 'memory-limit'='256mb' - )"#; - let stmt = first_stmt(sql); - assert!(is_type(stmt.as_ref(), "CreateFunction")); + fn test_parse_select_yields_query_ast() { + let stmts = parse_sql("SELECT 1").unwrap(); + assert_eq!(stmts.len(), 1); + assert!(matches!(stmts[0], DFStatement::Query(_))); } } diff --git a/src/storage/mod.rs b/src/storage/mod.rs index ec32bdfa..d72e8096 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -14,8 +14,16 @@ use std::sync::Arc; use anyhow::Context; +// State backend sources live under `src/wasm_runtime/src/state_backend/`; compiled here for `crate::storage::state_backend`. +#[path = "../wasm_runtime/src/state_backend/mod.rs"] pub mod state_backend; + +// Stream catalog + task storage sources under `src/catalog_storage/src/{stream_catalog,task}/`; +// compiled here so `crate::storage::stream_catalog` / `crate::storage::task` keep resolving. +#[path = "../catalog_storage/src/stream_catalog/mod.rs"] pub mod stream_catalog; + +#[path = "../catalog_storage/src/task/mod.rs"] pub mod task; /// Install the process-global [`stream_catalog::CatalogManager`] from configuration. diff --git a/src/streaming_runtime/src/lib.rs b/src/streaming_runtime/src/lib.rs index a4643ba9..ad83a58c 100644 --- a/src/streaming_runtime/src/lib.rs +++ b/src/streaming_runtime/src/lib.rs @@ -1,3 +1,21 @@ +// 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. + //! Streaming execution runtime. +//! +//! The streaming engine and shared runtime helpers (`streaming/`, `util/`) are +//! implemented under [`src/streaming`] and [`src/util`] in this package. They are +//! currently **compiled as part of the `function-stream` crate** via `#[path]` in +//! `src/runtime/mod.rs`, so in-tree paths like `crate::sql` +//! keep working until SQL is split into its own crate. pub const CRATE_NAME: &str = "function-stream-streaming-runtime"; diff --git a/src/runtime/streaming/api/context.rs b/src/streaming_runtime/src/streaming/api/context.rs similarity index 100% rename from src/runtime/streaming/api/context.rs rename to src/streaming_runtime/src/streaming/api/context.rs diff --git a/src/runtime/streaming/api/mod.rs b/src/streaming_runtime/src/streaming/api/mod.rs similarity index 100% rename from src/runtime/streaming/api/mod.rs rename to src/streaming_runtime/src/streaming/api/mod.rs diff --git a/src/runtime/streaming/api/operator.rs b/src/streaming_runtime/src/streaming/api/operator.rs similarity index 100% rename from src/runtime/streaming/api/operator.rs rename to src/streaming_runtime/src/streaming/api/operator.rs diff --git a/src/runtime/streaming/api/source.rs b/src/streaming_runtime/src/streaming/api/source.rs similarity index 100% rename from src/runtime/streaming/api/source.rs rename to src/streaming_runtime/src/streaming/api/source.rs diff --git a/src/runtime/streaming/error.rs b/src/streaming_runtime/src/streaming/error.rs similarity index 100% rename from src/runtime/streaming/error.rs rename to src/streaming_runtime/src/streaming/error.rs diff --git a/src/runtime/streaming/execution/mod.rs b/src/streaming_runtime/src/streaming/execution/mod.rs similarity index 100% rename from src/runtime/streaming/execution/mod.rs rename to src/streaming_runtime/src/streaming/execution/mod.rs diff --git a/src/runtime/streaming/execution/operator_chain.rs b/src/streaming_runtime/src/streaming/execution/operator_chain.rs similarity index 100% rename from src/runtime/streaming/execution/operator_chain.rs rename to src/streaming_runtime/src/streaming/execution/operator_chain.rs diff --git a/src/runtime/streaming/execution/pipeline.rs b/src/streaming_runtime/src/streaming/execution/pipeline.rs similarity index 100% rename from src/runtime/streaming/execution/pipeline.rs rename to src/streaming_runtime/src/streaming/execution/pipeline.rs diff --git a/src/runtime/streaming/execution/source_driver.rs b/src/streaming_runtime/src/streaming/execution/source_driver.rs similarity index 100% rename from src/runtime/streaming/execution/source_driver.rs rename to src/streaming_runtime/src/streaming/execution/source_driver.rs diff --git a/src/runtime/streaming/execution/tracker/barrier_aligner.rs b/src/streaming_runtime/src/streaming/execution/tracker/barrier_aligner.rs similarity index 100% rename from src/runtime/streaming/execution/tracker/barrier_aligner.rs rename to src/streaming_runtime/src/streaming/execution/tracker/barrier_aligner.rs diff --git a/src/runtime/streaming/execution/tracker/mod.rs b/src/streaming_runtime/src/streaming/execution/tracker/mod.rs similarity index 100% rename from src/runtime/streaming/execution/tracker/mod.rs rename to src/streaming_runtime/src/streaming/execution/tracker/mod.rs diff --git a/src/runtime/streaming/execution/tracker/watermark_tracker.rs b/src/streaming_runtime/src/streaming/execution/tracker/watermark_tracker.rs similarity index 100% rename from src/runtime/streaming/execution/tracker/watermark_tracker.rs rename to src/streaming_runtime/src/streaming/execution/tracker/watermark_tracker.rs diff --git a/src/runtime/streaming/factory/connector/delta.rs b/src/streaming_runtime/src/streaming/factory/connector/delta.rs similarity index 100% rename from src/runtime/streaming/factory/connector/delta.rs rename to src/streaming_runtime/src/streaming/factory/connector/delta.rs diff --git a/src/runtime/streaming/factory/connector/dispatchers.rs b/src/streaming_runtime/src/streaming/factory/connector/dispatchers.rs similarity index 100% rename from src/runtime/streaming/factory/connector/dispatchers.rs rename to src/streaming_runtime/src/streaming/factory/connector/dispatchers.rs diff --git a/src/runtime/streaming/factory/connector/filesystem.rs b/src/streaming_runtime/src/streaming/factory/connector/filesystem.rs similarity index 100% rename from src/runtime/streaming/factory/connector/filesystem.rs rename to src/streaming_runtime/src/streaming/factory/connector/filesystem.rs diff --git a/src/runtime/streaming/factory/connector/iceberg.rs b/src/streaming_runtime/src/streaming/factory/connector/iceberg.rs similarity index 100% rename from src/runtime/streaming/factory/connector/iceberg.rs rename to src/streaming_runtime/src/streaming/factory/connector/iceberg.rs diff --git a/src/runtime/streaming/factory/connector/kafka.rs b/src/streaming_runtime/src/streaming/factory/connector/kafka.rs similarity index 100% rename from src/runtime/streaming/factory/connector/kafka.rs rename to src/streaming_runtime/src/streaming/factory/connector/kafka.rs diff --git a/src/runtime/streaming/factory/connector/lancedb.rs b/src/streaming_runtime/src/streaming/factory/connector/lancedb.rs similarity index 100% rename from src/runtime/streaming/factory/connector/lancedb.rs rename to src/streaming_runtime/src/streaming/factory/connector/lancedb.rs diff --git a/src/runtime/streaming/factory/connector/mod.rs b/src/streaming_runtime/src/streaming/factory/connector/mod.rs similarity index 100% rename from src/runtime/streaming/factory/connector/mod.rs rename to src/streaming_runtime/src/streaming/factory/connector/mod.rs diff --git a/src/runtime/streaming/factory/connector/s3.rs b/src/streaming_runtime/src/streaming/factory/connector/s3.rs similarity index 100% rename from src/runtime/streaming/factory/connector/s3.rs rename to src/streaming_runtime/src/streaming/factory/connector/s3.rs diff --git a/src/runtime/streaming/factory/connector/sink_props_codec.rs b/src/streaming_runtime/src/streaming/factory/connector/sink_props_codec.rs similarity index 100% rename from src/runtime/streaming/factory/connector/sink_props_codec.rs rename to src/streaming_runtime/src/streaming/factory/connector/sink_props_codec.rs diff --git a/src/runtime/streaming/factory/global/mod.rs b/src/streaming_runtime/src/streaming/factory/global/mod.rs similarity index 100% rename from src/runtime/streaming/factory/global/mod.rs rename to src/streaming_runtime/src/streaming/factory/global/mod.rs diff --git a/src/runtime/streaming/factory/global/session_registry.rs b/src/streaming_runtime/src/streaming/factory/global/session_registry.rs similarity index 100% rename from src/runtime/streaming/factory/global/session_registry.rs rename to src/streaming_runtime/src/streaming/factory/global/session_registry.rs diff --git a/src/runtime/streaming/factory/mod.rs b/src/streaming_runtime/src/streaming/factory/mod.rs similarity index 100% rename from src/runtime/streaming/factory/mod.rs rename to src/streaming_runtime/src/streaming/factory/mod.rs diff --git a/src/runtime/streaming/factory/operator_constructor.rs b/src/streaming_runtime/src/streaming/factory/operator_constructor.rs similarity index 100% rename from src/runtime/streaming/factory/operator_constructor.rs rename to src/streaming_runtime/src/streaming/factory/operator_constructor.rs diff --git a/src/runtime/streaming/factory/operator_factory.rs b/src/streaming_runtime/src/streaming/factory/operator_factory.rs similarity index 100% rename from src/runtime/streaming/factory/operator_factory.rs rename to src/streaming_runtime/src/streaming/factory/operator_factory.rs diff --git a/src/runtime/streaming/format/config.rs b/src/streaming_runtime/src/streaming/format/config.rs similarity index 100% rename from src/runtime/streaming/format/config.rs rename to src/streaming_runtime/src/streaming/format/config.rs diff --git a/src/runtime/streaming/format/deserializer.rs b/src/streaming_runtime/src/streaming/format/deserializer.rs similarity index 100% rename from src/runtime/streaming/format/deserializer.rs rename to src/streaming_runtime/src/streaming/format/deserializer.rs diff --git a/src/runtime/streaming/format/encoder.rs b/src/streaming_runtime/src/streaming/format/encoder.rs similarity index 100% rename from src/runtime/streaming/format/encoder.rs rename to src/streaming_runtime/src/streaming/format/encoder.rs diff --git a/src/runtime/streaming/format/json_encoder.rs b/src/streaming_runtime/src/streaming/format/json_encoder.rs similarity index 100% rename from src/runtime/streaming/format/json_encoder.rs rename to src/streaming_runtime/src/streaming/format/json_encoder.rs diff --git a/src/runtime/streaming/format/mod.rs b/src/streaming_runtime/src/streaming/format/mod.rs similarity index 100% rename from src/runtime/streaming/format/mod.rs rename to src/streaming_runtime/src/streaming/format/mod.rs diff --git a/src/runtime/streaming/format/serializer.rs b/src/streaming_runtime/src/streaming/format/serializer.rs similarity index 100% rename from src/runtime/streaming/format/serializer.rs rename to src/streaming_runtime/src/streaming/format/serializer.rs diff --git a/src/runtime/streaming/job/edge_manager.rs b/src/streaming_runtime/src/streaming/job/edge_manager.rs similarity index 100% rename from src/runtime/streaming/job/edge_manager.rs rename to src/streaming_runtime/src/streaming/job/edge_manager.rs diff --git a/src/runtime/streaming/job/job_manager.rs b/src/streaming_runtime/src/streaming/job/job_manager.rs similarity index 100% rename from src/runtime/streaming/job/job_manager.rs rename to src/streaming_runtime/src/streaming/job/job_manager.rs diff --git a/src/runtime/streaming/job/mod.rs b/src/streaming_runtime/src/streaming/job/mod.rs similarity index 100% rename from src/runtime/streaming/job/mod.rs rename to src/streaming_runtime/src/streaming/job/mod.rs diff --git a/src/runtime/streaming/job/models.rs b/src/streaming_runtime/src/streaming/job/models.rs similarity index 100% rename from src/runtime/streaming/job/models.rs rename to src/streaming_runtime/src/streaming/job/models.rs diff --git a/src/runtime/streaming/mod.rs b/src/streaming_runtime/src/streaming/mod.rs similarity index 100% rename from src/runtime/streaming/mod.rs rename to src/streaming_runtime/src/streaming/mod.rs diff --git a/src/runtime/streaming/network/endpoint.rs b/src/streaming_runtime/src/streaming/network/endpoint.rs similarity index 100% rename from src/runtime/streaming/network/endpoint.rs rename to src/streaming_runtime/src/streaming/network/endpoint.rs diff --git a/src/runtime/streaming/network/environment.rs b/src/streaming_runtime/src/streaming/network/environment.rs similarity index 100% rename from src/runtime/streaming/network/environment.rs rename to src/streaming_runtime/src/streaming/network/environment.rs diff --git a/src/runtime/streaming/network/mod.rs b/src/streaming_runtime/src/streaming/network/mod.rs similarity index 100% rename from src/runtime/streaming/network/mod.rs rename to src/streaming_runtime/src/streaming/network/mod.rs diff --git a/src/runtime/streaming/operators/grouping/incremental_aggregate.rs b/src/streaming_runtime/src/streaming/operators/grouping/incremental_aggregate.rs similarity index 100% rename from src/runtime/streaming/operators/grouping/incremental_aggregate.rs rename to src/streaming_runtime/src/streaming/operators/grouping/incremental_aggregate.rs diff --git a/src/runtime/streaming/operators/grouping/mod.rs b/src/streaming_runtime/src/streaming/operators/grouping/mod.rs similarity index 100% rename from src/runtime/streaming/operators/grouping/mod.rs rename to src/streaming_runtime/src/streaming/operators/grouping/mod.rs diff --git a/src/runtime/streaming/operators/grouping/updating_cache.rs b/src/streaming_runtime/src/streaming/operators/grouping/updating_cache.rs similarity index 100% rename from src/runtime/streaming/operators/grouping/updating_cache.rs rename to src/streaming_runtime/src/streaming/operators/grouping/updating_cache.rs diff --git a/src/runtime/streaming/operators/joins/join_instance.rs b/src/streaming_runtime/src/streaming/operators/joins/join_instance.rs similarity index 100% rename from src/runtime/streaming/operators/joins/join_instance.rs rename to src/streaming_runtime/src/streaming/operators/joins/join_instance.rs diff --git a/src/runtime/streaming/operators/joins/join_with_expiration.rs b/src/streaming_runtime/src/streaming/operators/joins/join_with_expiration.rs similarity index 100% rename from src/runtime/streaming/operators/joins/join_with_expiration.rs rename to src/streaming_runtime/src/streaming/operators/joins/join_with_expiration.rs diff --git a/src/runtime/streaming/operators/joins/mod.rs b/src/streaming_runtime/src/streaming/operators/joins/mod.rs similarity index 100% rename from src/runtime/streaming/operators/joins/mod.rs rename to src/streaming_runtime/src/streaming/operators/joins/mod.rs diff --git a/src/runtime/streaming/operators/key_by.rs b/src/streaming_runtime/src/streaming/operators/key_by.rs similarity index 100% rename from src/runtime/streaming/operators/key_by.rs rename to src/streaming_runtime/src/streaming/operators/key_by.rs diff --git a/src/runtime/streaming/operators/key_operator.rs b/src/streaming_runtime/src/streaming/operators/key_operator.rs similarity index 100% rename from src/runtime/streaming/operators/key_operator.rs rename to src/streaming_runtime/src/streaming/operators/key_operator.rs diff --git a/src/runtime/streaming/operators/mod.rs b/src/streaming_runtime/src/streaming/operators/mod.rs similarity index 100% rename from src/runtime/streaming/operators/mod.rs rename to src/streaming_runtime/src/streaming/operators/mod.rs diff --git a/src/runtime/streaming/operators/projection.rs b/src/streaming_runtime/src/streaming/operators/projection.rs similarity index 100% rename from src/runtime/streaming/operators/projection.rs rename to src/streaming_runtime/src/streaming/operators/projection.rs diff --git a/src/runtime/streaming/operators/sink/delta/mod.rs b/src/streaming_runtime/src/streaming/operators/sink/delta/mod.rs similarity index 100% rename from src/runtime/streaming/operators/sink/delta/mod.rs rename to src/streaming_runtime/src/streaming/operators/sink/delta/mod.rs diff --git a/src/runtime/streaming/operators/sink/filesystem/mod.rs b/src/streaming_runtime/src/streaming/operators/sink/filesystem/mod.rs similarity index 100% rename from src/runtime/streaming/operators/sink/filesystem/mod.rs rename to src/streaming_runtime/src/streaming/operators/sink/filesystem/mod.rs diff --git a/src/runtime/streaming/operators/sink/iceberg/mod.rs b/src/streaming_runtime/src/streaming/operators/sink/iceberg/mod.rs similarity index 100% rename from src/runtime/streaming/operators/sink/iceberg/mod.rs rename to src/streaming_runtime/src/streaming/operators/sink/iceberg/mod.rs diff --git a/src/runtime/streaming/operators/sink/kafka/mod.rs b/src/streaming_runtime/src/streaming/operators/sink/kafka/mod.rs similarity index 100% rename from src/runtime/streaming/operators/sink/kafka/mod.rs rename to src/streaming_runtime/src/streaming/operators/sink/kafka/mod.rs diff --git a/src/runtime/streaming/operators/sink/lancedb/mod.rs b/src/streaming_runtime/src/streaming/operators/sink/lancedb/mod.rs similarity index 100% rename from src/runtime/streaming/operators/sink/lancedb/mod.rs rename to src/streaming_runtime/src/streaming/operators/sink/lancedb/mod.rs diff --git a/src/runtime/streaming/operators/sink/mod.rs b/src/streaming_runtime/src/streaming/operators/sink/mod.rs similarity index 100% rename from src/runtime/streaming/operators/sink/mod.rs rename to src/streaming_runtime/src/streaming/operators/sink/mod.rs diff --git a/src/runtime/streaming/operators/sink/s3/mod.rs b/src/streaming_runtime/src/streaming/operators/sink/s3/mod.rs similarity index 100% rename from src/runtime/streaming/operators/sink/s3/mod.rs rename to src/streaming_runtime/src/streaming/operators/sink/s3/mod.rs diff --git a/src/runtime/streaming/operators/source/kafka/mod.rs b/src/streaming_runtime/src/streaming/operators/source/kafka/mod.rs similarity index 100% rename from src/runtime/streaming/operators/source/kafka/mod.rs rename to src/streaming_runtime/src/streaming/operators/source/kafka/mod.rs diff --git a/src/runtime/streaming/operators/source/mod.rs b/src/streaming_runtime/src/streaming/operators/source/mod.rs similarity index 100% rename from src/runtime/streaming/operators/source/mod.rs rename to src/streaming_runtime/src/streaming/operators/source/mod.rs diff --git a/src/runtime/streaming/operators/stateless_physical_executor.rs b/src/streaming_runtime/src/streaming/operators/stateless_physical_executor.rs similarity index 100% rename from src/runtime/streaming/operators/stateless_physical_executor.rs rename to src/streaming_runtime/src/streaming/operators/stateless_physical_executor.rs diff --git a/src/runtime/streaming/operators/value_execution.rs b/src/streaming_runtime/src/streaming/operators/value_execution.rs similarity index 100% rename from src/runtime/streaming/operators/value_execution.rs rename to src/streaming_runtime/src/streaming/operators/value_execution.rs diff --git a/src/runtime/streaming/operators/watermark/mod.rs b/src/streaming_runtime/src/streaming/operators/watermark/mod.rs similarity index 100% rename from src/runtime/streaming/operators/watermark/mod.rs rename to src/streaming_runtime/src/streaming/operators/watermark/mod.rs diff --git a/src/runtime/streaming/operators/watermark/watermark_generator.rs b/src/streaming_runtime/src/streaming/operators/watermark/watermark_generator.rs similarity index 100% rename from src/runtime/streaming/operators/watermark/watermark_generator.rs rename to src/streaming_runtime/src/streaming/operators/watermark/watermark_generator.rs diff --git a/src/runtime/streaming/operators/windows/mod.rs b/src/streaming_runtime/src/streaming/operators/windows/mod.rs similarity index 100% rename from src/runtime/streaming/operators/windows/mod.rs rename to src/streaming_runtime/src/streaming/operators/windows/mod.rs diff --git a/src/runtime/streaming/operators/windows/session_aggregating_window.rs b/src/streaming_runtime/src/streaming/operators/windows/session_aggregating_window.rs similarity index 100% rename from src/runtime/streaming/operators/windows/session_aggregating_window.rs rename to src/streaming_runtime/src/streaming/operators/windows/session_aggregating_window.rs diff --git a/src/runtime/streaming/operators/windows/sliding_aggregating_window.rs b/src/streaming_runtime/src/streaming/operators/windows/sliding_aggregating_window.rs similarity index 100% rename from src/runtime/streaming/operators/windows/sliding_aggregating_window.rs rename to src/streaming_runtime/src/streaming/operators/windows/sliding_aggregating_window.rs diff --git a/src/runtime/streaming/operators/windows/tumbling_aggregating_window.rs b/src/streaming_runtime/src/streaming/operators/windows/tumbling_aggregating_window.rs similarity index 100% rename from src/runtime/streaming/operators/windows/tumbling_aggregating_window.rs rename to src/streaming_runtime/src/streaming/operators/windows/tumbling_aggregating_window.rs diff --git a/src/runtime/streaming/operators/windows/window_function.rs b/src/streaming_runtime/src/streaming/operators/windows/window_function.rs similarity index 100% rename from src/runtime/streaming/operators/windows/window_function.rs rename to src/streaming_runtime/src/streaming/operators/windows/window_function.rs diff --git a/src/runtime/streaming/protocol/control.rs b/src/streaming_runtime/src/streaming/protocol/control.rs similarity index 100% rename from src/runtime/streaming/protocol/control.rs rename to src/streaming_runtime/src/streaming/protocol/control.rs diff --git a/src/runtime/streaming/protocol/event.rs b/src/streaming_runtime/src/streaming/protocol/event.rs similarity index 100% rename from src/runtime/streaming/protocol/event.rs rename to src/streaming_runtime/src/streaming/protocol/event.rs diff --git a/src/runtime/streaming/protocol/mod.rs b/src/streaming_runtime/src/streaming/protocol/mod.rs similarity index 100% rename from src/runtime/streaming/protocol/mod.rs rename to src/streaming_runtime/src/streaming/protocol/mod.rs diff --git a/src/runtime/streaming/state/error.rs b/src/streaming_runtime/src/streaming/state/error.rs similarity index 100% rename from src/runtime/streaming/state/error.rs rename to src/streaming_runtime/src/streaming/state/error.rs diff --git a/src/runtime/streaming/state/io_manager.rs b/src/streaming_runtime/src/streaming/state/io_manager.rs similarity index 100% rename from src/runtime/streaming/state/io_manager.rs rename to src/streaming_runtime/src/streaming/state/io_manager.rs diff --git a/src/runtime/streaming/state/metrics.rs b/src/streaming_runtime/src/streaming/state/metrics.rs similarity index 100% rename from src/runtime/streaming/state/metrics.rs rename to src/streaming_runtime/src/streaming/state/metrics.rs diff --git a/src/runtime/streaming/state/mod.rs b/src/streaming_runtime/src/streaming/state/mod.rs similarity index 100% rename from src/runtime/streaming/state/mod.rs rename to src/streaming_runtime/src/streaming/state/mod.rs diff --git a/src/runtime/streaming/state/operator_state.rs b/src/streaming_runtime/src/streaming/state/operator_state.rs similarity index 100% rename from src/runtime/streaming/state/operator_state.rs rename to src/streaming_runtime/src/streaming/state/operator_state.rs diff --git a/src/runtime/util/mod.rs b/src/streaming_runtime/src/util/mod.rs similarity index 100% rename from src/runtime/util/mod.rs rename to src/streaming_runtime/src/util/mod.rs diff --git a/src/runtime/util/physical_aggregate.rs b/src/streaming_runtime/src/util/physical_aggregate.rs similarity index 100% rename from src/runtime/util/physical_aggregate.rs rename to src/streaming_runtime/src/util/physical_aggregate.rs diff --git a/src/wasm_runtime/src/lib.rs b/src/wasm_runtime/src/lib.rs index 3eeefce9..d0b4073e 100644 --- a/src/wasm_runtime/src/lib.rs +++ b/src/wasm_runtime/src/lib.rs @@ -1,3 +1,22 @@ +// 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. + //! WebAssembly execution runtime. +//! +//! Implementation lives under `src/wasm/` in this package. It is currently **compiled as +//! part of the `function-stream` crate** via `#[path]` in `src/runtime/mod.rs`, so paths +//! like `crate::sql` and `crate::runtime::memory` keep resolving until further crate splits. +//! +//! Operator state storage (`state_backend/`) also lives in this package and is compiled via +//! `#[path]` from `src/storage/mod.rs` as `crate::storage::state_backend`. pub const CRATE_NAME: &str = "function-stream-wasm-runtime"; diff --git a/src/storage/state_backend/error.rs b/src/wasm_runtime/src/state_backend/error.rs similarity index 100% rename from src/storage/state_backend/error.rs rename to src/wasm_runtime/src/state_backend/error.rs diff --git a/src/storage/state_backend/factory.rs b/src/wasm_runtime/src/state_backend/factory.rs similarity index 100% rename from src/storage/state_backend/factory.rs rename to src/wasm_runtime/src/state_backend/factory.rs diff --git a/src/storage/state_backend/key_builder.rs b/src/wasm_runtime/src/state_backend/key_builder.rs similarity index 100% rename from src/storage/state_backend/key_builder.rs rename to src/wasm_runtime/src/state_backend/key_builder.rs diff --git a/src/storage/state_backend/memory/factory.rs b/src/wasm_runtime/src/state_backend/memory/factory.rs similarity index 100% rename from src/storage/state_backend/memory/factory.rs rename to src/wasm_runtime/src/state_backend/memory/factory.rs diff --git a/src/storage/state_backend/memory/mod.rs b/src/wasm_runtime/src/state_backend/memory/mod.rs similarity index 100% rename from src/storage/state_backend/memory/mod.rs rename to src/wasm_runtime/src/state_backend/memory/mod.rs diff --git a/src/storage/state_backend/memory/store.rs b/src/wasm_runtime/src/state_backend/memory/store.rs similarity index 100% rename from src/storage/state_backend/memory/store.rs rename to src/wasm_runtime/src/state_backend/memory/store.rs diff --git a/src/storage/state_backend/mod.rs b/src/wasm_runtime/src/state_backend/mod.rs similarity index 100% rename from src/storage/state_backend/mod.rs rename to src/wasm_runtime/src/state_backend/mod.rs diff --git a/src/storage/state_backend/rocksdb/factory.rs b/src/wasm_runtime/src/state_backend/rocksdb/factory.rs similarity index 100% rename from src/storage/state_backend/rocksdb/factory.rs rename to src/wasm_runtime/src/state_backend/rocksdb/factory.rs diff --git a/src/storage/state_backend/rocksdb/mod.rs b/src/wasm_runtime/src/state_backend/rocksdb/mod.rs similarity index 100% rename from src/storage/state_backend/rocksdb/mod.rs rename to src/wasm_runtime/src/state_backend/rocksdb/mod.rs diff --git a/src/storage/state_backend/rocksdb/store.rs b/src/wasm_runtime/src/state_backend/rocksdb/store.rs similarity index 100% rename from src/storage/state_backend/rocksdb/store.rs rename to src/wasm_runtime/src/state_backend/rocksdb/store.rs diff --git a/src/storage/state_backend/server.rs b/src/wasm_runtime/src/state_backend/server.rs similarity index 100% rename from src/storage/state_backend/server.rs rename to src/wasm_runtime/src/state_backend/server.rs diff --git a/src/storage/state_backend/store.rs b/src/wasm_runtime/src/state_backend/store.rs similarity index 100% rename from src/storage/state_backend/store.rs rename to src/wasm_runtime/src/state_backend/store.rs diff --git a/src/runtime/wasm/buffer_and_event/buffer_or_event.rs b/src/wasm_runtime/src/wasm/buffer_and_event/buffer_or_event.rs similarity index 100% rename from src/runtime/wasm/buffer_and_event/buffer_or_event.rs rename to src/wasm_runtime/src/wasm/buffer_and_event/buffer_or_event.rs diff --git a/src/runtime/wasm/buffer_and_event/mod.rs b/src/wasm_runtime/src/wasm/buffer_and_event/mod.rs similarity index 100% rename from src/runtime/wasm/buffer_and_event/mod.rs rename to src/wasm_runtime/src/wasm/buffer_and_event/mod.rs diff --git a/src/runtime/wasm/buffer_and_event/stream_element/mod.rs b/src/wasm_runtime/src/wasm/buffer_and_event/stream_element/mod.rs similarity index 100% rename from src/runtime/wasm/buffer_and_event/stream_element/mod.rs rename to src/wasm_runtime/src/wasm/buffer_and_event/stream_element/mod.rs diff --git a/src/runtime/wasm/buffer_and_event/stream_element/stream_element.rs b/src/wasm_runtime/src/wasm/buffer_and_event/stream_element/stream_element.rs similarity index 100% rename from src/runtime/wasm/buffer_and_event/stream_element/stream_element.rs rename to src/wasm_runtime/src/wasm/buffer_and_event/stream_element/stream_element.rs diff --git a/src/runtime/wasm/input/input_protocol.rs b/src/wasm_runtime/src/wasm/input/input_protocol.rs similarity index 100% rename from src/runtime/wasm/input/input_protocol.rs rename to src/wasm_runtime/src/wasm/input/input_protocol.rs diff --git a/src/runtime/wasm/input/input_provider.rs b/src/wasm_runtime/src/wasm/input/input_provider.rs similarity index 100% rename from src/runtime/wasm/input/input_provider.rs rename to src/wasm_runtime/src/wasm/input/input_provider.rs diff --git a/src/runtime/wasm/input/input_runner.rs b/src/wasm_runtime/src/wasm/input/input_runner.rs similarity index 100% rename from src/runtime/wasm/input/input_runner.rs rename to src/wasm_runtime/src/wasm/input/input_runner.rs diff --git a/src/runtime/wasm/input/interface.rs b/src/wasm_runtime/src/wasm/input/interface.rs similarity index 100% rename from src/runtime/wasm/input/interface.rs rename to src/wasm_runtime/src/wasm/input/interface.rs diff --git a/src/runtime/wasm/input/mod.rs b/src/wasm_runtime/src/wasm/input/mod.rs similarity index 100% rename from src/runtime/wasm/input/mod.rs rename to src/wasm_runtime/src/wasm/input/mod.rs diff --git a/src/runtime/wasm/input/protocol/kafka/config.rs b/src/wasm_runtime/src/wasm/input/protocol/kafka/config.rs similarity index 100% rename from src/runtime/wasm/input/protocol/kafka/config.rs rename to src/wasm_runtime/src/wasm/input/protocol/kafka/config.rs diff --git a/src/runtime/wasm/input/protocol/kafka/kafka_protocol.rs b/src/wasm_runtime/src/wasm/input/protocol/kafka/kafka_protocol.rs similarity index 100% rename from src/runtime/wasm/input/protocol/kafka/kafka_protocol.rs rename to src/wasm_runtime/src/wasm/input/protocol/kafka/kafka_protocol.rs diff --git a/src/runtime/wasm/input/protocol/kafka/mod.rs b/src/wasm_runtime/src/wasm/input/protocol/kafka/mod.rs similarity index 100% rename from src/runtime/wasm/input/protocol/kafka/mod.rs rename to src/wasm_runtime/src/wasm/input/protocol/kafka/mod.rs diff --git a/src/runtime/wasm/input/protocol/mod.rs b/src/wasm_runtime/src/wasm/input/protocol/mod.rs similarity index 100% rename from src/runtime/wasm/input/protocol/mod.rs rename to src/wasm_runtime/src/wasm/input/protocol/mod.rs diff --git a/src/runtime/wasm/mod.rs b/src/wasm_runtime/src/wasm/mod.rs similarity index 100% rename from src/runtime/wasm/mod.rs rename to src/wasm_runtime/src/wasm/mod.rs diff --git a/src/runtime/wasm/output/interface.rs b/src/wasm_runtime/src/wasm/output/interface.rs similarity index 100% rename from src/runtime/wasm/output/interface.rs rename to src/wasm_runtime/src/wasm/output/interface.rs diff --git a/src/runtime/wasm/output/mod.rs b/src/wasm_runtime/src/wasm/output/mod.rs similarity index 100% rename from src/runtime/wasm/output/mod.rs rename to src/wasm_runtime/src/wasm/output/mod.rs diff --git a/src/runtime/wasm/output/output_protocol.rs b/src/wasm_runtime/src/wasm/output/output_protocol.rs similarity index 100% rename from src/runtime/wasm/output/output_protocol.rs rename to src/wasm_runtime/src/wasm/output/output_protocol.rs diff --git a/src/runtime/wasm/output/output_provider.rs b/src/wasm_runtime/src/wasm/output/output_provider.rs similarity index 100% rename from src/runtime/wasm/output/output_provider.rs rename to src/wasm_runtime/src/wasm/output/output_provider.rs diff --git a/src/runtime/wasm/output/output_runner.rs b/src/wasm_runtime/src/wasm/output/output_runner.rs similarity index 100% rename from src/runtime/wasm/output/output_runner.rs rename to src/wasm_runtime/src/wasm/output/output_runner.rs diff --git a/src/runtime/wasm/output/protocol/kafka/kafka_protocol.rs b/src/wasm_runtime/src/wasm/output/protocol/kafka/kafka_protocol.rs similarity index 100% rename from src/runtime/wasm/output/protocol/kafka/kafka_protocol.rs rename to src/wasm_runtime/src/wasm/output/protocol/kafka/kafka_protocol.rs diff --git a/src/runtime/wasm/output/protocol/kafka/mod.rs b/src/wasm_runtime/src/wasm/output/protocol/kafka/mod.rs similarity index 100% rename from src/runtime/wasm/output/protocol/kafka/mod.rs rename to src/wasm_runtime/src/wasm/output/protocol/kafka/mod.rs diff --git a/src/runtime/wasm/output/protocol/kafka/producer_config.rs b/src/wasm_runtime/src/wasm/output/protocol/kafka/producer_config.rs similarity index 100% rename from src/runtime/wasm/output/protocol/kafka/producer_config.rs rename to src/wasm_runtime/src/wasm/output/protocol/kafka/producer_config.rs diff --git a/src/runtime/wasm/output/protocol/mod.rs b/src/wasm_runtime/src/wasm/output/protocol/mod.rs similarity index 100% rename from src/runtime/wasm/output/protocol/mod.rs rename to src/wasm_runtime/src/wasm/output/protocol/mod.rs diff --git a/src/runtime/wasm/processor/function_error.rs b/src/wasm_runtime/src/wasm/processor/function_error.rs similarity index 100% rename from src/runtime/wasm/processor/function_error.rs rename to src/wasm_runtime/src/wasm/processor/function_error.rs diff --git a/src/runtime/wasm/processor/mod.rs b/src/wasm_runtime/src/wasm/processor/mod.rs similarity index 100% rename from src/runtime/wasm/processor/mod.rs rename to src/wasm_runtime/src/wasm/processor/mod.rs diff --git a/src/runtime/wasm/processor/python/mod.rs b/src/wasm_runtime/src/wasm/processor/python/mod.rs similarity index 100% rename from src/runtime/wasm/processor/python/mod.rs rename to src/wasm_runtime/src/wasm/processor/python/mod.rs diff --git a/src/runtime/wasm/processor/python/python_host.rs b/src/wasm_runtime/src/wasm/processor/python/python_host.rs similarity index 100% rename from src/runtime/wasm/processor/python/python_host.rs rename to src/wasm_runtime/src/wasm/processor/python/python_host.rs diff --git a/src/runtime/wasm/processor/python/python_service.rs b/src/wasm_runtime/src/wasm/processor/python/python_service.rs similarity index 100% rename from src/runtime/wasm/processor/python/python_service.rs rename to src/wasm_runtime/src/wasm/processor/python/python_service.rs diff --git a/src/runtime/wasm/processor/wasm/input_strategy.rs b/src/wasm_runtime/src/wasm/processor/wasm/input_strategy.rs similarity index 100% rename from src/runtime/wasm/processor/wasm/input_strategy.rs rename to src/wasm_runtime/src/wasm/processor/wasm/input_strategy.rs diff --git a/src/runtime/wasm/processor/wasm/mod.rs b/src/wasm_runtime/src/wasm/processor/wasm/mod.rs similarity index 100% rename from src/runtime/wasm/processor/wasm/mod.rs rename to src/wasm_runtime/src/wasm/processor/wasm/mod.rs diff --git a/src/runtime/wasm/processor/wasm/thread_pool.rs b/src/wasm_runtime/src/wasm/processor/wasm/thread_pool.rs similarity index 100% rename from src/runtime/wasm/processor/wasm/thread_pool.rs rename to src/wasm_runtime/src/wasm/processor/wasm/thread_pool.rs diff --git a/src/runtime/wasm/processor/wasm/wasm_cache.rs b/src/wasm_runtime/src/wasm/processor/wasm/wasm_cache.rs similarity index 100% rename from src/runtime/wasm/processor/wasm/wasm_cache.rs rename to src/wasm_runtime/src/wasm/processor/wasm/wasm_cache.rs diff --git a/src/runtime/wasm/processor/wasm/wasm_host.rs b/src/wasm_runtime/src/wasm/processor/wasm/wasm_host.rs similarity index 100% rename from src/runtime/wasm/processor/wasm/wasm_host.rs rename to src/wasm_runtime/src/wasm/processor/wasm/wasm_host.rs diff --git a/src/runtime/wasm/processor/wasm/wasm_processor.rs b/src/wasm_runtime/src/wasm/processor/wasm/wasm_processor.rs similarity index 100% rename from src/runtime/wasm/processor/wasm/wasm_processor.rs rename to src/wasm_runtime/src/wasm/processor/wasm/wasm_processor.rs diff --git a/src/runtime/wasm/processor/wasm/wasm_processor_trait.rs b/src/wasm_runtime/src/wasm/processor/wasm/wasm_processor_trait.rs similarity index 100% rename from src/runtime/wasm/processor/wasm/wasm_processor_trait.rs rename to src/wasm_runtime/src/wasm/processor/wasm/wasm_processor_trait.rs diff --git a/src/runtime/wasm/processor/wasm/wasm_task.rs b/src/wasm_runtime/src/wasm/processor/wasm/wasm_task.rs similarity index 100% rename from src/runtime/wasm/processor/wasm/wasm_task.rs rename to src/wasm_runtime/src/wasm/processor/wasm/wasm_task.rs diff --git a/src/runtime/wasm/task/builder/mod.rs b/src/wasm_runtime/src/wasm/task/builder/mod.rs similarity index 100% rename from src/runtime/wasm/task/builder/mod.rs rename to src/wasm_runtime/src/wasm/task/builder/mod.rs diff --git a/src/runtime/wasm/task/builder/processor/mod.rs b/src/wasm_runtime/src/wasm/task/builder/processor/mod.rs similarity index 100% rename from src/runtime/wasm/task/builder/processor/mod.rs rename to src/wasm_runtime/src/wasm/task/builder/processor/mod.rs diff --git a/src/runtime/wasm/task/builder/python/mod.rs b/src/wasm_runtime/src/wasm/task/builder/python/mod.rs similarity index 100% rename from src/runtime/wasm/task/builder/python/mod.rs rename to src/wasm_runtime/src/wasm/task/builder/python/mod.rs diff --git a/src/runtime/wasm/task/builder/sink/mod.rs b/src/wasm_runtime/src/wasm/task/builder/sink/mod.rs similarity index 100% rename from src/runtime/wasm/task/builder/sink/mod.rs rename to src/wasm_runtime/src/wasm/task/builder/sink/mod.rs diff --git a/src/runtime/wasm/task/builder/source/mod.rs b/src/wasm_runtime/src/wasm/task/builder/source/mod.rs similarity index 100% rename from src/runtime/wasm/task/builder/source/mod.rs rename to src/wasm_runtime/src/wasm/task/builder/source/mod.rs diff --git a/src/runtime/wasm/task/builder/task_builder.rs b/src/wasm_runtime/src/wasm/task/builder/task_builder.rs similarity index 100% rename from src/runtime/wasm/task/builder/task_builder.rs rename to src/wasm_runtime/src/wasm/task/builder/task_builder.rs diff --git a/src/runtime/wasm/task/control_mailbox.rs b/src/wasm_runtime/src/wasm/task/control_mailbox.rs similarity index 100% rename from src/runtime/wasm/task/control_mailbox.rs rename to src/wasm_runtime/src/wasm/task/control_mailbox.rs diff --git a/src/runtime/wasm/task/lifecycle.rs b/src/wasm_runtime/src/wasm/task/lifecycle.rs similarity index 100% rename from src/runtime/wasm/task/lifecycle.rs rename to src/wasm_runtime/src/wasm/task/lifecycle.rs diff --git a/src/runtime/wasm/task/mod.rs b/src/wasm_runtime/src/wasm/task/mod.rs similarity index 100% rename from src/runtime/wasm/task/mod.rs rename to src/wasm_runtime/src/wasm/task/mod.rs diff --git a/src/runtime/wasm/task/processor_config.rs b/src/wasm_runtime/src/wasm/task/processor_config.rs similarity index 100% rename from src/runtime/wasm/task/processor_config.rs rename to src/wasm_runtime/src/wasm/task/processor_config.rs diff --git a/src/runtime/wasm/task/yaml_keys.rs b/src/wasm_runtime/src/wasm/task/yaml_keys.rs similarity index 100% rename from src/runtime/wasm/task/yaml_keys.rs rename to src/wasm_runtime/src/wasm/task/yaml_keys.rs diff --git a/src/runtime/wasm/taskexecutor/init_context.rs b/src/wasm_runtime/src/wasm/taskexecutor/init_context.rs similarity index 100% rename from src/runtime/wasm/taskexecutor/init_context.rs rename to src/wasm_runtime/src/wasm/taskexecutor/init_context.rs diff --git a/src/runtime/wasm/taskexecutor/mod.rs b/src/wasm_runtime/src/wasm/taskexecutor/mod.rs similarity index 100% rename from src/runtime/wasm/taskexecutor/mod.rs rename to src/wasm_runtime/src/wasm/taskexecutor/mod.rs diff --git a/src/runtime/wasm/taskexecutor/task_manager.rs b/src/wasm_runtime/src/wasm/taskexecutor/task_manager.rs similarity index 100% rename from src/runtime/wasm/taskexecutor/task_manager.rs rename to src/wasm_runtime/src/wasm/taskexecutor/task_manager.rs From 88d57b95a171b9599c032d359cd544787630fe2a Mon Sep 17 00:00:00 2001 From: luoluoyuyu Date: Tue, 12 May 2026 09:04:12 +0800 Subject: [PATCH 3/7] update --- Cargo.lock | 44 ++ Cargo.toml | 2 + src/catalog_storage/src/lib.rs | 2 +- src/coordinator/src/legacy/mod.rs | 2 +- src/coordinator/src/legacy/sql_classify.rs | 7 +- src/lib.rs | 2 +- src/main.rs | 2 +- src/runtime/mod.rs | 31 - src/runtime_common/Cargo.toml | 1 + src/runtime_common/src/lib.rs | 1 + src/sql/analysis/aggregate_rewriter.rs | 279 -------- src/sql/analysis/async_udf_rewriter.rs | 133 ---- src/sql/analysis/join_rewriter.rs | 234 ------- src/sql/analysis/mod.rs | 214 ------ src/sql/analysis/row_time_rewriter.rs | 49 -- src/sql/analysis/sink_input_rewriter.rs | 57 -- src/sql/analysis/source_metadata_visitor.rs | 73 -- src/sql/analysis/source_rewriter.rs | 305 --------- src/sql/analysis/stream_rewriter.rs | 234 ------- src/sql/analysis/streaming_window_analzer.rs | 219 ------ src/sql/analysis/time_window.rs | 83 --- src/sql/analysis/udafs.rs | 43 -- src/sql/analysis/unnest_rewriter.rs | 179 ----- src/sql/analysis/window_function_rewriter.rs | 204 ------ src/sql/api/checkpoints.rs | 108 --- src/sql/api/connections.rs | 620 ----------------- src/sql/api/metrics.rs | 53 -- src/sql/api/mod.rs | 46 -- src/sql/api/pipelines.rs | 168 ----- src/sql/api/public_ids.rs | 69 -- src/sql/api/schema_resolver.rs | 94 --- src/sql/api/udfs.rs | 68 -- src/sql/api/var_str.rs | 91 --- src/sql/common/arrow_ext.rs | 182 ----- src/sql/common/connector_options.rs | 449 ------------ src/sql/common/constants.rs | 294 -------- src/sql/common/control.rs | 164 ----- src/sql/common/converter.rs | 95 --- src/sql/common/date.rs | 86 --- src/sql/common/debezium.rs | 148 ---- src/sql/common/errors.rs | 92 --- src/sql/common/format_from_opts.rs | 182 ----- src/sql/common/formats.rs | 267 -------- src/sql/common/fs_schema.rs | 470 ------------- src/sql/common/kafka_catalog.rs | 116 ---- src/sql/common/mod.rs | 65 -- src/sql/common/operator_config.rs | 21 - src/sql/common/time_utils.rs | 74 -- src/sql/common/topology.rs | 295 -------- src/sql/common/with_option_keys.rs | 105 --- src/sql/connector/config.rs | 91 --- src/sql/connector/factory.rs | 67 -- src/sql/connector/mod.rs | 18 - src/sql/connector/provider.rs | 52 -- src/sql/connector/registry.rs | 86 --- src/sql/connector/sink/delta.rs | 60 -- src/sql/connector/sink/filesystem.rs | 60 -- src/sql/connector/sink/iceberg.rs | 57 -- src/sql/connector/sink/kafka.rs | 159 ----- src/sql/connector/sink/lancedb.rs | 61 -- src/sql/connector/sink/mod.rs | 20 - src/sql/connector/sink/runtime_config.rs | 137 ---- src/sql/connector/sink/s3.rs | 75 -- src/sql/connector/sink/utils.rs | 91 --- src/sql/connector/source/kafka.rs | 185 ----- src/sql/connector/source/mod.rs | 13 - src/sql/functions/mod.rs | 612 ----------------- src/sql/logical_node/aggregate.rs | 644 ------------------ src/sql/logical_node/async_udf.rs | 247 ------- src/sql/logical_node/debezium.rs | 393 ----------- src/sql/logical_node/extension_try_from.rs | 70 -- src/sql/logical_node/is_retract.rs | 82 --- src/sql/logical_node/join.rs | 211 ------ src/sql/logical_node/key_calculation.rs | 309 --------- .../logical_node/logical/dylib_udf_config.rs | 71 -- .../logical/fs_program_convert.rs | 200 ------ src/sql/logical_node/logical/logical_edge.rs | 102 --- src/sql/logical_node/logical/logical_graph.rs | 30 - src/sql/logical_node/logical/logical_node.rs | 87 --- .../logical_node/logical/logical_program.rs | 153 ----- src/sql/logical_node/logical/mod.rs | 30 - .../logical_node/logical/operator_chain.rs | 142 ---- src/sql/logical_node/logical/operator_name.rs | 82 --- .../logical_node/logical/program_config.rs | 33 - .../logical_node/logical/python_udf_config.rs | 23 - src/sql/logical_node/lookup.rs | 256 ------- src/sql/logical_node/macros.rs | 28 - src/sql/logical_node/mod.rs | 42 -- src/sql/logical_node/projection.rs | 239 ------- src/sql/logical_node/remote_table.rs | 190 ------ src/sql/logical_node/sink.rs | 247 ------- .../streaming_operator_blueprint.rs | 65 -- src/sql/logical_node/table_source.rs | 180 ----- src/sql/logical_node/timestamp_append.rs | 121 ---- src/sql/logical_node/updating_aggregate.rs | 245 ------- src/sql/logical_node/watermark_node.rs | 229 ------- src/sql/logical_node/windows_function.rs | 191 ------ src/sql/logical_planner/mod.rs | 16 - .../logical_planner/optimizers/chaining.rs | 200 ------ src/sql/logical_planner/optimizers/mod.rs | 20 - .../optimizers/optimized_plan.rs | 95 --- src/sql/logical_planner/streaming_planner.rs | 435 ------------ src/sql/mod.rs | 27 - src/sql/parse.rs | 83 --- src/sql/physical/cdc/encode.rs | 342 ---------- src/sql/physical/cdc/mod.rs | 17 - src/sql/physical/cdc/unroll.rs | 322 --------- src/sql/physical/codec.rs | 307 --------- src/sql/physical/meta.rs | 47 -- src/sql/physical/mod.rs | 23 - src/sql/physical/source_exec.rs | 400 ----------- src/sql/physical/udfs.rs | 138 ---- src/sql/planning_runtime.rs | 35 - src/sql/schema/catalog.rs | 609 ----------------- src/sql/schema/column_descriptor.rs | 144 ---- src/sql/schema/connection_type.rs | 31 - src/sql/schema/data_encoding_format.rs | 89 --- src/sql/schema/introspection/ddl_formatter.rs | 156 ----- src/sql/schema/introspection/mod.rs | 21 - .../schema/introspection/show_formatter.rs | 100 --- .../schema/introspection/stream_formatter.rs | 120 ---- src/sql/schema/mod.rs | 30 - src/sql/schema/schema_provider.rs | 469 ------------- src/sql/schema/table.rs | 163 ----- src/sql/schema/table_role.rs | 104 --- src/sql/schema/temporal_pipeline_config.rs | 58 -- src/sql/schema/utils.rs | 79 --- src/sql/types/data_type.rs | 158 ----- src/sql/types/df_field.rs | 181 ----- src/sql/types/mod.rs | 65 -- src/sql/types/placeholder_udf.rs | 79 --- src/sql/types/stream_schema.rs | 133 ---- src/sql/types/window.rs | 134 ---- src/storage/mod.rs | 61 -- src/streaming_runtime/src/lib.rs | 3 +- .../src/streaming/protocol/event.rs | 93 +-- .../src/streaming/protocol/mod.rs | 1 + src/wasm_runtime/src/lib.rs | 6 +- 138 files changed, 63 insertions(+), 19139 deletions(-) delete mode 100644 src/runtime/mod.rs delete mode 100644 src/sql/analysis/aggregate_rewriter.rs delete mode 100644 src/sql/analysis/async_udf_rewriter.rs delete mode 100644 src/sql/analysis/join_rewriter.rs delete mode 100644 src/sql/analysis/mod.rs delete mode 100644 src/sql/analysis/row_time_rewriter.rs delete mode 100644 src/sql/analysis/sink_input_rewriter.rs delete mode 100644 src/sql/analysis/source_metadata_visitor.rs delete mode 100644 src/sql/analysis/source_rewriter.rs delete mode 100644 src/sql/analysis/stream_rewriter.rs delete mode 100644 src/sql/analysis/streaming_window_analzer.rs delete mode 100644 src/sql/analysis/time_window.rs delete mode 100644 src/sql/analysis/udafs.rs delete mode 100644 src/sql/analysis/unnest_rewriter.rs delete mode 100644 src/sql/analysis/window_function_rewriter.rs delete mode 100644 src/sql/api/checkpoints.rs delete mode 100644 src/sql/api/connections.rs delete mode 100644 src/sql/api/metrics.rs delete mode 100644 src/sql/api/mod.rs delete mode 100644 src/sql/api/pipelines.rs delete mode 100644 src/sql/api/public_ids.rs delete mode 100644 src/sql/api/schema_resolver.rs delete mode 100644 src/sql/api/udfs.rs delete mode 100644 src/sql/api/var_str.rs delete mode 100644 src/sql/common/arrow_ext.rs delete mode 100644 src/sql/common/connector_options.rs delete mode 100644 src/sql/common/constants.rs delete mode 100644 src/sql/common/control.rs delete mode 100644 src/sql/common/converter.rs delete mode 100644 src/sql/common/date.rs delete mode 100644 src/sql/common/debezium.rs delete mode 100644 src/sql/common/errors.rs delete mode 100644 src/sql/common/format_from_opts.rs delete mode 100644 src/sql/common/formats.rs delete mode 100644 src/sql/common/fs_schema.rs delete mode 100644 src/sql/common/kafka_catalog.rs delete mode 100644 src/sql/common/mod.rs delete mode 100644 src/sql/common/operator_config.rs delete mode 100644 src/sql/common/time_utils.rs delete mode 100644 src/sql/common/topology.rs delete mode 100644 src/sql/common/with_option_keys.rs delete mode 100644 src/sql/connector/config.rs delete mode 100644 src/sql/connector/factory.rs delete mode 100644 src/sql/connector/mod.rs delete mode 100644 src/sql/connector/provider.rs delete mode 100644 src/sql/connector/registry.rs delete mode 100644 src/sql/connector/sink/delta.rs delete mode 100644 src/sql/connector/sink/filesystem.rs delete mode 100644 src/sql/connector/sink/iceberg.rs delete mode 100644 src/sql/connector/sink/kafka.rs delete mode 100644 src/sql/connector/sink/lancedb.rs delete mode 100644 src/sql/connector/sink/mod.rs delete mode 100644 src/sql/connector/sink/runtime_config.rs delete mode 100644 src/sql/connector/sink/s3.rs delete mode 100644 src/sql/connector/sink/utils.rs delete mode 100644 src/sql/connector/source/kafka.rs delete mode 100644 src/sql/connector/source/mod.rs delete mode 100644 src/sql/functions/mod.rs delete mode 100644 src/sql/logical_node/aggregate.rs delete mode 100644 src/sql/logical_node/async_udf.rs delete mode 100644 src/sql/logical_node/debezium.rs delete mode 100644 src/sql/logical_node/extension_try_from.rs delete mode 100644 src/sql/logical_node/is_retract.rs delete mode 100644 src/sql/logical_node/join.rs delete mode 100644 src/sql/logical_node/key_calculation.rs delete mode 100644 src/sql/logical_node/logical/dylib_udf_config.rs delete mode 100644 src/sql/logical_node/logical/fs_program_convert.rs delete mode 100644 src/sql/logical_node/logical/logical_edge.rs delete mode 100644 src/sql/logical_node/logical/logical_graph.rs delete mode 100644 src/sql/logical_node/logical/logical_node.rs delete mode 100644 src/sql/logical_node/logical/logical_program.rs delete mode 100644 src/sql/logical_node/logical/mod.rs delete mode 100644 src/sql/logical_node/logical/operator_chain.rs delete mode 100644 src/sql/logical_node/logical/operator_name.rs delete mode 100644 src/sql/logical_node/logical/program_config.rs delete mode 100644 src/sql/logical_node/logical/python_udf_config.rs delete mode 100644 src/sql/logical_node/lookup.rs delete mode 100644 src/sql/logical_node/macros.rs delete mode 100644 src/sql/logical_node/mod.rs delete mode 100644 src/sql/logical_node/projection.rs delete mode 100644 src/sql/logical_node/remote_table.rs delete mode 100644 src/sql/logical_node/sink.rs delete mode 100644 src/sql/logical_node/streaming_operator_blueprint.rs delete mode 100644 src/sql/logical_node/table_source.rs delete mode 100644 src/sql/logical_node/timestamp_append.rs delete mode 100644 src/sql/logical_node/updating_aggregate.rs delete mode 100644 src/sql/logical_node/watermark_node.rs delete mode 100644 src/sql/logical_node/windows_function.rs delete mode 100644 src/sql/logical_planner/mod.rs delete mode 100644 src/sql/logical_planner/optimizers/chaining.rs delete mode 100644 src/sql/logical_planner/optimizers/mod.rs delete mode 100644 src/sql/logical_planner/optimizers/optimized_plan.rs delete mode 100644 src/sql/logical_planner/streaming_planner.rs delete mode 100644 src/sql/mod.rs delete mode 100644 src/sql/parse.rs delete mode 100644 src/sql/physical/cdc/encode.rs delete mode 100644 src/sql/physical/cdc/mod.rs delete mode 100644 src/sql/physical/cdc/unroll.rs delete mode 100644 src/sql/physical/codec.rs delete mode 100644 src/sql/physical/meta.rs delete mode 100644 src/sql/physical/mod.rs delete mode 100644 src/sql/physical/source_exec.rs delete mode 100644 src/sql/physical/udfs.rs delete mode 100644 src/sql/planning_runtime.rs delete mode 100644 src/sql/schema/catalog.rs delete mode 100644 src/sql/schema/column_descriptor.rs delete mode 100644 src/sql/schema/connection_type.rs delete mode 100644 src/sql/schema/data_encoding_format.rs delete mode 100644 src/sql/schema/introspection/ddl_formatter.rs delete mode 100644 src/sql/schema/introspection/mod.rs delete mode 100644 src/sql/schema/introspection/show_formatter.rs delete mode 100644 src/sql/schema/introspection/stream_formatter.rs delete mode 100644 src/sql/schema/mod.rs delete mode 100644 src/sql/schema/schema_provider.rs delete mode 100644 src/sql/schema/table.rs delete mode 100644 src/sql/schema/table_role.rs delete mode 100644 src/sql/schema/temporal_pipeline_config.rs delete mode 100644 src/sql/schema/utils.rs delete mode 100644 src/sql/types/data_type.rs delete mode 100644 src/sql/types/df_field.rs delete mode 100644 src/sql/types/mod.rs delete mode 100644 src/sql/types/placeholder_udf.rs delete mode 100644 src/sql/types/stream_schema.rs delete mode 100644 src/sql/types/window.rs delete mode 100644 src/storage/mod.rs diff --git a/Cargo.lock b/Cargo.lock index 05587163..97388459 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3803,6 +3803,7 @@ dependencies = [ "function-stream-config", "function-stream-logger", "function-stream-runtime-common", + "function-stream-streaming-planner", "futures", "governor", "itertools 0.14.0", @@ -3910,6 +3911,7 @@ name = "function-stream-runtime-common" version = "0.6.0" dependencies = [ "arrow-array 55.2.0", + "bincode", "parking_lot", "serde", "tokio", @@ -3924,6 +3926,48 @@ version = "0.6.0" name = "function-stream-sqlparser" version = "0.6.0" +[[package]] +name = "function-stream-streaming-planner" +version = "0.6.0" +dependencies = [ + "ahash", + "anyhow", + "apache-avro", + "arrow 55.2.0", + "arrow-array 55.2.0", + "arrow-json 55.2.0", + "arrow-schema 55.2.0", + "async-trait", + "bincode", + "bytes", + "chrono", + "datafusion 48.0.1", + "datafusion-common 48.0.1", + "datafusion-execution 48.0.1", + "datafusion-expr 48.0.1", + "datafusion-physical-expr 48.0.1", + "datafusion-proto", + "function-stream-config", + "function-stream-runtime-common", + "futures", + "itertools 0.14.0", + "petgraph 0.7.1", + "prost 0.13.5", + "protocol", + "rand 0.8.5", + "serde", + "serde_json", + "serde_json_path", + "sqlparser 0.55.0", + "strum 0.26.3", + "thiserror 2.0.18", + "tokio", + "tokio-stream", + "tracing", + "unicase", + "xxhash-rust", +] + [[package]] name = "function-stream-streaming-runtime" version = "0.6.0" diff --git a/Cargo.toml b/Cargo.toml index 3dbd8f5c..71c795b6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,6 +14,7 @@ members = [ "src/servicer", "src/sqlparser", "src/streaming_runtime", + "src/streaming_planner", "src/wasm_runtime", ] @@ -50,6 +51,7 @@ protocol = { path = "./protocol" } function-stream-config = { path = "src/config" } function-stream-logger = { path = "src/logger" } function-stream-runtime-common = { path = "src/runtime_common" } +function-stream-streaming-planner = { path = "src/streaming_planner" } prost = "0.13" rdkafka = { version = "0.38", features = ["cmake-build", "ssl", "gssapi", "curl"] } crossbeam-channel = "0.5" diff --git a/src/catalog_storage/src/lib.rs b/src/catalog_storage/src/lib.rs index da6e8fc2..4aa4cf8e 100644 --- a/src/catalog_storage/src/lib.rs +++ b/src/catalog_storage/src/lib.rs @@ -1,7 +1,7 @@ //! Persistent catalog storage implementations. //! //! The stream catalog manager and task persistence (`stream_catalog/`, `task/`) live in this -//! package and are compiled as part of `function-stream` via `#[path]` in `src/storage/mod.rs`. +//! package and are compiled as part of `function-stream` via `#[path]` in `src/storage.rs`. pub mod memory; pub mod rocksdb; diff --git a/src/coordinator/src/legacy/mod.rs b/src/coordinator/src/legacy/mod.rs index 445ec105..640d7692 100644 --- a/src/coordinator/src/legacy/mod.rs +++ b/src/coordinator/src/legacy/mod.rs @@ -13,13 +13,13 @@ mod analyze; #[allow(clippy::module_inception)] mod coordinator; -mod sql_classify; mod dataset; mod execution; mod execution_context; mod plan; mod runtime_context; mod statement; +mod sql_classify; mod streaming_table_options; mod tool; diff --git a/src/coordinator/src/legacy/sql_classify.rs b/src/coordinator/src/legacy/sql_classify.rs index af96a66c..fa43a347 100644 --- a/src/coordinator/src/legacy/sql_classify.rs +++ b/src/coordinator/src/legacy/sql_classify.rs @@ -15,13 +15,13 @@ use std::collections::HashMap; -use datafusion::common::{Result, plan_err}; +use datafusion::common::{plan_err, Result}; use datafusion::error::DataFusionError; use datafusion::sql::sqlparser::ast::{ ObjectType, ShowCreateObject, SqlOption, Statement as DFStatement, }; -use super::statement::{ +use super::{ CreateFunction, CreateTable, DropFunction, DropStreamingTableStatement, DropTableStatement, ShowCatalogTables, ShowCreateStreamingTable, ShowCreateTable, ShowFunctions, ShowStreamingTables, StartFunction, Statement, StopFunction, StreamingTableStatement, @@ -298,8 +298,7 @@ mod tests { #[test] fn test_insert_not_supported() { - let mut stmts = - parse_sql("INSERT INTO sink SELECT * FROM src").unwrap(); + let mut stmts = parse_sql("INSERT INTO sink SELECT * FROM src").unwrap(); let err = classify_statement(stmts.remove(0)).unwrap_err(); let msg = err.to_string(); assert!( diff --git a/src/lib.rs b/src/lib.rs index 4c1af5d9..177d8f72 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -21,5 +21,5 @@ pub use function_stream_logger as logging; pub mod runtime; #[path = "servicer/src/legacy/mod.rs"] pub mod server; -pub mod sql; +pub use function_stream_streaming_planner as sql; pub mod storage; diff --git a/src/main.rs b/src/main.rs index 49978af3..eb5bcfec 100644 --- a/src/main.rs +++ b/src/main.rs @@ -19,7 +19,7 @@ pub use function_stream_logger as logging; mod runtime; #[path = "servicer/src/legacy/mod.rs"] mod server; -mod sql; +pub use function_stream_streaming_planner as sql; mod storage; use anyhow::{Context, Result}; diff --git a/src/runtime/mod.rs b/src/runtime/mod.rs deleted file mode 100644 index 28f62306..00000000 --- a/src/runtime/mod.rs +++ /dev/null @@ -1,31 +0,0 @@ -// 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. - -// Runtime module - -pub use function_stream_runtime_common::{common, memory}; - -// Implementation sources live under `src/streaming_runtime/src/{streaming,util}/` and are -// compiled here so `crate::sql` / `crate::runtime::memory` paths keep resolving. -#[path = "../streaming_runtime/src/streaming/mod.rs"] -pub mod streaming; - -#[path = "../streaming_runtime/src/util/mod.rs"] -pub mod util; - -// WASM runtime sources live under `src/wasm_runtime/src/wasm/`; compiled here for `crate::` paths. -#[path = "../wasm_runtime/src/wasm/mod.rs"] -pub mod wasm; - -pub use wasm::input; -pub use wasm::output; -pub use wasm::processor; diff --git a/src/runtime_common/Cargo.toml b/src/runtime_common/Cargo.toml index 0a910f18..b7747b2a 100644 --- a/src/runtime_common/Cargo.toml +++ b/src/runtime_common/Cargo.toml @@ -9,6 +9,7 @@ path = "src/lib.rs" [dependencies] arrow-array = "55" +bincode = { version = "2", features = ["serde"] } parking_lot = "0.12" serde = { version = "1.0", features = ["derive"] } tokio = { version = "1.0", features = ["sync"] } diff --git a/src/runtime_common/src/lib.rs b/src/runtime_common/src/lib.rs index 3fad765f..1c97062b 100644 --- a/src/runtime_common/src/lib.rs +++ b/src/runtime_common/src/lib.rs @@ -14,3 +14,4 @@ pub mod common; pub mod memory; +pub mod streaming_protocol; diff --git a/src/sql/analysis/aggregate_rewriter.rs b/src/sql/analysis/aggregate_rewriter.rs deleted file mode 100644 index ddcb0294..00000000 --- a/src/sql/analysis/aggregate_rewriter.rs +++ /dev/null @@ -1,279 +0,0 @@ -// 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 datafusion::common::tree_node::{Transformed, TreeNodeRewriter}; -use datafusion::common::{DFSchema, DataFusionError, Result, not_impl_err, plan_err}; -use datafusion::functions_aggregate::expr_fn::max; -use datafusion::logical_expr::{Aggregate, Expr, Extension, LogicalPlan, Projection}; -use datafusion::prelude::col; -use std::sync::Arc; - -use crate::sql::analysis::streaming_window_analzer::StreamingWindowAnalzer; -use crate::sql::logical_node::aggregate::StreamWindowAggregateNode; -use crate::sql::logical_node::key_calculation::{KeyExtractionNode, KeyExtractionStrategy}; -use crate::sql::logical_node::updating_aggregate::ContinuousAggregateNode; -use crate::sql::schema::StreamSchemaProvider; -use crate::sql::types::{ - QualifiedField, TIMESTAMP_FIELD, WindowBehavior, WindowType, build_df_schema_with_metadata, - extract_qualified_fields, extract_window_type, -}; - -/// AggregateRewriter transforms batch DataFusion aggregates into streaming stateful operators. -/// It handles windowing (Tumble/Hop/Session), watermarks, and continuous updating aggregates. -pub(crate) struct AggregateRewriter<'a> { - pub schema_provider: &'a StreamSchemaProvider, -} - -impl TreeNodeRewriter for AggregateRewriter<'_> { - type Node = LogicalPlan; - - fn f_up(&mut self, node: Self::Node) -> Result> { - let LogicalPlan::Aggregate(mut agg) = node else { - return Ok(Transformed::no(node)); - }; - - // 1. Identify windowing functions (e.g., tumble, hop) in GROUP BY. - let mut window_exprs: Vec<_> = agg - .group_expr - .iter() - .enumerate() - .filter_map(|(i, e)| { - extract_window_type(e) - .map(|opt| opt.map(|w| (i, w))) - .transpose() - }) - .collect::>>()?; - - if window_exprs.len() > 1 { - return not_impl_err!("Streaming aggregates support at most one window expression"); - } - - // 2. Prepare internal metadata for Key-based distribution. - let mut key_fields: Vec = extract_qualified_fields(&agg.schema) - .iter() - .take(agg.group_expr.len()) - .map(|f| { - QualifiedField::new( - f.qualifier().cloned(), - format!("_key_{}", f.name()), - f.data_type().clone(), - f.is_nullable(), - ) - }) - .collect(); - - // 3. Dispatch to ContinuousAggregateNode (UpdatingAggregate) if no windowing is detected. - let input_window = StreamingWindowAnalzer::get_window(&agg.input)?; - if window_exprs.is_empty() && input_window.is_none() { - return self.rewrite_as_continuous_updating_aggregate( - agg.input, - key_fields, - agg.group_expr, - agg.aggr_expr, - agg.schema, - ); - } - - // 4. Resolve Windowing Strategy (InData vs FromOperator). - let behavior = self.resolve_window_context( - &agg.input, - &mut agg.group_expr, - &agg.schema, - &mut window_exprs, - )?; - - // Adjust keys if windowing is handled by the operator. - if let WindowBehavior::FromOperator { window_index, .. } = &behavior { - key_fields.remove(*window_index); - } - - let key_count = key_fields.len(); - let keyed_input = - self.build_keyed_input(agg.input.clone(), &agg.group_expr, &key_fields)?; - - // 5. Build the final StreamWindowAggregateNode for the physical planner. - let mut internal_fields = extract_qualified_fields(&agg.schema); - if let WindowBehavior::FromOperator { window_index, .. } = &behavior { - internal_fields.remove(*window_index); - } - let internal_schema = Arc::new(build_df_schema_with_metadata( - &internal_fields, - agg.schema.metadata().clone(), - )?); - - let rewritten_agg = Aggregate::try_new_with_schema( - Arc::new(keyed_input), - agg.group_expr, - agg.aggr_expr, - internal_schema, - )?; - - let extension = StreamWindowAggregateNode::try_new( - behavior, - LogicalPlan::Aggregate(rewritten_agg), - (0..key_count).collect(), - )?; - - Ok(Transformed::yes(LogicalPlan::Extension(Extension { - node: Arc::new(extension), - }))) - } -} - -impl<'a> AggregateRewriter<'a> { - pub fn new(schema_provider: &'a StreamSchemaProvider) -> Self { - Self { schema_provider } - } - - /// [Internal] Builds the physical Key Calculation layer required for distributed Shuffling. - /// This wraps the input in a Projection and a KeyExtractionNode. - fn build_keyed_input( - &self, - input: Arc, - group_expr: &[Expr], - key_fields: &[QualifiedField], - ) -> Result { - let key_count = group_expr.len(); - let mut projection_fields = key_fields.to_vec(); - projection_fields.extend(extract_qualified_fields(input.schema())); - - let key_schema = Arc::new(build_df_schema_with_metadata( - &projection_fields, - input.schema().metadata().clone(), - )?); - - // Map group expressions to '_key_' aliases while passing through all original columns. - let mut exprs: Vec<_> = group_expr - .iter() - .zip(key_fields.iter()) - .map(|(expr, f)| expr.clone().alias(f.name().to_string())) - .collect(); - - exprs.extend( - extract_qualified_fields(input.schema()) - .iter() - .map(|f| Expr::Column(f.qualified_column())), - ); - - let projection = - LogicalPlan::Projection(Projection::try_new_with_schema(exprs, input, key_schema)?); - - Ok(LogicalPlan::Extension(Extension { - node: Arc::new(KeyExtractionNode::new( - projection, - KeyExtractionStrategy::ColumnIndices((0..key_count).collect()), - )), - })) - } - - /// [Strategy] Rewrites standard GROUP BY into a ContinuousAggregateNode with retraction semantics. - /// Injected max(_timestamp) ensures the streaming pulse (Watermark) continues to propagate. - fn rewrite_as_continuous_updating_aggregate( - &self, - input: Arc, - key_fields: Vec, - group_expr: Vec, - mut aggr_expr: Vec, - schema: Arc, - ) -> Result> { - let key_count = key_fields.len(); - let keyed_input = self.build_keyed_input(input, &group_expr, &key_fields)?; - - // Ensure the updating stream maintains time awareness. - let timestamp_col = keyed_input - .schema() - .qualified_field_with_unqualified_name(TIMESTAMP_FIELD) - .map_err(|_| { - DataFusionError::Plan( - "Required _timestamp field missing for updating aggregate".to_string(), - ) - })?; - - let timestamp_field: QualifiedField = timestamp_col.into(); - aggr_expr.push(max(col(timestamp_field.qualified_column())).alias(TIMESTAMP_FIELD)); - - let mut output_fields = extract_qualified_fields(&schema); - output_fields.push(timestamp_field); - - let output_schema = Arc::new(build_df_schema_with_metadata( - &output_fields, - schema.metadata().clone(), - )?); - - let base_aggregate = Aggregate::try_new_with_schema( - Arc::new(keyed_input), - group_expr, - aggr_expr, - output_schema, - )?; - - let continuous_node = ContinuousAggregateNode::try_new( - LogicalPlan::Aggregate(base_aggregate), - (0..key_count).collect(), - None, - self.schema_provider.planning_options.ttl, - )?; - - Ok(Transformed::yes(LogicalPlan::Extension(Extension { - node: Arc::new(continuous_node), - }))) - } - - /// [Strategy] Reconciles window definitions between the input stream and the current GROUP BY. - fn resolve_window_context( - &self, - input: &LogicalPlan, - group_expr: &mut Vec, - schema: &DFSchema, - window_expr_info: &mut Vec<(usize, WindowType)>, - ) -> Result { - let mut visitor = StreamingWindowAnalzer::default(); - input.visit_with_subqueries(&mut visitor)?; - - let input_window = visitor.window; - let has_group_window = !window_expr_info.is_empty(); - - match (input_window, has_group_window) { - (Some(i_win), true) => { - let (idx, g_win) = window_expr_info.pop().unwrap(); - if i_win != g_win { - return plan_err!("Inconsistent windowing detected"); - } - - if let Some(field) = visitor.fields.iter().next() { - group_expr[idx] = Expr::Column(field.qualified_column()); - Ok(WindowBehavior::InData) - } else { - group_expr.remove(idx); - Ok(WindowBehavior::FromOperator { - window: i_win, - window_field: schema.qualified_field(idx).into(), - window_index: idx, - is_nested: true, - }) - } - } - (None, true) => { - let (idx, g_win) = window_expr_info.pop().unwrap(); - group_expr.remove(idx); - Ok(WindowBehavior::FromOperator { - window: g_win, - window_field: schema.qualified_field(idx).into(), - window_index: idx, - is_nested: false, - }) - } - (Some(_), false) => Ok(WindowBehavior::InData), - _ => unreachable!("Handled by updating path"), - } - } -} diff --git a/src/sql/analysis/async_udf_rewriter.rs b/src/sql/analysis/async_udf_rewriter.rs deleted file mode 100644 index d6d9b54b..00000000 --- a/src/sql/analysis/async_udf_rewriter.rs +++ /dev/null @@ -1,133 +0,0 @@ -// 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::sql::common::constants::sql_field; -use crate::sql::logical_node::AsyncFunctionExecutionNode; -use crate::sql::logical_node::remote_table::RemoteTableBoundaryNode; -use crate::sql::schema::StreamSchemaProvider; -use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRewriter}; -use datafusion::common::{Column, Result as DFResult, TableReference, plan_err}; -use datafusion::logical_expr::expr::ScalarFunction; -use datafusion::logical_expr::{Expr, Extension, LogicalPlan}; -use std::sync::Arc; -use std::time::Duration; - -type AsyncSplitResult = (String, AsyncOptions, Vec); - -#[derive(Debug, Clone, Copy)] -pub struct AsyncOptions { - pub ordered: bool, - pub max_concurrency: usize, - pub timeout: Duration, -} - -pub struct AsyncUdfRewriter<'a> { - provider: &'a StreamSchemaProvider, -} - -impl<'a> AsyncUdfRewriter<'a> { - pub fn new(provider: &'a StreamSchemaProvider) -> Self { - Self { provider } - } - - fn split_async( - expr: Expr, - provider: &StreamSchemaProvider, - ) -> DFResult<(Expr, Option)> { - let mut found: Option<(String, AsyncOptions, Vec)> = None; - let expr = expr.transform_up(|e| { - if let Expr::ScalarFunction(ScalarFunction { func: udf, args }) = &e - && let Some(opts) = provider.get_async_udf_options(udf.name()) - { - if found - .replace((udf.name().to_string(), opts, args.clone())) - .is_some() - { - return plan_err!( - "multiple async calls in the same expression, which is not allowed" - ); - } - return Ok(Transformed::yes(Expr::Column(Column::new_unqualified( - sql_field::ASYNC_RESULT, - )))); - } - Ok(Transformed::no(e)) - })?; - - Ok((expr.data, found)) - } -} - -impl TreeNodeRewriter for AsyncUdfRewriter<'_> { - type Node = LogicalPlan; - - fn f_up(&mut self, node: Self::Node) -> DFResult> { - let LogicalPlan::Projection(mut projection) = node else { - for e in node.expressions() { - if let (_, Some((udf, _, _))) = Self::split_async(e.clone(), self.provider)? { - return plan_err!( - "async UDFs are only supported in projections, but {udf} was called in another context" - ); - } - } - return Ok(Transformed::no(node)); - }; - - let mut args = None; - for e in projection.expr.iter_mut() { - let (new_e, Some(udf)) = Self::split_async(e.clone(), self.provider)? else { - continue; - }; - if let Some((prev, _, _)) = args.replace(udf) { - return plan_err!( - "Projection contains multiple async UDFs, which is not supported \ - \n(hint: two async UDF calls, {} and {}, appear in the same SELECT statement)", - prev, - args.unwrap().0 - ); - } - *e = new_e; - } - - let Some((name, opts, arg_exprs)) = args else { - return Ok(Transformed::no(LogicalPlan::Projection(projection))); - }; - let udf = self.provider.dylib_udfs.get(&name).unwrap().clone(); - - let input = if matches!(*projection.input, LogicalPlan::Projection(..)) { - Arc::new(LogicalPlan::Extension(Extension { - node: Arc::new(RemoteTableBoundaryNode { - upstream_plan: (*projection.input).clone(), - table_identifier: TableReference::bare("subquery_projection"), - resolved_schema: projection.input.schema().clone(), - requires_materialization: false, - }), - })) - } else { - projection.input - }; - - Ok(Transformed::yes(LogicalPlan::Extension(Extension { - node: Arc::new(AsyncFunctionExecutionNode { - upstream_plan: input, - operator_name: name, - function_config: udf, - invocation_args: arg_exprs, - result_projections: projection.expr, - preserve_ordering: opts.ordered, - concurrency_limit: opts.max_concurrency, - execution_timeout: opts.timeout, - resolved_schema: projection.schema, - }), - }))) - } -} diff --git a/src/sql/analysis/join_rewriter.rs b/src/sql/analysis/join_rewriter.rs deleted file mode 100644 index 8a9e5280..00000000 --- a/src/sql/analysis/join_rewriter.rs +++ /dev/null @@ -1,234 +0,0 @@ -// 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::sql::analysis::streaming_window_analzer::StreamingWindowAnalzer; -use crate::sql::common::TIMESTAMP_FIELD; -use crate::sql::common::constants::mem_exec_join_side; -use crate::sql::logical_node::join::StreamingJoinNode; -use crate::sql::logical_node::key_calculation::KeyExtractionNode; -use crate::sql::schema::StreamSchemaProvider; -use crate::sql::types::{WindowType, build_df_schema_with_metadata, extract_qualified_fields}; -use datafusion::common::tree_node::{Transformed, TreeNodeRewriter}; -use datafusion::common::{ - JoinConstraint, JoinType, Result, ScalarValue, TableReference, not_impl_err, plan_err, -}; -use datafusion::logical_expr::{ - self, BinaryExpr, Case, Expr, Extension, Join, LogicalPlan, Projection, build_join_schema, -}; -use datafusion::prelude::coalesce; -use std::sync::Arc; - -/// JoinRewriter handles the transformation of standard SQL joins into streaming-capable joins. -/// It manages stateful "Updating Joins" and time-aligned "Instant Joins". -pub(crate) struct JoinRewriter<'a> { - pub schema_provider: &'a StreamSchemaProvider, -} - -impl<'a> JoinRewriter<'a> { - pub fn new(schema_provider: &'a StreamSchemaProvider) -> Self { - Self { schema_provider } - } - - /// [Validation] Ensures left and right streams have compatible windowing strategies. - fn validate_join_windows(&self, join: &Join) -> Result { - let left_win = StreamingWindowAnalzer::get_window(&join.left)?; - let right_win = StreamingWindowAnalzer::get_window(&join.right)?; - - match (left_win, right_win) { - (None, None) => { - if join.join_type == JoinType::Inner { - Ok(false) // Standard Updating Join (Inner) - } else { - plan_err!( - "Non-inner joins (e.g., LEFT/RIGHT) require windowing to bound state." - ) - } - } - (Some(l), Some(r)) => { - if l != r { - return plan_err!( - "Join window mismatch: left={:?}, right={:?}. Windows must match exactly.", - l, - r - ); - } - if let WindowType::Session { .. } = l { - return plan_err!( - "Session windows are currently not supported in streaming joins." - ); - } - Ok(true) // Instant Windowed Join - } - _ => plan_err!( - "Mixed windowing detected. Both sides of a join must be either windowed or non-windowed." - ), - } - } - - /// [Internal] Wraps a join input in a key-extraction layer to facilitate shuffle / key-by distribution. - fn build_keyed_side( - &self, - input: Arc, - keys: Vec, - side: &str, - ) -> Result { - let key_count = keys.len(); - - let projection_exprs = keys - .into_iter() - .enumerate() - .map(|(i, e)| { - e.alias_qualified(Some(TableReference::bare("_stream")), format!("_key_{i}")) - }) - .chain( - extract_qualified_fields(input.schema()) - .iter() - .map(|f| Expr::Column(f.qualified_column())), - ) - .collect(); - - let projection = Projection::try_new(projection_exprs, input)?; - let key_ext = KeyExtractionNode::try_new_with_projection( - LogicalPlan::Projection(projection), - (0..key_count).collect(), - side.to_string(), - )?; - - Ok(LogicalPlan::Extension(Extension { - node: Arc::new(key_ext), - })) - } - - /// [Strategy] Resolves the output timestamp of the join. - /// Streaming joins must output the 'max' of the two input timestamps to ensure Watermark progression. - fn apply_timestamp_resolution(&self, join_plan: LogicalPlan) -> Result { - let schema = join_plan.schema(); - let all_fields = extract_qualified_fields(schema); - - let timestamp_fields: Vec<_> = all_fields - .iter() - .filter(|f| f.name() == "_timestamp") - .cloned() - .collect(); - - if timestamp_fields.len() != 2 { - return plan_err!( - "Streaming join requires exactly two input timestamp fields to resolve output time." - ); - } - - // Project all fields except the two raw timestamps - let mut exprs: Vec<_> = all_fields - .iter() - .filter(|f| f.name() != "_timestamp") - .map(|f| Expr::Column(f.qualified_column())) - .collect(); - - // Calculate: GREATEST(left._timestamp, right._timestamp) - let left_ts = Expr::Column(timestamp_fields[0].qualified_column()); - let right_ts = Expr::Column(timestamp_fields[1].qualified_column()); - - let max_ts_expr = Expr::Case(Case { - expr: Some(Box::new(Expr::BinaryExpr(BinaryExpr { - left: Box::new(left_ts.clone()), - op: logical_expr::Operator::GtEq, - right: Box::new(right_ts.clone()), - }))), - when_then_expr: vec![ - ( - Box::new(Expr::Literal(ScalarValue::Boolean(Some(true)), None)), - Box::new(left_ts.clone()), - ), - ( - Box::new(Expr::Literal(ScalarValue::Boolean(Some(false)), None)), - Box::new(right_ts.clone()), - ), - ], - else_expr: Some(Box::new(coalesce(vec![left_ts, right_ts]))), - }) - .alias(TIMESTAMP_FIELD); - - exprs.push(max_ts_expr); - - let out_fields: Vec<_> = all_fields - .iter() - .filter(|f| f.name() != "_timestamp") - .cloned() - .chain(std::iter::once(timestamp_fields[0].clone())) - .collect(); - - let out_schema = Arc::new(build_df_schema_with_metadata( - &out_fields, - schema.metadata().clone(), - )?); - - Ok(LogicalPlan::Projection(Projection::try_new_with_schema( - exprs, - Arc::new(join_plan), - out_schema, - )?)) - } -} - -impl TreeNodeRewriter for JoinRewriter<'_> { - type Node = LogicalPlan; - - fn f_up(&mut self, node: Self::Node) -> Result> { - let LogicalPlan::Join(join) = node else { - return Ok(Transformed::no(node)); - }; - - // 1. Validate Streaming Context - let is_instant = self.validate_join_windows(&join)?; - if join.join_constraint != JoinConstraint::On { - return not_impl_err!("Only 'ON' join constraints are supported in streaming SQL."); - } - if join.on.is_empty() && !is_instant { - return plan_err!("Updating joins require at least one equality condition (Equijoin)."); - } - - // 2. Prepare Keyed Inputs for Shuffle - let (left_on, right_on): (Vec<_>, Vec<_>) = join.on.clone().into_iter().unzip(); - let keyed_left = self.build_keyed_side(join.left, left_on, mem_exec_join_side::LEFT)?; - let keyed_right = self.build_keyed_side(join.right, right_on, mem_exec_join_side::RIGHT)?; - - // 3. Assemble Rewritten Join Node - let join_schema = Arc::new(build_join_schema( - keyed_left.schema(), - keyed_right.schema(), - &join.join_type, - )?); - let rewritten_join = LogicalPlan::Join(Join { - left: Arc::new(keyed_left), - right: Arc::new(keyed_right), - on: join.on, - filter: join.filter, - join_type: join.join_type, - join_constraint: JoinConstraint::On, - schema: join_schema, - null_equals_null: false, - }); - - // 4. Resolve Output Watermark (Timestamp Projection) - let plan_with_timestamp = self.apply_timestamp_resolution(rewritten_join)?; - - // 5. Wrap in StreamingJoinNode for physical planning - let state_retention_ttl = - (!is_instant).then_some(self.schema_provider.planning_options.ttl); - let extension = - StreamingJoinNode::new(plan_with_timestamp, is_instant, state_retention_ttl); - - Ok(Transformed::yes(LogicalPlan::Extension(Extension { - node: Arc::new(extension), - }))) - } -} diff --git a/src/sql/analysis/mod.rs b/src/sql/analysis/mod.rs deleted file mode 100644 index 019d8bf1..00000000 --- a/src/sql/analysis/mod.rs +++ /dev/null @@ -1,214 +0,0 @@ -// 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. - -#![allow(clippy::new_without_default)] - -pub(crate) mod aggregate_rewriter; -pub(crate) mod join_rewriter; -pub(crate) mod row_time_rewriter; -pub(crate) mod stream_rewriter; -pub(crate) mod streaming_window_analzer; -pub(crate) mod window_function_rewriter; - -pub mod async_udf_rewriter; -pub mod sink_input_rewriter; -pub mod source_metadata_visitor; -pub mod source_rewriter; -pub mod time_window; -pub mod unnest_rewriter; - -pub use async_udf_rewriter::AsyncOptions; -pub use sink_input_rewriter::SinkInputRewriter; -pub use time_window::{TimeWindowNullCheckRemover, TimeWindowUdfChecker}; -pub use unnest_rewriter::UNNESTED_COL; - -pub use crate::sql::schema::schema_provider::StreamSchemaProvider; - -use std::collections::HashMap; -use std::sync::Arc; - -use datafusion::common::tree_node::{Transformed, TreeNode}; -use datafusion::common::{Result, plan_err}; -use datafusion::error::DataFusionError; -use datafusion::logical_expr::{Extension, LogicalPlan, UserDefinedLogicalNodeCore}; -use tracing::{debug, info, instrument}; - -use crate::sql::logical_node::StreamingOperatorBlueprint; -use crate::sql::logical_node::key_calculation::{KeyExtractionNode, KeyExtractionStrategy}; -use crate::sql::logical_node::projection::StreamProjectionNode; -use crate::sql::logical_node::sink::StreamEgressNode; -use crate::sql::logical_planner::planner::NamedNode; - -fn duration_from_sql_expr( - expr: &datafusion::sql::sqlparser::ast::Expr, -) -> Result { - use datafusion::sql::sqlparser::ast::Expr as SqlExpr; - use datafusion::sql::sqlparser::ast::Value as SqlValue; - use datafusion::sql::sqlparser::ast::ValueWithSpan; - - match expr { - SqlExpr::Interval(interval) => { - let value_str = match interval.value.as_ref() { - SqlExpr::Value(ValueWithSpan { - value: SqlValue::SingleQuotedString(s), - .. - }) => s.clone(), - other => return plan_err!("expected interval string literal, found {other}"), - }; - - parse_interval_to_duration(&value_str) - } - SqlExpr::Value(ValueWithSpan { - value: SqlValue::SingleQuotedString(s), - .. - }) => parse_interval_to_duration(s), - other => plan_err!("expected an interval expression, found {other}"), - } -} - -fn parse_interval_to_duration(s: &str) -> Result { - let parts: Vec<&str> = s.split_whitespace().collect(); - if parts.len() != 2 { - return plan_err!("invalid interval string '{s}'; expected ' '"); - } - let value: u64 = parts[0] - .parse() - .map_err(|_| DataFusionError::Plan(format!("invalid interval number: {}", parts[0])))?; - match parts[1].to_lowercase().as_str() { - "second" | "seconds" | "s" => Ok(std::time::Duration::from_secs(value)), - "minute" | "minutes" | "min" => Ok(std::time::Duration::from_secs(value * 60)), - "hour" | "hours" | "h" => Ok(std::time::Duration::from_secs(value * 3600)), - "day" | "days" | "d" => Ok(std::time::Duration::from_secs(value * 86400)), - unit => plan_err!("unsupported interval unit '{unit}'"), - } -} - -fn build_sink_inputs(extensions: &[LogicalPlan]) -> HashMap> { - let mut sink_inputs = HashMap::>::new(); - for extension in extensions.iter() { - if let LogicalPlan::Extension(ext) = extension - && let Some(sink_node) = ext.node.as_any().downcast_ref::() - && let Some(named_node) = sink_node.operator_identity() - { - let inputs = sink_node - .inputs() - .into_iter() - .cloned() - .collect::>(); - sink_inputs.entry(named_node).or_default().extend(inputs); - } - } - sink_inputs -} - -pub(crate) fn maybe_add_key_extension_to_sink(plan: LogicalPlan) -> Result { - let LogicalPlan::Extension(ref ext) = plan else { - return Ok(plan); - }; - - let Some(sink) = ext.node.as_any().downcast_ref::() else { - return Ok(plan); - }; - - let Some(partition_exprs) = sink.destination_table.partition_exprs() else { - return Ok(plan); - }; - - if partition_exprs.is_empty() { - return Ok(plan); - } - - let inputs = plan - .inputs() - .into_iter() - .map(|input| { - Ok(LogicalPlan::Extension(Extension { - node: Arc::new(KeyExtractionNode { - operator_label: Some("key-calc-partition".to_string()), - resolved_schema: input.schema().clone(), - upstream_plan: input.clone(), - extraction_strategy: KeyExtractionStrategy::CalculatedExpressions( - partition_exprs.clone(), - ), - }), - })) - }) - .collect::>()?; - - use datafusion::prelude::col; - let unkey = LogicalPlan::Extension(Extension { - node: Arc::new( - StreamProjectionNode::try_new( - inputs, - Some("unkey".to_string()), - sink.schema().iter().map(|(_, f)| col(f.name())).collect(), - )? - .with_shuffle_routing(), - ), - }); - - let node = sink.with_exprs_and_inputs(vec![], vec![unkey])?; - Ok(LogicalPlan::Extension(Extension { - node: Arc::new(node), - })) -} - -pub fn rewrite_sinks(extensions: Vec) -> Result> { - let mut sink_inputs = build_sink_inputs(&extensions); - let mut new_extensions = vec![]; - for extension in extensions { - let mut rewriter = SinkInputRewriter::new(&mut sink_inputs); - let result = extension.rewrite(&mut rewriter)?; - if !rewriter.was_removed { - new_extensions.push(result.data); - } - } - - new_extensions - .into_iter() - .map(maybe_add_key_extension_to_sink) - .collect() -} - -/// Entry point for transforming a standard DataFusion LogicalPlan into a -/// Streaming-aware LogicalPlan. -/// -/// This function coordinates multiple rewriting passes and ensures the -/// resulting plan satisfies streaming constraints. -#[instrument(skip_all, level = "debug")] -pub fn rewrite_plan( - plan: LogicalPlan, - schema_provider: &StreamSchemaProvider, -) -> Result { - info!("Starting streaming plan rewrite pipeline"); - - let Transformed { data: plan, .. } = - plan.rewrite_with_subqueries(&mut source_rewriter::SourceRewriter::new(schema_provider))?; - - let mut rewriter = stream_rewriter::StreamRewriter::new(schema_provider); - let Transformed { - data: rewritten_plan, - .. - } = plan.rewrite_with_subqueries(&mut rewriter)?; - - rewritten_plan.visit_with_subqueries(&mut TimeWindowUdfChecker {})?; - - if cfg!(debug_assertions) { - debug!( - "Streaming logical plan graphviz:\n{}", - rewritten_plan.display_graphviz() - ); - } - - info!("Streaming plan rewrite completed successfully"); - Ok(rewritten_plan) -} diff --git a/src/sql/analysis/row_time_rewriter.rs b/src/sql/analysis/row_time_rewriter.rs deleted file mode 100644 index 13e2a048..00000000 --- a/src/sql/analysis/row_time_rewriter.rs +++ /dev/null @@ -1,49 +0,0 @@ -// 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 datafusion::common::tree_node::{Transformed, TreeNodeRewriter}; -use datafusion::common::{Column, Result as DFResult}; -use datafusion::logical_expr::Expr; - -use crate::sql::common::constants::planning_placeholder_udf; -use crate::sql::types::TIMESTAMP_FIELD; - -/// Replaces the virtual `row_time()` scalar function with a physical reference to `_timestamp`. -/// -/// This is a critical mapping step that allows users to use a friendly SQL function -/// while the engine operates on the mandatory internal streaming timestamp. -pub struct RowTimeRewriter; - -impl TreeNodeRewriter for RowTimeRewriter { - type Node = Expr; - - fn f_down(&mut self, node: Self::Node) -> DFResult> { - // Use pattern matching to identify the `row_time` scalar function. - if let Expr::ScalarFunction(func) = &node - && func.name() == planning_placeholder_udf::ROW_TIME - { - // Map the virtual function to the physical internal timestamp column. - // We use .alias() to preserve the original name "row_time()" in the output schema, - // ensuring that user-facing column names do not change unexpectedly. - let physical_col = Expr::Column(Column { - relation: None, - name: TIMESTAMP_FIELD.to_string(), - spans: Default::default(), - }) - .alias("row_time()"); - - return Ok(Transformed::yes(physical_col)); - } - - Ok(Transformed::no(node)) - } -} diff --git a/src/sql/analysis/sink_input_rewriter.rs b/src/sql/analysis/sink_input_rewriter.rs deleted file mode 100644 index 201439cc..00000000 --- a/src/sql/analysis/sink_input_rewriter.rs +++ /dev/null @@ -1,57 +0,0 @@ -// 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::sql::logical_node::StreamingOperatorBlueprint; -use crate::sql::logical_node::sink::StreamEgressNode; -use crate::sql::logical_planner::planner::NamedNode; -use datafusion::common::Result as DFResult; -use datafusion::common::tree_node::{Transformed, TreeNodeRecursion, TreeNodeRewriter}; -use datafusion::logical_expr::{Extension, LogicalPlan, UserDefinedLogicalNodeCore}; -use std::collections::HashMap; -use std::sync::Arc; - -type SinkInputs = HashMap>; - -/// Merges inputs for sinks with the same name to avoid duplicate sinks in the plan. -pub struct SinkInputRewriter<'a> { - sink_inputs: &'a mut SinkInputs, - pub was_removed: bool, -} - -impl<'a> SinkInputRewriter<'a> { - pub(crate) fn new(sink_inputs: &'a mut SinkInputs) -> Self { - Self { - sink_inputs, - was_removed: false, - } - } -} - -impl TreeNodeRewriter for SinkInputRewriter<'_> { - type Node = LogicalPlan; - - fn f_down(&mut self, node: Self::Node) -> DFResult> { - if let LogicalPlan::Extension(extension) = &node - && let Some(sink_node) = extension.node.as_any().downcast_ref::() - && let Some(named_node) = sink_node.operator_identity() - { - if let Some(inputs) = self.sink_inputs.remove(&named_node) { - let new_node = LogicalPlan::Extension(Extension { - node: Arc::new(sink_node.with_exprs_and_inputs(vec![], inputs)?), - }); - return Ok(Transformed::new(new_node, true, TreeNodeRecursion::Jump)); - } - self.was_removed = true; - } - Ok(Transformed::no(node)) - } -} diff --git a/src/sql/analysis/source_metadata_visitor.rs b/src/sql/analysis/source_metadata_visitor.rs deleted file mode 100644 index b37d3a1b..00000000 --- a/src/sql/analysis/source_metadata_visitor.rs +++ /dev/null @@ -1,73 +0,0 @@ -// 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::sql::logical_node::sink::{STREAM_EGRESS_NODE_NAME, StreamEgressNode}; -use crate::sql::logical_node::table_source::{STREAM_INGESTION_NODE_NAME, StreamIngestionNode}; -use crate::sql::schema::StreamSchemaProvider; -use datafusion::common::Result as DFResult; -use datafusion::common::tree_node::{TreeNodeRecursion, TreeNodeVisitor}; -use datafusion::logical_expr::{Extension, LogicalPlan}; -use std::collections::HashSet; - -/// Collects connection IDs from source and sink nodes in the logical plan. -pub struct SourceMetadataVisitor<'a> { - schema_provider: &'a StreamSchemaProvider, - pub connection_ids: HashSet, -} - -impl<'a> SourceMetadataVisitor<'a> { - pub fn new(schema_provider: &'a StreamSchemaProvider) -> Self { - Self { - schema_provider, - connection_ids: HashSet::new(), - } - } - - fn get_connection_id(&self, node: &LogicalPlan) -> Option { - let LogicalPlan::Extension(Extension { node }) = node else { - return None; - }; - - let table_name = match node.name() { - name if name == STREAM_INGESTION_NODE_NAME => { - let ext = node.as_any().downcast_ref::()?; - ext.source_identifier.to_string() - } - name if name == STREAM_EGRESS_NODE_NAME => { - let ext = node.as_any().downcast_ref::()?; - ext.target_identifier.to_string() - } - _ => return None, - }; - - let table = self.schema_provider.get_catalog_table(&table_name)?; - match table { - crate::sql::schema::table::CatalogEntity::ExternalConnector(b) => match b.as_ref() { - crate::sql::schema::catalog::ExternalTable::Source(t) => t.registry_id, - crate::sql::schema::catalog::ExternalTable::Lookup(t) => t.registry_id, - _ => None, - }, - _ => None, - } - } -} - -impl TreeNodeVisitor<'_> for SourceMetadataVisitor<'_> { - type Node = LogicalPlan; - - fn f_down(&mut self, node: &Self::Node) -> DFResult { - if let Some(id) = self.get_connection_id(node) { - self.connection_ids.insert(id); - } - Ok(TreeNodeRecursion::Continue) - } -} diff --git a/src/sql/analysis/source_rewriter.rs b/src/sql/analysis/source_rewriter.rs deleted file mode 100644 index 620ea336..00000000 --- a/src/sql/analysis/source_rewriter.rs +++ /dev/null @@ -1,305 +0,0 @@ -// 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::sync::Arc; -use std::time::Duration; - -use datafusion::common::ScalarValue; -use datafusion::common::tree_node::{Transformed, TreeNodeRewriter}; -use datafusion::common::{Column, DataFusionError, Result as DFResult, TableReference, plan_err}; -use datafusion::logical_expr::{ - self, BinaryExpr, Expr, Extension, LogicalPlan, Projection, TableScan, -}; - -use crate::sql::common::UPDATING_META_FIELD; -use crate::sql::logical_node::debezium::UnrollDebeziumPayloadNode; -use crate::sql::logical_node::remote_table::RemoteTableBoundaryNode; -use crate::sql::logical_node::table_source::StreamIngestionNode; -use crate::sql::logical_node::watermark_node::EventTimeWatermarkNode; -use crate::sql::schema::ColumnDescriptor; -use crate::sql::schema::StreamSchemaProvider; -use crate::sql::schema::catalog::{ExternalTable, SourceTable}; -use crate::sql::schema::table::CatalogEntity; -use crate::sql::types::TIMESTAMP_FIELD; - -/// Rewrites table scans: projections are lifted out of scans into a dedicated projection node -/// (including virtual fields), using a connector table-source extension instead of a bare -/// `TableScan`, optionally with Debezium unrolling for updating sources, then remote boundary and -/// watermark. -pub struct SourceRewriter<'a> { - pub(crate) schema_provider: &'a StreamSchemaProvider, -} - -impl<'a> SourceRewriter<'a> { - pub fn new(schema_provider: &'a StreamSchemaProvider) -> Self { - Self { schema_provider } - } -} - -impl SourceRewriter<'_> { - fn projection_expr_for_column(col: &ColumnDescriptor, qualifier: &TableReference) -> Expr { - if let Some(logic) = col.computation_logic() { - logic.clone().alias_qualified( - Some(qualifier.clone()), - col.arrow_field().name().to_string(), - ) - } else { - Expr::Column(Column { - relation: Some(qualifier.clone()), - name: col.arrow_field().name().to_string(), - spans: Default::default(), - }) - } - } - - fn watermark_expression(table: &SourceTable) -> DFResult { - match table.temporal_config.watermark_strategy_column.clone() { - Some(watermark_field) => table - .schema_specs - .iter() - .find_map(|c| { - if c.arrow_field().name() == watermark_field.as_str() { - return if let Some(expr) = c.computation_logic() { - Some(expr.clone()) - } else { - Some(Expr::Column(Column { - relation: None, - name: c.arrow_field().name().to_string(), - spans: Default::default(), - })) - }; - } - None - }) - .ok_or_else(|| { - DataFusionError::Plan(format!("Watermark field {watermark_field} not found")) - }), - None => Ok(Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column { - relation: None, - name: TIMESTAMP_FIELD.to_string(), - spans: Default::default(), - })), - op: logical_expr::Operator::Minus, - right: Box::new(Expr::Literal( - ScalarValue::DurationNanosecond(Some(Duration::from_secs(1).as_nanos() as i64)), - None, - )), - })), - } - } - - fn projection_expressions( - table: &SourceTable, - qualifier: &TableReference, - projection: &Option>, - ) -> DFResult> { - let mut expressions: Vec = table - .schema_specs - .iter() - .map(|col| Self::projection_expr_for_column(col, qualifier)) - .collect(); - - if let Some(proj) = projection { - expressions = proj.iter().map(|i| expressions[*i].clone()).collect(); - } - - if let Some(event_time_field) = table.temporal_config.event_column.clone() { - let expr = table - .schema_specs - .iter() - .find_map(|c| { - if c.arrow_field().name() == event_time_field.as_str() { - return Some(Self::projection_expr_for_column(c, qualifier)); - } - None - }) - .ok_or_else(|| { - DataFusionError::Plan(format!("Event time field {event_time_field} not found")) - })?; - - expressions - .push(expr.alias_qualified(Some(qualifier.clone()), TIMESTAMP_FIELD.to_string())); - } else { - let has_ts = table - .schema_specs - .iter() - .any(|c| c.arrow_field().name() == TIMESTAMP_FIELD); - if !has_ts { - return plan_err!( - "Connector table '{}' has no `{}` column; declare WATERMARK FOR AS ... in CREATE TABLE", - table.table_identifier, - TIMESTAMP_FIELD - ); - } - expressions.push(Expr::Column(Column::new( - Some(qualifier.clone()), - TIMESTAMP_FIELD, - ))); - } - - if table.is_updating() { - expressions.push(Expr::Column(Column::new( - Some(qualifier.clone()), - UPDATING_META_FIELD, - ))); - } - - Ok(expressions) - } - - /// Connector path: `StreamIngestionNode` (table source) → optional `UnrollDebeziumPayloadNode` - /// → `Projection`, mirroring Arroyo `TableSourceExtension` + Debezium unroll + projection. - fn projection(&self, table_scan: &TableScan, table: &SourceTable) -> DFResult { - let qualifier = table_scan.table_name.clone(); - - let table_source = LogicalPlan::Extension(Extension { - node: Arc::new(StreamIngestionNode::try_new( - qualifier.clone(), - table.clone(), - )?), - }); - - let (projection_input, scan_projection) = if table.is_updating() { - if table.key_constraints.is_empty() { - return plan_err!( - "Updating connector table `{}` requires at least one PRIMARY KEY for CDC unrolling", - table.table_identifier - ); - } - let unrolled = LogicalPlan::Extension(Extension { - node: Arc::new(UnrollDebeziumPayloadNode::try_new( - table_source, - Arc::new(table.key_constraints.clone()), - )?), - }); - (unrolled, None) - } else { - (table_source, table_scan.projection.clone()) - }; - - Ok(LogicalPlan::Projection(Projection::try_new( - Self::projection_expressions(table, &qualifier, &scan_projection)?, - Arc::new(projection_input), - )?)) - } - - fn mutate_connector_table( - &self, - table_scan: &TableScan, - table: &SourceTable, - ) -> DFResult> { - let input = self.projection(table_scan, table)?; - - let schema = input.schema().clone(); - let remote = LogicalPlan::Extension(Extension { - node: Arc::new(RemoteTableBoundaryNode { - upstream_plan: input, - table_identifier: table_scan.table_name.to_owned(), - resolved_schema: schema, - requires_materialization: true, - }), - }); - - let watermark_node = EventTimeWatermarkNode::try_new( - remote, - table_scan.table_name.clone(), - Self::watermark_expression(table)?, - ) - .map_err(|err| { - DataFusionError::Internal(format!("failed to create watermark node: {err}")) - })?; - - Ok(Transformed::yes(LogicalPlan::Extension(Extension { - node: Arc::new(watermark_node), - }))) - } - - fn mutate_table_from_query( - &self, - table_scan: &TableScan, - logical_plan: &LogicalPlan, - ) -> DFResult> { - let column_expressions: Vec<_> = if let Some(projection) = &table_scan.projection { - logical_plan - .schema() - .columns() - .into_iter() - .enumerate() - .filter_map(|(i, col)| { - if projection.contains(&i) { - Some(Expr::Column(col)) - } else { - None - } - }) - .collect() - } else { - logical_plan - .schema() - .columns() - .into_iter() - .map(Expr::Column) - .collect() - }; - - let target_columns: Vec<_> = table_scan.projected_schema.columns().into_iter().collect(); - - let expressions = column_expressions - .into_iter() - .zip(target_columns) - .map(|(expr, col)| expr.alias_qualified(col.relation, col.name)) - .collect(); - - let projection = LogicalPlan::Projection(Projection::try_new_with_schema( - expressions, - Arc::new(logical_plan.clone()), - table_scan.projected_schema.clone(), - )?); - - Ok(Transformed::yes(projection)) - } -} - -impl TreeNodeRewriter for SourceRewriter<'_> { - type Node = LogicalPlan; - - fn f_up(&mut self, node: Self::Node) -> DFResult> { - let LogicalPlan::TableScan(table_scan) = node else { - return Ok(Transformed::no(node)); - }; - - let table_name = table_scan.table_name.table(); - let table = self - .schema_provider - .get_catalog_table(table_name) - .ok_or_else(|| DataFusionError::Plan(format!("Table {table_name} not found")))?; - - match table { - CatalogEntity::ExternalConnector(b) => match b.as_ref() { - ExternalTable::Source(source) => self.mutate_connector_table(&table_scan, source), - ExternalTable::Lookup(_) => { - // TODO: implement LookupSource extension - plan_err!("Lookup tables are not yet supported") - } - ExternalTable::Sink(sink) => plan_err!( - "Cannot SELECT from sink table '{}' (sinks are write-only)", - sink.name() - ), - }, - CatalogEntity::ComputedTable { - name: _, - logical_plan, - } => self.mutate_table_from_query(&table_scan, logical_plan.as_ref()), - } - } -} diff --git a/src/sql/analysis/stream_rewriter.rs b/src/sql/analysis/stream_rewriter.rs deleted file mode 100644 index a4393bd1..00000000 --- a/src/sql/analysis/stream_rewriter.rs +++ /dev/null @@ -1,234 +0,0 @@ -// 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::sync::Arc; - -use super::StreamSchemaProvider; -use crate::sql::analysis::TimeWindowNullCheckRemover; -use crate::sql::analysis::row_time_rewriter::RowTimeRewriter; -use crate::sql::analysis::{ - aggregate_rewriter::AggregateRewriter, join_rewriter::JoinRewriter, - window_function_rewriter::WindowFunctionRewriter, -}; -use crate::sql::logical_node::StreamingOperatorBlueprint; -use crate::sql::logical_node::remote_table::RemoteTableBoundaryNode; -use crate::sql::schema::utils::{add_timestamp_field, has_timestamp_field}; -use crate::sql::types::{QualifiedField, TIMESTAMP_FIELD}; -use datafusion::common::tree_node::{Transformed, TreeNodeRewriter}; -use datafusion::common::{Column, DataFusionError, Result, Spans, TableReference, plan_err}; -use datafusion::logical_expr::{ - Expr, Extension, Filter, LogicalPlan, Projection, SubqueryAlias, Union, -}; -use datafusion_common::tree_node::TreeNode; -use datafusion_expr::{Aggregate, Join}; - -pub struct StreamRewriter<'a> { - pub(crate) schema_provider: &'a StreamSchemaProvider, -} - -impl TreeNodeRewriter for StreamRewriter<'_> { - type Node = LogicalPlan; - - fn f_up(&mut self, node: Self::Node) -> Result> { - match node { - // Logic Delegation - LogicalPlan::Projection(p) => self.rewrite_projection(p), - LogicalPlan::Filter(f) => self.rewrite_filter(f), - LogicalPlan::Union(u) => self.rewrite_union(u), - - // Delegation to specialized sub-rewriters - LogicalPlan::Aggregate(agg) => self.rewrite_aggregate(agg), - LogicalPlan::Join(join) => self.rewrite_join(join), - LogicalPlan::Window(_) => self.rewrite_window(node), - LogicalPlan::SubqueryAlias(sa) => self.rewrite_subquery_alias(sa), - - // Explicitly Unsupported Operations - LogicalPlan::Sort(_) => self.unsupported_error("ORDER BY", &node), - LogicalPlan::Limit(_) => self.unsupported_error("LIMIT", &node), - LogicalPlan::Repartition(_) => self.unsupported_error("Repartitions", &node), - LogicalPlan::Explain(_) => self.unsupported_error("EXPLAIN", &node), - LogicalPlan::Analyze(_) => self.unsupported_error("ANALYZE", &node), - - _ => Ok(Transformed::no(node)), - } - } -} - -impl<'a> StreamRewriter<'a> { - pub fn new(schema_provider: &'a StreamSchemaProvider) -> Self { - Self { schema_provider } - } - - /// Delegates to AggregateRewriter to transform batch aggregates into streaming stateful operators. - fn rewrite_aggregate(&self, agg: Aggregate) -> Result> { - AggregateRewriter { - schema_provider: self.schema_provider, - } - .f_up(LogicalPlan::Aggregate(agg)) - } - - /// Delegates to JoinRewriter to handle streaming join semantics (e.g., TTL, state management). - fn rewrite_join(&self, join: Join) -> Result> { - JoinRewriter { - schema_provider: self.schema_provider, - } - .f_up(LogicalPlan::Join(join)) - } - - /// Delegates to WindowFunctionRewriter for stream-aware windowing logic. - fn rewrite_window(&self, node: LogicalPlan) -> Result> { - WindowFunctionRewriter {}.f_up(node) - } - - /// Refreshes SubqueryAlias metadata to align with potentially rewritten internal schemas. - fn rewrite_subquery_alias(&self, sa: SubqueryAlias) -> Result> { - // Since the inner 'sa.input' has been rewritten (bottom-up), we must re-create - // the alias node to ensure the outer schema correctly reflects internal changes. - let new_sa = SubqueryAlias::try_new(sa.input, sa.alias).map_err(|e| { - DataFusionError::Internal(format!("Failed to re-alias subquery: {}", e)) - })?; - - Ok(Transformed::yes(LogicalPlan::SubqueryAlias(new_sa))) - } - - /// Handles timestamp propagation and row_time() mapping for Projections - fn rewrite_projection(&self, mut projection: Projection) -> Result> { - // Check if the current projection already has a timestamp field; - // if not, we must inject it to maintain streaming heartbeats. - if !has_timestamp_field(&projection.schema) { - let input_schema = projection.input.schema(); - - // Resolve the timestamp field from the input schema using the global constant. - let timestamp_field: QualifiedField = input_schema - .qualified_field_with_unqualified_name(TIMESTAMP_FIELD) - .map_err(|_| { - DataFusionError::Plan(format!( - "No timestamp field found in projection input ({})", - projection.input.display() - )) - })? - .into(); - - // Update the logical schema to include the newly injected timestamp. - projection.schema = add_timestamp_field( - projection.schema.clone(), - timestamp_field.qualifier().cloned(), - ) - .expect("Failed to add timestamp to projection schema"); - - // Physically push the timestamp column into the expression list. - projection.expr.push(Expr::Column(Column { - relation: timestamp_field.qualifier().cloned(), - name: TIMESTAMP_FIELD.to_string(), - spans: Spans::default(), - })); - } - - // Map user-friendly row_time() function calls to internal _timestamp column references. - let rewritten = projection - .expr - .iter() - .map(|expr| expr.clone().rewrite(&mut RowTimeRewriter {})) - .collect::>>()?; - - // If any expressions were modified (e.g., row_time() was replaced), update the projection. - if rewritten.iter().any(|r| r.transformed) { - projection.expr = rewritten.into_iter().map(|r| r.data).collect(); - } - - // Return the updated plan node wrapped in a Transformed container. - Ok(Transformed::yes(LogicalPlan::Projection(projection))) - } - - /// Harmonizes schemas across Union branches and wraps them in RemoteTableBoundaryNodes. - /// - /// This ensures that all inputs to a UNION operation share the exact same schema metadata, - /// preventing "Schema Drift" where different branches have different field qualifiers. - fn rewrite_union(&self, mut union: Union) -> Result> { - // Industrial engines use the first branch as the "Master Schema" for the Union. - // We clone it once to ensure all subsequent branches are forced to comply. - let master_schema = union.inputs[0].schema().clone(); - union.schema = master_schema.clone(); - - for input in union.inputs.iter_mut() { - // Optimization: If the node is already a non-transparent Extension, - // we skip wrapping to avoid unnecessary nesting of logical nodes. - if let LogicalPlan::Extension(Extension { node }) = input.as_ref() { - let stream_ext: &dyn StreamingOperatorBlueprint = node.try_into().map_err(|e| { - DataFusionError::Internal(format!( - "Failed to resolve StreamingOperatorBlueprint: {}", - e - )) - })?; - - if !stream_ext.is_passthrough_boundary() { - continue; - } - } - - // Wrap each branch in a RemoteTableBoundaryNode. - // This acts as a logical "bridge" that forces the input to adopt the master_schema, - // effectively stripping away branch-specific qualifiers (e.g., table aliases). - let remote_ext = Arc::new(RemoteTableBoundaryNode { - upstream_plan: input.as_ref().clone(), - table_identifier: TableReference::bare("union_input"), - resolved_schema: master_schema.clone(), - requires_materialization: false, // Internal logical boundary only; does not require physical sink. - }); - - // Atomically replace the input with the wrapped version. - *input = Arc::new(LogicalPlan::Extension(Extension { node: remote_ext })); - } - - Ok(Transformed::yes(LogicalPlan::Union(union))) - } - - /// Optimizes Filter nodes by stripping redundant NULL checks on time window expressions. - /// - /// In streaming SQL, DataFusion often injects 'IS NOT NULL' guards for window functions - /// that are redundant or can interfere with watermark propagation. This rewriter - /// cleans those predicates to simplify the physical execution plan. - fn rewrite_filter(&self, filter: Filter) -> Result> { - // We attempt to rewrite the predicate using a specialized sub-rewriter. - // The TimeWindowNullCheckRemover specifically targets expressions like - // `tumble(...) IS NOT NULL` and simplifies them to `TRUE`. - let rewritten_expr = filter - .predicate - .clone() - .rewrite(&mut TimeWindowNullCheckRemover {})?; - - if !rewritten_expr.transformed { - return Ok(Transformed::no(LogicalPlan::Filter(filter))); - } - - // Industrial Guard: Re-validate the predicate against the input schema. - // 'Filter::try_new' ensures that the transformed expression is still semantically - // valid for the underlying data stream. - let new_filter = Filter::try_new(rewritten_expr.data, filter.input).map_err(|e| { - DataFusionError::Internal(format!( - "Failed to re-validate filtered predicate after NULL-check removal: {}", - e - )) - })?; - - Ok(Transformed::yes(LogicalPlan::Filter(new_filter))) - } - - /// Centralized error handler for unsupported streaming operations - fn unsupported_error(&self, op: &str, node: &LogicalPlan) -> Result> { - plan_err!( - "{} is not currently supported in streaming SQL ({})", - op, - node.display() - ) - } -} diff --git a/src/sql/analysis/streaming_window_analzer.rs b/src/sql/analysis/streaming_window_analzer.rs deleted file mode 100644 index b8a7f78f..00000000 --- a/src/sql/analysis/streaming_window_analzer.rs +++ /dev/null @@ -1,219 +0,0 @@ -// 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 std::sync::Arc; - -use datafusion::common::tree_node::{TreeNodeRecursion, TreeNodeVisitor}; -use datafusion::common::{Column, DFSchema, DataFusionError, Result}; -use datafusion::logical_expr::{Expr, Extension, LogicalPlan, expr::Alias}; - -use crate::sql::logical_node::aggregate::{STREAM_AGG_EXTENSION_NAME, StreamWindowAggregateNode}; -use crate::sql::logical_node::join::STREAM_JOIN_NODE_TYPE; -use crate::sql::types::{ - QualifiedField, WindowBehavior, WindowType, extract_qualified_fields, extract_window_type, -}; - -/// WindowDetectingVisitor identifies windowing strategies and tracks window-carrying fields -/// as they propagate upward through the logical plan tree. -#[derive(Debug, Default)] -pub(crate) struct StreamingWindowAnalzer { - /// The specific window type discovered (Tumble, Hop, etc.) - pub(crate) window: Option, - /// Set of fields in the current plan node that carry window semantics. - pub(crate) fields: HashSet, -} - -impl StreamingWindowAnalzer { - /// Entry point to resolve the WindowType of a given plan branch. - pub(crate) fn get_window(logical_plan: &LogicalPlan) -> Result> { - let mut visitor = Self::default(); - logical_plan.visit_with_subqueries(&mut visitor)?; - Ok(visitor.window) - } - - /// Resolves whether an expression is a reference to an existing window field - /// or a definition of a new window function. - fn resolve_window_from_expr( - &self, - expr: &Expr, - input_schema: &DFSchema, - ) -> Result> { - // 1. Check if the expression directly references a known window field. - if let Some(col) = extract_column(expr) { - let field = input_schema.field_with_name(col.relation.as_ref(), &col.name)?; - let df_field: QualifiedField = (col.relation.clone(), Arc::new(field.clone())).into(); - - if self.fields.contains(&df_field) { - return Ok(self.window.clone()); - } - } - - // 2. Otherwise, check if it's a new window function call (e.g., tumble(), hop()). - extract_window_type(expr) - } - - /// Updates the internal state with new window findings and maps them to the output schema. - fn update_state( - &mut self, - matched_windows: Vec<(usize, WindowType)>, - schema: &DFSchema, - ) -> Result<()> { - // Clear fields from the previous level to maintain schema strictly for the current node. - self.fields.clear(); - - for (index, window) in matched_windows { - if let Some(existing) = &self.window { - if existing != &window { - return Err(DataFusionError::Plan(format!( - "Conflicting windows in the same operator: expected {:?}, found {:?}", - existing, window - ))); - } - } else { - self.window = Some(window); - } - // Record this specific index in the schema as a window carrier. - self.fields.insert(schema.qualified_field(index).into()); - } - Ok(()) - } -} - -pub(crate) fn extract_column(expr: &Expr) -> Option<&Column> { - match expr { - Expr::Column(column) => Some(column), - Expr::Alias(Alias { expr, .. }) => extract_column(expr), - _ => None, - } -} - -impl TreeNodeVisitor<'_> for StreamingWindowAnalzer { - type Node = LogicalPlan; - - fn f_down(&mut self, node: &Self::Node) -> Result { - // Joins require cross-branch validation to ensure left and right sides align on time. - if let LogicalPlan::Extension(Extension { node }) = node - && node.name() == STREAM_JOIN_NODE_TYPE - { - let mut branch_windows = HashSet::new(); - for input in node.inputs() { - if let Some(w) = Self::get_window(input)? { - branch_windows.insert(w); - } - } - - if branch_windows.len() > 1 { - return Err(DataFusionError::Plan( - "Join inputs have mismatched windowing strategies.".into(), - )); - } - self.window = branch_windows.into_iter().next(); - - // Optimization: No need to recurse manually if we've resolved the join boundary. - return Ok(TreeNodeRecursion::Jump); - } - Ok(TreeNodeRecursion::Continue) - } - - fn f_up(&mut self, node: &Self::Node) -> Result { - match node { - LogicalPlan::Projection(p) => { - let windows = p - .expr - .iter() - .enumerate() - .filter_map(|(i, e)| { - self.resolve_window_from_expr(e, p.input.schema()) - .transpose() - .map(|res| res.map(|w| (i, w))) - }) - .collect::>>()?; - - self.update_state(windows, &p.schema)?; - } - - LogicalPlan::Aggregate(agg) => { - let windows = agg - .group_expr - .iter() - .enumerate() - .filter_map(|(i, e)| { - self.resolve_window_from_expr(e, agg.input.schema()) - .transpose() - .map(|res| res.map(|w| (i, w))) - }) - .collect::>>()?; - - self.update_state(windows, &agg.schema)?; - } - - LogicalPlan::SubqueryAlias(sa) => { - // Map fields through the alias layer by resolving column indices. - let input_schema = sa.input.schema(); - let mapped = self - .fields - .drain() - .map(|f| { - let idx = input_schema.index_of_column(&f.qualified_column())?; - Ok(sa.schema.qualified_field(idx).into()) - }) - .collect::>>()?; - - self.fields = mapped; - } - - LogicalPlan::Extension(Extension { node }) - if node.name() == STREAM_AGG_EXTENSION_NAME => - { - let ext = node - .as_any() - .downcast_ref::() - .ok_or_else(|| { - DataFusionError::Internal("StreamWindowAggregateNode is malformed".into()) - })?; - - match &ext.window_spec { - WindowBehavior::FromOperator { - window, - window_field, - is_nested, - .. - } => { - if self.window.is_some() && !*is_nested { - return Err(DataFusionError::Plan( - "Redundant window definition on an already windowed stream.".into(), - )); - } - self.window = Some(window.clone()); - self.fields.insert(window_field.clone()); - } - WindowBehavior::InData => { - let current_schema_fields: HashSet<_> = - extract_qualified_fields(node.schema()) - .into_iter() - .collect(); - self.fields.retain(|f| current_schema_fields.contains(f)); - - if self.fields.is_empty() { - return Err(DataFusionError::Plan( - "Windowed aggregate missing window metadata from its input.".into(), - )); - } - } - } - } - _ => {} - } - Ok(TreeNodeRecursion::Continue) - } -} diff --git a/src/sql/analysis/time_window.rs b/src/sql/analysis/time_window.rs deleted file mode 100644 index 104c0cca..00000000 --- a/src/sql/analysis/time_window.rs +++ /dev/null @@ -1,83 +0,0 @@ -// 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 datafusion::common::tree_node::{ - Transformed, TreeNodeRecursion, TreeNodeRewriter, TreeNodeVisitor, -}; -use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue, plan_err}; -use datafusion::logical_expr::expr::ScalarFunction; -use datafusion::logical_expr::{Expr, LogicalPlan}; - -/// Returns the time window function name if the expression is one (tumble/hop/session). -pub fn is_time_window(expr: &Expr) -> Option<&str> { - if let Expr::ScalarFunction(ScalarFunction { func, args: _ }) = expr { - match func.name() { - "tumble" | "hop" | "session" => return Some(func.name()), - _ => {} - } - } - None -} - -struct TimeWindowExprChecker {} - -impl TreeNodeVisitor<'_> for TimeWindowExprChecker { - type Node = Expr; - - fn f_down(&mut self, node: &Self::Node) -> DFResult { - if let Some(w) = is_time_window(node) { - return plan_err!( - "time window function {} is not allowed in this context. \ - Are you missing a GROUP BY clause?", - w - ); - } - Ok(TreeNodeRecursion::Continue) - } -} - -/// Visitor that checks an entire LogicalPlan for misplaced time window UDFs. -pub struct TimeWindowUdfChecker {} - -impl TreeNodeVisitor<'_> for TimeWindowUdfChecker { - type Node = LogicalPlan; - - fn f_down(&mut self, node: &Self::Node) -> DFResult { - use datafusion::common::tree_node::TreeNode; - node.expressions().iter().try_for_each(|expr| { - let mut checker = TimeWindowExprChecker {}; - expr.visit(&mut checker)?; - Ok::<(), DataFusionError>(()) - })?; - Ok(TreeNodeRecursion::Continue) - } -} - -/// Removes `IS NOT NULL` checks wrapping time window functions, -/// replacing them with `true` since time windows are never null. -pub struct TimeWindowNullCheckRemover {} - -impl TreeNodeRewriter for TimeWindowNullCheckRemover { - type Node = Expr; - - fn f_down(&mut self, node: Self::Node) -> DFResult> { - if let Expr::IsNotNull(expr) = &node - && is_time_window(expr).is_some() - { - return Ok(Transformed::yes(Expr::Literal( - ScalarValue::Boolean(Some(true)), - None, - ))); - } - Ok(Transformed::no(node)) - } -} diff --git a/src/sql/analysis/udafs.rs b/src/sql/analysis/udafs.rs deleted file mode 100644 index 73fc062c..00000000 --- a/src/sql/analysis/udafs.rs +++ /dev/null @@ -1,43 +0,0 @@ -// 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 datafusion::arrow::array::ArrayRef; -use datafusion::error::Result; -use datafusion::physical_plan::Accumulator; -use datafusion::scalar::ScalarValue; -use std::fmt::Debug; - -/// Fake UDAF used just for plan-time placeholder. -#[derive(Debug)] -pub struct EmptyUdaf {} - -impl Accumulator for EmptyUdaf { - fn update_batch(&mut self, _: &[ArrayRef]) -> Result<()> { - unreachable!() - } - - fn evaluate(&self) -> Result { - unreachable!() - } - - fn size(&self) -> usize { - unreachable!() - } - - fn state(&self) -> Result> { - unreachable!() - } - - fn merge_batch(&mut self, _: &[ArrayRef]) -> Result<()> { - unreachable!() - } -} diff --git a/src/sql/analysis/unnest_rewriter.rs b/src/sql/analysis/unnest_rewriter.rs deleted file mode 100644 index 147b1f49..00000000 --- a/src/sql/analysis/unnest_rewriter.rs +++ /dev/null @@ -1,179 +0,0 @@ -// 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::sync::Arc; - -use datafusion::arrow::datatypes::DataType; -use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRewriter}; -use datafusion::common::{Column, Result as DFResult, plan_err}; -use datafusion::logical_expr::expr::ScalarFunction; -use datafusion::logical_expr::{ColumnUnnestList, Expr, LogicalPlan, Projection, Unnest}; - -use crate::sql::common::constants::planning_placeholder_udf; -use crate::sql::types::{QualifiedField, build_df_schema, extract_qualified_fields}; - -pub const UNNESTED_COL: &str = "__unnested"; - -/// Rewrites projections containing `unnest()` calls into proper Unnest logical plans. -pub struct UnnestRewriter {} - -impl UnnestRewriter { - fn split_unnest(expr: Expr) -> DFResult<(Expr, Option)> { - let mut captured: Option = None; - - let expr = expr.transform_up(|e| { - if let Expr::ScalarFunction(ScalarFunction { func: udf, args }) = &e - && udf.name() == planning_placeholder_udf::UNNEST - { - match args.len() { - 1 => { - if captured.replace(args[0].clone()).is_some() { - return plan_err!( - "Multiple unnests in expression, which is not allowed" - ); - } - return Ok(Transformed::yes(Expr::Column(Column::new_unqualified( - UNNESTED_COL, - )))); - } - n => { - panic!("Unnest has wrong number of arguments (expected 1, found {n})"); - } - } - } - Ok(Transformed::no(e)) - })?; - - Ok((expr.data, captured)) - } -} - -impl TreeNodeRewriter for UnnestRewriter { - type Node = LogicalPlan; - - fn f_up(&mut self, node: Self::Node) -> DFResult> { - let LogicalPlan::Projection(projection) = &node else { - if node.expressions().iter().any(|e| { - let e = Self::split_unnest(e.clone()); - e.is_err() || e.unwrap().1.is_some() - }) { - return plan_err!("unnest is only supported in SELECT statements"); - } - return Ok(Transformed::no(node)); - }; - - let mut unnest = None; - let exprs = projection - .expr - .clone() - .into_iter() - .enumerate() - .map(|(i, expr)| { - let (expr, opt) = Self::split_unnest(expr)?; - let is_unnest = if let Some(e) = opt { - if let Some(prev) = unnest.replace((e, i)) - && &prev != unnest.as_ref().unwrap() - { - return plan_err!( - "Projection contains multiple unnests, which is not currently supported" - ); - } - true - } else { - false - }; - - Ok((expr, is_unnest)) - }) - .collect::>>()?; - - if let Some((unnest_inner, unnest_idx)) = unnest { - let produce_list = Arc::new(LogicalPlan::Projection(Projection::try_new( - exprs - .iter() - .cloned() - .map(|(e, is_unnest)| { - if is_unnest { - unnest_inner.clone().alias(UNNESTED_COL) - } else { - e - } - }) - .collect(), - projection.input.clone(), - )?)); - - let unnest_fields = extract_qualified_fields(produce_list.schema()) - .iter() - .enumerate() - .map(|(i, f)| { - if i == unnest_idx { - let DataType::List(inner) = f.data_type() else { - return plan_err!( - "Argument '{}' to unnest is not a List", - f.qualified_name() - ); - }; - Ok(QualifiedField::new_unqualified( - UNNESTED_COL, - inner.data_type().clone(), - inner.is_nullable(), - )) - } else { - Ok((*f).clone()) - } - }) - .collect::>>()?; - - let unnest_node = LogicalPlan::Unnest(Unnest { - exec_columns: vec![ - QualifiedField::from(produce_list.schema().qualified_field(unnest_idx)) - .qualified_column(), - ], - input: produce_list, - list_type_columns: vec![( - unnest_idx, - ColumnUnnestList { - output_column: Column::new_unqualified(UNNESTED_COL), - depth: 1, - }, - )], - struct_type_columns: vec![], - dependency_indices: vec![], - schema: Arc::new(build_df_schema(&unnest_fields)?), - options: Default::default(), - }); - - let output_node = LogicalPlan::Projection(Projection::try_new( - exprs - .iter() - .enumerate() - .map(|(i, (expr, has_unnest))| { - if *has_unnest { - expr.clone() - } else { - Expr::Column( - QualifiedField::from(unnest_node.schema().qualified_field(i)) - .qualified_column(), - ) - } - }) - .collect(), - Arc::new(unnest_node), - )?); - - Ok(Transformed::yes(output_node)) - } else { - Ok(Transformed::no(LogicalPlan::Projection(projection.clone()))) - } - } -} diff --git a/src/sql/analysis/window_function_rewriter.rs b/src/sql/analysis/window_function_rewriter.rs deleted file mode 100644 index c1e3396d..00000000 --- a/src/sql/analysis/window_function_rewriter.rs +++ /dev/null @@ -1,204 +0,0 @@ -// 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 datafusion::common::tree_node::Transformed; -use datafusion::common::{Result as DFResult, plan_err, tree_node::TreeNodeRewriter}; -use datafusion::logical_expr::{ - self, Expr, Extension, LogicalPlan, Projection, Sort, Window, expr::WindowFunction, - expr::WindowFunctionParams, -}; -use datafusion_common::DataFusionError; -use std::sync::Arc; -use tracing::debug; - -use crate::sql::analysis::streaming_window_analzer::{StreamingWindowAnalzer, extract_column}; -use crate::sql::logical_node::key_calculation::{KeyExtractionNode, KeyExtractionStrategy}; -use crate::sql::logical_node::windows_function::StreamingWindowFunctionNode; -use crate::sql::types::{WindowType, build_df_schema, extract_qualified_fields}; - -/// WindowFunctionRewriter transforms standard SQL Window functions into streaming-compatible -/// stateful operators, ensuring proper data partitioning and sorting for distributed execution. -pub(crate) struct WindowFunctionRewriter; - -impl WindowFunctionRewriter { - /// Recursively unwraps Aliases to find the underlying WindowFunction. - #[allow(clippy::only_used_in_recursion)] - fn resolve_window_function(&self, expr: &Expr) -> DFResult<(WindowFunction, String)> { - match expr { - Expr::Alias(alias) => { - let (func, _) = self.resolve_window_function(&alias.expr)?; - Ok((func, alias.name.clone())) - } - Expr::WindowFunction(wf) => Ok((wf.as_ref().clone(), expr.name_for_alias()?)), - _ => plan_err!("Expected WindowFunction or Alias, found: {:?}", expr), - } - } - - /// Identifies which field in the PARTITION BY clause corresponds to the streaming window. - fn identify_window_partition( - &self, - params: &WindowFunctionParams, - input: &LogicalPlan, - input_window_fields: &std::collections::HashSet, - ) -> DFResult { - let matched: Vec<_> = params - .partition_by - .iter() - .enumerate() - .filter_map(|(i, e)| { - let col = extract_column(e)?; - let field = input - .schema() - .field_with_name(col.relation.as_ref(), &col.name) - .ok()?; - let df_field = (col.relation.clone(), Arc::new(field.clone())).into(); - - if input_window_fields.contains(&df_field) { - Some(i) - } else { - None - } - }) - .collect(); - - if matched.len() != 1 { - return plan_err!( - "Streaming window functions require exactly one window column in PARTITION BY. Found: {}", - matched.len() - ); - } - Ok(matched[0]) - } - - /// Wraps the input in a Projection and KeyExtractionNode to handle data distribution. - fn build_keyed_input( - &self, - input: Arc, - partition_keys: &[Expr], - ) -> DFResult { - let key_count = partition_keys.len(); - - // 1. Build projection: [_key_0, _key_1, ..., original_columns] - let mut exprs: Vec<_> = partition_keys - .iter() - .enumerate() - .map(|(i, e)| e.clone().alias(format!("_key_{i}"))) - .collect(); - - exprs.extend( - extract_qualified_fields(input.schema()) - .iter() - .map(|f| Expr::Column(f.qualified_column())), - ); - - // 2. Derive the keyed schema - let mut keyed_fields = - extract_qualified_fields(&Projection::try_new(exprs.clone(), input.clone())?.schema) - .iter() - .take(key_count) - .cloned() - .collect::>(); - keyed_fields.extend(extract_qualified_fields(input.schema())); - - let keyed_schema = Arc::new(build_df_schema(&keyed_fields)?); - - let projection = - LogicalPlan::Projection(Projection::try_new_with_schema(exprs, input, keyed_schema)?); - - // 3. Wrap in KeyExtractionNode for the physical planner - Ok(LogicalPlan::Extension(Extension { - node: Arc::new(KeyExtractionNode::new( - projection, - KeyExtractionStrategy::ColumnIndices((0..key_count).collect()), - )), - })) - } -} - -impl TreeNodeRewriter for WindowFunctionRewriter { - type Node = LogicalPlan; - - fn f_up(&mut self, node: Self::Node) -> DFResult> { - let LogicalPlan::Window(window) = node else { - return Ok(Transformed::no(node)); - }; - - debug!("Rewriting window function for streaming: {:?}", window); - - // 1. Analyze input windowing context - let mut analyzer = StreamingWindowAnalzer::default(); - window.input.visit_with_subqueries(&mut analyzer)?; - - let input_window = analyzer.window.ok_or_else(|| { - DataFusionError::Plan( - "Window functions require a windowed input stream (e.g., TUMBLE/HOP)".into(), - ) - })?; - - if matches!(input_window, WindowType::Session { .. }) { - return plan_err!( - "Streaming window functions (OVER) are not supported on Session windows." - ); - } - - // 2. Validate window expression constraints - if window.window_expr.len() != 1 { - return plan_err!( - "Arroyo currently supports exactly one window expression per OVER clause." - ); - } - - let (mut wf, original_name) = self.resolve_window_function(&window.window_expr[0])?; - - // 3. Identify and extract the window column from PARTITION BY - let window_part_idx = - self.identify_window_partition(&wf.params, &window.input, &analyzer.fields)?; - let mut partition_keys = wf.params.partition_by.clone(); - partition_keys.remove(window_part_idx); - - // Update function params to exclude the window column from internal partitioning - // as the streaming engine handles window boundaries natively. - wf.params.partition_by = partition_keys.clone(); - let key_count = partition_keys.len(); - - // 4. Build the data-shuffling pipeline (Projection -> KeyCalc -> Sort) - let keyed_plan = self.build_keyed_input(window.input.clone(), &partition_keys)?; - - let mut sort_exprs: Vec<_> = partition_keys - .iter() - .map(|e| logical_expr::expr::Sort { - expr: e.clone(), - asc: true, - nulls_first: false, - }) - .collect(); - sort_exprs.extend(wf.params.order_by.clone()); - - let sorted_plan = LogicalPlan::Sort(Sort { - expr: sort_exprs, - input: Arc::new(keyed_plan), - fetch: None, - }); - - // 5. Final Assembly - let final_wf_expr = Expr::WindowFunction(Box::new(wf)).alias_if_changed(original_name)?; - let rewritten_window = - LogicalPlan::Window(Window::try_new(vec![final_wf_expr], Arc::new(sorted_plan))?); - - Ok(Transformed::yes(LogicalPlan::Extension(Extension { - node: Arc::new(StreamingWindowFunctionNode::new( - rewritten_window, - (0..key_count).collect(), - )), - }))) - } -} diff --git a/src/sql/api/checkpoints.rs b/src/sql/api/checkpoints.rs deleted file mode 100644 index d9bdc139..00000000 --- a/src/sql/api/checkpoints.rs +++ /dev/null @@ -1,108 +0,0 @@ -// 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::sql::common::to_micros; -use serde::{Deserialize, Serialize}; -use std::time::SystemTime; - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct Checkpoint { - pub epoch: u32, - pub backend: String, - pub start_time: u64, - pub finish_time: Option, - pub events: Vec, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct CheckpointEventSpan { - pub start_time: u64, - pub finish_time: u64, - pub event: String, - pub description: String, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct SubtaskCheckpointGroup { - pub index: u32, - pub bytes: u64, - pub event_spans: Vec, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct OperatorCheckpointGroup { - pub operator_id: String, - pub bytes: u64, - pub started_metadata_write: Option, - pub finish_time: Option, - pub subtasks: Vec, -} - -#[derive(Debug, Copy, Clone, Eq, PartialEq, Serialize, Deserialize)] -pub enum JobCheckpointEventType { - Checkpointing, - CheckpointingOperators, - WritingMetadata, - Compacting, - Committing, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct JobCheckpointSpan { - pub event: JobCheckpointEventType, - pub start_time: u64, - pub finish_time: Option, -} - -impl JobCheckpointSpan { - pub fn now(event: JobCheckpointEventType) -> Self { - Self { - event, - start_time: to_micros(SystemTime::now()), - finish_time: None, - } - } - - pub fn finish(&mut self) { - if self.finish_time.is_none() { - self.finish_time = Some(to_micros(SystemTime::now())); - } - } -} - -impl From for CheckpointEventSpan { - fn from(value: JobCheckpointSpan) -> Self { - let description = match value.event { - JobCheckpointEventType::Checkpointing => "The entire checkpointing process", - JobCheckpointEventType::CheckpointingOperators => { - "The time spent checkpointing operator states" - } - JobCheckpointEventType::WritingMetadata => "Writing the final checkpoint metadata", - JobCheckpointEventType::Compacting => "Compacting old checkpoints", - JobCheckpointEventType::Committing => { - "Running two-phase commit for transactional connectors" - } - } - .to_string(); - - Self { - start_time: value.start_time, - finish_time: value.finish_time.unwrap_or_default(), - event: format!("{:?}", value.event), - description, - } - } -} diff --git a/src/sql/api/connections.rs b/src/sql/api/connections.rs deleted file mode 100644 index 3c5caf76..00000000 --- a/src/sql/api/connections.rs +++ /dev/null @@ -1,620 +0,0 @@ -// 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::sql::common::formats::{BadData, Format, Framing}; -use crate::sql::common::{FsExtensionType, FsSchema}; -use datafusion::arrow::datatypes::{DataType, Field, Fields, TimeUnit}; -use serde::ser::SerializeMap; -use serde::{Deserialize, Serialize, Serializer}; -use std::collections::{BTreeMap, HashMap, HashSet}; -use std::fmt::{Display, Formatter}; -use std::sync::Arc; - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct Connector { - pub id: String, - pub name: String, - pub icon: String, - pub description: String, - pub table_config: String, - pub enabled: bool, - pub source: bool, - pub sink: bool, - pub custom_schemas: bool, - pub testing: bool, - pub hidden: bool, - pub connection_config: Option, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct ConnectionProfile { - pub id: String, - pub name: String, - pub connector: String, - pub config: serde_json::Value, - pub description: String, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct ConnectionProfilePost { - pub name: String, - pub connector: String, - pub config: serde_json::Value, -} - -#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)] -#[serde(rename_all = "snake_case")] -pub enum ConnectionType { - Source, - Sink, - Lookup, -} - -impl Display for ConnectionType { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - match self { - ConnectionType::Source => write!(f, "SOURCE"), - ConnectionType::Sink => write!(f, "SINK"), - ConnectionType::Lookup => write!(f, "LOOKUP"), - } - } -} - -impl TryFrom for ConnectionType { - type Error = String; - - fn try_from(value: String) -> Result { - match value.to_lowercase().as_str() { - "source" => Ok(ConnectionType::Source), - "sink" => Ok(ConnectionType::Sink), - "lookup" => Ok(ConnectionType::Lookup), - _ => Err(format!("Invalid connection type: {value}")), - } - } -} - -// ─────────────────── Field Types ─────────────────── - -#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] -#[serde(tag = "type", rename_all = "snake_case")] -pub enum FieldType { - Int32, - Int64, - Uint32, - Uint64, - #[serde(alias = "f32")] - Float32, - #[serde(alias = "f64")] - Float64, - Decimal128(DecimalField), - Bool, - #[serde(alias = "utf8")] - String, - #[serde(alias = "binary")] - Bytes, - Timestamp(TimestampField), - Json, - Struct(StructField), - List(ListField), -} - -impl FieldType { - pub fn sql_type(&self) -> String { - match self { - FieldType::Int32 => "INTEGER".into(), - FieldType::Int64 => "BIGINT".into(), - FieldType::Uint32 => "INTEGER UNSIGNED".into(), - FieldType::Uint64 => "BIGINT UNSIGNED".into(), - FieldType::Float32 => "FLOAT".into(), - FieldType::Float64 => "DOUBLE".into(), - FieldType::Decimal128(f) => format!("DECIMAL({}, {})", f.precision, f.scale), - FieldType::Bool => "BOOLEAN".into(), - FieldType::String => "TEXT".into(), - FieldType::Bytes => "BINARY".into(), - FieldType::Timestamp(t) => format!("TIMESTAMP({})", t.unit.precision()), - FieldType::Json => "JSON".into(), - FieldType::List(item) => format!("{}[]", item.items.field_type.sql_type()), - FieldType::Struct(StructField { fields, .. }) => { - format!( - "STRUCT <{}>", - fields - .iter() - .map(|f| format!("{} {}", f.name, f.field_type.sql_type())) - .collect::>() - .join(", ") - ) - } - } - } -} - -#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, Default)] -#[serde(rename_all = "snake_case")] -pub enum TimestampUnit { - #[serde(alias = "s")] - Second, - #[default] - #[serde(alias = "ms")] - Millisecond, - #[serde(alias = "µs", alias = "us")] - Microsecond, - #[serde(alias = "ns")] - Nanosecond, -} - -impl TimestampUnit { - pub fn precision(&self) -> u8 { - match self { - TimestampUnit::Second => 0, - TimestampUnit::Millisecond => 3, - TimestampUnit::Microsecond => 6, - TimestampUnit::Nanosecond => 9, - } - } -} - -#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] -#[serde(rename_all = "snake_case")] -pub struct TimestampField { - #[serde(default)] - pub unit: TimestampUnit, -} - -#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] -#[serde(rename_all = "snake_case")] -pub struct DecimalField { - pub precision: u8, - pub scale: i8, -} - -#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] -#[serde(rename_all = "snake_case")] -pub struct StructField { - pub fields: Vec, -} - -#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] -#[serde(rename_all = "snake_case")] -pub struct ListField { - pub items: Box, -} - -fn default_item_name() -> String { - "item".to_string() -} - -#[derive(Deserialize, Clone, Debug, PartialEq, Eq)] -#[serde(rename_all = "snake_case")] -pub struct ListFieldItem { - #[serde(default = "default_item_name")] - pub name: String, - #[serde(flatten)] - pub field_type: FieldType, - #[serde(default)] - pub required: bool, - #[serde(default)] - pub sql_name: Option, -} - -impl From for Field { - fn from(value: ListFieldItem) -> Self { - SourceField { - name: value.name, - field_type: value.field_type, - required: value.required, - sql_name: None, - metadata_key: None, - } - .into() - } -} - -impl Serialize for ListFieldItem { - fn serialize(&self, s: S) -> Result - where - S: Serializer, - { - let mut f = Serializer::serialize_map(s, None)?; - f.serialize_entry("name", &self.name)?; - serialize_field_type_flat(&self.field_type, &mut f)?; - f.serialize_entry("required", &self.required)?; - f.serialize_entry("sql_name", &self.field_type.sql_type())?; - f.end() - } -} - -impl TryFrom for ListFieldItem { - type Error = String; - - fn try_from(value: Field) -> Result { - let source_field: SourceField = value.try_into()?; - Ok(Self { - name: source_field.name, - field_type: source_field.field_type, - required: source_field.required, - sql_name: None, - }) - } -} - -fn serialize_field_type_flat(ft: &FieldType, map: &mut M) -> Result<(), M::Error> { - let type_tag = match ft { - FieldType::Int32 => "int32", - FieldType::Int64 => "int64", - FieldType::Uint32 => "uint32", - FieldType::Uint64 => "uint64", - FieldType::Float32 => "float32", - FieldType::Float64 => "float64", - FieldType::Decimal128(_) => "decimal128", - FieldType::Bool => "bool", - FieldType::String => "string", - FieldType::Bytes => "bytes", - FieldType::Timestamp(_) => "timestamp", - FieldType::Json => "json", - FieldType::Struct(_) => "struct", - FieldType::List(_) => "list", - }; - map.serialize_entry("type", type_tag)?; - - match ft { - FieldType::Decimal128(d) => { - map.serialize_entry("precision", &d.precision)?; - map.serialize_entry("scale", &d.scale)?; - } - FieldType::Timestamp(t) => { - map.serialize_entry("unit", &t.unit)?; - } - FieldType::Struct(s) => { - map.serialize_entry("fields", &s.fields)?; - } - FieldType::List(l) => { - map.serialize_entry("items", &l.items)?; - } - _ => {} - } - Ok(()) -} - -// ─────────────────── Source Field ─────────────────── - -#[derive(Deserialize, Clone, Debug, PartialEq, Eq)] -#[serde(rename_all = "snake_case")] -pub struct SourceField { - pub name: String, - #[serde(flatten)] - pub field_type: FieldType, - #[serde(default)] - pub required: bool, - #[serde(default)] - pub sql_name: Option, - #[serde(default)] - pub metadata_key: Option, -} - -impl Serialize for SourceField { - fn serialize(&self, s: S) -> Result - where - S: Serializer, - { - let mut f = Serializer::serialize_map(s, None)?; - f.serialize_entry("name", &self.name)?; - serialize_field_type_flat(&self.field_type, &mut f)?; - f.serialize_entry("required", &self.required)?; - if let Some(metadata_key) = &self.metadata_key { - f.serialize_entry("metadata_key", metadata_key)?; - } - f.serialize_entry("sql_name", &self.field_type.sql_type())?; - f.end() - } -} - -impl From for Field { - fn from(f: SourceField) -> Self { - let (t, ext) = match f.field_type { - FieldType::Int32 => (DataType::Int32, None), - FieldType::Int64 => (DataType::Int64, None), - FieldType::Uint32 => (DataType::UInt32, None), - FieldType::Uint64 => (DataType::UInt64, None), - FieldType::Float32 => (DataType::Float32, None), - FieldType::Float64 => (DataType::Float64, None), - FieldType::Bool => (DataType::Boolean, None), - FieldType::String => (DataType::Utf8, None), - FieldType::Bytes => (DataType::Binary, None), - FieldType::Decimal128(d) => (DataType::Decimal128(d.precision, d.scale), None), - FieldType::Timestamp(TimestampField { - unit: TimestampUnit::Second, - }) => (DataType::Timestamp(TimeUnit::Second, None), None), - FieldType::Timestamp(TimestampField { - unit: TimestampUnit::Millisecond, - }) => (DataType::Timestamp(TimeUnit::Millisecond, None), None), - FieldType::Timestamp(TimestampField { - unit: TimestampUnit::Microsecond, - }) => (DataType::Timestamp(TimeUnit::Microsecond, None), None), - FieldType::Timestamp(TimestampField { - unit: TimestampUnit::Nanosecond, - }) => (DataType::Timestamp(TimeUnit::Nanosecond, None), None), - FieldType::Json => (DataType::Utf8, Some(FsExtensionType::JSON)), - FieldType::Struct(s) => ( - DataType::Struct(Fields::from( - s.fields - .into_iter() - .map(|t| t.into()) - .collect::>(), - )), - None, - ), - FieldType::List(t) => (DataType::List(Arc::new((*t.items).into())), None), - }; - - FsExtensionType::add_metadata(ext, Field::new(f.name, t, !f.required)) - } -} - -impl TryFrom for SourceField { - type Error = String; - - fn try_from(f: Field) -> Result { - let field_type = match (f.data_type(), FsExtensionType::from_map(f.metadata())) { - (DataType::Boolean, None) => FieldType::Bool, - (DataType::Int32, None) => FieldType::Int32, - (DataType::Int64, None) => FieldType::Int64, - (DataType::UInt32, None) => FieldType::Uint32, - (DataType::UInt64, None) => FieldType::Uint64, - (DataType::Float32, None) => FieldType::Float32, - (DataType::Float64, None) => FieldType::Float64, - (DataType::Decimal128(p, s), None) => FieldType::Decimal128(DecimalField { - precision: *p, - scale: *s, - }), - (DataType::Binary | DataType::LargeBinary | DataType::BinaryView, None) => { - FieldType::Bytes - } - (DataType::Timestamp(TimeUnit::Second, _), None) => { - FieldType::Timestamp(TimestampField { - unit: TimestampUnit::Second, - }) - } - (DataType::Timestamp(TimeUnit::Millisecond, _), None) => { - FieldType::Timestamp(TimestampField { - unit: TimestampUnit::Millisecond, - }) - } - (DataType::Timestamp(TimeUnit::Microsecond, _), None) => { - FieldType::Timestamp(TimestampField { - unit: TimestampUnit::Microsecond, - }) - } - (DataType::Timestamp(TimeUnit::Nanosecond, _), None) => { - FieldType::Timestamp(TimestampField { - unit: TimestampUnit::Nanosecond, - }) - } - (DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View, None) => FieldType::String, - ( - DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View, - Some(FsExtensionType::JSON), - ) => FieldType::Json, - (DataType::Struct(fields), None) => { - let fields: Result<_, String> = fields - .into_iter() - .map(|f| (**f).clone().try_into()) - .collect(); - FieldType::Struct(StructField { fields: fields? }) - } - (DataType::List(item), None) => FieldType::List(ListField { - items: Box::new((**item).clone().try_into()?), - }), - dt => return Err(format!("Unsupported data type {dt:?}")), - }; - - Ok(SourceField { - name: f.name().clone(), - field_type, - required: !f.is_nullable(), - sql_name: None, - metadata_key: None, - }) - } -} - -// ─────────────────── Schema Definitions ─────────────────── - -#[allow(clippy::enum_variant_names)] -#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)] -#[serde(rename_all = "snake_case", tag = "type")] -pub enum SchemaDefinition { - JsonSchema { - schema: String, - }, - ProtobufSchema { - schema: String, - #[serde(default)] - dependencies: HashMap, - }, - AvroSchema { - schema: String, - }, -} - -#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)] -#[serde(rename_all = "snake_case")] -pub struct ConnectionSchema { - pub format: Option, - #[serde(default)] - pub bad_data: Option, - #[serde(default)] - pub framing: Option, - #[serde(default)] - pub fields: Vec, - #[serde(default)] - pub definition: Option, - #[serde(default)] - pub inferred: Option, - #[serde(default)] - pub primary_keys: HashSet, -} - -impl ConnectionSchema { - pub fn try_new( - format: Option, - bad_data: Option, - framing: Option, - fields: Vec, - definition: Option, - inferred: Option, - primary_keys: HashSet, - ) -> anyhow::Result { - let s = ConnectionSchema { - format, - bad_data, - framing, - fields, - definition, - inferred, - primary_keys, - }; - s.validate() - } - - pub fn validate(self) -> anyhow::Result { - let non_metadata_fields: Vec<_> = self - .fields - .iter() - .filter(|f| f.metadata_key.is_none()) - .collect(); - - if let Some(Format::RawString(_)) = &self.format - && (non_metadata_fields.len() != 1 - || non_metadata_fields.first().unwrap().field_type != FieldType::String - || non_metadata_fields.first().unwrap().name != "value") - { - anyhow::bail!( - "raw_string format requires a schema with a single field called `value` of type TEXT" - ); - } - - if let Some(Format::Json(json_format)) = &self.format - && json_format.unstructured - && (non_metadata_fields.len() != 1 - || non_metadata_fields.first().unwrap().field_type != FieldType::Json - || non_metadata_fields.first().unwrap().name != "value") - { - anyhow::bail!( - "json format with unstructured flag enabled requires a schema with a single field called `value` of type JSON" - ); - } - - Ok(self) - } - - pub fn fs_schema(&self) -> Arc { - let fields: Vec = self.fields.iter().map(|f| f.clone().into()).collect(); - Arc::new(FsSchema::from_fields(fields)) - } -} - -impl From for FsSchema { - fn from(val: ConnectionSchema) -> Self { - let fields: Vec = val.fields.into_iter().map(|f| f.into()).collect(); - FsSchema::from_fields(fields) - } -} - -// ─────────────────── Connection Table ─────────────────── - -#[derive(Serialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct ConnectionTable { - #[serde(skip_serializing)] - pub id: i64, - #[serde(rename = "id")] - pub pub_id: String, - pub name: String, - pub created_at: u64, - pub connector: String, - pub connection_profile: Option, - pub table_type: ConnectionType, - pub config: serde_json::Value, - pub schema: ConnectionSchema, - pub consumers: u32, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct ConnectionTablePost { - pub name: String, - pub connector: String, - pub connection_profile_id: Option, - pub config: serde_json::Value, - pub schema: Option, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct ConnectionAutocompleteResp { - pub values: BTreeMap>, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct TestSourceMessage { - pub error: bool, - pub done: bool, - pub message: String, -} - -impl TestSourceMessage { - pub fn info(message: impl Into) -> Self { - Self { - error: false, - done: false, - message: message.into(), - } - } - pub fn error(message: impl Into) -> Self { - Self { - error: true, - done: false, - message: message.into(), - } - } - pub fn done(message: impl Into) -> Self { - Self { - error: false, - done: true, - message: message.into(), - } - } - pub fn fail(message: impl Into) -> Self { - Self { - error: true, - done: true, - message: message.into(), - } - } -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct ConfluentSchema { - pub schema: String, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct ConfluentSchemaQueryParams { - pub endpoint: String, - pub topic: String, -} diff --git a/src/sql/api/metrics.rs b/src/sql/api/metrics.rs deleted file mode 100644 index 671b52f6..00000000 --- a/src/sql/api/metrics.rs +++ /dev/null @@ -1,53 +0,0 @@ -// 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 serde::{Deserialize, Serialize}; - -#[derive(Serialize, Deserialize, Copy, Clone, Debug, Hash, PartialEq, Eq)] -#[serde(rename_all = "snake_case")] -pub enum MetricName { - BytesRecv, - BytesSent, - MessagesRecv, - MessagesSent, - Backpressure, - TxQueueSize, - TxQueueRem, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct Metric { - pub time: u64, - pub value: f64, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct SubtaskMetrics { - pub index: u32, - pub metrics: Vec, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct MetricGroup { - pub name: MetricName, - pub subtasks: Vec, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct OperatorMetricGroup { - pub node_id: u32, - pub metric_groups: Vec, -} diff --git a/src/sql/api/mod.rs b/src/sql/api/mod.rs deleted file mode 100644 index 9fc6b23f..00000000 --- a/src/sql/api/mod.rs +++ /dev/null @@ -1,46 +0,0 @@ -// 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. - -//! REST/RPC API types for the FunctionStream system. -//! -//! Adapted from Arroyo's `arroyo-rpc/src/api_types` and utility modules. - -pub mod checkpoints; -pub mod connections; -pub mod metrics; -pub mod pipelines; -pub mod public_ids; -pub mod schema_resolver; -pub mod udfs; -pub mod var_str; - -use serde::{Deserialize, Serialize}; - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "camelCase")] -pub struct PaginatedCollection { - pub data: Vec, - pub has_more: bool, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "camelCase")] -pub struct NonPaginatedCollection { - pub data: Vec, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct PaginationQueryParams { - pub starting_after: Option, - pub limit: Option, -} diff --git a/src/sql/api/pipelines.rs b/src/sql/api/pipelines.rs deleted file mode 100644 index d6cc5253..00000000 --- a/src/sql/api/pipelines.rs +++ /dev/null @@ -1,168 +0,0 @@ -// 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 super::udfs::Udf; -use crate::sql::common::control::ErrorDomain; -use serde::{Deserialize, Serialize}; - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct ValidateQueryPost { - pub query: String, - pub udfs: Option>, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct QueryValidationResult { - pub graph: Option, - pub errors: Vec, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct PipelinePost { - pub name: String, - pub query: String, - pub udfs: Option>, - pub parallelism: u64, - pub checkpoint_interval_micros: Option, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct PreviewPost { - pub query: String, - pub udfs: Option>, - #[serde(default)] - pub enable_sinks: bool, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct PipelinePatch { - pub parallelism: Option, - pub checkpoint_interval_micros: Option, - pub stop: Option, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct PipelineRestart { - pub force: Option, - pub ignore_state: Option, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct Pipeline { - pub id: String, - pub name: String, - pub query: String, - pub udfs: Vec, - pub checkpoint_interval_micros: u64, - pub stop: StopType, - pub created_at: u64, - pub action: Option, - pub action_text: String, - pub action_in_progress: bool, - pub graph: PipelineGraph, - pub preview: bool, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct PipelineGraph { - pub nodes: Vec, - pub edges: Vec, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct PipelineNode { - pub node_id: u32, - pub operator: String, - pub description: String, - pub parallelism: u32, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct PipelineEdge { - pub src_id: u32, - pub dest_id: u32, - pub key_type: String, - pub value_type: String, - pub edge_type: String, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub enum StopType { - None, - Checkpoint, - Graceful, - Immediate, - Force, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct FailureReason { - pub error: String, - pub domain: ErrorDomain, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct Job { - pub id: String, - pub running_desired: bool, - pub state: String, - pub run_id: u64, - pub start_time: Option, - pub finish_time: Option, - pub tasks: Option, - pub failure_reason: Option, - pub created_at: u64, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub enum JobLogLevel { - Info, - Warn, - Error, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct JobLogMessage { - pub id: String, - pub created_at: u64, - pub operator_id: Option, - pub task_index: Option, - pub level: JobLogLevel, - pub message: String, - pub details: String, - pub error_domain: Option, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct OutputData { - pub operator_id: String, - pub subtask_idx: u32, - pub timestamps: Vec, - pub start_id: u64, - pub batch: String, -} diff --git a/src/sql/api/public_ids.rs b/src/sql/api/public_ids.rs deleted file mode 100644 index 33aa6427..00000000 --- a/src/sql/api/public_ids.rs +++ /dev/null @@ -1,69 +0,0 @@ -// 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::time::{SystemTime, UNIX_EPOCH}; - -const ID_LENGTH: usize = 10; - -const ALPHABET: &[u8; 62] = b"0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"; - -pub enum IdTypes { - ApiKey, - ConnectionProfile, - Schema, - Pipeline, - JobConfig, - Checkpoint, - JobStatus, - ClusterInfo, - JobLogMessage, - ConnectionTable, - ConnectionTablePipeline, - Udf, -} - -/// Generates a unique identifier with a type-specific prefix. -/// -/// Uses a simple time + random approach instead of nanoid to avoid an extra dependency. -pub fn generate_id(id_type: IdTypes) -> String { - let prefix = match id_type { - IdTypes::ApiKey => "ak", - IdTypes::ConnectionProfile => "cp", - IdTypes::Schema => "sch", - IdTypes::Pipeline => "pl", - IdTypes::JobConfig => "job", - IdTypes::Checkpoint => "chk", - IdTypes::JobStatus => "js", - IdTypes::ClusterInfo => "ci", - IdTypes::JobLogMessage => "jlm", - IdTypes::ConnectionTable => "ct", - IdTypes::ConnectionTablePipeline => "ctp", - IdTypes::Udf => "udf", - }; - - let nanos = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_nanos(); - - let mut id = String::with_capacity(ID_LENGTH); - let mut seed = nanos; - for _ in 0..ID_LENGTH { - seed ^= seed - .wrapping_mul(6364136223846793005) - .wrapping_add(1442695040888963407); - let idx = (seed % ALPHABET.len() as u128) as usize; - id.push(ALPHABET[idx] as char); - } - - format!("{prefix}_{id}") -} diff --git a/src/sql/api/schema_resolver.rs b/src/sql/api/schema_resolver.rs deleted file mode 100644 index 57d3d702..00000000 --- a/src/sql/api/schema_resolver.rs +++ /dev/null @@ -1,94 +0,0 @@ -// 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 async_trait::async_trait; - -/// Trait for resolving schemas by ID (e.g., from a schema registry). -#[async_trait] -pub trait SchemaResolver: Send { - async fn resolve_schema(&self, id: u32) -> Result, String>; -} - -/// A resolver that always fails — used when no schema registry is configured. -pub struct FailingSchemaResolver; - -impl Default for FailingSchemaResolver { - fn default() -> Self { - Self - } -} - -#[async_trait] -impl SchemaResolver for FailingSchemaResolver { - async fn resolve_schema(&self, id: u32) -> Result, String> { - Err(format!( - "Schema with id {id} not available, and no schema registry configured" - )) - } -} - -/// A resolver that returns a fixed schema for a known ID. -pub struct FixedSchemaResolver { - id: u32, - schema: String, -} - -impl FixedSchemaResolver { - pub fn new(id: u32, schema: String) -> Self { - FixedSchemaResolver { id, schema } - } -} - -#[async_trait] -impl SchemaResolver for FixedSchemaResolver { - async fn resolve_schema(&self, id: u32) -> Result, String> { - if id == self.id { - Ok(Some(self.schema.clone())) - } else { - Err(format!("Unexpected schema id {}, expected {}", id, self.id)) - } - } -} - -/// A caching wrapper around any `SchemaResolver`. -pub struct CachingSchemaResolver { - inner: R, - cache: tokio::sync::RwLock>, -} - -impl CachingSchemaResolver { - pub fn new(inner: R) -> Self { - Self { - inner, - cache: tokio::sync::RwLock::new(std::collections::HashMap::new()), - } - } -} - -#[async_trait] -impl SchemaResolver for CachingSchemaResolver { - async fn resolve_schema(&self, id: u32) -> Result, String> { - { - let cache = self.cache.read().await; - if let Some(schema) = cache.get(&id) { - return Ok(Some(schema.clone())); - } - } - - let result = self.inner.resolve_schema(id).await?; - if let Some(ref schema) = result { - let mut cache = self.cache.write().await; - cache.insert(id, schema.clone()); - } - Ok(result) - } -} diff --git a/src/sql/api/udfs.rs b/src/sql/api/udfs.rs deleted file mode 100644 index 781d5b07..00000000 --- a/src/sql/api/udfs.rs +++ /dev/null @@ -1,68 +0,0 @@ -// 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 serde::{Deserialize, Serialize}; - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct Udf { - pub definition: String, - #[serde(default)] - pub language: UdfLanguage, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct ValidateUdfPost { - pub definition: String, - #[serde(default)] - pub language: UdfLanguage, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct UdfValidationResult { - pub udf_name: Option, - pub errors: Vec, -} - -#[derive(Serialize, Deserialize, Copy, Clone, Debug, Default, Eq, PartialEq)] -#[serde(rename_all = "snake_case")] -pub enum UdfLanguage { - Python, - #[default] - Rust, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct UdfPost { - pub prefix: String, - #[serde(default)] - pub language: UdfLanguage, - pub definition: String, - pub description: Option, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(rename_all = "snake_case")] -pub struct GlobalUdf { - pub id: String, - pub prefix: String, - pub name: String, - pub language: UdfLanguage, - pub created_at: u64, - pub updated_at: u64, - pub definition: String, - pub description: Option, - pub dylib_url: Option, -} diff --git a/src/sql/api/var_str.rs b/src/sql/api/var_str.rs deleted file mode 100644 index 2638cd06..00000000 --- a/src/sql/api/var_str.rs +++ /dev/null @@ -1,91 +0,0 @@ -// 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 serde::{Deserialize, Serialize}; -use std::env; - -/// A string that may contain `{{ VAR }}` placeholders for environment variable substitution. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -#[serde(transparent)] -pub struct VarStr { - raw_val: String, -} - -impl VarStr { - pub fn new(raw_val: String) -> Self { - VarStr { raw_val } - } - - pub fn raw(&self) -> &str { - &self.raw_val - } - - /// Substitute `{{ VAR_NAME }}` patterns with the corresponding environment variable values. - pub fn sub_env_vars(&self) -> anyhow::Result { - let mut result = self.raw_val.clone(); - let mut start = 0; - - while let Some(open) = result[start..].find("{{") { - let open_abs = start + open; - let Some(close) = result[open_abs..].find("}}") else { - break; - }; - let close_abs = open_abs + close; - - let var_name = result[open_abs + 2..close_abs].trim(); - if var_name.is_empty() { - start = close_abs + 2; - continue; - } - - match env::var(var_name) { - Ok(value) => { - let full_match = &result[open_abs..close_abs + 2]; - let full_match_owned = full_match.to_string(); - result = result.replacen(&full_match_owned, &value, 1); - start = open_abs + value.len(); - } - Err(_) => { - anyhow::bail!("Environment variable {} not found", var_name); - } - } - } - - Ok(result) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_no_placeholders() { - let input = "This is a test string with no placeholders"; - assert_eq!( - VarStr::new(input.to_string()).sub_env_vars().unwrap(), - input - ); - } - - #[test] - fn test_with_placeholders() { - unsafe { env::set_var("FS_TEST_VAR", "environment variable") }; - let input = "This is a {{ FS_TEST_VAR }}"; - let expected = "This is a environment variable"; - assert_eq!( - VarStr::new(input.to_string()).sub_env_vars().unwrap(), - expected - ); - unsafe { env::remove_var("FS_TEST_VAR") }; - } -} diff --git a/src/sql/common/arrow_ext.rs b/src/sql/common/arrow_ext.rs deleted file mode 100644 index f041ec6f..00000000 --- a/src/sql/common/arrow_ext.rs +++ /dev/null @@ -1,182 +0,0 @@ -// 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::fmt::{Display, Formatter}; -use std::time::SystemTime; - -use datafusion::arrow::datatypes::{DataType, Field, TimeUnit}; - -pub struct DisplayAsSql<'a>(pub &'a DataType); - -impl Display for DisplayAsSql<'_> { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - match self.0 { - DataType::Boolean => write!(f, "BOOLEAN"), - DataType::Int8 | DataType::Int16 | DataType::Int32 => write!(f, "INT"), - DataType::Int64 => write!(f, "BIGINT"), - DataType::UInt8 | DataType::UInt16 | DataType::UInt32 => write!(f, "INT UNSIGNED"), - DataType::UInt64 => write!(f, "BIGINT UNSIGNED"), - DataType::Float16 | DataType::Float32 => write!(f, "FLOAT"), - DataType::Float64 => write!(f, "DOUBLE"), - DataType::Timestamp(_, _) => write!(f, "TIMESTAMP"), - DataType::Date32 => write!(f, "DATE"), - DataType::Date64 => write!(f, "DATETIME"), - DataType::Time32(_) => write!(f, "TIME"), - DataType::Time64(_) => write!(f, "TIME"), - DataType::Duration(_) => write!(f, "INTERVAL"), - DataType::Interval(_) => write!(f, "INTERVAL"), - DataType::Binary | DataType::FixedSizeBinary(_) | DataType::LargeBinary => { - write!(f, "BYTEA") - } - DataType::Utf8 | DataType::LargeUtf8 => write!(f, "TEXT"), - DataType::List(inner) => { - write!(f, "{}[]", DisplayAsSql(inner.data_type())) - } - dt => write!(f, "{dt}"), - } - } -} - -/// Arrow extension type markers for FunctionStream-specific semantics. -#[allow(clippy::upper_case_acronyms)] -#[derive(Copy, Clone, Eq, PartialEq, Debug)] -pub enum FsExtensionType { - JSON, -} - -impl FsExtensionType { - pub fn from_map(map: &HashMap) -> Option { - match map.get("ARROW:extension:name")?.as_str() { - "functionstream.json" => Some(Self::JSON), - _ => None, - } - } - - pub fn add_metadata(v: Option, field: Field) -> Field { - if let Some(v) = v { - let mut m = HashMap::new(); - match v { - FsExtensionType::JSON => { - m.insert( - "ARROW:extension:name".to_string(), - "functionstream.json".to_string(), - ); - } - } - field.with_metadata(m) - } else { - field - } - } -} - -pub trait GetArrowType { - fn arrow_type() -> DataType; -} - -pub trait GetArrowSchema { - fn arrow_schema() -> datafusion::arrow::datatypes::Schema; -} - -impl GetArrowType for T -where - T: GetArrowSchema, -{ - fn arrow_type() -> DataType { - DataType::Struct(Self::arrow_schema().fields.clone()) - } -} - -impl GetArrowType for bool { - fn arrow_type() -> DataType { - DataType::Boolean - } -} - -impl GetArrowType for i8 { - fn arrow_type() -> DataType { - DataType::Int8 - } -} - -impl GetArrowType for i16 { - fn arrow_type() -> DataType { - DataType::Int16 - } -} - -impl GetArrowType for i32 { - fn arrow_type() -> DataType { - DataType::Int32 - } -} - -impl GetArrowType for i64 { - fn arrow_type() -> DataType { - DataType::Int64 - } -} - -impl GetArrowType for u8 { - fn arrow_type() -> DataType { - DataType::UInt8 - } -} - -impl GetArrowType for u16 { - fn arrow_type() -> DataType { - DataType::UInt16 - } -} - -impl GetArrowType for u32 { - fn arrow_type() -> DataType { - DataType::UInt32 - } -} - -impl GetArrowType for u64 { - fn arrow_type() -> DataType { - DataType::UInt64 - } -} - -impl GetArrowType for f32 { - fn arrow_type() -> DataType { - DataType::Float32 - } -} - -impl GetArrowType for f64 { - fn arrow_type() -> DataType { - DataType::Float64 - } -} - -impl GetArrowType for String { - fn arrow_type() -> DataType { - DataType::Utf8 - } -} - -impl GetArrowType for Vec { - fn arrow_type() -> DataType { - DataType::Binary - } -} - -impl GetArrowType for SystemTime { - fn arrow_type() -> DataType { - DataType::Timestamp(TimeUnit::Nanosecond, None) - } -} diff --git a/src/sql/common/connector_options.rs b/src/sql/common/connector_options.rs deleted file mode 100644 index 0702d945..00000000 --- a/src/sql/common/connector_options.rs +++ /dev/null @@ -1,449 +0,0 @@ -// 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, HashMap}; -use std::num::{NonZero, NonZeroU64}; -use std::str::FromStr; -use std::time::Duration; - -use datafusion::common::{Result as DFResult, plan_datafusion_err}; -use datafusion::error::DataFusionError; -use datafusion::sql::sqlparser::ast::{Expr, Ident, SqlOption, Value as SqlValue, ValueWithSpan}; -use tracing::warn; - -use super::constants::{interval_duration_unit, with_opt_bool_str}; - -pub trait FromOpts: Sized { - fn from_opts(opts: &mut ConnectorOptions) -> DFResult; -} - -pub struct ConnectorOptions { - options: HashMap, - partitions: Vec, -} - -fn sql_expr_to_catalog_string(e: &Expr) -> String { - match e { - Expr::Value(ValueWithSpan { value, .. }) => match value { - SqlValue::SingleQuotedString(s) | SqlValue::DoubleQuotedString(s) => s.clone(), - SqlValue::NationalStringLiteral(s) => s.clone(), - SqlValue::HexStringLiteral(s) => s.clone(), - SqlValue::Number(n, _) => n.clone(), - SqlValue::Boolean(b) => b.to_string(), - SqlValue::Null => "NULL".to_string(), - other => other.to_string(), - }, - Expr::Identifier(ident) => ident.value.clone(), - other => other.to_string(), - } -} - -impl ConnectorOptions { - /// Build options from persisted catalog string maps (same semantics as SQL `WITH` literals). - pub fn from_flat_string_map(map: HashMap) -> DFResult { - let mut options = HashMap::with_capacity(map.len()); - for (k, v) in map { - options.insert( - k, - Expr::Value(SqlValue::SingleQuotedString(v).with_empty_span()), - ); - } - Ok(Self { - options, - partitions: Vec::new(), - }) - } - - pub fn new(sql_opts: &[SqlOption], partition_by: &Option>) -> DFResult { - let mut options = HashMap::new(); - - for option in sql_opts { - let SqlOption::KeyValue { key, value } = option else { - return Err(plan_datafusion_err!( - "invalid with option: '{}'; expected an `=` delimited key-value pair", - option - )); - }; - - options.insert(key.value.clone(), value.clone()); - } - - Ok(Self { - options, - partitions: partition_by.clone().unwrap_or_default(), - }) - } - - pub fn partitions(&self) -> &[Expr] { - &self.partitions - } - - pub fn pull_struct(&mut self) -> DFResult { - T::from_opts(self) - } - - pub fn pull_opt_str(&mut self, name: &str) -> DFResult> { - match self.options.remove(name) { - Some(Expr::Value(ValueWithSpan { - value: SqlValue::SingleQuotedString(s), - span: _, - })) => Ok(Some(s)), - Some(e) => Err(plan_datafusion_err!( - "expected with option '{}' to be a single-quoted string, but it was `{:?}`", - name, - e - )), - None => Ok(None), - } - } - - pub fn peek_opt_str(&self, name: &str) -> DFResult> { - match self.options.get(name) { - Some(Expr::Value(ValueWithSpan { - value: SqlValue::SingleQuotedString(s), - span: _, - })) => Ok(Some(s.clone())), - Some(e) => Err(plan_datafusion_err!( - "expected with option '{}' to be a single-quoted string, but it was `{:?}`", - name, - e - )), - None => Ok(None), - } - } - - pub fn pull_str(&mut self, name: &str) -> DFResult { - self.pull_opt_str(name)? - .ok_or_else(|| plan_datafusion_err!("required option '{}' not set", name)) - } - - pub fn pull_opt_bool(&mut self, name: &str) -> DFResult> { - match self.options.remove(name) { - Some(Expr::Value(ValueWithSpan { - value: SqlValue::Boolean(b), - span: _, - })) => Ok(Some(b)), - Some(Expr::Value(ValueWithSpan { - value: SqlValue::SingleQuotedString(s), - span: _, - })) => match s.as_str() { - with_opt_bool_str::TRUE | with_opt_bool_str::YES => Ok(Some(true)), - with_opt_bool_str::FALSE | with_opt_bool_str::NO => Ok(Some(false)), - _ => Err(plan_datafusion_err!( - "expected with option '{}' to be a boolean, but it was `'{}'`", - name, - s - )), - }, - Some(e) => Err(plan_datafusion_err!( - "expected with option '{}' to be a boolean, but it was `{:?}`", - name, - e - )), - None => Ok(None), - } - } - - pub fn pull_opt_u64(&mut self, name: &str) -> DFResult> { - match self.options.remove(name) { - Some(Expr::Value(ValueWithSpan { - value: SqlValue::Number(s, _), - span: _, - })) - | Some(Expr::Value(ValueWithSpan { - value: SqlValue::SingleQuotedString(s), - span: _, - })) => s.parse::().map(Some).map_err(|_| { - plan_datafusion_err!( - "expected with option '{}' to be an unsigned integer, but it was `{}`", - name, - s - ) - }), - Some(e) => Err(plan_datafusion_err!( - "expected with option '{}' to be an unsigned integer, but it was `{:?}`", - name, - e - )), - None => Ok(None), - } - } - - pub fn pull_opt_nonzero_u64(&mut self, name: &str) -> DFResult>> { - match self.pull_opt_u64(name)? { - Some(0) => Err(plan_datafusion_err!( - "expected with option '{name}' to be greater than 0, but it was 0" - )), - Some(i) => Ok(Some(NonZeroU64::new(i).unwrap())), - None => Ok(None), - } - } - - pub fn pull_opt_data_size_bytes(&mut self, name: &str) -> DFResult> { - self.pull_opt_str(name)? - .map(|s| { - s.parse::().map_err(|_| { - plan_datafusion_err!( - "expected with option '{}' to be a size in bytes (unsigned integer), but it was `{}`", - name, - s - ) - }) - }) - .transpose() - } - - pub fn pull_opt_i64(&mut self, name: &str) -> DFResult> { - match self.options.remove(name) { - Some(Expr::Value(ValueWithSpan { - value: SqlValue::Number(s, _), - span: _, - })) - | Some(Expr::Value(ValueWithSpan { - value: SqlValue::SingleQuotedString(s), - span: _, - })) => s.parse::().map(Some).map_err(|_| { - plan_datafusion_err!( - "expected with option '{}' to be an integer, but it was `{}`", - name, - s - ) - }), - Some(e) => Err(plan_datafusion_err!( - "expected with option '{}' to be an integer, but it was `{:?}`", - name, - e - )), - None => Ok(None), - } - } - - pub fn pull_i64(&mut self, name: &str) -> DFResult { - self.pull_opt_i64(name)? - .ok_or_else(|| plan_datafusion_err!("required option '{}' not set", name)) - } - - pub fn pull_u64(&mut self, name: &str) -> DFResult { - self.pull_opt_u64(name)? - .ok_or_else(|| plan_datafusion_err!("required option '{}' not set", name)) - } - - pub fn pull_opt_f64(&mut self, name: &str) -> DFResult> { - match self.options.remove(name) { - Some(Expr::Value(ValueWithSpan { - value: SqlValue::Number(s, _), - span: _, - })) - | Some(Expr::Value(ValueWithSpan { - value: SqlValue::SingleQuotedString(s), - span: _, - })) => s.parse::().map(Some).map_err(|_| { - plan_datafusion_err!( - "expected with option '{}' to be a double, but it was `{}`", - name, - s - ) - }), - Some(e) => Err(plan_datafusion_err!( - "expected with option '{}' to be a double, but it was `{:?}`", - name, - e - )), - None => Ok(None), - } - } - - pub fn pull_f64(&mut self, name: &str) -> DFResult { - self.pull_opt_f64(name)? - .ok_or_else(|| plan_datafusion_err!("required option '{}' not set", name)) - } - - pub fn pull_bool(&mut self, name: &str) -> DFResult { - self.pull_opt_bool(name)? - .ok_or_else(|| plan_datafusion_err!("required option '{}' not set", name)) - } - - pub fn pull_opt_duration(&mut self, name: &str) -> DFResult> { - match self.options.remove(name) { - Some(e) => Ok(Some(duration_from_sql_expr(&e).map_err(|e| { - plan_datafusion_err!("in with clause '{name}': {}", e) - })?)), - None => Ok(None), - } - } - - pub fn pull_opt_field(&mut self, name: &str) -> DFResult> { - match self.options.remove(name) { - Some(Expr::Value(ValueWithSpan { - value: SqlValue::SingleQuotedString(s), - span: _, - })) => { - warn!( - "Referred to a field in `{name}` with a string—this is deprecated and will be unsupported after Arroyo 0.14" - ); - Ok(Some(s)) - } - Some(Expr::Identifier(Ident { value, .. })) => Ok(Some(value)), - Some(e) => Err(plan_datafusion_err!( - "expected with option '{}' to be a field, but it was `{:?}`", - name, - e - )), - None => Ok(None), - } - } - - pub fn pull_opt_array(&mut self, name: &str) -> Option> { - Some(match self.options.remove(name)? { - Expr::Value(ValueWithSpan { - value: SqlValue::SingleQuotedString(s), - span, - }) => s - .split(',') - .map(|p| { - Expr::Value(ValueWithSpan { - value: SqlValue::SingleQuotedString(p.to_string()), - span, - }) - }) - .collect(), - Expr::Array(a) => a.elem, - e => vec![e], - }) - } - - pub fn pull_opt_parsed(&mut self, name: &str) -> DFResult> { - Ok(match self.pull_opt_str(name)? { - Some(s) => Some( - s.parse() - .map_err(|_| plan_datafusion_err!("invalid value '{s}' for {name}"))?, - ), - None => None, - }) - } - - pub fn keys(&self) -> impl Iterator { - self.options.keys() - } - - pub fn keys_with_prefix<'a, 'b>( - &'a self, - prefix: &'b str, - ) -> impl Iterator + 'b - where - 'a: 'b, - { - self.options.keys().filter(move |k| k.starts_with(prefix)) - } - - pub fn insert_str( - &mut self, - name: impl Into, - value: impl Into, - ) -> DFResult> { - let name = name.into(); - let value = value.into(); - let existing = self.pull_opt_str(&name)?; - self.options.insert( - name, - Expr::Value(SqlValue::SingleQuotedString(value).with_empty_span()), - ); - Ok(existing) - } - - pub fn is_empty(&self) -> bool { - self.options.is_empty() - } - - pub fn contains_key(&self, key: &str) -> bool { - self.options.contains_key(key) - } - - /// Drain all remaining options into string values (for connector runtime config). - pub fn drain_remaining_string_values(&mut self) -> DFResult> { - let taken = std::mem::take(&mut self.options); - let mut out = HashMap::with_capacity(taken.len()); - for (k, v) in taken { - out.insert(k, format!("{v}")); - } - Ok(out) - } - - /// Snapshot of all current `WITH` key/value pairs for catalog persistence (`SHOW CREATE TABLE`). - /// Call before any `pull_*` consumes options. - pub fn snapshot_for_catalog(&self) -> BTreeMap { - self.options - .iter() - .map(|(k, v)| (k.clone(), sql_expr_to_catalog_string(v))) - .collect() - } -} - -fn duration_from_sql_expr(expr: &Expr) -> Result { - match expr { - Expr::Interval(interval) => { - let s = match interval.value.as_ref() { - Expr::Value(ValueWithSpan { - value: SqlValue::SingleQuotedString(s), - .. - }) => s.clone(), - other => { - return Err(DataFusionError::Plan(format!( - "expected interval string literal, found {other}" - ))); - } - }; - parse_interval_to_duration(&s) - } - Expr::Value(ValueWithSpan { - value: SqlValue::SingleQuotedString(s), - .. - }) => parse_interval_to_duration(s), - other => Err(DataFusionError::Plan(format!( - "expected an interval expression, found {other}" - ))), - } -} - -fn parse_interval_to_duration(s: &str) -> Result { - let parts: Vec<&str> = s.split_whitespace().collect(); - if parts.len() != 2 { - return Err(DataFusionError::Plan(format!( - "invalid interval string '{s}'; expected ' '" - ))); - } - let value: u64 = parts[0] - .parse() - .map_err(|_| DataFusionError::Plan(format!("invalid interval number: {}", parts[0])))?; - let unit_lc = parts[1].to_lowercase(); - let unit = unit_lc.as_str(); - let duration = match unit { - interval_duration_unit::SECOND - | interval_duration_unit::SECONDS - | interval_duration_unit::S => Duration::from_secs(value), - interval_duration_unit::MINUTE - | interval_duration_unit::MINUTES - | interval_duration_unit::MIN => Duration::from_secs(value * 60), - interval_duration_unit::HOUR - | interval_duration_unit::HOURS - | interval_duration_unit::H => Duration::from_secs(value * 3600), - interval_duration_unit::DAY | interval_duration_unit::DAYS | interval_duration_unit::D => { - Duration::from_secs(value * 86400) - } - unit => { - return Err(DataFusionError::Plan(format!( - "unsupported interval unit '{unit}'" - ))); - } - }; - Ok(duration) -} diff --git a/src/sql/common/constants.rs b/src/sql/common/constants.rs deleted file mode 100644 index 8cdb68e3..00000000 --- a/src/sql/common/constants.rs +++ /dev/null @@ -1,294 +0,0 @@ -// 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. - -pub mod scalar_fn { - pub const GET_FIRST_JSON_OBJECT: &str = "get_first_json_object"; - pub const EXTRACT_JSON: &str = "extract_json"; - pub const EXTRACT_JSON_STRING: &str = "extract_json_string"; - pub const SERIALIZE_JSON_UNION: &str = "serialize_json_union"; - pub const MULTI_HASH: &str = "multi_hash"; -} - -pub mod window_fn { - pub const HOP: &str = "hop"; - pub const TUMBLE: &str = "tumble"; - pub const SESSION: &str = "session"; -} - -pub mod planning_placeholder_udf { - pub const UNNEST: &str = "unnest"; - pub const ROW_TIME: &str = "row_time"; - pub const LIST_ELEMENT_FIELD: &str = "field"; -} - -pub mod operator_feature { - pub const ASYNC_UDF: &str = "async-udf"; - pub const JOIN_WITH_EXPIRATION: &str = "join-with-expiration"; - pub const WINDOWED_JOIN: &str = "windowed-join"; - pub const SQL_WINDOW_FUNCTION: &str = "sql-window-function"; - pub const LOOKUP_JOIN: &str = "lookup-join"; - pub const SQL_TUMBLING_WINDOW_AGGREGATE: &str = "sql-tumbling-window-aggregate"; - pub const SQL_SLIDING_WINDOW_AGGREGATE: &str = "sql-sliding-window-aggregate"; - pub const SQL_SESSION_WINDOW_AGGREGATE: &str = "sql-session-window-aggregate"; - pub const SQL_UPDATING_AGGREGATE: &str = "sql-updating-aggregate"; - pub const KEY_BY_ROUTING: &str = "key-by-routing"; - pub const CONNECTOR_SOURCE: &str = "connector-source"; - pub const CONNECTOR_SINK: &str = "connector-sink"; -} - -pub mod extension_node { - pub const STREAM_WINDOW_AGGREGATE: &str = "StreamWindowAggregateNode"; - pub const STREAMING_WINDOW_FUNCTION: &str = "StreamingWindowFunctionNode"; - pub const EVENT_TIME_WATERMARK: &str = "EventTimeWatermarkNode"; - pub const CONTINUOUS_AGGREGATE: &str = "ContinuousAggregateNode"; - pub const SYSTEM_TIMESTAMP_INJECTOR: &str = "SystemTimestampInjectorNode"; - pub const STREAM_INGESTION: &str = "StreamIngestionNode"; - pub const STREAM_EGRESS: &str = "StreamEgressNode"; - pub const STREAM_PROJECTION: &str = "StreamProjectionNode"; - pub const REMOTE_TABLE_BOUNDARY: &str = "RemoteTableBoundaryNode"; - pub const REFERENCE_TABLE_SOURCE: &str = "ReferenceTableSource"; - pub const STREAM_REFERENCE_JOIN: &str = "StreamReferenceJoin"; - pub const KEY_EXTRACTION: &str = "KeyExtractionNode"; - pub const STREAMING_JOIN: &str = "StreamingJoinNode"; - pub const ASYNC_FUNCTION_EXECUTION: &str = "AsyncFunctionExecutionNode"; - pub const UNROLL_DEBEZIUM_PAYLOAD: &str = "UnrollDebeziumPayloadNode"; - pub const PACK_DEBEZIUM_ENVELOPE: &str = "PackDebeziumEnvelopeNode"; -} - -pub mod proto_operator_name { - pub const TUMBLING_WINDOW: &str = "TumblingWindow"; - pub const UPDATING_AGGREGATE: &str = "UpdatingAggregate"; - pub const WINDOW_FUNCTION: &str = "WindowFunction"; - pub const SLIDING_WINDOW_LABEL: &str = "sliding window"; - pub const INSTANT_WINDOW: &str = "InstantWindow"; - pub const INSTANT_WINDOW_LABEL: &str = "instant window"; -} - -pub mod runtime_operator_kind { - pub const STREAMING_JOIN: &str = "streaming_join"; - pub const WATERMARK_GENERATOR: &str = "watermark_generator"; - pub const STREAMING_WINDOW_EVALUATOR: &str = "streaming_window_evaluator"; -} - -pub mod factory_operator_name { - pub const CONNECTOR_SOURCE: &str = "ConnectorSource"; - pub const CONNECTOR_SINK: &str = "ConnectorSink"; - pub const KAFKA_SOURCE: &str = "KafkaSource"; - pub const KAFKA_SINK: &str = "KafkaSink"; -} - -pub mod cdc { - pub const BEFORE: &str = "before"; - pub const AFTER: &str = "after"; - pub const OP: &str = "op"; -} - -pub mod updating_state_field { - pub const IS_RETRACT: &str = "is_retract"; - pub const ID: &str = "id"; -} - -pub mod sql_field { - pub const ASYNC_RESULT: &str = "__async_result"; - pub const DEFAULT_KEY_LABEL: &str = "key"; - pub const DEFAULT_PROJECTION_LABEL: &str = "projection"; - pub const COMPUTED_WATERMARK: &str = "__watermark"; - pub const TIMESTAMP_FIELD: &str = "_timestamp"; - pub const UPDATING_META_FIELD: &str = "_updating_meta"; -} - -pub mod sql_planning_default { - pub const DEFAULT_PARALLELISM: usize = 1; - /// Default physical parallelism for `KeyBy` / key-extraction pipelines (configurable via YAML). - pub const DEFAULT_KEY_BY_PARALLELISM: usize = 1; - /// Parallelism for aggregations that run after `KeyBy` / shuffle on non-empty routing keys. - pub const KEYED_AGGREGATE_DEFAULT_PARALLELISM: usize = 8; - pub const PLANNING_TTL_SECS: u64 = 24 * 60 * 60; -} - -pub mod with_opt_bool_str { - pub const TRUE: &str = "true"; - pub const YES: &str = "yes"; - pub const FALSE: &str = "false"; - pub const NO: &str = "no"; -} - -pub mod interval_duration_unit { - pub const SECOND: &str = "second"; - pub const SECONDS: &str = "seconds"; - pub const S: &str = "s"; - pub const MINUTE: &str = "minute"; - pub const MINUTES: &str = "minutes"; - pub const MIN: &str = "min"; - pub const HOUR: &str = "hour"; - pub const HOURS: &str = "hours"; - pub const H: &str = "h"; - pub const DAY: &str = "day"; - pub const DAYS: &str = "days"; - pub const D: &str = "d"; -} - -pub mod connection_format_value { - pub const JSON: &str = "json"; - pub const CSV: &str = "csv"; - pub const JSONL: &str = "jsonl"; - pub const NDJSON: &str = "ndjson"; - pub const LANCE: &str = "lance"; - pub const ORC: &str = "orc"; - pub const DEBEZIUM_JSON: &str = "debezium_json"; - pub const AVRO: &str = "avro"; - pub const PARQUET: &str = "parquet"; - pub const PROTOBUF: &str = "protobuf"; - pub const RAW_STRING: &str = "raw_string"; - pub const RAW_BYTES: &str = "raw_bytes"; -} - -pub mod framing_method_value { - pub const NEWLINE: &str = "newline"; - pub const NEWLINE_DELIMITED: &str = "newline_delimited"; -} - -pub mod bad_data_value { - pub const FAIL: &str = "fail"; - pub const DROP: &str = "drop"; -} - -pub mod timestamp_format_value { - pub const RFC3339_SNAKE: &str = "rfc3339"; - pub const RFC3339_UPPER: &str = "RFC3339"; - pub const UNIX_MILLIS_SNAKE: &str = "unix_millis"; - pub const UNIX_MILLIS_PASCAL: &str = "UnixMillis"; -} - -pub mod decimal_encoding_value { - pub const NUMBER: &str = "number"; - pub const STRING: &str = "string"; - pub const BYTES: &str = "bytes"; -} - -pub mod json_compression_value { - pub const UNCOMPRESSED: &str = "uncompressed"; - pub const GZIP: &str = "gzip"; -} - -pub mod parquet_compression_value { - pub const UNCOMPRESSED: &str = "uncompressed"; - pub const SNAPPY: &str = "snappy"; - pub const GZIP: &str = "gzip"; - pub const ZSTD: &str = "zstd"; - pub const LZ4: &str = "lz4"; - pub const LZ4_RAW: &str = "lz4_raw"; -} - -pub mod date_part_keyword { - pub const YEAR: &str = "year"; - pub const MONTH: &str = "month"; - pub const WEEK: &str = "week"; - pub const DAY: &str = "day"; - pub const HOUR: &str = "hour"; - pub const MINUTE: &str = "minute"; - pub const SECOND: &str = "second"; - pub const MILLISECOND: &str = "millisecond"; - pub const MICROSECOND: &str = "microsecond"; - pub const NANOSECOND: &str = "nanosecond"; - pub const DOW: &str = "dow"; - pub const DOY: &str = "doy"; -} - -pub mod date_trunc_keyword { - pub const YEAR: &str = "year"; - pub const QUARTER: &str = "quarter"; - pub const MONTH: &str = "month"; - pub const WEEK: &str = "week"; - pub const DAY: &str = "day"; - pub const HOUR: &str = "hour"; - pub const MINUTE: &str = "minute"; - pub const SECOND: &str = "second"; -} - -pub mod mem_exec_join_side { - pub const LEFT: &str = "left"; - pub const RIGHT: &str = "right"; -} - -pub mod physical_plan_node_name { - pub const RW_LOCK_READER: &str = "rw_lock_reader"; - pub const UNBOUNDED_READER: &str = "unbounded_reader"; - pub const VEC_READER: &str = "vec_reader"; - pub const MEM_EXEC: &str = "mem_exec"; - pub const DEBEZIUM_UNROLLING_EXEC: &str = "debezium_unrolling_exec"; - pub const TO_DEBEZIUM_EXEC: &str = "to_debezium_exec"; -} - -pub mod window_function_udf { - pub const NAME: &str = "window"; -} - -pub mod window_interval_field { - pub const START: &str = "start"; - pub const END: &str = "end"; -} - -pub mod debezium_op_short { - pub const CREATE: &str = "c"; - pub const READ: &str = "r"; - pub const UPDATE: &str = "u"; - pub const DELETE: &str = "d"; -} - -pub mod connector_type { - pub const KAFKA: &str = "kafka"; - pub const KINESIS: &str = "kinesis"; - pub const FILESYSTEM: &str = "filesystem"; - pub const DELTA: &str = "delta"; - pub const ICEBERG: &str = "iceberg"; - pub const LANCE_DB: &str = "lanceDB"; - pub const S3: &str = "s3"; - pub const PULSAR: &str = "pulsar"; - pub const NATS: &str = "nats"; - pub const REDIS: &str = "redis"; - pub const MQTT: &str = "mqtt"; - pub const WEBSOCKET: &str = "websocket"; - pub const SSE: &str = "sse"; - pub const NEXMARK: &str = "nexmark"; - pub const BLACKHOLE: &str = "blackhole"; - pub const MEMORY: &str = "memory"; - pub const POSTGRES: &str = "postgres"; -} - -pub mod connection_table_role { - pub const SOURCE: &str = "source"; - pub const SINK: &str = "sink"; - pub const LOOKUP: &str = "lookup"; -} - -pub const SUPPORTED_CONNECTOR_ADAPTERS: &[&str] = &[ - connector_type::KAFKA, - connector_type::FILESYSTEM, - connector_type::S3, - connector_type::DELTA, - connector_type::ICEBERG, - connector_type::LANCE_DB, -]; - -pub mod kafka_with_value { - pub const SCAN_LATEST: &str = "latest"; - pub const SCAN_EARLIEST: &str = "earliest"; - pub const SCAN_GROUP_OFFSETS: &str = "group-offsets"; - pub const SCAN_GROUP: &str = "group"; - pub const ISOLATION_READ_COMMITTED: &str = "read_committed"; - pub const ISOLATION_READ_UNCOMMITTED: &str = "read_uncommitted"; - pub const SINK_COMMIT_EXACTLY_ONCE_HYPHEN: &str = "exactly-once"; - pub const SINK_COMMIT_EXACTLY_ONCE_UNDERSCORE: &str = "exactly_once"; - pub const SINK_COMMIT_AT_LEAST_ONCE_HYPHEN: &str = "at-least-once"; - pub const SINK_COMMIT_AT_LEAST_ONCE_UNDERSCORE: &str = "at_least_once"; -} diff --git a/src/sql/common/control.rs b/src/sql/common/control.rs deleted file mode 100644 index eba88596..00000000 --- a/src/sql/common/control.rs +++ /dev/null @@ -1,164 +0,0 @@ -// 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::time::SystemTime; - -use crate::runtime::streaming::protocol::CheckpointBarrier; - -/// Control messages sent from the controller to worker tasks. -#[derive(Debug, Clone)] -pub enum ControlMessage { - Checkpoint(CheckpointBarrier), - Stop { - mode: StopMode, - }, - Commit { - epoch: u32, - commit_data: HashMap>>, - }, - LoadCompacted { - compacted: CompactionResult, - }, - NoOp, -} - -#[derive(Debug, Copy, Clone, PartialEq, Eq)] -pub enum StopMode { - Graceful, - Immediate, -} - -#[derive(Debug, Clone)] -pub struct CompactionResult { - pub operator_id: String, - pub compacted_tables: HashMap, -} - -#[derive(Debug, Clone)] -pub struct TableCheckpointMetadata { - pub table_type: TableType, - pub data: Vec, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum TableType { - GlobalKeyValue, - ExpiringKeyedTimeTable, -} - -/// Responses sent from worker tasks back to the controller. -#[derive(Debug, Clone)] -pub enum ControlResp { - CheckpointEvent(CheckpointEvent), - CheckpointCompleted(CheckpointCompleted), - TaskStarted { - node_id: u32, - task_index: usize, - start_time: SystemTime, - }, - TaskFinished { - node_id: u32, - task_index: usize, - }, - TaskFailed { - node_id: u32, - task_index: usize, - error: TaskError, - }, - Error { - node_id: u32, - operator_id: String, - task_index: usize, - message: String, - details: String, - }, -} - -#[derive(Debug, Clone)] -pub struct CheckpointCompleted { - pub checkpoint_epoch: u32, - pub node_id: u32, - pub operator_id: String, - pub subtask_metadata: SubtaskCheckpointMetadata, -} - -#[derive(Debug, Clone)] -pub struct SubtaskCheckpointMetadata { - pub subtask_index: u32, - pub start_time: u64, - pub finish_time: u64, - pub watermark: Option, - pub bytes: u64, - pub table_metadata: HashMap, - pub table_configs: HashMap, -} - -#[derive(Debug, Clone)] -pub struct TableSubtaskCheckpointMetadata { - pub subtask_index: u32, - pub table_type: TableType, - pub data: Vec, -} - -#[derive(Debug, Clone)] -pub struct TableConfig { - pub table_type: TableType, - pub config: Vec, - pub state_version: u32, -} - -#[derive(Debug, Clone)] -pub struct CheckpointEvent { - pub checkpoint_epoch: u32, - pub node_id: u32, - pub operator_id: String, - pub subtask_index: u32, - pub time: SystemTime, - pub event_type: TaskCheckpointEventType, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum TaskCheckpointEventType { - StartedAlignment, - StartedCheckpointing, - FinishedOperatorSetup, - FinishedSync, - FinishedCommit, -} - -#[derive(Debug, Clone)] -pub struct TaskError { - pub job_id: String, - pub node_id: u32, - pub operator_id: String, - pub operator_subtask: u64, - pub error: String, - pub error_domain: ErrorDomain, - pub retry_hint: RetryHint, - pub details: String, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ErrorDomain { - User, - Internal, - External, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum RetryHint { - NoRetry, - WithBackoff, -} diff --git a/src/sql/common/converter.rs b/src/sql/common/converter.rs deleted file mode 100644 index a9023342..00000000 --- a/src/sql/common/converter.rs +++ /dev/null @@ -1,95 +0,0 @@ -// 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 arrow::row::{OwnedRow, RowConverter, RowParser, Rows, SortField}; -use arrow_array::{Array, ArrayRef, BooleanArray}; -use arrow_schema::{ArrowError, DataType}; -use std::sync::Arc; - -// need to handle the empty case as a row converter without sort fields emits empty Rows. -#[derive(Debug)] -pub enum Converter { - RowConverter(RowConverter), - Empty(RowConverter, Arc), -} - -impl Converter { - pub fn new(sort_fields: Vec) -> Result { - if sort_fields.is_empty() { - let array = Arc::new(BooleanArray::from(vec![false])); - Ok(Self::Empty( - RowConverter::new(vec![SortField::new(DataType::Boolean)])?, - array, - )) - } else { - Ok(Self::RowConverter(RowConverter::new(sort_fields)?)) - } - } - - pub fn convert_columns(&self, columns: &[Arc]) -> Result { - match self { - Converter::RowConverter(row_converter) => { - Ok(row_converter.convert_columns(columns)?.row(0).owned()) - } - Converter::Empty(row_converter, array) => Ok(row_converter - .convert_columns(std::slice::from_ref(array))? - .row(0) - .owned()), - } - } - - pub fn convert_all_columns( - &self, - columns: &[Arc], - num_rows: usize, - ) -> Result { - match self { - Converter::RowConverter(row_converter) => Ok(row_converter.convert_columns(columns)?), - Converter::Empty(row_converter, _array) => { - let array = Arc::new(BooleanArray::from(vec![false; num_rows])); - Ok(row_converter.convert_columns(&[array])?) - } - } - } - - pub fn convert_rows( - &self, - rows: Vec>, - ) -> Result, ArrowError> { - match self { - Converter::RowConverter(row_converter) => Ok(row_converter.convert_rows(rows)?), - Converter::Empty(_row_converter, _array) => Ok(vec![]), - } - } - - pub fn convert_raw_rows(&self, row_bytes: Vec<&[u8]>) -> Result, ArrowError> { - match self { - Converter::RowConverter(row_converter) => { - let parser = row_converter.parser(); - let mut row_list = vec![]; - for bytes in row_bytes { - let row = parser.parse(bytes); - row_list.push(row); - } - Ok(row_converter.convert_rows(row_list)?) - } - Converter::Empty(_row_converter, _array) => Ok(vec![]), - } - } - - pub fn parser(&self) -> Option { - match self { - Converter::RowConverter(r) => Some(r.parser()), - Converter::Empty(_, _) => None, - } - } -} diff --git a/src/sql/common/date.rs b/src/sql/common/date.rs deleted file mode 100644 index ec310326..00000000 --- a/src/sql/common/date.rs +++ /dev/null @@ -1,86 +0,0 @@ -// 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 serde::Serialize; -use std::convert::TryFrom; - -use super::constants::{date_part_keyword, date_trunc_keyword}; - -#[derive(Debug, Clone, Copy, Eq, PartialEq, PartialOrd, Hash, Serialize)] -pub enum DatePart { - Year, - Month, - Week, - Day, - Hour, - Minute, - Second, - Millisecond, - Microsecond, - Nanosecond, - DayOfWeek, - DayOfYear, -} - -impl TryFrom<&str> for DatePart { - type Error = String; - - fn try_from(value: &str) -> Result { - let v = value.to_lowercase(); - match v.as_str() { - date_part_keyword::YEAR => Ok(DatePart::Year), - date_part_keyword::MONTH => Ok(DatePart::Month), - date_part_keyword::WEEK => Ok(DatePart::Week), - date_part_keyword::DAY => Ok(DatePart::Day), - date_part_keyword::HOUR => Ok(DatePart::Hour), - date_part_keyword::MINUTE => Ok(DatePart::Minute), - date_part_keyword::SECOND => Ok(DatePart::Second), - date_part_keyword::MILLISECOND => Ok(DatePart::Millisecond), - date_part_keyword::MICROSECOND => Ok(DatePart::Microsecond), - date_part_keyword::NANOSECOND => Ok(DatePart::Nanosecond), - date_part_keyword::DOW => Ok(DatePart::DayOfWeek), - date_part_keyword::DOY => Ok(DatePart::DayOfYear), - _ => Err(format!("'{value}' is not a valid DatePart")), - } - } -} - -#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, PartialOrd, Serialize)] -pub enum DateTruncPrecision { - Year, - Quarter, - Month, - Week, - Day, - Hour, - Minute, - Second, -} - -impl TryFrom<&str> for DateTruncPrecision { - type Error = String; - - fn try_from(value: &str) -> Result { - let v = value.to_lowercase(); - match v.as_str() { - date_trunc_keyword::YEAR => Ok(DateTruncPrecision::Year), - date_trunc_keyword::QUARTER => Ok(DateTruncPrecision::Quarter), - date_trunc_keyword::MONTH => Ok(DateTruncPrecision::Month), - date_trunc_keyword::WEEK => Ok(DateTruncPrecision::Week), - date_trunc_keyword::DAY => Ok(DateTruncPrecision::Day), - date_trunc_keyword::HOUR => Ok(DateTruncPrecision::Hour), - date_trunc_keyword::MINUTE => Ok(DateTruncPrecision::Minute), - date_trunc_keyword::SECOND => Ok(DateTruncPrecision::Second), - _ => Err(format!("'{value}' is not a valid DateTruncPrecision")), - } - } -} diff --git a/src/sql/common/debezium.rs b/src/sql/common/debezium.rs deleted file mode 100644 index 9dbc401f..00000000 --- a/src/sql/common/debezium.rs +++ /dev/null @@ -1,148 +0,0 @@ -// 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 bincode::{Decode, Encode}; -use serde::{Deserialize, Serialize}; -use std::convert::TryFrom; -use std::fmt::Debug; - -pub trait Key: - Debug + Clone + Encode + Decode<()> + std::hash::Hash + PartialEq + Eq + Send + 'static -{ -} -impl + std::hash::Hash + PartialEq + Eq + Send + 'static> Key - for T -{ -} - -pub trait Data: Debug + Clone + Encode + Decode<()> + Send + 'static {} -impl + Send + 'static> Data for T {} - -#[derive(Debug, Clone, PartialEq, Encode, Decode, Serialize, Deserialize)] -pub enum UpdatingData { - Retract(T), - Update { old: T, new: T }, - Append(T), -} - -impl UpdatingData { - pub fn lower(&self) -> T { - match self { - UpdatingData::Retract(_) => panic!("cannot lower retractions"), - UpdatingData::Update { new, .. } => new.clone(), - UpdatingData::Append(t) => t.clone(), - } - } - - pub fn unwrap_append(&self) -> &T { - match self { - UpdatingData::Append(t) => t, - _ => panic!("UpdatingData is not an append"), - } - } -} - -#[derive(Clone, Encode, Decode, Debug, Serialize, Deserialize, PartialEq)] -#[serde(try_from = "DebeziumShadow")] -pub struct Debezium { - pub before: Option, - pub after: Option, - pub op: DebeziumOp, -} - -#[derive(Clone, Encode, Decode, Debug, Serialize, Deserialize, PartialEq)] -struct DebeziumShadow { - before: Option, - after: Option, - op: DebeziumOp, -} - -impl TryFrom> for Debezium { - type Error = &'static str; - - fn try_from(value: DebeziumShadow) -> Result { - match (value.op, &value.before, &value.after) { - (DebeziumOp::Create, _, None) => { - Err("`after` must be set for Debezium create messages") - } - (DebeziumOp::Update, None, _) => { - Err("`before` must be set for Debezium update messages") - } - (DebeziumOp::Update, _, None) => { - Err("`after` must be set for Debezium update messages") - } - (DebeziumOp::Delete, None, _) => { - Err("`before` must be set for Debezium delete messages") - } - _ => Ok(Debezium { - before: value.before, - after: value.after, - op: value.op, - }), - } - } -} - -#[derive(Copy, Clone, Encode, Decode, Debug, PartialEq)] -pub enum DebeziumOp { - Create, - Update, - Delete, -} - -#[allow(clippy::to_string_trait_impl)] -impl ToString for DebeziumOp { - fn to_string(&self) -> String { - match self { - DebeziumOp::Create => "c", - DebeziumOp::Update => "u", - DebeziumOp::Delete => "d", - } - .to_string() - } -} - -impl<'de> Deserialize<'de> for DebeziumOp { - fn deserialize(deserializer: D) -> Result - where - D: serde::Deserializer<'de>, - { - let s = String::deserialize(deserializer)?; - match s.as_str() { - "c" | "r" => Ok(DebeziumOp::Create), - "u" => Ok(DebeziumOp::Update), - "d" => Ok(DebeziumOp::Delete), - _ => Err(serde::de::Error::custom(format!("Invalid DebeziumOp {s}"))), - } - } -} - -impl Serialize for DebeziumOp { - fn serialize(&self, serializer: S) -> Result - where - S: serde::Serializer, - { - match self { - DebeziumOp::Create => serializer.serialize_str("c"), - DebeziumOp::Update => serializer.serialize_str("u"), - DebeziumOp::Delete => serializer.serialize_str("d"), - } - } -} - -#[derive(Copy, Clone, Encode, Decode, Debug, PartialEq, Serialize, Deserialize)] -pub enum JoinType { - Inner, - Left, - Right, - Full, -} diff --git a/src/sql/common/errors.rs b/src/sql/common/errors.rs deleted file mode 100644 index fa4a722e..00000000 --- a/src/sql/common/errors.rs +++ /dev/null @@ -1,92 +0,0 @@ -// 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::fmt; - -/// Result type for streaming operators and collectors. -pub type DataflowResult = std::result::Result; - -/// Unified error type for streaming dataflow operations. -#[derive(Debug)] -pub enum DataflowError { - Arrow(arrow_schema::ArrowError), - DataFusion(datafusion::error::DataFusionError), - Operator(String), - State(String), - Connector(String), - Internal(String), -} - -impl fmt::Display for DataflowError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - DataflowError::Arrow(e) => write!(f, "Arrow error: {e}"), - DataflowError::DataFusion(e) => write!(f, "DataFusion error: {e}"), - DataflowError::Operator(msg) => write!(f, "Operator error: {msg}"), - DataflowError::State(msg) => write!(f, "State error: {msg}"), - DataflowError::Connector(msg) => write!(f, "Connector error: {msg}"), - DataflowError::Internal(msg) => write!(f, "Internal error: {msg}"), - } - } -} - -impl std::error::Error for DataflowError {} - -impl DataflowError { - pub fn with_operator(self, operator_id: impl Into) -> Self { - let id = operator_id.into(); - match self { - DataflowError::Operator(m) => DataflowError::Operator(format!("{id}: {m}")), - other => DataflowError::Operator(format!("{id}: {other}")), - } - } -} - -impl From for DataflowError { - fn from(e: arrow_schema::ArrowError) -> Self { - DataflowError::Arrow(e) - } -} - -impl From for DataflowError { - fn from(e: datafusion::error::DataFusionError) -> Self { - DataflowError::DataFusion(e) - } -} - -/// Macro for creating connector errors. -#[macro_export] -macro_rules! connector_err { - ($($arg:tt)*) => { - $crate::sql::common::errors::DataflowError::Connector(format!($($arg)*)) - }; -} - -/// State-related errors. -#[derive(Debug)] -pub enum StateError { - KeyNotFound(String), - SerializationError(String), - BackendError(String), -} - -impl fmt::Display for StateError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - StateError::KeyNotFound(key) => write!(f, "Key not found: {key}"), - StateError::SerializationError(msg) => write!(f, "Serialization error: {msg}"), - StateError::BackendError(msg) => write!(f, "State backend error: {msg}"), - } - } -} - -impl std::error::Error for StateError {} diff --git a/src/sql/common/format_from_opts.rs b/src/sql/common/format_from_opts.rs deleted file mode 100644 index ffd29572..00000000 --- a/src/sql/common/format_from_opts.rs +++ /dev/null @@ -1,182 +0,0 @@ -// 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 `WITH` clause format / framing / bad-data options (Arroyo-compatible keys). - -use std::str::FromStr; - -use datafusion::common::{Result as DFResult, plan_datafusion_err, plan_err}; - -use super::connector_options::ConnectorOptions; -use super::constants::{bad_data_value, connection_format_value, framing_method_value}; -use super::formats::{ - AvroFormat, BadData, CsvFormat, DecimalEncoding, Format, Framing, JsonCompression, JsonFormat, - LanceFormat, NewlineDelimitedFraming, ParquetCompression, ParquetFormat, ProtobufFormat, - RawBytesFormat, RawStringFormat, TimestampFormat, -}; -use super::with_option_keys as opt; - -impl JsonFormat { - pub fn from_opts(opts: &mut ConnectorOptions) -> DFResult { - let mut j = JsonFormat::default(); - if let Some(v) = opts.pull_opt_bool(opt::JSON_CONFLUENT_SCHEMA_REGISTRY)? { - j.confluent_schema_registry = v; - } - if let Some(v) = opts.pull_opt_u64(opt::JSON_CONFLUENT_SCHEMA_VERSION)? { - j.schema_id = Some(v as u32); - } - if let Some(v) = opts.pull_opt_bool(opt::JSON_INCLUDE_SCHEMA)? { - j.include_schema = v; - } - if let Some(v) = opts.pull_opt_bool(opt::JSON_DEBEZIUM)? { - j.debezium = v; - } - if let Some(v) = opts.pull_opt_bool(opt::JSON_UNSTRUCTURED)? { - j.unstructured = v; - } - if let Some(s) = opts.pull_opt_str(opt::JSON_TIMESTAMP_FORMAT)? { - j.timestamp_format = TimestampFormat::try_from(s.as_str()) - .map_err(|_| plan_datafusion_err!("invalid json.timestamp_format '{}'", s))?; - } - if let Some(s) = opts.pull_opt_str(opt::JSON_DECIMAL_ENCODING)? { - j.decimal_encoding = DecimalEncoding::try_from(s.as_str()) - .map_err(|_| plan_datafusion_err!("invalid json.decimal_encoding '{s}'"))?; - } - if let Some(s) = opts.pull_opt_str(opt::JSON_COMPRESSION)? { - j.compression = JsonCompression::from_str(&s) - .map_err(|e| plan_datafusion_err!("invalid json.compression: {e}"))?; - } - Ok(j) - } -} - -impl Format { - pub fn from_opts(opts: &mut ConnectorOptions) -> DFResult> { - let Some(name) = opts.peek_opt_str(opt::FORMAT)? else { - return Ok(None); - }; - let n = name.to_lowercase(); - match n.as_str() { - connection_format_value::JSON => Ok(Some(Format::Json(JsonFormat::from_opts(opts)?))), - connection_format_value::CSV => Ok(Some(Format::Csv(CsvFormat {}))), - connection_format_value::LANCE => Ok(Some(Format::Lance(LanceFormat {}))), - connection_format_value::DEBEZIUM_JSON => { - let mut j = JsonFormat::from_opts(opts)?; - j.debezium = true; - Ok(Some(Format::Json(j))) - } - connection_format_value::AVRO => Ok(Some(Format::Avro(AvroFormat::from_opts(opts)?))), - connection_format_value::PARQUET => { - Ok(Some(Format::Parquet(ParquetFormat::from_opts(opts)?))) - } - connection_format_value::PROTOBUF => { - Ok(Some(Format::Protobuf(ProtobufFormat::from_opts(opts)?))) - } - connection_format_value::RAW_STRING => Ok(Some(Format::RawString(RawStringFormat {}))), - connection_format_value::RAW_BYTES => Ok(Some(Format::RawBytes(RawBytesFormat {}))), - _ => plan_err!("unknown format '{name}'"), - } - } -} - -impl AvroFormat { - fn from_opts(opts: &mut ConnectorOptions) -> DFResult { - let mut a = AvroFormat { - confluent_schema_registry: false, - raw_datums: false, - into_unstructured_json: false, - schema_id: None, - }; - if let Some(v) = opts.pull_opt_bool(opt::AVRO_CONFLUENT_SCHEMA_REGISTRY)? { - a.confluent_schema_registry = v; - } - if let Some(v) = opts.pull_opt_bool(opt::AVRO_RAW_DATUMS)? { - a.raw_datums = v; - } - if let Some(v) = opts.pull_opt_bool(opt::AVRO_INTO_UNSTRUCTURED_JSON)? { - a.into_unstructured_json = v; - } - if let Some(v) = opts.pull_opt_u64(opt::AVRO_SCHEMA_ID)? { - a.schema_id = Some(v as u32); - } - Ok(a) - } -} - -impl ParquetFormat { - fn from_opts(opts: &mut ConnectorOptions) -> DFResult { - let mut p = ParquetFormat::default(); - if let Some(s) = opts.pull_opt_str(opt::PARQUET_COMPRESSION)? { - p.compression = ParquetCompression::from_str(&s) - .map_err(|e| plan_datafusion_err!("invalid parquet.compression: {e}"))?; - } - if let Some(v) = opts.pull_opt_u64(opt::PARQUET_ROW_GROUP_BYTES)? { - p.row_group_bytes = Some(v); - } - Ok(p) - } -} - -impl ProtobufFormat { - fn from_opts(opts: &mut ConnectorOptions) -> DFResult { - let mut p = ProtobufFormat { - into_unstructured_json: false, - message_name: None, - compiled_schema: None, - confluent_schema_registry: false, - length_delimited: false, - }; - if let Some(v) = opts.pull_opt_bool(opt::PROTOBUF_INTO_UNSTRUCTURED_JSON)? { - p.into_unstructured_json = v; - } - if let Some(s) = opts.pull_opt_str(opt::PROTOBUF_MESSAGE_NAME)? { - p.message_name = Some(s); - } - if let Some(v) = opts.pull_opt_bool(opt::PROTOBUF_CONFLUENT_SCHEMA_REGISTRY)? { - p.confluent_schema_registry = v; - } - if let Some(v) = opts.pull_opt_bool(opt::PROTOBUF_LENGTH_DELIMITED)? { - p.length_delimited = v; - } - Ok(p) - } -} - -impl Framing { - pub fn from_opts(opts: &mut ConnectorOptions) -> DFResult> { - let method = opts.pull_opt_str(opt::FRAMING_METHOD)?; - match method.as_deref() { - None => Ok(None), - Some(framing_method_value::NEWLINE) | Some(framing_method_value::NEWLINE_DELIMITED) => { - let max = opts.pull_opt_u64(opt::FRAMING_MAX_LINE_LENGTH)?; - Ok(Some(Framing::Newline(NewlineDelimitedFraming { - max_line_length: max, - }))) - } - Some(other) => plan_err!("unknown framing.method '{other}'"), - } - } -} - -impl BadData { - pub fn from_opts(opts: &mut ConnectorOptions) -> DFResult { - let Some(s) = opts.pull_opt_str(opt::BAD_DATA)? else { - return Ok(BadData::Fail {}); - }; - let v = s.to_lowercase(); - match v.as_str() { - bad_data_value::FAIL => Ok(BadData::Fail {}), - bad_data_value::DROP => Ok(BadData::Drop {}), - _ => plan_err!("invalid bad_data '{s}'"), - } - } -} diff --git a/src/sql/common/formats.rs b/src/sql/common/formats.rs deleted file mode 100644 index a47d93cf..00000000 --- a/src/sql/common/formats.rs +++ /dev/null @@ -1,267 +0,0 @@ -// 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 serde::{Deserialize, Serialize}; -use std::convert::TryFrom; -use std::fmt::{Display, Formatter}; -use std::str::FromStr; - -use super::constants::{ - connection_format_value, decimal_encoding_value, json_compression_value, - parquet_compression_value, timestamp_format_value, -}; - -#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Default, Hash, PartialOrd)] -#[serde(rename_all = "snake_case")] -pub enum TimestampFormat { - #[default] - #[serde(rename = "rfc3339")] - RFC3339, - UnixMillis, -} - -impl TryFrom<&str> for TimestampFormat { - type Error = (); - - fn try_from(value: &str) -> Result { - match value { - timestamp_format_value::RFC3339_UPPER | timestamp_format_value::RFC3339_SNAKE => { - Ok(TimestampFormat::RFC3339) - } - timestamp_format_value::UNIX_MILLIS_PASCAL - | timestamp_format_value::UNIX_MILLIS_SNAKE => Ok(TimestampFormat::UnixMillis), - _ => Err(()), - } - } -} - -#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Default, Hash, PartialOrd)] -#[serde(rename_all = "snake_case")] -pub enum DecimalEncoding { - #[default] - Number, - String, - Bytes, -} - -impl TryFrom<&str> for DecimalEncoding { - type Error = (); - - fn try_from(s: &str) -> Result { - match s { - decimal_encoding_value::NUMBER => Ok(Self::Number), - decimal_encoding_value::STRING => Ok(Self::String), - decimal_encoding_value::BYTES => Ok(Self::Bytes), - _ => Err(()), - } - } -} - -#[derive(Serialize, Deserialize, Default, Copy, Clone, Debug, PartialEq, Eq, Hash, PartialOrd)] -#[serde(rename_all = "snake_case")] -pub enum JsonCompression { - #[default] - Uncompressed, - Gzip, -} - -impl FromStr for JsonCompression { - type Err = String; - - fn from_str(s: &str) -> Result { - match s { - json_compression_value::UNCOMPRESSED => Ok(JsonCompression::Uncompressed), - json_compression_value::GZIP => Ok(JsonCompression::Gzip), - _ => Err(format!("invalid json compression '{s}'")), - } - } -} - -#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Default, Hash, PartialOrd)] -#[serde(rename_all = "snake_case")] -pub struct JsonFormat { - #[serde(default)] - pub confluent_schema_registry: bool, - #[serde(default, alias = "confluent_schema_version")] - pub schema_id: Option, - #[serde(default)] - pub include_schema: bool, - #[serde(default)] - pub debezium: bool, - #[serde(default)] - pub unstructured: bool, - #[serde(default)] - pub timestamp_format: TimestampFormat, - #[serde(default)] - pub decimal_encoding: DecimalEncoding, - #[serde(default)] - pub compression: JsonCompression, -} - -#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash, PartialOrd)] -#[serde(rename_all = "snake_case")] -pub struct RawStringFormat {} - -#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash, PartialOrd)] -#[serde(rename_all = "snake_case")] -pub struct RawBytesFormat {} - -#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash, PartialOrd)] -#[serde(rename_all = "snake_case")] -pub struct CsvFormat {} - -#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash, PartialOrd)] -#[serde(rename_all = "snake_case")] -pub struct LanceFormat {} - -#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash, PartialOrd)] -#[serde(rename_all = "snake_case")] -pub struct AvroFormat { - #[serde(default)] - pub confluent_schema_registry: bool, - #[serde(default)] - pub raw_datums: bool, - #[serde(default)] - pub into_unstructured_json: bool, - #[serde(default)] - pub schema_id: Option, -} - -impl AvroFormat { - pub fn new( - confluent_schema_registry: bool, - raw_datums: bool, - into_unstructured_json: bool, - ) -> Self { - Self { - confluent_schema_registry, - raw_datums, - into_unstructured_json, - schema_id: None, - } - } -} - -#[derive(Serialize, Deserialize, Copy, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Default)] -#[serde(rename_all = "snake_case")] -pub enum ParquetCompression { - Uncompressed, - Snappy, - Gzip, - #[default] - Zstd, - Lz4, - Lz4Raw, -} - -impl FromStr for ParquetCompression { - type Err = String; - - fn from_str(s: &str) -> Result { - match s { - parquet_compression_value::UNCOMPRESSED => Ok(ParquetCompression::Uncompressed), - parquet_compression_value::SNAPPY => Ok(ParquetCompression::Snappy), - parquet_compression_value::GZIP => Ok(ParquetCompression::Gzip), - parquet_compression_value::ZSTD => Ok(ParquetCompression::Zstd), - parquet_compression_value::LZ4 => Ok(ParquetCompression::Lz4), - parquet_compression_value::LZ4_RAW => Ok(ParquetCompression::Lz4Raw), - _ => Err(format!("invalid parquet compression '{s}'")), - } - } -} - -#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash, PartialOrd, Default)] -#[serde(rename_all = "snake_case")] -pub struct ParquetFormat { - #[serde(default)] - pub compression: ParquetCompression, - #[serde(default)] - pub row_group_bytes: Option, -} - -#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash, PartialOrd)] -#[serde(rename_all = "snake_case")] -pub struct ProtobufFormat { - #[serde(default)] - pub into_unstructured_json: bool, - #[serde(default)] - pub message_name: Option, - #[serde(default)] - pub compiled_schema: Option>, - #[serde(default)] - pub confluent_schema_registry: bool, - #[serde(default)] - pub length_delimited: bool, -} - -#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, Hash, PartialOrd)] -#[serde(rename_all = "snake_case", tag = "type")] -pub enum Format { - Json(JsonFormat), - Csv(CsvFormat), - Lance(LanceFormat), - Avro(AvroFormat), - Protobuf(ProtobufFormat), - Parquet(ParquetFormat), - RawString(RawStringFormat), - RawBytes(RawBytesFormat), -} - -impl Display for Format { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - f.write_str(self.name()) - } -} - -impl Format { - pub fn name(&self) -> &'static str { - match self { - Format::Json(_) => connection_format_value::JSON, - Format::Csv(_) => connection_format_value::CSV, - Format::Lance(_) => connection_format_value::LANCE, - Format::Avro(_) => connection_format_value::AVRO, - Format::Protobuf(_) => connection_format_value::PROTOBUF, - Format::Parquet(_) => connection_format_value::PARQUET, - Format::RawString(_) => connection_format_value::RAW_STRING, - Format::RawBytes(_) => connection_format_value::RAW_BYTES, - } - } - - pub fn is_updating(&self) -> bool { - matches!(self, Format::Json(JsonFormat { debezium: true, .. })) - } -} - -#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, Hash, PartialOrd)] -#[serde(rename_all = "snake_case", tag = "behavior")] -pub enum BadData { - Fail {}, - Drop {}, -} - -impl Default for BadData { - fn default() -> Self { - BadData::Fail {} - } -} - -#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, Hash, PartialOrd)] -#[serde(rename_all = "snake_case", tag = "method")] -pub enum Framing { - Newline(NewlineDelimitedFraming), -} - -#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, Hash, PartialOrd)] -#[serde(rename_all = "snake_case")] -pub struct NewlineDelimitedFraming { - pub max_line_length: Option, -} diff --git a/src/sql/common/fs_schema.rs b/src/sql/common/fs_schema.rs deleted file mode 100644 index 76a08537..00000000 --- a/src/sql/common/fs_schema.rs +++ /dev/null @@ -1,470 +0,0 @@ -// 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 super::{TIMESTAMP_FIELD, to_nanos}; -use crate::sql::common::converter::Converter; -use arrow::compute::kernels::cmp::gt_eq; -use arrow::compute::kernels::numeric::div; -use arrow::compute::{SortColumn, filter_record_batch, lexsort_to_indices, partition, take}; -use arrow::row::SortField; -use arrow_array::types::UInt64Type; -use arrow_array::{PrimitiveArray, UInt64Array}; -use datafusion::arrow::array::builder::{ArrayBuilder, make_builder}; -use datafusion::arrow::array::{RecordBatch, TimestampNanosecondArray}; -use datafusion::arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaBuilder, TimeUnit}; -use datafusion::arrow::error::ArrowError; -use datafusion::common::{DataFusionError, Result as DFResult}; -use protocol::function_stream_graph; -use serde::{Deserialize, Serialize}; -use std::ops::Range; -use std::sync::Arc; -use std::time::SystemTime; - -#[derive(Debug, Copy, Clone)] -pub enum FieldValueType<'a> { - Int64(Option), - UInt64(Option), - Int32(Option), - String(Option<&'a str>), - Bytes(Option<&'a [u8]>), -} - -pub type FsSchemaRef = Arc; - -#[derive(Debug, Clone, Eq, PartialEq, Hash, Serialize, Deserialize)] -pub struct FsSchema { - pub schema: Arc, - pub timestamp_index: usize, - key_indices: Option>, - /// If defined, these indices are used for routing (i.e., which subtask gets which piece of data) - routing_key_indices: Option>, -} - -impl TryFrom for FsSchema { - type Error = DataFusionError; - fn try_from(schema_proto: function_stream_graph::FsSchema) -> Result { - let schema: Schema = serde_json::from_str(&schema_proto.arrow_schema) - .map_err(|e| DataFusionError::Plan(format!("Invalid arrow schema: {e}")))?; - let timestamp_index = schema_proto.timestamp_index as usize; - - let key_indices = schema_proto.has_keys.then(|| { - schema_proto - .key_indices - .into_iter() - .map(|index| index as usize) - .collect() - }); - - let routing_key_indices = schema_proto.has_routing_keys.then(|| { - schema_proto - .routing_key_indices - .into_iter() - .map(|index| index as usize) - .collect() - }); - - Ok(Self { - schema: Arc::new(schema), - timestamp_index, - key_indices, - routing_key_indices, - }) - } -} - -impl From for function_stream_graph::FsSchema { - fn from(schema: FsSchema) -> Self { - let arrow_schema = serde_json::to_string(schema.schema.as_ref()).unwrap(); - let timestamp_index = schema.timestamp_index as u32; - - let has_keys = schema.key_indices.is_some(); - let key_indices = schema - .key_indices - .map(|ks| ks.into_iter().map(|index| index as u32).collect()) - .unwrap_or_default(); - - let has_routing_keys = schema.routing_key_indices.is_some(); - let routing_key_indices = schema - .routing_key_indices - .map(|ks| ks.into_iter().map(|index| index as u32).collect()) - .unwrap_or_default(); - - Self { - arrow_schema, - timestamp_index, - key_indices, - has_keys, - routing_key_indices, - has_routing_keys, - } - } -} - -impl FsSchema { - pub fn new( - schema: Arc, - timestamp_index: usize, - key_indices: Option>, - routing_key_indices: Option>, - ) -> Self { - Self { - schema, - timestamp_index, - key_indices, - routing_key_indices, - } - } - pub fn new_unkeyed(schema: Arc, timestamp_index: usize) -> Self { - Self { - schema, - timestamp_index, - key_indices: None, - routing_key_indices: None, - } - } - pub fn new_keyed(schema: Arc, timestamp_index: usize, key_indices: Vec) -> Self { - Self { - schema, - timestamp_index, - key_indices: Some(key_indices), - routing_key_indices: None, - } - } - - pub fn from_fields(mut fields: Vec) -> Self { - if !fields.iter().any(|f| f.name() == TIMESTAMP_FIELD) { - fields.push(Field::new( - TIMESTAMP_FIELD, - DataType::Timestamp(TimeUnit::Nanosecond, None), - false, - )); - } - - Self::from_schema_keys(Arc::new(Schema::new(fields)), vec![]).unwrap() - } - - pub fn from_schema_unkeyed(schema: Arc) -> DFResult { - let timestamp_index = schema - .column_with_name(TIMESTAMP_FIELD) - .ok_or_else(|| { - DataFusionError::Plan(format!( - "no {TIMESTAMP_FIELD} field in schema, schema is {schema:?}" - )) - })? - .0; - - Ok(Self { - schema, - timestamp_index, - key_indices: None, - routing_key_indices: None, - }) - } - - pub fn from_schema_keys(schema: Arc, key_indices: Vec) -> DFResult { - let timestamp_index = schema - .column_with_name(TIMESTAMP_FIELD) - .ok_or_else(|| { - DataFusionError::Plan(format!( - "no {TIMESTAMP_FIELD} field in schema, schema is {schema:?}" - )) - })? - .0; - - Ok(Self { - schema, - timestamp_index, - key_indices: Some(key_indices), - routing_key_indices: None, - }) - } - - pub fn schema_without_timestamp(&self) -> Schema { - let mut builder = SchemaBuilder::from(self.schema.fields()); - builder.remove(self.timestamp_index); - builder.finish() - } - - pub fn remove_timestamp_column(&self, batch: &mut RecordBatch) { - batch.remove_column(self.timestamp_index); - } - - pub fn builders(&self) -> Vec> { - self.schema - .fields - .iter() - .map(|f| make_builder(f.data_type(), 8)) - .collect() - } - - pub fn timestamp_column<'a>(&self, batch: &'a RecordBatch) -> &'a TimestampNanosecondArray { - batch - .column(self.timestamp_index) - .as_any() - .downcast_ref::() - .unwrap() - } - - pub fn has_routing_keys(&self) -> bool { - self.routing_keys().map(|k| !k.is_empty()).unwrap_or(false) - } - - pub fn routing_keys(&self) -> Option<&Vec> { - self.routing_key_indices - .as_ref() - .or(self.key_indices.as_ref()) - } - - pub fn storage_keys(&self) -> Option<&Vec> { - self.key_indices.as_ref() - } - - pub fn clone_storage_key_indices(&self) -> Option> { - self.key_indices.clone() - } - - pub fn clone_routing_key_indices(&self) -> Option> { - self.routing_key_indices.clone() - } - - pub fn filter_by_time( - &self, - batch: RecordBatch, - cutoff: Option, - ) -> Result { - let Some(cutoff) = cutoff else { - // no watermark, so we just return the same batch. - return Ok(batch); - }; - // filter out late data - let timestamp_column = batch - .column(self.timestamp_index) - .as_any() - .downcast_ref::() - .ok_or_else(|| ArrowError::CastError( - format!("failed to downcast column {} of {:?} to timestamp. Schema is supposed to be {:?}", - self.timestamp_index, batch, self.schema)))?; - let cutoff_scalar = TimestampNanosecondArray::new_scalar(to_nanos(cutoff) as i64); - let on_time = gt_eq(timestamp_column, &cutoff_scalar)?; - filter_record_batch(&batch, &on_time) - } - - pub fn sort_columns(&self, batch: &RecordBatch, with_timestamp: bool) -> Vec { - let mut columns = vec![]; - if let Some(keys) = &self.key_indices { - columns.extend(keys.iter().map(|index| SortColumn { - values: batch.column(*index).clone(), - options: None, - })); - } - if with_timestamp { - columns.push(SortColumn { - values: batch.column(self.timestamp_index).clone(), - options: None, - }); - } - columns - } - - pub fn sort_fields(&self, with_timestamp: bool) -> Vec { - let mut sort_fields = vec![]; - if let Some(keys) = &self.key_indices { - sort_fields.extend(keys.iter()); - } - if with_timestamp { - sort_fields.push(self.timestamp_index); - } - self.sort_fields_by_indices(&sort_fields) - } - - fn sort_fields_by_indices(&self, indices: &[usize]) -> Vec { - indices - .iter() - .map(|index| SortField::new(self.schema.field(*index).data_type().clone())) - .collect() - } - - pub fn converter(&self, with_timestamp: bool) -> Result { - Converter::new(self.sort_fields(with_timestamp)) - } - - pub fn value_converter( - &self, - with_timestamp: bool, - generation_index: usize, - ) -> Result { - match &self.key_indices { - None => { - let mut indices = (0..self.schema.fields().len()).collect::>(); - indices.remove(generation_index); - if !with_timestamp { - indices.remove(self.timestamp_index); - } - Converter::new(self.sort_fields_by_indices(&indices)) - } - Some(keys) => { - let indices = (0..self.schema.fields().len()) - .filter(|index| { - !keys.contains(index) - && (with_timestamp || *index != self.timestamp_index) - && *index != generation_index - }) - .collect::>(); - Converter::new(self.sort_fields_by_indices(&indices)) - } - } - } - - pub fn value_indices(&self, with_timestamp: bool) -> Vec { - let field_count = self.schema.fields().len(); - match &self.key_indices { - None => { - let mut indices = (0..field_count).collect::>(); - - if !with_timestamp { - indices.remove(self.timestamp_index); - } - indices - } - Some(keys) => (0..field_count) - .filter(|index| { - !keys.contains(index) && (with_timestamp || *index != self.timestamp_index) - }) - .collect::>(), - } - } - - pub fn sort( - &self, - batch: RecordBatch, - with_timestamp: bool, - ) -> Result { - if self.key_indices.is_none() && !with_timestamp { - return Ok(batch); - } - let sort_columns = self.sort_columns(&batch, with_timestamp); - let sort_indices = lexsort_to_indices(&sort_columns, None).expect("should be able to sort"); - let columns = batch - .columns() - .iter() - .map(|c| take(c, &sort_indices, None).unwrap()) - .collect(); - - RecordBatch::try_new(batch.schema(), columns) - } - - pub fn partition( - &self, - batch: &RecordBatch, - with_timestamp: bool, - ) -> Result>, ArrowError> { - if self.key_indices.is_none() && !with_timestamp { - #[allow(clippy::single_range_in_vec_init)] - return Ok(vec![0..batch.num_rows()]); - } - - let mut partition_columns = vec![]; - - if let Some(keys) = &self.routing_keys() { - partition_columns.extend(keys.iter().map(|index| batch.column(*index).clone())); - } - if with_timestamp { - partition_columns.push(batch.column(self.timestamp_index).clone()); - } - - Ok(partition(&partition_columns)?.ranges()) - } - - pub fn unkeyed_batch(&self, batch: &RecordBatch) -> Result { - if self.key_indices.is_none() { - return Ok(batch.clone()); - } - let columns: Vec<_> = (0..batch.num_columns()) - .filter(|index| !self.key_indices.as_ref().unwrap().contains(index)) - .collect(); - batch.project(&columns) - } - - pub fn schema_without_keys(&self) -> Result { - if self.key_indices.is_none() { - return Ok(self.clone()); - } - let key_indices = self.key_indices.as_ref().unwrap(); - let unkeyed_schema = Schema::new( - self.schema - .fields() - .iter() - .enumerate() - .filter(|(index, _field)| !key_indices.contains(index)) - .map(|(_, field)| field.as_ref().clone()) - .collect::>(), - ); - let timestamp_index = unkeyed_schema.index_of(TIMESTAMP_FIELD)?; - Ok(Self { - schema: Arc::new(unkeyed_schema), - timestamp_index, - key_indices: None, - routing_key_indices: None, - }) - } - - pub fn with_fields(&self, fields: Vec) -> Result { - let schema = Arc::new(Schema::new_with_metadata( - fields, - self.schema.metadata.clone(), - )); - - let timestamp_index = schema.index_of(TIMESTAMP_FIELD)?; - let max_index = *[&self.key_indices, &self.routing_key_indices] - .iter() - .map(|indices| indices.as_ref().and_then(|k| k.iter().max())) - .max() - .flatten() - .unwrap_or(&0); - - if schema.fields.len() - 1 < max_index { - return Err(ArrowError::InvalidArgumentError(format!( - "expected at least {} fields, but were only {}", - max_index + 1, - schema.fields.len() - ))); - } - - Ok(Self { - schema, - timestamp_index, - key_indices: self.key_indices.clone(), - routing_key_indices: self.routing_key_indices.clone(), - }) - } - - pub fn with_additional_fields( - &self, - new_fields: impl Iterator, - ) -> Result { - let mut fields = self.schema.fields.to_vec(); - fields.extend(new_fields.map(Arc::new)); - - self.with_fields(fields) - } -} - -pub fn server_for_hash_array( - hash: &PrimitiveArray, - n: usize, -) -> Result, ArrowError> { - let range_size = u64::MAX / (n as u64) + 1; - let range_scalar = UInt64Array::new_scalar(range_size); - let division = div(hash, &range_scalar)?; - let result: &PrimitiveArray = division.as_any().downcast_ref().unwrap(); - Ok(result.clone()) -} diff --git a/src/sql/common/kafka_catalog.rs b/src/sql/common/kafka_catalog.rs deleted file mode 100644 index 51ceee67..00000000 --- a/src/sql/common/kafka_catalog.rs +++ /dev/null @@ -1,116 +0,0 @@ -// 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 serde::{Deserialize, Serialize}; -use std::collections::HashMap; - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] -pub struct KafkaTable { - pub topic: String, - #[serde(flatten)] - pub kind: TableType, - #[serde(default)] - pub client_configs: HashMap, - pub value_subject: Option, -} - -impl KafkaTable { - pub fn subject(&self) -> String { - self.value_subject - .clone() - .unwrap_or_else(|| format!("{}-value", self.topic)) - } -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] -#[serde(tag = "type", rename_all = "snake_case")] -pub enum TableType { - Source { - offset: KafkaTableSourceOffset, - read_mode: Option, - group_id: Option, - group_id_prefix: Option, - }, - Sink { - commit_mode: SinkCommitMode, - key_field: Option, - timestamp_field: Option, - }, -} - -#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)] -#[serde(rename_all = "snake_case")] -pub enum KafkaTableSourceOffset { - Latest, - Earliest, - #[default] - Group, -} - -#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "snake_case")] -pub enum ReadMode { - ReadUncommitted, - ReadCommitted, -} - -#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)] -#[serde(rename_all = "snake_case")] -pub enum SinkCommitMode { - #[default] - AtLeastOnce, - ExactlyOnce, -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] -#[serde(rename_all = "camelCase")] -pub struct KafkaConfig { - pub bootstrap_servers: String, - #[serde(default)] - pub authentication: KafkaConfigAuthentication, - #[serde(default)] - pub schema_registry_enum: Option, - #[serde(default)] - pub connection_properties: HashMap, -} - -#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)] -#[serde(tag = "type")] -pub enum KafkaConfigAuthentication { - #[default] - #[serde(rename = "None")] - None, - #[serde(rename = "AWS_MSK_IAM")] - AwsMskIam { region: String }, - #[serde(rename = "SASL")] - Sasl { - protocol: String, - mechanism: String, - username: String, - password: String, - }, -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] -#[serde(tag = "type")] -pub enum SchemaRegistryConfig { - #[serde(rename = "None")] - None, - #[serde(rename = "Confluent Schema Registry")] - ConfluentSchemaRegistry { - endpoint: String, - #[serde(rename = "apiKey")] - api_key: Option, - #[serde(rename = "apiSecret")] - api_secret: Option, - }, -} diff --git a/src/sql/common/mod.rs b/src/sql/common/mod.rs deleted file mode 100644 index af44cb0f..00000000 --- a/src/sql/common/mod.rs +++ /dev/null @@ -1,65 +0,0 @@ -// 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. - -//! Shared core types and constants for FunctionStream (`crate::sql::common`). -//! -//! Used by the runtime, SQL planner, coordinator, and other subsystems — -//! analogous to `arroyo-types` + `arroyo-rpc` in Arroyo. - -pub mod arrow_ext; -pub mod connector_options; -pub mod constants; -pub mod control; -pub mod converter; -pub mod date; -pub mod debezium; -pub mod errors; -pub mod format_from_opts; -pub mod formats; -pub mod fs_schema; -pub mod kafka_catalog; -pub mod operator_config; -pub mod time_utils; -pub mod topology; -pub mod with_option_keys; - -// ── Re-exports from existing modules ── -pub use crate::runtime::streaming::protocol::{CheckpointBarrier, Watermark}; -pub use arrow_ext::FsExtensionType; -pub use time_utils::{from_nanos, to_micros, to_millis, to_nanos}; - -// ── Re-exports from new modules ── -pub use connector_options::ConnectorOptions; -pub use formats::{BadData, Format, Framing, JsonCompression, JsonFormat}; -pub use fs_schema::{FsSchema, FsSchemaRef}; -pub use operator_config::MetadataField; - -// ── Well-known column names ── -pub use constants::sql_field::{TIMESTAMP_FIELD, UPDATING_META_FIELD}; -pub use topology::render_program_topology; - -// ── Environment variables ── -pub const JOB_ID_ENV: &str = "JOB_ID"; -pub const RUN_ID_ENV: &str = "RUN_ID"; - -// ── Metric names ── -pub const MESSAGES_RECV: &str = "fs_worker_messages_recv"; -pub const MESSAGES_SENT: &str = "fs_worker_messages_sent"; -pub const BYTES_RECV: &str = "fs_worker_bytes_recv"; -pub const BYTES_SENT: &str = "fs_worker_bytes_sent"; -pub const BATCHES_RECV: &str = "fs_worker_batches_recv"; -pub const BATCHES_SENT: &str = "fs_worker_batches_sent"; -pub const TX_QUEUE_SIZE: &str = "fs_worker_tx_queue_size"; -pub const TX_QUEUE_REM: &str = "fs_worker_tx_queue_rem"; -pub const DESERIALIZATION_ERRORS: &str = "fs_worker_deserialization_errors"; - -pub const LOOKUP_KEY_INDEX_FIELD: &str = "__lookup_key_index"; diff --git a/src/sql/common/operator_config.rs b/src/sql/common/operator_config.rs deleted file mode 100644 index 209bee48..00000000 --- a/src/sql/common/operator_config.rs +++ /dev/null @@ -1,21 +0,0 @@ -// 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 serde::{Deserialize, Serialize}; - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct MetadataField { - pub field_name: String, - pub key: String, - #[serde(default)] - pub data_type: Option, -} diff --git a/src/sql/common/time_utils.rs b/src/sql/common/time_utils.rs deleted file mode 100644 index 323445cd..00000000 --- a/src/sql/common/time_utils.rs +++ /dev/null @@ -1,74 +0,0 @@ -// 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::hash::Hash; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; - -pub fn to_millis(time: SystemTime) -> u64 { - time.duration_since(UNIX_EPOCH).unwrap().as_millis() as u64 -} - -pub fn to_micros(time: SystemTime) -> u64 { - time.duration_since(UNIX_EPOCH).unwrap().as_micros() as u64 -} - -pub fn from_millis(ts: u64) -> SystemTime { - UNIX_EPOCH + Duration::from_millis(ts) -} - -pub fn from_micros(ts: u64) -> SystemTime { - UNIX_EPOCH + Duration::from_micros(ts) -} - -pub fn to_nanos(time: SystemTime) -> u128 { - time.duration_since(UNIX_EPOCH).unwrap().as_nanos() -} - -pub fn from_nanos(ts: u128) -> SystemTime { - UNIX_EPOCH - + Duration::from_secs((ts / 1_000_000_000) as u64) - + Duration::from_nanos((ts % 1_000_000_000) as u64) -} - -pub fn print_time(time: SystemTime) -> String { - chrono::DateTime::::from(time) - .format("%Y-%m-%d %H:%M:%S%.3f") - .to_string() -} - -/// Returns the number of days since the UNIX epoch (for Avro serialization). -pub fn days_since_epoch(time: SystemTime) -> i32 { - time.duration_since(UNIX_EPOCH) - .unwrap() - .as_secs() - .div_euclid(86400) as i32 -} - -pub fn single_item_hash_map, K: Hash + Eq, V>(key: I, value: V) -> HashMap { - let mut map = HashMap::new(); - map.insert(key.into(), value); - map -} - -pub fn string_to_map(s: &str, pair_delimiter: char) -> Option> { - if s.trim().is_empty() { - return Some(HashMap::new()); - } - - s.split(',') - .map(|s| { - let mut kv = s.trim().split(pair_delimiter); - Some((kv.next()?.trim().to_string(), kv.next()?.trim().to_string())) - }) - .collect() -} diff --git a/src/sql/common/topology.rs b/src/sql/common/topology.rs deleted file mode 100644 index 3b4f892f..00000000 --- a/src/sql/common/topology.rs +++ /dev/null @@ -1,295 +0,0 @@ -// 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, VecDeque}; -use std::fmt::Write; - -use protocol::function_stream_graph::FsProgram; - -fn edge_type_label(edge_type: i32) -> &'static str { - match edge_type { - 1 => "Forward", - 2 => "Shuffle", - 3 => "LeftJoin", - 4 => "RightJoin", - _ => "Unknown", - } -} - -pub fn render_program_topology(program: &FsProgram) -> String { - if program.nodes.is_empty() { - return "(empty topology)".to_string(); - } - - struct EdgeInfo { - target: i32, - edge_type: i32, - } - struct InputInfo { - source: i32, - edge_type: i32, - } - - let node_map: BTreeMap = - program.nodes.iter().map(|n| (n.node_index, n)).collect(); - - let mut downstream: BTreeMap> = BTreeMap::new(); - let mut upstream: BTreeMap> = BTreeMap::new(); - let mut in_degree: BTreeMap = BTreeMap::new(); - - for idx in node_map.keys() { - in_degree.entry(*idx).or_insert(0); - } - for edge in &program.edges { - downstream.entry(edge.source).or_default().push(EdgeInfo { - target: edge.target, - edge_type: edge.edge_type, - }); - upstream.entry(edge.target).or_default().push(InputInfo { - source: edge.source, - edge_type: edge.edge_type, - }); - *in_degree.entry(edge.target).or_insert(0) += 1; - } - - // Kahn's topological sort - let mut queue: VecDeque = in_degree - .iter() - .filter(|(_, deg)| **deg == 0) - .map(|(idx, _)| *idx) - .collect(); - let mut topo_order: Vec = Vec::with_capacity(node_map.len()); - let mut remaining = in_degree.clone(); - while let Some(idx) = queue.pop_front() { - topo_order.push(idx); - if let Some(edges) = downstream.get(&idx) { - for e in edges { - if let Some(deg) = remaining.get_mut(&e.target) { - *deg -= 1; - if *deg == 0 { - queue.push_back(e.target); - } - } - } - } - } - for idx in node_map.keys() { - if !topo_order.contains(idx) { - topo_order.push(*idx); - } - } - - let is_source = |idx: &i32| upstream.get(idx).is_none_or(|v| v.is_empty()); - let is_sink = |idx: &i32| downstream.get(idx).is_none_or(|v| v.is_empty()); - - let mut out = String::new(); - let _ = writeln!( - out, - "Pipeline Topology ({} nodes, {} edges)", - program.nodes.len(), - program.edges.len(), - ); - let _ = writeln!(out, "{}", "=".repeat(50)); - - for (pos, &node_idx) in topo_order.iter().enumerate() { - let Some(node) = node_map.get(&node_idx) else { - continue; - }; - - let op_chain: String = node - .operators - .iter() - .map(|op| op.operator_name.as_str()) - .collect::>() - .join(" -> "); - - let role = if is_source(&node_idx) { - "Source" - } else if is_sink(&node_idx) { - "Sink" - } else { - "Operator" - }; - - let _ = writeln!(out); - let _ = writeln!( - out, - "[{role}] Node {node_idx} parallelism = {}", - node.parallelism, - ); - let _ = writeln!(out, " operators: {op_chain}"); - - if !node.description.is_empty() { - let _ = writeln!(out, " desc: {}", node.description); - } - - if let Some(inputs) = upstream.get(&node_idx) { - if inputs.len() == 1 { - let i = &inputs[0]; - let _ = writeln!( - out, - " input: <-- [{}] Node {}", - edge_type_label(i.edge_type), - i.source, - ); - } else if inputs.len() > 1 { - let _ = writeln!(out, " inputs:"); - for i in inputs { - let _ = writeln!( - out, - " <-- [{}] Node {}", - edge_type_label(i.edge_type), - i.source, - ); - } - } - } - - if let Some(outputs) = downstream.get(&node_idx) { - if outputs.len() == 1 { - let e = &outputs[0]; - let _ = writeln!( - out, - " output: --> [{}] Node {}", - edge_type_label(e.edge_type), - e.target, - ); - } else if outputs.len() > 1 { - let _ = writeln!(out, " outputs:"); - for e in outputs { - let _ = writeln!( - out, - " --> [{}] Node {}", - edge_type_label(e.edge_type), - e.target, - ); - } - } - } - - if pos < topo_order.len() - 1 { - let single_out = downstream.get(&node_idx).is_some_and(|v| v.len() == 1); - let next_idx = topo_order.get(pos + 1).copied(); - let is_direct = single_out - && next_idx - .is_some_and(|n| downstream.get(&node_idx).is_some_and(|v| v[0].target == n)); - let next_single_in = next_idx - .and_then(|n| upstream.get(&n)) - .is_some_and(|v| v.len() == 1); - - if is_direct && next_single_in { - let etype = downstream.get(&node_idx).unwrap()[0].edge_type; - let _ = writeln!(out, " |"); - let _ = writeln!(out, " | {}", edge_type_label(etype)); - let _ = writeln!(out, " v"); - } - } - } - - out.trim_end().to_string() -} - -#[cfg(test)] -mod tests { - use super::*; - use protocol::function_stream_graph::{ChainedOperator, FsEdge, FsNode, FsProgram}; - - fn make_node( - node_index: i32, - operators: Vec<(&str, &str)>, - desc: &str, - parallelism: u32, - ) -> FsNode { - FsNode { - node_index, - node_id: node_index as u32, - parallelism, - description: desc.to_string(), - operators: operators - .into_iter() - .map(|(id, name)| ChainedOperator { - operator_id: id.to_string(), - operator_name: name.to_string(), - operator_config: Vec::new(), - }) - .collect(), - edges: Vec::new(), - } - } - - fn make_edge(source: i32, target: i32, edge_type: i32) -> FsEdge { - FsEdge { - source, - target, - schema: None, - edge_type, - } - } - - #[test] - fn empty_program_renders_placeholder() { - let program = FsProgram { - nodes: vec![], - edges: vec![], - program_config: None, - }; - assert_eq!(render_program_topology(&program), "(empty topology)"); - } - - #[test] - fn linear_pipeline_renders_correctly() { - let program = FsProgram { - nodes: vec![ - make_node(0, vec![("src_0", "ConnectorSource")], "", 1), - make_node( - 1, - vec![("val_1", "Value"), ("wm_2", "ExpressionWatermark")], - "source -> watermark", - 1, - ), - make_node(2, vec![("sink_3", "ConnectorSink")], "sink (kafka)", 1), - ], - edges: vec![make_edge(0, 1, 1), make_edge(1, 2, 1)], - program_config: None, - }; - let result = render_program_topology(&program); - assert!(result.contains("[Source] Node 0")); - assert!(result.contains("[Operator] Node 1")); - assert!(result.contains("[Sink] Node 2")); - assert!(result.contains("ConnectorSource")); - assert!(result.contains("Value -> ExpressionWatermark")); - assert!(result.contains("Forward")); - } - - #[test] - fn join_topology_shows_multiple_inputs() { - let program = FsProgram { - nodes: vec![ - make_node(0, vec![("src_a", "ConnectorSource")], "source A", 1), - make_node(1, vec![("src_b", "ConnectorSource")], "source B", 1), - make_node(2, vec![("join_0", "WindowJoin")], "join node", 2), - make_node(3, vec![("sink_0", "ConnectorSink")], "sink", 1), - ], - edges: vec![ - make_edge(0, 2, 3), // LeftJoin - make_edge(1, 2, 4), // RightJoin - make_edge(2, 3, 1), // Forward - ], - program_config: None, - }; - let result = render_program_topology(&program); - assert!(result.contains("inputs:")); - assert!(result.contains("LeftJoin")); - assert!(result.contains("RightJoin")); - assert!(result.contains("[Operator] Node 2")); - } -} diff --git a/src/sql/common/with_option_keys.rs b/src/sql/common/with_option_keys.rs deleted file mode 100644 index 21bfa691..00000000 --- a/src/sql/common/with_option_keys.rs +++ /dev/null @@ -1,105 +0,0 @@ -// 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. - -pub const CONNECTOR: &str = "connector"; -pub const TYPE: &str = "type"; -pub const FORMAT: &str = "format"; -pub const DEFAULT_FORMAT_VALUE: &str = "json"; -pub const BAD_DATA: &str = "bad_data"; -pub const PARTITION_BY: &str = "partition_by"; -pub const PATH: &str = "path"; -pub const SINK_PATH: &str = "sink.path"; - -pub const EVENT_TIME_FIELD: &str = "event_time_field"; -pub const WATERMARK_FIELD: &str = "watermark_field"; - -pub const IDLE_MICROS: &str = "idle_micros"; -pub const IDLE_TIME: &str = "idle_time"; - -pub const LOOKUP_CACHE_MAX_BYTES: &str = "lookup.cache.max_bytes"; -pub const LOOKUP_CACHE_TTL: &str = "lookup.cache.ttl"; - -pub const CONNECTION_SCHEMA: &str = "connection_schema"; - -pub const ADAPTER: &str = "adapter"; - -// ── Kafka ───────────────────────────────────────────────────────────────── - -pub const KAFKA_BOOTSTRAP_SERVERS: &str = "bootstrap.servers"; -pub const KAFKA_BOOTSTRAP_SERVERS_LEGACY: &str = "bootstrap_servers"; -pub const KAFKA_TOPIC: &str = "topic"; -pub const KAFKA_RATE_LIMIT_MESSAGES_PER_SECOND: &str = "rate_limit.messages_per_second"; -pub const KAFKA_VALUE_SUBJECT: &str = "value.subject"; -pub const KAFKA_SCAN_STARTUP_MODE: &str = "scan.startup.mode"; -pub const KAFKA_ISOLATION_LEVEL: &str = "isolation.level"; -pub const KAFKA_GROUP_ID: &str = "group.id"; -pub const KAFKA_GROUP_ID_LEGACY: &str = "group_id"; -pub const KAFKA_GROUP_ID_PREFIX: &str = "group.id.prefix"; -pub const KAFKA_SINK_COMMIT_MODE: &str = "sink.commit.mode"; -pub const KAFKA_SINK_KEY_FIELD: &str = "sink.key.field"; -pub const KAFKA_KEY_FIELD_LEGACY: &str = "key.field"; -pub const KAFKA_SINK_TIMESTAMP_FIELD: &str = "sink.timestamp.field"; -pub const KAFKA_TIMESTAMP_FIELD_LEGACY: &str = "timestamp.field"; - -// ── JSON format ─────────────────────────────────────────────────────────── - -pub const JSON_CONFLUENT_SCHEMA_REGISTRY: &str = "json.confluent_schema_registry"; -pub const JSON_CONFLUENT_SCHEMA_VERSION: &str = "json.confluent_schema_version"; -pub const JSON_INCLUDE_SCHEMA: &str = "json.include_schema"; -pub const JSON_DEBEZIUM: &str = "json.debezium"; -pub const JSON_UNSTRUCTURED: &str = "json.unstructured"; -pub const JSON_TIMESTAMP_FORMAT: &str = "json.timestamp_format"; -pub const JSON_DECIMAL_ENCODING: &str = "json.decimal_encoding"; -pub const JSON_COMPRESSION: &str = "json.compression"; - -// ── Avro ────────────────────────────────────────────────────────────────── - -pub const AVRO_CONFLUENT_SCHEMA_REGISTRY: &str = "avro.confluent_schema_registry"; -pub const AVRO_RAW_DATUMS: &str = "avro.raw_datums"; -pub const AVRO_INTO_UNSTRUCTURED_JSON: &str = "avro.into_unstructured_json"; -pub const AVRO_SCHEMA_ID: &str = "avro.schema_id"; - -// ── Parquet ─────────────────────────────────────────────────────────────── - -pub const PARQUET_COMPRESSION: &str = "parquet.compression"; -pub const PARQUET_ROW_GROUP_BYTES: &str = "parquet.row_group_bytes"; - -// ── S3 ──────────────────────────────────────────────────────────────────── - -pub const S3_BUCKET: &str = "s3.bucket"; -pub const S3_REGION: &str = "s3.region"; -pub const S3_ENDPOINT: &str = "s3.endpoint"; -pub const S3_ACCESS_KEY_ID: &str = "s3.access_key_id"; -pub const S3_SECRET_ACCESS_KEY: &str = "s3.secret_access_key"; -pub const S3_SESSION_TOKEN: &str = "s3.session_token"; - -// ── Protobuf ──────────────────────────────────────────────────────────────── - -pub const PROTOBUF_INTO_UNSTRUCTURED_JSON: &str = "protobuf.into_unstructured_json"; -pub const PROTOBUF_MESSAGE_NAME: &str = "protobuf.message_name"; -pub const PROTOBUF_CONFLUENT_SCHEMA_REGISTRY: &str = "protobuf.confluent_schema_registry"; -pub const PROTOBUF_LENGTH_DELIMITED: &str = "protobuf.length_delimited"; - -// ── Framing ───────────────────────────────────────────────────────────────── - -pub const FRAMING_METHOD: &str = "framing.method"; -pub const FRAMING_MAX_LINE_LENGTH: &str = "framing.max_line_length"; - -pub const FORMAT_DEBEZIUM_FLAG: &str = "format.debezium"; - -// ── Streaming runtime common options ─────────────────────────────────────── - -pub const CHECKPOINT_INTERVAL_MS: &str = "checkpoint.interval.ms"; -pub const PIPELINE_PARALLELISM: &str = "pipeline.parallelism"; -pub const KEY_BY_PARALLELISM: &str = "key_by.parallelism"; -pub const OPERATOR_MEMORY_BYTES: &str = "operator.memory.bytes"; -pub const SINK_MEMORY_BYTES: &str = "sink.memory.bytes"; diff --git a/src/sql/connector/config.rs b/src/sql/connector/config.rs deleted file mode 100644 index c3eaf00f..00000000 --- a/src/sql/connector/config.rs +++ /dev/null @@ -1,91 +0,0 @@ -// 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 protocol::function_stream_graph::{ - DeltaSinkConfig, FilesystemSinkConfig, IcebergSinkConfig, KafkaSinkConfig, KafkaSourceConfig, - LanceDbSinkConfig, S3SinkConfig, connector_op, -}; - -#[derive(Debug, Clone)] -pub enum ConnectorConfig { - KafkaSource(KafkaSourceConfig), - KafkaSink(KafkaSinkConfig), - FilesystemSink(FilesystemSinkConfig), - DeltaSink(DeltaSinkConfig), - IcebergSink(IcebergSinkConfig), - S3Sink(S3SinkConfig), - LanceDbSink(LanceDbSinkConfig), -} - -impl ConnectorConfig { - pub fn to_proto_config(&self) -> connector_op::Config { - match self { - ConnectorConfig::KafkaSource(cfg) => connector_op::Config::KafkaSource(cfg.clone()), - ConnectorConfig::KafkaSink(cfg) => connector_op::Config::KafkaSink(cfg.clone()), - ConnectorConfig::FilesystemSink(cfg) => { - connector_op::Config::FilesystemSink(cfg.clone()) - } - ConnectorConfig::DeltaSink(cfg) => connector_op::Config::DeltaSink(cfg.clone()), - ConnectorConfig::IcebergSink(cfg) => connector_op::Config::IcebergSink(cfg.clone()), - ConnectorConfig::S3Sink(cfg) => connector_op::Config::S3Sink(cfg.clone()), - ConnectorConfig::LanceDbSink(cfg) => connector_op::Config::LancedbSink(cfg.clone()), - } - } -} - -impl PartialEq for ConnectorConfig { - fn eq(&self, other: &Self) -> bool { - use prost::Message; - match (self, other) { - (ConnectorConfig::KafkaSource(a), ConnectorConfig::KafkaSource(b)) => { - a.encode_to_vec() == b.encode_to_vec() - } - (ConnectorConfig::KafkaSink(a), ConnectorConfig::KafkaSink(b)) => { - a.encode_to_vec() == b.encode_to_vec() - } - (ConnectorConfig::FilesystemSink(a), ConnectorConfig::FilesystemSink(b)) => { - a.encode_to_vec() == b.encode_to_vec() - } - (ConnectorConfig::DeltaSink(a), ConnectorConfig::DeltaSink(b)) => { - a.encode_to_vec() == b.encode_to_vec() - } - (ConnectorConfig::IcebergSink(a), ConnectorConfig::IcebergSink(b)) => { - a.encode_to_vec() == b.encode_to_vec() - } - (ConnectorConfig::S3Sink(a), ConnectorConfig::S3Sink(b)) => { - a.encode_to_vec() == b.encode_to_vec() - } - (ConnectorConfig::LanceDbSink(a), ConnectorConfig::LanceDbSink(b)) => { - a.encode_to_vec() == b.encode_to_vec() - } - _ => false, - } - } -} - -impl Eq for ConnectorConfig {} - -impl std::hash::Hash for ConnectorConfig { - fn hash(&self, state: &mut H) { - use prost::Message; - std::mem::discriminant(self).hash(state); - match self { - ConnectorConfig::KafkaSource(cfg) => cfg.encode_to_vec().hash(state), - ConnectorConfig::KafkaSink(cfg) => cfg.encode_to_vec().hash(state), - ConnectorConfig::FilesystemSink(cfg) => cfg.encode_to_vec().hash(state), - ConnectorConfig::DeltaSink(cfg) => cfg.encode_to_vec().hash(state), - ConnectorConfig::IcebergSink(cfg) => cfg.encode_to_vec().hash(state), - ConnectorConfig::S3Sink(cfg) => cfg.encode_to_vec().hash(state), - ConnectorConfig::LanceDbSink(cfg) => cfg.encode_to_vec().hash(state), - } - } -} diff --git a/src/sql/connector/factory.rs b/src/sql/connector/factory.rs deleted file mode 100644 index 8c37a15a..00000000 --- a/src/sql/connector/factory.rs +++ /dev/null @@ -1,67 +0,0 @@ -// 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 datafusion::arrow::datatypes::Schema; -use datafusion::common::Result; - -use super::config::ConnectorConfig; -use super::registry::REGISTRY; -use super::sink::runtime_config::SinkRuntimeConfig; -use crate::sql::common::connector_options::ConnectorOptions; -use crate::sql::common::formats::{BadData, Format}; -use crate::sql::schema::table_role::TableRole; - -pub fn build_connector_config( - connector_name: &str, - role: TableRole, - options: &mut ConnectorOptions, - format: &Option, - bad_data: BadData, -) -> Result { - let runtime_opts_map = options.snapshot_for_catalog().into_iter().collect(); - let runtime_props = - SinkRuntimeConfig::from_options_map(&runtime_opts_map)?.to_runtime_properties(); - match role { - TableRole::Ingestion | TableRole::Reference => REGISTRY - .get_source(connector_name)? - .build_source_config(options, format, bad_data), - TableRole::Egress => { - REGISTRY - .get_sink(connector_name)? - .build_sink_config(options, format, &runtime_props) - } - } -} - -pub fn build_connector_config_from_options( - connector_name: &str, - role: TableRole, - options: &mut ConnectorOptions, - format: &Option, - bad_data: BadData, -) -> Result { - build_connector_config(connector_name, role, options, format, bad_data) -} - -pub fn build_connector_config_from_catalog( - connector_name: &str, - role: TableRole, - opts: HashMap, - _physical_schema: &Schema, -) -> Result { - let mut options = ConnectorOptions::from_flat_string_map(opts)?; - let format = Format::from_opts(&mut options)?; - let bad_data = BadData::from_opts(&mut options)?; - build_connector_config(connector_name, role, &mut options, &format, bad_data) -} diff --git a/src/sql/connector/mod.rs b/src/sql/connector/mod.rs deleted file mode 100644 index f477c976..00000000 --- a/src/sql/connector/mod.rs +++ /dev/null @@ -1,18 +0,0 @@ -// 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. - -pub mod config; -pub mod factory; -pub mod provider; -pub mod registry; -pub mod sink; -pub mod source; diff --git a/src/sql/connector/provider.rs b/src/sql/connector/provider.rs deleted file mode 100644 index 8875ee0c..00000000 --- a/src/sql/connector/provider.rs +++ /dev/null @@ -1,52 +0,0 @@ -// 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 datafusion::common::{DataFusionError, Result}; - -use super::config::ConnectorConfig; -use super::sink::runtime_config::SinkRuntimeProperties; -use crate::sql::common::connector_options::ConnectorOptions; -use crate::sql::common::formats::{BadData, Format}; - -pub trait SourceProvider: Send + Sync { - fn name(&self) -> &'static str; - - fn build_source_config( - &self, - options: &mut ConnectorOptions, - format: &Option, - bad_data: BadData, - ) -> Result; -} - -pub trait SinkProvider: Send + Sync { - fn name(&self) -> &'static str; - fn build_sink_config( - &self, - options: &mut ConnectorOptions, - format: &Option, - runtime_props: &SinkRuntimeProperties, - ) -> Result; -} - -pub fn require_option( - options: &mut ConnectorOptions, - key: &str, - connector_name: &str, -) -> Result { - options.pull_opt_str(key)?.ok_or_else(|| { - DataFusionError::Plan(format!( - "Connector '{}' requires option '{}' to be set", - connector_name, key - )) - }) -} diff --git a/src/sql/connector/registry.rs b/src/sql/connector/registry.rs deleted file mode 100644 index 4a8a8c1c..00000000 --- a/src/sql/connector/registry.rs +++ /dev/null @@ -1,86 +0,0 @@ -// 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::{Arc, LazyLock}; - -use datafusion::common::{DataFusionError, Result}; - -use super::provider::{SinkProvider, SourceProvider}; -use super::sink::delta::DeltaSinkConnector; -use super::sink::filesystem::FilesystemSinkConnector; -use super::sink::iceberg::IcebergSinkConnector; -use super::sink::kafka::KafkaSinkConnector; -use super::sink::lancedb::LanceDbSinkConnector; -use super::sink::s3::S3SinkConnector; -use super::source::kafka::KafkaSourceConnector; - -pub struct ConnectorRegistry { - sources: HashMap>, - sinks: HashMap>, -} - -impl ConnectorRegistry { - fn new() -> Self { - let mut registry = Self { - sources: HashMap::new(), - sinks: HashMap::new(), - }; - - registry.register_source(Arc::new(KafkaSourceConnector)); - - registry.register_sink(Arc::new(KafkaSinkConnector)); - registry.register_sink(Arc::new(S3SinkConnector)); - registry.register_sink(Arc::new(FilesystemSinkConnector)); - registry.register_sink(Arc::new(DeltaSinkConnector)); - registry.register_sink(Arc::new(IcebergSinkConnector)); - registry.register_sink(Arc::new(LanceDbSinkConnector)); - - registry - } - - pub fn register_source(&mut self, provider: Arc) { - self.sources - .insert(provider.name().to_ascii_lowercase(), provider); - } - - pub fn register_sink(&mut self, provider: Arc) { - self.sinks - .insert(provider.name().to_ascii_lowercase(), provider); - } - - pub fn get_source(&self, connector_name: &str) -> Result> { - self.sources - .get(&connector_name.to_ascii_lowercase()) - .cloned() - .ok_or_else(|| { - DataFusionError::Plan(format!( - "Connector '{}' is not registered or does not support being used as a SOURCE", - connector_name - )) - }) - } - - pub fn get_sink(&self, connector_name: &str) -> Result> { - self.sinks - .get(&connector_name.to_ascii_lowercase()) - .cloned() - .ok_or_else(|| { - DataFusionError::Plan(format!( - "Connector '{}' is not registered or does not support being used as a SINK", - connector_name - )) - }) - } -} - -pub static REGISTRY: LazyLock = LazyLock::new(ConnectorRegistry::new); diff --git a/src/sql/connector/sink/delta.rs b/src/sql/connector/sink/delta.rs deleted file mode 100644 index cd86660d..00000000 --- a/src/sql/connector/sink/delta.rs +++ /dev/null @@ -1,60 +0,0 @@ -// 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 datafusion::common::Result; -use protocol::function_stream_graph::{DeltaSinkConfig, SinkFormatProto}; - -use crate::sql::common::Format; -use crate::sql::common::connector_options::ConnectorOptions; -use crate::sql::common::constants::connector_type; -use crate::sql::connector::config::ConnectorConfig; -use crate::sql::connector::provider::SinkProvider; -use crate::sql::connector::sink::runtime_config::SinkRuntimeProperties; -use crate::sql::connector::sink::utils::SinkUtils; - -pub struct DeltaSinkConnector; - -impl SinkProvider for DeltaSinkConnector { - fn name(&self) -> &'static str { - connector_type::DELTA - } - - fn build_sink_config( - &self, - options: &mut ConnectorOptions, - format: &Option, - runtime_props: &SinkRuntimeProperties, - ) -> Result { - let path = SinkUtils::require_path(options)?; - let parquet_compression = SinkUtils::extract_parquet_compression(options)?; - let format_proto = SinkUtils::resolve_sink_format( - format, - self.name(), - &[ - SinkFormatProto::SinkFormatCsv, - SinkFormatProto::SinkFormatJsonl, - SinkFormatProto::SinkFormatAvro, - SinkFormatProto::SinkFormatParquet, - SinkFormatProto::SinkFormatOrc, - ], - )?; - let extra_properties = options.drain_remaining_string_values()?; - - Ok(ConnectorConfig::DeltaSink(DeltaSinkConfig { - path, - format: format_proto, - parquet_compression, - extra_properties, - runtime_properties: runtime_props.clone(), - })) - } -} diff --git a/src/sql/connector/sink/filesystem.rs b/src/sql/connector/sink/filesystem.rs deleted file mode 100644 index 224b1805..00000000 --- a/src/sql/connector/sink/filesystem.rs +++ /dev/null @@ -1,60 +0,0 @@ -// 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 datafusion::common::Result; -use protocol::function_stream_graph::{FilesystemSinkConfig, SinkFormatProto}; - -use crate::sql::common::Format; -use crate::sql::common::connector_options::ConnectorOptions; -use crate::sql::common::constants::connector_type; -use crate::sql::connector::config::ConnectorConfig; -use crate::sql::connector::provider::SinkProvider; -use crate::sql::connector::sink::runtime_config::SinkRuntimeProperties; -use crate::sql::connector::sink::utils::SinkUtils; - -pub struct FilesystemSinkConnector; - -impl SinkProvider for FilesystemSinkConnector { - fn name(&self) -> &'static str { - connector_type::FILESYSTEM - } - - fn build_sink_config( - &self, - options: &mut ConnectorOptions, - format: &Option, - runtime_props: &SinkRuntimeProperties, - ) -> Result { - let path = SinkUtils::require_path(options)?; - let parquet_compression = SinkUtils::extract_parquet_compression(options)?; - let format_proto = SinkUtils::resolve_sink_format( - format, - self.name(), - &[ - SinkFormatProto::SinkFormatCsv, - SinkFormatProto::SinkFormatJsonl, - SinkFormatProto::SinkFormatAvro, - SinkFormatProto::SinkFormatParquet, - SinkFormatProto::SinkFormatOrc, - ], - )?; - let extra_properties = options.drain_remaining_string_values()?; - - Ok(ConnectorConfig::FilesystemSink(FilesystemSinkConfig { - path, - format: format_proto, - parquet_compression, - extra_properties, - runtime_properties: runtime_props.clone(), - })) - } -} diff --git a/src/sql/connector/sink/iceberg.rs b/src/sql/connector/sink/iceberg.rs deleted file mode 100644 index 12f0d378..00000000 --- a/src/sql/connector/sink/iceberg.rs +++ /dev/null @@ -1,57 +0,0 @@ -// 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 datafusion::common::Result; -use protocol::function_stream_graph::{IcebergSinkConfig, SinkFormatProto}; - -use crate::sql::common::Format; -use crate::sql::common::connector_options::ConnectorOptions; -use crate::sql::common::constants::connector_type; -use crate::sql::connector::config::ConnectorConfig; -use crate::sql::connector::provider::SinkProvider; -use crate::sql::connector::sink::runtime_config::SinkRuntimeProperties; -use crate::sql::connector::sink::utils::SinkUtils; - -pub struct IcebergSinkConnector; - -impl SinkProvider for IcebergSinkConnector { - fn name(&self) -> &'static str { - connector_type::ICEBERG - } - - fn build_sink_config( - &self, - options: &mut ConnectorOptions, - format: &Option, - runtime_props: &SinkRuntimeProperties, - ) -> Result { - let path = SinkUtils::require_path(options)?; - let parquet_compression = SinkUtils::extract_parquet_compression(options)?; - let format_proto = SinkUtils::resolve_sink_format( - format, - self.name(), - &[ - SinkFormatProto::SinkFormatCsv, - SinkFormatProto::SinkFormatParquet, - ], - )?; - let extra_properties = options.drain_remaining_string_values()?; - - Ok(ConnectorConfig::IcebergSink(IcebergSinkConfig { - path, - format: format_proto, - parquet_compression, - extra_properties, - runtime_properties: runtime_props.clone(), - })) - } -} diff --git a/src/sql/connector/sink/kafka.rs b/src/sql/connector/sink/kafka.rs deleted file mode 100644 index a6fd115c..00000000 --- a/src/sql/connector/sink/kafka.rs +++ /dev/null @@ -1,159 +0,0 @@ -// 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 datafusion::common::{Result, plan_datafusion_err, plan_err}; -use protocol::function_stream_graph::{ - DecimalEncodingProto, FormatConfig, JsonFormatConfig, KafkaAuthConfig, KafkaAuthNone, - KafkaSinkCommitMode, KafkaSinkConfig, RawBytesFormatConfig, RawStringFormatConfig, - TimestampFormatProto, format_config, kafka_auth_config, -}; - -use crate::sql::common::connector_options::ConnectorOptions; -use crate::sql::common::constants::{connector_type, kafka_with_value}; -use crate::sql::common::formats::{ - DecimalEncoding as SqlDecimalEncoding, Format as SqlFormat, - TimestampFormat as SqlTimestampFormat, -}; -use crate::sql::common::with_option_keys as opt; -use crate::sql::connector::config::ConnectorConfig; -use crate::sql::connector::provider::SinkProvider; -use crate::sql::connector::sink::runtime_config::SinkRuntimeProperties; - -pub struct KafkaSinkConnector; - -impl KafkaSinkConnector { - fn sql_format_to_proto(fmt: &SqlFormat) -> Result { - match fmt { - SqlFormat::Json(j) => Ok(FormatConfig { - format: Some(format_config::Format::Json(JsonFormatConfig { - timestamp_format: match j.timestamp_format { - SqlTimestampFormat::RFC3339 => { - TimestampFormatProto::TimestampRfc3339 as i32 - } - SqlTimestampFormat::UnixMillis => { - TimestampFormatProto::TimestampUnixMillis as i32 - } - }, - decimal_encoding: match j.decimal_encoding { - SqlDecimalEncoding::Number => DecimalEncodingProto::DecimalNumber as i32, - SqlDecimalEncoding::String => DecimalEncodingProto::DecimalString as i32, - SqlDecimalEncoding::Bytes => DecimalEncodingProto::DecimalBytes as i32, - }, - include_schema: j.include_schema, - confluent_schema_registry: j.confluent_schema_registry, - schema_id: j.schema_id, - debezium: j.debezium, - unstructured: j.unstructured, - })), - }), - SqlFormat::RawString(_) => Ok(FormatConfig { - format: Some(format_config::Format::RawString(RawStringFormatConfig {})), - }), - SqlFormat::RawBytes(_) => Ok(FormatConfig { - format: Some(format_config::Format::RawBytes(RawBytesFormatConfig {})), - }), - other => plan_err!( - "Kafka sink connector: format '{}' is not supported", - other.name() - ), - } - } -} - -impl SinkProvider for KafkaSinkConnector { - fn name(&self) -> &'static str { - connector_type::KAFKA - } - - fn build_sink_config( - &self, - options: &mut ConnectorOptions, - format: &Option, - _runtime_props: &SinkRuntimeProperties, - ) -> Result { - let bootstrap_servers = match options.pull_opt_str(opt::KAFKA_BOOTSTRAP_SERVERS)? { - Some(s) => s, - None => options - .pull_opt_str(opt::KAFKA_BOOTSTRAP_SERVERS_LEGACY)? - .ok_or_else(|| { - plan_datafusion_err!( - "Kafka connector requires 'bootstrap.servers' in the WITH clause" - ) - })?, - }; - - let topic = options.pull_opt_str(opt::KAFKA_TOPIC)?.ok_or_else(|| { - plan_datafusion_err!("Kafka connector requires 'topic' in the WITH clause") - })?; - - let sql_format = format.as_ref().ok_or_else(|| { - plan_datafusion_err!( - "Kafka sink requires 'format' in the WITH clause (e.g. format = 'json')" - ) - })?; - let proto_format = Self::sql_format_to_proto(sql_format)?; - - let value_subject = options.pull_opt_str(opt::KAFKA_VALUE_SUBJECT)?; - - let commit_mode = match options - .pull_opt_str(opt::KAFKA_SINK_COMMIT_MODE)? - .as_deref() - { - Some(s) - if s == kafka_with_value::SINK_COMMIT_EXACTLY_ONCE_HYPHEN - || s == kafka_with_value::SINK_COMMIT_EXACTLY_ONCE_UNDERSCORE => - { - KafkaSinkCommitMode::KafkaSinkExactlyOnce as i32 - } - Some(s) - if s == kafka_with_value::SINK_COMMIT_AT_LEAST_ONCE_HYPHEN - || s == kafka_with_value::SINK_COMMIT_AT_LEAST_ONCE_UNDERSCORE => - { - KafkaSinkCommitMode::KafkaSinkAtLeastOnce as i32 - } - None => KafkaSinkCommitMode::KafkaSinkAtLeastOnce as i32, - Some(other) => return plan_err!("invalid sink.commit.mode '{other}'"), - }; - - let key_field = match options.pull_opt_str(opt::KAFKA_SINK_KEY_FIELD)? { - Some(s) => Some(s), - None => options.pull_opt_str(opt::KAFKA_KEY_FIELD_LEGACY)?, - }; - let timestamp_field = match options.pull_opt_str(opt::KAFKA_SINK_TIMESTAMP_FIELD)? { - Some(s) => Some(s), - None => options.pull_opt_str(opt::KAFKA_TIMESTAMP_FIELD_LEGACY)?, - }; - - let _ = options.pull_opt_str(opt::TYPE)?; - let _ = options.pull_opt_str(opt::CONNECTOR)?; - - let mut client_configs = options.drain_remaining_string_values()?; - client_configs.remove(opt::CHECKPOINT_INTERVAL_MS); - client_configs.remove(opt::PIPELINE_PARALLELISM); - client_configs.remove(opt::KEY_BY_PARALLELISM); - client_configs.remove(opt::FORMAT); - - Ok(ConnectorConfig::KafkaSink(KafkaSinkConfig { - topic, - bootstrap_servers, - commit_mode, - key_field, - timestamp_field, - auth: Some(KafkaAuthConfig { - auth: Some(kafka_auth_config::Auth::None(KafkaAuthNone {})), - }), - client_configs, - format: Some(proto_format), - value_subject, - })) - } -} diff --git a/src/sql/connector/sink/lancedb.rs b/src/sql/connector/sink/lancedb.rs deleted file mode 100644 index aee79735..00000000 --- a/src/sql/connector/sink/lancedb.rs +++ /dev/null @@ -1,61 +0,0 @@ -// 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 datafusion::common::Result; -use protocol::function_stream_graph::{LanceDbSinkConfig, SinkFormatProto}; - -use crate::sql::common::Format; -use crate::sql::common::connector_options::ConnectorOptions; -use crate::sql::common::with_option_keys as opt; -use crate::sql::connector::config::ConnectorConfig; -use crate::sql::connector::provider::SinkProvider; -use crate::sql::connector::sink::runtime_config::SinkRuntimeProperties; -use crate::sql::connector::sink::utils::SinkUtils; - -pub struct LanceDbSinkConnector; - -impl SinkProvider for LanceDbSinkConnector { - fn name(&self) -> &'static str { - "lancedb" - } - - fn build_sink_config( - &self, - options: &mut ConnectorOptions, - _format: &Option, - runtime_props: &SinkRuntimeProperties, - ) -> Result { - let path = SinkUtils::require_path(options)?; - - let s3_bucket = options.pull_opt_str(opt::S3_BUCKET)?; - let s3_region = options.pull_opt_str(opt::S3_REGION)?; - let s3_endpoint = options.pull_opt_str(opt::S3_ENDPOINT)?; - let s3_access_key_id = options.pull_opt_str(opt::S3_ACCESS_KEY_ID)?; - let s3_secret_access_key = options.pull_opt_str(opt::S3_SECRET_ACCESS_KEY)?; - let s3_session_token = options.pull_opt_str(opt::S3_SESSION_TOKEN)?; - - let extra_properties = options.drain_remaining_string_values()?; - - Ok(ConnectorConfig::LanceDbSink(LanceDbSinkConfig { - path, - format: SinkFormatProto::SinkFormatLance as i32, - s3_bucket, - s3_region, - s3_endpoint, - s3_access_key_id, - s3_secret_access_key, - s3_session_token, - extra_properties, - runtime_properties: runtime_props.clone(), - })) - } -} diff --git a/src/sql/connector/sink/mod.rs b/src/sql/connector/sink/mod.rs deleted file mode 100644 index b7d645ca..00000000 --- a/src/sql/connector/sink/mod.rs +++ /dev/null @@ -1,20 +0,0 @@ -// 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. - -pub mod delta; -pub mod filesystem; -pub mod iceberg; -pub mod kafka; -pub mod lancedb; -pub mod runtime_config; -pub mod s3; -pub mod utils; diff --git a/src/sql/connector/sink/runtime_config.rs b/src/sql/connector/sink/runtime_config.rs deleted file mode 100644 index e0ffaeee..00000000 --- a/src/sql/connector/sink/runtime_config.rs +++ /dev/null @@ -1,137 +0,0 @@ -// 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 datafusion::common::{DataFusionError, Result, plan_err}; - -use crate::config::global_config::{ - DEFAULT_OPERATOR_STATE_STORE_MEMORY_BYTES, DEFAULT_SINK_BUFFER_MEMORY_BYTES, -}; -use crate::config::streaming_job::DEFAULT_CHECKPOINT_INTERVAL_MS; -use crate::sql::common::connector_options::ConnectorOptions; -use crate::sql::common::with_option_keys as opt; - -#[derive(Debug, Clone, PartialEq, Eq, Default)] -pub struct SinkRuntimeConfig { - pub pipeline_parallelism: Option, - pub key_by_parallelism: Option, - pub checkpoint_interval_ms: u64, - pub operator_memory_bytes: u64, - pub sink_memory_bytes: u64, -} - -pub type SinkRuntimeProperties = HashMap; - -impl SinkRuntimeConfig { - pub fn extract_from_options(options: &mut ConnectorOptions) -> Result { - let pipeline_parallelism = options - .pull_opt_u64(opt::PIPELINE_PARALLELISM)? - .map(|v| v as u32); - let key_by_parallelism = options - .pull_opt_u64(opt::KEY_BY_PARALLELISM)? - .map(|v| v as u32); - let checkpoint_interval_ms = options - .pull_opt_u64(opt::CHECKPOINT_INTERVAL_MS)? - .unwrap_or(DEFAULT_CHECKPOINT_INTERVAL_MS); - let operator_memory_bytes = options - .pull_opt_u64(opt::OPERATOR_MEMORY_BYTES)? - .unwrap_or(DEFAULT_OPERATOR_STATE_STORE_MEMORY_BYTES); - let sink_memory_bytes = options - .pull_opt_u64(opt::SINK_MEMORY_BYTES)? - .unwrap_or(DEFAULT_SINK_BUFFER_MEMORY_BYTES); - Ok(Self { - pipeline_parallelism, - key_by_parallelism, - checkpoint_interval_ms, - operator_memory_bytes, - sink_memory_bytes, - }) - } - - pub fn from_options_map(opts: &HashMap) -> Result { - let pipeline_parallelism = parse_opt_u32(opts, opt::PIPELINE_PARALLELISM)?; - let key_by_parallelism = parse_opt_u32(opts, opt::KEY_BY_PARALLELISM)?; - let checkpoint_interval_ms = parse_opt_u64(opts, opt::CHECKPOINT_INTERVAL_MS)? - .unwrap_or(DEFAULT_CHECKPOINT_INTERVAL_MS); - let operator_memory_bytes = parse_opt_u64(opts, opt::OPERATOR_MEMORY_BYTES)? - .unwrap_or(DEFAULT_OPERATOR_STATE_STORE_MEMORY_BYTES); - let sink_memory_bytes = parse_opt_u64(opts, opt::SINK_MEMORY_BYTES)? - .unwrap_or(DEFAULT_SINK_BUFFER_MEMORY_BYTES); - Ok(Self { - pipeline_parallelism, - key_by_parallelism, - checkpoint_interval_ms, - operator_memory_bytes, - sink_memory_bytes, - }) - } - - pub fn to_runtime_properties(&self) -> HashMap { - let mut out = HashMap::new(); - if let Some(v) = self.pipeline_parallelism { - out.insert(opt::PIPELINE_PARALLELISM.to_string(), v.to_string()); - } - if let Some(v) = self.key_by_parallelism { - out.insert(opt::KEY_BY_PARALLELISM.to_string(), v.to_string()); - } - out.insert( - opt::CHECKPOINT_INTERVAL_MS.to_string(), - self.checkpoint_interval_ms.to_string(), - ); - out.insert( - opt::OPERATOR_MEMORY_BYTES.to_string(), - self.operator_memory_bytes.to_string(), - ); - out.insert( - opt::SINK_MEMORY_BYTES.to_string(), - self.sink_memory_bytes.to_string(), - ); - out - } -} - -fn parse_opt_u32(opts: &HashMap, key: &str) -> Result> { - let Some(raw) = opts.get(key) else { - return Ok(None); - }; - let normalized = normalize_numeric_option(raw); - let parsed = normalized.parse::().map_err(|_| { - DataFusionError::Plan(format!( - "WITH option '{key}' expects unsigned integer, got '{raw}'" - )) - })?; - if parsed == 0 { - return plan_err!("WITH option '{key}' must be > 0"); - } - Ok(Some(parsed)) -} - -fn parse_opt_u64(opts: &HashMap, key: &str) -> Result> { - let Some(raw) = opts.get(key) else { - return Ok(None); - }; - let normalized = normalize_numeric_option(raw); - let parsed = normalized.parse::().map_err(|_| { - DataFusionError::Plan(format!( - "WITH option '{key}' expects unsigned integer, got '{raw}'" - )) - })?; - if parsed == 0 { - return plan_err!("WITH option '{key}' must be > 0"); - } - Ok(Some(parsed)) -} - -fn normalize_numeric_option(raw: &str) -> &str { - raw.trim().trim_matches('\'').trim_matches('"').trim() -} diff --git a/src/sql/connector/sink/s3.rs b/src/sql/connector/sink/s3.rs deleted file mode 100644 index 5d04ce46..00000000 --- a/src/sql/connector/sink/s3.rs +++ /dev/null @@ -1,75 +0,0 @@ -// 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 datafusion::common::Result; -use protocol::function_stream_graph::{S3SinkConfig, SinkFormatProto}; - -use crate::sql::common::Format; -use crate::sql::common::connector_options::ConnectorOptions; -use crate::sql::common::constants::connector_type; -use crate::sql::common::with_option_keys as opt; -use crate::sql::connector::config::ConnectorConfig; -use crate::sql::connector::provider::SinkProvider; -use crate::sql::connector::sink::runtime_config::SinkRuntimeProperties; -use crate::sql::connector::sink::utils::SinkUtils; - -pub struct S3SinkConnector; - -impl SinkProvider for S3SinkConnector { - fn name(&self) -> &'static str { - connector_type::S3 - } - - fn build_sink_config( - &self, - options: &mut ConnectorOptions, - format: &Option, - runtime_props: &SinkRuntimeProperties, - ) -> Result { - let path = SinkUtils::require_path(options)?; - - let format_proto = SinkUtils::resolve_sink_format( - format, - self.name(), - &[ - SinkFormatProto::SinkFormatCsv, - SinkFormatProto::SinkFormatParquet, - ], - )?; - - let bucket = SinkUtils::require_str(options, opt::S3_BUCKET, self.name())?; - let region = options - .pull_opt_str(opt::S3_REGION)? - .unwrap_or_else(|| "us-east-1".to_string()); - let endpoint = options.pull_opt_str(opt::S3_ENDPOINT)?; - let access_key_id = options.pull_opt_str(opt::S3_ACCESS_KEY_ID)?; - let secret_access_key = options.pull_opt_str(opt::S3_SECRET_ACCESS_KEY)?; - let session_token = options.pull_opt_str(opt::S3_SESSION_TOKEN)?; - - let parquet_compression = SinkUtils::extract_parquet_compression(options)?; - let extra_properties = options.drain_remaining_string_values()?; - - Ok(ConnectorConfig::S3Sink(S3SinkConfig { - path, - format: format_proto, - bucket, - region, - endpoint, - access_key_id, - secret_access_key, - session_token, - parquet_compression, - extra_properties, - runtime_properties: runtime_props.clone(), - })) - } -} diff --git a/src/sql/connector/sink/utils.rs b/src/sql/connector/sink/utils.rs deleted file mode 100644 index e61cd870..00000000 --- a/src/sql/connector/sink/utils.rs +++ /dev/null @@ -1,91 +0,0 @@ -// 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 datafusion::common::{DataFusionError, Result, plan_err}; -use protocol::function_stream_graph::{ParquetCompressionProto, SinkFormatProto}; - -use crate::sql::common::Format; -use crate::sql::common::connector_options::ConnectorOptions; -use crate::sql::common::constants::parquet_compression_value; -use crate::sql::common::with_option_keys as opt; - -pub struct SinkUtils; - -impl SinkUtils { - pub fn require_path(options: &mut ConnectorOptions) -> Result { - if let Some(v) = options.pull_opt_str(opt::PATH)? { - return Ok(v); - } - if let Some(v) = options.pull_opt_str(opt::SINK_PATH)? { - return Ok(v); - } - plan_err!("Missing required WITH option 'path' (or 'sink.path')") - } - - pub fn extract_parquet_compression(options: &mut ConnectorOptions) -> Result> { - let Some(v) = options.pull_opt_str(opt::PARQUET_COMPRESSION)? else { - return Ok(None); - }; - let parsed = match v.to_ascii_lowercase().as_str() { - parquet_compression_value::UNCOMPRESSED => { - ParquetCompressionProto::ParquetCompressionUncompressed - } - parquet_compression_value::SNAPPY => ParquetCompressionProto::ParquetCompressionSnappy, - parquet_compression_value::GZIP => ParquetCompressionProto::ParquetCompressionGzip, - parquet_compression_value::ZSTD => ParquetCompressionProto::ParquetCompressionZstd, - parquet_compression_value::LZ4 => ParquetCompressionProto::ParquetCompressionLz4, - parquet_compression_value::LZ4_RAW => ParquetCompressionProto::ParquetCompressionLz4Raw, - other => return plan_err!("Unsupported parquet.compression '{other}'"), - }; - Ok(Some(parsed as i32)) - } - - pub fn require_str( - options: &mut ConnectorOptions, - key: &str, - connector: &str, - ) -> Result { - options.pull_opt_str(key)?.ok_or_else(|| { - DataFusionError::Plan(format!( - "Connector '{connector}' requires WITH option '{key}'" - )) - }) - } - - pub fn resolve_sink_format( - format: &Option, - connector_name: &str, - supported_formats: &[SinkFormatProto], - ) -> Result { - let proto_format = match format { - Some(Format::Csv(_)) => SinkFormatProto::SinkFormatCsv, - Some(Format::Json(_)) => SinkFormatProto::SinkFormatJsonl, - Some(Format::Avro(_)) => SinkFormatProto::SinkFormatAvro, - Some(Format::Parquet(_)) => SinkFormatProto::SinkFormatParquet, - Some(Format::Lance(_)) => SinkFormatProto::SinkFormatLance, - Some(f) => { - return plan_err!("Format '{f:?}' cannot be mapped to a sink format"); - } - None => { - return plan_err!("Connector '{connector_name}' requires a format to be specified"); - } - }; - - if !supported_formats.contains(&proto_format) { - return plan_err!( - "Format {proto_format:?} is not supported by connector '{connector_name}'" - ); - } - - Ok(proto_format as i32) - } -} diff --git a/src/sql/connector/source/kafka.rs b/src/sql/connector/source/kafka.rs deleted file mode 100644 index 0bc220d9..00000000 --- a/src/sql/connector/source/kafka.rs +++ /dev/null @@ -1,185 +0,0 @@ -// 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 datafusion::common::{Result, plan_datafusion_err, plan_err}; -use protocol::function_stream_graph::{ - BadDataPolicy, DecimalEncodingProto, FormatConfig, JsonFormatConfig, KafkaAuthConfig, - KafkaAuthNone, KafkaOffsetMode, KafkaReadMode, KafkaSourceConfig, RawBytesFormatConfig, - RawStringFormatConfig, TimestampFormatProto, format_config, kafka_auth_config, -}; - -use crate::sql::common::connector_options::ConnectorOptions; -use crate::sql::common::constants::{connector_type, kafka_with_value}; -use crate::sql::common::formats::{ - BadData, DecimalEncoding as SqlDecimalEncoding, Format as SqlFormat, - TimestampFormat as SqlTimestampFormat, -}; -use crate::sql::common::with_option_keys as opt; -use crate::sql::connector::config::ConnectorConfig; -use crate::sql::connector::provider::SourceProvider; - -pub struct KafkaSourceConnector; - -impl KafkaSourceConnector { - fn sql_format_to_proto(fmt: &SqlFormat) -> Result { - match fmt { - SqlFormat::Json(j) => Ok(FormatConfig { - format: Some(format_config::Format::Json(JsonFormatConfig { - timestamp_format: match j.timestamp_format { - SqlTimestampFormat::RFC3339 => { - TimestampFormatProto::TimestampRfc3339 as i32 - } - SqlTimestampFormat::UnixMillis => { - TimestampFormatProto::TimestampUnixMillis as i32 - } - }, - decimal_encoding: match j.decimal_encoding { - SqlDecimalEncoding::Number => DecimalEncodingProto::DecimalNumber as i32, - SqlDecimalEncoding::String => DecimalEncodingProto::DecimalString as i32, - SqlDecimalEncoding::Bytes => DecimalEncodingProto::DecimalBytes as i32, - }, - include_schema: j.include_schema, - confluent_schema_registry: j.confluent_schema_registry, - schema_id: j.schema_id, - debezium: j.debezium, - unstructured: j.unstructured, - })), - }), - SqlFormat::RawString(_) => Ok(FormatConfig { - format: Some(format_config::Format::RawString(RawStringFormatConfig {})), - }), - SqlFormat::RawBytes(_) => Ok(FormatConfig { - format: Some(format_config::Format::RawBytes(RawBytesFormatConfig {})), - }), - other => plan_err!( - "Kafka source connector: format '{}' is not supported", - other.name() - ), - } - } - - fn bad_data_to_proto(bad: &BadData) -> i32 { - match bad { - BadData::Fail {} => BadDataPolicy::BadDataFail as i32, - BadData::Drop {} => BadDataPolicy::BadDataDrop as i32, - } - } -} - -impl SourceProvider for KafkaSourceConnector { - fn name(&self) -> &'static str { - connector_type::KAFKA - } - - fn build_source_config( - &self, - options: &mut ConnectorOptions, - format: &Option, - bad_data: BadData, - ) -> Result { - let bootstrap_servers = match options.pull_opt_str(opt::KAFKA_BOOTSTRAP_SERVERS)? { - Some(s) => s, - None => options - .pull_opt_str(opt::KAFKA_BOOTSTRAP_SERVERS_LEGACY)? - .ok_or_else(|| { - plan_datafusion_err!( - "Kafka connector requires 'bootstrap.servers' in the WITH clause" - ) - })?, - }; - - let topic = options.pull_opt_str(opt::KAFKA_TOPIC)?.ok_or_else(|| { - plan_datafusion_err!("Kafka connector requires 'topic' in the WITH clause") - })?; - - let sql_format = format.as_ref().ok_or_else(|| { - plan_datafusion_err!( - "Kafka source requires 'format' in the WITH clause (e.g. format = 'json')" - ) - })?; - let proto_format = Self::sql_format_to_proto(sql_format)?; - - let rate_limit = options - .pull_opt_u64(opt::KAFKA_RATE_LIMIT_MESSAGES_PER_SECOND)? - .map(|v| v.clamp(1, u32::MAX as u64) as u32) - .unwrap_or(0); - - let value_subject = options.pull_opt_str(opt::KAFKA_VALUE_SUBJECT)?; - - let offset_mode = match options - .pull_opt_str(opt::KAFKA_SCAN_STARTUP_MODE)? - .as_deref() - { - Some(s) if s == kafka_with_value::SCAN_LATEST => { - KafkaOffsetMode::KafkaOffsetLatest as i32 - } - Some(s) if s == kafka_with_value::SCAN_EARLIEST => { - KafkaOffsetMode::KafkaOffsetEarliest as i32 - } - Some(s) - if s == kafka_with_value::SCAN_GROUP_OFFSETS - || s == kafka_with_value::SCAN_GROUP => - { - KafkaOffsetMode::KafkaOffsetGroup as i32 - } - None => KafkaOffsetMode::KafkaOffsetGroup as i32, - Some(other) => { - return plan_err!( - "invalid scan.startup.mode '{other}'; expected latest, earliest, or group-offsets" - ); - } - }; - - let read_mode = match options.pull_opt_str(opt::KAFKA_ISOLATION_LEVEL)?.as_deref() { - Some(s) if s == kafka_with_value::ISOLATION_READ_COMMITTED => { - KafkaReadMode::KafkaReadCommitted as i32 - } - Some(s) if s == kafka_with_value::ISOLATION_READ_UNCOMMITTED => { - KafkaReadMode::KafkaReadUncommitted as i32 - } - None => KafkaReadMode::KafkaReadDefault as i32, - Some(other) => return plan_err!("invalid isolation.level '{other}'"), - }; - - let group_id = match options.pull_opt_str(opt::KAFKA_GROUP_ID)? { - Some(s) => Some(s), - None => options.pull_opt_str(opt::KAFKA_GROUP_ID_LEGACY)?, - }; - let group_id_prefix = options.pull_opt_str(opt::KAFKA_GROUP_ID_PREFIX)?; - - let _ = options.pull_opt_str(opt::TYPE)?; - let _ = options.pull_opt_str(opt::CONNECTOR)?; - - let mut client_configs = options.drain_remaining_string_values()?; - client_configs.remove(opt::CHECKPOINT_INTERVAL_MS); - client_configs.remove(opt::PIPELINE_PARALLELISM); - client_configs.remove(opt::KEY_BY_PARALLELISM); - client_configs.remove(opt::FORMAT); - - Ok(ConnectorConfig::KafkaSource(KafkaSourceConfig { - topic, - bootstrap_servers, - group_id, - group_id_prefix, - offset_mode, - read_mode, - auth: Some(KafkaAuthConfig { - auth: Some(kafka_auth_config::Auth::None(KafkaAuthNone {})), - }), - client_configs, - format: Some(proto_format), - bad_data_policy: Self::bad_data_to_proto(&bad_data), - rate_limit_msgs_per_sec: rate_limit, - value_subject, - })) - } -} diff --git a/src/sql/connector/source/mod.rs b/src/sql/connector/source/mod.rs deleted file mode 100644 index b9574391..00000000 --- a/src/sql/connector/source/mod.rs +++ /dev/null @@ -1,13 +0,0 @@ -// 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. - -pub mod kafka; diff --git a/src/sql/functions/mod.rs b/src/sql/functions/mod.rs deleted file mode 100644 index b78f5d2a..00000000 --- a/src/sql/functions/mod.rs +++ /dev/null @@ -1,612 +0,0 @@ -// 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::sql::schema::StreamSchemaProvider; -use datafusion::arrow::array::{ - Array, ArrayRef, StringArray, UnionArray, - builder::{FixedSizeBinaryBuilder, ListBuilder, StringBuilder}, - cast::{AsArray, as_string_array}, - types::{Float64Type, Int64Type}, -}; -use datafusion::arrow::datatypes::{DataType, Field, UnionFields, UnionMode}; -use datafusion::arrow::row::{RowConverter, SortField}; -use datafusion::common::{DataFusionError, ScalarValue}; -use datafusion::common::{Result, TableReference}; -use datafusion::execution::FunctionRegistry; -use datafusion::logical_expr::expr::{Alias, ScalarFunction}; -use datafusion::logical_expr::{ - ColumnarValue, LogicalPlan, Projection, ScalarFunctionArgs, ScalarUDFImpl, Signature, - TypeSignature, Volatility, create_udf, -}; -use datafusion::prelude::{Expr, col}; -use serde_json_path::JsonPath; -use std::any::Any; -use std::collections::HashMap; -use std::fmt::{Debug, Write}; -use std::sync::{Arc, OnceLock}; - -use crate::sql::common::constants::scalar_fn; - -/// Borrowed from DataFusion -/// -/// Creates a singleton `ScalarUDF` of the `$UDF` function named `$GNAME` and a -/// function named `$NAME` which returns that function named $NAME. -/// -/// This is used to ensure creating the list of `ScalarUDF` only happens once. -#[macro_export] -macro_rules! make_udf_function { - ($UDF:ty, $GNAME:ident, $NAME:ident) => { - /// Singleton instance of the function - static $GNAME: std::sync::OnceLock> = - std::sync::OnceLock::new(); - - /// Return a [`ScalarUDF`] for [`$UDF`] - /// - /// [`ScalarUDF`]: datafusion_expr::ScalarUDF - pub fn $NAME() -> std::sync::Arc { - $GNAME - .get_or_init(|| { - std::sync::Arc::new(datafusion::logical_expr::ScalarUDF::new_from_impl( - <$UDF>::default(), - )) - }) - .clone() - } - }; -} - -make_udf_function!(MultiHashFunction, MULTI_HASH, multi_hash); - -pub fn register_all(registry: &mut dyn FunctionRegistry) { - registry - .register_udf(Arc::new(create_udf( - scalar_fn::GET_FIRST_JSON_OBJECT, - vec![DataType::Utf8, DataType::Utf8], - DataType::Utf8, - Volatility::Immutable, - Arc::new(get_first_json_object), - ))) - .unwrap(); - - registry - .register_udf(Arc::new(create_udf( - scalar_fn::EXTRACT_JSON, - vec![DataType::Utf8, DataType::Utf8], - DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))), - Volatility::Immutable, - Arc::new(extract_json), - ))) - .unwrap(); - - registry - .register_udf(Arc::new(create_udf( - scalar_fn::EXTRACT_JSON_STRING, - vec![DataType::Utf8, DataType::Utf8], - DataType::Utf8, - Volatility::Immutable, - Arc::new(extract_json_string), - ))) - .unwrap(); - - registry - .register_udf(Arc::new(create_udf( - scalar_fn::SERIALIZE_JSON_UNION, - vec![DataType::Union(union_fields(), UnionMode::Sparse)], - DataType::Utf8, - Volatility::Immutable, - Arc::new(serialize_json_union), - ))) - .unwrap(); - - registry.register_udf(multi_hash()).unwrap(); -} - -fn parse_path(name: &str, path: &ScalarValue) -> Result> { - let path = match path { - ScalarValue::Utf8(Some(s)) => JsonPath::parse(s) - .map_err(|e| DataFusionError::Execution(format!("Invalid json path '{s}': {e:?}")))?, - ScalarValue::Utf8(None) => { - return Err(DataFusionError::Execution(format!( - "The path argument to {name} cannot be null" - ))); - } - _ => { - return Err(DataFusionError::Execution(format!( - "The path argument to {name} must be of type TEXT" - ))); - } - }; - - Ok(Arc::new(path)) -} - -// Hash function that can take any number of arguments and produces a fast (non-cryptographic) -// 128-bit hash from their string representations -#[derive(Debug)] -pub struct MultiHashFunction { - signature: Signature, -} - -impl MultiHashFunction { - pub fn invoke(&self, args: &[ColumnarValue]) -> Result { - let mut hasher = xxhash_rust::xxh3::Xxh3::new(); - - let all_scalar = args.iter().all(|a| matches!(a, ColumnarValue::Scalar(_))); - - let length = args - .iter() - .map(|t| match t { - ColumnarValue::Scalar(_) => 1, - ColumnarValue::Array(a) => a.len(), - }) - .max() - .ok_or_else(|| { - DataFusionError::Plan("multi_hash must have at least one argument".to_string()) - })?; - - let row_builder = RowConverter::new( - args.iter() - .map(|t| SortField::new(t.data_type().clone())) - .collect(), - )?; - - let arrays = args - .iter() - .map(|c| c.clone().into_array(length)) - .collect::>>()?; - let rows = row_builder.convert_columns(&arrays)?; - - if all_scalar { - hasher.update(rows.row(0).as_ref()); - let result = hasher.digest128().to_be_bytes().to_vec(); - hasher.reset(); - Ok(ColumnarValue::Scalar(ScalarValue::FixedSizeBinary( - size_of::() as i32, - Some(result), - ))) - } else { - let mut builder = - FixedSizeBinaryBuilder::with_capacity(length, size_of::() as i32); - - for row in rows.iter() { - hasher.update(row.as_ref()); - builder.append_value(hasher.digest128().to_be_bytes())?; - hasher.reset(); - } - - Ok(ColumnarValue::Array(Arc::new(builder.finish()))) - } - } -} - -impl Default for MultiHashFunction { - fn default() -> Self { - Self { - signature: Signature::new(TypeSignature::VariadicAny, Volatility::Immutable), - } - } -} - -impl ScalarUDFImpl for MultiHashFunction { - fn as_any(&self) -> &dyn Any { - self - } - - fn name(&self) -> &str { - scalar_fn::MULTI_HASH - } - - fn signature(&self) -> &Signature { - &self.signature - } - - fn return_type(&self, _arg_types: &[DataType]) -> Result { - Ok(DataType::FixedSizeBinary(size_of::() as i32)) - } - - fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - self.invoke(&args.args) - } -} - -fn json_function( - name: &str, - f: F, - to_scalar: ToS, - args: &[ColumnarValue], -) -> Result -where - ArrayT: Array + FromIterator> + 'static, - F: Fn(serde_json::Value, &JsonPath) -> Option, - ToS: Fn(Option) -> ScalarValue, -{ - assert_eq!(args.len(), 2); - Ok(match (&args[0], &args[1]) { - (ColumnarValue::Array(values), ColumnarValue::Scalar(path)) => { - let path = parse_path(name, path)?; - let vs = as_string_array(values); - ColumnarValue::Array(Arc::new( - vs.iter() - .map(|s| s.and_then(|s| f(serde_json::from_str(s).ok()?, &path))) - .collect::(), - ) as ArrayRef) - } - (ColumnarValue::Scalar(value), ColumnarValue::Scalar(path)) => { - let path = parse_path(name, path)?; - let ScalarValue::Utf8(value) = value else { - return Err(DataFusionError::Execution(format!( - "The value argument to {name} must be of type TEXT" - ))); - }; - - let result = value - .as_ref() - .and_then(|v| f(serde_json::from_str(v).ok()?, &path)); - ColumnarValue::Scalar(to_scalar(result)) - } - _ => { - return Err(DataFusionError::Execution( - "The path argument to {name} must be a literal".to_string(), - )); - } - }) -} - -pub fn extract_json(args: &[ColumnarValue]) -> Result { - assert_eq!(args.len(), 2); - - let inner = |s, path: &JsonPath| { - Some( - path.query(&serde_json::from_str(s).ok()?) - .iter() - .map(|v| Some(v.to_string())) - .collect::>>(), - ) - }; - - Ok(match (&args[0], &args[1]) { - (ColumnarValue::Array(values), ColumnarValue::Scalar(path)) => { - let path = parse_path("extract_json", path)?; - let values = as_string_array(values); - - let mut builder = ListBuilder::with_capacity(StringBuilder::new(), values.len()); - - let queried = values.iter().map(|s| s.and_then(|s| inner(s, &path))); - - for v in queried { - builder.append_option(v); - } - - ColumnarValue::Array(Arc::new(builder.finish())) - } - (ColumnarValue::Scalar(value), ColumnarValue::Scalar(path)) => { - let path = parse_path("extract_json", path)?; - let ScalarValue::Utf8(v) = value else { - return Err(DataFusionError::Execution( - "The value argument to extract_json must be of type TEXT".to_string(), - )); - }; - - let mut builder = ListBuilder::with_capacity(StringBuilder::new(), 1); - let result = v.as_ref().and_then(|s| inner(s, &path)); - builder.append_option(result); - - ColumnarValue::Scalar(ScalarValue::List(Arc::new(builder.finish()))) - } - _ => { - return Err(DataFusionError::Execution( - "The path argument to extract_json must be a literal".to_string(), - )); - } - }) -} - -pub fn get_first_json_object(args: &[ColumnarValue]) -> Result { - json_function::( - "get_first_json_object", - |s, path| path.query(&s).first().map(|v| v.to_string()), - |s| s.as_deref().into(), - args, - ) -} - -pub fn extract_json_string(args: &[ColumnarValue]) -> Result { - json_function::( - "extract_json_string", - |s, path| { - path.query(&s) - .first() - .and_then(|v| v.as_str().map(|s| s.to_string())) - }, - |s| s.as_deref().into(), - args, - ) -} - -// This code is vendored from -// https://github.com/datafusion-contrib/datafusion-functions-json/blob/main/src/common_union.rs -// as the `is_json_union` function is not public. It should be kept in sync with that code so -// that we are able to detect JSON unions and rewrite them to serialized JSON for sinks. -pub(crate) fn is_json_union(data_type: &DataType) -> bool { - match data_type { - DataType::Union(fields, UnionMode::Sparse) => fields == &union_fields(), - _ => false, - } -} - -pub(crate) const TYPE_ID_NULL: i8 = 0; -const TYPE_ID_BOOL: i8 = 1; -const TYPE_ID_INT: i8 = 2; -const TYPE_ID_FLOAT: i8 = 3; -const TYPE_ID_STR: i8 = 4; -const TYPE_ID_ARRAY: i8 = 5; -const TYPE_ID_OBJECT: i8 = 6; - -fn union_fields() -> UnionFields { - static FIELDS: OnceLock = OnceLock::new(); - FIELDS - .get_or_init(|| { - let json_metadata: HashMap = - HashMap::from_iter(vec![("is_json".to_string(), "true".to_string())]); - UnionFields::from_iter([ - ( - TYPE_ID_NULL, - Arc::new(Field::new("null", DataType::Null, true)), - ), - ( - TYPE_ID_BOOL, - Arc::new(Field::new("bool", DataType::Boolean, false)), - ), - ( - TYPE_ID_INT, - Arc::new(Field::new("int", DataType::Int64, false)), - ), - ( - TYPE_ID_FLOAT, - Arc::new(Field::new("float", DataType::Float64, false)), - ), - ( - TYPE_ID_STR, - Arc::new(Field::new("str", DataType::Utf8, false)), - ), - ( - TYPE_ID_ARRAY, - Arc::new( - Field::new("array", DataType::Utf8, false) - .with_metadata(json_metadata.clone()), - ), - ), - ( - TYPE_ID_OBJECT, - Arc::new( - Field::new("object", DataType::Utf8, false) - .with_metadata(json_metadata.clone()), - ), - ), - ]) - }) - .clone() -} -// End vendored code - -pub fn serialize_json_union(args: &[ColumnarValue]) -> Result { - assert_eq!(args.len(), 1); - let array = match args.first().unwrap() { - ColumnarValue::Array(a) => a.clone(), - ColumnarValue::Scalar(s) => s.to_array_of_size(1)?, - }; - - let mut b = StringBuilder::with_capacity(array.len(), array.get_array_memory_size()); - - write_union(&mut b, &array)?; - - Ok(ColumnarValue::Array(Arc::new(b.finish()))) -} - -fn write_union(b: &mut StringBuilder, array: &ArrayRef) -> Result<(), std::fmt::Error> { - assert!( - is_json_union(array.data_type()), - "array item is not a valid JSON union" - ); - let json_union = array.as_any().downcast_ref::().unwrap(); - - for i in 0..json_union.len() { - if json_union.is_null(i) { - b.append_null(); - } else { - write_value(b, json_union.type_id(i), &json_union.value(i))?; - b.append_value(""); - } - } - - Ok(()) -} - -fn write_value(b: &mut StringBuilder, id: i8, a: &ArrayRef) -> Result<(), std::fmt::Error> { - match id { - TYPE_ID_NULL => write!(b, "null")?, - TYPE_ID_BOOL => write!(b, "{}", a.as_boolean().value(0))?, - TYPE_ID_INT => write!(b, "{}", a.as_primitive::().value(0))?, - TYPE_ID_FLOAT => write!(b, "{}", a.as_primitive::().value(0))?, - TYPE_ID_STR => { - // assumes that this is already a valid (escaped) json string as the only way to - // construct these values are by parsing (valid) JSON - b.write_char('"')?; - b.write_str(a.as_string::().value(0))?; - b.write_char('"')?; - } - TYPE_ID_ARRAY => { - b.write_str(a.as_string::().value(0))?; - } - TYPE_ID_OBJECT => { - b.write_str(a.as_string::().value(0))?; - } - _ => unreachable!("invalid union type in JSON union: {}", id), - } - - Ok(()) -} - -pub(crate) fn serialize_outgoing_json( - registry: &StreamSchemaProvider, - node: Arc, -) -> LogicalPlan { - let exprs = node - .schema() - .fields() - .iter() - .map(|f| { - if is_json_union(f.data_type()) { - Expr::Alias(Alias::new( - Expr::ScalarFunction(ScalarFunction::new_udf( - registry.udf(scalar_fn::SERIALIZE_JSON_UNION).unwrap(), - vec![col(f.name())], - )), - Option::::None, - f.name(), - )) - } else { - col(f.name()) - } - }) - .collect(); - - LogicalPlan::Projection(Projection::try_new(exprs, node).unwrap()) -} - -#[cfg(test)] -mod test { - use datafusion::arrow::array::StringArray; - use datafusion::arrow::array::builder::{ListBuilder, StringBuilder}; - use datafusion::common::ScalarValue; - use std::sync::Arc; - - #[test] - fn test_extract_json() { - let input = Arc::new(StringArray::from(vec![ - r#"{"a": 1, "b": 2, "c": { "d": "hello" }}"#, - r#"{"a": 3, "b": 4}"#, - r#"{"a": 5, "b": 6}"#, - ])); - - let path = "$.c.d"; - - let result = super::extract_json(&[ - super::ColumnarValue::Array(input), - super::ColumnarValue::Scalar(path.into()), - ]) - .unwrap(); - - let mut expected = ListBuilder::new(StringBuilder::new()); - expected.append_value(vec![Some("\"hello\"".to_string())]); - expected.append_value(Vec::>::new()); - expected.append_value(Vec::>::new()); - if let super::ColumnarValue::Array(result) = result { - assert_eq!(*result, expected.finish()); - } else { - panic!("Expected array, got scalar"); - } - - let result = super::extract_json(&[ - super::ColumnarValue::Scalar(r#"{"a": 1, "b": 2, "c": { "d": "hello" }}"#.into()), - super::ColumnarValue::Scalar(path.into()), - ]) - .unwrap(); - - let mut expected = ListBuilder::with_capacity(StringBuilder::new(), 1); - expected.append_value(vec![Some("\"hello\"".to_string())]); - - if let super::ColumnarValue::Scalar(ScalarValue::List(result)) = result { - assert_eq!(*result, expected.finish()); - } else { - panic!("Expected scalar list"); - } - } - - #[test] - fn test_get_first_json_object() { - let input = Arc::new(StringArray::from(vec![ - r#"{"a": 1, "b": 2}"#, - r#"{"a": 3}"#, - r#"{"a": 5, "b": 6}"#, - ])); - - let path = "$.b"; - - let result = super::get_first_json_object(&[ - super::ColumnarValue::Array(input), - super::ColumnarValue::Scalar(path.into()), - ]) - .unwrap(); - - let expected = StringArray::from(vec![Some("2"), None, Some("6")]); - - if let super::ColumnarValue::Array(result) = result { - assert_eq!(*result, expected); - } else { - panic!("Expected array, got scalar"); - } - - let result = super::get_first_json_object(&[ - super::ColumnarValue::Scalar(r#"{"a": 1, "b": 2, "c": { "d": "hello" }}"#.into()), - super::ColumnarValue::Scalar("$.c.d".into()), - ]) - .unwrap(); - - let expected = ScalarValue::Utf8(Some("\"hello\"".to_string())); - - if let super::ColumnarValue::Scalar(result) = result { - assert_eq!(result, expected); - } else { - panic!("Expected scalar"); - } - } - - #[test] - fn test_extract_json_string() { - let input = Arc::new(StringArray::from(vec![ - r#"{"a": 1, "b": 2, "c": { "d": "hello" }}"#, - r#"{"a": 3, "b": 4}"#, - r#"{"a": 5, "b": 6}"#, - ])); - - let path = "$.c.d"; - - let result = super::extract_json_string(&[ - super::ColumnarValue::Array(input), - super::ColumnarValue::Scalar(path.into()), - ]) - .unwrap(); - - let expected = StringArray::from(vec![Some("hello"), None, None]); - - if let super::ColumnarValue::Array(result) = result { - assert_eq!(*result, expected); - } else { - panic!("Expected array, got scalar"); - } - - let result = super::extract_json_string(&[ - super::ColumnarValue::Scalar(r#"{"a": 1, "b": 2, "c": { "d": "hello" }}"#.into()), - super::ColumnarValue::Scalar(path.into()), - ]) - .unwrap(); - - let expected = ScalarValue::Utf8(Some("hello".to_string())); - - if let super::ColumnarValue::Scalar(result) = result { - assert_eq!(result, expected); - } else { - panic!("Expected scalar"); - } - } -} diff --git a/src/sql/logical_node/aggregate.rs b/src/sql/logical_node/aggregate.rs deleted file mode 100644 index 1e288ab5..00000000 --- a/src/sql/logical_node/aggregate.rs +++ /dev/null @@ -1,644 +0,0 @@ -// 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::fmt::Formatter; -use std::sync::Arc; -use std::time::Duration; - -use arrow_array::types::IntervalMonthDayNanoType; -use datafusion::common::{Column, DFSchemaRef, Result, ScalarValue, internal_err}; -use datafusion::logical_expr::{ - self, BinaryExpr, Expr, Extension, LogicalPlan, UserDefinedLogicalNodeCore, - expr::ScalarFunction, -}; -use datafusion_common::{DFSchema, DataFusionError, plan_err}; -use datafusion_expr::Aggregate; -use datafusion_proto::physical_plan::to_proto::serialize_physical_expr; -use datafusion_proto::physical_plan::{AsExecutionPlan, DefaultPhysicalExtensionCodec}; -use datafusion_proto::protobuf::PhysicalPlanNode; -use prost::Message; -use protocol::function_stream_graph::{ - SessionWindowAggregateOperator, SlidingWindowAggregateOperator, TumblingWindowAggregateOperator, -}; - -use crate::multifield_partial_ord; -use crate::sql::common::constants::{extension_node, proto_operator_name}; -use crate::sql::common::{FsSchema, FsSchemaRef}; -use crate::sql::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; -use crate::sql::logical_node::{ - CompiledTopologyNode, StreamingOperatorBlueprint, SystemTimestampInjectorNode, -}; -use crate::sql::logical_planner::planner::{NamedNode, Planner, SplitPlanOutput}; -use crate::sql::physical::{StreamingExtensionCodec, window}; -use crate::sql::types::{ - QualifiedField, TIMESTAMP_FIELD, WindowBehavior, WindowType, build_df_schema, - build_df_schema_with_metadata, extract_qualified_fields, -}; - -pub(crate) const STREAM_AGG_EXTENSION_NAME: &str = extension_node::STREAM_WINDOW_AGGREGATE; - -/// Represents a streaming windowed aggregation node in the logical plan. -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub(crate) struct StreamWindowAggregateNode { - pub(crate) window_spec: WindowBehavior, - pub(crate) base_agg_plan: LogicalPlan, - pub(crate) output_schema: DFSchemaRef, - pub(crate) partition_keys: Vec, - pub(crate) post_aggregation_plan: LogicalPlan, -} - -multifield_partial_ord!( - StreamWindowAggregateNode, - base_agg_plan, - partition_keys, - post_aggregation_plan -); - -impl StreamWindowAggregateNode { - /// This node is only emitted after `KeyExtractionNode` in streaming rewrites; `partition_keys` - /// may be empty when GROUP BY is only a window call (window column stripped from key list), - /// but the pipeline still consumes a shuffle — use keyed aggregate parallelism. - fn parallelism_after_keyed_shuffle(&self, planner: &Planner) -> usize { - planner.keyed_aggregate_parallelism() - } - - /// Safely constructs a new node, computing the final projection without panicking. - pub fn try_new( - window_spec: WindowBehavior, - base_agg_plan: LogicalPlan, - partition_keys: Vec, - ) -> Result { - let post_aggregation_plan = - WindowBoundaryMath::build_post_aggregation(&base_agg_plan, window_spec.clone())?; - - Ok(Self { - window_spec, - base_agg_plan, - output_schema: post_aggregation_plan.schema().clone(), - partition_keys, - post_aggregation_plan, - }) - } - - fn build_tumbling_operator( - &self, - planner: &Planner, - node_id: usize, - input_schema: DFSchemaRef, - duration: Duration, - ) -> Result { - let binning_expr = planner.binning_function_proto(duration, input_schema.clone())?; - - let SplitPlanOutput { - partial_aggregation_plan, - partial_schema, - finish_plan, - } = planner.split_physical_plan(self.partition_keys.clone(), &self.base_agg_plan, true)?; - - let final_physical = planner.sync_plan(&self.post_aggregation_plan)?; - let final_physical_proto = PhysicalPlanNode::try_from_physical_plan( - final_physical, - &StreamingExtensionCodec::default(), - )?; - - let operator_config = TumblingWindowAggregateOperator { - name: proto_operator_name::TUMBLING_WINDOW.to_string(), - width_micros: duration.as_micros() as u64, - binning_function: binning_expr.encode_to_vec(), - input_schema: Some( - FsSchema::from_schema_keys( - Arc::new(input_schema.as_ref().into()), - self.partition_keys.clone(), - )? - .into(), - ), - partial_schema: Some(partial_schema.into()), - partial_aggregation_plan: partial_aggregation_plan.encode_to_vec(), - final_aggregation_plan: finish_plan.encode_to_vec(), - final_projection: Some(final_physical_proto.encode_to_vec()), - }; - - Ok(LogicalNode::single( - node_id as u32, - format!("tumbling_{node_id}"), - OperatorName::TumblingWindowAggregate, - operator_config.encode_to_vec(), - format!("TumblingWindow<{}>", operator_config.name), - self.parallelism_after_keyed_shuffle(planner), - )) - } - - fn build_sliding_operator( - &self, - planner: &Planner, - node_id: usize, - input_schema: DFSchemaRef, - duration: Duration, - slide_interval: Duration, - ) -> Result { - let binning_expr = planner.binning_function_proto(slide_interval, input_schema.clone())?; - - let SplitPlanOutput { - partial_aggregation_plan, - partial_schema, - finish_plan, - } = planner.split_physical_plan(self.partition_keys.clone(), &self.base_agg_plan, true)?; - - let final_physical = planner.sync_plan(&self.post_aggregation_plan)?; - let final_physical_proto = PhysicalPlanNode::try_from_physical_plan( - final_physical, - &StreamingExtensionCodec::default(), - )?; - - let operator_config = SlidingWindowAggregateOperator { - name: format!("SlidingWindow<{duration:?}>"), - width_micros: duration.as_micros() as u64, - slide_micros: slide_interval.as_micros() as u64, - binning_function: binning_expr.encode_to_vec(), - input_schema: Some( - FsSchema::from_schema_keys( - Arc::new(input_schema.as_ref().into()), - self.partition_keys.clone(), - )? - .into(), - ), - partial_schema: Some(partial_schema.into()), - partial_aggregation_plan: partial_aggregation_plan.encode_to_vec(), - final_aggregation_plan: finish_plan.encode_to_vec(), - final_projection: final_physical_proto.encode_to_vec(), - }; - - Ok(LogicalNode::single( - node_id as u32, - format!("sliding_window_{node_id}"), - OperatorName::SlidingWindowAggregate, - operator_config.encode_to_vec(), - proto_operator_name::SLIDING_WINDOW_LABEL.to_string(), - self.parallelism_after_keyed_shuffle(planner), - )) - } - - fn build_session_operator( - &self, - planner: &Planner, - node_id: usize, - input_schema: DFSchemaRef, - ) -> Result { - let WindowBehavior::FromOperator { - window: WindowType::Session { gap }, - window_index, - window_field, - is_nested: false, - } = &self.window_spec - else { - return plan_err!("Expected standard session window configuration"); - }; - - let output_fields = extract_qualified_fields(self.base_agg_plan.schema()); - let LogicalPlan::Aggregate(base_agg) = self.base_agg_plan.clone() else { - return plan_err!("Base plan must be an Aggregate node"); - }; - - let key_count = self.partition_keys.len(); - let unkeyed_schema = Arc::new(build_df_schema_with_metadata( - &output_fields[key_count..], - self.base_agg_plan.schema().metadata().clone(), - )?); - - let unkeyed_agg_node = Aggregate::try_new_with_schema( - base_agg.input.clone(), - vec![], - base_agg.aggr_expr.clone(), - unkeyed_schema, - )?; - - let physical_agg = planner.sync_plan(&LogicalPlan::Aggregate(unkeyed_agg_node))?; - let physical_agg_proto = PhysicalPlanNode::try_from_physical_plan( - physical_agg, - &StreamingExtensionCodec::default(), - )?; - - let operator_config = SessionWindowAggregateOperator { - name: format!("session_window_{node_id}"), - gap_micros: gap.as_micros() as u64, - window_field_name: window_field.name().to_string(), - window_index: *window_index as u64, - input_schema: Some( - FsSchema::from_schema_keys( - Arc::new(input_schema.as_ref().into()), - self.partition_keys.clone(), - )? - .into(), - ), - unkeyed_aggregate_schema: None, - partial_aggregation_plan: vec![], - final_aggregation_plan: physical_agg_proto.encode_to_vec(), - }; - - Ok(LogicalNode::single( - node_id as u32, - format!("SessionWindow<{gap:?}>"), - OperatorName::SessionWindowAggregate, - operator_config.encode_to_vec(), - operator_config.name.clone(), - self.parallelism_after_keyed_shuffle(planner), - )) - } - - fn build_instant_operator( - &self, - planner: &Planner, - node_id: usize, - input_schema: DFSchemaRef, - apply_final_projection: bool, - ) -> Result { - let ts_column_expr = Expr::Column(Column::new_unqualified(TIMESTAMP_FIELD.to_string())); - let binning_expr = planner.create_physical_expr(&ts_column_expr, &input_schema)?; - let binning_proto = - serialize_physical_expr(&binning_expr, &DefaultPhysicalExtensionCodec {})?; - - let final_projection_payload = if apply_final_projection { - let physical_plan = planner.sync_plan(&self.post_aggregation_plan)?; - let proto_node = PhysicalPlanNode::try_from_physical_plan( - physical_plan, - &StreamingExtensionCodec::default(), - )?; - Some(proto_node.encode_to_vec()) - } else { - None - }; - - let SplitPlanOutput { - partial_aggregation_plan, - partial_schema, - finish_plan, - } = planner.split_physical_plan(self.partition_keys.clone(), &self.base_agg_plan, true)?; - - let operator_config = TumblingWindowAggregateOperator { - name: proto_operator_name::INSTANT_WINDOW.to_string(), - width_micros: 0, - binning_function: binning_proto.encode_to_vec(), - input_schema: Some( - FsSchema::from_schema_keys( - Arc::new(input_schema.as_ref().into()), - self.partition_keys.clone(), - )? - .into(), - ), - partial_schema: Some(partial_schema.into()), - partial_aggregation_plan: partial_aggregation_plan.encode_to_vec(), - final_aggregation_plan: finish_plan.encode_to_vec(), - final_projection: final_projection_payload, - }; - - Ok(LogicalNode::single( - node_id as u32, - format!("instant_window_{node_id}"), - OperatorName::TumblingWindowAggregate, - operator_config.encode_to_vec(), - proto_operator_name::INSTANT_WINDOW_LABEL.to_string(), - self.parallelism_after_keyed_shuffle(planner), - )) - } -} - -impl StreamingOperatorBlueprint for StreamWindowAggregateNode { - fn operator_identity(&self) -> Option { - None - } - - fn compile_to_graph_node( - &self, - planner: &Planner, - node_id: usize, - mut input_schemas: Vec, - ) -> Result { - if input_schemas.len() != 1 { - return plan_err!("StreamWindowAggregateNode requires exactly one input schema"); - } - - let raw_schema = input_schemas.remove(0); - let df_schema = Arc::new(DFSchema::try_from(raw_schema.schema.as_ref().clone())?); - - let logical_operator = match &self.window_spec { - WindowBehavior::FromOperator { - window, is_nested, .. - } => { - if *is_nested { - self.build_instant_operator(planner, node_id, df_schema, true)? - } else { - match window { - WindowType::Tumbling { width } => { - self.build_tumbling_operator(planner, node_id, df_schema, *width)? - } - WindowType::Sliding { width, slide } => self - .build_sliding_operator(planner, node_id, df_schema, *width, *slide)?, - WindowType::Session { .. } => { - self.build_session_operator(planner, node_id, df_schema)? - } - WindowType::Instant => { - return plan_err!( - "Instant window is invalid within standard operator context" - ); - } - } - } - } - WindowBehavior::InData => self - .build_instant_operator(planner, node_id, df_schema, false) - .map_err(|e| e.context("Failed compiling instant window"))?, - }; - - let link = LogicalEdge::project_all(LogicalEdgeType::Shuffle, (*raw_schema).clone()); - Ok(CompiledTopologyNode { - execution_unit: logical_operator, - routing_edges: vec![link], - }) - } - - fn yielded_schema(&self) -> FsSchema { - let schema_ref = (*self.output_schema).clone().into(); - FsSchema::from_schema_unkeyed(Arc::new(schema_ref)) - .expect("StreamWindowAggregateNode output schema must contain timestamp column") - } -} - -impl UserDefinedLogicalNodeCore for StreamWindowAggregateNode { - fn name(&self) -> &str { - STREAM_AGG_EXTENSION_NAME - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - vec![&self.base_agg_plan] - } - - fn schema(&self) -> &DFSchemaRef { - &self.output_schema - } - - fn expressions(&self) -> Vec { - vec![] - } - - fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { - let spec_desc = match &self.window_spec { - WindowBehavior::InData => "InData".to_string(), - WindowBehavior::FromOperator { window, .. } => format!("FromOperator({window:?})"), - }; - write!( - f, - "StreamWindowAggregate: {} | spec: {}", - self.schema(), - spec_desc - ) - } - - fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { - if inputs.len() != 1 { - return internal_err!("StreamWindowAggregateNode expects exactly 1 input"); - } - Self::try_new( - self.window_spec.clone(), - inputs[0].clone(), - self.partition_keys.clone(), - ) - } -} - -// ----------------------------------------------------------------------------- -// Dedicated boundary math for window bin / post-aggregation projection -// ----------------------------------------------------------------------------- - -struct WindowBoundaryMath; - -impl WindowBoundaryMath { - fn interval_nanos(nanos: i64) -> Expr { - Expr::Literal( - ScalarValue::IntervalMonthDayNano(Some(IntervalMonthDayNanoType::make_value( - 0, 0, nanos, - ))), - None, - ) - } - - fn build_post_aggregation( - agg_plan: &LogicalPlan, - window_spec: WindowBehavior, - ) -> Result { - let ts_field: QualifiedField = agg_plan - .inputs() - .first() - .ok_or_else(|| DataFusionError::Plan("Aggregate has no inputs".into()))? - .schema() - .qualified_field_with_unqualified_name(TIMESTAMP_FIELD)? - .into(); - - let plan_with_ts = LogicalPlan::Extension(Extension { - node: Arc::new(SystemTimestampInjectorNode::try_new( - agg_plan.clone(), - ts_field.qualifier().cloned(), - )?), - }); - - let (win_field, win_index, duration, is_nested) = match window_spec { - WindowBehavior::InData => return Ok(plan_with_ts), - WindowBehavior::FromOperator { - window, - window_field, - window_index, - is_nested, - } => match window { - WindowType::Tumbling { width } | WindowType::Sliding { width, .. } => { - (window_field, window_index, width, is_nested) - } - WindowType::Session { .. } => { - return Ok(LogicalPlan::Extension(Extension { - node: Arc::new(InjectWindowFieldNode::try_new( - plan_with_ts, - window_field, - window_index, - )?), - })); - } - WindowType::Instant => return Ok(plan_with_ts), - }, - }; - - if is_nested { - return Self::build_nested_projection(plan_with_ts, win_field, win_index, duration); - } - - let mut output_fields = extract_qualified_fields(agg_plan.schema()); - let mut projections: Vec<_> = output_fields - .iter() - .map(|f| Expr::Column(f.qualified_column())) - .collect(); - - let ts_col_expr = Expr::Column(Column::new(ts_field.qualifier().cloned(), ts_field.name())); - - output_fields.insert(win_index, win_field.clone()); - - let win_func_expr = Expr::ScalarFunction(ScalarFunction { - func: window(), - args: vec![ - ts_col_expr.clone(), - Expr::BinaryExpr(BinaryExpr { - left: Box::new(ts_col_expr.clone()), - op: logical_expr::Operator::Plus, - right: Box::new(Self::interval_nanos(duration.as_nanos() as i64)), - }), - ], - }); - - projections.insert( - win_index, - win_func_expr.alias_qualified(win_field.qualifier().cloned(), win_field.name()), - ); - - output_fields.push(ts_field); - - let bin_end_expr = Expr::BinaryExpr(BinaryExpr { - left: Box::new(ts_col_expr), - op: logical_expr::Operator::Plus, - right: Box::new(Self::interval_nanos((duration.as_nanos() - 1) as i64)), - }); - projections.push(bin_end_expr); - - Ok(LogicalPlan::Projection( - logical_expr::Projection::try_new_with_schema( - projections, - Arc::new(plan_with_ts), - Arc::new(build_df_schema(&output_fields)?), - )?, - )) - } - - fn build_nested_projection( - plan: LogicalPlan, - win_field: QualifiedField, - win_index: usize, - duration: Duration, - ) -> Result { - let ts_field: QualifiedField = plan - .schema() - .qualified_field_with_unqualified_name(TIMESTAMP_FIELD)? - .into(); - let ts_col_expr = Expr::Column(Column::new(ts_field.qualifier().cloned(), ts_field.name())); - - let mut output_fields = extract_qualified_fields(plan.schema()); - let mut projections: Vec<_> = output_fields - .iter() - .map(|f| Expr::Column(f.qualified_column())) - .collect(); - - output_fields.insert(win_index, win_field.clone()); - - let win_func_expr = Expr::ScalarFunction(ScalarFunction { - func: window(), - args: vec![ - Expr::BinaryExpr(BinaryExpr { - left: Box::new(ts_col_expr.clone()), - op: logical_expr::Operator::Minus, - right: Box::new(Self::interval_nanos(duration.as_nanos() as i64 - 1)), - }), - Expr::BinaryExpr(BinaryExpr { - left: Box::new(ts_col_expr), - op: logical_expr::Operator::Plus, - right: Box::new(Self::interval_nanos(1)), - }), - ], - }); - - projections.insert( - win_index, - win_func_expr.alias_qualified(win_field.qualifier().cloned(), win_field.name()), - ); - - Ok(LogicalPlan::Projection( - logical_expr::Projection::try_new_with_schema( - projections, - Arc::new(plan), - Arc::new(build_df_schema(&output_fields)?), - )?, - )) - } -} - -// ----------------------------------------------------------------------------- -// Field injection node (session window column placement) -// ----------------------------------------------------------------------------- - -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -struct InjectWindowFieldNode { - pub(crate) upstream_plan: LogicalPlan, - pub(crate) target_field: QualifiedField, - pub(crate) insertion_index: usize, - pub(crate) new_schema: DFSchemaRef, -} - -multifield_partial_ord!(InjectWindowFieldNode, upstream_plan, insertion_index); - -impl InjectWindowFieldNode { - fn try_new( - upstream_plan: LogicalPlan, - target_field: QualifiedField, - insertion_index: usize, - ) -> Result { - let mut fields = extract_qualified_fields(upstream_plan.schema()); - fields.insert(insertion_index, target_field.clone()); - let meta = upstream_plan.schema().metadata().clone(); - - Ok(Self { - upstream_plan, - target_field, - insertion_index, - new_schema: Arc::new(build_df_schema_with_metadata(&fields, meta)?), - }) - } -} - -impl UserDefinedLogicalNodeCore for InjectWindowFieldNode { - fn name(&self) -> &str { - "InjectWindowFieldNode" - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - vec![&self.upstream_plan] - } - - fn schema(&self) -> &DFSchemaRef { - &self.new_schema - } - - fn expressions(&self) -> Vec { - vec![] - } - - fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - write!( - f, - "InjectWindowField: insert {:?} at offset {}", - self.target_field, self.insertion_index - ) - } - - fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { - if inputs.len() != 1 { - return internal_err!("InjectWindowFieldNode expects exactly 1 input"); - } - Self::try_new( - inputs[0].clone(), - self.target_field.clone(), - self.insertion_index, - ) - } -} diff --git a/src/sql/logical_node/async_udf.rs b/src/sql/logical_node/async_udf.rs deleted file mode 100644 index 1c35398e..00000000 --- a/src/sql/logical_node/async_udf.rs +++ /dev/null @@ -1,247 +0,0 @@ -// 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::fmt::Formatter; -use std::sync::Arc; -use std::time::Duration; - -use datafusion::common::{DFSchemaRef, Result}; -use datafusion::logical_expr::{ - Expr, LogicalPlan, UserDefinedLogicalNode, UserDefinedLogicalNodeCore, -}; -use datafusion_common::{internal_err, plan_err}; -use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec; -use datafusion_proto::physical_plan::to_proto::serialize_physical_expr; -use prost::Message; -use protocol::function_stream_graph::{AsyncUdfOperator, AsyncUdfOrdering}; - -use crate::multifield_partial_ord; -use crate::sql::common::constants::extension_node; -use crate::sql::common::constants::sql_field; -use crate::sql::common::{FsSchema, FsSchemaRef}; -use crate::sql::logical_node::logical::{ - DylibUdfConfig, LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName, -}; -use crate::sql::logical_node::streaming_operator_blueprint::{ - CompiledTopologyNode, StreamingOperatorBlueprint, -}; -use crate::sql::logical_planner::planner::{NamedNode, Planner}; -use crate::sql::types::{QualifiedField, build_df_schema, extract_qualified_fields}; - -pub(crate) const NODE_TYPE_NAME: &str = extension_node::ASYNC_FUNCTION_EXECUTION; - -/// Represents a logical node that executes an external asynchronous function (UDF) -/// and projects the final results into the streaming pipeline. -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub(crate) struct AsyncFunctionExecutionNode { - pub(crate) upstream_plan: Arc, - pub(crate) operator_name: String, - pub(crate) function_config: DylibUdfConfig, - pub(crate) invocation_args: Vec, - pub(crate) result_projections: Vec, - pub(crate) preserve_ordering: bool, - pub(crate) concurrency_limit: usize, - pub(crate) execution_timeout: Duration, - pub(crate) resolved_schema: DFSchemaRef, -} - -multifield_partial_ord!( - AsyncFunctionExecutionNode, - upstream_plan, - operator_name, - function_config, - invocation_args, - result_projections, - preserve_ordering, - concurrency_limit, - execution_timeout -); - -impl AsyncFunctionExecutionNode { - /// Compiles logical expressions into serialized physical protobuf bytes. - fn compile_physical_expressions( - &self, - planner: &Planner, - expressions: &[Expr], - schema_context: &DFSchemaRef, - ) -> Result>> { - expressions - .iter() - .map(|logical_expr| { - let physical_expr = planner.create_physical_expr(logical_expr, schema_context)?; - let serialized = - serialize_physical_expr(&physical_expr, &DefaultPhysicalExtensionCodec {})?; - Ok(serialized.encode_to_vec()) - }) - .collect() - } - - /// Computes the intermediate schema which bridges the upstream output - /// and the raw asynchronous result injected by the UDF execution. - fn compute_intermediate_schema(&self) -> Result { - let mut fields = extract_qualified_fields(self.upstream_plan.schema()); - - let raw_result_field = QualifiedField::new( - None, - sql_field::ASYNC_RESULT, - self.function_config.return_type.clone(), - true, - ); - fields.push(raw_result_field); - - Ok(Arc::new(build_df_schema(&fields)?)) - } - - fn to_protobuf_config( - &self, - compiled_args: Vec>, - compiled_projections: Vec>, - ) -> AsyncUdfOperator { - let ordering_strategy = if self.preserve_ordering { - AsyncUdfOrdering::Ordered - } else { - AsyncUdfOrdering::Unordered - }; - - AsyncUdfOperator { - name: self.operator_name.clone(), - udf: Some(self.function_config.clone().into()), - arg_exprs: compiled_args, - final_exprs: compiled_projections, - ordering: ordering_strategy as i32, - max_concurrency: self.concurrency_limit as u32, - timeout_micros: self.execution_timeout.as_micros() as u64, - } - } -} - -impl StreamingOperatorBlueprint for AsyncFunctionExecutionNode { - fn operator_identity(&self) -> Option { - None - } - - fn compile_to_graph_node( - &self, - planner: &Planner, - node_index: usize, - mut input_schemas: Vec, - ) -> Result { - if input_schemas.len() != 1 { - return plan_err!("AsyncFunctionExecutionNode requires exactly one input schema"); - } - - let compiled_args = self.compile_physical_expressions( - planner, - &self.invocation_args, - self.upstream_plan.schema(), - )?; - - let intermediate_schema = self.compute_intermediate_schema()?; - let compiled_projections = self.compile_physical_expressions( - planner, - &self.result_projections, - &intermediate_schema, - )?; - - let operator_config = self.to_protobuf_config(compiled_args, compiled_projections); - - let logical_node = LogicalNode::single( - node_index as u32, - format!("async_udf_{node_index}"), - OperatorName::AsyncUdf, - operator_config.encode_to_vec(), - format!("AsyncUdf<{}>", self.operator_name), - planner.default_parallelism(), - ); - - let upstream_schema = input_schemas.remove(0); - let data_edge = - LogicalEdge::project_all(LogicalEdgeType::Forward, (*upstream_schema).clone()); - - Ok(CompiledTopologyNode { - execution_unit: logical_node, - routing_edges: vec![data_edge], - }) - } - - fn yielded_schema(&self) -> FsSchema { - let arrow_fields: Vec<_> = self - .resolved_schema - .fields() - .iter() - .map(|f| (**f).clone()) - .collect(); - - FsSchema::from_fields(arrow_fields) - } -} - -impl UserDefinedLogicalNodeCore for AsyncFunctionExecutionNode { - fn name(&self) -> &str { - NODE_TYPE_NAME - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - vec![&self.upstream_plan] - } - - fn schema(&self) -> &DFSchemaRef { - &self.resolved_schema - } - - fn expressions(&self) -> Vec { - self.invocation_args - .iter() - .chain(self.result_projections.iter()) - .cloned() - .collect() - } - - fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { - write!( - f, - "AsyncFunctionExecution<{}>: Concurrency={}, Ordered={}", - self.operator_name, self.concurrency_limit, self.preserve_ordering - ) - } - - fn with_exprs_and_inputs( - &self, - exprs: Vec, - mut inputs: Vec, - ) -> Result { - if inputs.len() != 1 { - return internal_err!( - "AsyncFunctionExecutionNode expects exactly 1 input, but received {}", - inputs.len() - ); - } - - if UserDefinedLogicalNode::expressions(self) != exprs { - return internal_err!( - "Attempted to mutate async UDF expressions during logical planning, which is not supported." - ); - } - - Ok(Self { - upstream_plan: Arc::new(inputs.remove(0)), - operator_name: self.operator_name.clone(), - function_config: self.function_config.clone(), - invocation_args: self.invocation_args.clone(), - result_projections: self.result_projections.clone(), - preserve_ordering: self.preserve_ordering, - concurrency_limit: self.concurrency_limit, - execution_timeout: self.execution_timeout, - resolved_schema: self.resolved_schema.clone(), - }) - } -} diff --git a/src/sql/logical_node/debezium.rs b/src/sql/logical_node/debezium.rs deleted file mode 100644 index 8d69c6ec..00000000 --- a/src/sql/logical_node/debezium.rs +++ /dev/null @@ -1,393 +0,0 @@ -// 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::sync::Arc; - -use arrow_schema::{DataType, Field, Schema}; -use datafusion::common::{ - DFSchema, DFSchemaRef, DataFusionError, Result, TableReference, internal_err, plan_err, -}; -use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; -use datafusion::physical_plan::DisplayAs; - -use crate::multifield_partial_ord; -use crate::sql::common::constants::{cdc, extension_node}; -use crate::sql::common::{FsSchema, FsSchemaRef, UPDATING_META_FIELD}; -use crate::sql::logical_planner::planner::{NamedNode, Planner}; -use crate::sql::physical::updating_meta_field; -use crate::sql::types::TIMESTAMP_FIELD; - -use super::{CompiledTopologyNode, StreamingOperatorBlueprint}; - -// ----------------------------------------------------------------------------- -// Constants & Identifiers -// ----------------------------------------------------------------------------- - -pub(crate) const UNROLL_NODE_NAME: &str = extension_node::UNROLL_DEBEZIUM_PAYLOAD; -pub(crate) const PACK_NODE_NAME: &str = extension_node::PACK_DEBEZIUM_ENVELOPE; - -// ----------------------------------------------------------------------------- -// Core Schema Codec -// ----------------------------------------------------------------------------- - -/// Transforms between flat schemas and Debezium CDC envelopes. -pub(crate) struct DebeziumSchemaCodec; - -impl DebeziumSchemaCodec { - /// Wraps a flat physical schema into a Debezium CDC envelope structure. - pub(crate) fn wrap_into_envelope( - flat_schema: &DFSchemaRef, - qualifier_override: Option, - ) -> Result { - let ts_field = if flat_schema.has_column_with_unqualified_name(TIMESTAMP_FIELD) { - Some( - flat_schema - .field_with_unqualified_name(TIMESTAMP_FIELD)? - .clone(), - ) - } else { - None - }; - - let payload_fields: Vec<_> = flat_schema - .fields() - .iter() - .filter(|f| f.name() != TIMESTAMP_FIELD && f.name() != UPDATING_META_FIELD) - .cloned() - .collect(); - - let payload_struct_type = DataType::Struct(payload_fields.into()); - - let mut envelope_fields = vec![ - Arc::new(Field::new(cdc::BEFORE, payload_struct_type.clone(), true)), - Arc::new(Field::new(cdc::AFTER, payload_struct_type, true)), - Arc::new(Field::new(cdc::OP, DataType::Utf8, true)), - ]; - - if let Some(ts) = ts_field { - envelope_fields.push(Arc::new(ts)); - } - - let arrow_schema = Schema::new(envelope_fields); - let final_schema = match qualifier_override { - Some(qualifier) => DFSchema::try_from_qualified_schema(qualifier, &arrow_schema)?, - None => DFSchema::try_from(arrow_schema)?, - }; - - Ok(Arc::new(final_schema)) - } -} - -// ----------------------------------------------------------------------------- -// Logical Node: Unroll Debezium Payload -// ----------------------------------------------------------------------------- - -/// Decodes an incoming Debezium envelope into a flat, updating stream representation. -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct UnrollDebeziumPayloadNode { - upstream_plan: LogicalPlan, - resolved_schema: DFSchemaRef, - pub pk_indices: Vec, - pk_names: Arc>, -} - -multifield_partial_ord!( - UnrollDebeziumPayloadNode, - upstream_plan, - pk_indices, - pk_names -); - -impl UnrollDebeziumPayloadNode { - pub fn try_new(upstream_plan: LogicalPlan, pk_names: Arc>) -> Result { - let input_schema = upstream_plan.schema(); - - let (before_idx, after_idx) = Self::validate_envelope_structure(input_schema)?; - - let payload_fields = Self::extract_payload_fields(input_schema, before_idx)?; - - let pk_indices = Self::map_primary_keys(payload_fields, &pk_names)?; - - let qualifier = Self::resolve_schema_qualifier(input_schema, before_idx, after_idx)?; - - let resolved_schema = - Self::compile_unrolled_schema(input_schema, payload_fields, qualifier)?; - - Ok(Self { - upstream_plan, - resolved_schema, - pk_indices, - pk_names, - }) - } - - fn validate_envelope_structure(schema: &DFSchemaRef) -> Result<(usize, usize)> { - let before_idx = schema - .index_of_column_by_name(None, cdc::BEFORE) - .ok_or_else(|| { - DataFusionError::Plan("Missing 'before' state column in CDC stream".into()) - })?; - - let after_idx = schema - .index_of_column_by_name(None, cdc::AFTER) - .ok_or_else(|| { - DataFusionError::Plan("Missing 'after' state column in CDC stream".into()) - })?; - - let op_idx = schema - .index_of_column_by_name(None, cdc::OP) - .ok_or_else(|| { - DataFusionError::Plan("Missing 'op' operation column in CDC stream".into()) - })?; - - let before_type = schema.field(before_idx).data_type(); - let after_type = schema.field(after_idx).data_type(); - - if before_type != after_type { - return plan_err!( - "State column type mismatch: 'before' is {before_type}, but 'after' is {after_type}" - ); - } - - if *schema.field(op_idx).data_type() != DataType::Utf8 { - return plan_err!("The '{}' column must be of type Utf8", cdc::OP); - } - - Ok((before_idx, after_idx)) - } - - fn extract_payload_fields( - schema: &DFSchemaRef, - state_idx: usize, - ) -> Result<&arrow_schema::Fields> { - match schema.field(state_idx).data_type() { - DataType::Struct(fields) => Ok(fields), - other => plan_err!("State columns must be of type Struct, found {other}"), - } - } - - fn map_primary_keys(fields: &arrow_schema::Fields, pk_names: &[String]) -> Result> { - pk_names - .iter() - .map(|pk| fields.find(pk).map(|(idx, _)| idx)) - .collect::>>() - .ok_or_else(|| { - DataFusionError::Plan("Specified primary key not found in payload schema".into()) - }) - } - - fn resolve_schema_qualifier( - schema: &DFSchemaRef, - before_idx: usize, - after_idx: usize, - ) -> Result> { - let before_qualifier = schema.qualified_field(before_idx).0; - let after_qualifier = schema.qualified_field(after_idx).0; - - match (before_qualifier, after_qualifier) { - (Some(bq), Some(aq)) if bq == aq => Ok(Some(bq.clone())), - (None, None) => Ok(None), - _ => plan_err!("'before' and 'after' columns must share the same namespace/qualifier"), - } - } - - fn compile_unrolled_schema( - original_schema: &DFSchemaRef, - payload_fields: &arrow_schema::Fields, - qualifier: Option, - ) -> Result { - let mut flat_fields = payload_fields.to_vec(); - - flat_fields.push(updating_meta_field()); - - let ts_idx = original_schema - .index_of_column_by_name(None, TIMESTAMP_FIELD) - .ok_or_else(|| { - DataFusionError::Plan(format!( - "Required event time field '{TIMESTAMP_FIELD}' is missing" - )) - })?; - - flat_fields.push(Arc::new(original_schema.field(ts_idx).clone())); - - let arrow_schema = Schema::new(flat_fields); - let compiled_schema = match qualifier { - Some(q) => DFSchema::try_from_qualified_schema(q, &arrow_schema)?, - None => DFSchema::try_from(arrow_schema)?, - }; - - Ok(Arc::new(compiled_schema)) - } -} - -impl UserDefinedLogicalNodeCore for UnrollDebeziumPayloadNode { - fn name(&self) -> &str { - UNROLL_NODE_NAME - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - vec![&self.upstream_plan] - } - - fn schema(&self) -> &DFSchemaRef { - &self.resolved_schema - } - - fn expressions(&self) -> Vec { - vec![] - } - - fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - write!(f, "UnrollDebeziumPayload") - } - - fn with_exprs_and_inputs( - &self, - _exprs: Vec, - mut inputs: Vec, - ) -> Result { - if inputs.len() != 1 { - return internal_err!( - "UnrollDebeziumPayloadNode expects exactly 1 input, got {}", - inputs.len() - ); - } - Self::try_new(inputs.remove(0), self.pk_names.clone()) - } -} - -impl StreamingOperatorBlueprint for UnrollDebeziumPayloadNode { - fn operator_identity(&self) -> Option { - None - } - - fn is_passthrough_boundary(&self) -> bool { - true - } - - fn compile_to_graph_node( - &self, - _: &Planner, - _: usize, - _: Vec, - ) -> Result { - plan_err!( - "UnrollDebeziumPayloadNode is a logical boundary and should not be physically planned" - ) - } - - fn yielded_schema(&self) -> FsSchema { - FsSchema::from_schema_unkeyed(Arc::new(self.resolved_schema.as_ref().into())) - .unwrap_or_else(|_| { - panic!("Failed to extract physical schema for {}", UNROLL_NODE_NAME) - }) - } -} - -// ----------------------------------------------------------------------------- -// Logical Node: Pack Debezium Envelope -// ----------------------------------------------------------------------------- - -/// Encodes a flat updating stream back into a Debezium CDC envelope representation. -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub(crate) struct PackDebeziumEnvelopeNode { - upstream_plan: Arc, - envelope_schema: DFSchemaRef, -} - -multifield_partial_ord!(PackDebeziumEnvelopeNode, upstream_plan); - -impl PackDebeziumEnvelopeNode { - pub(crate) fn try_new(upstream_plan: LogicalPlan) -> Result { - let envelope_schema = DebeziumSchemaCodec::wrap_into_envelope(upstream_plan.schema(), None) - .map_err(|e| { - DataFusionError::Plan(format!("Failed to compile Debezium envelope schema: {e}")) - })?; - - Ok(Self { - upstream_plan: Arc::new(upstream_plan), - envelope_schema, - }) - } -} - -impl DisplayAs for PackDebeziumEnvelopeNode { - fn fmt_as( - &self, - _t: datafusion::physical_plan::DisplayFormatType, - f: &mut std::fmt::Formatter, - ) -> std::fmt::Result { - write!(f, "PackDebeziumEnvelope") - } -} - -impl UserDefinedLogicalNodeCore for PackDebeziumEnvelopeNode { - fn name(&self) -> &str { - PACK_NODE_NAME - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - vec![&self.upstream_plan] - } - - fn schema(&self) -> &DFSchemaRef { - &self.envelope_schema - } - - fn expressions(&self) -> Vec { - vec![] - } - - fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - write!(f, "PackDebeziumEnvelope") - } - - fn with_exprs_and_inputs( - &self, - _exprs: Vec, - mut inputs: Vec, - ) -> Result { - if inputs.len() != 1 { - return internal_err!( - "PackDebeziumEnvelopeNode expects exactly 1 input, got {}", - inputs.len() - ); - } - Self::try_new(inputs.remove(0)) - } -} - -impl StreamingOperatorBlueprint for PackDebeziumEnvelopeNode { - fn operator_identity(&self) -> Option { - None - } - - fn is_passthrough_boundary(&self) -> bool { - true - } - - fn compile_to_graph_node( - &self, - _: &Planner, - _: usize, - _: Vec, - ) -> Result { - internal_err!( - "PackDebeziumEnvelopeNode is a logical boundary and should not be physically planned" - ) - } - - fn yielded_schema(&self) -> FsSchema { - FsSchema::from_schema_unkeyed(Arc::new(self.envelope_schema.as_ref().into())) - .unwrap_or_else(|_| panic!("Failed to extract physical schema for {}", PACK_NODE_NAME)) - } -} diff --git a/src/sql/logical_node/extension_try_from.rs b/src/sql/logical_node/extension_try_from.rs deleted file mode 100644 index 32b12d6c..00000000 --- a/src/sql/logical_node/extension_try_from.rs +++ /dev/null @@ -1,70 +0,0 @@ -// 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::sync::Arc; - -use datafusion::common::{DataFusionError, Result}; -use datafusion::logical_expr::UserDefinedLogicalNode; - -use crate::sql::logical_node::aggregate::StreamWindowAggregateNode; -use crate::sql::logical_node::async_udf::AsyncFunctionExecutionNode; -use crate::sql::logical_node::debezium::{PackDebeziumEnvelopeNode, UnrollDebeziumPayloadNode}; -use crate::sql::logical_node::join::StreamingJoinNode; -use crate::sql::logical_node::key_calculation::KeyExtractionNode; -use crate::sql::logical_node::lookup::StreamReferenceJoinNode; -use crate::sql::logical_node::projection::StreamProjectionNode; -use crate::sql::logical_node::remote_table::RemoteTableBoundaryNode; -use crate::sql::logical_node::sink::StreamEgressNode; -use crate::sql::logical_node::streaming_operator_blueprint::StreamingOperatorBlueprint; -use crate::sql::logical_node::table_source::StreamIngestionNode; -use crate::sql::logical_node::updating_aggregate::ContinuousAggregateNode; -use crate::sql::logical_node::watermark_node::EventTimeWatermarkNode; -use crate::sql::logical_node::windows_function::StreamingWindowFunctionNode; - -fn try_from_t( - node: &dyn UserDefinedLogicalNode, -) -> std::result::Result<&dyn StreamingOperatorBlueprint, ()> { - node.as_any() - .downcast_ref::() - .map(|t| t as &dyn StreamingOperatorBlueprint) - .ok_or(()) -} - -impl<'a> TryFrom<&'a dyn UserDefinedLogicalNode> for &'a dyn StreamingOperatorBlueprint { - type Error = DataFusionError; - - fn try_from(node: &'a dyn UserDefinedLogicalNode) -> Result { - try_from_t::(node) - .or_else(|_| try_from_t::(node)) - .or_else(|_| try_from_t::(node)) - .or_else(|_| try_from_t::(node)) - .or_else(|_| try_from_t::(node)) - .or_else(|_| try_from_t::(node)) - .or_else(|_| try_from_t::(node)) - .or_else(|_| try_from_t::(node)) - .or_else(|_| try_from_t::(node)) - .or_else(|_| try_from_t::(node)) - .or_else(|_| try_from_t::(node)) - .or_else(|_| try_from_t::(node)) - .or_else(|_| try_from_t::(node)) - .or_else(|_| try_from_t::(node)) - .map_err(|_| DataFusionError::Plan(format!("unexpected node: {}", node.name()))) - } -} - -impl<'a> TryFrom<&'a Arc> for &'a dyn StreamingOperatorBlueprint { - type Error = DataFusionError; - - fn try_from(node: &'a Arc) -> Result { - TryFrom::try_from(node.as_ref()) - } -} diff --git a/src/sql/logical_node/is_retract.rs b/src/sql/logical_node/is_retract.rs deleted file mode 100644 index 4370f6ae..00000000 --- a/src/sql/logical_node/is_retract.rs +++ /dev/null @@ -1,82 +0,0 @@ -// 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::sync::Arc; - -use datafusion::arrow::datatypes::{DataType, TimeUnit}; -use datafusion::common::{DFSchemaRef, Result, TableReference}; -use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; - -use crate::multifield_partial_ord; -use crate::sql::physical::updating_meta_field; -use crate::sql::types::{ - QualifiedField, TIMESTAMP_FIELD, build_df_schema, extract_qualified_fields, -}; - -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub(crate) struct IsRetractExtension { - pub(crate) input: LogicalPlan, - pub(crate) schema: DFSchemaRef, - pub(crate) timestamp_qualifier: Option, -} - -multifield_partial_ord!(IsRetractExtension, input, timestamp_qualifier); - -impl IsRetractExtension { - pub(crate) fn new(input: LogicalPlan, timestamp_qualifier: Option) -> Self { - let mut output_fields = extract_qualified_fields(input.schema()); - - let timestamp_index = output_fields.len() - 1; - output_fields[timestamp_index] = QualifiedField::new( - timestamp_qualifier.clone(), - TIMESTAMP_FIELD, - DataType::Timestamp(TimeUnit::Nanosecond, None), - false, - ); - output_fields.push((timestamp_qualifier.clone(), updating_meta_field()).into()); - let schema = Arc::new(build_df_schema(&output_fields).unwrap()); - Self { - input, - schema, - timestamp_qualifier, - } - } -} - -impl UserDefinedLogicalNodeCore for IsRetractExtension { - fn name(&self) -> &str { - "IsRetractExtension" - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - vec![&self.input] - } - - fn schema(&self) -> &DFSchemaRef { - &self.schema - } - - fn expressions(&self) -> Vec { - vec![] - } - - fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - write!(f, "IsRetractExtension") - } - - fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { - Ok(Self::new( - inputs[0].clone(), - self.timestamp_qualifier.clone(), - )) - } -} diff --git a/src/sql/logical_node/join.rs b/src/sql/logical_node/join.rs deleted file mode 100644 index 15631f1f..00000000 --- a/src/sql/logical_node/join.rs +++ /dev/null @@ -1,211 +0,0 @@ -// 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::fmt::Formatter; -use std::time::Duration; - -use datafusion::common::{DFSchemaRef, Result}; -use datafusion::logical_expr::expr::Expr; -use datafusion::logical_expr::{LogicalPlan, UserDefinedLogicalNodeCore}; -use datafusion_common::plan_err; -use datafusion_proto::physical_plan::AsExecutionPlan; -use datafusion_proto::protobuf::PhysicalPlanNode; -use prost::Message; -use protocol::function_stream_graph::JoinOperator; - -use crate::sql::common::constants::{extension_node, runtime_operator_kind}; -use crate::sql::common::{FsSchema, FsSchemaRef}; -use crate::sql::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; -use crate::sql::logical_node::{CompiledTopologyNode, StreamingOperatorBlueprint}; -use crate::sql::logical_planner::planner::{NamedNode, Planner}; -use crate::sql::physical::StreamingExtensionCodec; - -// ----------------------------------------------------------------------------- -// Constants -// ----------------------------------------------------------------------------- - -pub(crate) const STREAM_JOIN_NODE_TYPE: &str = extension_node::STREAMING_JOIN; - -// ----------------------------------------------------------------------------- -// Logical Node Definition -// ----------------------------------------------------------------------------- - -/// A logical plan node representing a streaming join operation. -/// It bridges the DataFusion logical plan with the physical streaming execution engine. -#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd)] -pub struct StreamingJoinNode { - pub(crate) underlying_plan: LogicalPlan, - pub(crate) instant_execution_mode: bool, - pub(crate) state_retention_ttl: Option, -} - -impl StreamingJoinNode { - /// Creates a new instance of the streaming join node. - pub fn new( - underlying_plan: LogicalPlan, - instant_execution_mode: bool, - state_retention_ttl: Option, - ) -> Self { - Self { - underlying_plan, - instant_execution_mode, - state_retention_ttl, - } - } - - /// Compiles the physical execution plan and serializes it into a Protobuf configuration payload. - fn compile_operator_config( - &self, - planner: &Planner, - node_identifier: &str, - left_schema: FsSchemaRef, - right_schema: FsSchemaRef, - ) -> Result { - let physical_plan = planner.sync_plan(&self.underlying_plan)?; - - let proto_node = PhysicalPlanNode::try_from_physical_plan( - physical_plan, - &StreamingExtensionCodec::default(), - )?; - - Ok(JoinOperator { - name: node_identifier.to_string(), - left_schema: Some(left_schema.as_ref().clone().into()), - right_schema: Some(right_schema.as_ref().clone().into()), - output_schema: Some(self.extract_fs_schema().into()), - join_plan: proto_node.encode_to_vec(), - ttl_micros: self.state_retention_ttl.map(|ttl| ttl.as_micros() as u64), - }) - } - - fn determine_operator_type(&self) -> OperatorName { - if self.instant_execution_mode { - OperatorName::InstantJoin - } else { - OperatorName::Join - } - } - - fn extract_fs_schema(&self) -> FsSchema { - FsSchema::from_schema_unkeyed(self.underlying_plan.schema().inner().clone()) - .expect("Fatal: Failed to convert internal join schema to FsSchema without keys") - } -} - -// ----------------------------------------------------------------------------- -// DataFusion Logical Node Core Implementation -// ----------------------------------------------------------------------------- - -impl UserDefinedLogicalNodeCore for StreamingJoinNode { - fn name(&self) -> &str { - STREAM_JOIN_NODE_TYPE - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - vec![&self.underlying_plan] - } - - fn schema(&self) -> &DFSchemaRef { - self.underlying_plan.schema() - } - - fn expressions(&self) -> Vec { - vec![] - } - - fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { - write!( - f, - "StreamingJoinNode: Schema={}, InstantMode={}, TTL={:?}", - self.schema(), - self.instant_execution_mode, - self.state_retention_ttl - ) - } - - fn with_exprs_and_inputs( - &self, - _exprs: Vec, - mut inputs: Vec, - ) -> Result { - if inputs.len() != 1 { - return plan_err!( - "StreamingJoinNode expects exactly 1 underlying logical plan during recreation" - ); - } - - Ok(Self::new( - inputs.remove(0), - self.instant_execution_mode, - self.state_retention_ttl, - )) - } -} - -// ----------------------------------------------------------------------------- -// Streaming Graph Extension Implementation -// ----------------------------------------------------------------------------- - -impl StreamingOperatorBlueprint for StreamingJoinNode { - fn operator_identity(&self) -> Option { - None - } - - fn compile_to_graph_node( - &self, - planner: &Planner, - node_index: usize, - mut input_schemas: Vec, - ) -> Result { - if input_schemas.len() != 2 { - return plan_err!( - "Invalid topology: StreamingJoinNode requires exactly two upstream inputs, received {}", - input_schemas.len() - ); - } - - let right_schema = input_schemas.pop().unwrap(); - let left_schema = input_schemas.pop().unwrap(); - - let node_identifier = format!("stream_join_{node_index}"); - - let operator_config = self.compile_operator_config( - planner, - &node_identifier, - left_schema.clone(), - right_schema.clone(), - )?; - - let logical_node = LogicalNode::single( - node_index as u32, - node_identifier.clone(), - self.determine_operator_type(), - operator_config.encode_to_vec(), - runtime_operator_kind::STREAMING_JOIN.to_string(), - planner.default_parallelism(), - ); - - let left_edge = - LogicalEdge::project_all(LogicalEdgeType::LeftJoin, left_schema.as_ref().clone()); - let right_edge = - LogicalEdge::project_all(LogicalEdgeType::RightJoin, right_schema.as_ref().clone()); - - Ok(CompiledTopologyNode { - execution_unit: logical_node, - routing_edges: vec![left_edge, right_edge], - }) - } - - fn yielded_schema(&self) -> FsSchema { - self.extract_fs_schema() - } -} diff --git a/src/sql/logical_node/key_calculation.rs b/src/sql/logical_node/key_calculation.rs deleted file mode 100644 index ec83e108..00000000 --- a/src/sql/logical_node/key_calculation.rs +++ /dev/null @@ -1,309 +0,0 @@ -// 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::fmt::Formatter; -use std::sync::Arc; - -use datafusion::arrow::datatypes::{Field, Schema}; -use datafusion::common::{DFSchemaRef, Result, internal_err, plan_err}; -use datafusion::logical_expr::{Expr, ExprSchemable, LogicalPlan, UserDefinedLogicalNodeCore}; -use datafusion_common::DFSchema; -use datafusion_expr::col; -use datafusion_proto::physical_plan::to_proto::serialize_physical_expr; -use datafusion_proto::physical_plan::{AsExecutionPlan, DefaultPhysicalExtensionCodec}; -use datafusion_proto::protobuf::PhysicalPlanNode; -use itertools::Itertools; -use prost::Message; - -use protocol::function_stream_graph::{KeyPlanOperator, ProjectionOperator}; - -use crate::multifield_partial_ord; -use crate::sql::common::constants::{extension_node, sql_field}; -use crate::sql::common::{FsSchema, FsSchemaRef}; -use crate::sql::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; -use crate::sql::logical_node::{CompiledTopologyNode, StreamingOperatorBlueprint}; -use crate::sql::logical_planner::planner::{NamedNode, Planner}; -use crate::sql::physical::StreamingExtensionCodec; -use crate::sql::types::{build_df_schema_with_metadata, extract_qualified_fields}; - -pub(crate) const EXTENSION_NODE_IDENTIFIER: &str = extension_node::KEY_EXTRACTION; - -/// Routing strategy for shuffling data across the stream topology. -#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd)] -pub enum KeyExtractionStrategy { - ColumnIndices(Vec), - CalculatedExpressions(Vec), -} - -/// Logical node that computes or extracts routing keys. -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub(crate) struct KeyExtractionNode { - pub(crate) operator_label: Option, - pub(crate) upstream_plan: LogicalPlan, - pub(crate) extraction_strategy: KeyExtractionStrategy, - pub(crate) resolved_schema: DFSchemaRef, -} - -multifield_partial_ord!( - KeyExtractionNode, - operator_label, - upstream_plan, - extraction_strategy -); - -impl KeyExtractionNode { - /// Extracts keys and hides them from the downstream projection. - pub fn try_new_with_projection( - upstream_plan: LogicalPlan, - target_indices: Vec, - label: String, - ) -> Result { - let projected_fields: Vec<_> = extract_qualified_fields(upstream_plan.schema()) - .into_iter() - .enumerate() - .filter(|(idx, _)| !target_indices.contains(idx)) - .map(|(_, field)| field) - .collect(); - - let metadata = upstream_plan.schema().metadata().clone(); - let resolved_schema = build_df_schema_with_metadata(&projected_fields, metadata)?; - - Ok(Self { - operator_label: Some(label), - upstream_plan, - extraction_strategy: KeyExtractionStrategy::ColumnIndices(target_indices), - resolved_schema: Arc::new(resolved_schema), - }) - } - - /// Creates a node using an explicit strategy without changing the visible schema. - pub fn new(upstream_plan: LogicalPlan, strategy: KeyExtractionStrategy) -> Self { - let resolved_schema = upstream_plan.schema().clone(); - Self { - operator_label: None, - upstream_plan, - extraction_strategy: strategy, - resolved_schema, - } - } - - fn compile_index_router( - &self, - physical_plan_proto: PhysicalPlanNode, - indices: &[usize], - ) -> (Vec, OperatorName) { - let operator_config = KeyPlanOperator { - name: sql_field::DEFAULT_KEY_LABEL.into(), - physical_plan: physical_plan_proto.encode_to_vec(), - key_fields: indices.iter().map(|&idx| idx as u64).collect(), - }; - - (operator_config.encode_to_vec(), OperatorName::KeyBy) - } - - fn compile_expression_router( - &self, - planner: &Planner, - expressions: &[Expr], - input_schema_ref: &FsSchemaRef, - input_df_schema: &DFSchemaRef, - ) -> Result<(Vec, OperatorName)> { - let mut target_exprs = expressions.to_vec(); - - for field in input_schema_ref.schema.fields.iter() { - target_exprs.push(col(field.name())); - } - - let output_fs_schema = self.generate_fs_schema()?; - - for (compiled_expr, expected_field) in - target_exprs.iter().zip(output_fs_schema.schema.fields()) - { - let (expr_type, expr_nullable) = - compiled_expr.data_type_and_nullable(input_df_schema)?; - if expr_type != *expected_field.data_type() - || expr_nullable != expected_field.is_nullable() - { - return plan_err!( - "Type mismatch in key calculation: Expected {} (nullable: {}), got {} (nullable: {})", - expected_field.data_type(), - expected_field.is_nullable(), - expr_type, - expr_nullable - ); - } - } - - let mut physical_expr_payloads = Vec::with_capacity(target_exprs.len()); - for logical_expr in target_exprs { - let physical_expr = planner - .create_physical_expr(&logical_expr, input_df_schema) - .map_err(|e| e.context("Failed to physicalize PARTITION BY expression"))?; - - let serialized_expr = - serialize_physical_expr(&physical_expr, &DefaultPhysicalExtensionCodec {})?; - physical_expr_payloads.push(serialized_expr.encode_to_vec()); - } - - let operator_config = ProjectionOperator { - name: self - .operator_label - .as_deref() - .unwrap_or(sql_field::DEFAULT_KEY_LABEL) - .to_string(), - input_schema: Some(input_schema_ref.as_ref().clone().into()), - output_schema: Some(output_fs_schema.into()), - exprs: physical_expr_payloads, - }; - - Ok((operator_config.encode_to_vec(), OperatorName::Projection)) - } - - fn generate_fs_schema(&self) -> Result { - let base_arrow_schema = self.upstream_plan.schema().as_ref(); - - match &self.extraction_strategy { - KeyExtractionStrategy::ColumnIndices(indices) => { - FsSchema::from_schema_keys(Arc::new(base_arrow_schema.into()), indices.clone()) - } - KeyExtractionStrategy::CalculatedExpressions(expressions) => { - let mut composite_fields = - Vec::with_capacity(expressions.len() + base_arrow_schema.fields().len()); - - for (idx, expr) in expressions.iter().enumerate() { - let (data_type, nullable) = expr.data_type_and_nullable(base_arrow_schema)?; - composite_fields - .push(Field::new(format!("__key_{idx}"), data_type, nullable).into()); - } - - for field in base_arrow_schema.fields().iter() { - composite_fields.push(field.clone()); - } - - let final_schema = Arc::new(Schema::new(composite_fields)); - let key_mapping = (1..=expressions.len()).collect_vec(); - FsSchema::from_schema_keys(final_schema, key_mapping) - } - } - } -} - -impl StreamingOperatorBlueprint for KeyExtractionNode { - fn operator_identity(&self) -> Option { - None - } - - fn compile_to_graph_node( - &self, - planner: &Planner, - node_index: usize, - mut input_schemas: Vec, - ) -> Result { - if input_schemas.len() != 1 { - return plan_err!("KeyExtractionNode requires exactly one upstream input schema"); - } - - let input_schema_ref = input_schemas.remove(0); - let input_df_schema = Arc::new(DFSchema::try_from( - input_schema_ref.schema.as_ref().clone(), - )?); - - let physical_plan = planner.sync_plan(&self.upstream_plan)?; - let physical_plan_proto = PhysicalPlanNode::try_from_physical_plan( - physical_plan, - &StreamingExtensionCodec::default(), - )?; - - let (protobuf_payload, engine_operator_name) = match &self.extraction_strategy { - KeyExtractionStrategy::ColumnIndices(indices) => { - self.compile_index_router(physical_plan_proto, indices) - } - KeyExtractionStrategy::CalculatedExpressions(exprs) => { - self.compile_expression_router(planner, exprs, &input_schema_ref, &input_df_schema)? - } - }; - - let logical_node = LogicalNode::single( - node_index as u32, - format!("key_{node_index}"), - engine_operator_name, - protobuf_payload, - format!("Key<{}>", self.operator_label.as_deref().unwrap_or("_")), - planner.key_by_parallelism(), - ); - - let data_edge = - LogicalEdge::project_all(LogicalEdgeType::Forward, (*input_schema_ref).clone()); - - Ok(CompiledTopologyNode { - execution_unit: logical_node, - routing_edges: vec![data_edge], - }) - } - - fn yielded_schema(&self) -> FsSchema { - self.generate_fs_schema() - .expect("Fatal: Failed to generate output schema for KeyExtractionNode") - } -} - -impl UserDefinedLogicalNodeCore for KeyExtractionNode { - fn name(&self) -> &str { - EXTENSION_NODE_IDENTIFIER - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - vec![&self.upstream_plan] - } - - fn schema(&self) -> &DFSchemaRef { - &self.resolved_schema - } - - fn expressions(&self) -> Vec { - vec![] - } - - fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { - write!( - f, - "KeyExtractionNode: Strategy={:?} | Schema={}", - self.extraction_strategy, self.resolved_schema - ) - } - - fn with_exprs_and_inputs( - &self, - exprs: Vec, - mut inputs: Vec, - ) -> Result { - if inputs.len() != 1 { - return internal_err!("KeyExtractionNode requires exactly 1 input logical plan"); - } - - let strategy = match &self.extraction_strategy { - KeyExtractionStrategy::ColumnIndices(indices) => { - KeyExtractionStrategy::ColumnIndices(indices.clone()) - } - KeyExtractionStrategy::CalculatedExpressions(_) => { - KeyExtractionStrategy::CalculatedExpressions(exprs) - } - }; - - Ok(Self { - operator_label: self.operator_label.clone(), - upstream_plan: inputs.remove(0), - extraction_strategy: strategy, - resolved_schema: self.resolved_schema.clone(), - }) - } -} diff --git a/src/sql/logical_node/logical/dylib_udf_config.rs b/src/sql/logical_node/logical/dylib_udf_config.rs deleted file mode 100644 index 9bf3368f..00000000 --- a/src/sql/logical_node/logical/dylib_udf_config.rs +++ /dev/null @@ -1,71 +0,0 @@ -// 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 datafusion::arrow::datatypes::DataType; -use datafusion_proto::protobuf::ArrowType; -use prost::Message; -use protocol::function_stream_graph; - -#[derive(Clone, Debug, Eq, PartialEq, Hash, PartialOrd)] -pub struct DylibUdfConfig { - pub dylib_path: String, - pub arg_types: Vec, - pub return_type: DataType, - pub aggregate: bool, - pub is_async: bool, -} - -impl From for function_stream_graph::DylibUdfConfig { - fn from(from: DylibUdfConfig) -> Self { - function_stream_graph::DylibUdfConfig { - dylib_path: from.dylib_path, - arg_types: from - .arg_types - .iter() - .map(|t| { - ArrowType::try_from(t) - .expect("unsupported data type") - .encode_to_vec() - }) - .collect(), - return_type: ArrowType::try_from(&from.return_type) - .expect("unsupported data type") - .encode_to_vec(), - aggregate: from.aggregate, - is_async: from.is_async, - } - } -} - -impl From for DylibUdfConfig { - fn from(from: function_stream_graph::DylibUdfConfig) -> Self { - DylibUdfConfig { - dylib_path: from.dylib_path, - arg_types: from - .arg_types - .iter() - .map(|t| { - DataType::try_from( - &ArrowType::decode(&mut t.as_slice()).expect("invalid arrow type"), - ) - .expect("invalid arrow type") - }) - .collect(), - return_type: DataType::try_from( - &ArrowType::decode(&mut from.return_type.as_slice()).unwrap(), - ) - .expect("invalid arrow type"), - aggregate: from.aggregate, - is_async: from.is_async, - } - } -} diff --git a/src/sql/logical_node/logical/fs_program_convert.rs b/src/sql/logical_node/logical/fs_program_convert.rs deleted file mode 100644 index b05d68f5..00000000 --- a/src/sql/logical_node/logical/fs_program_convert.rs +++ /dev/null @@ -1,200 +0,0 @@ -// 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::str::FromStr; -use std::sync::Arc; - -use datafusion::common::{DataFusionError, Result as DFResult}; -use petgraph::graph::DiGraph; -use petgraph::prelude::EdgeRef; -use protocol::function_stream_graph::{ - ChainedOperator, EdgeType as ProtoEdgeType, FsEdge, FsNode, FsProgram, - FsSchema as ProtoFsSchema, -}; - -use crate::sql::api::pipelines::{PipelineEdge, PipelineGraph, PipelineNode}; -use crate::sql::common::FsSchema; - -use super::logical_edge::logical_edge_type_from_proto_i32; -use super::operator_chain::{ChainedLogicalOperator, OperatorChain}; -use super::operator_name::OperatorName; -use super::{LogicalEdge, LogicalNode, LogicalProgram, ProgramConfig}; - -impl TryFrom for LogicalProgram { - type Error = DataFusionError; - - fn try_from(value: FsProgram) -> DFResult { - let mut graph = DiGraph::new(); - let mut id_map = HashMap::with_capacity(value.nodes.len()); - - for node in value.nodes { - let operators = node - .operators - .into_iter() - .map(|op| { - let ChainedOperator { - operator_id, - operator_name: name_str, - operator_config, - } = op; - let operator_name = OperatorName::from_str(&name_str).map_err(|_| { - DataFusionError::Plan(format!("Invalid operator name: {name_str}")) - })?; - Ok(ChainedLogicalOperator { - operator_id, - operator_name, - operator_config, - }) - }) - .collect::>>()?; - - let edges = node - .edges - .into_iter() - .map(|e| { - let fs: FsSchema = e.try_into()?; - Ok(Arc::new(fs)) - }) - .collect::>>()?; - - let logical_node = LogicalNode { - node_id: node.node_id, - description: node.description, - operator_chain: OperatorChain { operators, edges }, - parallelism: node.parallelism as usize, - }; - - id_map.insert(node.node_index, graph.add_node(logical_node)); - } - - for edge in value.edges { - let source = *id_map.get(&edge.source).ok_or_else(|| { - DataFusionError::Plan("Graph integrity error: Missing source node".into()) - })?; - let target = *id_map.get(&edge.target).ok_or_else(|| { - DataFusionError::Plan("Graph integrity error: Missing target node".into()) - })?; - let schema = edge.schema.ok_or_else(|| { - DataFusionError::Plan("Graph integrity error: Missing edge schema".into()) - })?; - let edge_type = logical_edge_type_from_proto_i32(edge.edge_type)?; - - graph.add_edge( - source, - target, - LogicalEdge { - edge_type, - schema: Arc::new(FsSchema::try_from(schema)?), - }, - ); - } - - let program_config = value - .program_config - .map(ProgramConfig::from) - .unwrap_or_default(); - - Ok(LogicalProgram::new(graph, program_config)) - } -} - -impl From for FsProgram { - fn from(value: LogicalProgram) -> Self { - let nodes = value - .graph - .node_indices() - .filter_map(|idx| value.graph.node_weight(idx).map(|node| (idx, node))) - .map(|(idx, node)| FsNode { - node_index: idx.index() as i32, - node_id: node.node_id, - parallelism: node.parallelism as u32, - description: node.description.clone(), - operators: node - .operator_chain - .operators - .iter() - .map(|op| ChainedOperator { - operator_id: op.operator_id.clone(), - operator_name: op.operator_name.to_string(), - operator_config: op.operator_config.clone(), - }) - .collect(), - edges: node - .operator_chain - .edges - .iter() - .map(|edge| ProtoFsSchema::from((**edge).clone())) - .collect(), - }) - .collect(); - - let edges = value - .graph - .edge_indices() - .filter_map(|eidx| { - let edge = value.graph.edge_weight(eidx)?; - let (source, target) = value.graph.edge_endpoints(eidx)?; - Some(FsEdge { - source: source.index() as i32, - target: target.index() as i32, - schema: Some(ProtoFsSchema::from((*edge.schema).clone())), - edge_type: ProtoEdgeType::from(edge.edge_type) as i32, - }) - }) - .collect(); - - FsProgram { - nodes, - edges, - program_config: Some(value.program_config.into()), - } - } -} - -impl TryFrom for PipelineGraph { - type Error = DataFusionError; - - fn try_from(value: LogicalProgram) -> DFResult { - let nodes = value - .graph - .node_weights() - .map(|node| { - Ok(PipelineNode { - node_id: node.node_id, - operator: node.resolve_pipeline_operator_name()?, - description: node.description.clone(), - parallelism: node.parallelism as u32, - }) - }) - .collect::>>()?; - - let edges = value - .graph - .edge_references() - .filter_map(|edge| { - let src = value.graph.node_weight(edge.source())?; - let target = value.graph.node_weight(edge.target())?; - Some(PipelineEdge { - src_id: src.node_id, - dest_id: target.node_id, - key_type: "()".to_string(), - value_type: "()".to_string(), - edge_type: format!("{:?}", edge.weight().edge_type), - }) - }) - .collect(); - - Ok(PipelineGraph { nodes, edges }) - } -} diff --git a/src/sql/logical_node/logical/logical_edge.rs b/src/sql/logical_node/logical/logical_edge.rs deleted file mode 100644 index 87950e70..00000000 --- a/src/sql/logical_node/logical/logical_edge.rs +++ /dev/null @@ -1,102 +0,0 @@ -// 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::fmt::{Display, Formatter}; -use std::sync::Arc; - -use datafusion::common::{DataFusionError, Result}; -use protocol::function_stream_graph::EdgeType as ProtoEdgeType; -use serde::{Deserialize, Serialize}; - -use crate::sql::common::FsSchema; - -#[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Ord, Serialize, Deserialize)] -pub enum LogicalEdgeType { - Forward, - Shuffle, - LeftJoin, - RightJoin, -} - -impl Display for LogicalEdgeType { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - let symbol = match self { - LogicalEdgeType::Forward => "→", - LogicalEdgeType::Shuffle => "⤨", - LogicalEdgeType::LeftJoin => "-[left]⤨", - LogicalEdgeType::RightJoin => "-[right]⤨", - }; - write!(f, "{symbol}") - } -} - -impl From for LogicalEdgeType { - fn from(value: ProtoEdgeType) -> Self { - match value { - ProtoEdgeType::Unused => { - panic!("Critical: Invalid EdgeType 'Unused' encountered") - } - ProtoEdgeType::Forward => Self::Forward, - ProtoEdgeType::Shuffle => Self::Shuffle, - ProtoEdgeType::LeftJoin => Self::LeftJoin, - ProtoEdgeType::RightJoin => Self::RightJoin, - } - } -} - -impl From for ProtoEdgeType { - fn from(value: LogicalEdgeType) -> Self { - match value { - LogicalEdgeType::Forward => Self::Forward, - LogicalEdgeType::Shuffle => Self::Shuffle, - LogicalEdgeType::LeftJoin => Self::LeftJoin, - LogicalEdgeType::RightJoin => Self::RightJoin, - } - } -} - -pub(crate) fn logical_edge_type_from_proto_i32(i: i32) -> Result { - let e = ProtoEdgeType::try_from(i).map_err(|_| { - DataFusionError::Plan(format!("invalid protobuf EdgeType discriminant {i}")) - })?; - match e { - ProtoEdgeType::Unused => Err(DataFusionError::Plan( - "Critical: Invalid EdgeType 'Unused' encountered".into(), - )), - ProtoEdgeType::Forward => Ok(LogicalEdgeType::Forward), - ProtoEdgeType::Shuffle => Ok(LogicalEdgeType::Shuffle), - ProtoEdgeType::LeftJoin => Ok(LogicalEdgeType::LeftJoin), - ProtoEdgeType::RightJoin => Ok(LogicalEdgeType::RightJoin), - } -} - -#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] -pub struct LogicalEdge { - pub edge_type: LogicalEdgeType, - pub schema: Arc, -} - -impl LogicalEdge { - pub fn new(edge_type: LogicalEdgeType, schema: FsSchema) -> Self { - LogicalEdge { - edge_type, - schema: Arc::new(schema), - } - } - - pub fn project_all(edge_type: LogicalEdgeType, schema: FsSchema) -> Self { - LogicalEdge { - edge_type, - schema: Arc::new(schema), - } - } -} diff --git a/src/sql/logical_node/logical/logical_graph.rs b/src/sql/logical_node/logical/logical_graph.rs deleted file mode 100644 index b877e2a0..00000000 --- a/src/sql/logical_node/logical/logical_graph.rs +++ /dev/null @@ -1,30 +0,0 @@ -// 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 petgraph::graph::DiGraph; - -use super::logical_edge::LogicalEdge; -use super::logical_node::LogicalNode; - -pub type LogicalGraph = DiGraph; - -pub trait Optimizer { - fn optimize_once(&self, plan: &mut LogicalGraph) -> bool; - - fn optimize(&self, plan: &mut LogicalGraph) { - loop { - if !self.optimize_once(plan) { - break; - } - } - } -} diff --git a/src/sql/logical_node/logical/logical_node.rs b/src/sql/logical_node/logical/logical_node.rs deleted file mode 100644 index 5f00dc4b..00000000 --- a/src/sql/logical_node/logical/logical_node.rs +++ /dev/null @@ -1,87 +0,0 @@ -// 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::fmt::{Debug, Display, Formatter}; - -use datafusion::common::{DataFusionError, Result}; -use itertools::Itertools; -use serde::{Deserialize, Serialize}; - -use super::operator_chain::{ChainedLogicalOperator, OperatorChain}; -use super::operator_name::OperatorName; - -#[derive(Clone, Serialize, Deserialize)] -pub struct LogicalNode { - pub node_id: u32, - pub description: String, - pub operator_chain: OperatorChain, - pub parallelism: usize, -} - -impl LogicalNode { - pub fn single( - id: u32, - operator_id: String, - name: OperatorName, - config: Vec, - description: String, - parallelism: usize, - ) -> Self { - Self { - node_id: id, - description, - operator_chain: OperatorChain { - operators: vec![ChainedLogicalOperator { - operator_id, - operator_name: name, - operator_config: config, - }], - edges: vec![], - }, - parallelism, - } - } - - pub fn resolve_pipeline_operator_name(&self) -> Result { - let first_op = self.operator_chain.operators.first().ok_or_else(|| { - DataFusionError::Plan("Invalid LogicalNode: Operator chain is empty".into()) - })?; - - if let Some(connector_name) = first_op.extract_connector_name() { - return Ok(connector_name); - } - - if self.operator_chain.len() == 1 { - return Ok(first_op.operator_id.clone()); - } - - Ok("chained_op".to_string()) - } -} - -impl Display for LogicalNode { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "{}", self.description) - } -} - -impl Debug for LogicalNode { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - let chain_path = self - .operator_chain - .operators - .iter() - .map(|op| op.operator_id.as_str()) - .join(" -> "); - write!(f, "{chain_path}[{}]", self.parallelism) - } -} diff --git a/src/sql/logical_node/logical/logical_program.rs b/src/sql/logical_node/logical/logical_program.rs deleted file mode 100644 index 119ac469..00000000 --- a/src/sql/logical_node/logical/logical_program.rs +++ /dev/null @@ -1,153 +0,0 @@ -// 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::hash_map::DefaultHasher; -use std::collections::{HashMap, HashSet}; -use std::hash::Hasher; -use std::sync::Arc; - -use datafusion::arrow::datatypes::Schema; -use datafusion::common::{DataFusionError, Result as DFResult}; -use petgraph::Direction; -use petgraph::dot::Dot; -use prost::Message; -use protocol::function_stream_graph::FsProgram; -use rand::distributions::Alphanumeric; -use rand::rngs::SmallRng; -use rand::{Rng, SeedableRng}; - -use super::logical_graph::{LogicalGraph, Optimizer}; -use super::operator_name::OperatorName; -use super::program_config::ProgramConfig; - -#[derive(Clone, Debug, Default)] -pub struct LogicalProgram { - pub graph: LogicalGraph, - pub program_config: ProgramConfig, -} - -impl LogicalProgram { - pub fn new(graph: LogicalGraph, program_config: ProgramConfig) -> Self { - Self { - graph, - program_config, - } - } - - pub fn optimize(&mut self, optimizer: &dyn Optimizer) { - optimizer.optimize(&mut self.graph); - } - - pub fn update_parallelism(&mut self, overrides: &HashMap) { - for node in self.graph.node_weights_mut() { - if let Some(&p) = overrides.get(&node.node_id) { - node.parallelism = p; - } - } - } - - pub fn dot(&self) -> String { - format!("{:?}", Dot::with_config(&self.graph, &[])) - } - - pub fn task_count(&self) -> usize { - self.graph.node_weights().map(|nw| nw.parallelism).sum() - } - - pub fn sources(&self) -> HashSet { - self.graph - .externals(Direction::Incoming) - .filter_map(|idx| self.graph.node_weight(idx)) - .map(|node| node.node_id) - .collect() - } - - pub fn get_hash(&self) -> String { - let mut hasher = DefaultHasher::new(); - let program_bytes = FsProgram::from(self.clone()).encode_to_vec(); - hasher.write(&program_bytes); - let rng = SmallRng::seed_from_u64(hasher.finish()); - rng.sample_iter(&Alphanumeric) - .take(16) - .map(|c| (c as char).to_ascii_lowercase()) - .collect() - } - - pub fn tasks_per_operator(&self) -> HashMap { - self.graph - .node_weights() - .flat_map(|node| { - node.operator_chain - .operators - .iter() - .map(move |op| (op.operator_id.clone(), node.parallelism)) - }) - .collect() - } - - pub fn operator_names_by_id(&self) -> HashMap { - self.graph - .node_weights() - .flat_map(|node| &node.operator_chain.operators) - .map(|op| { - let resolved_name = op - .extract_connector_name() - .unwrap_or_else(|| op.operator_name.to_string()); - (op.operator_id.clone(), resolved_name) - }) - .collect() - } - - pub fn tasks_per_node(&self) -> HashMap { - self.graph - .node_weights() - .map(|node| (node.node_id, node.parallelism)) - .collect() - } - - pub fn features(&self) -> HashSet { - self.graph - .node_weights() - .flat_map(|node| &node.operator_chain.operators) - .filter_map(|op| op.extract_feature()) - .collect() - } - - /// Arrow schema carried on edges into the connector-sink node, if present. - pub fn egress_arrow_schema(&self) -> Option> { - for idx in self.graph.node_indices() { - let node = self.graph.node_weight(idx)?; - if node - .operator_chain - .operators - .iter() - .any(|op| op.operator_name == OperatorName::ConnectorSink) - { - let e = self.graph.edges_directed(idx, Direction::Incoming).next()?; - return Some(Arc::clone(&e.weight().schema.schema)); - } - } - None - } - - pub fn encode_for_catalog(&self) -> DFResult> { - Ok(FsProgram::from(self.clone()).encode_to_vec()) - } - - pub fn decode_for_catalog(bytes: &[u8]) -> DFResult { - let proto = FsProgram::decode(bytes).map_err(|e| { - DataFusionError::Execution(format!("FsProgram catalog decode failed: {e}")) - })?; - LogicalProgram::try_from(proto) - } -} diff --git a/src/sql/logical_node/logical/mod.rs b/src/sql/logical_node/logical/mod.rs deleted file mode 100644 index d2e9a327..00000000 --- a/src/sql/logical_node/logical/mod.rs +++ /dev/null @@ -1,30 +0,0 @@ -// 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. - -mod dylib_udf_config; -mod fs_program_convert; -mod logical_edge; -mod logical_graph; -mod logical_node; -mod logical_program; -mod operator_chain; -mod operator_name; -mod program_config; -mod python_udf_config; - -pub use dylib_udf_config::DylibUdfConfig; -pub use logical_edge::{LogicalEdge, LogicalEdgeType}; -pub use logical_graph::{LogicalGraph, Optimizer}; -pub use logical_node::LogicalNode; -pub use logical_program::LogicalProgram; -pub use operator_name::OperatorName; -pub use program_config::ProgramConfig; diff --git a/src/sql/logical_node/logical/operator_chain.rs b/src/sql/logical_node/logical/operator_chain.rs deleted file mode 100644 index 2aecddd6..00000000 --- a/src/sql/logical_node/logical/operator_chain.rs +++ /dev/null @@ -1,142 +0,0 @@ -// 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::sync::Arc; - -use itertools::{EitherOrBoth, Itertools}; -use prost::Message; -use protocol::function_stream_graph::ConnectorOp; -use serde::{Deserialize, Serialize}; - -use super::operator_name::OperatorName; -use crate::sql::common::FsSchema; - -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct ChainedLogicalOperator { - pub operator_id: String, - pub operator_name: OperatorName, - pub operator_config: Vec, -} - -impl ChainedLogicalOperator { - pub fn extract_connector_name(&self) -> Option { - if matches!( - self.operator_name, - OperatorName::ConnectorSource | OperatorName::ConnectorSink - ) { - ConnectorOp::decode(self.operator_config.as_slice()) - .ok() - .map(|op| op.connector) - } else { - None - } - } - - pub fn extract_feature(&self) -> Option { - match self.operator_name { - OperatorName::AsyncUdf => Some("async-udf".to_string()), - OperatorName::Join => Some("join-with-expiration".to_string()), - OperatorName::InstantJoin => Some("windowed-join".to_string()), - OperatorName::WindowFunction => Some("sql-window-function".to_string()), - OperatorName::LookupJoin => Some("lookup-join".to_string()), - OperatorName::TumblingWindowAggregate => { - Some("sql-tumbling-window-aggregate".to_string()) - } - OperatorName::SlidingWindowAggregate => { - Some("sql-sliding-window-aggregate".to_string()) - } - OperatorName::SessionWindowAggregate => { - Some("sql-session-window-aggregate".to_string()) - } - OperatorName::UpdatingAggregate => Some("sql-updating-aggregate".to_string()), - OperatorName::ConnectorSource => { - self.extract_connector_name().map(|c| format!("{c}-source")) - } - OperatorName::ConnectorSink => { - self.extract_connector_name().map(|c| format!("{c}-sink")) - } - _ => None, - } - } -} - -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct OperatorChain { - pub(crate) operators: Vec, - pub(crate) edges: Vec>, -} - -impl OperatorChain { - pub fn new(operator: ChainedLogicalOperator) -> Self { - Self { - operators: vec![operator], - edges: vec![], - } - } - - pub fn iter(&self) -> impl Iterator>)> { - self.operators - .iter() - .zip_longest(&self.edges) - .filter_map(|e| match e { - EitherOrBoth::Both(op, edge) => Some((op, Some(edge))), - EitherOrBoth::Left(op) => Some((op, None)), - EitherOrBoth::Right(_) => None, - }) - } - - pub fn iter_mut( - &mut self, - ) -> impl Iterator>)> { - self.operators - .iter_mut() - .zip_longest(&self.edges) - .filter_map(|e| match e { - EitherOrBoth::Both(op, edge) => Some((op, Some(edge))), - EitherOrBoth::Left(op) => Some((op, None)), - EitherOrBoth::Right(_) => None, - }) - } - - pub fn first(&self) -> &ChainedLogicalOperator { - self.operators - .first() - .expect("OperatorChain must contain at least one operator") - } - - pub fn len(&self) -> usize { - self.operators.len() - } - - pub fn is_empty(&self) -> bool { - self.operators.is_empty() - } - - pub fn is_source(&self) -> bool { - self.operators[0].operator_name == OperatorName::ConnectorSource - } - - pub fn is_sink(&self) -> bool { - self.operators[0].operator_name == OperatorName::ConnectorSink - } - - /// Operators safe to run at a higher upstream `TaskContext::parallelism` when fused after a - /// stateful node (e.g. window aggregate @ 8 → projection @ 1). - pub fn is_parallelism_upstream_expandable(&self) -> bool { - self.operators.iter().all(|op| { - matches!( - op.operator_name, - OperatorName::Projection | OperatorName::Value | OperatorName::ExpressionWatermark - ) - }) - } -} diff --git a/src/sql/logical_node/logical/operator_name.rs b/src/sql/logical_node/logical/operator_name.rs deleted file mode 100644 index 57f53f90..00000000 --- a/src/sql/logical_node/logical/operator_name.rs +++ /dev/null @@ -1,82 +0,0 @@ -// 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::str::FromStr; - -use serde::{Deserialize, Deserializer, Serialize, Serializer}; -use strum::{Display, EnumString, IntoStaticStr}; - -use crate::sql::common::constants::operator_feature; - -#[derive(Clone, Copy, Debug, Eq, PartialEq, EnumString, Display, IntoStaticStr)] -pub enum OperatorName { - ExpressionWatermark, - Value, - KeyBy, - Projection, - AsyncUdf, - Join, - InstantJoin, - LookupJoin, - WindowFunction, - TumblingWindowAggregate, - SlidingWindowAggregate, - SessionWindowAggregate, - UpdatingAggregate, - ConnectorSource, - ConnectorSink, -} - -impl OperatorName { - /// Registry / worker lookup key; matches [`Display`] and protobuf operator names. - #[inline] - pub fn as_registry_key(self) -> &'static str { - self.into() - } - - pub fn feature_tag(self) -> Option<&'static str> { - match self { - Self::ExpressionWatermark | Self::Value | Self::Projection => None, - Self::AsyncUdf => Some(operator_feature::ASYNC_UDF), - Self::Join => Some(operator_feature::JOIN_WITH_EXPIRATION), - Self::InstantJoin => Some(operator_feature::WINDOWED_JOIN), - Self::WindowFunction => Some(operator_feature::SQL_WINDOW_FUNCTION), - Self::LookupJoin => Some(operator_feature::LOOKUP_JOIN), - Self::TumblingWindowAggregate => Some(operator_feature::SQL_TUMBLING_WINDOW_AGGREGATE), - Self::SlidingWindowAggregate => Some(operator_feature::SQL_SLIDING_WINDOW_AGGREGATE), - Self::SessionWindowAggregate => Some(operator_feature::SQL_SESSION_WINDOW_AGGREGATE), - Self::UpdatingAggregate => Some(operator_feature::SQL_UPDATING_AGGREGATE), - Self::KeyBy => Some(operator_feature::KEY_BY_ROUTING), - Self::ConnectorSource => Some(operator_feature::CONNECTOR_SOURCE), - Self::ConnectorSink => Some(operator_feature::CONNECTOR_SINK), - } - } -} - -impl Serialize for OperatorName { - fn serialize(&self, serializer: S) -> Result - where - S: Serializer, - { - serializer.serialize_str(&self.to_string()) - } -} - -impl<'de> Deserialize<'de> for OperatorName { - fn deserialize(deserializer: D) -> Result - where - D: Deserializer<'de>, - { - let s = String::deserialize(deserializer)?; - Self::from_str(&s).map_err(serde::de::Error::custom) - } -} diff --git a/src/sql/logical_node/logical/program_config.rs b/src/sql/logical_node/logical/program_config.rs deleted file mode 100644 index 177326f4..00000000 --- a/src/sql/logical_node/logical/program_config.rs +++ /dev/null @@ -1,33 +0,0 @@ -// 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 protocol::function_stream_graph::FsProgramConfig; - -/// Placeholder program-level config (UDF tables live elsewhere; wire maps stay empty). -#[derive(Clone, Debug, Default)] -pub struct ProgramConfig {} - -impl From for FsProgramConfig { - fn from(_: ProgramConfig) -> Self { - Self { - udf_dylibs: Default::default(), - python_udfs: Default::default(), - } - } -} - -impl From for ProgramConfig { - fn from(_: FsProgramConfig) -> Self { - Self::default() - } -} diff --git a/src/sql/logical_node/logical/python_udf_config.rs b/src/sql/logical_node/logical/python_udf_config.rs deleted file mode 100644 index 6e7d5c66..00000000 --- a/src/sql/logical_node/logical/python_udf_config.rs +++ /dev/null @@ -1,23 +0,0 @@ -// 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::sync::Arc; - -use datafusion::arrow::datatypes::DataType; - -#[derive(Clone, Debug, Eq, PartialEq, Hash)] -pub struct PythonUdfConfig { - pub arg_types: Vec, - pub return_type: DataType, - pub name: Arc, - pub definition: Arc, -} diff --git a/src/sql/logical_node/lookup.rs b/src/sql/logical_node/lookup.rs deleted file mode 100644 index d2817c85..00000000 --- a/src/sql/logical_node/lookup.rs +++ /dev/null @@ -1,256 +0,0 @@ -// 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::fmt::Formatter; -use std::sync::Arc; - -use datafusion::common::{Column, DFSchemaRef, JoinType, Result, internal_err, plan_err}; -use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; -use datafusion::sql::TableReference; -use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec; -use datafusion_proto::physical_plan::to_proto::serialize_physical_expr; -use prost::Message; - -use protocol::function_stream_graph; -use protocol::function_stream_graph::{ConnectorOp, LookupJoinCondition, LookupJoinOperator}; - -use crate::multifield_partial_ord; -use crate::sql::common::constants::extension_node; -use crate::sql::common::{FsSchema, FsSchemaRef}; -use crate::sql::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; -use crate::sql::logical_node::{CompiledTopologyNode, StreamingOperatorBlueprint}; -use crate::sql::logical_planner::planner::{NamedNode, Planner}; -use crate::sql::schema::LookupTable; -use crate::sql::schema::utils::add_timestamp_field_arrow; - -pub const DICTIONARY_SOURCE_NODE_NAME: &str = extension_node::REFERENCE_TABLE_SOURCE; -pub const STREAM_DICTIONARY_JOIN_NODE_NAME: &str = extension_node::STREAM_REFERENCE_JOIN; - -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct ReferenceTableSourceNode { - pub(crate) source_definition: LookupTable, - pub(crate) resolved_schema: DFSchemaRef, -} - -multifield_partial_ord!(ReferenceTableSourceNode, source_definition); - -impl UserDefinedLogicalNodeCore for ReferenceTableSourceNode { - fn name(&self) -> &str { - DICTIONARY_SOURCE_NODE_NAME - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - vec![] - } - - fn schema(&self) -> &DFSchemaRef { - &self.resolved_schema - } - - fn expressions(&self) -> Vec { - vec![] - } - - fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { - write!(f, "ReferenceTableSource: Schema={}", self.resolved_schema) - } - - fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { - if !inputs.is_empty() { - return internal_err!( - "ReferenceTableSource is a leaf node and cannot accept upstream inputs" - ); - } - - Ok(Self { - source_definition: self.source_definition.clone(), - resolved_schema: self.resolved_schema.clone(), - }) - } -} - -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct StreamReferenceJoinNode { - pub(crate) upstream_stream_plan: LogicalPlan, - pub(crate) output_schema: DFSchemaRef, - pub(crate) external_dictionary: LookupTable, - pub(crate) equijoin_conditions: Vec<(Expr, Column)>, - pub(crate) post_join_filter: Option, - pub(crate) namespace_alias: Option, - pub(crate) join_semantics: JoinType, -} - -multifield_partial_ord!( - StreamReferenceJoinNode, - upstream_stream_plan, - external_dictionary, - equijoin_conditions, - post_join_filter, - namespace_alias -); - -impl StreamReferenceJoinNode { - fn compile_join_conditions(&self, planner: &Planner) -> Result> { - self.equijoin_conditions - .iter() - .map(|(logical_left_expr, right_column)| { - let physical_expr = - planner.create_physical_expr(logical_left_expr, &self.output_schema)?; - let serialized_expr = - serialize_physical_expr(&physical_expr, &DefaultPhysicalExtensionCodec {})?; - - Ok(LookupJoinCondition { - left_expr: serialized_expr.encode_to_vec(), - right_key: right_column.name.clone(), - }) - }) - .collect() - } - - fn map_api_join_type(&self) -> Result { - match self.join_semantics { - JoinType::Inner => Ok(function_stream_graph::JoinType::Inner as i32), - JoinType::Left => Ok(function_stream_graph::JoinType::Left as i32), - unsupported => plan_err!( - "Unsupported join type '{unsupported}' for dictionary lookups. Only INNER and LEFT joins are permitted." - ), - } - } - - fn build_engine_operator( - &self, - planner: &Planner, - _upstream_schema: &FsSchemaRef, - ) -> Result { - let internal_input_schema = - FsSchema::from_schema_unkeyed(Arc::new(self.output_schema.as_ref().into()))?; - let dictionary_physical_schema = self.external_dictionary.produce_physical_schema(); - let lookup_fs_schema = - FsSchema::from_schema_unkeyed(add_timestamp_field_arrow(dictionary_physical_schema))?; - - Ok(LookupJoinOperator { - input_schema: Some(internal_input_schema.into()), - lookup_schema: Some(lookup_fs_schema.clone().into()), - connector: Some(ConnectorOp { - connector: self.external_dictionary.adapter_type.clone(), - fs_schema: Some(lookup_fs_schema.into()), - name: self.external_dictionary.table_identifier.clone(), - description: self.external_dictionary.description.clone(), - config: Some(self.external_dictionary.connector_config.to_proto_config()), - }), - key_exprs: self.compile_join_conditions(planner)?, - join_type: self.map_api_join_type()?, - ttl_micros: self - .external_dictionary - .lookup_cache_ttl - .map(|t| t.as_micros() as u64), - max_capacity_bytes: self.external_dictionary.lookup_cache_max_bytes, - }) - } -} - -impl StreamingOperatorBlueprint for StreamReferenceJoinNode { - fn operator_identity(&self) -> Option { - None - } - - fn compile_to_graph_node( - &self, - planner: &Planner, - node_index: usize, - mut input_schemas: Vec, - ) -> Result { - if input_schemas.len() != 1 { - return plan_err!("StreamReferenceJoinNode requires exactly one upstream stream input"); - } - let upstream_schema = input_schemas.remove(0); - - let operator_config = self.build_engine_operator(planner, &upstream_schema)?; - - let logical_node = LogicalNode::single( - node_index as u32, - format!("lookup_join_{node_index}"), - OperatorName::LookupJoin, - operator_config.encode_to_vec(), - format!( - "DictionaryJoin<{}>", - self.external_dictionary.table_identifier - ), - planner.default_parallelism(), - ); - - let incoming_edge = - LogicalEdge::project_all(LogicalEdgeType::Shuffle, (*upstream_schema).clone()); - - Ok(CompiledTopologyNode { - execution_unit: logical_node, - routing_edges: vec![incoming_edge], - }) - } - - fn yielded_schema(&self) -> FsSchema { - FsSchema::from_schema_unkeyed(self.output_schema.inner().clone()) - .expect("Failed to convert lookup join output schema to FsSchema") - } -} - -impl UserDefinedLogicalNodeCore for StreamReferenceJoinNode { - fn name(&self) -> &str { - STREAM_DICTIONARY_JOIN_NODE_NAME - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - vec![&self.upstream_stream_plan] - } - - fn schema(&self) -> &DFSchemaRef { - &self.output_schema - } - - fn expressions(&self) -> Vec { - let mut exprs: Vec<_> = self - .equijoin_conditions - .iter() - .map(|(l, _)| l.clone()) - .collect(); - if let Some(filter) = &self.post_join_filter { - exprs.push(filter.clone()); - } - exprs - } - - fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { - write!( - f, - "StreamReferenceJoin: join_type={:?} | {}", - self.join_semantics, self.output_schema - ) - } - - fn with_exprs_and_inputs(&self, _: Vec, inputs: Vec) -> Result { - if inputs.len() != 1 { - return internal_err!( - "StreamReferenceJoinNode expects exactly 1 upstream plan, got {}", - inputs.len() - ); - } - Ok(Self { - upstream_stream_plan: inputs[0].clone(), - output_schema: self.output_schema.clone(), - external_dictionary: self.external_dictionary.clone(), - equijoin_conditions: self.equijoin_conditions.clone(), - post_join_filter: self.post_join_filter.clone(), - namespace_alias: self.namespace_alias.clone(), - join_semantics: self.join_semantics, - }) - } -} diff --git a/src/sql/logical_node/macros.rs b/src/sql/logical_node/macros.rs deleted file mode 100644 index 4ce649c2..00000000 --- a/src/sql/logical_node/macros.rs +++ /dev/null @@ -1,28 +0,0 @@ -// 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. - -#[macro_export] -macro_rules! multifield_partial_ord { - ($ty:ty, $($field:tt), *) => { - impl PartialOrd for $ty { - fn partial_cmp(&self, other: &Self) -> Option { - $( - let cmp = self.$field.partial_cmp(&other.$field)?; - if cmp != std::cmp::Ordering::Equal { - return Some(cmp); - } - )* - Some(std::cmp::Ordering::Equal) - } - } - }; -} diff --git a/src/sql/logical_node/mod.rs b/src/sql/logical_node/mod.rs deleted file mode 100644 index 2da1b8a4..00000000 --- a/src/sql/logical_node/mod.rs +++ /dev/null @@ -1,42 +0,0 @@ -// 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. - -pub mod logical; - -mod macros; - -pub(crate) mod streaming_operator_blueprint; -pub(crate) use streaming_operator_blueprint::{CompiledTopologyNode, StreamingOperatorBlueprint}; - -pub(crate) mod aggregate; -pub(crate) mod debezium; -pub(crate) mod join; -pub(crate) mod key_calculation; -pub(crate) mod lookup; -pub(crate) mod projection; -pub(crate) mod remote_table; -pub(crate) mod sink; -pub(crate) mod table_source; -pub(crate) mod updating_aggregate; -pub(crate) mod watermark_node; -pub(crate) mod windows_function; - -pub(crate) mod timestamp_append; -pub(crate) use timestamp_append::SystemTimestampInjectorNode; - -pub(crate) mod async_udf; -pub(crate) use async_udf::AsyncFunctionExecutionNode; - -pub(crate) mod is_retract; -pub(crate) use is_retract::IsRetractExtension; - -mod extension_try_from; diff --git a/src/sql/logical_node/projection.rs b/src/sql/logical_node/projection.rs deleted file mode 100644 index 3c5cfccb..00000000 --- a/src/sql/logical_node/projection.rs +++ /dev/null @@ -1,239 +0,0 @@ -// 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::fmt::Formatter; -use std::sync::Arc; - -use datafusion::common::{DFSchema, DFSchemaRef, Result, internal_err}; -use datafusion::logical_expr::{Expr, ExprSchemable, LogicalPlan, UserDefinedLogicalNodeCore}; -use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec; -use datafusion_proto::physical_plan::to_proto::serialize_physical_expr; -use prost::Message; - -use protocol::function_stream_graph::ProjectionOperator; - -use crate::multifield_partial_ord; -use crate::sql::common::constants::{extension_node, sql_field}; -use crate::sql::common::{FsSchema, FsSchemaRef}; -use crate::sql::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; -use crate::sql::logical_node::{CompiledTopologyNode, StreamingOperatorBlueprint}; -use crate::sql::logical_planner::planner::{NamedNode, Planner}; -use crate::sql::types::{QualifiedField, build_df_schema}; - -// ----------------------------------------------------------------------------- -// Constants & Identifiers -// ----------------------------------------------------------------------------- - -pub(crate) const STREAM_PROJECTION_NODE_NAME: &str = extension_node::STREAM_PROJECTION; -const DEFAULT_PROJECTION_LABEL: &str = sql_field::DEFAULT_PROJECTION_LABEL; - -// ----------------------------------------------------------------------------- -// Logical Node Definition -// ----------------------------------------------------------------------------- - -/// Projection within a streaming execution topology. -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub(crate) struct StreamProjectionNode { - pub(crate) upstream_plans: Vec, - pub(crate) operator_label: Option, - pub(crate) projection_exprs: Vec, - pub(crate) resolved_schema: DFSchemaRef, - pub(crate) requires_shuffle: bool, -} - -multifield_partial_ord!(StreamProjectionNode, operator_label, projection_exprs); - -impl StreamProjectionNode { - pub(crate) fn try_new( - upstream_plans: Vec, - operator_label: Option, - projection_exprs: Vec, - ) -> Result { - if upstream_plans.is_empty() { - return internal_err!("StreamProjectionNode requires at least one upstream plan"); - } - let primary_input = &upstream_plans[0]; - let upstream_schema = primary_input.schema(); - - let mut projected_fields = Vec::with_capacity(projection_exprs.len()); - for logical_expr in &projection_exprs { - let arrow_field = logical_expr.to_field(upstream_schema)?; - projected_fields.push(QualifiedField::from(arrow_field)); - } - - let resolved_schema = Arc::new(build_df_schema(&projected_fields)?); - - Ok(Self { - upstream_plans, - operator_label, - projection_exprs, - resolved_schema, - requires_shuffle: false, - }) - } - - pub(crate) fn with_shuffle_routing(mut self) -> Self { - self.requires_shuffle = true; - self - } - - fn validate_uniform_schemas(input_schemas: &[FsSchemaRef]) -> Result { - if input_schemas.is_empty() { - return internal_err!("No input schemas provided to projection planner"); - } - let primary_schema = input_schemas[0].clone(); - - for schema in input_schemas.iter().skip(1) { - if **schema != *primary_schema { - return internal_err!( - "Schema mismatch: All upstream inputs to a projection node must share the identical schema topology." - ); - } - } - - Ok(primary_schema) - } - - fn compile_physical_expressions( - &self, - planner: &Planner, - input_df_schema: &DFSchemaRef, - ) -> Result>> { - self.projection_exprs - .iter() - .map(|logical_expr| { - let physical_expr = planner - .create_physical_expr(logical_expr, input_df_schema) - .map_err(|e| e.context("Failed to compile physical projection expression"))?; - - let serialized_expr = - serialize_physical_expr(&physical_expr, &DefaultPhysicalExtensionCodec {})?; - - Ok(serialized_expr.encode_to_vec()) - }) - .collect() - } -} - -// ----------------------------------------------------------------------------- -// Stream Extension Trait Implementation -// ----------------------------------------------------------------------------- - -impl StreamingOperatorBlueprint for StreamProjectionNode { - fn operator_identity(&self) -> Option { - None - } - - fn compile_to_graph_node( - &self, - planner: &Planner, - node_index: usize, - input_schemas: Vec, - ) -> Result { - let unified_input_schema = Self::validate_uniform_schemas(&input_schemas)?; - let input_df_schema = Arc::new(DFSchema::try_from( - unified_input_schema.schema.as_ref().clone(), - )?); - - let compiled_expr_payloads = - self.compile_physical_expressions(planner, &input_df_schema)?; - - let operator_config = ProjectionOperator { - name: self - .operator_label - .as_deref() - .unwrap_or(DEFAULT_PROJECTION_LABEL) - .to_string(), - input_schema: Some(unified_input_schema.as_ref().clone().into()), - output_schema: Some(self.yielded_schema().into()), - exprs: compiled_expr_payloads, - }; - - let node_identifier = format!("projection_{node_index}"); - let label = format!( - "ArrowProjection<{}>", - self.operator_label.as_deref().unwrap_or("_") - ); - - let logical_node = LogicalNode::single( - node_index as u32, - node_identifier, - OperatorName::Projection, - operator_config.encode_to_vec(), - label, - planner.default_parallelism(), - ); - - let routing_strategy = if self.requires_shuffle { - LogicalEdgeType::Shuffle - } else { - LogicalEdgeType::Forward - }; - - let outgoing_edge = - LogicalEdge::project_all(routing_strategy, (*unified_input_schema).clone()); - - Ok(CompiledTopologyNode { - execution_unit: logical_node, - routing_edges: vec![outgoing_edge], - }) - } - - fn yielded_schema(&self) -> FsSchema { - FsSchema::from_schema_unkeyed(Arc::new(self.resolved_schema.as_arrow().clone())) - .expect("Fatal: Failed to generate unkeyed output schema for projection") - } -} - -// ----------------------------------------------------------------------------- -// DataFusion Logical Node Hooks -// ----------------------------------------------------------------------------- - -impl UserDefinedLogicalNodeCore for StreamProjectionNode { - fn name(&self) -> &str { - STREAM_PROJECTION_NODE_NAME - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - self.upstream_plans.iter().collect() - } - - fn schema(&self) -> &DFSchemaRef { - &self.resolved_schema - } - - fn expressions(&self) -> Vec { - vec![] - } - - fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { - write!( - f, - "StreamProjectionNode: RequiresShuffle={}, Schema={}", - self.requires_shuffle, self.resolved_schema - ) - } - - fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { - let mut new_node = Self::try_new( - inputs, - self.operator_label.clone(), - self.projection_exprs.clone(), - )?; - - if self.requires_shuffle { - new_node = new_node.with_shuffle_routing(); - } - - Ok(new_node) - } -} diff --git a/src/sql/logical_node/remote_table.rs b/src/sql/logical_node/remote_table.rs deleted file mode 100644 index bde1d47f..00000000 --- a/src/sql/logical_node/remote_table.rs +++ /dev/null @@ -1,190 +0,0 @@ -// 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::fmt::Formatter; -use std::sync::Arc; - -use datafusion::common::{DFSchemaRef, Result, TableReference, internal_err, plan_err}; -use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; -use datafusion_proto::physical_plan::AsExecutionPlan; -use datafusion_proto::protobuf::PhysicalPlanNode; -use prost::Message; - -use protocol::function_stream_graph::ValuePlanOperator; - -use crate::multifield_partial_ord; -use crate::sql::common::constants::extension_node; -use crate::sql::common::{FsSchema, FsSchemaRef}; -use crate::sql::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; -use crate::sql::logical_node::{CompiledTopologyNode, StreamingOperatorBlueprint}; -use crate::sql::logical_planner::planner::{NamedNode, Planner}; -use crate::sql::physical::StreamingExtensionCodec; - -// ----------------------------------------------------------------------------- -// Constants & Identifiers -// ----------------------------------------------------------------------------- - -pub(crate) const REMOTE_TABLE_NODE_NAME: &str = extension_node::REMOTE_TABLE_BOUNDARY; - -// ----------------------------------------------------------------------------- -// Logical Node Definition -// ----------------------------------------------------------------------------- - -/// Segments the execution graph and merges nodes sharing the same identifier; acts as a boundary. -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub(crate) struct RemoteTableBoundaryNode { - pub(crate) upstream_plan: LogicalPlan, - pub(crate) table_identifier: TableReference, - pub(crate) resolved_schema: DFSchemaRef, - pub(crate) requires_materialization: bool, -} - -multifield_partial_ord!( - RemoteTableBoundaryNode, - upstream_plan, - table_identifier, - requires_materialization -); - -impl RemoteTableBoundaryNode { - fn compile_engine_operator(&self, planner: &Planner) -> Result> { - let physical_plan = planner.sync_plan(&self.upstream_plan)?; - - let physical_plan_proto = PhysicalPlanNode::try_from_physical_plan( - physical_plan, - &StreamingExtensionCodec::default(), - )?; - - let operator_config = ValuePlanOperator { - name: format!("value_calculation({})", self.table_identifier), - physical_plan: physical_plan_proto.encode_to_vec(), - }; - - Ok(operator_config.encode_to_vec()) - } - - fn validate_uniform_schemas(input_schemas: &[FsSchemaRef]) -> Result<()> { - if input_schemas.len() <= 1 { - return Ok(()); - } - - let primary_schema = &input_schemas[0]; - for schema in input_schemas.iter().skip(1) { - if *schema != *primary_schema { - return plan_err!( - "Topology error: Multiple input streams routed to the same remote table must share an identical schema structure." - ); - } - } - - Ok(()) - } -} - -// ----------------------------------------------------------------------------- -// Stream Extension Trait Implementation -// ----------------------------------------------------------------------------- - -impl StreamingOperatorBlueprint for RemoteTableBoundaryNode { - fn operator_identity(&self) -> Option { - if self.requires_materialization { - Some(NamedNode::RemoteTable(self.table_identifier.clone())) - } else { - None - } - } - - fn compile_to_graph_node( - &self, - planner: &Planner, - node_index: usize, - input_schemas: Vec, - ) -> Result { - Self::validate_uniform_schemas(&input_schemas)?; - - let operator_payload = self.compile_engine_operator(planner)?; - - let logical_node = LogicalNode::single( - node_index as u32, - format!("value_{node_index}"), - OperatorName::Value, - operator_payload, - self.table_identifier.to_string(), - planner.default_parallelism(), - ); - - let routing_edges: Vec = input_schemas - .into_iter() - .map(|schema| LogicalEdge::project_all(LogicalEdgeType::Forward, (*schema).clone())) - .collect(); - - Ok(CompiledTopologyNode { - execution_unit: logical_node, - routing_edges, - }) - } - - fn yielded_schema(&self) -> FsSchema { - FsSchema::from_schema_keys(Arc::new(self.resolved_schema.as_ref().into()), vec![]) - .expect("Fatal: Failed to generate output schema for remote table boundary") - } -} - -// ----------------------------------------------------------------------------- -// DataFusion Logical Node Hooks -// ----------------------------------------------------------------------------- - -impl UserDefinedLogicalNodeCore for RemoteTableBoundaryNode { - fn name(&self) -> &str { - REMOTE_TABLE_NODE_NAME - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - vec![&self.upstream_plan] - } - - fn schema(&self) -> &DFSchemaRef { - &self.resolved_schema - } - - fn expressions(&self) -> Vec { - vec![] - } - - fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { - write!( - f, - "RemoteTableBoundaryNode: Identifier={}, Materialized={}, Schema={}", - self.table_identifier, self.requires_materialization, self.resolved_schema - ) - } - - fn with_exprs_and_inputs( - &self, - _exprs: Vec, - mut inputs: Vec, - ) -> Result { - if inputs.len() != 1 { - return internal_err!( - "RemoteTableBoundaryNode expects exactly 1 upstream logical plan, but received {}", - inputs.len() - ); - } - - Ok(Self { - upstream_plan: inputs.remove(0), - table_identifier: self.table_identifier.clone(), - resolved_schema: self.resolved_schema.clone(), - requires_materialization: self.requires_materialization, - }) - } -} diff --git a/src/sql/logical_node/sink.rs b/src/sql/logical_node/sink.rs deleted file mode 100644 index d767afe3..00000000 --- a/src/sql/logical_node/sink.rs +++ /dev/null @@ -1,247 +0,0 @@ -// 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::fmt::Formatter; -use std::sync::Arc; - -use datafusion::common::{DFSchemaRef, Result, TableReference, plan_err}; -use datafusion::logical_expr::{Expr, Extension, LogicalPlan, UserDefinedLogicalNodeCore}; -use prost::Message; - -use crate::multifield_partial_ord; -use crate::sql::common::constants::extension_node; -use crate::sql::common::{FsSchema, FsSchemaRef, UPDATING_META_FIELD}; -use crate::sql::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; -use crate::sql::logical_node::{CompiledTopologyNode, StreamingOperatorBlueprint}; -use crate::sql::logical_planner::planner::{NamedNode, Planner}; -use crate::sql::schema::CatalogEntity; -use crate::sql::schema::catalog::ExternalTable; - -use super::debezium::PackDebeziumEnvelopeNode; -use super::remote_table::RemoteTableBoundaryNode; - -// ----------------------------------------------------------------------------- -// Constants & Identifiers -// ----------------------------------------------------------------------------- - -pub(crate) const STREAM_EGRESS_NODE_NAME: &str = extension_node::STREAM_EGRESS; - -// ----------------------------------------------------------------------------- -// Logical Node Definition -// ----------------------------------------------------------------------------- - -/// Terminal node routing processed data into an external sink (e.g. Kafka, PostgreSQL). -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub(crate) struct StreamEgressNode { - pub(crate) target_identifier: TableReference, - pub(crate) destination_table: CatalogEntity, - pub(crate) egress_schema: DFSchemaRef, - upstream_plans: Arc>, -} - -multifield_partial_ord!(StreamEgressNode, target_identifier, upstream_plans); - -impl StreamEgressNode { - pub fn try_new( - target_identifier: TableReference, - destination_table: CatalogEntity, - initial_schema: DFSchemaRef, - upstream_plan: LogicalPlan, - ) -> Result { - let (mut processed_plan, mut resolved_schema) = - Self::apply_cdc_transformations(upstream_plan, initial_schema, &destination_table)?; - - Self::enforce_computational_boundary(&mut resolved_schema, &mut processed_plan); - - Ok(Self { - target_identifier, - destination_table, - egress_schema: resolved_schema, - upstream_plans: Arc::new(vec![processed_plan]), - }) - } - - fn apply_cdc_transformations( - plan: LogicalPlan, - schema: DFSchemaRef, - destination: &CatalogEntity, - ) -> Result<(LogicalPlan, DFSchemaRef)> { - let is_upstream_updating = plan - .schema() - .has_column_with_unqualified_name(UPDATING_META_FIELD); - - match destination { - CatalogEntity::ExternalConnector(b) => match b.as_ref() { - ExternalTable::Sink(sink) => { - let is_sink_updating = sink.is_updating(); - - match (is_upstream_updating, is_sink_updating) { - (_, true) => { - let debezium_encoder = PackDebeziumEnvelopeNode::try_new(plan)?; - let wrapped_plan = LogicalPlan::Extension(Extension { - node: Arc::new(debezium_encoder), - }); - let new_schema = wrapped_plan.schema().clone(); - - Ok((wrapped_plan, new_schema)) - } - (true, false) => { - plan_err!( - "Topology Mismatch: The upstream is producing an updating stream (CDC), \ - but the target sink '{}' is not configured to accept updates. \ - Hint: set `format = 'debezium_json'` in the WITH clause.", - sink.name() - ) - } - (false, false) => Ok((plan, schema)), - } - } - ExternalTable::Source(source) => { - let is_sink_updating = source.is_updating(); - match (is_upstream_updating, is_sink_updating) { - (_, true) => { - let debezium_encoder = PackDebeziumEnvelopeNode::try_new(plan)?; - let wrapped_plan = LogicalPlan::Extension(Extension { - node: Arc::new(debezium_encoder), - }); - let new_schema = wrapped_plan.schema().clone(); - Ok((wrapped_plan, new_schema)) - } - (true, false) => plan_err!( - "Topology Mismatch: upstream produces CDC but target '{}' is a non-updating source table", - source.name() - ), - (false, false) => Ok((plan, schema)), - } - } - ExternalTable::Lookup(_) => plan_err!( - "Topology Violation: A Lookup Table cannot be used as a streaming data sink." - ), - }, - CatalogEntity::ComputedTable { .. } => Ok((plan, schema)), - } - } - - fn enforce_computational_boundary(schema: &mut DFSchemaRef, plan: &mut LogicalPlan) { - let requires_boundary = if let LogicalPlan::Extension(extension) = plan { - let stream_ext: &dyn StreamingOperatorBlueprint = (&extension.node) - .try_into() - .expect("Fatal: Egress node encountered an extension that does not implement StreamingOperatorBlueprint"); - - stream_ext.is_passthrough_boundary() - } else { - true - }; - - if requires_boundary { - let boundary_node = RemoteTableBoundaryNode { - upstream_plan: plan.clone(), - table_identifier: TableReference::bare("sink projection"), - resolved_schema: schema.clone(), - requires_materialization: false, - }; - - *plan = LogicalPlan::Extension(Extension { - node: Arc::new(boundary_node), - }); - } - } -} - -// ----------------------------------------------------------------------------- -// Stream Extension Trait Implementation -// ----------------------------------------------------------------------------- - -impl StreamingOperatorBlueprint for StreamEgressNode { - fn operator_identity(&self) -> Option { - Some(NamedNode::Sink(self.target_identifier.clone())) - } - - fn compile_to_graph_node( - &self, - planner: &Planner, - node_index: usize, - input_schemas: Vec, - ) -> Result { - let connector_operator = self - .destination_table - .connector_op() - .map_err(|e| e.context("Failed to generate connector operation payload"))?; - - let operator_description = connector_operator.description.clone(); - let operator_payload = connector_operator.encode_to_vec(); - - let logical_node = LogicalNode::single( - node_index as u32, - format!("sink_{}_{node_index}", self.target_identifier), - OperatorName::ConnectorSink, - operator_payload, - operator_description, - planner.default_parallelism(), - ); - - let routing_edges: Vec = input_schemas - .into_iter() - .map(|input_schema| { - LogicalEdge::project_all(LogicalEdgeType::Forward, (*input_schema).clone()) - }) - .collect(); - - Ok(CompiledTopologyNode { - execution_unit: logical_node, - routing_edges, - }) - } - - fn yielded_schema(&self) -> FsSchema { - FsSchema::from_fields(vec![]) - } -} - -// ----------------------------------------------------------------------------- -// DataFusion Logical Node Hooks -// ----------------------------------------------------------------------------- - -impl UserDefinedLogicalNodeCore for StreamEgressNode { - fn name(&self) -> &str { - STREAM_EGRESS_NODE_NAME - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - self.upstream_plans.iter().collect() - } - - fn schema(&self) -> &DFSchemaRef { - &self.egress_schema - } - - fn expressions(&self) -> Vec { - vec![] - } - - fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { - write!( - f, - "StreamEgressNode({:?}): Schema={}", - self.target_identifier, self.egress_schema - ) - } - - fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { - Ok(Self { - target_identifier: self.target_identifier.clone(), - destination_table: self.destination_table.clone(), - egress_schema: self.egress_schema.clone(), - upstream_plans: Arc::new(inputs), - }) - } -} diff --git a/src/sql/logical_node/streaming_operator_blueprint.rs b/src/sql/logical_node/streaming_operator_blueprint.rs deleted file mode 100644 index d3f9d459..00000000 --- a/src/sql/logical_node/streaming_operator_blueprint.rs +++ /dev/null @@ -1,65 +0,0 @@ -// 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::fmt::Debug; - -use datafusion::common::Result; - -use crate::sql::common::{FsSchema, FsSchemaRef}; -use crate::sql::logical_node::logical::{LogicalEdge, LogicalNode}; -use crate::sql::logical_planner::planner::{NamedNode, Planner}; - -// ----------------------------------------------------------------------------- -// Core Execution Blueprint -// ----------------------------------------------------------------------------- - -/// Atomic unit within a streaming execution topology: translates streaming SQL into graph nodes. -pub(crate) trait StreamingOperatorBlueprint: Debug { - /// Canonical named identity for this operator, if any (sources, sinks, etc.). - fn operator_identity(&self) -> Option; - - /// Compiles this operator into a graph vertex and its incoming routing edges. - fn compile_to_graph_node( - &self, - compiler_context: &Planner, - node_id_sequence: usize, - upstream_schemas: Vec, - ) -> Result; - - /// Schema of records this operator yields downstream. - fn yielded_schema(&self) -> FsSchema; - - /// Logical passthrough boundary (no physical state change); default is stateful / materializing. - fn is_passthrough_boundary(&self) -> bool { - false - } -} - -// ----------------------------------------------------------------------------- -// Graph Topology Structures -// ----------------------------------------------------------------------------- - -/// Compiled vertex: execution unit plus upstream routing edges. -#[derive(Debug, Clone)] -pub(crate) struct CompiledTopologyNode { - pub execution_unit: LogicalNode, - pub routing_edges: Vec, -} - -impl CompiledTopologyNode { - pub fn new(execution_unit: LogicalNode, routing_edges: Vec) -> Self { - Self { - execution_unit, - routing_edges, - } - } -} diff --git a/src/sql/logical_node/table_source.rs b/src/sql/logical_node/table_source.rs deleted file mode 100644 index b1c6bfdd..00000000 --- a/src/sql/logical_node/table_source.rs +++ /dev/null @@ -1,180 +0,0 @@ -// 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::fmt::Formatter; -use std::sync::Arc; - -use datafusion::common::{DFSchemaRef, Result, TableReference, plan_err}; -use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; -use prost::Message; - -use crate::multifield_partial_ord; -use crate::sql::common::constants::extension_node; -use crate::sql::common::{FsSchema, FsSchemaRef}; -use crate::sql::logical_node::debezium::DebeziumSchemaCodec; -use crate::sql::logical_node::logical::{LogicalNode, OperatorName}; -use crate::sql::logical_planner::planner::{NamedNode, Planner}; -use crate::sql::schema::SourceTable; -use crate::sql::schema::utils::add_timestamp_field; -use crate::sql::types::build_df_schema; - -use super::{CompiledTopologyNode, StreamingOperatorBlueprint}; - -// ----------------------------------------------------------------------------- -// Constants & Identifiers -// ----------------------------------------------------------------------------- - -pub(crate) const STREAM_INGESTION_NODE_NAME: &str = extension_node::STREAM_INGESTION; - -// ----------------------------------------------------------------------------- -// Logical Node Definition -// ----------------------------------------------------------------------------- - -/// Foundational ingestion point: connects to external systems and injects raw or CDC data. -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub(crate) struct StreamIngestionNode { - pub(crate) source_identifier: TableReference, - pub(crate) source_definition: SourceTable, - pub(crate) resolved_schema: DFSchemaRef, -} - -multifield_partial_ord!(StreamIngestionNode, source_identifier, source_definition); - -impl StreamIngestionNode { - pub fn try_new( - source_identifier: TableReference, - source_definition: SourceTable, - ) -> Result { - let resolved_schema = Self::build_ingestion_schema(&source_identifier, &source_definition)?; - - Ok(Self { - source_identifier, - source_definition, - resolved_schema, - }) - } - - fn build_ingestion_schema( - identifier: &TableReference, - definition: &SourceTable, - ) -> Result { - let physical_fields: Vec<_> = definition - .schema_specs - .iter() - .filter(|col| !col.is_computed()) - .map(|col| { - ( - Some(identifier.clone()), - Arc::new(col.arrow_field().clone()), - ) - .into() - }) - .collect(); - - let base_schema = Arc::new(build_df_schema(&physical_fields)?); - - let enveloped_schema = if definition.is_updating() { - DebeziumSchemaCodec::wrap_into_envelope(&base_schema, Some(identifier.clone()))? - } else { - base_schema - }; - - add_timestamp_field(enveloped_schema, Some(identifier.clone())) - } -} - -// ----------------------------------------------------------------------------- -// DataFusion Logical Node Hooks -// ----------------------------------------------------------------------------- - -impl UserDefinedLogicalNodeCore for StreamIngestionNode { - fn name(&self) -> &str { - STREAM_INGESTION_NODE_NAME - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - vec![] - } - - fn schema(&self) -> &DFSchemaRef { - &self.resolved_schema - } - - fn expressions(&self) -> Vec { - vec![] - } - - fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { - write!( - f, - "StreamIngestionNode({}): Schema={}", - self.source_identifier, self.resolved_schema - ) - } - - fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { - if !inputs.is_empty() { - return plan_err!( - "StreamIngestionNode acts as a leaf boundary and cannot accept upstream inputs." - ); - } - - Ok(Self { - source_identifier: self.source_identifier.clone(), - source_definition: self.source_definition.clone(), - resolved_schema: self.resolved_schema.clone(), - }) - } -} - -// ----------------------------------------------------------------------------- -// Core Execution Blueprint Implementation -// ----------------------------------------------------------------------------- - -impl StreamingOperatorBlueprint for StreamIngestionNode { - fn operator_identity(&self) -> Option { - Some(NamedNode::Source(self.source_identifier.clone())) - } - - fn compile_to_graph_node( - &self, - compiler_context: &Planner, - node_id_sequence: usize, - upstream_schemas: Vec, - ) -> Result { - if !upstream_schemas.is_empty() { - return plan_err!( - "Topology Violation: StreamIngestionNode is a source origin and cannot process upstream routing edges." - ); - } - - let sql_source = self.source_definition.as_sql_source()?; - let connector_payload = sql_source.source.config.encode_to_vec(); - let operator_description = sql_source.source.config.description.clone(); - - let execution_unit = LogicalNode::single( - node_id_sequence as u32, - format!("source_{}_{node_id_sequence}", self.source_identifier), - OperatorName::ConnectorSource, - connector_payload, - operator_description, - compiler_context.default_parallelism(), - ); - - Ok(CompiledTopologyNode::new(execution_unit, vec![])) - } - - fn yielded_schema(&self) -> FsSchema { - FsSchema::from_schema_keys(Arc::new(self.resolved_schema.as_ref().into()), vec![]) - .expect("Fatal: Failed to generate output schema for stream ingestion") - } -} diff --git a/src/sql/logical_node/timestamp_append.rs b/src/sql/logical_node/timestamp_append.rs deleted file mode 100644 index 630e5a66..00000000 --- a/src/sql/logical_node/timestamp_append.rs +++ /dev/null @@ -1,121 +0,0 @@ -// 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::fmt::Formatter; - -use datafusion::common::{DFSchemaRef, Result, TableReference, internal_err}; -use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; - -use crate::multifield_partial_ord; -use crate::sql::common::constants::extension_node; -use crate::sql::schema::utils::{add_timestamp_field, has_timestamp_field}; - -// ----------------------------------------------------------------------------- -// Constants & Identifiers -// ----------------------------------------------------------------------------- - -pub(crate) const TIMESTAMP_INJECTOR_NODE_NAME: &str = extension_node::SYSTEM_TIMESTAMP_INJECTOR; - -// ----------------------------------------------------------------------------- -// Logical Node Definition -// ----------------------------------------------------------------------------- - -/// Injects the mandatory system `_timestamp` field into the upstream streaming schema. -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub(crate) struct SystemTimestampInjectorNode { - pub(crate) upstream_plan: LogicalPlan, - pub(crate) target_qualifier: Option, - pub(crate) resolved_schema: DFSchemaRef, -} - -multifield_partial_ord!(SystemTimestampInjectorNode, upstream_plan, target_qualifier); - -impl SystemTimestampInjectorNode { - pub(crate) fn try_new( - upstream_plan: LogicalPlan, - target_qualifier: Option, - ) -> Result { - let upstream_schema = upstream_plan.schema(); - - if has_timestamp_field(upstream_schema) { - return internal_err!( - "Topology Violation: Attempted to inject a system timestamp into an upstream plan \ - that already contains one. \ - \nPlan:\n {:?} \nSchema:\n {:?}", - upstream_plan, - upstream_schema - ); - } - - let resolved_schema = - add_timestamp_field(upstream_schema.clone(), target_qualifier.clone())?; - - Ok(Self { - upstream_plan, - target_qualifier, - resolved_schema, - }) - } -} - -// ----------------------------------------------------------------------------- -// DataFusion Logical Node Hooks -// ----------------------------------------------------------------------------- - -impl UserDefinedLogicalNodeCore for SystemTimestampInjectorNode { - fn name(&self) -> &str { - TIMESTAMP_INJECTOR_NODE_NAME - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - vec![&self.upstream_plan] - } - - fn schema(&self) -> &DFSchemaRef { - &self.resolved_schema - } - - fn expressions(&self) -> Vec { - vec![] - } - - fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { - let field_names = self - .resolved_schema - .fields() - .iter() - .map(|field| field.name().to_string()) - .collect::>() - .join(", "); - - write!( - f, - "SystemTimestampInjector(Qualifier={:?}): [{}]", - self.target_qualifier, field_names - ) - } - - fn with_exprs_and_inputs( - &self, - _exprs: Vec, - mut inputs: Vec, - ) -> Result { - if inputs.len() != 1 { - return internal_err!( - "SystemTimestampInjectorNode requires exactly 1 upstream logical plan, but received {}", - inputs.len() - ); - } - - Self::try_new(inputs.remove(0), self.target_qualifier.clone()) - } -} diff --git a/src/sql/logical_node/updating_aggregate.rs b/src/sql/logical_node/updating_aggregate.rs deleted file mode 100644 index 0ddb2b28..00000000 --- a/src/sql/logical_node/updating_aggregate.rs +++ /dev/null @@ -1,245 +0,0 @@ -// 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::sync::Arc; -use std::time::Duration; - -use datafusion::common::{DFSchemaRef, Result, TableReference, ToDFSchema, internal_err, plan_err}; -use datafusion::logical_expr::expr::ScalarFunction; -use datafusion::logical_expr::{ - Expr, Extension, LogicalPlan, UserDefinedLogicalNodeCore, col, lit, -}; -use datafusion::prelude::named_struct; -use datafusion::scalar::ScalarValue; -use datafusion_proto::physical_plan::AsExecutionPlan; -use datafusion_proto::protobuf::PhysicalPlanNode; -use prost::Message; -use protocol::function_stream_graph::UpdatingAggregateOperator; - -use crate::sql::common::constants::{extension_node, proto_operator_name, updating_state_field}; -use crate::sql::common::{FsSchema, FsSchemaRef}; -use crate::sql::functions::multi_hash; -use crate::sql::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; -use crate::sql::logical_node::{ - CompiledTopologyNode, IsRetractExtension, StreamingOperatorBlueprint, -}; -use crate::sql::logical_planner::planner::{NamedNode, Planner}; -use crate::sql::physical::StreamingExtensionCodec; - -// ----------------------------------------------------------------------------- -// Constants & Configuration -// ----------------------------------------------------------------------------- - -pub(crate) const CONTINUOUS_AGGREGATE_NODE_NAME: &str = extension_node::CONTINUOUS_AGGREGATE; - -const DEFAULT_FLUSH_INTERVAL_MICROS: u64 = 10_000_000; - -const STATIC_HASH_SIZE_BYTES: i32 = 16; - -// ----------------------------------------------------------------------------- -// Logical Node Definition -// ----------------------------------------------------------------------------- - -/// Stateful continuous aggregation: running aggregates with updating / retraction semantics. -#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd)] -pub(crate) struct ContinuousAggregateNode { - pub(crate) base_aggregate_plan: LogicalPlan, - pub(crate) partition_key_indices: Vec, - pub(crate) retract_injected_plan: LogicalPlan, - pub(crate) namespace_qualifier: Option, - pub(crate) state_retention_ttl: Duration, -} - -impl ContinuousAggregateNode { - pub fn try_new( - base_aggregate_plan: LogicalPlan, - partition_key_indices: Vec, - namespace_qualifier: Option, - state_retention_ttl: Duration, - ) -> Result { - let retract_injected_plan = LogicalPlan::Extension(Extension { - node: Arc::new(IsRetractExtension::new( - base_aggregate_plan.clone(), - namespace_qualifier.clone(), - )), - }); - - Ok(Self { - base_aggregate_plan, - partition_key_indices, - retract_injected_plan, - namespace_qualifier, - state_retention_ttl, - }) - } - - fn construct_state_metadata_expr(&self, upstream_schema: &FsSchemaRef) -> Expr { - let routing_keys: Vec = self - .partition_key_indices - .iter() - .map(|&idx| col(upstream_schema.schema.field(idx).name())) - .collect(); - - let state_id_hash = if routing_keys.is_empty() { - Expr::Literal( - ScalarValue::FixedSizeBinary( - STATIC_HASH_SIZE_BYTES, - Some(vec![0; STATIC_HASH_SIZE_BYTES as usize]), - ), - None, - ) - } else { - Expr::ScalarFunction(ScalarFunction { - func: multi_hash(), - args: routing_keys, - }) - }; - - named_struct(vec![ - lit(updating_state_field::IS_RETRACT), - lit(false), - lit(updating_state_field::ID), - state_id_hash, - ]) - } - - fn compile_operator_config( - &self, - planner: &Planner, - upstream_schema: &FsSchemaRef, - ) -> Result { - let upstream_df_schema = upstream_schema.schema.clone().to_dfschema()?; - - let physical_agg_plan = planner.sync_plan(&self.base_aggregate_plan)?; - let compiled_agg_payload = PhysicalPlanNode::try_from_physical_plan( - physical_agg_plan, - &StreamingExtensionCodec::default(), - )? - .encode_to_vec(); - - let meta_expr = self.construct_state_metadata_expr(upstream_schema); - let compiled_meta_expr = - planner.serialize_as_physical_expr(&meta_expr, &upstream_df_schema)?; - - Ok(UpdatingAggregateOperator { - name: proto_operator_name::UPDATING_AGGREGATE.to_string(), - input_schema: Some((**upstream_schema).clone().into()), - final_schema: Some(self.yielded_schema().into()), - aggregate_exec: compiled_agg_payload, - metadata_expr: compiled_meta_expr, - flush_interval_micros: DEFAULT_FLUSH_INTERVAL_MICROS, - ttl_micros: self.state_retention_ttl.as_micros() as u64, - }) - } -} - -// ----------------------------------------------------------------------------- -// DataFusion Logical Node Hooks -// ----------------------------------------------------------------------------- - -impl UserDefinedLogicalNodeCore for ContinuousAggregateNode { - fn name(&self) -> &str { - CONTINUOUS_AGGREGATE_NODE_NAME - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - vec![&self.base_aggregate_plan] - } - - fn schema(&self) -> &DFSchemaRef { - self.retract_injected_plan.schema() - } - - fn expressions(&self) -> Vec { - vec![] - } - - fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - write!( - f, - "ContinuousAggregateNode(TTL={:?})", - self.state_retention_ttl - ) - } - - fn with_exprs_and_inputs( - &self, - _exprs: Vec, - mut inputs: Vec, - ) -> Result { - if inputs.len() != 1 { - return internal_err!( - "ContinuousAggregateNode requires exactly 1 upstream input, got {}", - inputs.len() - ); - } - - Self::try_new( - inputs.remove(0), - self.partition_key_indices.clone(), - self.namespace_qualifier.clone(), - self.state_retention_ttl, - ) - } -} - -// ----------------------------------------------------------------------------- -// Core Execution Blueprint Implementation -// ----------------------------------------------------------------------------- - -impl StreamingOperatorBlueprint for ContinuousAggregateNode { - fn operator_identity(&self) -> Option { - None - } - - fn compile_to_graph_node( - &self, - planner: &Planner, - node_index: usize, - mut upstream_schemas: Vec, - ) -> Result { - if upstream_schemas.len() != 1 { - return plan_err!( - "Topology Violation: ContinuousAggregateNode requires exactly 1 upstream input, received {}", - upstream_schemas.len() - ); - } - - let upstream_schema = upstream_schemas.remove(0); - - let operator_config = self.compile_operator_config(planner, &upstream_schema)?; - - let parallelism = planner.keyed_aggregate_parallelism(); - - let logical_node = LogicalNode::single( - node_index as u32, - format!("updating_aggregate_{node_index}"), - OperatorName::UpdatingAggregate, - operator_config.encode_to_vec(), - proto_operator_name::UPDATING_AGGREGATE.to_string(), - parallelism, - ); - - let shuffle_edge = - LogicalEdge::project_all(LogicalEdgeType::Shuffle, (*upstream_schema).clone()); - - Ok(CompiledTopologyNode { - execution_unit: logical_node, - routing_edges: vec![shuffle_edge], - }) - } - - fn yielded_schema(&self) -> FsSchema { - FsSchema::from_schema_unkeyed(Arc::new(self.schema().as_ref().into())) - .expect("Fatal: Failed to generate unkeyed output schema for continuous aggregate") - } -} diff --git a/src/sql/logical_node/watermark_node.rs b/src/sql/logical_node/watermark_node.rs deleted file mode 100644 index 9a8fc9d6..00000000 --- a/src/sql/logical_node/watermark_node.rs +++ /dev/null @@ -1,229 +0,0 @@ -// 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::fmt::Formatter; -use std::sync::Arc; - -use datafusion::common::{DFSchemaRef, Result, TableReference, internal_err, plan_err}; -use datafusion::error::DataFusionError; -use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; -use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec; -use datafusion_proto::physical_plan::to_proto::serialize_physical_expr; -use prost::Message; -use protocol::function_stream_graph::ExpressionWatermarkConfig; - -use crate::multifield_partial_ord; -use crate::sql::common::constants::{extension_node, runtime_operator_kind}; -use crate::sql::common::{FsSchema, FsSchemaRef}; -use crate::sql::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; -use crate::sql::logical_node::{CompiledTopologyNode, StreamingOperatorBlueprint}; -use crate::sql::logical_planner::planner::{NamedNode, Planner}; -use crate::sql::schema::utils::add_timestamp_field; -use crate::sql::types::TIMESTAMP_FIELD; - -// ----------------------------------------------------------------------------- -// Constants & Identifiers -// ----------------------------------------------------------------------------- - -pub(crate) const EVENT_TIME_WATERMARK_NODE_NAME: &str = extension_node::EVENT_TIME_WATERMARK; - -const DEFAULT_WATERMARK_EMISSION_PERIOD_MICROS: u64 = 1_000_000; - -// ----------------------------------------------------------------------------- -// Logical Node Definition -// ----------------------------------------------------------------------------- - -/// Event-time watermark from a user strategy; drives time progress in stateful operators. -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub(crate) struct EventTimeWatermarkNode { - pub(crate) upstream_plan: LogicalPlan, - pub(crate) namespace_qualifier: TableReference, - pub(crate) watermark_strategy_expr: Expr, - pub(crate) resolved_schema: DFSchemaRef, - pub(crate) internal_timestamp_offset: usize, -} - -multifield_partial_ord!( - EventTimeWatermarkNode, - upstream_plan, - namespace_qualifier, - watermark_strategy_expr, - internal_timestamp_offset -); - -impl EventTimeWatermarkNode { - pub(crate) fn try_new( - upstream_plan: LogicalPlan, - namespace_qualifier: TableReference, - watermark_strategy_expr: Expr, - ) -> Result { - let resolved_schema = add_timestamp_field( - upstream_plan.schema().clone(), - Some(namespace_qualifier.clone()), - )?; - - let internal_timestamp_offset = resolved_schema - .index_of_column_by_name(None, TIMESTAMP_FIELD) - .ok_or_else(|| { - DataFusionError::Plan(format!( - "Fatal: Failed to resolve mandatory temporal column '{}'", - TIMESTAMP_FIELD - )) - })?; - - Ok(Self { - upstream_plan, - namespace_qualifier, - watermark_strategy_expr, - resolved_schema, - internal_timestamp_offset, - }) - } - - pub(crate) fn generate_fs_schema(&self) -> FsSchema { - FsSchema::new_unkeyed( - Arc::new(self.resolved_schema.as_ref().into()), - self.internal_timestamp_offset, - ) - } - - fn compile_operator_config(&self, planner: &Planner) -> Result { - let physical_expr = - planner.create_physical_expr(&self.watermark_strategy_expr, &self.resolved_schema)?; - - let serialized_expr = - serialize_physical_expr(&physical_expr, &DefaultPhysicalExtensionCodec {})?; - - Ok(ExpressionWatermarkConfig { - period_micros: DEFAULT_WATERMARK_EMISSION_PERIOD_MICROS, - idle_time_micros: None, - expression: serialized_expr.encode_to_vec(), - input_schema: Some(self.generate_fs_schema().into()), - }) - } -} - -// ----------------------------------------------------------------------------- -// DataFusion Logical Node Hooks -// ----------------------------------------------------------------------------- - -impl UserDefinedLogicalNodeCore for EventTimeWatermarkNode { - fn name(&self) -> &str { - EVENT_TIME_WATERMARK_NODE_NAME - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - vec![&self.upstream_plan] - } - - fn schema(&self) -> &DFSchemaRef { - &self.resolved_schema - } - - fn expressions(&self) -> Vec { - vec![self.watermark_strategy_expr.clone()] - } - - fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { - write!( - f, - "EventTimeWatermarkNode({}): Schema={}", - self.namespace_qualifier, self.resolved_schema - ) - } - - fn with_exprs_and_inputs( - &self, - mut exprs: Vec, - mut inputs: Vec, - ) -> Result { - if inputs.len() != 1 { - return internal_err!( - "EventTimeWatermarkNode requires exactly 1 upstream logical plan, but received {}", - inputs.len() - ); - } - if exprs.len() != 1 { - return internal_err!( - "EventTimeWatermarkNode requires exactly 1 watermark strategy expression, but received {}", - exprs.len() - ); - } - - let internal_timestamp_offset = self - .resolved_schema - .index_of_column_by_name(Some(&self.namespace_qualifier), TIMESTAMP_FIELD) - .ok_or_else(|| { - DataFusionError::Plan(format!( - "Optimizer Error: Lost tracking of temporal column '{}'", - TIMESTAMP_FIELD - )) - })?; - - Ok(Self { - upstream_plan: inputs.remove(0), - namespace_qualifier: self.namespace_qualifier.clone(), - watermark_strategy_expr: exprs.remove(0), - resolved_schema: self.resolved_schema.clone(), - internal_timestamp_offset, - }) - } -} - -// ----------------------------------------------------------------------------- -// Core Execution Blueprint Implementation -// ----------------------------------------------------------------------------- - -impl StreamingOperatorBlueprint for EventTimeWatermarkNode { - fn operator_identity(&self) -> Option { - Some(NamedNode::Watermark(self.namespace_qualifier.clone())) - } - - fn compile_to_graph_node( - &self, - planner: &Planner, - node_index: usize, - mut upstream_schemas: Vec, - ) -> Result { - if upstream_schemas.len() != 1 { - return plan_err!( - "Topology Violation: EventTimeWatermarkNode requires exactly 1 upstream input, received {}", - upstream_schemas.len() - ); - } - - let operator_config = self.compile_operator_config(planner)?; - - let execution_unit = LogicalNode::single( - node_index as u32, - format!("watermark_{node_index}"), - OperatorName::ExpressionWatermark, - operator_config.encode_to_vec(), - runtime_operator_kind::WATERMARK_GENERATOR.to_string(), - planner.default_parallelism(), - ); - - let incoming_edge = LogicalEdge::project_all( - LogicalEdgeType::Forward, - (*upstream_schemas.remove(0)).clone(), - ); - - Ok(CompiledTopologyNode { - execution_unit, - routing_edges: vec![incoming_edge], - }) - } - - fn yielded_schema(&self) -> FsSchema { - self.generate_fs_schema() - } -} diff --git a/src/sql/logical_node/windows_function.rs b/src/sql/logical_node/windows_function.rs deleted file mode 100644 index 9be37382..00000000 --- a/src/sql/logical_node/windows_function.rs +++ /dev/null @@ -1,191 +0,0 @@ -// 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::fmt::Formatter; -use std::sync::Arc; - -use datafusion::common::{Column, DFSchema, DFSchemaRef, Result, internal_err, plan_err}; -use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; -use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec; -use datafusion_proto::physical_plan::to_proto::serialize_physical_expr; -use datafusion_proto::{physical_plan::AsExecutionPlan, protobuf::PhysicalPlanNode}; -use prost::Message; -use protocol::function_stream_graph::WindowFunctionOperator; - -use crate::sql::common::constants::{extension_node, proto_operator_name, runtime_operator_kind}; -use crate::sql::common::{FsSchema, FsSchemaRef}; -use crate::sql::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; -use crate::sql::logical_planner::planner::{NamedNode, Planner}; -use crate::sql::physical::StreamingExtensionCodec; -use crate::sql::types::TIMESTAMP_FIELD; - -use super::{CompiledTopologyNode, StreamingOperatorBlueprint}; - -// ----------------------------------------------------------------------------- -// Constants & Identifiers -// ----------------------------------------------------------------------------- - -pub(crate) const STREAMING_WINDOW_NODE_NAME: &str = extension_node::STREAMING_WINDOW_FUNCTION; - -// ----------------------------------------------------------------------------- -// Logical Node Definition -// ----------------------------------------------------------------------------- - -/// Stateful streaming window: temporal binning plus underlying window evaluation plan. -#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd)] -pub(crate) struct StreamingWindowFunctionNode { - pub(crate) underlying_evaluation_plan: LogicalPlan, - pub(crate) partition_key_indices: Vec, -} - -impl StreamingWindowFunctionNode { - pub fn new(underlying_evaluation_plan: LogicalPlan, partition_key_indices: Vec) -> Self { - Self { - underlying_evaluation_plan, - partition_key_indices, - } - } - - fn compile_temporal_binning_function( - &self, - planner: &Planner, - input_df_schema: &DFSchema, - ) -> Result> { - let timestamp_column = Expr::Column(Column::new_unqualified(TIMESTAMP_FIELD.to_string())); - - let physical_binning_expr = - planner.create_physical_expr(×tamp_column, input_df_schema)?; - - let serialized_expr = - serialize_physical_expr(&physical_binning_expr, &DefaultPhysicalExtensionCodec {})?; - - Ok(serialized_expr.encode_to_vec()) - } - - fn compile_physical_evaluation_plan(&self, planner: &Planner) -> Result> { - let physical_window_plan = planner.sync_plan(&self.underlying_evaluation_plan)?; - - let proto_plan_node = PhysicalPlanNode::try_from_physical_plan( - physical_window_plan, - &StreamingExtensionCodec::default(), - )?; - - Ok(proto_plan_node.encode_to_vec()) - } -} - -// ----------------------------------------------------------------------------- -// DataFusion Logical Node Hooks -// ----------------------------------------------------------------------------- - -impl UserDefinedLogicalNodeCore for StreamingWindowFunctionNode { - fn name(&self) -> &str { - STREAMING_WINDOW_NODE_NAME - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - vec![&self.underlying_evaluation_plan] - } - - fn schema(&self) -> &DFSchemaRef { - self.underlying_evaluation_plan.schema() - } - - fn expressions(&self) -> Vec { - vec![] - } - - fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { - write!(f, "StreamingWindowFunction: Schema={}", self.schema()) - } - - fn with_exprs_and_inputs( - &self, - _exprs: Vec, - mut inputs: Vec, - ) -> Result { - if inputs.len() != 1 { - return internal_err!( - "StreamingWindowFunctionNode requires exactly 1 upstream input, got {}", - inputs.len() - ); - } - - Ok(Self::new( - inputs.remove(0), - self.partition_key_indices.clone(), - )) - } -} - -// ----------------------------------------------------------------------------- -// Core Execution Blueprint Implementation -// ----------------------------------------------------------------------------- - -impl StreamingOperatorBlueprint for StreamingWindowFunctionNode { - fn operator_identity(&self) -> Option { - None - } - - fn compile_to_graph_node( - &self, - planner: &Planner, - node_index: usize, - mut input_schemas: Vec, - ) -> Result { - if input_schemas.len() != 1 { - return plan_err!( - "Topology Violation: StreamingWindowFunctionNode requires exactly 1 upstream input schema, received {}", - input_schemas.len() - ); - } - - let input_schema = input_schemas.remove(0); - - let input_df_schema = DFSchema::try_from(input_schema.schema.as_ref().clone())?; - - let binning_payload = self.compile_temporal_binning_function(planner, &input_df_schema)?; - let evaluation_plan_payload = self.compile_physical_evaluation_plan(planner)?; - - let operator_config = WindowFunctionOperator { - name: proto_operator_name::WINDOW_FUNCTION.to_string(), - input_schema: Some(input_schema.as_ref().clone().into()), - binning_function: binning_payload, - window_function_plan: evaluation_plan_payload, - }; - - let parallelism = planner.keyed_aggregate_parallelism(); - - let logical_node = LogicalNode::single( - node_index as u32, - format!("window_function_{node_index}"), - OperatorName::WindowFunction, - operator_config.encode_to_vec(), - runtime_operator_kind::STREAMING_WINDOW_EVALUATOR.to_string(), - parallelism, - ); - - let routing_edge = - LogicalEdge::project_all(LogicalEdgeType::Shuffle, (*input_schema).clone()); - - Ok(CompiledTopologyNode { - execution_unit: logical_node, - routing_edges: vec![routing_edge], - }) - } - - fn yielded_schema(&self) -> FsSchema { - FsSchema::from_schema_unkeyed(Arc::new(self.schema().as_ref().clone().into())).expect( - "Fatal: Failed to generate unkeyed output schema for StreamingWindowFunctionNode", - ) - } -} diff --git a/src/sql/logical_planner/mod.rs b/src/sql/logical_planner/mod.rs deleted file mode 100644 index 9ecfb676..00000000 --- a/src/sql/logical_planner/mod.rs +++ /dev/null @@ -1,16 +0,0 @@ -// 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. - -pub mod optimizers; - -pub(crate) mod streaming_planner; -pub(crate) use streaming_planner as planner; diff --git a/src/sql/logical_planner/optimizers/chaining.rs b/src/sql/logical_planner/optimizers/chaining.rs deleted file mode 100644 index ea7bd885..00000000 --- a/src/sql/logical_planner/optimizers/chaining.rs +++ /dev/null @@ -1,200 +0,0 @@ -// 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 petgraph::prelude::*; -use petgraph::visit::NodeIndexable; -use tracing::debug; - -use crate::sql::logical_node::logical::{LogicalEdgeType, LogicalGraph, Optimizer}; - -pub struct ChainingOptimizer {} - -impl Optimizer for ChainingOptimizer { - fn optimize_once(&self, plan: &mut LogicalGraph) -> bool { - let mut match_found = None; - - for node_idx in plan.node_indices() { - let mut outgoing = plan.edges_directed(node_idx, Outgoing); - let first_out = outgoing.next(); - if first_out.is_none() || outgoing.next().is_some() { - continue; - } - let edge = first_out.unwrap(); - - if edge.weight().edge_type != LogicalEdgeType::Forward { - continue; - } - - let target_idx = edge.target(); - - let mut incoming = plan.edges_directed(target_idx, Incoming); - let first_in = incoming.next(); - if first_in.is_none() || incoming.next().is_some() { - continue; - } - - let source_node = plan.node_weight(node_idx).expect("Source node missing"); - let target_node = plan.node_weight(target_idx).expect("Target node missing"); - - let parallelism_ok = source_node.parallelism == target_node.parallelism - || target_node - .operator_chain - .is_parallelism_upstream_expandable(); - - if source_node.operator_chain.is_source() - || target_node.operator_chain.is_sink() - || !parallelism_ok - { - continue; - } - - match_found = Some((node_idx, target_idx, edge.id())); - break; - } - - if let Some((source_idx, target_idx, edge_id)) = match_found { - let edge_weight = plan.remove_edge(edge_id).expect("Edge should exist"); - - let target_outgoing: Vec<_> = plan - .edges_directed(target_idx, Outgoing) - .map(|e| (e.id(), e.target())) - .collect(); - - for (e_id, next_target_idx) in target_outgoing { - let weight = plan.remove_edge(e_id).expect("Outgoing edge missing"); - plan.add_edge(source_idx, next_target_idx, weight); - } - - let is_source_last = source_idx.index() == plan.node_bound() - 1; - - let target_node = plan - .remove_node(target_idx) - .expect("Target node should exist"); - - let actual_source_idx = if is_source_last { - target_idx - } else { - source_idx - }; - - let source_node = plan - .node_weight_mut(actual_source_idx) - .expect("Source node missing"); - - debug!( - "Chaining Optimizer: Fusing '{}' -> '{}'", - source_node.description, target_node.description - ); - - source_node.description = - format!("{} -> {}", source_node.description, target_node.description); - - source_node.parallelism = source_node.parallelism.max(target_node.parallelism); - - source_node - .operator_chain - .operators - .extend(target_node.operator_chain.operators); - source_node.operator_chain.edges.push(edge_weight.schema); - - return true; - } - - false - } -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit}; - - use crate::sql::common::FsSchema; - use crate::sql::logical_node::logical::{ - LogicalEdge, LogicalEdgeType, LogicalGraph, LogicalNode, OperatorName, Optimizer, - }; - - use super::ChainingOptimizer; - - fn forward_edge() -> LogicalEdge { - let s = Arc::new(Schema::new(vec![Field::new( - "_timestamp", - DataType::Timestamp(TimeUnit::Nanosecond, None), - false, - )])); - LogicalEdge::new(LogicalEdgeType::Forward, FsSchema::new_unkeyed(s, 0)) - } - - fn proj_node(id: u32, label: &str) -> LogicalNode { - LogicalNode::single( - id, - format!("op_{label}"), - OperatorName::Projection, - vec![], - label.to_string(), - 1, - ) - } - - fn source_node() -> LogicalNode { - LogicalNode::single( - 0, - "src".into(), - OperatorName::ConnectorSource, - vec![], - "source".into(), - 1, - ) - } - - /// Window aggregate at higher default parallelism may forward into projection @ 1: still fuse - /// so each branch does not reserve a separate global state-memory block for the same sub-chain. - #[test] - fn fusion_stateful_high_parallelism_into_expandable_low() { - let mut g = LogicalGraph::new(); - let n0 = g.add_node(source_node()); - let n1 = g.add_node(proj_node(1, "tumble")); - let n2 = g.add_node(proj_node(2, "proj")); - let n1w = g.node_weight_mut(n1).unwrap(); - n1w.parallelism = 8; - let e = forward_edge(); - g.add_edge(n0, n1, e.clone()); - g.add_edge(n1, n2, e); - - let changed = ChainingOptimizer {}.optimize_once(&mut g); - assert!(changed); - assert_eq!(g.node_count(), 2); - let fused = g - .node_weights() - .find(|n| n.description.contains("->")) - .unwrap(); - assert_eq!(fused.parallelism, 8); - assert_eq!(fused.operator_chain.len(), 2); - } - - /// Regression: upstream at last `NodeIndex` + remove non-last downstream swaps indices. - #[test] - fn fusion_remaps_when_upstream_was_last_node_index() { - let mut g = LogicalGraph::new(); - let n0 = g.add_node(source_node()); - let n1 = g.add_node(proj_node(1, "downstream")); - let n2 = g.add_node(proj_node(2, "upstream_last_index")); - let e = forward_edge(); - g.add_edge(n0, n2, e.clone()); - g.add_edge(n2, n1, e); - - let changed = ChainingOptimizer {}.optimize_once(&mut g); - assert!(changed); - assert_eq!(g.node_count(), 2); - } -} diff --git a/src/sql/logical_planner/optimizers/mod.rs b/src/sql/logical_planner/optimizers/mod.rs deleted file mode 100644 index c7981313..00000000 --- a/src/sql/logical_planner/optimizers/mod.rs +++ /dev/null @@ -1,20 +0,0 @@ -// 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. - -//! Logical planner optimizers: graph-level chaining ([`ChainingOptimizer`]) and -//! DataFusion SQL logical-plan rules ([`produce_optimized_plan`]). - -mod chaining; -mod optimized_plan; - -pub use chaining::ChainingOptimizer; -pub use optimized_plan::produce_optimized_plan; diff --git a/src/sql/logical_planner/optimizers/optimized_plan.rs b/src/sql/logical_planner/optimizers/optimized_plan.rs deleted file mode 100644 index fbb64845..00000000 --- a/src/sql/logical_planner/optimizers/optimized_plan.rs +++ /dev/null @@ -1,95 +0,0 @@ -// 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::sync::Arc; - -use datafusion::common::Result; -use datafusion::common::config::ConfigOptions; -use datafusion::logical_expr::LogicalPlan; -use datafusion::optimizer::OptimizerContext; -use datafusion::optimizer::OptimizerRule; -use datafusion::optimizer::common_subexpr_eliminate::CommonSubexprEliminate; -use datafusion::optimizer::decorrelate_lateral_join::DecorrelateLateralJoin; -use datafusion::optimizer::decorrelate_predicate_subquery::DecorrelatePredicateSubquery; -use datafusion::optimizer::eliminate_cross_join::EliminateCrossJoin; -use datafusion::optimizer::eliminate_duplicated_expr::EliminateDuplicatedExpr; -use datafusion::optimizer::eliminate_filter::EliminateFilter; -use datafusion::optimizer::eliminate_group_by_constant::EliminateGroupByConstant; -use datafusion::optimizer::eliminate_join::EliminateJoin; -use datafusion::optimizer::eliminate_limit::EliminateLimit; -use datafusion::optimizer::eliminate_nested_union::EliminateNestedUnion; -use datafusion::optimizer::eliminate_one_union::EliminateOneUnion; -use datafusion::optimizer::eliminate_outer_join::EliminateOuterJoin; -use datafusion::optimizer::extract_equijoin_predicate::ExtractEquijoinPredicate; -use datafusion::optimizer::filter_null_join_keys::FilterNullJoinKeys; -use datafusion::optimizer::optimizer::Optimizer; -use datafusion::optimizer::propagate_empty_relation::PropagateEmptyRelation; -use datafusion::optimizer::push_down_filter::PushDownFilter; -use datafusion::optimizer::push_down_limit::PushDownLimit; -use datafusion::optimizer::replace_distinct_aggregate::ReplaceDistinctWithAggregate; -use datafusion::optimizer::scalar_subquery_to_join::ScalarSubqueryToJoin; -use datafusion::optimizer::simplify_expressions::SimplifyExpressions; -use datafusion::sql::planner::SqlToRel; -use datafusion::sql::sqlparser::ast::Statement; - -use crate::sql::schema::StreamSchemaProvider; - -/// Converts a SQL statement into an optimized DataFusion logical plan. -/// -/// Applies the DataFusion analyzer followed by a curated set of optimizer rules -/// suitable for streaming SQL (some rules like OptimizeProjections are excluded -/// because they can drop event-time calculation fields). -pub fn produce_optimized_plan( - statement: &Statement, - schema_provider: &StreamSchemaProvider, -) -> Result { - let sql_to_rel = SqlToRel::new(schema_provider); - let plan = sql_to_rel.sql_statement_to_plan(statement.clone())?; - - let analyzed_plan = schema_provider.analyzer.execute_and_check( - plan, - &ConfigOptions::default(), - |_plan, _rule| {}, - )?; - - let rules: Vec> = vec![ - Arc::new(EliminateNestedUnion::new()), - Arc::new(SimplifyExpressions::new()), - Arc::new(ReplaceDistinctWithAggregate::new()), - Arc::new(EliminateJoin::new()), - Arc::new(DecorrelatePredicateSubquery::new()), - Arc::new(ScalarSubqueryToJoin::new()), - Arc::new(DecorrelateLateralJoin::new()), - Arc::new(ExtractEquijoinPredicate::new()), - Arc::new(EliminateDuplicatedExpr::new()), - Arc::new(EliminateFilter::new()), - Arc::new(EliminateCrossJoin::new()), - Arc::new(EliminateLimit::new()), - Arc::new(PropagateEmptyRelation::new()), - Arc::new(EliminateOneUnion::new()), - Arc::new(FilterNullJoinKeys::default()), - Arc::new(EliminateOuterJoin::new()), - Arc::new(PushDownLimit::new()), - Arc::new(PushDownFilter::new()), - Arc::new(EliminateGroupByConstant::new()), - Arc::new(CommonSubexprEliminate::new()), - ]; - - let optimizer = Optimizer::with_rules(rules); - let optimized = optimizer.optimize( - analyzed_plan, - &OptimizerContext::default(), - |_plan, _rule| {}, - )?; - - Ok(optimized) -} diff --git a/src/sql/logical_planner/streaming_planner.rs b/src/sql/logical_planner/streaming_planner.rs deleted file mode 100644 index 1e999c2a..00000000 --- a/src/sql/logical_planner/streaming_planner.rs +++ /dev/null @@ -1,435 +0,0 @@ -// 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::Arc; -use std::thread; -use std::time::Duration; - -use datafusion::arrow::datatypes::IntervalMonthDayNanoType; -use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor}; -use datafusion::common::{ - DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, Spans, plan_err, -}; -use datafusion::execution::context::SessionState; -use datafusion::execution::runtime_env::RuntimeEnvBuilder; -use datafusion::functions::datetime::date_bin; -use datafusion::logical_expr::{Expr, Extension, LogicalPlan, UserDefinedLogicalNode}; -use datafusion::physical_expr::PhysicalExpr; -use datafusion::physical_plan::ExecutionPlan; -use datafusion::physical_planner::{DefaultPhysicalPlanner, ExtensionPlanner, PhysicalPlanner}; -use datafusion_proto::protobuf::{PhysicalExprNode, PhysicalPlanNode}; -use datafusion_proto::{ - physical_plan::AsExecutionPlan, - protobuf::{AggregateMode, physical_plan_node::PhysicalPlanType}, -}; -use petgraph::graph::{DiGraph, NodeIndex}; -use prost::Message; -use tokio::runtime::Builder; -use tokio::sync::oneshot; - -use async_trait::async_trait; -use datafusion_common::TableReference; -use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec; -use datafusion_proto::physical_plan::to_proto::serialize_physical_expr; - -use crate::sql::common::constants::sql_planning_default; -use crate::sql::common::{FsSchema, FsSchemaRef}; -use crate::sql::logical_node::debezium::{ - PACK_NODE_NAME, UNROLL_NODE_NAME, UnrollDebeziumPayloadNode, -}; -use crate::sql::logical_node::key_calculation::KeyExtractionNode; -use crate::sql::logical_node::logical::{LogicalEdge, LogicalGraph, LogicalNode}; -use crate::sql::logical_node::{CompiledTopologyNode, StreamingOperatorBlueprint}; -use crate::sql::physical::{ - CdcDebeziumPackExec, CdcDebeziumUnrollExec, FsMemExec, StreamingDecodingContext, - StreamingExtensionCodec, -}; -use crate::sql::schema::StreamSchemaProvider; -use crate::sql::schema::utils::add_timestamp_field_arrow; - -pub(crate) struct SplitPlanOutput { - pub(crate) partial_aggregation_plan: PhysicalPlanNode, - pub(crate) partial_schema: FsSchema, - pub(crate) finish_plan: PhysicalPlanNode, -} -#[derive(Eq, Hash, PartialEq, Debug)] -pub(crate) enum NamedNode { - Source(TableReference), - Watermark(TableReference), - RemoteTable(TableReference), - Sink(TableReference), -} - -pub(crate) struct PlanToGraphVisitor<'a> { - graph: DiGraph, - output_schemas: HashMap, - named_nodes: HashMap, - traversal: Vec>, - planner: Planner<'a>, -} - -impl<'a> PlanToGraphVisitor<'a> { - pub fn new(schema_provider: &'a StreamSchemaProvider, session_state: &'a SessionState) -> Self { - Self { - graph: Default::default(), - output_schemas: Default::default(), - named_nodes: Default::default(), - traversal: vec![], - planner: Planner::new(schema_provider, session_state), - } - } -} - -pub(crate) struct Planner<'a> { - schema_provider: &'a StreamSchemaProvider, - planner: DefaultPhysicalPlanner, - session_state: &'a SessionState, -} - -impl<'a> Planner<'a> { - #[inline] - pub(crate) fn default_parallelism(&self) -> usize { - self.schema_provider.default_parallelism() - } - - #[inline] - pub(crate) fn key_by_parallelism(&self) -> usize { - self.schema_provider.key_by_parallelism() - } - - /// Parallelism for operators that consume a keyed shuffle (non-empty partition keys). - #[inline] - pub(crate) fn keyed_aggregate_parallelism(&self) -> usize { - sql_planning_default::KEYED_AGGREGATE_DEFAULT_PARALLELISM - } - - pub(crate) fn new( - schema_provider: &'a StreamSchemaProvider, - session_state: &'a SessionState, - ) -> Self { - let planner = - DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new(FsExtensionPlanner {})]); - Self { - schema_provider, - planner, - session_state, - } - } - - pub(crate) fn sync_plan(&self, plan: &LogicalPlan) -> Result> { - let fut = self.planner.create_physical_plan(plan, self.session_state); - let (tx, mut rx) = oneshot::channel(); - thread::scope(|s| { - let builder = thread::Builder::new(); - let builder = if cfg!(debug_assertions) { - builder.stack_size(10_000_000) - } else { - builder - }; - builder - .spawn_scoped(s, move || { - let rt = Builder::new_current_thread().enable_all().build().unwrap(); - rt.block_on(async { - let plan = fut.await; - tx.send(plan).unwrap(); - }); - }) - .unwrap(); - }); - - rx.try_recv().unwrap() - } - - pub(crate) fn create_physical_expr( - &self, - expr: &Expr, - input_dfschema: &DFSchema, - ) -> Result> { - self.planner - .create_physical_expr(expr, input_dfschema, self.session_state) - } - - pub(crate) fn serialize_as_physical_expr( - &self, - expr: &Expr, - schema: &DFSchema, - ) -> Result> { - let physical = self.create_physical_expr(expr, schema)?; - let proto = serialize_physical_expr(&physical, &DefaultPhysicalExtensionCodec {})?; - Ok(proto.encode_to_vec()) - } - - pub(crate) fn split_physical_plan( - &self, - key_indices: Vec, - aggregate: &LogicalPlan, - add_timestamp_field: bool, - ) -> Result { - let physical_plan = self.sync_plan(aggregate)?; - let codec = StreamingExtensionCodec { - context: StreamingDecodingContext::Planning, - }; - let mut physical_plan_node = - PhysicalPlanNode::try_from_physical_plan(physical_plan.clone(), &codec)?; - let PhysicalPlanType::Aggregate(mut final_aggregate_proto) = physical_plan_node - .physical_plan_type - .take() - .ok_or_else(|| DataFusionError::Plan("missing physical plan type".to_string()))? - else { - return plan_err!("unexpected physical plan type"); - }; - let AggregateMode::Final = final_aggregate_proto.mode() else { - return plan_err!("unexpected physical plan type"); - }; - - let partial_aggregation_plan = *final_aggregate_proto - .input - .take() - .ok_or_else(|| DataFusionError::Plan("missing input".to_string()))?; - - let partial_aggregation_exec_plan = partial_aggregation_plan.try_into_physical_plan( - self.schema_provider, - &RuntimeEnvBuilder::new().build().unwrap(), - &codec, - )?; - - let partial_schema = partial_aggregation_exec_plan.schema(); - let final_input_table_provider = FsMemExec::new("partial".into(), partial_schema.clone()); - - final_aggregate_proto.input = Some(Box::new(PhysicalPlanNode::try_from_physical_plan( - Arc::new(final_input_table_provider), - &codec, - )?)); - - let finish_plan = PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Aggregate(final_aggregate_proto)), - }; - - let (partial_schema, timestamp_index) = if add_timestamp_field { - ( - add_timestamp_field_arrow((*partial_schema).clone()), - partial_schema.fields().len(), - ) - } else { - (partial_schema.clone(), partial_schema.fields().len() - 1) - }; - - let partial_schema = FsSchema::new_keyed(partial_schema, timestamp_index, key_indices); - - Ok(SplitPlanOutput { - partial_aggregation_plan, - partial_schema, - finish_plan, - }) - } - - pub fn binning_function_proto( - &self, - width: Duration, - input_schema: DFSchemaRef, - ) -> Result { - let date_bin = date_bin().call(vec![ - Expr::Literal( - ScalarValue::IntervalMonthDayNano(Some(IntervalMonthDayNanoType::make_value( - 0, - 0, - width.as_nanos() as i64, - ))), - None, - ), - Expr::Column(datafusion::common::Column { - relation: None, - name: "_timestamp".into(), - spans: Spans::new(), - }), - ]); - - let binning_function = self.create_physical_expr(&date_bin, &input_schema)?; - serialize_physical_expr(&binning_function, &DefaultPhysicalExtensionCodec {}) - } -} - -struct FsExtensionPlanner {} - -#[async_trait] -impl ExtensionPlanner for FsExtensionPlanner { - async fn plan_extension( - &self, - _planner: &dyn PhysicalPlanner, - node: &dyn UserDefinedLogicalNode, - _logical_inputs: &[&LogicalPlan], - physical_inputs: &[Arc], - _session_state: &SessionState, - ) -> Result>> { - let schema = node.schema().as_ref().into(); - if let Ok::<&dyn StreamingOperatorBlueprint, _>(stream_extension) = node.try_into() - && stream_extension.is_passthrough_boundary() - { - match node.name() { - UNROLL_NODE_NAME => { - let node = node - .as_any() - .downcast_ref::() - .unwrap(); - let input = physical_inputs[0].clone(); - return Ok(Some(Arc::new(CdcDebeziumUnrollExec::try_new( - input, - node.pk_indices.clone(), - )?))); - } - PACK_NODE_NAME => { - let input = physical_inputs[0].clone(); - return Ok(Some(Arc::new(CdcDebeziumPackExec::try_new(input)?))); - } - _ => return Ok(None), - } - } - let name = if let Some(key_extension) = node.as_any().downcast_ref::() { - key_extension.operator_label.clone() - } else { - None - }; - Ok(Some(Arc::new(FsMemExec::new( - name.unwrap_or("memory".to_string()), - Arc::new(schema), - )))) - } -} - -impl PlanToGraphVisitor<'_> { - fn add_index_to_traversal(&mut self, index: NodeIndex) { - if let Some(last) = self.traversal.last_mut() { - last.push(index); - } - } - - pub(crate) fn add_plan(&mut self, plan: LogicalPlan) -> Result<()> { - self.traversal.clear(); - plan.visit(self)?; - Ok(()) - } - - pub fn into_graph(self) -> LogicalGraph { - self.graph - } - - pub fn build_extension( - &mut self, - input_nodes: Vec, - extension: &dyn StreamingOperatorBlueprint, - ) -> Result<()> { - if let Some(node_name) = extension.operator_identity() - && self.named_nodes.contains_key(&node_name) - { - return plan_err!( - "extension {:?} has already been planned, shouldn't try again.", - node_name - ); - } - - let input_schemas = input_nodes - .iter() - .map(|index| { - Ok(self - .output_schemas - .get(index) - .ok_or_else(|| DataFusionError::Plan("missing input node".to_string()))? - .clone()) - }) - .collect::>>()?; - - let CompiledTopologyNode { - execution_unit, - routing_edges, - } = extension - .compile_to_graph_node(&self.planner, self.graph.node_count(), input_schemas) - .map_err(|e| e.context(format!("planning operator {extension:?}")))?; - - let node_index = self.graph.add_node(execution_unit); - self.add_index_to_traversal(node_index); - - for (source, edge) in input_nodes.into_iter().zip(routing_edges) { - self.graph.add_edge(source, node_index, edge); - } - - self.output_schemas - .insert(node_index, extension.yielded_schema().into()); - - if let Some(node_name) = extension.operator_identity() { - self.named_nodes.insert(node_name, node_index); - } - Ok(()) - } -} - -impl TreeNodeVisitor<'_> for PlanToGraphVisitor<'_> { - type Node = LogicalPlan; - - fn f_down(&mut self, node: &Self::Node) -> Result { - let LogicalPlan::Extension(Extension { node }) = node else { - return Ok(TreeNodeRecursion::Continue); - }; - - let stream_extension: &dyn StreamingOperatorBlueprint = node - .try_into() - .map_err(|e: DataFusionError| e.context("converting extension"))?; - if stream_extension.is_passthrough_boundary() { - return Ok(TreeNodeRecursion::Continue); - } - - if let Some(name) = stream_extension.operator_identity() - && let Some(node_index) = self.named_nodes.get(&name) - { - self.add_index_to_traversal(*node_index); - return Ok(TreeNodeRecursion::Jump); - } - - if !node.inputs().is_empty() { - self.traversal.push(vec![]); - } - - Ok(TreeNodeRecursion::Continue) - } - - fn f_up(&mut self, node: &Self::Node) -> Result { - let LogicalPlan::Extension(Extension { node }) = node else { - return Ok(TreeNodeRecursion::Continue); - }; - - let stream_extension: &dyn StreamingOperatorBlueprint = node - .try_into() - .map_err(|e: DataFusionError| e.context("planning extension"))?; - - if stream_extension.is_passthrough_boundary() { - return Ok(TreeNodeRecursion::Continue); - } - - if let Some(name) = stream_extension.operator_identity() - && self.named_nodes.contains_key(&name) - { - return Ok(TreeNodeRecursion::Continue); - } - - let input_nodes = if !node.inputs().is_empty() { - self.traversal.pop().unwrap_or_default() - } else { - vec![] - }; - let stream_extension: &dyn StreamingOperatorBlueprint = node - .try_into() - .map_err(|e: DataFusionError| e.context("converting extension"))?; - self.build_extension(input_nodes, stream_extension)?; - - Ok(TreeNodeRecursion::Continue) - } -} diff --git a/src/sql/mod.rs b/src/sql/mod.rs deleted file mode 100644 index f4a0eef6..00000000 --- a/src/sql/mod.rs +++ /dev/null @@ -1,27 +0,0 @@ -// 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. - -pub mod api; -pub mod common; - -pub mod analysis; -pub mod connector; -pub mod functions; -pub mod logical_node; -pub mod logical_planner; -pub mod parse; -pub mod physical; -pub(crate) mod planning_runtime; -pub mod schema; -pub mod types; - -pub use analysis::rewrite_plan; diff --git a/src/sql/parse.rs b/src/sql/parse.rs deleted file mode 100644 index 8c9d4bb0..00000000 --- a/src/sql/parse.rs +++ /dev/null @@ -1,83 +0,0 @@ -// 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. - -//! FunctionStream SQL parsing (`parse_sql`). -//! -//! This module only performs lexical/syntactic parsing into sqlparser -//! [`Statement`](datafusion::sql::sqlparser::ast::Statement) values using -//! [`FunctionStreamDialect`]. Mapping those AST nodes to coordinator -//! [`Statement`](crate::coordinator::Statement) implementations is done by -//! [`crate::coordinator::classify_statement`]. -//! -//! **Data-definition / pipeline shape (supported forms in the dialect)** -//! - **`CREATE TABLE ... (cols [, WATERMARK FOR ...]) WITH (...)`** — connector-backed source DDL -//! - **`CREATE TABLE ...`** other forms (including `AS SELECT` where the dialect accepts it) -//! - **`CREATE STREAMING TABLE ... WITH (...) AS SELECT ...`** -//! - **`DROP TABLE`** / **`DROP STREAMING TABLE`** -//! - **`SHOW TABLES`**, **`SHOW STREAMING TABLE(S)`**, **`SHOW CREATE TABLE`**, **`SHOW CREATE STREAMING TABLE`** -//! -//! **`INSERT` is not supported** at the coordinator layer — use `CREATE TABLE ... AS SELECT` or -//! `CREATE STREAMING TABLE ... AS SELECT` instead (see coordinator classification). - -use datafusion::common::{Result, plan_err}; -use datafusion::error::DataFusionError; -use datafusion::sql::sqlparser::ast::Statement as DFStatement; -use datafusion::sql::sqlparser::dialect::FunctionStreamDialect; -use datafusion::sql::sqlparser::parser::Parser; - -/// Parse SQL text into zero or more dialect [`Statement`](DFStatement) nodes. -pub fn parse_sql(query: &str) -> Result> { - let trimmed = query.trim(); - if trimmed.is_empty() { - return plan_err!("Query is empty"); - } - - let dialect = FunctionStreamDialect {}; - let statements = Parser::parse_sql(&dialect, trimmed) - .map_err(|e| DataFusionError::Plan(format!("SQL parse error: {e}")))?; - - if statements.is_empty() { - return plan_err!("No SQL statements found"); - } - - Ok(statements) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_parse_multiple_statements_ast() { - let sql = concat!( - "CREATE TABLE t1 (id INT); ", - "CREATE STREAMING TABLE sk WITH ('connector' = 'kafka') AS SELECT id FROM t1", - ); - let stmts = parse_sql(sql).unwrap(); - assert_eq!(stmts.len(), 2); - assert!(matches!(stmts[0], DFStatement::CreateTable(_))); - assert!(matches!(stmts[1], DFStatement::CreateStreamingTable { .. })); - } - - #[test] - fn test_parse_empty() { - assert!(parse_sql("").is_err()); - assert!(parse_sql(" ").is_err()); - } - - #[test] - fn test_parse_select_yields_query_ast() { - let stmts = parse_sql("SELECT 1").unwrap(); - assert_eq!(stmts.len(), 1); - assert!(matches!(stmts[0], DFStatement::Query(_))); - } -} diff --git a/src/sql/physical/cdc/encode.rs b/src/sql/physical/cdc/encode.rs deleted file mode 100644 index 65ec758d..00000000 --- a/src/sql/physical/cdc/encode.rs +++ /dev/null @@ -1,342 +0,0 @@ -// 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::any::Any; -use std::collections::HashMap; -use std::pin::Pin; -use std::sync::Arc; -use std::task::{Context, Poll}; - -use datafusion::arrow::array::AsArray; -use datafusion::arrow::array::{ - Array, BooleanArray, FixedSizeBinaryArray, PrimitiveArray, RecordBatch, StringArray, - StructArray, TimestampNanosecondBuilder, UInt32Array, UInt32Builder, -}; -use datafusion::arrow::compute::take; -use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimestampNanosecondType}; -use datafusion::common::{DataFusionError, Result}; -use datafusion::execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext}; -use datafusion::physical_plan::{DisplayAs, ExecutionPlan, PlanProperties}; -use futures::{StreamExt, ready, stream::Stream}; - -use crate::sql::common::constants::{cdc, debezium_op_short, physical_plan_node_name}; -use crate::sql::common::{TIMESTAMP_FIELD, UPDATING_META_FIELD}; -use crate::sql::physical::source_exec::make_stream_properties; - -// ============================================================================ -// CdcDebeziumPackExec (Execution Plan Node) -// ============================================================================ - -/// Packs internal flat changelog rows into Debezium-style `before` / `after` / `op` / timestamp. -/// -/// Intended as the last physical node before a sink that expects Debezium CDC envelopes. -#[derive(Debug)] -pub struct CdcDebeziumPackExec { - input: Arc, - schema: SchemaRef, - properties: PlanProperties, -} - -impl CdcDebeziumPackExec { - pub fn try_new(input: Arc) -> Result { - let input_schema = input.schema(); - let timestamp_index = input_schema.index_of(TIMESTAMP_FIELD)?; - - let struct_fields: Vec<_> = input_schema - .fields() - .iter() - .enumerate() - .filter_map(|(index, field)| { - if field.name() == UPDATING_META_FIELD || index == timestamp_index { - None - } else { - Some(field.clone()) - } - }) - .collect(); - - let payload_struct_type = DataType::Struct(struct_fields.into()); - - let before_field = Arc::new(Field::new(cdc::BEFORE, payload_struct_type.clone(), true)); - let after_field = Arc::new(Field::new(cdc::AFTER, payload_struct_type, true)); - let op_field = Arc::new(Field::new(cdc::OP, DataType::Utf8, false)); - let timestamp_field = Arc::new(input_schema.field(timestamp_index).clone()); - - let output_schema = Arc::new(Schema::new(vec![ - before_field, - after_field, - op_field, - timestamp_field, - ])); - - Ok(Self { - input, - schema: output_schema.clone(), - properties: make_stream_properties(output_schema), - }) - } - - pub(crate) fn from_decoded_parts(input: Arc, schema: SchemaRef) -> Self { - Self { - properties: make_stream_properties(schema.clone()), - input, - schema, - } - } -} - -impl DisplayAs for CdcDebeziumPackExec { - fn fmt_as( - &self, - _t: datafusion::physical_plan::DisplayFormatType, - f: &mut std::fmt::Formatter, - ) -> std::fmt::Result { - write!(f, "CdcDebeziumPackExec") - } -} - -impl ExecutionPlan for CdcDebeziumPackExec { - fn name(&self) -> &str { - physical_plan_node_name::TO_DEBEZIUM_EXEC - } - - fn as_any(&self) -> &dyn Any { - self - } - fn schema(&self) -> SchemaRef { - self.schema.clone() - } - fn properties(&self) -> &PlanProperties { - &self.properties - } - fn children(&self) -> Vec<&Arc> { - vec![&self.input] - } - - fn with_new_children( - self: Arc, - children: Vec>, - ) -> Result> { - if children.len() != 1 { - return Err(DataFusionError::Internal( - "CdcDebeziumPackExec expects exactly 1 child".into(), - )); - } - Ok(Arc::new(Self::try_new(children[0].clone())?)) - } - - fn execute( - &self, - partition: usize, - context: Arc, - ) -> Result { - let updating_meta_index = self.input.schema().index_of(UPDATING_META_FIELD).ok(); - let timestamp_index = self.input.schema().index_of(TIMESTAMP_FIELD)?; - - let struct_projection = (0..self.input.schema().fields().len()) - .filter(|index| (updating_meta_index != Some(*index)) && *index != timestamp_index) - .collect(); - - Ok(Box::pin(CdcDebeziumPackStream { - input: self.input.execute(partition, context)?, - schema: self.schema.clone(), - updating_meta_index, - timestamp_index, - struct_projection, - })) - } - - fn reset(&self) -> Result<()> { - self.input.reset() - } -} - -// ============================================================================ -// CdcDebeziumPackStream (Physical Stream Execution) -// ============================================================================ - -struct CdcDebeziumPackStream { - input: SendableRecordBatchStream, - schema: SchemaRef, - updating_meta_index: Option, - timestamp_index: usize, - struct_projection: Vec, -} - -#[derive(Debug)] -struct RowCompactionState { - first_idx: usize, - last_idx: usize, - first_is_create: bool, - last_is_create: bool, - max_timestamp: i64, -} - -impl CdcDebeziumPackStream { - fn compact_changelog<'a>( - num_rows: usize, - is_retract: &'a BooleanArray, - id_array: &'a FixedSizeBinaryArray, - timestamps: &'a PrimitiveArray, - ) -> (Vec<&'a [u8]>, HashMap<&'a [u8], RowCompactionState>) { - let mut state_map: HashMap<&[u8], RowCompactionState> = HashMap::new(); - let mut unique_order = vec![]; - - for i in 0..num_rows { - let row_id = id_array.value(i); - let is_create = !is_retract.value(i); - let timestamp = timestamps.value(i); - - state_map - .entry(row_id) - .and_modify(|state| { - state.last_idx = i; - state.last_is_create = is_create; - state.max_timestamp = state.max_timestamp.max(timestamp); - }) - .or_insert_with(|| { - unique_order.push(row_id); - RowCompactionState { - first_idx: i, - last_idx: i, - first_is_create: is_create, - last_is_create: is_create, - max_timestamp: timestamp, - } - }); - } - (unique_order, state_map) - } - - fn as_debezium_batch(&mut self, batch: &RecordBatch) -> Result { - let value_struct = batch.project(&self.struct_projection)?; - let timestamps = batch - .column(self.timestamp_index) - .as_primitive::(); - - let columns: Vec> = if let Some(meta_index) = self.updating_meta_index { - let metadata = batch.column(meta_index).as_struct(); - let is_retract = metadata.column(0).as_boolean(); - let row_ids = metadata.column(1).as_fixed_size_binary(); - - let (ordered_ids, state_map) = - Self::compact_changelog(batch.num_rows(), is_retract, row_ids, timestamps); - - let mut before_builder = UInt32Builder::with_capacity(state_map.len()); - let mut after_builder = UInt32Builder::with_capacity(state_map.len()); - let mut op_vec = Vec::with_capacity(state_map.len()); - let mut ts_builder = TimestampNanosecondBuilder::with_capacity(state_map.len()); - - for row_id in ordered_ids { - let state = state_map - .get(row_id) - .expect("row id from order must exist in map"); - - match (state.first_is_create, state.last_is_create) { - (true, true) => { - before_builder.append_null(); - after_builder.append_value(state.last_idx as u32); - op_vec.push(debezium_op_short::CREATE); - } - (false, false) => { - before_builder.append_value(state.first_idx as u32); - after_builder.append_null(); - op_vec.push(debezium_op_short::DELETE); - } - (false, true) => { - before_builder.append_value(state.first_idx as u32); - after_builder.append_value(state.last_idx as u32); - op_vec.push(debezium_op_short::UPDATE); - } - (true, false) => { - continue; - } - } - ts_builder.append_value(state.max_timestamp); - } - - let before_indices = before_builder.finish(); - let after_indices = after_builder.finish(); - - let before_array = Self::take_struct_columns(&value_struct, &before_indices)?; - let after_array = Self::take_struct_columns(&value_struct, &after_indices)?; - let op_array = StringArray::from(op_vec); - - vec![ - Arc::new(before_array), - Arc::new(after_array), - Arc::new(op_array), - Arc::new(ts_builder.finish()), - ] - } else { - let num_rows = value_struct.num_rows(); - - let after_array = StructArray::try_new( - value_struct.schema().fields().clone(), - value_struct.columns().to_vec(), - None, - )?; - let before_array = - StructArray::new_null(value_struct.schema().fields().clone(), num_rows); - - let op_array = StringArray::from_iter_values(std::iter::repeat_n( - debezium_op_short::CREATE, - num_rows, - )); - - vec![ - Arc::new(before_array), - Arc::new(after_array), - Arc::new(op_array), - batch.column(self.timestamp_index).clone(), - ] - }; - - Ok(RecordBatch::try_new(self.schema.clone(), columns)?) - } - - fn take_struct_columns( - value_struct: &RecordBatch, - indices: &UInt32Array, - ) -> Result { - let mut arrays: Vec> = Vec::with_capacity(value_struct.num_columns()); - - for col in value_struct.columns() { - arrays.push(take(col.as_ref(), indices, None)?); - } - - Ok(StructArray::try_new( - value_struct.schema().fields().clone(), - arrays, - indices.nulls().cloned(), - )?) - } -} - -impl Stream for CdcDebeziumPackStream { - type Item = Result; - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let this = self.as_mut().get_mut(); - match ready!(this.input.poll_next_unpin(cx)) { - Some(Ok(batch)) => Poll::Ready(Some(this.as_debezium_batch(&batch))), - Some(Err(e)) => Poll::Ready(Some(Err(e))), - None => Poll::Ready(None), - } - } -} - -impl RecordBatchStream for CdcDebeziumPackStream { - fn schema(&self) -> SchemaRef { - self.schema.clone() - } -} diff --git a/src/sql/physical/cdc/mod.rs b/src/sql/physical/cdc/mod.rs deleted file mode 100644 index 216dd4c1..00000000 --- a/src/sql/physical/cdc/mod.rs +++ /dev/null @@ -1,17 +0,0 @@ -// 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. - -mod encode; -mod unroll; - -pub use encode::CdcDebeziumPackExec; -pub use unroll::CdcDebeziumUnrollExec; diff --git a/src/sql/physical/cdc/unroll.rs b/src/sql/physical/cdc/unroll.rs deleted file mode 100644 index 10c62c6c..00000000 --- a/src/sql/physical/cdc/unroll.rs +++ /dev/null @@ -1,322 +0,0 @@ -// 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::any::Any; -use std::pin::Pin; -use std::sync::Arc; -use std::task::{Context, Poll}; - -use datafusion::arrow::array::AsArray; -use datafusion::arrow::array::{ - BooleanBuilder, RecordBatch, StructArray, TimestampNanosecondBuilder, UInt32Builder, -}; -use datafusion::arrow::compute::{concat, take}; -use datafusion::arrow::datatypes::{ - DataType, Field, Schema, SchemaRef, TimeUnit, TimestampNanosecondType, -}; -use datafusion::common::{DataFusionError, Result, plan_err}; -use datafusion::execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext}; -use datafusion::logical_expr::ColumnarValue; -use datafusion::physical_plan::{DisplayAs, ExecutionPlan, PlanProperties}; -use futures::{StreamExt, ready, stream::Stream}; - -use crate::sql::common::TIMESTAMP_FIELD; -use crate::sql::common::constants::{cdc, debezium_op_short, physical_plan_node_name}; -use crate::sql::functions::MultiHashFunction; -use crate::sql::physical::meta::{updating_meta_field, updating_meta_fields}; -use crate::sql::physical::source_exec::make_stream_properties; - -// ============================================================================ -// CdcDebeziumUnrollExec (Execution Plan Node) -// ============================================================================ - -/// Physical node that unrolls Debezium CDC payloads (`before` / `after` / `op`) into a flat -/// changelog stream with retract metadata. -/// -/// - `c` / `r` → emit `after` (`is_retract = false`) -/// - `d` → emit `before` (`is_retract = true`) -/// - `u` → emit `before` (retract) then `after` (insert) -#[derive(Debug)] -pub struct CdcDebeziumUnrollExec { - input: Arc, - schema: SchemaRef, - properties: PlanProperties, - primary_key_indices: Vec, -} - -impl CdcDebeziumUnrollExec { - /// Builds the node and validates Debezium payload schema constraints. - pub fn try_new(input: Arc, primary_key_indices: Vec) -> Result { - let input_schema = input.schema(); - - let before_index = input_schema.index_of(cdc::BEFORE)?; - let after_index = input_schema.index_of(cdc::AFTER)?; - let op_index = input_schema.index_of(cdc::OP)?; - let _timestamp_index = input_schema.index_of(TIMESTAMP_FIELD)?; - - let before_type = input_schema.field(before_index).data_type(); - let after_type = input_schema.field(after_index).data_type(); - - if before_type != after_type { - return Err(DataFusionError::Plan( - "CDC 'before' and 'after' columns must share the exact same DataType".to_string(), - )); - } - - if *input_schema.field(op_index).data_type() != DataType::Utf8 { - return Err(DataFusionError::Plan( - "CDC 'op' (operation) column must be of type Utf8 (String)".to_string(), - )); - } - - let DataType::Struct(fields) = before_type else { - return Err(DataFusionError::Plan( - "CDC 'before' and 'after' payload columns must be Structs".to_string(), - )); - }; - - let mut unrolled_fields = fields.to_vec(); - unrolled_fields.push(updating_meta_field()); - unrolled_fields.push(Arc::new(Field::new( - TIMESTAMP_FIELD, - DataType::Timestamp(TimeUnit::Nanosecond, None), - false, - ))); - - let schema = Arc::new(Schema::new(unrolled_fields)); - - Ok(Self { - input, - schema: schema.clone(), - properties: make_stream_properties(schema), - primary_key_indices, - }) - } - - /// Used when deserializing a plan with a pre-baked output schema (see [`StreamingExtensionCodec`]). - pub(crate) fn from_decoded_parts( - input: Arc, - schema: SchemaRef, - primary_key_indices: Vec, - ) -> Self { - Self { - properties: make_stream_properties(schema.clone()), - input, - schema, - primary_key_indices, - } - } - - pub fn primary_key_indices(&self) -> &[usize] { - &self.primary_key_indices - } -} - -impl DisplayAs for CdcDebeziumUnrollExec { - fn fmt_as( - &self, - _t: datafusion::physical_plan::DisplayFormatType, - f: &mut std::fmt::Formatter, - ) -> std::fmt::Result { - write!(f, "CdcDebeziumUnrollExec") - } -} - -impl ExecutionPlan for CdcDebeziumUnrollExec { - fn name(&self) -> &str { - physical_plan_node_name::DEBEZIUM_UNROLLING_EXEC - } - - fn as_any(&self) -> &dyn Any { - self - } - fn schema(&self) -> SchemaRef { - self.schema.clone() - } - fn properties(&self) -> &PlanProperties { - &self.properties - } - fn children(&self) -> Vec<&Arc> { - vec![&self.input] - } - - fn with_new_children( - self: Arc, - children: Vec>, - ) -> Result> { - if children.len() != 1 { - return Err(DataFusionError::Internal( - "CdcDebeziumUnrollExec expects exactly one child".to_string(), - )); - } - Ok(Arc::new(Self { - input: children[0].clone(), - schema: self.schema.clone(), - properties: self.properties.clone(), - primary_key_indices: self.primary_key_indices.clone(), - })) - } - - fn execute( - &self, - partition: usize, - context: Arc, - ) -> Result { - Ok(Box::pin(CdcDebeziumUnrollStream::try_new( - self.input.execute(partition, context)?, - self.schema.clone(), - self.primary_key_indices.clone(), - )?)) - } - - fn reset(&self) -> Result<()> { - self.input.reset() - } -} - -// ============================================================================ -// CdcDebeziumUnrollStream (Physical Stream Execution) -// ============================================================================ - -struct CdcDebeziumUnrollStream { - input: SendableRecordBatchStream, - schema: SchemaRef, - before_index: usize, - after_index: usize, - op_index: usize, - timestamp_index: usize, - primary_key_indices: Vec, -} - -impl CdcDebeziumUnrollStream { - fn try_new( - input: SendableRecordBatchStream, - schema: SchemaRef, - primary_key_indices: Vec, - ) -> Result { - if primary_key_indices.is_empty() { - return plan_err!( - "A CDC source requires at least one primary key to maintain state correctly." - ); - } - - let input_schema = input.schema(); - Ok(Self { - input, - schema, - before_index: input_schema.index_of(cdc::BEFORE)?, - after_index: input_schema.index_of(cdc::AFTER)?, - op_index: input_schema.index_of(cdc::OP)?, - timestamp_index: input_schema.index_of(TIMESTAMP_FIELD)?, - primary_key_indices, - }) - } - - fn unroll_batch(&self, batch: &RecordBatch) -> Result { - let num_rows = batch.num_rows(); - if num_rows == 0 { - return Ok(RecordBatch::new_empty(self.schema.clone())); - } - - let before_col = batch.column(self.before_index); - let after_col = batch.column(self.after_index); - - let op_array = batch.column(self.op_index).as_string::(); - let timestamp_array = batch - .column(self.timestamp_index) - .as_primitive::(); - - let max_capacity = num_rows * 2; - let mut take_indices = UInt32Builder::with_capacity(max_capacity); - let mut is_retract_builder = BooleanBuilder::with_capacity(max_capacity); - let mut timestamp_builder = TimestampNanosecondBuilder::with_capacity(max_capacity); - - for i in 0..num_rows { - let op = op_array.value(i); - let ts = timestamp_array.value(i); - - match op { - debezium_op_short::CREATE | debezium_op_short::READ => { - take_indices.append_value((i + num_rows) as u32); - is_retract_builder.append_value(false); - timestamp_builder.append_value(ts); - } - debezium_op_short::DELETE => { - take_indices.append_value(i as u32); - is_retract_builder.append_value(true); - timestamp_builder.append_value(ts); - } - debezium_op_short::UPDATE => { - take_indices.append_value(i as u32); - is_retract_builder.append_value(true); - timestamp_builder.append_value(ts); - - take_indices.append_value((i + num_rows) as u32); - is_retract_builder.append_value(false); - timestamp_builder.append_value(ts); - } - _ => { - return Err(DataFusionError::Execution(format!( - "Encountered unexpected Debezium operation code: '{op}'" - ))); - } - } - } - - let take_indices = take_indices.finish(); - let unrolled_row_count = take_indices.len(); - - let combined_array = concat(&[before_col.as_ref(), after_col.as_ref()])?; - let unrolled_array = take(&combined_array, &take_indices, None)?; - - let mut final_columns = unrolled_array.as_struct().columns().to_vec(); - - let pk_columns: Vec = self - .primary_key_indices - .iter() - .map(|&idx| ColumnarValue::Array(Arc::clone(&final_columns[idx]))) - .collect(); - - let hash_column = MultiHashFunction::default().invoke(&pk_columns)?; - let ids_array = hash_column.into_array(unrolled_row_count)?; - - let meta_struct = StructArray::try_new( - updating_meta_fields(), - vec![Arc::new(is_retract_builder.finish()), ids_array], - None, - )?; - - final_columns.push(Arc::new(meta_struct)); - final_columns.push(Arc::new(timestamp_builder.finish())); - - Ok(RecordBatch::try_new(self.schema.clone(), final_columns)?) - } -} - -impl Stream for CdcDebeziumUnrollStream { - type Item = Result; - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let this = self.as_mut().get_mut(); - match ready!(this.input.poll_next_unpin(cx)) { - Some(Ok(batch)) => Poll::Ready(Some(this.unroll_batch(&batch))), - Some(Err(e)) => Poll::Ready(Some(Err(e))), - None => Poll::Ready(None), - } - } -} - -impl RecordBatchStream for CdcDebeziumUnrollStream { - fn schema(&self) -> SchemaRef { - self.schema.clone() - } -} diff --git a/src/sql/physical/codec.rs b/src/sql/physical/codec.rs deleted file mode 100644 index 1b96a9d6..00000000 --- a/src/sql/physical/codec.rs +++ /dev/null @@ -1,307 +0,0 @@ -// 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::sync::Arc; - -use datafusion::arrow::array::RecordBatch; -use datafusion::arrow::datatypes::Schema; -use datafusion::common::{DataFusionError, Result, UnnestOptions, not_impl_err}; -use datafusion::execution::FunctionRegistry; -use datafusion::logical_expr::ScalarUDF; -use datafusion::physical_plan::ExecutionPlan; -use datafusion::physical_plan::unnest::{ListUnnest, UnnestExec}; -use datafusion_proto::physical_plan::PhysicalExtensionCodec; -use prost::Message; -use protocol::function_stream_graph::{ - DebeziumDecodeNode, DebeziumEncodeNode, FsExecNode, MemExecNode, UnnestExecNode, - fs_exec_node::Node, -}; -use tokio::sync::mpsc::UnboundedReceiver; - -use crate::sql::analysis::UNNESTED_COL; -use crate::sql::common::constants::{mem_exec_join_side, window_function_udf}; -use crate::sql::physical::cdc::{CdcDebeziumPackExec, CdcDebeziumUnrollExec}; -use crate::sql::physical::source_exec::{ - BufferedBatchesExec, InjectableSingleBatchExec, MpscReceiverStreamExec, PlanningPlaceholderExec, -}; -use crate::sql::physical::udfs::window; - -// ============================================================================ -// StreamingExtensionCodec & StreamingDecodingContext -// ============================================================================ - -/// Worker-side context used when deserializing a physical plan from the coordinator. -/// -/// Planning uses [`PlanningPlaceholderExec`]; at runtime this selects the real source -/// implementation (locked batch, MPSC stream, join sides, etc.). -#[derive(Debug)] -pub enum StreamingDecodingContext { - None, - Planning, - SingleLockedBatch(Arc>>), - UnboundedBatchStream(Arc>>>), - LockedBatchVec(Arc>>), - LockedJoinPair { - left: Arc>>, - right: Arc>>, - }, - LockedJoinStream { - left: Arc>>>, - right: Arc>>>, - }, -} - -/// Codec for custom streaming physical extension nodes (`FsExecNode` protobuf). -#[derive(Debug)] -pub struct StreamingExtensionCodec { - pub context: StreamingDecodingContext, -} - -impl Default for StreamingExtensionCodec { - fn default() -> Self { - Self { - context: StreamingDecodingContext::None, - } - } -} - -impl PhysicalExtensionCodec for StreamingExtensionCodec { - fn try_decode( - &self, - buf: &[u8], - inputs: &[Arc], - _registry: &dyn FunctionRegistry, - ) -> Result> { - let exec: FsExecNode = Message::decode(buf).map_err(|err| { - DataFusionError::Internal(format!("Failed to deserialize FsExecNode protobuf: {err}")) - })?; - - let node = exec.node.ok_or_else(|| { - DataFusionError::Internal("Decoded FsExecNode contains no inner node data".to_string()) - })?; - - match node { - Node::MemExec(mem) => self.decode_placeholder_exec(mem), - Node::UnnestExec(unnest) => decode_unnest_exec(unnest, inputs), - Node::DebeziumDecode(debezium) => decode_debezium_unroll(debezium, inputs), - Node::DebeziumEncode(debezium) => decode_debezium_pack(debezium, inputs), - } - } - - fn try_encode(&self, node: Arc, buf: &mut Vec) -> Result<()> { - let mut proto = None; - - if let Some(table) = node.as_any().downcast_ref::() { - let schema_json = serde_json::to_string(&table.schema).map_err(|e| { - DataFusionError::Internal(format!("Failed to serialize schema to JSON: {e}")) - })?; - - proto = Some(FsExecNode { - node: Some(Node::MemExec(MemExecNode { - table_name: table.table_name.clone(), - schema: schema_json, - })), - }); - } else if let Some(unnest) = node.as_any().downcast_ref::() { - let schema_json = serde_json::to_string(&unnest.schema()).map_err(|e| { - DataFusionError::Internal(format!("Failed to serialize unnest schema to JSON: {e}")) - })?; - - proto = Some(FsExecNode { - node: Some(Node::UnnestExec(UnnestExecNode { - schema: schema_json, - })), - }); - } else if let Some(decode) = node.as_any().downcast_ref::() { - let schema_json = serde_json::to_string(decode.schema().as_ref()).map_err(|e| { - DataFusionError::Internal(format!("Failed to serialize CDC unroll schema: {e}")) - })?; - - proto = Some(FsExecNode { - node: Some(Node::DebeziumDecode(DebeziumDecodeNode { - schema: schema_json, - primary_keys: decode - .primary_key_indices() - .iter() - .map(|&c| c as u64) - .collect(), - })), - }); - } else if let Some(encode) = node.as_any().downcast_ref::() { - let schema_json = serde_json::to_string(encode.schema().as_ref()).map_err(|e| { - DataFusionError::Internal(format!("Failed to serialize CDC pack schema: {e}")) - })?; - - proto = Some(FsExecNode { - node: Some(Node::DebeziumEncode(DebeziumEncodeNode { - schema: schema_json, - })), - }); - } - - if let Some(proto_node) = proto { - proto_node.encode(buf).map_err(|err| { - DataFusionError::Internal(format!("Failed to encode protobuf node: {err}")) - })?; - Ok(()) - } else { - Err(DataFusionError::Internal(format!( - "Cannot serialize unknown physical plan node: {node:?}" - ))) - } - } - - fn try_decode_udf(&self, name: &str, _buf: &[u8]) -> Result> { - if name == window_function_udf::NAME { - return Ok(window()); - } - not_impl_err!("PhysicalExtensionCodec does not support scalar function '{name}'") - } -} - -impl StreamingExtensionCodec { - fn decode_placeholder_exec(&self, mem_exec: MemExecNode) -> Result> { - let schema: Schema = serde_json::from_str(&mem_exec.schema).map_err(|e| { - DataFusionError::Internal(format!("Invalid schema JSON in exec codec: {e:?}")) - })?; - let schema = Arc::new(schema); - - match &self.context { - StreamingDecodingContext::SingleLockedBatch(single_batch) => Ok(Arc::new( - InjectableSingleBatchExec::new(schema, single_batch.clone()), - )), - StreamingDecodingContext::UnboundedBatchStream(unbounded_stream) => Ok(Arc::new( - MpscReceiverStreamExec::new(schema, unbounded_stream.clone()), - )), - StreamingDecodingContext::LockedBatchVec(locked_batches) => Ok(Arc::new( - BufferedBatchesExec::new(schema, locked_batches.clone()), - )), - StreamingDecodingContext::Planning => Ok(Arc::new(PlanningPlaceholderExec::new( - mem_exec.table_name, - schema, - ))), - StreamingDecodingContext::None => Err(DataFusionError::Internal( - "A valid StreamingDecodingContext is required to decode placeholders into execution streams.".into(), - )), - StreamingDecodingContext::LockedJoinPair { left, right } => { - match mem_exec.table_name.as_str() { - mem_exec_join_side::LEFT => Ok(Arc::new(InjectableSingleBatchExec::new( - schema, - left.clone(), - ))), - mem_exec_join_side::RIGHT => Ok(Arc::new(InjectableSingleBatchExec::new( - schema, - right.clone(), - ))), - _ => Err(DataFusionError::Internal(format!( - "Unknown join side table name: {}", - mem_exec.table_name - ))), - } - } - StreamingDecodingContext::LockedJoinStream { left, right } => { - match mem_exec.table_name.as_str() { - mem_exec_join_side::LEFT => Ok(Arc::new(MpscReceiverStreamExec::new( - schema, - left.clone(), - ))), - mem_exec_join_side::RIGHT => Ok(Arc::new(MpscReceiverStreamExec::new( - schema, - right.clone(), - ))), - _ => Err(DataFusionError::Internal(format!( - "Unknown join side table name: {}", - mem_exec.table_name - ))), - } - } - } - } -} - -fn decode_unnest_exec( - unnest: UnnestExecNode, - inputs: &[Arc], -) -> Result> { - let schema: Schema = serde_json::from_str(&unnest.schema) - .map_err(|e| DataFusionError::Internal(format!("Invalid unnest schema JSON: {e:?}")))?; - - let column = schema.index_of(UNNESTED_COL).map_err(|_| { - DataFusionError::Internal(format!( - "Unnest schema missing required column: {UNNESTED_COL}" - )) - })?; - - let input = inputs.first().ok_or_else(|| { - DataFusionError::Internal("UnnestExec requires exactly one input plan".to_string()) - })?; - - Ok(Arc::new(UnnestExec::new( - input.clone(), - vec![ListUnnest { - index_in_input_schema: column, - depth: 1, - }], - vec![], - Arc::new(schema), - UnnestOptions::default(), - ))) -} - -fn decode_debezium_unroll( - debezium: DebeziumDecodeNode, - inputs: &[Arc], -) -> Result> { - let schema = Arc::new( - serde_json::from_str::(&debezium.schema).map_err(|e| { - DataFusionError::Internal(format!("Invalid DebeziumDecode schema JSON: {e:?}")) - })?, - ); - - let input = inputs.first().ok_or_else(|| { - DataFusionError::Internal( - "CdcDebeziumUnrollExec requires exactly one input plan".to_string(), - ) - })?; - - let primary_keys = debezium - .primary_keys - .into_iter() - .map(|c| c as usize) - .collect(); - - Ok(Arc::new(CdcDebeziumUnrollExec::from_decoded_parts( - input.clone(), - schema, - primary_keys, - ))) -} - -fn decode_debezium_pack( - debezium: DebeziumEncodeNode, - inputs: &[Arc], -) -> Result> { - let schema = Arc::new( - serde_json::from_str::(&debezium.schema).map_err(|e| { - DataFusionError::Internal(format!("Invalid DebeziumEncode schema JSON: {e:?}")) - })?, - ); - - let input = inputs.first().ok_or_else(|| { - DataFusionError::Internal("CdcDebeziumPackExec requires exactly one input plan".to_string()) - })?; - - Ok(Arc::new(CdcDebeziumPackExec::from_decoded_parts( - input.clone(), - schema, - ))) -} diff --git a/src/sql/physical/meta.rs b/src/sql/physical/meta.rs deleted file mode 100644 index 1387482c..00000000 --- a/src/sql/physical/meta.rs +++ /dev/null @@ -1,47 +0,0 @@ -// 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::sync::{Arc, OnceLock}; - -use datafusion::arrow::datatypes::{DataType, Field, Fields}; - -use crate::sql::common::UPDATING_META_FIELD; -use crate::sql::common::constants::updating_state_field; - -pub fn updating_meta_fields() -> Fields { - static FIELDS: OnceLock = OnceLock::new(); - FIELDS - .get_or_init(|| { - Fields::from(vec![ - Field::new(updating_state_field::IS_RETRACT, DataType::Boolean, true), - Field::new( - updating_state_field::ID, - DataType::FixedSizeBinary(16), - true, - ), - ]) - }) - .clone() -} - -pub fn updating_meta_field() -> Arc { - static FIELD: OnceLock> = OnceLock::new(); - FIELD - .get_or_init(|| { - Arc::new(Field::new( - UPDATING_META_FIELD, - DataType::Struct(updating_meta_fields()), - false, - )) - }) - .clone() -} diff --git a/src/sql/physical/mod.rs b/src/sql/physical/mod.rs deleted file mode 100644 index 77f5c107..00000000 --- a/src/sql/physical/mod.rs +++ /dev/null @@ -1,23 +0,0 @@ -// 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. - -mod cdc; -mod codec; -mod meta; -mod source_exec; -mod udfs; - -pub use cdc::{CdcDebeziumPackExec, CdcDebeziumUnrollExec}; -pub use codec::{StreamingDecodingContext, StreamingExtensionCodec}; -pub use meta::{updating_meta_field, updating_meta_fields}; -pub use source_exec::FsMemExec; -pub use udfs::window; diff --git a/src/sql/physical/source_exec.rs b/src/sql/physical/source_exec.rs deleted file mode 100644 index fa65cbfd..00000000 --- a/src/sql/physical/source_exec.rs +++ /dev/null @@ -1,400 +0,0 @@ -// 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::any::Any; -use std::mem; -use std::sync::Arc; - -use datafusion::arrow::array::RecordBatch; -use datafusion::arrow::datatypes::SchemaRef; -use datafusion::catalog::memory::MemorySourceConfig; -use datafusion::common::{DataFusionError, Result, Statistics, not_impl_err, plan_err}; -use datafusion::datasource::memory::DataSourceExec; -use datafusion::execution::{SendableRecordBatchStream, TaskContext}; -use datafusion::physical_expr::EquivalenceProperties; -use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; -use datafusion::physical_plan::memory::MemoryStream; -use datafusion::physical_plan::stream::RecordBatchStreamAdapter; -use datafusion::physical_plan::{DisplayAs, ExecutionPlan, Partitioning, PlanProperties}; -use futures::StreamExt; -use tokio::sync::mpsc::UnboundedReceiver; -use tokio_stream::wrappers::UnboundedReceiverStream; - -use crate::sql::common::constants::physical_plan_node_name; - -/// Standard [`PlanProperties`] for a continuous, unbounded stream: incremental emission, -/// unknown partitioning, and unbounded boundedness (without requiring infinite memory). -pub(crate) fn create_unbounded_stream_properties(schema: SchemaRef) -> PlanProperties { - PlanProperties::new( - EquivalenceProperties::new(schema), - Partitioning::UnknownPartitioning(1), - EmissionType::Incremental, - Boundedness::Unbounded { - requires_infinite_memory: false, - }, - ) -} - -/// Alias for call sites that still use the older name. -pub(crate) fn make_stream_properties(schema: SchemaRef) -> PlanProperties { - create_unbounded_stream_properties(schema) -} - -// ============================================================================ -// InjectableSingleBatchExec (formerly RwLockRecordBatchReader) -// ============================================================================ - -/// Yields exactly one [`RecordBatch`], injected via a lock before `execute` runs. -/// -/// For event-driven loops that receive a single batch from the network and run a DataFusion -/// plan over it, the batch is stored in the lock until execution starts. -#[derive(Debug)] -pub(crate) struct InjectableSingleBatchExec { - schema: SchemaRef, - injected_batch: Arc>>, - properties: PlanProperties, -} - -impl InjectableSingleBatchExec { - pub(crate) fn new( - schema: SchemaRef, - injected_batch: Arc>>, - ) -> Self { - Self { - schema: schema.clone(), - injected_batch, - properties: create_unbounded_stream_properties(schema), - } - } -} - -impl DisplayAs for InjectableSingleBatchExec { - fn fmt_as( - &self, - _t: datafusion::physical_plan::DisplayFormatType, - f: &mut std::fmt::Formatter, - ) -> std::fmt::Result { - write!(f, "InjectableSingleBatchExec") - } -} - -impl ExecutionPlan for InjectableSingleBatchExec { - fn as_any(&self) -> &dyn Any { - self - } - fn schema(&self) -> SchemaRef { - self.schema.clone() - } - fn children(&self) -> Vec<&Arc> { - vec![] - } - fn properties(&self) -> &PlanProperties { - &self.properties - } - fn name(&self) -> &str { - physical_plan_node_name::RW_LOCK_READER - } - - fn with_new_children( - self: Arc, - _children: Vec>, - ) -> Result> { - Err(DataFusionError::Internal( - "InjectableSingleBatchExec does not support children".into(), - )) - } - - fn execute( - &self, - _partition: usize, - _context: Arc, - ) -> Result { - let mut guard = self.injected_batch.write().map_err(|e| { - DataFusionError::Execution(format!("Failed to acquire write lock: {e}")) - })?; - - let batch = guard.take().ok_or_else(|| { - DataFusionError::Execution( - "Execution triggered, but no RecordBatch was injected into the node.".into(), - ) - })?; - - Ok(Box::pin(MemoryStream::try_new( - vec![batch], - self.schema.clone(), - None, - )?)) - } - - fn statistics(&self) -> Result { - Ok(Statistics::new_unknown(&self.schema)) - } - - fn reset(&self) -> Result<()> { - Ok(()) - } -} - -// ============================================================================ -// MpscReceiverStreamExec (formerly UnboundedRecordBatchReader) -// ============================================================================ - -/// Unbounded streaming source backed by a Tokio `mpsc` receiver. -/// -/// Bridges async producers (e.g. network threads) into a DataFusion pipeline. -#[derive(Debug)] -pub(crate) struct MpscReceiverStreamExec { - schema: SchemaRef, - channel_receiver: Arc>>>, - properties: PlanProperties, -} - -impl MpscReceiverStreamExec { - pub(crate) fn new( - schema: SchemaRef, - channel_receiver: Arc>>>, - ) -> Self { - Self { - schema: schema.clone(), - channel_receiver, - properties: create_unbounded_stream_properties(schema), - } - } -} - -impl DisplayAs for MpscReceiverStreamExec { - fn fmt_as( - &self, - _t: datafusion::physical_plan::DisplayFormatType, - f: &mut std::fmt::Formatter, - ) -> std::fmt::Result { - write!(f, "MpscReceiverStreamExec") - } -} - -impl ExecutionPlan for MpscReceiverStreamExec { - fn as_any(&self) -> &dyn Any { - self - } - fn schema(&self) -> SchemaRef { - self.schema.clone() - } - fn children(&self) -> Vec<&Arc> { - vec![] - } - fn properties(&self) -> &PlanProperties { - &self.properties - } - fn name(&self) -> &str { - physical_plan_node_name::UNBOUNDED_READER - } - - fn with_new_children( - self: Arc, - _children: Vec>, - ) -> Result> { - Err(DataFusionError::Internal( - "MpscReceiverStreamExec does not support children".into(), - )) - } - - fn execute( - &self, - _partition: usize, - _context: Arc, - ) -> Result { - let mut guard = self.channel_receiver.write().map_err(|e| { - DataFusionError::Execution(format!("Failed to acquire lock for MPSC receiver: {e}")) - })?; - - let receiver = guard.take().ok_or_else(|| { - DataFusionError::Execution( - "The MPSC receiver was already consumed by a previous execution.".into(), - ) - })?; - - Ok(Box::pin(RecordBatchStreamAdapter::new( - self.schema.clone(), - UnboundedReceiverStream::new(receiver).map(Ok), - ))) - } - - fn statistics(&self) -> Result { - Ok(Statistics::new_unknown(&self.schema)) - } - - fn reset(&self) -> Result<()> { - Ok(()) - } -} - -// ============================================================================ -// BufferedBatchesExec (formerly RecordBatchVecReader) -// ============================================================================ - -/// Drains a growable, locked `Vec` when `execute` runs (micro-batching). -#[derive(Debug)] -pub(crate) struct BufferedBatchesExec { - schema: SchemaRef, - buffered_batches: Arc>>, - properties: PlanProperties, -} - -impl BufferedBatchesExec { - pub(crate) fn new( - schema: SchemaRef, - buffered_batches: Arc>>, - ) -> Self { - Self { - schema: schema.clone(), - buffered_batches, - properties: create_unbounded_stream_properties(schema), - } - } -} - -impl DisplayAs for BufferedBatchesExec { - fn fmt_as( - &self, - _t: datafusion::physical_plan::DisplayFormatType, - f: &mut std::fmt::Formatter, - ) -> std::fmt::Result { - write!(f, "BufferedBatchesExec") - } -} - -impl ExecutionPlan for BufferedBatchesExec { - fn as_any(&self) -> &dyn Any { - self - } - fn schema(&self) -> SchemaRef { - self.schema.clone() - } - fn children(&self) -> Vec<&Arc> { - vec![] - } - fn properties(&self) -> &PlanProperties { - &self.properties - } - fn name(&self) -> &str { - physical_plan_node_name::VEC_READER - } - - fn with_new_children( - self: Arc, - _children: Vec>, - ) -> Result> { - Err(DataFusionError::Internal( - "BufferedBatchesExec does not support children".into(), - )) - } - - fn execute( - &self, - partition: usize, - context: Arc, - ) -> Result { - let mut guard = self.buffered_batches.write().map_err(|e| { - DataFusionError::Execution(format!("Failed to acquire lock for buffered batches: {e}")) - })?; - - let accumulated_batches = mem::take(&mut *guard); - - let memory_config = - MemorySourceConfig::try_new(&[accumulated_batches], self.schema.clone(), None)?; - - DataSourceExec::new(Arc::new(memory_config)).execute(partition, context) - } - - fn statistics(&self) -> Result { - Ok(Statistics::new_unknown(&self.schema)) - } - - fn reset(&self) -> Result<()> { - Ok(()) - } -} - -// ============================================================================ -#[derive(Debug, Clone)] -pub struct PlanningPlaceholderExec { - pub table_name: String, - pub schema: SchemaRef, - properties: PlanProperties, -} - -impl PlanningPlaceholderExec { - pub fn new(table_name: String, schema: SchemaRef) -> Self { - Self { - schema: schema.clone(), - table_name, - properties: create_unbounded_stream_properties(schema), - } - } -} - -impl DisplayAs for PlanningPlaceholderExec { - fn fmt_as( - &self, - _t: datafusion::physical_plan::DisplayFormatType, - f: &mut std::fmt::Formatter, - ) -> std::fmt::Result { - write!(f, "PlanningPlaceholderExec: schema={}", self.schema) - } -} - -impl ExecutionPlan for PlanningPlaceholderExec { - fn as_any(&self) -> &dyn Any { - self - } - fn schema(&self) -> SchemaRef { - self.schema.clone() - } - fn children(&self) -> Vec<&Arc> { - vec![] - } - fn properties(&self) -> &PlanProperties { - &self.properties - } - fn name(&self) -> &str { - physical_plan_node_name::MEM_EXEC - } - - fn with_new_children( - self: Arc, - _children: Vec>, - ) -> Result> { - not_impl_err!("PlanningPlaceholderExec does not accept children.") - } - - fn execute( - &self, - _partition: usize, - _context: Arc, - ) -> Result { - plan_err!("PlanningPlaceholderExec cannot be executed; swap for a real source before run.") - } - - fn statistics(&self) -> Result { - Ok(Statistics::new_unknown(&self.schema)) - } - - fn reset(&self) -> Result<()> { - Ok(()) - } -} - -// Backward-compatible aliases -pub type FsMemExec = PlanningPlaceholderExec; -pub type RwLockRecordBatchReader = InjectableSingleBatchExec; -pub type UnboundedRecordBatchReader = MpscReceiverStreamExec; -pub type RecordBatchVecReader = BufferedBatchesExec; diff --git a/src/sql/physical/udfs.rs b/src/sql/physical/udfs.rs deleted file mode 100644 index 767abf06..00000000 --- a/src/sql/physical/udfs.rs +++ /dev/null @@ -1,138 +0,0 @@ -// 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::any::Any; -use std::sync::Arc; - -use datafusion::arrow::array::StructArray; -use datafusion::arrow::datatypes::{DataType, TimeUnit}; -use datafusion::common::{Result, ScalarValue, exec_err}; -use datafusion::logical_expr::{ - ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility, -}; - -use crate::make_udf_function; -use crate::sql::common::constants::window_function_udf; -use crate::sql::schema::utils::window_arrow_struct; - -// ============================================================================ -// WindowFunctionUdf (User-Defined Scalar Function) -// ============================================================================ - -/// UDF that packs two nanosecond timestamps into the canonical window `Struct` type. -/// -/// Stream SQL uses a single struct column `[start, end)` for tumbling/hopping windows; -/// this keeps `GROUP BY` and physical codec alignment on one Arrow shape. -#[derive(Debug)] -pub struct WindowFunctionUdf { - signature: Signature, -} - -impl Default for WindowFunctionUdf { - fn default() -> Self { - Self { - signature: Signature::new( - TypeSignature::Exact(vec![ - DataType::Timestamp(TimeUnit::Nanosecond, None), - DataType::Timestamp(TimeUnit::Nanosecond, None), - ]), - Volatility::Immutable, - ), - } - } -} - -impl ScalarUDFImpl for WindowFunctionUdf { - fn as_any(&self) -> &dyn Any { - self - } - - fn name(&self) -> &str { - window_function_udf::NAME - } - - fn signature(&self) -> &Signature { - &self.signature - } - - fn return_type(&self, _arg_types: &[DataType]) -> Result { - Ok(window_arrow_struct()) - } - - fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - let columns = args.args; - - if columns.len() != 2 { - return exec_err!( - "Window UDF expected exactly 2 arguments, but received {}", - columns.len() - ); - } - - let DataType::Struct(fields) = window_arrow_struct() else { - return exec_err!( - "Internal Engine Error: window_arrow_struct() must return a Struct DataType" - ); - }; - - let start_val = &columns[0]; - let end_val = &columns[1]; - - if !matches!( - start_val.data_type(), - DataType::Timestamp(TimeUnit::Nanosecond, _) - ) { - return exec_err!("Window UDF expected first argument to be a Nanosecond Timestamp"); - } - if !matches!( - end_val.data_type(), - DataType::Timestamp(TimeUnit::Nanosecond, _) - ) { - return exec_err!("Window UDF expected second argument to be a Nanosecond Timestamp"); - } - - match (start_val, end_val) { - (ColumnarValue::Array(start_arr), ColumnarValue::Array(end_arr)) => { - let struct_array = - StructArray::try_new(fields, vec![start_arr.clone(), end_arr.clone()], None)?; - Ok(ColumnarValue::Array(Arc::new(struct_array))) - } - - (ColumnarValue::Array(start_arr), ColumnarValue::Scalar(end_scalar)) => { - let end_arr = end_scalar.to_array_of_size(start_arr.len())?; - let struct_array = - StructArray::try_new(fields, vec![start_arr.clone(), end_arr], None)?; - Ok(ColumnarValue::Array(Arc::new(struct_array))) - } - - (ColumnarValue::Scalar(start_scalar), ColumnarValue::Array(end_arr)) => { - let start_arr = start_scalar.to_array_of_size(end_arr.len())?; - let struct_array = - StructArray::try_new(fields, vec![start_arr, end_arr.clone()], None)?; - Ok(ColumnarValue::Array(Arc::new(struct_array))) - } - - (ColumnarValue::Scalar(start_scalar), ColumnarValue::Scalar(end_scalar)) => { - let struct_array = StructArray::try_new( - fields, - vec![start_scalar.to_array()?, end_scalar.to_array()?], - None, - )?; - Ok(ColumnarValue::Scalar(ScalarValue::Struct(Arc::new( - struct_array, - )))) - } - } - } -} - -make_udf_function!(WindowFunctionUdf, WINDOW_FUNCTION, window); diff --git a/src/sql/planning_runtime.rs b/src/sql/planning_runtime.rs deleted file mode 100644 index dc4749ad..00000000 --- a/src/sql/planning_runtime.rs +++ /dev/null @@ -1,35 +0,0 @@ -// 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. - -//! Runtime-installed SQL planning defaults (from `GlobalConfig` / `conf/config.yaml`). - -use std::sync::OnceLock; - -use crate::config::streaming_job::ResolvedStreamingJobConfig; -use crate::sql::common::constants::sql_planning_default; -use crate::sql::types::SqlConfig; - -static SQL_PLANNING: OnceLock = OnceLock::new(); - -/// Installs [`SqlConfig`] derived from resolved streaming job YAML (KeyBy parallelism, etc.). -/// Safe to call once at bootstrap; later calls are ignored if already set. -pub fn install_sql_planning_from_streaming_job(job: &ResolvedStreamingJobConfig) { - let cfg = SqlConfig { - default_parallelism: sql_planning_default::DEFAULT_PARALLELISM, - key_by_parallelism: job.key_by_parallelism as usize, - }; - let _ = SQL_PLANNING.set(cfg).ok(); -} - -pub(crate) fn sql_planning_snapshot() -> SqlConfig { - SQL_PLANNING.get().cloned().unwrap_or_default() -} diff --git a/src/sql/schema/catalog.rs b/src/sql/schema/catalog.rs deleted file mode 100644 index 479df682..00000000 --- a/src/sql/schema/catalog.rs +++ /dev/null @@ -1,609 +0,0 @@ -// 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. - -//! External connector catalog: [`ExternalTable`] as [`SourceTable`] | [`SinkTable`] | [`LookupTable`]. - -use std::collections::BTreeMap; -use std::sync::Arc; -use std::time::Duration; - -use datafusion::arrow::datatypes::{DataType, Field, FieldRef, Schema}; -use datafusion::common::{Column, Result, plan_err}; -use datafusion::error::DataFusionError; -use datafusion::logical_expr::Expr; -use protocol::function_stream_graph::ConnectorOp; - -use super::column_descriptor::ColumnDescriptor; -use super::data_encoding_format::DataEncodingFormat; -use super::table::SqlSource; -use super::temporal_pipeline_config::TemporalPipelineConfig; -use crate::multifield_partial_ord; -use crate::sql::common::constants::sql_field; -use crate::sql::common::{Format, FsSchema}; -use crate::sql::connector::config::ConnectorConfig; -use crate::sql::types::ProcessingMode; - -#[derive(Debug, Clone)] -pub struct EngineDescriptor { - pub engine_type: String, - pub raw_payload: String, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub enum SyncMode { - AppendOnly, - Incremental, -} - -#[derive(Debug, Clone)] -pub struct TableExecutionUnit { - pub label: String, - pub engine_meta: EngineDescriptor, - pub sync_mode: SyncMode, - pub temporal_offset: TemporalPipelineConfig, -} - -/// The only legal shape an external-connector catalog row can take. -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub enum ExternalTable { - Source(SourceTable), - Sink(SinkTable), - Lookup(LookupTable), -} - -impl ExternalTable { - #[inline] - pub fn name(&self) -> &str { - match self { - ExternalTable::Source(t) => t.table_identifier.as_str(), - ExternalTable::Sink(t) => t.table_identifier.as_str(), - ExternalTable::Lookup(t) => t.table_identifier.as_str(), - } - } - - #[inline] - pub fn adapter_type(&self) -> &str { - match self { - ExternalTable::Source(t) => t.adapter_type.as_str(), - ExternalTable::Sink(t) => t.adapter_type.as_str(), - ExternalTable::Lookup(t) => t.adapter_type.as_str(), - } - } - - #[inline] - pub fn description(&self) -> &str { - match self { - ExternalTable::Source(t) => t.description.as_str(), - ExternalTable::Sink(t) => t.description.as_str(), - ExternalTable::Lookup(t) => t.description.as_str(), - } - } - - #[inline] - pub fn schema_specs(&self) -> &[ColumnDescriptor] { - match self { - ExternalTable::Source(t) => &t.schema_specs, - ExternalTable::Sink(t) => &t.schema_specs, - ExternalTable::Lookup(t) => &t.schema_specs, - } - } - - #[inline] - pub fn connector_config(&self) -> &ConnectorConfig { - match self { - ExternalTable::Source(t) => &t.connector_config, - ExternalTable::Sink(t) => &t.connector_config, - ExternalTable::Lookup(t) => &t.connector_config, - } - } - - #[inline] - pub fn key_constraints(&self) -> &[String] { - match self { - ExternalTable::Source(t) => &t.key_constraints, - ExternalTable::Sink(t) => &t.key_constraints, - ExternalTable::Lookup(t) => &t.key_constraints, - } - } - - #[inline] - pub fn connection_format(&self) -> Option<&Format> { - match self { - ExternalTable::Source(t) => t.connection_format.as_ref(), - ExternalTable::Sink(t) => t.connection_format.as_ref(), - ExternalTable::Lookup(t) => t.connection_format.as_ref(), - } - } - - #[inline] - pub fn catalog_with_options(&self) -> &BTreeMap { - match self { - ExternalTable::Source(t) => &t.catalog_with_options, - ExternalTable::Sink(t) => &t.catalog_with_options, - ExternalTable::Lookup(t) => &t.catalog_with_options, - } - } - - pub fn produce_physical_schema(&self) -> Schema { - Schema::new( - self.schema_specs() - .iter() - .filter(|c| !c.is_computed()) - .map(|c| c.arrow_field().clone()) - .collect::>(), - ) - } - - pub fn connector_op(&self) -> ConnectorOp { - let physical = self.produce_physical_schema(); - let fields: Vec = physical - .fields() - .iter() - .map(|f| f.as_ref().clone()) - .collect(); - let fs_schema = FsSchema::from_fields(fields); - - ConnectorOp { - connector: self.adapter_type().to_string(), - fs_schema: Some(fs_schema.into()), - name: self.name().to_string(), - description: self.description().to_string(), - config: Some(self.connector_config().to_proto_config()), - } - } - - #[inline] - pub fn is_updating(&self) -> bool { - match self { - ExternalTable::Source(t) => t.is_updating(), - ExternalTable::Sink(t) => t - .connection_format - .as_ref() - .is_some_and(|f| f.is_updating()), - ExternalTable::Lookup(_) => false, - } - } - - /// Variant-agnostic view of "persisted Arrow fields post-planning". - /// Only Source / Lookup track inferred schema — Sinks derive theirs from the upstream plan. - pub fn effective_fields(&self) -> Vec { - match self { - ExternalTable::Source(t) => t.effective_fields(), - ExternalTable::Sink(t) => t.effective_fields(), - ExternalTable::Lookup(t) => t.effective_fields(), - } - } - - #[inline] - pub fn as_source(&self) -> Option<&SourceTable> { - match self { - ExternalTable::Source(t) => Some(t), - _ => None, - } - } - - #[inline] - pub fn as_sink(&self) -> Option<&SinkTable> { - match self { - ExternalTable::Sink(t) => Some(t), - _ => None, - } - } - - #[inline] - pub fn as_lookup(&self) -> Option<&LookupTable> { - match self { - ExternalTable::Lookup(t) => Some(t), - _ => None, - } - } -} - -/// Ingress external connector (`CREATE TABLE ... WITH (type='source', ...)`). -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct SourceTable { - pub table_identifier: String, - pub adapter_type: String, - pub schema_specs: Vec, - pub connector_config: ConnectorConfig, - pub temporal_config: TemporalPipelineConfig, - pub key_constraints: Vec, - pub payload_format: Option, - pub connection_format: Option, - pub description: String, - pub catalog_with_options: BTreeMap, - - // Planner / catalog; not in SQL text. - pub registry_id: Option, - pub inferred_fields: Option>, -} - -multifield_partial_ord!( - SourceTable, - registry_id, - adapter_type, - table_identifier, - description, - key_constraints, - connection_format, - catalog_with_options -); - -impl SourceTable { - #[inline] - pub fn name(&self) -> &str { - self.table_identifier.as_str() - } - - #[inline] - pub fn connector(&self) -> &str { - self.adapter_type.as_str() - } - - pub fn event_time_field(&self) -> Option<&str> { - self.temporal_config.event_column.as_deref() - } - - pub fn watermark_field(&self) -> Option<&str> { - self.temporal_config.watermark_strategy_column.as_deref() - } - - /// Watermark column safe to persist to the stream catalog. Omits the - /// generated `__watermark` column — that is only resolvable at compile - /// time, the catalog round-trip cannot reconstruct it. - pub fn stream_catalog_watermark_field(&self) -> Option { - self.temporal_config - .watermark_strategy_column - .as_deref() - .filter(|w| *w != sql_field::COMPUTED_WATERMARK) - .map(str::to_string) - } - - #[inline] - pub fn catalog_with_options(&self) -> &BTreeMap { - &self.catalog_with_options - } - - pub fn idle_time(&self) -> Option { - self.temporal_config.liveness_timeout - } - - pub fn produce_physical_schema(&self) -> Schema { - Schema::new( - self.schema_specs - .iter() - .filter(|c| !c.is_computed()) - .map(|c| c.arrow_field().clone()) - .collect::>(), - ) - } - - #[inline] - pub fn physical_schema(&self) -> Schema { - self.produce_physical_schema() - } - - pub fn effective_fields(&self) -> Vec { - self.inferred_fields.clone().unwrap_or_else(|| { - self.schema_specs - .iter() - .map(|c| Arc::new(c.arrow_field().clone())) - .collect() - }) - } - - pub fn convert_to_execution_unit(&self) -> Result { - if self.is_cdc_enabled() && self.schema_specs.iter().any(|c| c.is_computed()) { - return plan_err!("CDC cannot be mixed with computed columns natively"); - } - - let mode = if self.is_cdc_enabled() { - SyncMode::Incremental - } else { - SyncMode::AppendOnly - }; - - Ok(TableExecutionUnit { - label: self.table_identifier.clone(), - engine_meta: EngineDescriptor { - engine_type: self.adapter_type.clone(), - raw_payload: String::new(), - }, - sync_mode: mode, - temporal_offset: self.temporal_config.clone(), - }) - } - - #[inline] - pub fn to_execution_unit(&self) -> Result { - self.convert_to_execution_unit() - } - - fn is_cdc_enabled(&self) -> bool { - self.payload_format - .as_ref() - .is_some_and(|f| f.supports_delta_updates()) - } - - pub fn has_virtual_fields(&self) -> bool { - self.schema_specs.iter().any(|c| c.is_computed()) - } - - pub fn is_updating(&self) -> bool { - self.connection_format - .as_ref() - .is_some_and(|f| f.is_updating()) - || self.payload_format == Some(DataEncodingFormat::DebeziumJson) - } - - pub fn connector_op(&self) -> ConnectorOp { - let physical = self.produce_physical_schema(); - let fields: Vec = physical - .fields() - .iter() - .map(|f| f.as_ref().clone()) - .collect(); - let fs_schema = FsSchema::from_fields(fields); - - ConnectorOp { - connector: self.adapter_type.clone(), - fs_schema: Some(fs_schema.into()), - name: self.table_identifier.clone(), - description: self.description.clone(), - config: Some(self.connector_config.to_proto_config()), - } - } - - pub fn processing_mode(&self) -> ProcessingMode { - if self.is_updating() { - ProcessingMode::Update - } else { - ProcessingMode::Append - } - } - - pub fn timestamp_override(&self) -> Result> { - if let Some(field_name) = self.temporal_config.event_column.clone() { - if self.is_updating() { - return plan_err!("can't use event_time_field with update mode"); - } - let _field = self.get_time_column(&field_name)?; - Ok(Some(Expr::Column(Column::from_name(field_name.as_str())))) - } else { - Ok(None) - } - } - - fn get_time_column(&self, field_name: &str) -> Result<&ColumnDescriptor> { - self.schema_specs - .iter() - .find(|c| { - c.arrow_field().name() == field_name - && matches!(c.arrow_field().data_type(), DataType::Timestamp(..)) - }) - .ok_or_else(|| { - DataFusionError::Plan(format!("field {field_name} not found or not a timestamp")) - }) - } - - pub fn watermark_column(&self) -> Result> { - if let Some(field_name) = self.temporal_config.watermark_strategy_column.clone() { - let _field = self.get_time_column(&field_name)?; - Ok(Some(Expr::Column(Column::from_name(field_name.as_str())))) - } else { - Ok(None) - } - } - - pub fn as_sql_source(&self) -> Result { - if self.is_updating() && self.has_virtual_fields() { - return plan_err!("can't read from a source with virtual fields and update mode."); - } - - let timestamp_override = self.timestamp_override()?; - let watermark_column = self.watermark_column()?; - - let source = SqlSource { - id: self.registry_id, - struct_def: self - .schema_specs - .iter() - .filter(|c| !c.is_computed()) - .map(|c| Arc::new(c.arrow_field().clone())) - .collect(), - config: self.connector_op(), - processing_mode: self.processing_mode(), - idle_time: self.temporal_config.liveness_timeout, - }; - - Ok(SourceOperator { - name: self.table_identifier.clone(), - source, - timestamp_override, - watermark_column, - }) - } -} - -/// Egress external connector, or the sink of `CREATE STREAMING TABLE ... AS SELECT`. -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct SinkTable { - pub table_identifier: String, - pub adapter_type: String, - pub schema_specs: Vec, - pub connector_config: ConnectorConfig, - pub partition_exprs: Arc>>, - pub key_constraints: Vec, - pub connection_format: Option, - pub description: String, - pub catalog_with_options: BTreeMap, -} - -multifield_partial_ord!( - SinkTable, - adapter_type, - table_identifier, - description, - key_constraints, - connection_format, - catalog_with_options -); - -impl SinkTable { - #[inline] - pub fn name(&self) -> &str { - self.table_identifier.as_str() - } - - #[inline] - pub fn connector(&self) -> &str { - self.adapter_type.as_str() - } - - #[inline] - pub fn catalog_with_options(&self) -> &BTreeMap { - &self.catalog_with_options - } - - pub fn produce_physical_schema(&self) -> Schema { - Schema::new( - self.schema_specs - .iter() - .filter(|c| !c.is_computed()) - .map(|c| c.arrow_field().clone()) - .collect::>(), - ) - } - - pub fn effective_fields(&self) -> Vec { - self.schema_specs - .iter() - .map(|c| Arc::new(c.arrow_field().clone())) - .collect() - } - - pub fn is_updating(&self) -> bool { - self.connection_format - .as_ref() - .is_some_and(|f| f.is_updating()) - } - - pub fn connector_op(&self) -> ConnectorOp { - let physical = self.produce_physical_schema(); - let fields: Vec = physical - .fields() - .iter() - .map(|f| f.as_ref().clone()) - .collect(); - let fs_schema = FsSchema::from_fields(fields); - - ConnectorOp { - connector: self.adapter_type.clone(), - fs_schema: Some(fs_schema.into()), - name: self.table_identifier.clone(), - description: self.description.clone(), - config: Some(self.connector_config.to_proto_config()), - } - } -} - -/// Lookup-join only; not a scan source. -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct LookupTable { - pub table_identifier: String, - pub adapter_type: String, - pub schema_specs: Vec, - pub connector_config: ConnectorConfig, - pub key_constraints: Vec, - pub lookup_cache_max_bytes: Option, - pub lookup_cache_ttl: Option, - pub connection_format: Option, - pub description: String, - pub catalog_with_options: BTreeMap, - - pub registry_id: Option, - pub inferred_fields: Option>, -} - -multifield_partial_ord!( - LookupTable, - registry_id, - adapter_type, - table_identifier, - description, - key_constraints, - connection_format, - catalog_with_options -); - -impl LookupTable { - #[inline] - pub fn name(&self) -> &str { - self.table_identifier.as_str() - } - - #[inline] - pub fn connector(&self) -> &str { - self.adapter_type.as_str() - } - - #[inline] - pub fn catalog_with_options(&self) -> &BTreeMap { - &self.catalog_with_options - } - - pub fn produce_physical_schema(&self) -> Schema { - Schema::new( - self.schema_specs - .iter() - .filter(|c| !c.is_computed()) - .map(|c| c.arrow_field().clone()) - .collect::>(), - ) - } - - pub fn effective_fields(&self) -> Vec { - self.inferred_fields.clone().unwrap_or_else(|| { - self.schema_specs - .iter() - .map(|c| Arc::new(c.arrow_field().clone())) - .collect() - }) - } - - pub fn connector_op(&self) -> ConnectorOp { - let physical = self.produce_physical_schema(); - let fields: Vec = physical - .fields() - .iter() - .map(|f| f.as_ref().clone()) - .collect(); - let fs_schema = FsSchema::from_fields(fields); - - ConnectorOp { - connector: self.adapter_type.clone(), - fs_schema: Some(fs_schema.into()), - name: self.table_identifier.clone(), - description: self.description.clone(), - config: Some(self.connector_config.to_proto_config()), - } - } -} - -/// [`SourceTable`] as an ingestion logical node input. -#[derive(Debug, Clone)] -pub struct SourceOperator { - pub name: String, - pub source: SqlSource, - pub timestamp_override: Option, - pub watermark_column: Option, -} diff --git a/src/sql/schema/column_descriptor.rs b/src/sql/schema/column_descriptor.rs deleted file mode 100644 index 4228816f..00000000 --- a/src/sql/schema/column_descriptor.rs +++ /dev/null @@ -1,144 +0,0 @@ -// 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 datafusion::arrow::datatypes::{DataType, Field, TimeUnit}; -use datafusion::logical_expr::Expr; - -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub enum ColumnDescriptor { - Physical(Field), - SystemMeta { field: Field, meta_key: String }, - Computed { field: Field, logic: Box }, -} - -impl ColumnDescriptor { - #[inline] - pub fn new_physical(field: Field) -> Self { - Self::Physical(field) - } - - #[inline] - pub fn new_system_meta(field: Field, meta_key: impl Into) -> Self { - Self::SystemMeta { - field, - meta_key: meta_key.into(), - } - } - - #[inline] - pub fn new_computed(field: Field, logic: Expr) -> Self { - Self::Computed { - field, - logic: Box::new(logic), - } - } - - #[inline] - pub fn arrow_field(&self) -> &Field { - match self { - Self::Physical(f) => f, - Self::SystemMeta { field: f, .. } => f, - Self::Computed { field: f, .. } => f, - } - } - - #[inline] - pub fn into_arrow_field(self) -> Field { - match self { - Self::Physical(f) => f, - Self::SystemMeta { field: f, .. } => f, - Self::Computed { field: f, .. } => f, - } - } - - #[inline] - pub fn is_computed(&self) -> bool { - matches!(self, Self::Computed { .. }) - } - - #[inline] - pub fn is_physical(&self) -> bool { - matches!(self, Self::Physical(_)) - } - - #[inline] - pub fn system_meta_key(&self) -> Option<&str> { - if let Self::SystemMeta { meta_key, .. } = self { - Some(meta_key.as_str()) - } else { - None - } - } - - #[inline] - pub fn computation_logic(&self) -> Option<&Expr> { - if let Self::Computed { logic, .. } = self { - Some(logic) - } else { - None - } - } - - #[inline] - pub fn data_type(&self) -> &DataType { - self.arrow_field().data_type() - } - - pub fn set_nullable(&mut self, nullable: bool) { - let f = match self { - Self::Physical(f) => f, - Self::SystemMeta { field, .. } => field, - Self::Computed { field, .. } => field, - }; - *f = Field::new(f.name(), f.data_type().clone(), nullable) - .with_metadata(f.metadata().clone()); - } - - pub fn force_precision(&mut self, unit: TimeUnit) { - match self { - Self::Physical(f) => { - if let DataType::Timestamp(_, tz) = f.data_type() { - *f = Field::new( - f.name(), - DataType::Timestamp(unit, tz.clone()), - f.is_nullable(), - ); - } - } - Self::SystemMeta { field, .. } => { - if let DataType::Timestamp(_, tz) = field.data_type() { - *field = Field::new( - field.name(), - DataType::Timestamp(unit, tz.clone()), - field.is_nullable(), - ); - } - } - Self::Computed { field, .. } => { - if let DataType::Timestamp(_, tz) = field.data_type() { - *field = Field::new( - field.name(), - DataType::Timestamp(unit, tz.clone()), - field.is_nullable(), - ); - } - } - } - } -} - -impl From for ColumnDescriptor { - #[inline] - fn from(field: Field) -> Self { - Self::Physical(field) - } -} diff --git a/src/sql/schema/connection_type.rs b/src/sql/schema/connection_type.rs deleted file mode 100644 index 06a3df92..00000000 --- a/src/sql/schema/connection_type.rs +++ /dev/null @@ -1,31 +0,0 @@ -// 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::fmt; - -/// Describes the role of a connection in the streaming pipeline. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub enum ConnectionType { - Source, - Sink, - Lookup, -} - -impl fmt::Display for ConnectionType { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - ConnectionType::Source => write!(f, "source"), - ConnectionType::Sink => write!(f, "sink"), - ConnectionType::Lookup => write!(f, "lookup"), - } - } -} diff --git a/src/sql/schema/data_encoding_format.rs b/src/sql/schema/data_encoding_format.rs deleted file mode 100644 index 0b6f5e1d..00000000 --- a/src/sql/schema/data_encoding_format.rs +++ /dev/null @@ -1,89 +0,0 @@ -// 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 datafusion::arrow::datatypes::{DataType, Field}; -use datafusion::common::{Result, plan_err}; - -use super::column_descriptor::ColumnDescriptor; -use crate::sql::common::Format; -use crate::sql::common::constants::cdc; - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] -pub enum DataEncodingFormat { - #[default] - Raw, - StandardJson, - DebeziumJson, - Avro, - Parquet, - Csv, - JsonL, - Orc, - Protobuf, -} - -impl DataEncodingFormat { - pub fn from_format(format: Option<&Format>) -> Self { - match format { - Some(Format::Json(j)) if j.debezium => Self::DebeziumJson, - Some(Format::Json(_)) => Self::StandardJson, - Some(Format::Avro(_)) => Self::Avro, - Some(Format::Parquet(_)) => Self::Parquet, - Some(Format::Csv(_)) => Self::Csv, - Some(Format::Protobuf(_)) => Self::Protobuf, - Some(Format::RawString(_)) | Some(Format::RawBytes(_)) | None => Self::Raw, - Some(_) => Self::Raw, - } - } - - pub fn is_cdc_format(&self) -> bool { - matches!(self, Self::DebeziumJson) - } - - #[inline] - pub fn supports_delta_updates(&self) -> bool { - self.is_cdc_format() - } - - pub fn apply_envelope( - &self, - logical_columns: Vec, - ) -> Result> { - if !self.is_cdc_format() { - return Ok(logical_columns); - } - - if logical_columns.is_empty() { - return Ok(logical_columns); - } - - if logical_columns.iter().any(|c| c.is_computed()) { - return plan_err!( - "Computed/Virtual columns are not supported directly inside a CDC source table; \ - define computed columns in a downstream VIEW or AS SELECT streaming query" - ); - } - - let inner_fields: Vec = logical_columns - .into_iter() - .map(|c| c.into_arrow_field()) - .collect(); - - let row_struct_type = DataType::Struct(inner_fields.into()); - - Ok(vec![ - ColumnDescriptor::new_physical(Field::new(cdc::BEFORE, row_struct_type.clone(), true)), - ColumnDescriptor::new_physical(Field::new(cdc::AFTER, row_struct_type, true)), - ColumnDescriptor::new_physical(Field::new(cdc::OP, DataType::Utf8, true)), - ]) - } -} diff --git a/src/sql/schema/introspection/ddl_formatter.rs b/src/sql/schema/introspection/ddl_formatter.rs deleted file mode 100644 index f4ce36e6..00000000 --- a/src/sql/schema/introspection/ddl_formatter.rs +++ /dev/null @@ -1,156 +0,0 @@ -// 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 std::fmt::{self, Write}; - -use datafusion::arrow::datatypes::{DataType, Schema, TimeUnit}; - -use crate::sql::common::constants::sql_field; - -pub struct DdlBuilder<'a> { - table_name: &'a str, - schema: &'a Schema, - watermark_column: Option<&'a str>, - primary_keys: &'a [String], - options: BTreeMap, -} - -impl<'a> DdlBuilder<'a> { - pub fn new(table_name: &'a str, schema: &'a Schema) -> Self { - Self { - table_name, - schema, - watermark_column: None, - primary_keys: &[], - options: BTreeMap::new(), - } - } - - pub fn with_watermark(mut self, watermark: Option<&'a str>) -> Self { - self.watermark_column = watermark; - self - } - - pub fn with_primary_keys(mut self, keys: &'a [String]) -> Self { - self.primary_keys = keys; - self - } - - pub fn with_options( - mut self, - opts: &BTreeMap, - role: &str, - connector: &str, - ) -> Self { - self.options = opts.clone(); - self.options - .entry("type".to_string()) - .or_insert_with(|| role.to_string()); - self.options - .entry("connector".to_string()) - .or_insert_with(|| connector.to_string()); - self - } -} - -impl<'a> fmt::Display for DdlBuilder<'a> { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - writeln!(f, "CREATE TABLE {} (", self.table_name)?; - - let mut rows: Vec = Vec::new(); - for field in self.schema.fields() { - let null_constraint = if field.is_nullable() { "" } else { " NOT NULL" }; - rows.push(format!( - " {} {}{}", - field.name(), - format_data_type(field.data_type()), - null_constraint - )); - } - - if let Some(wm) = self.watermark_column - && wm != sql_field::COMPUTED_WATERMARK - { - rows.push(format!(" WATERMARK FOR {wm}")); - } - - if !self.primary_keys.is_empty() { - rows.push(format!(" PRIMARY KEY ({})", self.primary_keys.join(", "))); - } - - writeln!(f, "{}", rows.join(",\n"))?; - write!(f, ")")?; - - if !self.options.is_empty() { - writeln!(f)?; - writeln!(f, "WITH (")?; - let mut opt_lines: Vec = Vec::with_capacity(self.options.len()); - for (k, v) in &self.options { - let k_esc = k.replace('\'', "''"); - let v_esc = v.replace('\'', "''"); - opt_lines.push(format!(" '{k_esc}' = '{v_esc}'")); - } - write!(f, "{}\n);", opt_lines.join(",\n"))?; - } else { - write!(f, ";")?; - } - - Ok(()) - } -} - -pub fn format_data_type(dt: &DataType) -> String { - match dt { - DataType::Null => "NULL".to_string(), - DataType::Boolean => "BOOLEAN".to_string(), - DataType::Int8 => "TINYINT".to_string(), - DataType::Int16 => "SMALLINT".to_string(), - DataType::Int32 => "INT".to_string(), - DataType::Int64 => "BIGINT".to_string(), - DataType::UInt8 => "TINYINT UNSIGNED".to_string(), - DataType::UInt16 => "SMALLINT UNSIGNED".to_string(), - DataType::UInt32 => "INT UNSIGNED".to_string(), - DataType::UInt64 => "BIGINT UNSIGNED".to_string(), - DataType::Float16 => "FLOAT".to_string(), - DataType::Float32 => "REAL".to_string(), - DataType::Float64 => "DOUBLE".to_string(), - DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => "VARCHAR".to_string(), - DataType::Binary | DataType::LargeBinary => "VARBINARY".to_string(), - DataType::Date32 | DataType::Date64 => "DATE".to_string(), - DataType::Timestamp(unit, tz) => match (unit, tz) { - (TimeUnit::Second, None) => "TIMESTAMP(0)".to_string(), - (TimeUnit::Millisecond, None) => "TIMESTAMP(3)".to_string(), - (TimeUnit::Microsecond, None) => "TIMESTAMP(6)".to_string(), - (TimeUnit::Nanosecond, None) => "TIMESTAMP(9)".to_string(), - (_, Some(_)) => "TIMESTAMP WITH TIME ZONE".to_string(), - }, - DataType::Decimal128(p, s) | DataType::Decimal256(p, s) => format!("DECIMAL({p}, {s})"), - _ => dt.to_string(), - } -} - -pub fn schema_columns_one_line(schema: &Schema) -> String { - let mut buf = String::new(); - for (idx, field) in schema.fields().iter().enumerate() { - if idx > 0 { - buf.push_str(", "); - } - let _ = write!( - buf, - "{}:{}", - field.name(), - format_data_type(field.data_type()) - ); - } - buf -} diff --git a/src/sql/schema/introspection/mod.rs b/src/sql/schema/introspection/mod.rs deleted file mode 100644 index 1ba9c816..00000000 --- a/src/sql/schema/introspection/mod.rs +++ /dev/null @@ -1,21 +0,0 @@ -// 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. - -pub mod ddl_formatter; -pub mod show_formatter; -pub mod stream_formatter; - -#[allow(unused_imports)] -pub use ddl_formatter::{DdlBuilder, format_data_type, schema_columns_one_line}; -pub use show_formatter::{catalog_table_row_detail, show_create_catalog_table}; -#[allow(unused_imports)] -pub use stream_formatter::{show_create_stream_table, stream_table_row_detail}; diff --git a/src/sql/schema/introspection/show_formatter.rs b/src/sql/schema/introspection/show_formatter.rs deleted file mode 100644 index 28a81ae9..00000000 --- a/src/sql/schema/introspection/show_formatter.rs +++ /dev/null @@ -1,100 +0,0 @@ -// 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::sql::common::constants::connection_table_role; -use crate::sql::schema::catalog::ExternalTable; -use crate::sql::schema::table::CatalogEntity; - -use super::ddl_formatter::DdlBuilder; - -impl ExternalTable { - pub fn to_ddl_string(&self) -> String { - match self { - ExternalTable::Source(source) => { - let schema = source.produce_physical_schema(); - DdlBuilder::new(&source.table_identifier, &schema) - .with_watermark(source.temporal_config.watermark_strategy_column.as_deref()) - .with_primary_keys(&source.key_constraints) - .with_options( - &source.catalog_with_options, - connection_table_role::SOURCE, - &source.adapter_type, - ) - .to_string() - } - ExternalTable::Sink(sink) => { - let schema = sink.produce_physical_schema(); - DdlBuilder::new(&sink.table_identifier, &schema) - .with_primary_keys(&sink.key_constraints) - .with_options( - &sink.catalog_with_options, - connection_table_role::SINK, - &sink.adapter_type, - ) - .to_string() - } - ExternalTable::Lookup(lookup) => { - let schema = lookup.produce_physical_schema(); - DdlBuilder::new(&lookup.table_identifier, &schema) - .with_primary_keys(&lookup.key_constraints) - .with_options( - &lookup.catalog_with_options, - connection_table_role::LOOKUP, - &lookup.adapter_type, - ) - .to_string() - } - } - } - - pub fn to_row_detail(&self) -> String { - match self { - ExternalTable::Source(s) => format!( - "{{ kind: 'source', connector: '{}', watermark: '{}', options_count: {} }}", - s.adapter_type, - s.temporal_config - .watermark_strategy_column - .as_deref() - .unwrap_or("none"), - s.catalog_with_options.len() - ), - ExternalTable::Sink(s) => format!( - "{{ kind: 'sink', connector: '{}', partitioned: {}, options_count: {} }}", - s.adapter_type, - s.partition_exprs.as_ref().is_some(), - s.catalog_with_options.len() - ), - ExternalTable::Lookup(s) => format!( - "{{ kind: 'lookup', connector: '{}', cache_ttl_secs: {}, options_count: {} }}", - s.adapter_type, - s.lookup_cache_ttl.map(|d| d.as_secs()).unwrap_or(0), - s.catalog_with_options.len() - ), - } - } -} - -pub fn show_create_catalog_table(table: &CatalogEntity) -> String { - match table { - CatalogEntity::ExternalConnector(ext) => ext.to_ddl_string(), - CatalogEntity::ComputedTable { name, .. } => { - format!("-- Logical query view\nCREATE VIEW {name} AS SELECT ...;") - } - } -} - -pub fn catalog_table_row_detail(table: &CatalogEntity) -> String { - match table { - CatalogEntity::ExternalConnector(ext) => ext.to_row_detail(), - CatalogEntity::ComputedTable { .. } => "{ kind: 'logical_view' }".to_string(), - } -} diff --git a/src/sql/schema/introspection/stream_formatter.rs b/src/sql/schema/introspection/stream_formatter.rs deleted file mode 100644 index ebb02330..00000000 --- a/src/sql/schema/introspection/stream_formatter.rs +++ /dev/null @@ -1,120 +0,0 @@ -// 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::sync::Arc; - -use datafusion::arrow::datatypes::Schema; - -use crate::sql::common::constants::connection_table_role; -use crate::sql::logical_node::logical::LogicalProgram; -use crate::sql::schema::schema_provider::StreamTable; - -use super::ddl_formatter::DdlBuilder; - -impl StreamTable { - pub fn to_ddl_string(&self) -> String { - match self { - StreamTable::Source { - name, - connector, - schema, - event_time_field: _, - watermark_field, - with_options, - } => DdlBuilder::new(name, schema) - .with_watermark(watermark_field.as_deref()) - .with_options(with_options, connection_table_role::SOURCE, connector) - .to_string(), - StreamTable::Sink { name, program } => { - let schema: Arc = program - .egress_arrow_schema() - .unwrap_or_else(|| Arc::new(Schema::empty())); - - let mut ddl = format!("CREATE STREAMING TABLE {name} AS SELECT ...\n\n"); - ddl.push_str("/* === SINK SCHEMA === */\n"); - let schema_ddl = DdlBuilder::new(name, &schema).to_string(); - ddl.push_str(&schema_ddl); - ddl.push_str("\n\n/* === STREAMING TOPOLOGY === */\n"); - ddl.push_str(&format_pipeline(program)); - ddl - } - } - } - - pub fn to_row_detail(&self) -> String { - match self { - StreamTable::Source { - connector, - event_time_field, - watermark_field, - with_options, - .. - } => format!( - "{{ kind: 'stream_source', connector: '{}', event_time: '{}', watermark: '{}', options_count: {} }}", - connector, - event_time_field.as_deref().unwrap_or("none"), - watermark_field.as_deref().unwrap_or("none"), - with_options.len() - ), - StreamTable::Sink { program, .. } => format!( - "{{ kind: 'streaming_sink', tasks: {}, nodes: {} }}", - program.task_count(), - program.graph.node_count() - ), - } - } -} - -pub fn show_create_stream_table(table: &StreamTable) -> String { - table.to_ddl_string() -} - -pub fn stream_table_row_detail(table: &StreamTable) -> String { - table.to_row_detail() -} - -fn format_pipeline(program: &LogicalProgram) -> String { - let mut lines: Vec = Vec::new(); - lines.push(format!("Pipeline Hash : {}", program.get_hash())); - lines.push(format!("Total Tasks : {}", program.task_count())); - lines.push(format!("Node Count : {}", program.graph.node_count())); - lines.push(String::from("Operator Chains:")); - - for nw in program.graph.node_weights() { - let chain = nw - .operator_chain - .operators - .iter() - .map(|op| format!("{}[{}]", op.operator_name, op.operator_id)) - .collect::>() - .join(" -> "); - - lines.push(format!( - " Node {:<3} | Parallelism {:<3} | {}", - nw.node_id, nw.parallelism, chain - )); - } - - let dot = program.dot(); - const MAX_DOT: usize = 5_000; - if dot.len() > MAX_DOT { - lines.push(format!( - "\nGraphviz DOT (truncated, {} bytes omitted):\n{}...", - dot.len() - MAX_DOT, - &dot[..MAX_DOT] - )); - } else { - lines.push(format!("\nGraphviz DOT:\n{dot}")); - } - - lines.join("\n") -} diff --git a/src/sql/schema/mod.rs b/src/sql/schema/mod.rs deleted file mode 100644 index e11e4808..00000000 --- a/src/sql/schema/mod.rs +++ /dev/null @@ -1,30 +0,0 @@ -// 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. - -pub mod catalog; -pub mod column_descriptor; -pub mod connection_type; -pub mod data_encoding_format; -pub mod introspection; -pub mod schema_provider; -pub mod table; -pub mod table_role; -pub mod temporal_pipeline_config; -pub mod utils; - -pub use catalog::{ExternalTable, LookupTable, SinkTable, SourceTable}; -pub use column_descriptor::ColumnDescriptor; -pub use introspection::{ - catalog_table_row_detail, schema_columns_one_line, show_create_catalog_table, -}; -pub use schema_provider::{ObjectName, StreamPlanningContext, StreamSchemaProvider, StreamTable}; -pub use table::CatalogEntity; diff --git a/src/sql/schema/schema_provider.rs b/src/sql/schema/schema_provider.rs deleted file mode 100644 index 15cd58ee..00000000 --- a/src/sql/schema/schema_provider.rs +++ /dev/null @@ -1,469 +0,0 @@ -// 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, HashMap, HashSet}; -use std::sync::Arc; - -use datafusion::arrow::datatypes::{self as datatypes, DataType, Field, Schema}; -use datafusion::common::{DataFusionError, Result as DataFusionResult}; -use datafusion::datasource::{DefaultTableSource, TableProvider, TableType}; -use datafusion::execution::{FunctionRegistry, SessionStateDefaults}; -use datafusion::logical_expr::expr_rewriter::FunctionRewrite; -use datafusion::logical_expr::planner::ExprPlanner; -use datafusion::logical_expr::{AggregateUDF, Expr, ScalarUDF, TableSource, WindowUDF}; -use datafusion::optimizer::Analyzer; -use datafusion::sql::TableReference; -use datafusion::sql::planner::ContextProvider; -use thiserror::Error; -use tracing::{debug, error, info}; -use unicase::UniCase; - -use crate::sql::common::constants::{planning_placeholder_udf, window_fn}; -use crate::sql::logical_node::logical::{DylibUdfConfig, LogicalProgram}; -use crate::sql::schema::table::CatalogEntity; -use crate::sql::schema::utils::window_arrow_struct; -use crate::sql::types::{PlanningOptions, PlanningPlaceholderUdf, SqlConfig}; - -pub type ObjectName = UniCase; - -#[inline] -fn object_name(s: impl Into) -> ObjectName { - UniCase::new(s.into()) -} - -#[derive(Error, Debug)] -pub enum PlanningError { - #[error("Catalog table not found: {0}")] - TableNotFound(String), - #[error("Planning init failed: {0}")] - InitError(String), - #[error("Engine error: {0}")] - Engine(#[from] DataFusionError), -} - -impl From for DataFusionError { - fn from(err: PlanningError) -> Self { - match err { - PlanningError::Engine(inner) => inner, - other => DataFusionError::Plan(other.to_string()), - } - } -} - -#[derive(Clone, Debug)] -pub enum StreamTable { - Source { - name: String, - connector: String, - schema: Arc, - event_time_field: Option, - watermark_field: Option, - with_options: BTreeMap, - }, - Sink { - name: String, - program: LogicalProgram, - }, -} - -impl StreamTable { - pub fn name(&self) -> &str { - match self { - Self::Source { name, .. } | Self::Sink { name, .. } => name, - } - } - - pub fn schema(&self) -> Arc { - match self { - Self::Source { schema, .. } => Arc::clone(schema), - Self::Sink { program, .. } => program - .egress_arrow_schema() - .unwrap_or_else(|| Arc::new(Schema::empty())), - } - } -} - -#[derive(Debug, Clone)] -pub struct LogicalBatchInput { - pub table_name: String, - pub schema: Arc, -} - -#[async_trait::async_trait] -impl TableProvider for LogicalBatchInput { - fn as_any(&self) -> &dyn std::any::Any { - self - } - - fn schema(&self) -> Arc { - Arc::clone(&self.schema) - } - - fn table_type(&self) -> TableType { - TableType::Temporary - } - - async fn scan( - &self, - _state: &dyn datafusion::catalog::Session, - _projection: Option<&Vec>, - _filters: &[Expr], - _limit: Option, - ) -> DataFusionResult> { - Ok(Arc::new(crate::sql::physical::FsMemExec::new( - self.table_name.clone(), - Arc::clone(&self.schema), - ))) - } -} - -#[derive(Clone, Default)] -pub struct FunctionCatalog { - pub scalars: HashMap>, - pub aggregates: HashMap>, - pub windows: HashMap>, - pub planners: Vec>, -} - -#[derive(Clone, Default)] -pub struct TableCatalog { - pub streams: HashMap>, - pub catalogs: HashMap>, - pub source_defs: HashMap, -} - -#[derive(Clone, Default)] -pub struct StreamPlanningContext { - pub tables: TableCatalog, - pub functions: FunctionCatalog, - pub dylib_udfs: HashMap, - pub config_options: datafusion::config::ConfigOptions, - pub planning_options: PlanningOptions, - pub analyzer: Analyzer, - pub sql_config: SqlConfig, -} - -pub type StreamSchemaProvider = StreamPlanningContext; - -impl StreamPlanningContext { - pub fn builder() -> StreamPlanningContextBuilder { - StreamPlanningContextBuilder::default() - } - - #[inline] - pub fn default_parallelism(&self) -> usize { - self.sql_config.default_parallelism - } - - #[inline] - pub fn key_by_parallelism(&self) -> usize { - self.sql_config.key_by_parallelism - } - - pub fn try_new(config: SqlConfig) -> Result { - info!("Initializing StreamPlanningContext"); - let mut builder = StreamPlanningContextBuilder::default(); - builder - .with_streaming_extensions()? - .with_default_functions()?; - let mut ctx = builder.build(); - ctx.sql_config = config; - Ok(ctx) - } - - pub fn new() -> Self { - let config = crate::sql::planning_runtime::sql_planning_snapshot(); - Self::try_new(config).expect("StreamPlanningContext bootstrap") - } - - pub fn register_stream_table(&mut self, table: StreamTable) { - let key = object_name(table.name().to_string()); - debug!(table = %key, "register stream table"); - self.tables.streams.insert(key, Arc::new(table)); - } - - pub fn get_stream_table(&self, name: &str) -> Option> { - self.tables - .streams - .get(&object_name(name.to_string())) - .cloned() - } - - pub fn register_catalog_table(&mut self, table: CatalogEntity) { - let key = object_name(table.name().to_string()); - debug!(table = %key, "register catalog table"); - self.tables.catalogs.insert(key, Arc::new(table)); - } - - pub fn get_catalog_table(&self, table_name: impl AsRef) -> Option<&CatalogEntity> { - self.tables - .catalogs - .get(&object_name(table_name.as_ref().to_string())) - .map(|t| t.as_ref()) - } - - pub fn get_catalog_table_mut( - &mut self, - table_name: impl AsRef, - ) -> Option<&mut CatalogEntity> { - self.tables - .catalogs - .get_mut(&object_name(table_name.as_ref().to_string())) - .map(Arc::make_mut) - } - - pub fn add_source_table( - &mut self, - name: String, - schema: Arc, - event_time_field: Option, - watermark_field: Option, - ) { - self.register_stream_table(StreamTable::Source { - name, - connector: "stream_catalog".to_string(), - schema, - event_time_field, - watermark_field, - with_options: BTreeMap::new(), - }); - } - - pub fn add_sink_table(&mut self, name: String, program: LogicalProgram) { - self.register_stream_table(StreamTable::Sink { name, program }); - } - - pub fn insert_table(&mut self, table: StreamTable) { - self.register_stream_table(table); - } - - pub fn insert_catalog_table(&mut self, table: CatalogEntity) { - self.register_catalog_table(table); - } - - pub fn get_table(&self, table_name: impl AsRef) -> Option<&StreamTable> { - self.tables - .streams - .get(&object_name(table_name.as_ref().to_string())) - .map(|a| a.as_ref()) - } - - pub fn get_table_mut(&mut self, table_name: impl AsRef) -> Option<&mut StreamTable> { - self.tables - .streams - .get_mut(&object_name(table_name.as_ref().to_string())) - .map(Arc::make_mut) - } - - pub fn get_async_udf_options(&self, _name: &str) -> Option { - None - } - - fn create_table_source(name: String, schema: Arc) -> Arc { - let provider = LogicalBatchInput { - table_name: name, - schema, - }; - Arc::new(DefaultTableSource::new(Arc::new(provider))) - } -} - -impl ContextProvider for StreamPlanningContext { - fn get_table_source(&self, name: TableReference) -> DataFusionResult> { - let name_str = name.table(); - match self.get_stream_table(name_str) { - Some(table) => Ok(Self::create_table_source(name.to_string(), table.schema())), - None => { - error!(table = %name_str, "stream table lookup failed"); - Err(DataFusionError::Plan(format!("Table {} not found", name))) - } - } - } - - fn get_function_meta(&self, name: &str) -> Option> { - self.functions.scalars.get(name).cloned() - } - - fn get_aggregate_meta(&self, name: &str) -> Option> { - self.functions.aggregates.get(name).cloned() - } - - fn get_window_meta(&self, name: &str) -> Option> { - self.functions.windows.get(name).cloned() - } - - fn get_variable_type(&self, _variable_names: &[String]) -> Option { - None - } - - fn options(&self) -> &datafusion::config::ConfigOptions { - &self.config_options - } - - fn udf_names(&self) -> Vec { - self.functions.scalars.keys().cloned().collect() - } - - fn udaf_names(&self) -> Vec { - self.functions.aggregates.keys().cloned().collect() - } - - fn udwf_names(&self) -> Vec { - self.functions.windows.keys().cloned().collect() - } - - fn get_expr_planners(&self) -> &[Arc] { - &self.functions.planners - } -} - -impl FunctionRegistry for StreamPlanningContext { - fn udfs(&self) -> HashSet { - self.functions.scalars.keys().cloned().collect() - } - - fn udf(&self, name: &str) -> DataFusionResult> { - self.functions - .scalars - .get(name) - .cloned() - .ok_or_else(|| DataFusionError::Plan(format!("No UDF with name {name}"))) - } - - fn udaf(&self, name: &str) -> DataFusionResult> { - self.functions - .aggregates - .get(name) - .cloned() - .ok_or_else(|| DataFusionError::Plan(format!("No UDAF with name {name}"))) - } - - fn udwf(&self, name: &str) -> DataFusionResult> { - self.functions - .windows - .get(name) - .cloned() - .ok_or_else(|| DataFusionError::Plan(format!("No UDWF with name {name}"))) - } - - fn register_function_rewrite( - &mut self, - rewrite: Arc, - ) -> DataFusionResult<()> { - self.analyzer.add_function_rewrite(rewrite); - Ok(()) - } - - fn register_udf(&mut self, udf: Arc) -> DataFusionResult>> { - Ok(self.functions.scalars.insert(udf.name().to_string(), udf)) - } - - fn register_udaf( - &mut self, - udaf: Arc, - ) -> DataFusionResult>> { - Ok(self - .functions - .aggregates - .insert(udaf.name().to_string(), udaf)) - } - - fn register_udwf(&mut self, udwf: Arc) -> DataFusionResult>> { - Ok(self.functions.windows.insert(udwf.name().to_string(), udwf)) - } - - fn register_expr_planner( - &mut self, - expr_planner: Arc, - ) -> DataFusionResult<()> { - self.functions.planners.push(expr_planner); - Ok(()) - } - - fn expr_planners(&self) -> Vec> { - self.functions.planners.clone() - } -} - -#[derive(Default)] -pub struct StreamPlanningContextBuilder { - context: StreamPlanningContext, -} - -impl StreamPlanningContextBuilder { - pub fn new() -> Self { - Self::default() - } - - pub fn with_default_functions(&mut self) -> Result<&mut Self, PlanningError> { - for p in SessionStateDefaults::default_scalar_functions() { - self.context.register_udf(p)?; - } - for p in SessionStateDefaults::default_aggregate_functions() { - self.context.register_udaf(p)?; - } - for p in SessionStateDefaults::default_window_functions() { - self.context.register_udwf(p)?; - } - for p in SessionStateDefaults::default_expr_planners() { - self.context.register_expr_planner(p)?; - } - Ok(self) - } - - pub fn with_streaming_extensions(&mut self) -> Result<&mut Self, PlanningError> { - let extensions = vec![ - PlanningPlaceholderUdf::new_with_return( - window_fn::HOP, - vec![ - DataType::Interval(datatypes::IntervalUnit::MonthDayNano), - DataType::Interval(datatypes::IntervalUnit::MonthDayNano), - ], - window_arrow_struct(), - ), - PlanningPlaceholderUdf::new_with_return( - window_fn::TUMBLE, - vec![DataType::Interval(datatypes::IntervalUnit::MonthDayNano)], - window_arrow_struct(), - ), - PlanningPlaceholderUdf::new_with_return( - window_fn::SESSION, - vec![DataType::Interval(datatypes::IntervalUnit::MonthDayNano)], - window_arrow_struct(), - ), - PlanningPlaceholderUdf::new_with_return( - planning_placeholder_udf::UNNEST, - vec![DataType::List(Arc::new(Field::new( - planning_placeholder_udf::LIST_ELEMENT_FIELD, - DataType::Utf8, - true, - )))], - DataType::Utf8, - ), - PlanningPlaceholderUdf::new_with_return( - planning_placeholder_udf::ROW_TIME, - vec![], - DataType::Timestamp(datatypes::TimeUnit::Nanosecond, None), - ), - ]; - - for ext in extensions { - self.context.register_udf(ext)?; - } - - Ok(self) - } - - pub fn build(self) -> StreamPlanningContext { - self.context - } -} diff --git a/src/sql/schema/table.rs b/src/sql/schema/table.rs deleted file mode 100644 index 6c001d9c..00000000 --- a/src/sql/schema/table.rs +++ /dev/null @@ -1,163 +0,0 @@ -// 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::sql::analysis::rewrite_plan; -use crate::sql::logical_node::remote_table::RemoteTableBoundaryNode; -use crate::sql::logical_planner::optimizers::produce_optimized_plan; -use crate::sql::schema::StreamSchemaProvider; -use crate::sql::schema::catalog::ExternalTable; -use crate::sql::types::{ProcessingMode, QualifiedField}; -use datafusion::arrow::datatypes::FieldRef; -use datafusion::common::{Result, plan_err}; -use datafusion::logical_expr::{Extension, LogicalPlan}; -use datafusion::sql::sqlparser::ast::Statement; -use protocol::function_stream_graph::ConnectorOp; -use std::sync::Arc; -use std::time::Duration; - -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub enum CatalogEntity { - /// Both payload variants are boxed so the enum is not padded to the largest field. - ExternalConnector(Box), - ComputedTable { - name: String, - logical_plan: Box, - }, -} - -impl CatalogEntity { - #[inline] - pub fn external(table: ExternalTable) -> Self { - Self::ExternalConnector(Box::new(table)) - } - - pub fn try_from_statement( - statement: &Statement, - schema_provider: &StreamSchemaProvider, - ) -> Result> { - use datafusion::logical_expr::{CreateMemoryTable, CreateView, DdlStatement}; - use datafusion::sql::sqlparser::ast::CreateTable; - - if let Statement::CreateTable(CreateTable { query: None, .. }) = statement { - return plan_err!( - "CREATE TABLE without AS SELECT is not supported; use CREATE TABLE ... AS SELECT or a connector table" - ); - } - - match produce_optimized_plan(statement, schema_provider) { - Ok(LogicalPlan::Ddl(DdlStatement::CreateView(CreateView { name, input, .. }))) - | Ok(LogicalPlan::Ddl(DdlStatement::CreateMemoryTable(CreateMemoryTable { - name, - input, - .. - }))) => { - let rewritten = rewrite_plan(input.as_ref().clone(), schema_provider)?; - let schema = rewritten.schema().clone(); - let remote = RemoteTableBoundaryNode { - upstream_plan: rewritten, - table_identifier: name.to_owned(), - resolved_schema: schema, - requires_materialization: true, - }; - Ok(Some(CatalogEntity::ComputedTable { - name: name.to_string(), - logical_plan: Box::new(LogicalPlan::Extension(Extension { - node: Arc::new(remote), - })), - })) - } - _ => Ok(None), - } - } - - pub fn name(&self) -> &str { - match self { - CatalogEntity::ComputedTable { name, .. } => name.as_str(), - CatalogEntity::ExternalConnector(e) => e.name(), - } - } - - pub fn get_fields(&self) -> Vec { - match self { - CatalogEntity::ExternalConnector(e) => e.effective_fields(), - CatalogEntity::ComputedTable { logical_plan, .. } => { - logical_plan.schema().fields().iter().cloned().collect() - } - } - } - - pub fn set_inferred_fields(&mut self, fields: Vec) -> Result<()> { - let CatalogEntity::ExternalConnector(ext) = self else { - return Ok(()); - }; - let ExternalTable::Source(t) = ext.as_mut() else { - return Ok(()); - }; - - if !t.schema_specs.is_empty() { - return Ok(()); - } - - if let Some(existing) = &t.inferred_fields { - let matches = existing.len() == fields.len() - && existing - .iter() - .zip(&fields) - .all(|(a, b)| a.name() == b.name() && a.data_type() == b.data_type()); - - if !matches { - return plan_err!("all inserts into a table must share the same schema"); - } - } - - let fields: Vec<_> = fields.into_iter().map(|f| f.field().clone()).collect(); - t.inferred_fields.replace(fields); - - Ok(()) - } - - pub fn connector_op(&self) -> Result { - match self { - CatalogEntity::ExternalConnector(e) => Ok(e.connector_op()), - CatalogEntity::ComputedTable { .. } => { - plan_err!("can't write to a query-defined table") - } - } - } - - pub fn partition_exprs(&self) -> Option<&Vec> { - let CatalogEntity::ExternalConnector(ext) = self else { - return None; - }; - let ExternalTable::Sink(s) = ext.as_ref() else { - return None; - }; - (*s.partition_exprs).as_ref() - } - - #[inline] - pub fn as_external(&self) -> Option<&ExternalTable> { - match self { - CatalogEntity::ExternalConnector(e) => Some(e.as_ref()), - _ => None, - } - } -} - -#[derive(Clone, Debug)] -pub struct SqlSource { - pub id: Option, - pub struct_def: Vec, - pub config: ConnectorOp, - pub processing_mode: ProcessingMode, - pub idle_time: Option, -} diff --git a/src/sql/schema/table_role.rs b/src/sql/schema/table_role.rs deleted file mode 100644 index 7d301f9d..00000000 --- a/src/sql/schema/table_role.rs +++ /dev/null @@ -1,104 +0,0 @@ -// 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 datafusion::arrow::datatypes::{DataType, TimeUnit}; -use datafusion::common::{Result, plan_err}; -use datafusion::error::DataFusionError; - -use super::column_descriptor::ColumnDescriptor; -use super::connection_type::ConnectionType; -use crate::sql::common::constants::{ - SUPPORTED_CONNECTOR_ADAPTERS, connection_table_role, connector_type, -}; -use crate::sql::common::with_option_keys as opt; - -/// Role of a connector-backed table in the pipeline (ingest / egress / lookup). -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub enum TableRole { - Ingestion, - Egress, - Reference, -} - -impl From for ConnectionType { - fn from(r: TableRole) -> Self { - match r { - TableRole::Ingestion => ConnectionType::Source, - TableRole::Egress => ConnectionType::Sink, - TableRole::Reference => ConnectionType::Lookup, - } - } -} - -impl From for TableRole { - fn from(c: ConnectionType) -> Self { - match c { - ConnectionType::Source => TableRole::Ingestion, - ConnectionType::Sink => TableRole::Egress, - ConnectionType::Lookup => TableRole::Reference, - } - } -} - -pub fn validate_adapter_availability(adapter: &str) -> Result<()> { - if !SUPPORTED_CONNECTOR_ADAPTERS.contains(&adapter) { - return Err(DataFusionError::Plan(format!( - "Unknown adapter '{adapter}'" - ))); - } - Ok(()) -} - -pub fn apply_adapter_specific_rules( - adapter: &str, - mut cols: Vec, -) -> Vec { - match adapter { - a if a == connector_type::DELTA || a == connector_type::ICEBERG => { - for c in &mut cols { - if matches!(c.data_type(), DataType::Timestamp(_, _)) { - c.force_precision(TimeUnit::Microsecond); - } - } - cols - } - _ => cols, - } -} - -pub fn deduce_role(options: &HashMap) -> Result { - match options.get(opt::TYPE).map(|s| s.as_str()) { - None | Some(connection_table_role::SOURCE) => Ok(TableRole::Ingestion), - Some(connection_table_role::SINK) => Ok(TableRole::Egress), - Some(connection_table_role::LOOKUP) => Ok(TableRole::Reference), - Some(other) => plan_err!("Invalid role '{other}'"), - } -} - -pub fn serialize_backend_params( - adapter: &str, - options: &HashMap, -) -> Result { - let mut payload = serde_json::Map::new(); - payload.insert( - opt::ADAPTER.to_string(), - serde_json::Value::String(adapter.to_string()), - ); - - for (k, v) in options { - payload.insert(k.clone(), serde_json::Value::String(v.clone())); - } - - serde_json::to_string(&payload).map_err(|e| DataFusionError::Plan(e.to_string())) -} diff --git a/src/sql/schema/temporal_pipeline_config.rs b/src/sql/schema/temporal_pipeline_config.rs deleted file mode 100644 index f672e552..00000000 --- a/src/sql/schema/temporal_pipeline_config.rs +++ /dev/null @@ -1,58 +0,0 @@ -// 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::time::Duration; - -use datafusion::common::{Result, plan_err}; -use datafusion::logical_expr::Expr; - -use super::column_descriptor::ColumnDescriptor; -use crate::sql::common::constants::sql_field; - -/// Event-time and watermark configuration for streaming tables. -#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] -pub struct TemporalPipelineConfig { - pub event_column: Option, - pub watermark_strategy_column: Option, - pub liveness_timeout: Option, -} - -#[derive(Debug, Clone)] -pub struct TemporalSpec { - pub time_field: String, - pub watermark_expr: Option, -} - -pub fn resolve_temporal_logic( - columns: &[ColumnDescriptor], - time_meta: Option, -) -> Result { - let mut config = TemporalPipelineConfig::default(); - - if let Some(meta) = time_meta { - let field_exists = columns - .iter() - .any(|c| c.arrow_field().name() == meta.time_field.as_str()); - if !field_exists { - return plan_err!("Temporal field {} does not exist", meta.time_field); - } - config.event_column = Some(meta.time_field.clone()); - - if meta.watermark_expr.is_some() { - config.watermark_strategy_column = Some(sql_field::COMPUTED_WATERMARK.to_string()); - } else { - config.watermark_strategy_column = Some(meta.time_field); - } - } - - Ok(config) -} diff --git a/src/sql/schema/utils.rs b/src/sql/schema/utils.rs deleted file mode 100644 index 45254e5f..00000000 --- a/src/sql/schema/utils.rs +++ /dev/null @@ -1,79 +0,0 @@ -// 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::Arc; - -use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; -use datafusion::common::{DFSchema, DFSchemaRef, Result as DFResult, TableReference}; - -use crate::sql::common::constants::window_interval_field; -use crate::sql::types::{QualifiedField, TIMESTAMP_FIELD}; - -/// Returns the Arrow struct type for a window (start, end) pair. -pub fn window_arrow_struct() -> DataType { - DataType::Struct( - vec![ - Arc::new(Field::new( - window_interval_field::START, - DataType::Timestamp(TimeUnit::Nanosecond, None), - false, - )), - Arc::new(Field::new( - window_interval_field::END, - DataType::Timestamp(TimeUnit::Nanosecond, None), - false, - )), - ] - .into(), - ) -} - -/// Adds a `_timestamp` field to a DFSchema if it doesn't already have one. -pub fn add_timestamp_field( - schema: DFSchemaRef, - qualifier: Option, -) -> DFResult { - if has_timestamp_field(&schema) { - return Ok(schema); - } - - let timestamp_field = QualifiedField::new( - qualifier, - TIMESTAMP_FIELD, - DataType::Timestamp(TimeUnit::Nanosecond, None), - false, - ); - Ok(Arc::new(schema.join(&DFSchema::new_with_metadata( - vec![timestamp_field.into()], - HashMap::new(), - )?)?)) -} - -/// Checks whether a DFSchema contains a `_timestamp` field. -pub fn has_timestamp_field(schema: &DFSchemaRef) -> bool { - schema - .fields() - .iter() - .any(|field| field.name() == TIMESTAMP_FIELD) -} - -/// Adds a `_timestamp` field to an Arrow Schema, returning a new SchemaRef. -pub fn add_timestamp_field_arrow(schema: Schema) -> SchemaRef { - let mut fields = schema.fields().to_vec(); - fields.push(Arc::new(Field::new( - TIMESTAMP_FIELD, - DataType::Timestamp(TimeUnit::Nanosecond, None), - false, - ))); - Arc::new(Schema::new(fields)) -} diff --git a/src/sql/types/data_type.rs b/src/sql/types/data_type.rs deleted file mode 100644 index 387a4190..00000000 --- a/src/sql/types/data_type.rs +++ /dev/null @@ -1,158 +0,0 @@ -// 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::sync::Arc; - -use datafusion::arrow::datatypes::{ - DECIMAL_DEFAULT_SCALE, DECIMAL128_MAX_PRECISION, DataType, Field, IntervalUnit, TimeUnit, -}; -use datafusion::common::{Result, plan_datafusion_err, plan_err}; - -use crate::sql::common::FsExtensionType; -use crate::sql::common::constants::planning_placeholder_udf; - -pub fn convert_data_type( - sql_type: &datafusion::sql::sqlparser::ast::DataType, -) -> Result<(DataType, Option)> { - use datafusion::sql::sqlparser::ast::ArrayElemTypeDef; - use datafusion::sql::sqlparser::ast::DataType as SQLDataType; - - match sql_type { - SQLDataType::Array(ArrayElemTypeDef::AngleBracket(inner_sql_type)) - | SQLDataType::Array(ArrayElemTypeDef::SquareBracket(inner_sql_type, _)) => { - let (data_type, extension) = convert_simple_data_type(inner_sql_type)?; - - Ok(( - DataType::List(Arc::new(FsExtensionType::add_metadata( - extension, - Field::new( - planning_placeholder_udf::LIST_ELEMENT_FIELD, - data_type, - true, - ), - ))), - None, - )) - } - SQLDataType::Array(ArrayElemTypeDef::None) => { - plan_err!("Arrays with unspecified type is not supported") - } - other => convert_simple_data_type(other), - } -} - -fn convert_simple_data_type( - sql_type: &datafusion::sql::sqlparser::ast::DataType, -) -> Result<(DataType, Option)> { - use datafusion::sql::sqlparser::ast::DataType as SQLDataType; - use datafusion::sql::sqlparser::ast::{ExactNumberInfo, TimezoneInfo}; - - if matches!(sql_type, SQLDataType::JSON) { - return Ok((DataType::Utf8, Some(FsExtensionType::JSON))); - } - - let dt = match sql_type { - SQLDataType::Boolean | SQLDataType::Bool => Ok(DataType::Boolean), - SQLDataType::TinyInt(_) => Ok(DataType::Int8), - SQLDataType::SmallInt(_) | SQLDataType::Int2(_) => Ok(DataType::Int16), - SQLDataType::Int(_) | SQLDataType::Integer(_) | SQLDataType::Int4(_) => Ok(DataType::Int32), - SQLDataType::BigInt(_) | SQLDataType::Int8(_) => Ok(DataType::Int64), - SQLDataType::TinyIntUnsigned(_) => Ok(DataType::UInt8), - SQLDataType::SmallIntUnsigned(_) | SQLDataType::Int2Unsigned(_) => Ok(DataType::UInt16), - SQLDataType::IntUnsigned(_) - | SQLDataType::UnsignedInteger - | SQLDataType::Int4Unsigned(_) => Ok(DataType::UInt32), - SQLDataType::BigIntUnsigned(_) | SQLDataType::Int8Unsigned(_) => Ok(DataType::UInt64), - SQLDataType::Float(_) => Ok(DataType::Float32), - SQLDataType::Real | SQLDataType::Float4 => Ok(DataType::Float32), - SQLDataType::Double(_) | SQLDataType::DoublePrecision | SQLDataType::Float8 => { - Ok(DataType::Float64) - } - SQLDataType::Char(_) - | SQLDataType::Varchar(_) - | SQLDataType::Text - | SQLDataType::String(_) => Ok(DataType::Utf8), - SQLDataType::Timestamp(None, TimezoneInfo::None) | SQLDataType::Datetime(_) => { - Ok(DataType::Timestamp(TimeUnit::Nanosecond, None)) - } - SQLDataType::Timestamp(Some(precision), TimezoneInfo::None) => match *precision { - 0 => Ok(DataType::Timestamp(TimeUnit::Second, None)), - 3 => Ok(DataType::Timestamp(TimeUnit::Millisecond, None)), - 6 => Ok(DataType::Timestamp(TimeUnit::Microsecond, None)), - 9 => Ok(DataType::Timestamp(TimeUnit::Nanosecond, None)), - _ => { - return plan_err!( - "unsupported precision {} -- supported precisions are 0 (seconds), \ - 3 (milliseconds), 6 (microseconds), and 9 (nanoseconds)", - precision - ); - } - }, - SQLDataType::Date => Ok(DataType::Date32), - SQLDataType::Time(None, tz_info) - if matches!(tz_info, TimezoneInfo::None) - || matches!(tz_info, TimezoneInfo::WithoutTimeZone) => - { - Ok(DataType::Time64(TimeUnit::Nanosecond)) - } - SQLDataType::Numeric(exact_number_info) | SQLDataType::Decimal(exact_number_info) => { - let (precision, scale) = match *exact_number_info { - ExactNumberInfo::None => (None, None), - ExactNumberInfo::Precision(precision) => (Some(precision), None), - ExactNumberInfo::PrecisionAndScale(precision, scale) => { - (Some(precision), Some(scale)) - } - }; - make_decimal_type(precision, scale) - } - SQLDataType::Bytea => Ok(DataType::Binary), - SQLDataType::Interval => Ok(DataType::Interval(IntervalUnit::MonthDayNano)), - SQLDataType::Struct(fields, _) => { - let fields: Vec<_> = fields - .iter() - .map(|f| { - Ok::<_, datafusion::error::DataFusionError>(Arc::new(Field::new( - f.field_name - .as_ref() - .ok_or_else(|| { - plan_datafusion_err!("anonymous struct fields are not allowed") - })? - .to_string(), - convert_data_type(&f.field_type)?.0, - true, - ))) - }) - .collect::>()?; - Ok(DataType::Struct(fields.into())) - } - _ => return plan_err!("Unsupported SQL type {sql_type:?}"), - }; - - Ok((dt?, None)) -} - -fn make_decimal_type(precision: Option, scale: Option) -> Result { - let (precision, scale) = match (precision, scale) { - (Some(p), Some(s)) => (p as u8, s as i8), - (Some(p), None) => (p as u8, 0), - (None, Some(_)) => return plan_err!("Cannot specify only scale for decimal data type"), - (None, None) => (DECIMAL128_MAX_PRECISION, DECIMAL_DEFAULT_SCALE), - }; - - if precision == 0 || precision > DECIMAL128_MAX_PRECISION || scale.unsigned_abs() > precision { - plan_err!( - "Decimal(precision = {precision}, scale = {scale}) should satisfy `0 < precision <= 38`, and `scale <= precision`." - ) - } else { - Ok(DataType::Decimal128(precision, scale)) - } -} diff --git a/src/sql/types/df_field.rs b/src/sql/types/df_field.rs deleted file mode 100644 index a32d7bc8..00000000 --- a/src/sql/types/df_field.rs +++ /dev/null @@ -1,181 +0,0 @@ -// 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::Arc; - -use datafusion::arrow::datatypes::{DataType, Field, FieldRef}; -use datafusion::common::{Column, DFSchema, Result, TableReference}; - -// ============================================================================ -// QualifiedField (Strongly-typed Field Wrapper) -// ============================================================================ - -/// Arrow [`Field`] plus optional SQL [`TableReference`] qualifier. -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct QualifiedField { - qualifier: Option, - field: FieldRef, -} - -// ============================================================================ -// Type Conversions (Interoperability with DataFusion) -// ============================================================================ - -impl From<(Option, FieldRef)> for QualifiedField { - fn from((qualifier, field): (Option, FieldRef)) -> Self { - Self { qualifier, field } - } -} - -impl From<(Option<&TableReference>, &Field)> for QualifiedField { - fn from((qualifier, field): (Option<&TableReference>, &Field)) -> Self { - Self { - qualifier: qualifier.cloned(), - field: Arc::new(field.clone()), - } - } -} - -impl From for (Option, FieldRef) { - fn from(value: QualifiedField) -> Self { - (value.qualifier, value.field) - } -} - -// ============================================================================ -// Core API -// ============================================================================ - -impl QualifiedField { - pub fn new( - qualifier: Option, - name: impl Into, - data_type: DataType, - nullable: bool, - ) -> Self { - Self { - qualifier, - field: Arc::new(Field::new(name, data_type, nullable)), - } - } - - pub fn new_unqualified(name: &str, data_type: DataType, nullable: bool) -> Self { - Self { - qualifier: None, - field: Arc::new(Field::new(name, data_type, nullable)), - } - } - - #[inline] - pub fn name(&self) -> &str { - self.field.name() - } - - #[inline] - pub fn data_type(&self) -> &DataType { - self.field.data_type() - } - - #[inline] - pub fn is_nullable(&self) -> bool { - self.field.is_nullable() - } - - #[inline] - pub fn metadata(&self) -> &HashMap { - self.field.metadata() - } - - #[inline] - pub fn qualifier(&self) -> Option<&TableReference> { - self.qualifier.as_ref() - } - - #[inline] - pub fn field(&self) -> &FieldRef { - &self.field - } - - pub fn qualified_name(&self) -> String { - match &self.qualifier { - Some(qualifier) => format!("{}.{}", qualifier, self.field.name()), - None => self.field.name().to_owned(), - } - } - - pub fn qualified_column(&self) -> Column { - Column { - relation: self.qualifier.clone(), - name: self.field.name().to_string(), - spans: Default::default(), - } - } - - pub fn unqualified_column(&self) -> Column { - Column { - relation: None, - name: self.field.name().to_string(), - spans: Default::default(), - } - } - - pub fn strip_qualifier(mut self) -> Self { - self.qualifier = None; - self - } - - pub fn with_nullable(mut self, nullable: bool) -> Self { - if self.field.is_nullable() == nullable { - return self; - } - let field = Arc::try_unwrap(self.field).unwrap_or_else(|arc| (*arc).clone()); - self.field = Arc::new(field.with_nullable(nullable)); - self - } - - pub fn with_metadata(mut self, metadata: HashMap) -> Self { - let field = Arc::try_unwrap(self.field).unwrap_or_else(|arc| (*arc).clone()); - self.field = Arc::new(field.with_metadata(metadata)); - self - } -} - -// ============================================================================ -// Schema Collection Helpers -// ============================================================================ - -pub fn extract_qualified_fields(schema: &DFSchema) -> Vec { - schema - .fields() - .iter() - .enumerate() - .map(|(i, field)| { - let (qualifier, _) = schema.qualified_field(i); - QualifiedField { - qualifier: qualifier.cloned(), - field: field.clone(), - } - }) - .collect() -} - -pub fn build_df_schema(fields: &[QualifiedField]) -> Result { - build_df_schema_with_metadata(fields, HashMap::new()) -} - -pub fn build_df_schema_with_metadata( - fields: &[QualifiedField], - metadata: HashMap, -) -> Result { - DFSchema::new_with_metadata(fields.iter().map(|f| f.clone().into()).collect(), metadata) -} diff --git a/src/sql/types/mod.rs b/src/sql/types/mod.rs deleted file mode 100644 index d5124bcc..00000000 --- a/src/sql/types/mod.rs +++ /dev/null @@ -1,65 +0,0 @@ -// 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. - -mod data_type; -mod df_field; -pub(crate) mod placeholder_udf; -mod stream_schema; -mod window; - -use std::time::Duration; - -use crate::sql::common::constants::sql_planning_default; - -pub use df_field::{ - QualifiedField, build_df_schema, build_df_schema_with_metadata, extract_qualified_fields, -}; -pub(crate) use placeholder_udf::PlanningPlaceholderUdf; -pub(crate) use window::WindowBehavior; -pub use window::{WindowType, extract_window_type}; - -pub use crate::sql::common::constants::sql_field::TIMESTAMP_FIELD; - -#[derive(Clone, Debug, Eq, PartialEq)] -pub enum ProcessingMode { - Append, - Update, -} - -#[derive(Clone, Debug)] -pub struct SqlConfig { - pub default_parallelism: usize, - /// Physical pipeline parallelism for [`KeyExtractionNode`](crate::sql::logical_node::key_calculation::KeyExtractionNode) / KeyBy. - pub key_by_parallelism: usize, -} - -impl Default for SqlConfig { - fn default() -> Self { - Self { - default_parallelism: sql_planning_default::DEFAULT_PARALLELISM, - key_by_parallelism: sql_planning_default::DEFAULT_KEY_BY_PARALLELISM, - } - } -} - -#[derive(Clone)] -pub struct PlanningOptions { - pub ttl: Duration, -} - -impl Default for PlanningOptions { - fn default() -> Self { - Self { - ttl: Duration::from_secs(sql_planning_default::PLANNING_TTL_SECS), - } - } -} diff --git a/src/sql/types/placeholder_udf.rs b/src/sql/types/placeholder_udf.rs deleted file mode 100644 index 059637e9..00000000 --- a/src/sql/types/placeholder_udf.rs +++ /dev/null @@ -1,79 +0,0 @@ -// 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::any::Any; -use std::fmt::{Debug, Formatter}; -use std::sync::Arc; - -use datafusion::arrow::datatypes::DataType; -use datafusion::common::{Result, internal_err}; -use datafusion::logical_expr::{ - ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, -}; - -// ============================================================================ -// PlanningPlaceholderUdf -// ============================================================================ - -/// Logical-planning-only UDF: satisfies type checking until real functions are wired in. -pub(crate) struct PlanningPlaceholderUdf { - name: String, - signature: Signature, - return_type: DataType, -} - -impl Debug for PlanningPlaceholderUdf { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "PlanningPlaceholderUDF<{}>", self.name) - } -} - -impl ScalarUDFImpl for PlanningPlaceholderUdf { - fn as_any(&self) -> &dyn Any { - self - } - - fn name(&self) -> &str { - &self.name - } - - fn signature(&self) -> &Signature { - &self.signature - } - - fn return_type(&self, _arg_types: &[DataType]) -> Result { - Ok(self.return_type.clone()) - } - - fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { - internal_err!( - "PlanningPlaceholderUDF '{}' was invoked during physical execution. \ - This indicates a bug in the stream query compiler: placeholders must be \ - swapped with actual physical UDFs before execution begins.", - self.name - ) - } -} - -impl PlanningPlaceholderUdf { - pub fn new_with_return( - name: impl Into, - args: Vec, - return_type: DataType, - ) -> Arc { - Arc::new(ScalarUDF::new_from_impl(Self { - name: name.into(), - signature: Signature::exact(args, Volatility::Volatile), - return_type, - })) - } -} diff --git a/src/sql/types/stream_schema.rs b/src/sql/types/stream_schema.rs deleted file mode 100644 index c973386e..00000000 --- a/src/sql/types/stream_schema.rs +++ /dev/null @@ -1,133 +0,0 @@ -// 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::sync::Arc; - -use datafusion::arrow::datatypes::{Field, Schema, SchemaRef}; -use datafusion::common::{DataFusionError, Result}; - -use super::TIMESTAMP_FIELD; - -// ============================================================================ -// StreamSchema -// ============================================================================ - -/// Schema wrapper for continuous streaming: requires event-time (`TIMESTAMP_FIELD`) for watermarks -/// and optionally tracks key column indices for partitioned state / shuffle. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct StreamSchema { - schema: SchemaRef, - timestamp_index: usize, - key_indices: Option>, -} - -impl StreamSchema { - // ======================================================================== - // Raw Constructors (When indices are strictly known in advance) - // ======================================================================== - - /// Keyed stream when indices are already verified. - pub fn new_keyed(schema: SchemaRef, timestamp_index: usize, key_indices: Vec) -> Self { - Self { - schema, - timestamp_index, - key_indices: Some(key_indices), - } - } - - /// Unkeyed stream when `timestamp_index` is already verified. - pub fn new_unkeyed(schema: SchemaRef, timestamp_index: usize) -> Self { - Self { - schema, - timestamp_index, - key_indices: None, - } - } - - // ======================================================================== - // Safe Builders (Dynamically resolves and validates indices) - // ======================================================================== - - /// Unkeyed stream from a field list. Replaces the old `unwrap_or(0)` default when the timestamp - /// column was missing (silent wrong index / corruption). - pub fn try_from_fields(fields: impl Into>) -> Result { - let schema = Arc::new(Schema::new(fields.into())); - Self::try_from_schema_unkeyed(schema) - } - - /// Keyed stream from `SchemaRef`; resolves and validates the mandatory timestamp column. - pub fn try_from_schema_keyed(schema: SchemaRef, key_indices: Vec) -> Result { - let timestamp_index = schema - .column_with_name(TIMESTAMP_FIELD) - .ok_or_else(|| { - DataFusionError::Plan(format!( - "Streaming Topology Error: Mandatory event-time field '{}' is missing in the schema. \ - Current schema fields: {:?}", - TIMESTAMP_FIELD, - schema.fields() - )) - })? - .0; - - Ok(Self { - schema, - timestamp_index, - key_indices: Some(key_indices), - }) - } - - /// Unkeyed stream from `SchemaRef`; resolves and validates the mandatory timestamp column. - pub fn try_from_schema_unkeyed(schema: SchemaRef) -> Result { - let timestamp_index = schema - .column_with_name(TIMESTAMP_FIELD) - .ok_or_else(|| { - DataFusionError::Plan(format!( - "Streaming Topology Error: Mandatory event-time field '{}' is missing.", - TIMESTAMP_FIELD - )) - })? - .0; - - Ok(Self { - schema, - timestamp_index, - key_indices: None, - }) - } - - // ======================================================================== - // Zero-cost Getters - // ======================================================================== - - /// Underlying Arrow schema. - #[inline] - pub fn arrow_schema(&self) -> &SchemaRef { - &self.schema - } - - /// Physical column index used as event time / watermark driver. - #[inline] - pub fn timestamp_index(&self) -> usize { - self.timestamp_index - } - - /// Key column indices for shuffle / state, if keyed. - #[inline] - pub fn key_indices(&self) -> Option<&[usize]> { - self.key_indices.as_deref() - } - - #[inline] - pub fn is_keyed(&self) -> bool { - self.key_indices.is_some() - } -} diff --git a/src/sql/types/window.rs b/src/sql/types/window.rs deleted file mode 100644 index 1aa05f42..00000000 --- a/src/sql/types/window.rs +++ /dev/null @@ -1,134 +0,0 @@ -// 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::time::Duration; - -use datafusion::common::{Result, ScalarValue, not_impl_err, plan_err}; -use datafusion::logical_expr::Expr; -use datafusion::logical_expr::expr::{Alias, ScalarFunction}; - -use crate::sql::common::constants::window_fn; - -use super::QualifiedField; - -// ============================================================================ -// Window Definitions -// ============================================================================ - -/// Temporal windowing semantics for streaming aggregations. -#[derive(Clone, Debug, PartialEq, Eq, Hash)] -pub enum WindowType { - Tumbling { width: Duration }, - Sliding { width: Duration, slide: Duration }, - Session { gap: Duration }, - Instant, -} - -/// How windowing is represented in the physical plan. -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub(crate) enum WindowBehavior { - FromOperator { - window: WindowType, - window_field: QualifiedField, - window_index: usize, - is_nested: bool, - }, - InData, -} - -// ============================================================================ -// Logical Expression Parsers -// ============================================================================ - -pub fn extract_duration(expression: &Expr) -> Result { - match expression { - Expr::Literal(ScalarValue::IntervalDayTime(Some(val)), _) => { - let secs = (val.days as u64) * 24 * 60 * 60; - let millis = val.milliseconds as u64; - Ok(Duration::from_secs(secs) + Duration::from_millis(millis)) - } - Expr::Literal(ScalarValue::IntervalMonthDayNano(Some(val)), _) => { - if val.months != 0 { - return not_impl_err!( - "Streaming engine does not support window durations specified in months due to variable month lengths." - ); - } - let secs = (val.days as u64) * 24 * 60 * 60; - let nanos = val.nanoseconds as u64; - Ok(Duration::from_secs(secs) + Duration::from_nanos(nanos)) - } - _ => plan_err!( - "Unsupported window duration expression. Expected an interval literal (e.g., INTERVAL '1' MINUTE), got: {}", - expression - ), - } -} - -pub fn extract_window_type(expression: &Expr) -> Result> { - match expression { - Expr::ScalarFunction(ScalarFunction { func, args }) => match func.name() { - name if name == window_fn::HOP => { - if args.len() != 2 { - return plan_err!( - "hop() window function expects exactly 2 arguments (slide, width), got {}", - args.len() - ); - } - - let slide = extract_duration(&args[0])?; - let width = extract_duration(&args[1])?; - - if width.as_nanos() % slide.as_nanos() != 0 { - return plan_err!( - "Streaming Topology Error: hop() window width {:?} must be a perfect multiple of slide {:?}", - width, - slide - ); - } - - if slide == width { - Ok(Some(WindowType::Tumbling { width })) - } else { - Ok(Some(WindowType::Sliding { width, slide })) - } - } - - name if name == window_fn::TUMBLE => { - if args.len() != 1 { - return plan_err!( - "tumble() window function expects exactly 1 argument (width), got {}", - args.len() - ); - } - let width = extract_duration(&args[0])?; - Ok(Some(WindowType::Tumbling { width })) - } - - name if name == window_fn::SESSION => { - if args.len() != 1 { - return plan_err!( - "session() window function expects exactly 1 argument (gap), got {}", - args.len() - ); - } - let gap = extract_duration(&args[0])?; - Ok(Some(WindowType::Session { gap })) - } - - _ => Ok(None), - }, - - Expr::Alias(Alias { expr, .. }) => extract_window_type(expr), - - _ => Ok(None), - } -} diff --git a/src/storage/mod.rs b/src/storage/mod.rs deleted file mode 100644 index d72e8096..00000000 --- a/src/storage/mod.rs +++ /dev/null @@ -1,61 +0,0 @@ -// 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::sync::Arc; - -use anyhow::Context; - -// State backend sources live under `src/wasm_runtime/src/state_backend/`; compiled here for `crate::storage::state_backend`. -#[path = "../wasm_runtime/src/state_backend/mod.rs"] -pub mod state_backend; - -// Stream catalog + task storage sources under `src/catalog_storage/src/{stream_catalog,task}/`; -// compiled here so `crate::storage::stream_catalog` / `crate::storage::task` keep resolving. -#[path = "../catalog_storage/src/stream_catalog/mod.rs"] -pub mod stream_catalog; - -#[path = "../catalog_storage/src/task/mod.rs"] -pub mod task; - -/// Install the process-global [`stream_catalog::CatalogManager`] from configuration. -/// In-memory when `config.stream_catalog.persist` is `false`, otherwise a durable -/// [`stream_catalog::RocksDbMetaStore`] (default path: `{data_dir}/catalog.db`). -pub fn initialize_stream_catalog(config: &crate::config::GlobalConfig) -> anyhow::Result<()> { - use stream_catalog::{CatalogManager, InMemoryMetaStore, MetaStore, RocksDbMetaStore}; - - let store: Arc = if !config.stream_catalog.persist { - Arc::new(InMemoryMetaStore::new()) - } else { - let path = config - .stream_catalog - .db_path - .as_ref() - .map(|p| crate::config::resolve_path(p)) - .unwrap_or_else(|| crate::config::get_data_dir().join("catalog.db")); - - std::fs::create_dir_all(&path).with_context(|| { - format!( - "Failed to create stream catalog RocksDB directory {}", - path.display() - ) - })?; - - Arc::new(RocksDbMetaStore::open(&path).with_context(|| { - format!( - "Failed to open stream catalog RocksDB at {}", - path.display() - ) - })?) - }; - - CatalogManager::init_global(store).context("Stream catalog (CatalogManager) global init failed") -} diff --git a/src/streaming_runtime/src/lib.rs b/src/streaming_runtime/src/lib.rs index ad83a58c..c6c2c94d 100644 --- a/src/streaming_runtime/src/lib.rs +++ b/src/streaming_runtime/src/lib.rs @@ -15,7 +15,6 @@ //! The streaming engine and shared runtime helpers (`streaming/`, `util/`) are //! implemented under [`src/streaming`] and [`src/util`] in this package. They are //! currently **compiled as part of the `function-stream` crate** via `#[path]` in -//! `src/runtime/mod.rs`, so in-tree paths like `crate::sql` -//! keep working until SQL is split into its own crate. +//! `src/runtime.rs`, sharing the root `crate::sql` name (re-exported streaming planner crate). pub const CRATE_NAME: &str = "function-stream-streaming-runtime"; diff --git a/src/streaming_runtime/src/streaming/protocol/event.rs b/src/streaming_runtime/src/streaming/protocol/event.rs index 093d99ba..8c0a4989 100644 --- a/src/streaming_runtime/src/streaming/protocol/event.rs +++ b/src/streaming_runtime/src/streaming/protocol/event.rs @@ -10,28 +10,15 @@ // See the License for the specific language governing permissions and // limitations under the License. -use bincode::{Decode, Encode}; -use serde::{Deserialize, Serialize}; use std::sync::Arc; -use std::time::SystemTime; use arrow_array::RecordBatch; use crate::runtime::memory::MemoryTicket; -#[derive(Debug, Copy, Clone, PartialEq, Eq, Encode, Decode, Serialize, Deserialize)] -pub enum Watermark { - EventTime(SystemTime), - Idle, -} - -#[derive(Debug, Copy, Clone, PartialEq, Eq, Encode, Decode, Serialize, Deserialize)] -pub struct CheckpointBarrier { - pub epoch: u64, - pub min_epoch: u64, - pub timestamp: SystemTime, - pub then_stop: bool, -} +pub use function_stream_runtime_common::streaming_protocol::{ + CheckpointBarrier, Watermark, merge_watermarks, watermark_strictly_advances, +}; #[derive(Debug, Clone)] pub enum StreamEvent { @@ -70,77 +57,3 @@ impl TrackedEvent { } } } - -pub fn merge_watermarks(per_input: &[Option]) -> Option { - if per_input.iter().any(|w| w.is_none()) { - return None; - } - - let mut min_event: Option = None; - let mut all_idle = true; - - for w in per_input.iter().flatten() { - match w { - Watermark::Idle => {} - Watermark::EventTime(t) => { - all_idle = false; - min_event = Some(match min_event { - None => *t, - Some(m) => m.min(*t), - }); - } - } - } - - if all_idle { - Some(Watermark::Idle) - } else { - Some(Watermark::EventTime(min_event.expect( - "non-idle alignment must have at least one EventTime", - ))) - } -} - -pub fn watermark_strictly_advances(new: Watermark, previous: Option) -> bool { - match previous { - None => true, - Some(prev) => match (new, prev) { - (Watermark::EventTime(tn), Watermark::EventTime(tp)) => tn > tp, - (Watermark::Idle, Watermark::Idle) => false, - (Watermark::Idle, Watermark::EventTime(_)) => true, - (Watermark::EventTime(_), Watermark::Idle) => true, - }, - } -} - -#[cfg(test)] -mod watermark_tests { - use super::*; - use std::time::Duration; - - #[test] - fn merge_waits_for_all_channels() { - let wms = vec![Some(Watermark::EventTime(SystemTime::UNIX_EPOCH)), None]; - assert!(merge_watermarks(&wms).is_none()); - } - - #[test] - fn merge_min_event_time_ignores_idle() { - let t1 = SystemTime::UNIX_EPOCH + Duration::from_secs(10); - let t2 = SystemTime::UNIX_EPOCH + Duration::from_secs(5); - let wms = vec![Some(Watermark::EventTime(t1)), Some(Watermark::Idle)]; - assert_eq!(merge_watermarks(&wms), Some(Watermark::EventTime(t1))); - - let wms = vec![ - Some(Watermark::EventTime(t1)), - Some(Watermark::EventTime(t2)), - ]; - assert_eq!(merge_watermarks(&wms), Some(Watermark::EventTime(t2))); - } - - #[test] - fn merge_all_idle() { - let wms = vec![Some(Watermark::Idle), Some(Watermark::Idle)]; - assert_eq!(merge_watermarks(&wms), Some(Watermark::Idle)); - } -} diff --git a/src/streaming_runtime/src/streaming/protocol/mod.rs b/src/streaming_runtime/src/streaming/protocol/mod.rs index 28fd85a4..d41d9fcf 100644 --- a/src/streaming_runtime/src/streaming/protocol/mod.rs +++ b/src/streaming_runtime/src/streaming/protocol/mod.rs @@ -15,4 +15,5 @@ pub mod event; #[allow(unused_imports)] pub use control::{ControlCommand, JobMasterEvent, StopMode}; +#[allow(unused_imports)] pub use event::{CheckpointBarrier, StreamOutput, Watermark}; diff --git a/src/wasm_runtime/src/lib.rs b/src/wasm_runtime/src/lib.rs index d0b4073e..dc036f7d 100644 --- a/src/wasm_runtime/src/lib.rs +++ b/src/wasm_runtime/src/lib.rs @@ -13,10 +13,10 @@ //! WebAssembly execution runtime. //! //! Implementation lives under `src/wasm/` in this package. It is currently **compiled as -//! part of the `function-stream` crate** via `#[path]` in `src/runtime/mod.rs`, so paths -//! like `crate::sql` and `crate::runtime::memory` keep resolving until further crate splits. +//! part of the `function-stream` crate** via `#[path]` in `src/runtime.rs`, so paths +//! like `crate::sql` (streaming planner dependency) and `crate::runtime::memory` keep resolving. //! //! Operator state storage (`state_backend/`) also lives in this package and is compiled via -//! `#[path]` from `src/storage/mod.rs` as `crate::storage::state_backend`. +//! `#[path]` from `src/storage.rs` as `crate::storage::state_backend`. pub const CRATE_NAME: &str = "function-stream-wasm-runtime"; From f9206590adc57f097c81224a1645fa9cf062a403 Mon Sep 17 00:00:00 2001 From: luoluoyuyu Date: Tue, 12 May 2026 09:04:23 +0800 Subject: [PATCH 4/7] update --- src/runtime.rs | 27 + src/runtime_common/src/streaming_protocol.rs | 105 +++ src/storage.rs | 61 ++ src/streaming_planner/Cargo.toml | 49 ++ .../src/analysis/aggregate_rewriter.rs | 279 ++++++++ .../src/analysis/async_udf_rewriter.rs | 133 ++++ .../src/analysis/join_rewriter.rs | 234 +++++++ src/streaming_planner/src/analysis/mod.rs | 214 ++++++ .../src/analysis/row_time_rewriter.rs | 49 ++ .../src/analysis/sink_input_rewriter.rs | 57 ++ .../src/analysis/source_metadata_visitor.rs | 73 ++ .../src/analysis/source_rewriter.rs | 305 +++++++++ .../src/analysis/stream_rewriter.rs | 234 +++++++ .../src/analysis/streaming_window_analzer.rs | 219 ++++++ .../src/analysis/time_window.rs | 83 +++ src/streaming_planner/src/analysis/udafs.rs | 43 ++ .../src/analysis/unnest_rewriter.rs | 179 +++++ .../src/analysis/window_function_rewriter.rs | 204 ++++++ src/streaming_planner/src/api/checkpoints.rs | 108 +++ src/streaming_planner/src/api/connections.rs | 620 +++++++++++++++++ src/streaming_planner/src/api/metrics.rs | 53 ++ src/streaming_planner/src/api/mod.rs | 46 ++ src/streaming_planner/src/api/pipelines.rs | 168 +++++ src/streaming_planner/src/api/public_ids.rs | 69 ++ .../src/api/schema_resolver.rs | 94 +++ src/streaming_planner/src/api/udfs.rs | 68 ++ src/streaming_planner/src/api/var_str.rs | 91 +++ src/streaming_planner/src/common/arrow_ext.rs | 182 +++++ .../src/common/connector_options.rs | 449 ++++++++++++ src/streaming_planner/src/common/constants.rs | 294 ++++++++ src/streaming_planner/src/common/control.rs | 164 +++++ src/streaming_planner/src/common/converter.rs | 95 +++ src/streaming_planner/src/common/date.rs | 86 +++ src/streaming_planner/src/common/debezium.rs | 148 ++++ src/streaming_planner/src/common/errors.rs | 92 +++ .../src/common/format_from_opts.rs | 182 +++++ src/streaming_planner/src/common/formats.rs | 267 ++++++++ src/streaming_planner/src/common/fs_schema.rs | 470 +++++++++++++ .../src/common/kafka_catalog.rs | 116 ++++ src/streaming_planner/src/common/mod.rs | 65 ++ .../src/common/operator_config.rs | 21 + .../src/common/time_utils.rs | 74 ++ src/streaming_planner/src/common/topology.rs | 295 ++++++++ .../src/common/with_option_keys.rs | 105 +++ src/streaming_planner/src/connector/config.rs | 91 +++ .../src/connector/factory.rs | 67 ++ src/streaming_planner/src/connector/mod.rs | 18 + .../src/connector/provider.rs | 52 ++ .../src/connector/registry.rs | 86 +++ .../src/connector/sink/delta.rs | 60 ++ .../src/connector/sink/filesystem.rs | 60 ++ .../src/connector/sink/iceberg.rs | 57 ++ .../src/connector/sink/kafka.rs | 159 +++++ .../src/connector/sink/lancedb.rs | 61 ++ .../src/connector/sink/mod.rs | 20 + .../src/connector/sink/runtime_config.rs | 137 ++++ .../src/connector/sink/s3.rs | 75 ++ .../src/connector/sink/utils.rs | 91 +++ .../src/connector/source/kafka.rs | 185 +++++ .../src/connector/source/mod.rs | 13 + src/streaming_planner/src/functions/mod.rs | 612 +++++++++++++++++ src/streaming_planner/src/lib.rs | 32 + .../src/logical_node/aggregate.rs | 644 ++++++++++++++++++ .../src/logical_node/async_udf.rs | 247 +++++++ .../src/logical_node/debezium.rs | 393 +++++++++++ .../src/logical_node/extension_try_from.rs | 70 ++ .../src/logical_node/is_retract.rs | 82 +++ .../src/logical_node/join.rs | 211 ++++++ .../src/logical_node/key_calculation.rs | 309 +++++++++ .../logical_node/logical/dylib_udf_config.rs | 71 ++ .../logical/fs_program_convert.rs | 200 ++++++ .../src/logical_node/logical/logical_edge.rs | 102 +++ .../src/logical_node/logical/logical_graph.rs | 30 + .../src/logical_node/logical/logical_node.rs | 87 +++ .../logical_node/logical/logical_program.rs | 153 +++++ .../src/logical_node/logical/mod.rs | 30 + .../logical_node/logical/operator_chain.rs | 142 ++++ .../src/logical_node/logical/operator_name.rs | 82 +++ .../logical_node/logical/program_config.rs | 33 + .../logical_node/logical/python_udf_config.rs | 23 + .../src/logical_node/lookup.rs | 256 +++++++ .../src/logical_node/macros.rs | 28 + src/streaming_planner/src/logical_node/mod.rs | 42 ++ .../src/logical_node/projection.rs | 239 +++++++ .../src/logical_node/remote_table.rs | 190 ++++++ .../src/logical_node/sink.rs | 247 +++++++ .../streaming_operator_blueprint.rs | 65 ++ .../src/logical_node/table_source.rs | 180 +++++ .../src/logical_node/timestamp_append.rs | 121 ++++ .../src/logical_node/updating_aggregate.rs | 245 +++++++ .../src/logical_node/watermark_node.rs | 229 +++++++ .../src/logical_node/windows_function.rs | 191 ++++++ .../src/logical_planner/mod.rs | 16 + .../logical_planner/optimizers/chaining.rs | 200 ++++++ .../src/logical_planner/optimizers/mod.rs | 20 + .../optimizers/optimized_plan.rs | 95 +++ .../src/logical_planner/streaming_planner.rs | 435 ++++++++++++ src/streaming_planner/src/parse.rs | 82 +++ .../src/physical/cdc/encode.rs | 342 ++++++++++ src/streaming_planner/src/physical/cdc/mod.rs | 17 + .../src/physical/cdc/unroll.rs | 322 +++++++++ src/streaming_planner/src/physical/codec.rs | 307 +++++++++ src/streaming_planner/src/physical/meta.rs | 47 ++ src/streaming_planner/src/physical/mod.rs | 23 + .../src/physical/source_exec.rs | 400 +++++++++++ src/streaming_planner/src/physical/udfs.rs | 138 ++++ src/streaming_planner/src/planning_runtime.rs | 35 + src/streaming_planner/src/schema/catalog.rs | 609 +++++++++++++++++ .../src/schema/column_descriptor.rs | 144 ++++ .../src/schema/connection_type.rs | 31 + .../src/schema/data_encoding_format.rs | 89 +++ .../src/schema/introspection/ddl_formatter.rs | 156 +++++ .../src/schema/introspection/mod.rs | 21 + .../schema/introspection/show_formatter.rs | 100 +++ .../schema/introspection/stream_formatter.rs | 120 ++++ src/streaming_planner/src/schema/mod.rs | 30 + .../src/schema/schema_provider.rs | 469 +++++++++++++ src/streaming_planner/src/schema/table.rs | 163 +++++ .../src/schema/table_role.rs | 104 +++ .../src/schema/temporal_pipeline_config.rs | 58 ++ src/streaming_planner/src/schema/utils.rs | 79 +++ src/streaming_planner/src/types/data_type.rs | 158 +++++ src/streaming_planner/src/types/df_field.rs | 181 +++++ src/streaming_planner/src/types/mod.rs | 65 ++ .../src/types/placeholder_udf.rs | 79 +++ .../src/types/stream_schema.rs | 133 ++++ src/streaming_planner/src/types/window.rs | 134 ++++ 127 files changed, 19190 insertions(+) create mode 100644 src/runtime.rs create mode 100644 src/runtime_common/src/streaming_protocol.rs create mode 100644 src/storage.rs create mode 100644 src/streaming_planner/Cargo.toml create mode 100644 src/streaming_planner/src/analysis/aggregate_rewriter.rs create mode 100644 src/streaming_planner/src/analysis/async_udf_rewriter.rs create mode 100644 src/streaming_planner/src/analysis/join_rewriter.rs create mode 100644 src/streaming_planner/src/analysis/mod.rs create mode 100644 src/streaming_planner/src/analysis/row_time_rewriter.rs create mode 100644 src/streaming_planner/src/analysis/sink_input_rewriter.rs create mode 100644 src/streaming_planner/src/analysis/source_metadata_visitor.rs create mode 100644 src/streaming_planner/src/analysis/source_rewriter.rs create mode 100644 src/streaming_planner/src/analysis/stream_rewriter.rs create mode 100644 src/streaming_planner/src/analysis/streaming_window_analzer.rs create mode 100644 src/streaming_planner/src/analysis/time_window.rs create mode 100644 src/streaming_planner/src/analysis/udafs.rs create mode 100644 src/streaming_planner/src/analysis/unnest_rewriter.rs create mode 100644 src/streaming_planner/src/analysis/window_function_rewriter.rs create mode 100644 src/streaming_planner/src/api/checkpoints.rs create mode 100644 src/streaming_planner/src/api/connections.rs create mode 100644 src/streaming_planner/src/api/metrics.rs create mode 100644 src/streaming_planner/src/api/mod.rs create mode 100644 src/streaming_planner/src/api/pipelines.rs create mode 100644 src/streaming_planner/src/api/public_ids.rs create mode 100644 src/streaming_planner/src/api/schema_resolver.rs create mode 100644 src/streaming_planner/src/api/udfs.rs create mode 100644 src/streaming_planner/src/api/var_str.rs create mode 100644 src/streaming_planner/src/common/arrow_ext.rs create mode 100644 src/streaming_planner/src/common/connector_options.rs create mode 100644 src/streaming_planner/src/common/constants.rs create mode 100644 src/streaming_planner/src/common/control.rs create mode 100644 src/streaming_planner/src/common/converter.rs create mode 100644 src/streaming_planner/src/common/date.rs create mode 100644 src/streaming_planner/src/common/debezium.rs create mode 100644 src/streaming_planner/src/common/errors.rs create mode 100644 src/streaming_planner/src/common/format_from_opts.rs create mode 100644 src/streaming_planner/src/common/formats.rs create mode 100644 src/streaming_planner/src/common/fs_schema.rs create mode 100644 src/streaming_planner/src/common/kafka_catalog.rs create mode 100644 src/streaming_planner/src/common/mod.rs create mode 100644 src/streaming_planner/src/common/operator_config.rs create mode 100644 src/streaming_planner/src/common/time_utils.rs create mode 100644 src/streaming_planner/src/common/topology.rs create mode 100644 src/streaming_planner/src/common/with_option_keys.rs create mode 100644 src/streaming_planner/src/connector/config.rs create mode 100644 src/streaming_planner/src/connector/factory.rs create mode 100644 src/streaming_planner/src/connector/mod.rs create mode 100644 src/streaming_planner/src/connector/provider.rs create mode 100644 src/streaming_planner/src/connector/registry.rs create mode 100644 src/streaming_planner/src/connector/sink/delta.rs create mode 100644 src/streaming_planner/src/connector/sink/filesystem.rs create mode 100644 src/streaming_planner/src/connector/sink/iceberg.rs create mode 100644 src/streaming_planner/src/connector/sink/kafka.rs create mode 100644 src/streaming_planner/src/connector/sink/lancedb.rs create mode 100644 src/streaming_planner/src/connector/sink/mod.rs create mode 100644 src/streaming_planner/src/connector/sink/runtime_config.rs create mode 100644 src/streaming_planner/src/connector/sink/s3.rs create mode 100644 src/streaming_planner/src/connector/sink/utils.rs create mode 100644 src/streaming_planner/src/connector/source/kafka.rs create mode 100644 src/streaming_planner/src/connector/source/mod.rs create mode 100644 src/streaming_planner/src/functions/mod.rs create mode 100644 src/streaming_planner/src/lib.rs create mode 100644 src/streaming_planner/src/logical_node/aggregate.rs create mode 100644 src/streaming_planner/src/logical_node/async_udf.rs create mode 100644 src/streaming_planner/src/logical_node/debezium.rs create mode 100644 src/streaming_planner/src/logical_node/extension_try_from.rs create mode 100644 src/streaming_planner/src/logical_node/is_retract.rs create mode 100644 src/streaming_planner/src/logical_node/join.rs create mode 100644 src/streaming_planner/src/logical_node/key_calculation.rs create mode 100644 src/streaming_planner/src/logical_node/logical/dylib_udf_config.rs create mode 100644 src/streaming_planner/src/logical_node/logical/fs_program_convert.rs create mode 100644 src/streaming_planner/src/logical_node/logical/logical_edge.rs create mode 100644 src/streaming_planner/src/logical_node/logical/logical_graph.rs create mode 100644 src/streaming_planner/src/logical_node/logical/logical_node.rs create mode 100644 src/streaming_planner/src/logical_node/logical/logical_program.rs create mode 100644 src/streaming_planner/src/logical_node/logical/mod.rs create mode 100644 src/streaming_planner/src/logical_node/logical/operator_chain.rs create mode 100644 src/streaming_planner/src/logical_node/logical/operator_name.rs create mode 100644 src/streaming_planner/src/logical_node/logical/program_config.rs create mode 100644 src/streaming_planner/src/logical_node/logical/python_udf_config.rs create mode 100644 src/streaming_planner/src/logical_node/lookup.rs create mode 100644 src/streaming_planner/src/logical_node/macros.rs create mode 100644 src/streaming_planner/src/logical_node/mod.rs create mode 100644 src/streaming_planner/src/logical_node/projection.rs create mode 100644 src/streaming_planner/src/logical_node/remote_table.rs create mode 100644 src/streaming_planner/src/logical_node/sink.rs create mode 100644 src/streaming_planner/src/logical_node/streaming_operator_blueprint.rs create mode 100644 src/streaming_planner/src/logical_node/table_source.rs create mode 100644 src/streaming_planner/src/logical_node/timestamp_append.rs create mode 100644 src/streaming_planner/src/logical_node/updating_aggregate.rs create mode 100644 src/streaming_planner/src/logical_node/watermark_node.rs create mode 100644 src/streaming_planner/src/logical_node/windows_function.rs create mode 100644 src/streaming_planner/src/logical_planner/mod.rs create mode 100644 src/streaming_planner/src/logical_planner/optimizers/chaining.rs create mode 100644 src/streaming_planner/src/logical_planner/optimizers/mod.rs create mode 100644 src/streaming_planner/src/logical_planner/optimizers/optimized_plan.rs create mode 100644 src/streaming_planner/src/logical_planner/streaming_planner.rs create mode 100644 src/streaming_planner/src/parse.rs create mode 100644 src/streaming_planner/src/physical/cdc/encode.rs create mode 100644 src/streaming_planner/src/physical/cdc/mod.rs create mode 100644 src/streaming_planner/src/physical/cdc/unroll.rs create mode 100644 src/streaming_planner/src/physical/codec.rs create mode 100644 src/streaming_planner/src/physical/meta.rs create mode 100644 src/streaming_planner/src/physical/mod.rs create mode 100644 src/streaming_planner/src/physical/source_exec.rs create mode 100644 src/streaming_planner/src/physical/udfs.rs create mode 100644 src/streaming_planner/src/planning_runtime.rs create mode 100644 src/streaming_planner/src/schema/catalog.rs create mode 100644 src/streaming_planner/src/schema/column_descriptor.rs create mode 100644 src/streaming_planner/src/schema/connection_type.rs create mode 100644 src/streaming_planner/src/schema/data_encoding_format.rs create mode 100644 src/streaming_planner/src/schema/introspection/ddl_formatter.rs create mode 100644 src/streaming_planner/src/schema/introspection/mod.rs create mode 100644 src/streaming_planner/src/schema/introspection/show_formatter.rs create mode 100644 src/streaming_planner/src/schema/introspection/stream_formatter.rs create mode 100644 src/streaming_planner/src/schema/mod.rs create mode 100644 src/streaming_planner/src/schema/schema_provider.rs create mode 100644 src/streaming_planner/src/schema/table.rs create mode 100644 src/streaming_planner/src/schema/table_role.rs create mode 100644 src/streaming_planner/src/schema/temporal_pipeline_config.rs create mode 100644 src/streaming_planner/src/schema/utils.rs create mode 100644 src/streaming_planner/src/types/data_type.rs create mode 100644 src/streaming_planner/src/types/df_field.rs create mode 100644 src/streaming_planner/src/types/mod.rs create mode 100644 src/streaming_planner/src/types/placeholder_udf.rs create mode 100644 src/streaming_planner/src/types/stream_schema.rs create mode 100644 src/streaming_planner/src/types/window.rs diff --git a/src/runtime.rs b/src/runtime.rs new file mode 100644 index 00000000..d08cb5c2 --- /dev/null +++ b/src/runtime.rs @@ -0,0 +1,27 @@ +// 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-tree runtime: streaming engine, util helpers, and WASM task runtime. +// Paths are relative to `src/` (this file lives at `src/runtime.rs`). + +pub use function_stream_runtime_common::{common, memory}; + +#[path = "streaming_runtime/src/streaming/mod.rs"] +pub mod streaming; + +#[path = "streaming_runtime/src/util/mod.rs"] +pub mod util; + +#[path = "wasm_runtime/src/wasm/mod.rs"] +pub mod wasm; + +pub use wasm::{input, output, processor}; diff --git a/src/runtime_common/src/streaming_protocol.rs b/src/runtime_common/src/streaming_protocol.rs new file mode 100644 index 00000000..7c578f7a --- /dev/null +++ b/src/runtime_common/src/streaming_protocol.rs @@ -0,0 +1,105 @@ +// 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. + +//! Streaming control-plane types shared by the SQL planner and the execution runtime. + +use bincode::{Decode, Encode}; +use serde::{Deserialize, Serialize}; +use std::time::SystemTime; + +#[derive(Debug, Copy, Clone, PartialEq, Eq, Encode, Decode, Serialize, Deserialize)] +pub enum Watermark { + EventTime(SystemTime), + Idle, +} + +#[derive(Debug, Copy, Clone, PartialEq, Eq, Encode, Decode, Serialize, Deserialize)] +pub struct CheckpointBarrier { + pub epoch: u64, + pub min_epoch: u64, + pub timestamp: SystemTime, + pub then_stop: bool, +} + +pub fn merge_watermarks(per_input: &[Option]) -> Option { + if per_input.iter().any(|w| w.is_none()) { + return None; + } + + let mut min_event: Option = None; + let mut all_idle = true; + + for w in per_input.iter().flatten() { + match w { + Watermark::Idle => {} + Watermark::EventTime(t) => { + all_idle = false; + min_event = Some(match min_event { + None => *t, + Some(m) => m.min(*t), + }); + } + } + } + + if all_idle { + Some(Watermark::Idle) + } else { + Some(Watermark::EventTime(min_event.expect( + "non-idle alignment must have at least one EventTime", + ))) + } +} + +pub fn watermark_strictly_advances(new: Watermark, previous: Option) -> bool { + match previous { + None => true, + Some(prev) => match (new, prev) { + (Watermark::EventTime(tn), Watermark::EventTime(tp)) => tn > tp, + (Watermark::Idle, Watermark::Idle) => false, + (Watermark::Idle, Watermark::EventTime(_)) => true, + (Watermark::EventTime(_), Watermark::Idle) => true, + }, + } +} + +#[cfg(test)] +mod watermark_tests { + use super::*; + use std::time::Duration; + + #[test] + fn merge_waits_for_all_channels() { + let wms = vec![Some(Watermark::EventTime(SystemTime::UNIX_EPOCH)), None]; + assert!(merge_watermarks(&wms).is_none()); + } + + #[test] + fn merge_min_event_time_ignores_idle() { + let t1 = SystemTime::UNIX_EPOCH + Duration::from_secs(10); + let t2 = SystemTime::UNIX_EPOCH + Duration::from_secs(5); + let wms = vec![Some(Watermark::EventTime(t1)), Some(Watermark::Idle)]; + assert_eq!(merge_watermarks(&wms), Some(Watermark::EventTime(t1))); + + let wms = vec![ + Some(Watermark::EventTime(t1)), + Some(Watermark::EventTime(t2)), + ]; + assert_eq!(merge_watermarks(&wms), Some(Watermark::EventTime(t2))); + } + + #[test] + fn merge_all_idle() { + let wms = vec![Some(Watermark::Idle), Some(Watermark::Idle)]; + assert_eq!(merge_watermarks(&wms), Some(Watermark::Idle)); + } +} diff --git a/src/storage.rs b/src/storage.rs new file mode 100644 index 00000000..713b149e --- /dev/null +++ b/src/storage.rs @@ -0,0 +1,61 @@ +// 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. + +// Persistent / catalog-related modules compiled into the root crate. +// Paths are relative to `src/` (this file lives at `src/storage.rs`). + +use std::sync::Arc; + +use anyhow::Context; + +#[path = "wasm_runtime/src/state_backend/mod.rs"] +pub mod state_backend; + +#[path = "catalog_storage/src/stream_catalog/mod.rs"] +pub mod stream_catalog; + +#[path = "catalog_storage/src/task/mod.rs"] +pub mod task; + +/// Install the process-global [`stream_catalog::CatalogManager`] from configuration. +/// In-memory when `config.stream_catalog.persist` is `false`, otherwise a durable +/// [`stream_catalog::RocksDbMetaStore`] (default path: `{data_dir}/catalog.db`). +pub fn initialize_stream_catalog(config: &crate::config::GlobalConfig) -> anyhow::Result<()> { + use stream_catalog::{CatalogManager, InMemoryMetaStore, MetaStore, RocksDbMetaStore}; + + let store: Arc = if !config.stream_catalog.persist { + Arc::new(InMemoryMetaStore::new()) + } else { + let path = config + .stream_catalog + .db_path + .as_ref() + .map(|p| crate::config::resolve_path(p)) + .unwrap_or_else(|| crate::config::get_data_dir().join("catalog.db")); + + std::fs::create_dir_all(&path).with_context(|| { + format!( + "Failed to create stream catalog RocksDB directory {}", + path.display() + ) + })?; + + Arc::new(RocksDbMetaStore::open(&path).with_context(|| { + format!( + "Failed to open stream catalog RocksDB at {}", + path.display() + ) + })?) + }; + + CatalogManager::init_global(store).context("Stream catalog (CatalogManager) global init failed") +} diff --git a/src/streaming_planner/Cargo.toml b/src/streaming_planner/Cargo.toml new file mode 100644 index 00000000..d1039f97 --- /dev/null +++ b/src/streaming_planner/Cargo.toml @@ -0,0 +1,49 @@ +[package] +name = "function-stream-streaming-planner" +version = "0.6.0" +edition = "2024" + +[lib] +name = "function_stream_streaming_planner" +path = "src/lib.rs" + +[dependencies] +protocol = { path = "../../protocol" } +prost = "0.13" +function-stream-config = { path = "../config" } +function-stream-runtime-common = { path = "../runtime_common" } +tokio = { version = "1.0", features = ["macros", "rt-multi-thread", "sync", "time", "net"] } +tokio-stream = "0.1.18" +anyhow = "1.0" +xxhash-rust = { version = "0.8", features = ["xxh3"] } +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" +thiserror = "2" +tracing = "0.1" +async-trait = "0.1" +futures = "0.3" +itertools = "0.14" +petgraph = "0.7" +unicase = "2.7" +rand = { version = "0.8", features = ["small_rng"] } +bincode = { version = "2", features = ["serde"] } +chrono = "0.4" +bytes = "1" +ahash = "0.8" +strum = { version = "0.26", features = ["derive"] } +serde_json_path = "0.7" + +arrow = { version = "55", default-features = false } +arrow-array = "55" +arrow-schema = { version = "55", features = ["serde"] } +arrow-json = { version = "55.2.0" } +apache-avro = "0.21" + +datafusion = { git = "https://github.com/FunctionStream/datafusion", branch = "48.0.1/fs" } +datafusion-common = { git = "https://github.com/FunctionStream/datafusion", branch = "48.0.1/fs" } +datafusion-execution = { git = "https://github.com/FunctionStream/datafusion", branch = "48.0.1/fs" } +datafusion-expr = { git = "https://github.com/FunctionStream/datafusion", branch = "48.0.1/fs" } +datafusion-physical-expr = { git = "https://github.com/FunctionStream/datafusion", branch = "48.0.1/fs" } +datafusion-proto = { git = "https://github.com/FunctionStream/datafusion", branch = "48.0.1/fs" } + +sqlparser = { git = "https://github.com/FunctionStream/sqlparser-rs", branch = "0.58.0/fs" } diff --git a/src/streaming_planner/src/analysis/aggregate_rewriter.rs b/src/streaming_planner/src/analysis/aggregate_rewriter.rs new file mode 100644 index 00000000..22dcb03c --- /dev/null +++ b/src/streaming_planner/src/analysis/aggregate_rewriter.rs @@ -0,0 +1,279 @@ +// 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 datafusion::common::tree_node::{Transformed, TreeNodeRewriter}; +use datafusion::common::{DFSchema, DataFusionError, Result, not_impl_err, plan_err}; +use datafusion::functions_aggregate::expr_fn::max; +use datafusion::logical_expr::{Aggregate, Expr, Extension, LogicalPlan, Projection}; +use datafusion::prelude::col; +use std::sync::Arc; + +use crate::analysis::streaming_window_analzer::StreamingWindowAnalzer; +use crate::logical_node::aggregate::StreamWindowAggregateNode; +use crate::logical_node::key_calculation::{KeyExtractionNode, KeyExtractionStrategy}; +use crate::logical_node::updating_aggregate::ContinuousAggregateNode; +use crate::schema::StreamSchemaProvider; +use crate::types::{ + QualifiedField, TIMESTAMP_FIELD, WindowBehavior, WindowType, build_df_schema_with_metadata, + extract_qualified_fields, extract_window_type, +}; + +/// AggregateRewriter transforms batch DataFusion aggregates into streaming stateful operators. +/// It handles windowing (Tumble/Hop/Session), watermarks, and continuous updating aggregates. +pub struct AggregateRewriter<'a> { + pub schema_provider: &'a StreamSchemaProvider, +} + +impl TreeNodeRewriter for AggregateRewriter<'_> { + type Node = LogicalPlan; + + fn f_up(&mut self, node: Self::Node) -> Result> { + let LogicalPlan::Aggregate(mut agg) = node else { + return Ok(Transformed::no(node)); + }; + + // 1. Identify windowing functions (e.g., tumble, hop) in GROUP BY. + let mut window_exprs: Vec<_> = agg + .group_expr + .iter() + .enumerate() + .filter_map(|(i, e)| { + extract_window_type(e) + .map(|opt| opt.map(|w| (i, w))) + .transpose() + }) + .collect::>>()?; + + if window_exprs.len() > 1 { + return not_impl_err!("Streaming aggregates support at most one window expression"); + } + + // 2. Prepare internal metadata for Key-based distribution. + let mut key_fields: Vec = extract_qualified_fields(&agg.schema) + .iter() + .take(agg.group_expr.len()) + .map(|f| { + QualifiedField::new( + f.qualifier().cloned(), + format!("_key_{}", f.name()), + f.data_type().clone(), + f.is_nullable(), + ) + }) + .collect(); + + // 3. Dispatch to ContinuousAggregateNode (UpdatingAggregate) if no windowing is detected. + let input_window = StreamingWindowAnalzer::get_window(&agg.input)?; + if window_exprs.is_empty() && input_window.is_none() { + return self.rewrite_as_continuous_updating_aggregate( + agg.input, + key_fields, + agg.group_expr, + agg.aggr_expr, + agg.schema, + ); + } + + // 4. Resolve Windowing Strategy (InData vs FromOperator). + let behavior = self.resolve_window_context( + &agg.input, + &mut agg.group_expr, + &agg.schema, + &mut window_exprs, + )?; + + // Adjust keys if windowing is handled by the operator. + if let WindowBehavior::FromOperator { window_index, .. } = &behavior { + key_fields.remove(*window_index); + } + + let key_count = key_fields.len(); + let keyed_input = + self.build_keyed_input(agg.input.clone(), &agg.group_expr, &key_fields)?; + + // 5. Build the final StreamWindowAggregateNode for the physical planner. + let mut internal_fields = extract_qualified_fields(&agg.schema); + if let WindowBehavior::FromOperator { window_index, .. } = &behavior { + internal_fields.remove(*window_index); + } + let internal_schema = Arc::new(build_df_schema_with_metadata( + &internal_fields, + agg.schema.metadata().clone(), + )?); + + let rewritten_agg = Aggregate::try_new_with_schema( + Arc::new(keyed_input), + agg.group_expr, + agg.aggr_expr, + internal_schema, + )?; + + let extension = StreamWindowAggregateNode::try_new( + behavior, + LogicalPlan::Aggregate(rewritten_agg), + (0..key_count).collect(), + )?; + + Ok(Transformed::yes(LogicalPlan::Extension(Extension { + node: Arc::new(extension), + }))) + } +} + +impl<'a> AggregateRewriter<'a> { + pub fn new(schema_provider: &'a StreamSchemaProvider) -> Self { + Self { schema_provider } + } + + /// [Internal] Builds the physical Key Calculation layer required for distributed Shuffling. + /// This wraps the input in a Projection and a KeyExtractionNode. + fn build_keyed_input( + &self, + input: Arc, + group_expr: &[Expr], + key_fields: &[QualifiedField], + ) -> Result { + let key_count = group_expr.len(); + let mut projection_fields = key_fields.to_vec(); + projection_fields.extend(extract_qualified_fields(input.schema())); + + let key_schema = Arc::new(build_df_schema_with_metadata( + &projection_fields, + input.schema().metadata().clone(), + )?); + + // Map group expressions to '_key_' aliases while passing through all original columns. + let mut exprs: Vec<_> = group_expr + .iter() + .zip(key_fields.iter()) + .map(|(expr, f)| expr.clone().alias(f.name().to_string())) + .collect(); + + exprs.extend( + extract_qualified_fields(input.schema()) + .iter() + .map(|f| Expr::Column(f.qualified_column())), + ); + + let projection = + LogicalPlan::Projection(Projection::try_new_with_schema(exprs, input, key_schema)?); + + Ok(LogicalPlan::Extension(Extension { + node: Arc::new(KeyExtractionNode::new( + projection, + KeyExtractionStrategy::ColumnIndices((0..key_count).collect()), + )), + })) + } + + /// [Strategy] Rewrites standard GROUP BY into a ContinuousAggregateNode with retraction semantics. + /// Injected max(_timestamp) ensures the streaming pulse (Watermark) continues to propagate. + fn rewrite_as_continuous_updating_aggregate( + &self, + input: Arc, + key_fields: Vec, + group_expr: Vec, + mut aggr_expr: Vec, + schema: Arc, + ) -> Result> { + let key_count = key_fields.len(); + let keyed_input = self.build_keyed_input(input, &group_expr, &key_fields)?; + + // Ensure the updating stream maintains time awareness. + let timestamp_col = keyed_input + .schema() + .qualified_field_with_unqualified_name(TIMESTAMP_FIELD) + .map_err(|_| { + DataFusionError::Plan( + "Required _timestamp field missing for updating aggregate".to_string(), + ) + })?; + + let timestamp_field: QualifiedField = timestamp_col.into(); + aggr_expr.push(max(col(timestamp_field.qualified_column())).alias(TIMESTAMP_FIELD)); + + let mut output_fields = extract_qualified_fields(&schema); + output_fields.push(timestamp_field); + + let output_schema = Arc::new(build_df_schema_with_metadata( + &output_fields, + schema.metadata().clone(), + )?); + + let base_aggregate = Aggregate::try_new_with_schema( + Arc::new(keyed_input), + group_expr, + aggr_expr, + output_schema, + )?; + + let continuous_node = ContinuousAggregateNode::try_new( + LogicalPlan::Aggregate(base_aggregate), + (0..key_count).collect(), + None, + self.schema_provider.planning_options.ttl, + )?; + + Ok(Transformed::yes(LogicalPlan::Extension(Extension { + node: Arc::new(continuous_node), + }))) + } + + /// [Strategy] Reconciles window definitions between the input stream and the current GROUP BY. + fn resolve_window_context( + &self, + input: &LogicalPlan, + group_expr: &mut Vec, + schema: &DFSchema, + window_expr_info: &mut Vec<(usize, WindowType)>, + ) -> Result { + let mut visitor = StreamingWindowAnalzer::default(); + input.visit_with_subqueries(&mut visitor)?; + + let input_window = visitor.window; + let has_group_window = !window_expr_info.is_empty(); + + match (input_window, has_group_window) { + (Some(i_win), true) => { + let (idx, g_win) = window_expr_info.pop().unwrap(); + if i_win != g_win { + return plan_err!("Inconsistent windowing detected"); + } + + if let Some(field) = visitor.fields.iter().next() { + group_expr[idx] = Expr::Column(field.qualified_column()); + Ok(WindowBehavior::InData) + } else { + group_expr.remove(idx); + Ok(WindowBehavior::FromOperator { + window: i_win, + window_field: schema.qualified_field(idx).into(), + window_index: idx, + is_nested: true, + }) + } + } + (None, true) => { + let (idx, g_win) = window_expr_info.pop().unwrap(); + group_expr.remove(idx); + Ok(WindowBehavior::FromOperator { + window: g_win, + window_field: schema.qualified_field(idx).into(), + window_index: idx, + is_nested: false, + }) + } + (Some(_), false) => Ok(WindowBehavior::InData), + _ => unreachable!("Handled by updating path"), + } + } +} diff --git a/src/streaming_planner/src/analysis/async_udf_rewriter.rs b/src/streaming_planner/src/analysis/async_udf_rewriter.rs new file mode 100644 index 00000000..80f3828d --- /dev/null +++ b/src/streaming_planner/src/analysis/async_udf_rewriter.rs @@ -0,0 +1,133 @@ +// 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::common::constants::sql_field; +use crate::logical_node::AsyncFunctionExecutionNode; +use crate::logical_node::remote_table::RemoteTableBoundaryNode; +use crate::schema::StreamSchemaProvider; +use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRewriter}; +use datafusion::common::{Column, Result as DFResult, TableReference, plan_err}; +use datafusion::logical_expr::expr::ScalarFunction; +use datafusion::logical_expr::{Expr, Extension, LogicalPlan}; +use std::sync::Arc; +use std::time::Duration; + +type AsyncSplitResult = (String, AsyncOptions, Vec); + +#[derive(Debug, Clone, Copy)] +pub struct AsyncOptions { + pub ordered: bool, + pub max_concurrency: usize, + pub timeout: Duration, +} + +pub struct AsyncUdfRewriter<'a> { + provider: &'a StreamSchemaProvider, +} + +impl<'a> AsyncUdfRewriter<'a> { + pub fn new(provider: &'a StreamSchemaProvider) -> Self { + Self { provider } + } + + fn split_async( + expr: Expr, + provider: &StreamSchemaProvider, + ) -> DFResult<(Expr, Option)> { + let mut found: Option<(String, AsyncOptions, Vec)> = None; + let expr = expr.transform_up(|e| { + if let Expr::ScalarFunction(ScalarFunction { func: udf, args }) = &e + && let Some(opts) = provider.get_async_udf_options(udf.name()) + { + if found + .replace((udf.name().to_string(), opts, args.clone())) + .is_some() + { + return plan_err!( + "multiple async calls in the same expression, which is not allowed" + ); + } + return Ok(Transformed::yes(Expr::Column(Column::new_unqualified( + sql_field::ASYNC_RESULT, + )))); + } + Ok(Transformed::no(e)) + })?; + + Ok((expr.data, found)) + } +} + +impl TreeNodeRewriter for AsyncUdfRewriter<'_> { + type Node = LogicalPlan; + + fn f_up(&mut self, node: Self::Node) -> DFResult> { + let LogicalPlan::Projection(mut projection) = node else { + for e in node.expressions() { + if let (_, Some((udf, _, _))) = Self::split_async(e.clone(), self.provider)? { + return plan_err!( + "async UDFs are only supported in projections, but {udf} was called in another context" + ); + } + } + return Ok(Transformed::no(node)); + }; + + let mut args = None; + for e in projection.expr.iter_mut() { + let (new_e, Some(udf)) = Self::split_async(e.clone(), self.provider)? else { + continue; + }; + if let Some((prev, _, _)) = args.replace(udf) { + return plan_err!( + "Projection contains multiple async UDFs, which is not supported \ + \n(hint: two async UDF calls, {} and {}, appear in the same SELECT statement)", + prev, + args.unwrap().0 + ); + } + *e = new_e; + } + + let Some((name, opts, arg_exprs)) = args else { + return Ok(Transformed::no(LogicalPlan::Projection(projection))); + }; + let udf = self.provider.dylib_udfs.get(&name).unwrap().clone(); + + let input = if matches!(*projection.input, LogicalPlan::Projection(..)) { + Arc::new(LogicalPlan::Extension(Extension { + node: Arc::new(RemoteTableBoundaryNode { + upstream_plan: (*projection.input).clone(), + table_identifier: TableReference::bare("subquery_projection"), + resolved_schema: projection.input.schema().clone(), + requires_materialization: false, + }), + })) + } else { + projection.input + }; + + Ok(Transformed::yes(LogicalPlan::Extension(Extension { + node: Arc::new(AsyncFunctionExecutionNode { + upstream_plan: input, + operator_name: name, + function_config: udf, + invocation_args: arg_exprs, + result_projections: projection.expr, + preserve_ordering: opts.ordered, + concurrency_limit: opts.max_concurrency, + execution_timeout: opts.timeout, + resolved_schema: projection.schema, + }), + }))) + } +} diff --git a/src/streaming_planner/src/analysis/join_rewriter.rs b/src/streaming_planner/src/analysis/join_rewriter.rs new file mode 100644 index 00000000..07bc0006 --- /dev/null +++ b/src/streaming_planner/src/analysis/join_rewriter.rs @@ -0,0 +1,234 @@ +// 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::analysis::streaming_window_analzer::StreamingWindowAnalzer; +use crate::common::TIMESTAMP_FIELD; +use crate::common::constants::mem_exec_join_side; +use crate::logical_node::join::StreamingJoinNode; +use crate::logical_node::key_calculation::KeyExtractionNode; +use crate::schema::StreamSchemaProvider; +use crate::types::{WindowType, build_df_schema_with_metadata, extract_qualified_fields}; +use datafusion::common::tree_node::{Transformed, TreeNodeRewriter}; +use datafusion::common::{ + JoinConstraint, JoinType, Result, ScalarValue, TableReference, not_impl_err, plan_err, +}; +use datafusion::logical_expr::{ + self, BinaryExpr, Case, Expr, Extension, Join, LogicalPlan, Projection, build_join_schema, +}; +use datafusion::prelude::coalesce; +use std::sync::Arc; + +/// JoinRewriter handles the transformation of standard SQL joins into streaming-capable joins. +/// It manages stateful "Updating Joins" and time-aligned "Instant Joins". +pub struct JoinRewriter<'a> { + pub schema_provider: &'a StreamSchemaProvider, +} + +impl<'a> JoinRewriter<'a> { + pub fn new(schema_provider: &'a StreamSchemaProvider) -> Self { + Self { schema_provider } + } + + /// [Validation] Ensures left and right streams have compatible windowing strategies. + fn validate_join_windows(&self, join: &Join) -> Result { + let left_win = StreamingWindowAnalzer::get_window(&join.left)?; + let right_win = StreamingWindowAnalzer::get_window(&join.right)?; + + match (left_win, right_win) { + (None, None) => { + if join.join_type == JoinType::Inner { + Ok(false) // Standard Updating Join (Inner) + } else { + plan_err!( + "Non-inner joins (e.g., LEFT/RIGHT) require windowing to bound state." + ) + } + } + (Some(l), Some(r)) => { + if l != r { + return plan_err!( + "Join window mismatch: left={:?}, right={:?}. Windows must match exactly.", + l, + r + ); + } + if let WindowType::Session { .. } = l { + return plan_err!( + "Session windows are currently not supported in streaming joins." + ); + } + Ok(true) // Instant Windowed Join + } + _ => plan_err!( + "Mixed windowing detected. Both sides of a join must be either windowed or non-windowed." + ), + } + } + + /// [Internal] Wraps a join input in a key-extraction layer to facilitate shuffle / key-by distribution. + fn build_keyed_side( + &self, + input: Arc, + keys: Vec, + side: &str, + ) -> Result { + let key_count = keys.len(); + + let projection_exprs = keys + .into_iter() + .enumerate() + .map(|(i, e)| { + e.alias_qualified(Some(TableReference::bare("_stream")), format!("_key_{i}")) + }) + .chain( + extract_qualified_fields(input.schema()) + .iter() + .map(|f| Expr::Column(f.qualified_column())), + ) + .collect(); + + let projection = Projection::try_new(projection_exprs, input)?; + let key_ext = KeyExtractionNode::try_new_with_projection( + LogicalPlan::Projection(projection), + (0..key_count).collect(), + side.to_string(), + )?; + + Ok(LogicalPlan::Extension(Extension { + node: Arc::new(key_ext), + })) + } + + /// [Strategy] Resolves the output timestamp of the join. + /// Streaming joins must output the 'max' of the two input timestamps to ensure Watermark progression. + fn apply_timestamp_resolution(&self, join_plan: LogicalPlan) -> Result { + let schema = join_plan.schema(); + let all_fields = extract_qualified_fields(schema); + + let timestamp_fields: Vec<_> = all_fields + .iter() + .filter(|f| f.name() == "_timestamp") + .cloned() + .collect(); + + if timestamp_fields.len() != 2 { + return plan_err!( + "Streaming join requires exactly two input timestamp fields to resolve output time." + ); + } + + // Project all fields except the two raw timestamps + let mut exprs: Vec<_> = all_fields + .iter() + .filter(|f| f.name() != "_timestamp") + .map(|f| Expr::Column(f.qualified_column())) + .collect(); + + // Calculate: GREATEST(left._timestamp, right._timestamp) + let left_ts = Expr::Column(timestamp_fields[0].qualified_column()); + let right_ts = Expr::Column(timestamp_fields[1].qualified_column()); + + let max_ts_expr = Expr::Case(Case { + expr: Some(Box::new(Expr::BinaryExpr(BinaryExpr { + left: Box::new(left_ts.clone()), + op: logical_expr::Operator::GtEq, + right: Box::new(right_ts.clone()), + }))), + when_then_expr: vec![ + ( + Box::new(Expr::Literal(ScalarValue::Boolean(Some(true)), None)), + Box::new(left_ts.clone()), + ), + ( + Box::new(Expr::Literal(ScalarValue::Boolean(Some(false)), None)), + Box::new(right_ts.clone()), + ), + ], + else_expr: Some(Box::new(coalesce(vec![left_ts, right_ts]))), + }) + .alias(TIMESTAMP_FIELD); + + exprs.push(max_ts_expr); + + let out_fields: Vec<_> = all_fields + .iter() + .filter(|f| f.name() != "_timestamp") + .cloned() + .chain(std::iter::once(timestamp_fields[0].clone())) + .collect(); + + let out_schema = Arc::new(build_df_schema_with_metadata( + &out_fields, + schema.metadata().clone(), + )?); + + Ok(LogicalPlan::Projection(Projection::try_new_with_schema( + exprs, + Arc::new(join_plan), + out_schema, + )?)) + } +} + +impl TreeNodeRewriter for JoinRewriter<'_> { + type Node = LogicalPlan; + + fn f_up(&mut self, node: Self::Node) -> Result> { + let LogicalPlan::Join(join) = node else { + return Ok(Transformed::no(node)); + }; + + // 1. Validate Streaming Context + let is_instant = self.validate_join_windows(&join)?; + if join.join_constraint != JoinConstraint::On { + return not_impl_err!("Only 'ON' join constraints are supported in streaming SQL."); + } + if join.on.is_empty() && !is_instant { + return plan_err!("Updating joins require at least one equality condition (Equijoin)."); + } + + // 2. Prepare Keyed Inputs for Shuffle + let (left_on, right_on): (Vec<_>, Vec<_>) = join.on.clone().into_iter().unzip(); + let keyed_left = self.build_keyed_side(join.left, left_on, mem_exec_join_side::LEFT)?; + let keyed_right = self.build_keyed_side(join.right, right_on, mem_exec_join_side::RIGHT)?; + + // 3. Assemble Rewritten Join Node + let join_schema = Arc::new(build_join_schema( + keyed_left.schema(), + keyed_right.schema(), + &join.join_type, + )?); + let rewritten_join = LogicalPlan::Join(Join { + left: Arc::new(keyed_left), + right: Arc::new(keyed_right), + on: join.on, + filter: join.filter, + join_type: join.join_type, + join_constraint: JoinConstraint::On, + schema: join_schema, + null_equals_null: false, + }); + + // 4. Resolve Output Watermark (Timestamp Projection) + let plan_with_timestamp = self.apply_timestamp_resolution(rewritten_join)?; + + // 5. Wrap in StreamingJoinNode for physical planning + let state_retention_ttl = + (!is_instant).then_some(self.schema_provider.planning_options.ttl); + let extension = + StreamingJoinNode::new(plan_with_timestamp, is_instant, state_retention_ttl); + + Ok(Transformed::yes(LogicalPlan::Extension(Extension { + node: Arc::new(extension), + }))) + } +} diff --git a/src/streaming_planner/src/analysis/mod.rs b/src/streaming_planner/src/analysis/mod.rs new file mode 100644 index 00000000..7ac1d4e8 --- /dev/null +++ b/src/streaming_planner/src/analysis/mod.rs @@ -0,0 +1,214 @@ +// 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. + +#![allow(clippy::new_without_default)] + +pub mod aggregate_rewriter; +pub mod join_rewriter; +pub mod row_time_rewriter; +pub mod stream_rewriter; +pub mod streaming_window_analzer; +pub mod window_function_rewriter; + +pub mod async_udf_rewriter; +pub mod sink_input_rewriter; +pub mod source_metadata_visitor; +pub mod source_rewriter; +pub mod time_window; +pub mod unnest_rewriter; + +pub use async_udf_rewriter::AsyncOptions; +pub use sink_input_rewriter::SinkInputRewriter; +pub use time_window::{TimeWindowNullCheckRemover, TimeWindowUdfChecker}; +pub use unnest_rewriter::UNNESTED_COL; + +pub use crate::schema::schema_provider::StreamSchemaProvider; + +use std::collections::HashMap; +use std::sync::Arc; + +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::common::{Result, plan_err}; +use datafusion::error::DataFusionError; +use datafusion::logical_expr::{Extension, LogicalPlan, UserDefinedLogicalNodeCore}; +use tracing::{debug, info, instrument}; + +use crate::logical_node::StreamingOperatorBlueprint; +use crate::logical_node::key_calculation::{KeyExtractionNode, KeyExtractionStrategy}; +use crate::logical_node::projection::StreamProjectionNode; +use crate::logical_node::sink::StreamEgressNode; +use crate::logical_planner::planner::NamedNode; + +fn duration_from_sql_expr( + expr: &datafusion::sql::sqlparser::ast::Expr, +) -> Result { + use datafusion::sql::sqlparser::ast::Expr as SqlExpr; + use datafusion::sql::sqlparser::ast::Value as SqlValue; + use datafusion::sql::sqlparser::ast::ValueWithSpan; + + match expr { + SqlExpr::Interval(interval) => { + let value_str = match interval.value.as_ref() { + SqlExpr::Value(ValueWithSpan { + value: SqlValue::SingleQuotedString(s), + .. + }) => s.clone(), + other => return plan_err!("expected interval string literal, found {other}"), + }; + + parse_interval_to_duration(&value_str) + } + SqlExpr::Value(ValueWithSpan { + value: SqlValue::SingleQuotedString(s), + .. + }) => parse_interval_to_duration(s), + other => plan_err!("expected an interval expression, found {other}"), + } +} + +fn parse_interval_to_duration(s: &str) -> Result { + let parts: Vec<&str> = s.split_whitespace().collect(); + if parts.len() != 2 { + return plan_err!("invalid interval string '{s}'; expected ' '"); + } + let value: u64 = parts[0] + .parse() + .map_err(|_| DataFusionError::Plan(format!("invalid interval number: {}", parts[0])))?; + match parts[1].to_lowercase().as_str() { + "second" | "seconds" | "s" => Ok(std::time::Duration::from_secs(value)), + "minute" | "minutes" | "min" => Ok(std::time::Duration::from_secs(value * 60)), + "hour" | "hours" | "h" => Ok(std::time::Duration::from_secs(value * 3600)), + "day" | "days" | "d" => Ok(std::time::Duration::from_secs(value * 86400)), + unit => plan_err!("unsupported interval unit '{unit}'"), + } +} + +fn build_sink_inputs(extensions: &[LogicalPlan]) -> HashMap> { + let mut sink_inputs = HashMap::>::new(); + for extension in extensions.iter() { + if let LogicalPlan::Extension(ext) = extension + && let Some(sink_node) = ext.node.as_any().downcast_ref::() + && let Some(named_node) = sink_node.operator_identity() + { + let inputs = sink_node + .inputs() + .into_iter() + .cloned() + .collect::>(); + sink_inputs.entry(named_node).or_default().extend(inputs); + } + } + sink_inputs +} + +pub fn maybe_add_key_extension_to_sink(plan: LogicalPlan) -> Result { + let LogicalPlan::Extension(ref ext) = plan else { + return Ok(plan); + }; + + let Some(sink) = ext.node.as_any().downcast_ref::() else { + return Ok(plan); + }; + + let Some(partition_exprs) = sink.destination_table.partition_exprs() else { + return Ok(plan); + }; + + if partition_exprs.is_empty() { + return Ok(plan); + } + + let inputs = plan + .inputs() + .into_iter() + .map(|input| { + Ok(LogicalPlan::Extension(Extension { + node: Arc::new(KeyExtractionNode { + operator_label: Some("key-calc-partition".to_string()), + resolved_schema: input.schema().clone(), + upstream_plan: input.clone(), + extraction_strategy: KeyExtractionStrategy::CalculatedExpressions( + partition_exprs.clone(), + ), + }), + })) + }) + .collect::>()?; + + use datafusion::prelude::col; + let unkey = LogicalPlan::Extension(Extension { + node: Arc::new( + StreamProjectionNode::try_new( + inputs, + Some("unkey".to_string()), + sink.schema().iter().map(|(_, f)| col(f.name())).collect(), + )? + .with_shuffle_routing(), + ), + }); + + let node = sink.with_exprs_and_inputs(vec![], vec![unkey])?; + Ok(LogicalPlan::Extension(Extension { + node: Arc::new(node), + })) +} + +pub fn rewrite_sinks(extensions: Vec) -> Result> { + let mut sink_inputs = build_sink_inputs(&extensions); + let mut new_extensions = vec![]; + for extension in extensions { + let mut rewriter = SinkInputRewriter::new(&mut sink_inputs); + let result = extension.rewrite(&mut rewriter)?; + if !rewriter.was_removed { + new_extensions.push(result.data); + } + } + + new_extensions + .into_iter() + .map(maybe_add_key_extension_to_sink) + .collect() +} + +/// Entry point for transforming a standard DataFusion LogicalPlan into a +/// Streaming-aware LogicalPlan. +/// +/// This function coordinates multiple rewriting passes and ensures the +/// resulting plan satisfies streaming constraints. +#[instrument(skip_all, level = "debug")] +pub fn rewrite_plan( + plan: LogicalPlan, + schema_provider: &StreamSchemaProvider, +) -> Result { + info!("Starting streaming plan rewrite pipeline"); + + let Transformed { data: plan, .. } = + plan.rewrite_with_subqueries(&mut source_rewriter::SourceRewriter::new(schema_provider))?; + + let mut rewriter = stream_rewriter::StreamRewriter::new(schema_provider); + let Transformed { + data: rewritten_plan, + .. + } = plan.rewrite_with_subqueries(&mut rewriter)?; + + rewritten_plan.visit_with_subqueries(&mut TimeWindowUdfChecker {})?; + + if cfg!(debug_assertions) { + debug!( + "Streaming logical plan graphviz:\n{}", + rewritten_plan.display_graphviz() + ); + } + + info!("Streaming plan rewrite completed successfully"); + Ok(rewritten_plan) +} diff --git a/src/streaming_planner/src/analysis/row_time_rewriter.rs b/src/streaming_planner/src/analysis/row_time_rewriter.rs new file mode 100644 index 00000000..aa558696 --- /dev/null +++ b/src/streaming_planner/src/analysis/row_time_rewriter.rs @@ -0,0 +1,49 @@ +// 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 datafusion::common::tree_node::{Transformed, TreeNodeRewriter}; +use datafusion::common::{Column, Result as DFResult}; +use datafusion::logical_expr::Expr; + +use crate::common::constants::planning_placeholder_udf; +use crate::types::TIMESTAMP_FIELD; + +/// Replaces the virtual `row_time()` scalar function with a physical reference to `_timestamp`. +/// +/// This is a critical mapping step that allows users to use a friendly SQL function +/// while the engine operates on the mandatory internal streaming timestamp. +pub struct RowTimeRewriter; + +impl TreeNodeRewriter for RowTimeRewriter { + type Node = Expr; + + fn f_down(&mut self, node: Self::Node) -> DFResult> { + // Use pattern matching to identify the `row_time` scalar function. + if let Expr::ScalarFunction(func) = &node + && func.name() == planning_placeholder_udf::ROW_TIME + { + // Map the virtual function to the physical internal timestamp column. + // We use .alias() to preserve the original name "row_time()" in the output schema, + // ensuring that user-facing column names do not change unexpectedly. + let physical_col = Expr::Column(Column { + relation: None, + name: TIMESTAMP_FIELD.to_string(), + spans: Default::default(), + }) + .alias("row_time()"); + + return Ok(Transformed::yes(physical_col)); + } + + Ok(Transformed::no(node)) + } +} diff --git a/src/streaming_planner/src/analysis/sink_input_rewriter.rs b/src/streaming_planner/src/analysis/sink_input_rewriter.rs new file mode 100644 index 00000000..9f8fdcb7 --- /dev/null +++ b/src/streaming_planner/src/analysis/sink_input_rewriter.rs @@ -0,0 +1,57 @@ +// 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::logical_node::StreamingOperatorBlueprint; +use crate::logical_node::sink::StreamEgressNode; +use crate::logical_planner::planner::NamedNode; +use datafusion::common::Result as DFResult; +use datafusion::common::tree_node::{Transformed, TreeNodeRecursion, TreeNodeRewriter}; +use datafusion::logical_expr::{Extension, LogicalPlan, UserDefinedLogicalNodeCore}; +use std::collections::HashMap; +use std::sync::Arc; + +type SinkInputs = HashMap>; + +/// Merges inputs for sinks with the same name to avoid duplicate sinks in the plan. +pub struct SinkInputRewriter<'a> { + sink_inputs: &'a mut SinkInputs, + pub was_removed: bool, +} + +impl<'a> SinkInputRewriter<'a> { + pub fn new(sink_inputs: &'a mut SinkInputs) -> Self { + Self { + sink_inputs, + was_removed: false, + } + } +} + +impl TreeNodeRewriter for SinkInputRewriter<'_> { + type Node = LogicalPlan; + + fn f_down(&mut self, node: Self::Node) -> DFResult> { + if let LogicalPlan::Extension(extension) = &node + && let Some(sink_node) = extension.node.as_any().downcast_ref::() + && let Some(named_node) = sink_node.operator_identity() + { + if let Some(inputs) = self.sink_inputs.remove(&named_node) { + let new_node = LogicalPlan::Extension(Extension { + node: Arc::new(sink_node.with_exprs_and_inputs(vec![], inputs)?), + }); + return Ok(Transformed::new(new_node, true, TreeNodeRecursion::Jump)); + } + self.was_removed = true; + } + Ok(Transformed::no(node)) + } +} diff --git a/src/streaming_planner/src/analysis/source_metadata_visitor.rs b/src/streaming_planner/src/analysis/source_metadata_visitor.rs new file mode 100644 index 00000000..ea2821b9 --- /dev/null +++ b/src/streaming_planner/src/analysis/source_metadata_visitor.rs @@ -0,0 +1,73 @@ +// 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::logical_node::sink::{STREAM_EGRESS_NODE_NAME, StreamEgressNode}; +use crate::logical_node::table_source::{STREAM_INGESTION_NODE_NAME, StreamIngestionNode}; +use crate::schema::StreamSchemaProvider; +use datafusion::common::Result as DFResult; +use datafusion::common::tree_node::{TreeNodeRecursion, TreeNodeVisitor}; +use datafusion::logical_expr::{Extension, LogicalPlan}; +use std::collections::HashSet; + +/// Collects connection IDs from source and sink nodes in the logical plan. +pub struct SourceMetadataVisitor<'a> { + schema_provider: &'a StreamSchemaProvider, + pub connection_ids: HashSet, +} + +impl<'a> SourceMetadataVisitor<'a> { + pub fn new(schema_provider: &'a StreamSchemaProvider) -> Self { + Self { + schema_provider, + connection_ids: HashSet::new(), + } + } + + fn get_connection_id(&self, node: &LogicalPlan) -> Option { + let LogicalPlan::Extension(Extension { node }) = node else { + return None; + }; + + let table_name = match node.name() { + name if name == STREAM_INGESTION_NODE_NAME => { + let ext = node.as_any().downcast_ref::()?; + ext.source_identifier.to_string() + } + name if name == STREAM_EGRESS_NODE_NAME => { + let ext = node.as_any().downcast_ref::()?; + ext.target_identifier.to_string() + } + _ => return None, + }; + + let table = self.schema_provider.get_catalog_table(&table_name)?; + match table { + crate::schema::table::CatalogEntity::ExternalConnector(b) => match b.as_ref() { + crate::schema::catalog::ExternalTable::Source(t) => t.registry_id, + crate::schema::catalog::ExternalTable::Lookup(t) => t.registry_id, + _ => None, + }, + _ => None, + } + } +} + +impl TreeNodeVisitor<'_> for SourceMetadataVisitor<'_> { + type Node = LogicalPlan; + + fn f_down(&mut self, node: &Self::Node) -> DFResult { + if let Some(id) = self.get_connection_id(node) { + self.connection_ids.insert(id); + } + Ok(TreeNodeRecursion::Continue) + } +} diff --git a/src/streaming_planner/src/analysis/source_rewriter.rs b/src/streaming_planner/src/analysis/source_rewriter.rs new file mode 100644 index 00000000..e1a7c75a --- /dev/null +++ b/src/streaming_planner/src/analysis/source_rewriter.rs @@ -0,0 +1,305 @@ +// 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::sync::Arc; +use std::time::Duration; + +use datafusion::common::ScalarValue; +use datafusion::common::tree_node::{Transformed, TreeNodeRewriter}; +use datafusion::common::{Column, DataFusionError, Result as DFResult, TableReference, plan_err}; +use datafusion::logical_expr::{ + self, BinaryExpr, Expr, Extension, LogicalPlan, Projection, TableScan, +}; + +use crate::common::UPDATING_META_FIELD; +use crate::logical_node::debezium::UnrollDebeziumPayloadNode; +use crate::logical_node::remote_table::RemoteTableBoundaryNode; +use crate::logical_node::table_source::StreamIngestionNode; +use crate::logical_node::watermark_node::EventTimeWatermarkNode; +use crate::schema::ColumnDescriptor; +use crate::schema::StreamSchemaProvider; +use crate::schema::catalog::{ExternalTable, SourceTable}; +use crate::schema::table::CatalogEntity; +use crate::types::TIMESTAMP_FIELD; + +/// Rewrites table scans: projections are lifted out of scans into a dedicated projection node +/// (including virtual fields), using a connector table-source extension instead of a bare +/// `TableScan`, optionally with Debezium unrolling for updating sources, then remote boundary and +/// watermark. +pub struct SourceRewriter<'a> { + pub(crate) schema_provider: &'a StreamSchemaProvider, +} + +impl<'a> SourceRewriter<'a> { + pub fn new(schema_provider: &'a StreamSchemaProvider) -> Self { + Self { schema_provider } + } +} + +impl SourceRewriter<'_> { + fn projection_expr_for_column(col: &ColumnDescriptor, qualifier: &TableReference) -> Expr { + if let Some(logic) = col.computation_logic() { + logic.clone().alias_qualified( + Some(qualifier.clone()), + col.arrow_field().name().to_string(), + ) + } else { + Expr::Column(Column { + relation: Some(qualifier.clone()), + name: col.arrow_field().name().to_string(), + spans: Default::default(), + }) + } + } + + fn watermark_expression(table: &SourceTable) -> DFResult { + match table.temporal_config.watermark_strategy_column.clone() { + Some(watermark_field) => table + .schema_specs + .iter() + .find_map(|c| { + if c.arrow_field().name() == watermark_field.as_str() { + return if let Some(expr) = c.computation_logic() { + Some(expr.clone()) + } else { + Some(Expr::Column(Column { + relation: None, + name: c.arrow_field().name().to_string(), + spans: Default::default(), + })) + }; + } + None + }) + .ok_or_else(|| { + DataFusionError::Plan(format!("Watermark field {watermark_field} not found")) + }), + None => Ok(Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::Column(Column { + relation: None, + name: TIMESTAMP_FIELD.to_string(), + spans: Default::default(), + })), + op: logical_expr::Operator::Minus, + right: Box::new(Expr::Literal( + ScalarValue::DurationNanosecond(Some(Duration::from_secs(1).as_nanos() as i64)), + None, + )), + })), + } + } + + fn projection_expressions( + table: &SourceTable, + qualifier: &TableReference, + projection: &Option>, + ) -> DFResult> { + let mut expressions: Vec = table + .schema_specs + .iter() + .map(|col| Self::projection_expr_for_column(col, qualifier)) + .collect(); + + if let Some(proj) = projection { + expressions = proj.iter().map(|i| expressions[*i].clone()).collect(); + } + + if let Some(event_time_field) = table.temporal_config.event_column.clone() { + let expr = table + .schema_specs + .iter() + .find_map(|c| { + if c.arrow_field().name() == event_time_field.as_str() { + return Some(Self::projection_expr_for_column(c, qualifier)); + } + None + }) + .ok_or_else(|| { + DataFusionError::Plan(format!("Event time field {event_time_field} not found")) + })?; + + expressions + .push(expr.alias_qualified(Some(qualifier.clone()), TIMESTAMP_FIELD.to_string())); + } else { + let has_ts = table + .schema_specs + .iter() + .any(|c| c.arrow_field().name() == TIMESTAMP_FIELD); + if !has_ts { + return plan_err!( + "Connector table '{}' has no `{}` column; declare WATERMARK FOR AS ... in CREATE TABLE", + table.table_identifier, + TIMESTAMP_FIELD + ); + } + expressions.push(Expr::Column(Column::new( + Some(qualifier.clone()), + TIMESTAMP_FIELD, + ))); + } + + if table.is_updating() { + expressions.push(Expr::Column(Column::new( + Some(qualifier.clone()), + UPDATING_META_FIELD, + ))); + } + + Ok(expressions) + } + + /// Connector path: `StreamIngestionNode` (table source) → optional `UnrollDebeziumPayloadNode` + /// → `Projection`, mirroring Arroyo `TableSourceExtension` + Debezium unroll + projection. + fn projection(&self, table_scan: &TableScan, table: &SourceTable) -> DFResult { + let qualifier = table_scan.table_name.clone(); + + let table_source = LogicalPlan::Extension(Extension { + node: Arc::new(StreamIngestionNode::try_new( + qualifier.clone(), + table.clone(), + )?), + }); + + let (projection_input, scan_projection) = if table.is_updating() { + if table.key_constraints.is_empty() { + return plan_err!( + "Updating connector table `{}` requires at least one PRIMARY KEY for CDC unrolling", + table.table_identifier + ); + } + let unrolled = LogicalPlan::Extension(Extension { + node: Arc::new(UnrollDebeziumPayloadNode::try_new( + table_source, + Arc::new(table.key_constraints.clone()), + )?), + }); + (unrolled, None) + } else { + (table_source, table_scan.projection.clone()) + }; + + Ok(LogicalPlan::Projection(Projection::try_new( + Self::projection_expressions(table, &qualifier, &scan_projection)?, + Arc::new(projection_input), + )?)) + } + + fn mutate_connector_table( + &self, + table_scan: &TableScan, + table: &SourceTable, + ) -> DFResult> { + let input = self.projection(table_scan, table)?; + + let schema = input.schema().clone(); + let remote = LogicalPlan::Extension(Extension { + node: Arc::new(RemoteTableBoundaryNode { + upstream_plan: input, + table_identifier: table_scan.table_name.to_owned(), + resolved_schema: schema, + requires_materialization: true, + }), + }); + + let watermark_node = EventTimeWatermarkNode::try_new( + remote, + table_scan.table_name.clone(), + Self::watermark_expression(table)?, + ) + .map_err(|err| { + DataFusionError::Internal(format!("failed to create watermark node: {err}")) + })?; + + Ok(Transformed::yes(LogicalPlan::Extension(Extension { + node: Arc::new(watermark_node), + }))) + } + + fn mutate_table_from_query( + &self, + table_scan: &TableScan, + logical_plan: &LogicalPlan, + ) -> DFResult> { + let column_expressions: Vec<_> = if let Some(projection) = &table_scan.projection { + logical_plan + .schema() + .columns() + .into_iter() + .enumerate() + .filter_map(|(i, col)| { + if projection.contains(&i) { + Some(Expr::Column(col)) + } else { + None + } + }) + .collect() + } else { + logical_plan + .schema() + .columns() + .into_iter() + .map(Expr::Column) + .collect() + }; + + let target_columns: Vec<_> = table_scan.projected_schema.columns().into_iter().collect(); + + let expressions = column_expressions + .into_iter() + .zip(target_columns) + .map(|(expr, col)| expr.alias_qualified(col.relation, col.name)) + .collect(); + + let projection = LogicalPlan::Projection(Projection::try_new_with_schema( + expressions, + Arc::new(logical_plan.clone()), + table_scan.projected_schema.clone(), + )?); + + Ok(Transformed::yes(projection)) + } +} + +impl TreeNodeRewriter for SourceRewriter<'_> { + type Node = LogicalPlan; + + fn f_up(&mut self, node: Self::Node) -> DFResult> { + let LogicalPlan::TableScan(table_scan) = node else { + return Ok(Transformed::no(node)); + }; + + let table_name = table_scan.table_name.table(); + let table = self + .schema_provider + .get_catalog_table(table_name) + .ok_or_else(|| DataFusionError::Plan(format!("Table {table_name} not found")))?; + + match table { + CatalogEntity::ExternalConnector(b) => match b.as_ref() { + ExternalTable::Source(source) => self.mutate_connector_table(&table_scan, source), + ExternalTable::Lookup(_) => { + // TODO: implement LookupSource extension + plan_err!("Lookup tables are not yet supported") + } + ExternalTable::Sink(sink) => plan_err!( + "Cannot SELECT from sink table '{}' (sinks are write-only)", + sink.name() + ), + }, + CatalogEntity::ComputedTable { + name: _, + logical_plan, + } => self.mutate_table_from_query(&table_scan, logical_plan.as_ref()), + } + } +} diff --git a/src/streaming_planner/src/analysis/stream_rewriter.rs b/src/streaming_planner/src/analysis/stream_rewriter.rs new file mode 100644 index 00000000..efc4e40c --- /dev/null +++ b/src/streaming_planner/src/analysis/stream_rewriter.rs @@ -0,0 +1,234 @@ +// 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::sync::Arc; + +use super::StreamSchemaProvider; +use crate::analysis::TimeWindowNullCheckRemover; +use crate::analysis::row_time_rewriter::RowTimeRewriter; +use crate::analysis::{ + aggregate_rewriter::AggregateRewriter, join_rewriter::JoinRewriter, + window_function_rewriter::WindowFunctionRewriter, +}; +use crate::logical_node::StreamingOperatorBlueprint; +use crate::logical_node::remote_table::RemoteTableBoundaryNode; +use crate::schema::utils::{add_timestamp_field, has_timestamp_field}; +use crate::types::{QualifiedField, TIMESTAMP_FIELD}; +use datafusion::common::tree_node::{Transformed, TreeNodeRewriter}; +use datafusion::common::{Column, DataFusionError, Result, Spans, TableReference, plan_err}; +use datafusion::logical_expr::{ + Expr, Extension, Filter, LogicalPlan, Projection, SubqueryAlias, Union, +}; +use datafusion_common::tree_node::TreeNode; +use datafusion_expr::{Aggregate, Join}; + +pub struct StreamRewriter<'a> { + pub(crate) schema_provider: &'a StreamSchemaProvider, +} + +impl TreeNodeRewriter for StreamRewriter<'_> { + type Node = LogicalPlan; + + fn f_up(&mut self, node: Self::Node) -> Result> { + match node { + // Logic Delegation + LogicalPlan::Projection(p) => self.rewrite_projection(p), + LogicalPlan::Filter(f) => self.rewrite_filter(f), + LogicalPlan::Union(u) => self.rewrite_union(u), + + // Delegation to specialized sub-rewriters + LogicalPlan::Aggregate(agg) => self.rewrite_aggregate(agg), + LogicalPlan::Join(join) => self.rewrite_join(join), + LogicalPlan::Window(_) => self.rewrite_window(node), + LogicalPlan::SubqueryAlias(sa) => self.rewrite_subquery_alias(sa), + + // Explicitly Unsupported Operations + LogicalPlan::Sort(_) => self.unsupported_error("ORDER BY", &node), + LogicalPlan::Limit(_) => self.unsupported_error("LIMIT", &node), + LogicalPlan::Repartition(_) => self.unsupported_error("Repartitions", &node), + LogicalPlan::Explain(_) => self.unsupported_error("EXPLAIN", &node), + LogicalPlan::Analyze(_) => self.unsupported_error("ANALYZE", &node), + + _ => Ok(Transformed::no(node)), + } + } +} + +impl<'a> StreamRewriter<'a> { + pub fn new(schema_provider: &'a StreamSchemaProvider) -> Self { + Self { schema_provider } + } + + /// Delegates to AggregateRewriter to transform batch aggregates into streaming stateful operators. + fn rewrite_aggregate(&self, agg: Aggregate) -> Result> { + AggregateRewriter { + schema_provider: self.schema_provider, + } + .f_up(LogicalPlan::Aggregate(agg)) + } + + /// Delegates to JoinRewriter to handle streaming join semantics (e.g., TTL, state management). + fn rewrite_join(&self, join: Join) -> Result> { + JoinRewriter { + schema_provider: self.schema_provider, + } + .f_up(LogicalPlan::Join(join)) + } + + /// Delegates to WindowFunctionRewriter for stream-aware windowing logic. + fn rewrite_window(&self, node: LogicalPlan) -> Result> { + WindowFunctionRewriter {}.f_up(node) + } + + /// Refreshes SubqueryAlias metadata to align with potentially rewritten internal schemas. + fn rewrite_subquery_alias(&self, sa: SubqueryAlias) -> Result> { + // Since the inner 'sa.input' has been rewritten (bottom-up), we must re-create + // the alias node to ensure the outer schema correctly reflects internal changes. + let new_sa = SubqueryAlias::try_new(sa.input, sa.alias).map_err(|e| { + DataFusionError::Internal(format!("Failed to re-alias subquery: {}", e)) + })?; + + Ok(Transformed::yes(LogicalPlan::SubqueryAlias(new_sa))) + } + + /// Handles timestamp propagation and row_time() mapping for Projections + fn rewrite_projection(&self, mut projection: Projection) -> Result> { + // Check if the current projection already has a timestamp field; + // if not, we must inject it to maintain streaming heartbeats. + if !has_timestamp_field(&projection.schema) { + let input_schema = projection.input.schema(); + + // Resolve the timestamp field from the input schema using the global constant. + let timestamp_field: QualifiedField = input_schema + .qualified_field_with_unqualified_name(TIMESTAMP_FIELD) + .map_err(|_| { + DataFusionError::Plan(format!( + "No timestamp field found in projection input ({})", + projection.input.display() + )) + })? + .into(); + + // Update the logical schema to include the newly injected timestamp. + projection.schema = add_timestamp_field( + projection.schema.clone(), + timestamp_field.qualifier().cloned(), + ) + .expect("Failed to add timestamp to projection schema"); + + // Physically push the timestamp column into the expression list. + projection.expr.push(Expr::Column(Column { + relation: timestamp_field.qualifier().cloned(), + name: TIMESTAMP_FIELD.to_string(), + spans: Spans::default(), + })); + } + + // Map user-friendly row_time() function calls to internal _timestamp column references. + let rewritten = projection + .expr + .iter() + .map(|expr| expr.clone().rewrite(&mut RowTimeRewriter {})) + .collect::>>()?; + + // If any expressions were modified (e.g., row_time() was replaced), update the projection. + if rewritten.iter().any(|r| r.transformed) { + projection.expr = rewritten.into_iter().map(|r| r.data).collect(); + } + + // Return the updated plan node wrapped in a Transformed container. + Ok(Transformed::yes(LogicalPlan::Projection(projection))) + } + + /// Harmonizes schemas across Union branches and wraps them in RemoteTableBoundaryNodes. + /// + /// This ensures that all inputs to a UNION operation share the exact same schema metadata, + /// preventing "Schema Drift" where different branches have different field qualifiers. + fn rewrite_union(&self, mut union: Union) -> Result> { + // Industrial engines use the first branch as the "Master Schema" for the Union. + // We clone it once to ensure all subsequent branches are forced to comply. + let master_schema = union.inputs[0].schema().clone(); + union.schema = master_schema.clone(); + + for input in union.inputs.iter_mut() { + // Optimization: If the node is already a non-transparent Extension, + // we skip wrapping to avoid unnecessary nesting of logical nodes. + if let LogicalPlan::Extension(Extension { node }) = input.as_ref() { + let stream_ext: &dyn StreamingOperatorBlueprint = node.try_into().map_err(|e| { + DataFusionError::Internal(format!( + "Failed to resolve StreamingOperatorBlueprint: {}", + e + )) + })?; + + if !stream_ext.is_passthrough_boundary() { + continue; + } + } + + // Wrap each branch in a RemoteTableBoundaryNode. + // This acts as a logical "bridge" that forces the input to adopt the master_schema, + // effectively stripping away branch-specific qualifiers (e.g., table aliases). + let remote_ext = Arc::new(RemoteTableBoundaryNode { + upstream_plan: input.as_ref().clone(), + table_identifier: TableReference::bare("union_input"), + resolved_schema: master_schema.clone(), + requires_materialization: false, // Internal logical boundary only; does not require physical sink. + }); + + // Atomically replace the input with the wrapped version. + *input = Arc::new(LogicalPlan::Extension(Extension { node: remote_ext })); + } + + Ok(Transformed::yes(LogicalPlan::Union(union))) + } + + /// Optimizes Filter nodes by stripping redundant NULL checks on time window expressions. + /// + /// In streaming SQL, DataFusion often injects 'IS NOT NULL' guards for window functions + /// that are redundant or can interfere with watermark propagation. This rewriter + /// cleans those predicates to simplify the physical execution plan. + fn rewrite_filter(&self, filter: Filter) -> Result> { + // We attempt to rewrite the predicate using a specialized sub-rewriter. + // The TimeWindowNullCheckRemover specifically targets expressions like + // `tumble(...) IS NOT NULL` and simplifies them to `TRUE`. + let rewritten_expr = filter + .predicate + .clone() + .rewrite(&mut TimeWindowNullCheckRemover {})?; + + if !rewritten_expr.transformed { + return Ok(Transformed::no(LogicalPlan::Filter(filter))); + } + + // Industrial Guard: Re-validate the predicate against the input schema. + // 'Filter::try_new' ensures that the transformed expression is still semantically + // valid for the underlying data stream. + let new_filter = Filter::try_new(rewritten_expr.data, filter.input).map_err(|e| { + DataFusionError::Internal(format!( + "Failed to re-validate filtered predicate after NULL-check removal: {}", + e + )) + })?; + + Ok(Transformed::yes(LogicalPlan::Filter(new_filter))) + } + + /// Centralized error handler for unsupported streaming operations + fn unsupported_error(&self, op: &str, node: &LogicalPlan) -> Result> { + plan_err!( + "{} is not currently supported in streaming SQL ({})", + op, + node.display() + ) + } +} diff --git a/src/streaming_planner/src/analysis/streaming_window_analzer.rs b/src/streaming_planner/src/analysis/streaming_window_analzer.rs new file mode 100644 index 00000000..5e6d3e91 --- /dev/null +++ b/src/streaming_planner/src/analysis/streaming_window_analzer.rs @@ -0,0 +1,219 @@ +// 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 std::sync::Arc; + +use datafusion::common::tree_node::{TreeNodeRecursion, TreeNodeVisitor}; +use datafusion::common::{Column, DFSchema, DataFusionError, Result}; +use datafusion::logical_expr::{Expr, Extension, LogicalPlan, expr::Alias}; + +use crate::logical_node::aggregate::{STREAM_AGG_EXTENSION_NAME, StreamWindowAggregateNode}; +use crate::logical_node::join::STREAM_JOIN_NODE_TYPE; +use crate::types::{ + QualifiedField, WindowBehavior, WindowType, extract_qualified_fields, extract_window_type, +}; + +/// WindowDetectingVisitor identifies windowing strategies and tracks window-carrying fields +/// as they propagate upward through the logical plan tree. +#[derive(Debug, Default)] +pub struct StreamingWindowAnalzer { + /// The specific window type discovered (Tumble, Hop, etc.) + pub(crate) window: Option, + /// Set of fields in the current plan node that carry window semantics. + pub(crate) fields: HashSet, +} + +impl StreamingWindowAnalzer { + /// Entry point to resolve the WindowType of a given plan branch. + pub fn get_window(logical_plan: &LogicalPlan) -> Result> { + let mut visitor = Self::default(); + logical_plan.visit_with_subqueries(&mut visitor)?; + Ok(visitor.window) + } + + /// Resolves whether an expression is a reference to an existing window field + /// or a definition of a new window function. + fn resolve_window_from_expr( + &self, + expr: &Expr, + input_schema: &DFSchema, + ) -> Result> { + // 1. Check if the expression directly references a known window field. + if let Some(col) = extract_column(expr) { + let field = input_schema.field_with_name(col.relation.as_ref(), &col.name)?; + let df_field: QualifiedField = (col.relation.clone(), Arc::new(field.clone())).into(); + + if self.fields.contains(&df_field) { + return Ok(self.window.clone()); + } + } + + // 2. Otherwise, check if it's a new window function call (e.g., tumble(), hop()). + extract_window_type(expr) + } + + /// Updates the internal state with new window findings and maps them to the output schema. + fn update_state( + &mut self, + matched_windows: Vec<(usize, WindowType)>, + schema: &DFSchema, + ) -> Result<()> { + // Clear fields from the previous level to maintain schema strictly for the current node. + self.fields.clear(); + + for (index, window) in matched_windows { + if let Some(existing) = &self.window { + if existing != &window { + return Err(DataFusionError::Plan(format!( + "Conflicting windows in the same operator: expected {:?}, found {:?}", + existing, window + ))); + } + } else { + self.window = Some(window); + } + // Record this specific index in the schema as a window carrier. + self.fields.insert(schema.qualified_field(index).into()); + } + Ok(()) + } +} + +pub fn extract_column(expr: &Expr) -> Option<&Column> { + match expr { + Expr::Column(column) => Some(column), + Expr::Alias(Alias { expr, .. }) => extract_column(expr), + _ => None, + } +} + +impl TreeNodeVisitor<'_> for StreamingWindowAnalzer { + type Node = LogicalPlan; + + fn f_down(&mut self, node: &Self::Node) -> Result { + // Joins require cross-branch validation to ensure left and right sides align on time. + if let LogicalPlan::Extension(Extension { node }) = node + && node.name() == STREAM_JOIN_NODE_TYPE + { + let mut branch_windows = HashSet::new(); + for input in node.inputs() { + if let Some(w) = Self::get_window(input)? { + branch_windows.insert(w); + } + } + + if branch_windows.len() > 1 { + return Err(DataFusionError::Plan( + "Join inputs have mismatched windowing strategies.".into(), + )); + } + self.window = branch_windows.into_iter().next(); + + // Optimization: No need to recurse manually if we've resolved the join boundary. + return Ok(TreeNodeRecursion::Jump); + } + Ok(TreeNodeRecursion::Continue) + } + + fn f_up(&mut self, node: &Self::Node) -> Result { + match node { + LogicalPlan::Projection(p) => { + let windows = p + .expr + .iter() + .enumerate() + .filter_map(|(i, e)| { + self.resolve_window_from_expr(e, p.input.schema()) + .transpose() + .map(|res| res.map(|w| (i, w))) + }) + .collect::>>()?; + + self.update_state(windows, &p.schema)?; + } + + LogicalPlan::Aggregate(agg) => { + let windows = agg + .group_expr + .iter() + .enumerate() + .filter_map(|(i, e)| { + self.resolve_window_from_expr(e, agg.input.schema()) + .transpose() + .map(|res| res.map(|w| (i, w))) + }) + .collect::>>()?; + + self.update_state(windows, &agg.schema)?; + } + + LogicalPlan::SubqueryAlias(sa) => { + // Map fields through the alias layer by resolving column indices. + let input_schema = sa.input.schema(); + let mapped = self + .fields + .drain() + .map(|f| { + let idx = input_schema.index_of_column(&f.qualified_column())?; + Ok(sa.schema.qualified_field(idx).into()) + }) + .collect::>>()?; + + self.fields = mapped; + } + + LogicalPlan::Extension(Extension { node }) + if node.name() == STREAM_AGG_EXTENSION_NAME => + { + let ext = node + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Internal("StreamWindowAggregateNode is malformed".into()) + })?; + + match &ext.window_spec { + WindowBehavior::FromOperator { + window, + window_field, + is_nested, + .. + } => { + if self.window.is_some() && !*is_nested { + return Err(DataFusionError::Plan( + "Redundant window definition on an already windowed stream.".into(), + )); + } + self.window = Some(window.clone()); + self.fields.insert(window_field.clone()); + } + WindowBehavior::InData => { + let current_schema_fields: HashSet<_> = + extract_qualified_fields(node.schema()) + .into_iter() + .collect(); + self.fields.retain(|f| current_schema_fields.contains(f)); + + if self.fields.is_empty() { + return Err(DataFusionError::Plan( + "Windowed aggregate missing window metadata from its input.".into(), + )); + } + } + } + } + _ => {} + } + Ok(TreeNodeRecursion::Continue) + } +} diff --git a/src/streaming_planner/src/analysis/time_window.rs b/src/streaming_planner/src/analysis/time_window.rs new file mode 100644 index 00000000..104c0cca --- /dev/null +++ b/src/streaming_planner/src/analysis/time_window.rs @@ -0,0 +1,83 @@ +// 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 datafusion::common::tree_node::{ + Transformed, TreeNodeRecursion, TreeNodeRewriter, TreeNodeVisitor, +}; +use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue, plan_err}; +use datafusion::logical_expr::expr::ScalarFunction; +use datafusion::logical_expr::{Expr, LogicalPlan}; + +/// Returns the time window function name if the expression is one (tumble/hop/session). +pub fn is_time_window(expr: &Expr) -> Option<&str> { + if let Expr::ScalarFunction(ScalarFunction { func, args: _ }) = expr { + match func.name() { + "tumble" | "hop" | "session" => return Some(func.name()), + _ => {} + } + } + None +} + +struct TimeWindowExprChecker {} + +impl TreeNodeVisitor<'_> for TimeWindowExprChecker { + type Node = Expr; + + fn f_down(&mut self, node: &Self::Node) -> DFResult { + if let Some(w) = is_time_window(node) { + return plan_err!( + "time window function {} is not allowed in this context. \ + Are you missing a GROUP BY clause?", + w + ); + } + Ok(TreeNodeRecursion::Continue) + } +} + +/// Visitor that checks an entire LogicalPlan for misplaced time window UDFs. +pub struct TimeWindowUdfChecker {} + +impl TreeNodeVisitor<'_> for TimeWindowUdfChecker { + type Node = LogicalPlan; + + fn f_down(&mut self, node: &Self::Node) -> DFResult { + use datafusion::common::tree_node::TreeNode; + node.expressions().iter().try_for_each(|expr| { + let mut checker = TimeWindowExprChecker {}; + expr.visit(&mut checker)?; + Ok::<(), DataFusionError>(()) + })?; + Ok(TreeNodeRecursion::Continue) + } +} + +/// Removes `IS NOT NULL` checks wrapping time window functions, +/// replacing them with `true` since time windows are never null. +pub struct TimeWindowNullCheckRemover {} + +impl TreeNodeRewriter for TimeWindowNullCheckRemover { + type Node = Expr; + + fn f_down(&mut self, node: Self::Node) -> DFResult> { + if let Expr::IsNotNull(expr) = &node + && is_time_window(expr).is_some() + { + return Ok(Transformed::yes(Expr::Literal( + ScalarValue::Boolean(Some(true)), + None, + ))); + } + Ok(Transformed::no(node)) + } +} diff --git a/src/streaming_planner/src/analysis/udafs.rs b/src/streaming_planner/src/analysis/udafs.rs new file mode 100644 index 00000000..73fc062c --- /dev/null +++ b/src/streaming_planner/src/analysis/udafs.rs @@ -0,0 +1,43 @@ +// 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 datafusion::arrow::array::ArrayRef; +use datafusion::error::Result; +use datafusion::physical_plan::Accumulator; +use datafusion::scalar::ScalarValue; +use std::fmt::Debug; + +/// Fake UDAF used just for plan-time placeholder. +#[derive(Debug)] +pub struct EmptyUdaf {} + +impl Accumulator for EmptyUdaf { + fn update_batch(&mut self, _: &[ArrayRef]) -> Result<()> { + unreachable!() + } + + fn evaluate(&self) -> Result { + unreachable!() + } + + fn size(&self) -> usize { + unreachable!() + } + + fn state(&self) -> Result> { + unreachable!() + } + + fn merge_batch(&mut self, _: &[ArrayRef]) -> Result<()> { + unreachable!() + } +} diff --git a/src/streaming_planner/src/analysis/unnest_rewriter.rs b/src/streaming_planner/src/analysis/unnest_rewriter.rs new file mode 100644 index 00000000..e63878f7 --- /dev/null +++ b/src/streaming_planner/src/analysis/unnest_rewriter.rs @@ -0,0 +1,179 @@ +// 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::sync::Arc; + +use datafusion::arrow::datatypes::DataType; +use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRewriter}; +use datafusion::common::{Column, Result as DFResult, plan_err}; +use datafusion::logical_expr::expr::ScalarFunction; +use datafusion::logical_expr::{ColumnUnnestList, Expr, LogicalPlan, Projection, Unnest}; + +use crate::common::constants::planning_placeholder_udf; +use crate::types::{QualifiedField, build_df_schema, extract_qualified_fields}; + +pub const UNNESTED_COL: &str = "__unnested"; + +/// Rewrites projections containing `unnest()` calls into proper Unnest logical plans. +pub struct UnnestRewriter {} + +impl UnnestRewriter { + fn split_unnest(expr: Expr) -> DFResult<(Expr, Option)> { + let mut captured: Option = None; + + let expr = expr.transform_up(|e| { + if let Expr::ScalarFunction(ScalarFunction { func: udf, args }) = &e + && udf.name() == planning_placeholder_udf::UNNEST + { + match args.len() { + 1 => { + if captured.replace(args[0].clone()).is_some() { + return plan_err!( + "Multiple unnests in expression, which is not allowed" + ); + } + return Ok(Transformed::yes(Expr::Column(Column::new_unqualified( + UNNESTED_COL, + )))); + } + n => { + panic!("Unnest has wrong number of arguments (expected 1, found {n})"); + } + } + } + Ok(Transformed::no(e)) + })?; + + Ok((expr.data, captured)) + } +} + +impl TreeNodeRewriter for UnnestRewriter { + type Node = LogicalPlan; + + fn f_up(&mut self, node: Self::Node) -> DFResult> { + let LogicalPlan::Projection(projection) = &node else { + if node.expressions().iter().any(|e| { + let e = Self::split_unnest(e.clone()); + e.is_err() || e.unwrap().1.is_some() + }) { + return plan_err!("unnest is only supported in SELECT statements"); + } + return Ok(Transformed::no(node)); + }; + + let mut unnest = None; + let exprs = projection + .expr + .clone() + .into_iter() + .enumerate() + .map(|(i, expr)| { + let (expr, opt) = Self::split_unnest(expr)?; + let is_unnest = if let Some(e) = opt { + if let Some(prev) = unnest.replace((e, i)) + && &prev != unnest.as_ref().unwrap() + { + return plan_err!( + "Projection contains multiple unnests, which is not currently supported" + ); + } + true + } else { + false + }; + + Ok((expr, is_unnest)) + }) + .collect::>>()?; + + if let Some((unnest_inner, unnest_idx)) = unnest { + let produce_list = Arc::new(LogicalPlan::Projection(Projection::try_new( + exprs + .iter() + .cloned() + .map(|(e, is_unnest)| { + if is_unnest { + unnest_inner.clone().alias(UNNESTED_COL) + } else { + e + } + }) + .collect(), + projection.input.clone(), + )?)); + + let unnest_fields = extract_qualified_fields(produce_list.schema()) + .iter() + .enumerate() + .map(|(i, f)| { + if i == unnest_idx { + let DataType::List(inner) = f.data_type() else { + return plan_err!( + "Argument '{}' to unnest is not a List", + f.qualified_name() + ); + }; + Ok(QualifiedField::new_unqualified( + UNNESTED_COL, + inner.data_type().clone(), + inner.is_nullable(), + )) + } else { + Ok((*f).clone()) + } + }) + .collect::>>()?; + + let unnest_node = LogicalPlan::Unnest(Unnest { + exec_columns: vec![ + QualifiedField::from(produce_list.schema().qualified_field(unnest_idx)) + .qualified_column(), + ], + input: produce_list, + list_type_columns: vec![( + unnest_idx, + ColumnUnnestList { + output_column: Column::new_unqualified(UNNESTED_COL), + depth: 1, + }, + )], + struct_type_columns: vec![], + dependency_indices: vec![], + schema: Arc::new(build_df_schema(&unnest_fields)?), + options: Default::default(), + }); + + let output_node = LogicalPlan::Projection(Projection::try_new( + exprs + .iter() + .enumerate() + .map(|(i, (expr, has_unnest))| { + if *has_unnest { + expr.clone() + } else { + Expr::Column( + QualifiedField::from(unnest_node.schema().qualified_field(i)) + .qualified_column(), + ) + } + }) + .collect(), + Arc::new(unnest_node), + )?); + + Ok(Transformed::yes(output_node)) + } else { + Ok(Transformed::no(LogicalPlan::Projection(projection.clone()))) + } + } +} diff --git a/src/streaming_planner/src/analysis/window_function_rewriter.rs b/src/streaming_planner/src/analysis/window_function_rewriter.rs new file mode 100644 index 00000000..569a3c00 --- /dev/null +++ b/src/streaming_planner/src/analysis/window_function_rewriter.rs @@ -0,0 +1,204 @@ +// 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 datafusion::common::tree_node::Transformed; +use datafusion::common::{Result as DFResult, plan_err, tree_node::TreeNodeRewriter}; +use datafusion::logical_expr::{ + self, Expr, Extension, LogicalPlan, Projection, Sort, Window, expr::WindowFunction, + expr::WindowFunctionParams, +}; +use datafusion_common::DataFusionError; +use std::sync::Arc; +use tracing::debug; + +use crate::analysis::streaming_window_analzer::{StreamingWindowAnalzer, extract_column}; +use crate::logical_node::key_calculation::{KeyExtractionNode, KeyExtractionStrategy}; +use crate::logical_node::windows_function::StreamingWindowFunctionNode; +use crate::types::{WindowType, build_df_schema, extract_qualified_fields}; + +/// WindowFunctionRewriter transforms standard SQL Window functions into streaming-compatible +/// stateful operators, ensuring proper data partitioning and sorting for distributed execution. +pub struct WindowFunctionRewriter; + +impl WindowFunctionRewriter { + /// Recursively unwraps Aliases to find the underlying WindowFunction. + #[allow(clippy::only_used_in_recursion)] + fn resolve_window_function(&self, expr: &Expr) -> DFResult<(WindowFunction, String)> { + match expr { + Expr::Alias(alias) => { + let (func, _) = self.resolve_window_function(&alias.expr)?; + Ok((func, alias.name.clone())) + } + Expr::WindowFunction(wf) => Ok((wf.as_ref().clone(), expr.name_for_alias()?)), + _ => plan_err!("Expected WindowFunction or Alias, found: {:?}", expr), + } + } + + /// Identifies which field in the PARTITION BY clause corresponds to the streaming window. + fn identify_window_partition( + &self, + params: &WindowFunctionParams, + input: &LogicalPlan, + input_window_fields: &std::collections::HashSet, + ) -> DFResult { + let matched: Vec<_> = params + .partition_by + .iter() + .enumerate() + .filter_map(|(i, e)| { + let col = extract_column(e)?; + let field = input + .schema() + .field_with_name(col.relation.as_ref(), &col.name) + .ok()?; + let df_field = (col.relation.clone(), Arc::new(field.clone())).into(); + + if input_window_fields.contains(&df_field) { + Some(i) + } else { + None + } + }) + .collect(); + + if matched.len() != 1 { + return plan_err!( + "Streaming window functions require exactly one window column in PARTITION BY. Found: {}", + matched.len() + ); + } + Ok(matched[0]) + } + + /// Wraps the input in a Projection and KeyExtractionNode to handle data distribution. + fn build_keyed_input( + &self, + input: Arc, + partition_keys: &[Expr], + ) -> DFResult { + let key_count = partition_keys.len(); + + // 1. Build projection: [_key_0, _key_1, ..., original_columns] + let mut exprs: Vec<_> = partition_keys + .iter() + .enumerate() + .map(|(i, e)| e.clone().alias(format!("_key_{i}"))) + .collect(); + + exprs.extend( + extract_qualified_fields(input.schema()) + .iter() + .map(|f| Expr::Column(f.qualified_column())), + ); + + // 2. Derive the keyed schema + let mut keyed_fields = + extract_qualified_fields(&Projection::try_new(exprs.clone(), input.clone())?.schema) + .iter() + .take(key_count) + .cloned() + .collect::>(); + keyed_fields.extend(extract_qualified_fields(input.schema())); + + let keyed_schema = Arc::new(build_df_schema(&keyed_fields)?); + + let projection = + LogicalPlan::Projection(Projection::try_new_with_schema(exprs, input, keyed_schema)?); + + // 3. Wrap in KeyExtractionNode for the physical planner + Ok(LogicalPlan::Extension(Extension { + node: Arc::new(KeyExtractionNode::new( + projection, + KeyExtractionStrategy::ColumnIndices((0..key_count).collect()), + )), + })) + } +} + +impl TreeNodeRewriter for WindowFunctionRewriter { + type Node = LogicalPlan; + + fn f_up(&mut self, node: Self::Node) -> DFResult> { + let LogicalPlan::Window(window) = node else { + return Ok(Transformed::no(node)); + }; + + debug!("Rewriting window function for streaming: {:?}", window); + + // 1. Analyze input windowing context + let mut analyzer = StreamingWindowAnalzer::default(); + window.input.visit_with_subqueries(&mut analyzer)?; + + let input_window = analyzer.window.ok_or_else(|| { + DataFusionError::Plan( + "Window functions require a windowed input stream (e.g., TUMBLE/HOP)".into(), + ) + })?; + + if matches!(input_window, WindowType::Session { .. }) { + return plan_err!( + "Streaming window functions (OVER) are not supported on Session windows." + ); + } + + // 2. Validate window expression constraints + if window.window_expr.len() != 1 { + return plan_err!( + "Arroyo currently supports exactly one window expression per OVER clause." + ); + } + + let (mut wf, original_name) = self.resolve_window_function(&window.window_expr[0])?; + + // 3. Identify and extract the window column from PARTITION BY + let window_part_idx = + self.identify_window_partition(&wf.params, &window.input, &analyzer.fields)?; + let mut partition_keys = wf.params.partition_by.clone(); + partition_keys.remove(window_part_idx); + + // Update function params to exclude the window column from internal partitioning + // as the streaming engine handles window boundaries natively. + wf.params.partition_by = partition_keys.clone(); + let key_count = partition_keys.len(); + + // 4. Build the data-shuffling pipeline (Projection -> KeyCalc -> Sort) + let keyed_plan = self.build_keyed_input(window.input.clone(), &partition_keys)?; + + let mut sort_exprs: Vec<_> = partition_keys + .iter() + .map(|e| logical_expr::expr::Sort { + expr: e.clone(), + asc: true, + nulls_first: false, + }) + .collect(); + sort_exprs.extend(wf.params.order_by.clone()); + + let sorted_plan = LogicalPlan::Sort(Sort { + expr: sort_exprs, + input: Arc::new(keyed_plan), + fetch: None, + }); + + // 5. Final Assembly + let final_wf_expr = Expr::WindowFunction(Box::new(wf)).alias_if_changed(original_name)?; + let rewritten_window = + LogicalPlan::Window(Window::try_new(vec![final_wf_expr], Arc::new(sorted_plan))?); + + Ok(Transformed::yes(LogicalPlan::Extension(Extension { + node: Arc::new(StreamingWindowFunctionNode::new( + rewritten_window, + (0..key_count).collect(), + )), + }))) + } +} diff --git a/src/streaming_planner/src/api/checkpoints.rs b/src/streaming_planner/src/api/checkpoints.rs new file mode 100644 index 00000000..bd326b10 --- /dev/null +++ b/src/streaming_planner/src/api/checkpoints.rs @@ -0,0 +1,108 @@ +// 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::common::to_micros; +use serde::{Deserialize, Serialize}; +use std::time::SystemTime; + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct Checkpoint { + pub epoch: u32, + pub backend: String, + pub start_time: u64, + pub finish_time: Option, + pub events: Vec, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct CheckpointEventSpan { + pub start_time: u64, + pub finish_time: u64, + pub event: String, + pub description: String, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct SubtaskCheckpointGroup { + pub index: u32, + pub bytes: u64, + pub event_spans: Vec, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct OperatorCheckpointGroup { + pub operator_id: String, + pub bytes: u64, + pub started_metadata_write: Option, + pub finish_time: Option, + pub subtasks: Vec, +} + +#[derive(Debug, Copy, Clone, Eq, PartialEq, Serialize, Deserialize)] +pub enum JobCheckpointEventType { + Checkpointing, + CheckpointingOperators, + WritingMetadata, + Compacting, + Committing, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct JobCheckpointSpan { + pub event: JobCheckpointEventType, + pub start_time: u64, + pub finish_time: Option, +} + +impl JobCheckpointSpan { + pub fn now(event: JobCheckpointEventType) -> Self { + Self { + event, + start_time: to_micros(SystemTime::now()), + finish_time: None, + } + } + + pub fn finish(&mut self) { + if self.finish_time.is_none() { + self.finish_time = Some(to_micros(SystemTime::now())); + } + } +} + +impl From for CheckpointEventSpan { + fn from(value: JobCheckpointSpan) -> Self { + let description = match value.event { + JobCheckpointEventType::Checkpointing => "The entire checkpointing process", + JobCheckpointEventType::CheckpointingOperators => { + "The time spent checkpointing operator states" + } + JobCheckpointEventType::WritingMetadata => "Writing the final checkpoint metadata", + JobCheckpointEventType::Compacting => "Compacting old checkpoints", + JobCheckpointEventType::Committing => { + "Running two-phase commit for transactional connectors" + } + } + .to_string(); + + Self { + start_time: value.start_time, + finish_time: value.finish_time.unwrap_or_default(), + event: format!("{:?}", value.event), + description, + } + } +} diff --git a/src/streaming_planner/src/api/connections.rs b/src/streaming_planner/src/api/connections.rs new file mode 100644 index 00000000..d2d92bcb --- /dev/null +++ b/src/streaming_planner/src/api/connections.rs @@ -0,0 +1,620 @@ +// 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::common::formats::{BadData, Format, Framing}; +use crate::common::{FsExtensionType, FsSchema}; +use datafusion::arrow::datatypes::{DataType, Field, Fields, TimeUnit}; +use serde::ser::SerializeMap; +use serde::{Deserialize, Serialize, Serializer}; +use std::collections::{BTreeMap, HashMap, HashSet}; +use std::fmt::{Display, Formatter}; +use std::sync::Arc; + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct Connector { + pub id: String, + pub name: String, + pub icon: String, + pub description: String, + pub table_config: String, + pub enabled: bool, + pub source: bool, + pub sink: bool, + pub custom_schemas: bool, + pub testing: bool, + pub hidden: bool, + pub connection_config: Option, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct ConnectionProfile { + pub id: String, + pub name: String, + pub connector: String, + pub config: serde_json::Value, + pub description: String, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct ConnectionProfilePost { + pub name: String, + pub connector: String, + pub config: serde_json::Value, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)] +#[serde(rename_all = "snake_case")] +pub enum ConnectionType { + Source, + Sink, + Lookup, +} + +impl Display for ConnectionType { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + match self { + ConnectionType::Source => write!(f, "SOURCE"), + ConnectionType::Sink => write!(f, "SINK"), + ConnectionType::Lookup => write!(f, "LOOKUP"), + } + } +} + +impl TryFrom for ConnectionType { + type Error = String; + + fn try_from(value: String) -> Result { + match value.to_lowercase().as_str() { + "source" => Ok(ConnectionType::Source), + "sink" => Ok(ConnectionType::Sink), + "lookup" => Ok(ConnectionType::Lookup), + _ => Err(format!("Invalid connection type: {value}")), + } + } +} + +// ─────────────────── Field Types ─────────────────── + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum FieldType { + Int32, + Int64, + Uint32, + Uint64, + #[serde(alias = "f32")] + Float32, + #[serde(alias = "f64")] + Float64, + Decimal128(DecimalField), + Bool, + #[serde(alias = "utf8")] + String, + #[serde(alias = "binary")] + Bytes, + Timestamp(TimestampField), + Json, + Struct(StructField), + List(ListField), +} + +impl FieldType { + pub fn sql_type(&self) -> String { + match self { + FieldType::Int32 => "INTEGER".into(), + FieldType::Int64 => "BIGINT".into(), + FieldType::Uint32 => "INTEGER UNSIGNED".into(), + FieldType::Uint64 => "BIGINT UNSIGNED".into(), + FieldType::Float32 => "FLOAT".into(), + FieldType::Float64 => "DOUBLE".into(), + FieldType::Decimal128(f) => format!("DECIMAL({}, {})", f.precision, f.scale), + FieldType::Bool => "BOOLEAN".into(), + FieldType::String => "TEXT".into(), + FieldType::Bytes => "BINARY".into(), + FieldType::Timestamp(t) => format!("TIMESTAMP({})", t.unit.precision()), + FieldType::Json => "JSON".into(), + FieldType::List(item) => format!("{}[]", item.items.field_type.sql_type()), + FieldType::Struct(StructField { fields, .. }) => { + format!( + "STRUCT <{}>", + fields + .iter() + .map(|f| format!("{} {}", f.name, f.field_type.sql_type())) + .collect::>() + .join(", ") + ) + } + } + } +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, Default)] +#[serde(rename_all = "snake_case")] +pub enum TimestampUnit { + #[serde(alias = "s")] + Second, + #[default] + #[serde(alias = "ms")] + Millisecond, + #[serde(alias = "µs", alias = "us")] + Microsecond, + #[serde(alias = "ns")] + Nanosecond, +} + +impl TimestampUnit { + pub fn precision(&self) -> u8 { + match self { + TimestampUnit::Second => 0, + TimestampUnit::Millisecond => 3, + TimestampUnit::Microsecond => 6, + TimestampUnit::Nanosecond => 9, + } + } +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub struct TimestampField { + #[serde(default)] + pub unit: TimestampUnit, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub struct DecimalField { + pub precision: u8, + pub scale: i8, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub struct StructField { + pub fields: Vec, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub struct ListField { + pub items: Box, +} + +fn default_item_name() -> String { + "item".to_string() +} + +#[derive(Deserialize, Clone, Debug, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub struct ListFieldItem { + #[serde(default = "default_item_name")] + pub name: String, + #[serde(flatten)] + pub field_type: FieldType, + #[serde(default)] + pub required: bool, + #[serde(default)] + pub sql_name: Option, +} + +impl From for Field { + fn from(value: ListFieldItem) -> Self { + SourceField { + name: value.name, + field_type: value.field_type, + required: value.required, + sql_name: None, + metadata_key: None, + } + .into() + } +} + +impl Serialize for ListFieldItem { + fn serialize(&self, s: S) -> Result + where + S: Serializer, + { + let mut f = Serializer::serialize_map(s, None)?; + f.serialize_entry("name", &self.name)?; + serialize_field_type_flat(&self.field_type, &mut f)?; + f.serialize_entry("required", &self.required)?; + f.serialize_entry("sql_name", &self.field_type.sql_type())?; + f.end() + } +} + +impl TryFrom for ListFieldItem { + type Error = String; + + fn try_from(value: Field) -> Result { + let source_field: SourceField = value.try_into()?; + Ok(Self { + name: source_field.name, + field_type: source_field.field_type, + required: source_field.required, + sql_name: None, + }) + } +} + +fn serialize_field_type_flat(ft: &FieldType, map: &mut M) -> Result<(), M::Error> { + let type_tag = match ft { + FieldType::Int32 => "int32", + FieldType::Int64 => "int64", + FieldType::Uint32 => "uint32", + FieldType::Uint64 => "uint64", + FieldType::Float32 => "float32", + FieldType::Float64 => "float64", + FieldType::Decimal128(_) => "decimal128", + FieldType::Bool => "bool", + FieldType::String => "string", + FieldType::Bytes => "bytes", + FieldType::Timestamp(_) => "timestamp", + FieldType::Json => "json", + FieldType::Struct(_) => "struct", + FieldType::List(_) => "list", + }; + map.serialize_entry("type", type_tag)?; + + match ft { + FieldType::Decimal128(d) => { + map.serialize_entry("precision", &d.precision)?; + map.serialize_entry("scale", &d.scale)?; + } + FieldType::Timestamp(t) => { + map.serialize_entry("unit", &t.unit)?; + } + FieldType::Struct(s) => { + map.serialize_entry("fields", &s.fields)?; + } + FieldType::List(l) => { + map.serialize_entry("items", &l.items)?; + } + _ => {} + } + Ok(()) +} + +// ─────────────────── Source Field ─────────────────── + +#[derive(Deserialize, Clone, Debug, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub struct SourceField { + pub name: String, + #[serde(flatten)] + pub field_type: FieldType, + #[serde(default)] + pub required: bool, + #[serde(default)] + pub sql_name: Option, + #[serde(default)] + pub metadata_key: Option, +} + +impl Serialize for SourceField { + fn serialize(&self, s: S) -> Result + where + S: Serializer, + { + let mut f = Serializer::serialize_map(s, None)?; + f.serialize_entry("name", &self.name)?; + serialize_field_type_flat(&self.field_type, &mut f)?; + f.serialize_entry("required", &self.required)?; + if let Some(metadata_key) = &self.metadata_key { + f.serialize_entry("metadata_key", metadata_key)?; + } + f.serialize_entry("sql_name", &self.field_type.sql_type())?; + f.end() + } +} + +impl From for Field { + fn from(f: SourceField) -> Self { + let (t, ext) = match f.field_type { + FieldType::Int32 => (DataType::Int32, None), + FieldType::Int64 => (DataType::Int64, None), + FieldType::Uint32 => (DataType::UInt32, None), + FieldType::Uint64 => (DataType::UInt64, None), + FieldType::Float32 => (DataType::Float32, None), + FieldType::Float64 => (DataType::Float64, None), + FieldType::Bool => (DataType::Boolean, None), + FieldType::String => (DataType::Utf8, None), + FieldType::Bytes => (DataType::Binary, None), + FieldType::Decimal128(d) => (DataType::Decimal128(d.precision, d.scale), None), + FieldType::Timestamp(TimestampField { + unit: TimestampUnit::Second, + }) => (DataType::Timestamp(TimeUnit::Second, None), None), + FieldType::Timestamp(TimestampField { + unit: TimestampUnit::Millisecond, + }) => (DataType::Timestamp(TimeUnit::Millisecond, None), None), + FieldType::Timestamp(TimestampField { + unit: TimestampUnit::Microsecond, + }) => (DataType::Timestamp(TimeUnit::Microsecond, None), None), + FieldType::Timestamp(TimestampField { + unit: TimestampUnit::Nanosecond, + }) => (DataType::Timestamp(TimeUnit::Nanosecond, None), None), + FieldType::Json => (DataType::Utf8, Some(FsExtensionType::JSON)), + FieldType::Struct(s) => ( + DataType::Struct(Fields::from( + s.fields + .into_iter() + .map(|t| t.into()) + .collect::>(), + )), + None, + ), + FieldType::List(t) => (DataType::List(Arc::new((*t.items).into())), None), + }; + + FsExtensionType::add_metadata(ext, Field::new(f.name, t, !f.required)) + } +} + +impl TryFrom for SourceField { + type Error = String; + + fn try_from(f: Field) -> Result { + let field_type = match (f.data_type(), FsExtensionType::from_map(f.metadata())) { + (DataType::Boolean, None) => FieldType::Bool, + (DataType::Int32, None) => FieldType::Int32, + (DataType::Int64, None) => FieldType::Int64, + (DataType::UInt32, None) => FieldType::Uint32, + (DataType::UInt64, None) => FieldType::Uint64, + (DataType::Float32, None) => FieldType::Float32, + (DataType::Float64, None) => FieldType::Float64, + (DataType::Decimal128(p, s), None) => FieldType::Decimal128(DecimalField { + precision: *p, + scale: *s, + }), + (DataType::Binary | DataType::LargeBinary | DataType::BinaryView, None) => { + FieldType::Bytes + } + (DataType::Timestamp(TimeUnit::Second, _), None) => { + FieldType::Timestamp(TimestampField { + unit: TimestampUnit::Second, + }) + } + (DataType::Timestamp(TimeUnit::Millisecond, _), None) => { + FieldType::Timestamp(TimestampField { + unit: TimestampUnit::Millisecond, + }) + } + (DataType::Timestamp(TimeUnit::Microsecond, _), None) => { + FieldType::Timestamp(TimestampField { + unit: TimestampUnit::Microsecond, + }) + } + (DataType::Timestamp(TimeUnit::Nanosecond, _), None) => { + FieldType::Timestamp(TimestampField { + unit: TimestampUnit::Nanosecond, + }) + } + (DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View, None) => FieldType::String, + ( + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View, + Some(FsExtensionType::JSON), + ) => FieldType::Json, + (DataType::Struct(fields), None) => { + let fields: Result<_, String> = fields + .into_iter() + .map(|f| (**f).clone().try_into()) + .collect(); + FieldType::Struct(StructField { fields: fields? }) + } + (DataType::List(item), None) => FieldType::List(ListField { + items: Box::new((**item).clone().try_into()?), + }), + dt => return Err(format!("Unsupported data type {dt:?}")), + }; + + Ok(SourceField { + name: f.name().clone(), + field_type, + required: !f.is_nullable(), + sql_name: None, + metadata_key: None, + }) + } +} + +// ─────────────────── Schema Definitions ─────────────────── + +#[allow(clippy::enum_variant_names)] +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)] +#[serde(rename_all = "snake_case", tag = "type")] +pub enum SchemaDefinition { + JsonSchema { + schema: String, + }, + ProtobufSchema { + schema: String, + #[serde(default)] + dependencies: HashMap, + }, + AvroSchema { + schema: String, + }, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)] +#[serde(rename_all = "snake_case")] +pub struct ConnectionSchema { + pub format: Option, + #[serde(default)] + pub bad_data: Option, + #[serde(default)] + pub framing: Option, + #[serde(default)] + pub fields: Vec, + #[serde(default)] + pub definition: Option, + #[serde(default)] + pub inferred: Option, + #[serde(default)] + pub primary_keys: HashSet, +} + +impl ConnectionSchema { + pub fn try_new( + format: Option, + bad_data: Option, + framing: Option, + fields: Vec, + definition: Option, + inferred: Option, + primary_keys: HashSet, + ) -> anyhow::Result { + let s = ConnectionSchema { + format, + bad_data, + framing, + fields, + definition, + inferred, + primary_keys, + }; + s.validate() + } + + pub fn validate(self) -> anyhow::Result { + let non_metadata_fields: Vec<_> = self + .fields + .iter() + .filter(|f| f.metadata_key.is_none()) + .collect(); + + if let Some(Format::RawString(_)) = &self.format + && (non_metadata_fields.len() != 1 + || non_metadata_fields.first().unwrap().field_type != FieldType::String + || non_metadata_fields.first().unwrap().name != "value") + { + anyhow::bail!( + "raw_string format requires a schema with a single field called `value` of type TEXT" + ); + } + + if let Some(Format::Json(json_format)) = &self.format + && json_format.unstructured + && (non_metadata_fields.len() != 1 + || non_metadata_fields.first().unwrap().field_type != FieldType::Json + || non_metadata_fields.first().unwrap().name != "value") + { + anyhow::bail!( + "json format with unstructured flag enabled requires a schema with a single field called `value` of type JSON" + ); + } + + Ok(self) + } + + pub fn fs_schema(&self) -> Arc { + let fields: Vec = self.fields.iter().map(|f| f.clone().into()).collect(); + Arc::new(FsSchema::from_fields(fields)) + } +} + +impl From for FsSchema { + fn from(val: ConnectionSchema) -> Self { + let fields: Vec = val.fields.into_iter().map(|f| f.into()).collect(); + FsSchema::from_fields(fields) + } +} + +// ─────────────────── Connection Table ─────────────────── + +#[derive(Serialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct ConnectionTable { + #[serde(skip_serializing)] + pub id: i64, + #[serde(rename = "id")] + pub pub_id: String, + pub name: String, + pub created_at: u64, + pub connector: String, + pub connection_profile: Option, + pub table_type: ConnectionType, + pub config: serde_json::Value, + pub schema: ConnectionSchema, + pub consumers: u32, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct ConnectionTablePost { + pub name: String, + pub connector: String, + pub connection_profile_id: Option, + pub config: serde_json::Value, + pub schema: Option, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct ConnectionAutocompleteResp { + pub values: BTreeMap>, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct TestSourceMessage { + pub error: bool, + pub done: bool, + pub message: String, +} + +impl TestSourceMessage { + pub fn info(message: impl Into) -> Self { + Self { + error: false, + done: false, + message: message.into(), + } + } + pub fn error(message: impl Into) -> Self { + Self { + error: true, + done: false, + message: message.into(), + } + } + pub fn done(message: impl Into) -> Self { + Self { + error: false, + done: true, + message: message.into(), + } + } + pub fn fail(message: impl Into) -> Self { + Self { + error: true, + done: true, + message: message.into(), + } + } +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct ConfluentSchema { + pub schema: String, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct ConfluentSchemaQueryParams { + pub endpoint: String, + pub topic: String, +} diff --git a/src/streaming_planner/src/api/metrics.rs b/src/streaming_planner/src/api/metrics.rs new file mode 100644 index 00000000..671b52f6 --- /dev/null +++ b/src/streaming_planner/src/api/metrics.rs @@ -0,0 +1,53 @@ +// 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 serde::{Deserialize, Serialize}; + +#[derive(Serialize, Deserialize, Copy, Clone, Debug, Hash, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum MetricName { + BytesRecv, + BytesSent, + MessagesRecv, + MessagesSent, + Backpressure, + TxQueueSize, + TxQueueRem, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct Metric { + pub time: u64, + pub value: f64, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct SubtaskMetrics { + pub index: u32, + pub metrics: Vec, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct MetricGroup { + pub name: MetricName, + pub subtasks: Vec, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct OperatorMetricGroup { + pub node_id: u32, + pub metric_groups: Vec, +} diff --git a/src/streaming_planner/src/api/mod.rs b/src/streaming_planner/src/api/mod.rs new file mode 100644 index 00000000..9fc6b23f --- /dev/null +++ b/src/streaming_planner/src/api/mod.rs @@ -0,0 +1,46 @@ +// 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. + +//! REST/RPC API types for the FunctionStream system. +//! +//! Adapted from Arroyo's `arroyo-rpc/src/api_types` and utility modules. + +pub mod checkpoints; +pub mod connections; +pub mod metrics; +pub mod pipelines; +pub mod public_ids; +pub mod schema_resolver; +pub mod udfs; +pub mod var_str; + +use serde::{Deserialize, Serialize}; + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "camelCase")] +pub struct PaginatedCollection { + pub data: Vec, + pub has_more: bool, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "camelCase")] +pub struct NonPaginatedCollection { + pub data: Vec, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct PaginationQueryParams { + pub starting_after: Option, + pub limit: Option, +} diff --git a/src/streaming_planner/src/api/pipelines.rs b/src/streaming_planner/src/api/pipelines.rs new file mode 100644 index 00000000..990c1ba9 --- /dev/null +++ b/src/streaming_planner/src/api/pipelines.rs @@ -0,0 +1,168 @@ +// 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 super::udfs::Udf; +use crate::common::control::ErrorDomain; +use serde::{Deserialize, Serialize}; + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct ValidateQueryPost { + pub query: String, + pub udfs: Option>, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct QueryValidationResult { + pub graph: Option, + pub errors: Vec, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct PipelinePost { + pub name: String, + pub query: String, + pub udfs: Option>, + pub parallelism: u64, + pub checkpoint_interval_micros: Option, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct PreviewPost { + pub query: String, + pub udfs: Option>, + #[serde(default)] + pub enable_sinks: bool, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct PipelinePatch { + pub parallelism: Option, + pub checkpoint_interval_micros: Option, + pub stop: Option, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct PipelineRestart { + pub force: Option, + pub ignore_state: Option, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct Pipeline { + pub id: String, + pub name: String, + pub query: String, + pub udfs: Vec, + pub checkpoint_interval_micros: u64, + pub stop: StopType, + pub created_at: u64, + pub action: Option, + pub action_text: String, + pub action_in_progress: bool, + pub graph: PipelineGraph, + pub preview: bool, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct PipelineGraph { + pub nodes: Vec, + pub edges: Vec, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct PipelineNode { + pub node_id: u32, + pub operator: String, + pub description: String, + pub parallelism: u32, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct PipelineEdge { + pub src_id: u32, + pub dest_id: u32, + pub key_type: String, + pub value_type: String, + pub edge_type: String, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub enum StopType { + None, + Checkpoint, + Graceful, + Immediate, + Force, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct FailureReason { + pub error: String, + pub domain: ErrorDomain, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct Job { + pub id: String, + pub running_desired: bool, + pub state: String, + pub run_id: u64, + pub start_time: Option, + pub finish_time: Option, + pub tasks: Option, + pub failure_reason: Option, + pub created_at: u64, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub enum JobLogLevel { + Info, + Warn, + Error, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct JobLogMessage { + pub id: String, + pub created_at: u64, + pub operator_id: Option, + pub task_index: Option, + pub level: JobLogLevel, + pub message: String, + pub details: String, + pub error_domain: Option, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct OutputData { + pub operator_id: String, + pub subtask_idx: u32, + pub timestamps: Vec, + pub start_id: u64, + pub batch: String, +} diff --git a/src/streaming_planner/src/api/public_ids.rs b/src/streaming_planner/src/api/public_ids.rs new file mode 100644 index 00000000..33aa6427 --- /dev/null +++ b/src/streaming_planner/src/api/public_ids.rs @@ -0,0 +1,69 @@ +// 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::time::{SystemTime, UNIX_EPOCH}; + +const ID_LENGTH: usize = 10; + +const ALPHABET: &[u8; 62] = b"0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"; + +pub enum IdTypes { + ApiKey, + ConnectionProfile, + Schema, + Pipeline, + JobConfig, + Checkpoint, + JobStatus, + ClusterInfo, + JobLogMessage, + ConnectionTable, + ConnectionTablePipeline, + Udf, +} + +/// Generates a unique identifier with a type-specific prefix. +/// +/// Uses a simple time + random approach instead of nanoid to avoid an extra dependency. +pub fn generate_id(id_type: IdTypes) -> String { + let prefix = match id_type { + IdTypes::ApiKey => "ak", + IdTypes::ConnectionProfile => "cp", + IdTypes::Schema => "sch", + IdTypes::Pipeline => "pl", + IdTypes::JobConfig => "job", + IdTypes::Checkpoint => "chk", + IdTypes::JobStatus => "js", + IdTypes::ClusterInfo => "ci", + IdTypes::JobLogMessage => "jlm", + IdTypes::ConnectionTable => "ct", + IdTypes::ConnectionTablePipeline => "ctp", + IdTypes::Udf => "udf", + }; + + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_nanos(); + + let mut id = String::with_capacity(ID_LENGTH); + let mut seed = nanos; + for _ in 0..ID_LENGTH { + seed ^= seed + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + let idx = (seed % ALPHABET.len() as u128) as usize; + id.push(ALPHABET[idx] as char); + } + + format!("{prefix}_{id}") +} diff --git a/src/streaming_planner/src/api/schema_resolver.rs b/src/streaming_planner/src/api/schema_resolver.rs new file mode 100644 index 00000000..57d3d702 --- /dev/null +++ b/src/streaming_planner/src/api/schema_resolver.rs @@ -0,0 +1,94 @@ +// 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 async_trait::async_trait; + +/// Trait for resolving schemas by ID (e.g., from a schema registry). +#[async_trait] +pub trait SchemaResolver: Send { + async fn resolve_schema(&self, id: u32) -> Result, String>; +} + +/// A resolver that always fails — used when no schema registry is configured. +pub struct FailingSchemaResolver; + +impl Default for FailingSchemaResolver { + fn default() -> Self { + Self + } +} + +#[async_trait] +impl SchemaResolver for FailingSchemaResolver { + async fn resolve_schema(&self, id: u32) -> Result, String> { + Err(format!( + "Schema with id {id} not available, and no schema registry configured" + )) + } +} + +/// A resolver that returns a fixed schema for a known ID. +pub struct FixedSchemaResolver { + id: u32, + schema: String, +} + +impl FixedSchemaResolver { + pub fn new(id: u32, schema: String) -> Self { + FixedSchemaResolver { id, schema } + } +} + +#[async_trait] +impl SchemaResolver for FixedSchemaResolver { + async fn resolve_schema(&self, id: u32) -> Result, String> { + if id == self.id { + Ok(Some(self.schema.clone())) + } else { + Err(format!("Unexpected schema id {}, expected {}", id, self.id)) + } + } +} + +/// A caching wrapper around any `SchemaResolver`. +pub struct CachingSchemaResolver { + inner: R, + cache: tokio::sync::RwLock>, +} + +impl CachingSchemaResolver { + pub fn new(inner: R) -> Self { + Self { + inner, + cache: tokio::sync::RwLock::new(std::collections::HashMap::new()), + } + } +} + +#[async_trait] +impl SchemaResolver for CachingSchemaResolver { + async fn resolve_schema(&self, id: u32) -> Result, String> { + { + let cache = self.cache.read().await; + if let Some(schema) = cache.get(&id) { + return Ok(Some(schema.clone())); + } + } + + let result = self.inner.resolve_schema(id).await?; + if let Some(ref schema) = result { + let mut cache = self.cache.write().await; + cache.insert(id, schema.clone()); + } + Ok(result) + } +} diff --git a/src/streaming_planner/src/api/udfs.rs b/src/streaming_planner/src/api/udfs.rs new file mode 100644 index 00000000..781d5b07 --- /dev/null +++ b/src/streaming_planner/src/api/udfs.rs @@ -0,0 +1,68 @@ +// 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 serde::{Deserialize, Serialize}; + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct Udf { + pub definition: String, + #[serde(default)] + pub language: UdfLanguage, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct ValidateUdfPost { + pub definition: String, + #[serde(default)] + pub language: UdfLanguage, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct UdfValidationResult { + pub udf_name: Option, + pub errors: Vec, +} + +#[derive(Serialize, Deserialize, Copy, Clone, Debug, Default, Eq, PartialEq)] +#[serde(rename_all = "snake_case")] +pub enum UdfLanguage { + Python, + #[default] + Rust, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct UdfPost { + pub prefix: String, + #[serde(default)] + pub language: UdfLanguage, + pub definition: String, + pub description: Option, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct GlobalUdf { + pub id: String, + pub prefix: String, + pub name: String, + pub language: UdfLanguage, + pub created_at: u64, + pub updated_at: u64, + pub definition: String, + pub description: Option, + pub dylib_url: Option, +} diff --git a/src/streaming_planner/src/api/var_str.rs b/src/streaming_planner/src/api/var_str.rs new file mode 100644 index 00000000..2638cd06 --- /dev/null +++ b/src/streaming_planner/src/api/var_str.rs @@ -0,0 +1,91 @@ +// 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 serde::{Deserialize, Serialize}; +use std::env; + +/// A string that may contain `{{ VAR }}` placeholders for environment variable substitution. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(transparent)] +pub struct VarStr { + raw_val: String, +} + +impl VarStr { + pub fn new(raw_val: String) -> Self { + VarStr { raw_val } + } + + pub fn raw(&self) -> &str { + &self.raw_val + } + + /// Substitute `{{ VAR_NAME }}` patterns with the corresponding environment variable values. + pub fn sub_env_vars(&self) -> anyhow::Result { + let mut result = self.raw_val.clone(); + let mut start = 0; + + while let Some(open) = result[start..].find("{{") { + let open_abs = start + open; + let Some(close) = result[open_abs..].find("}}") else { + break; + }; + let close_abs = open_abs + close; + + let var_name = result[open_abs + 2..close_abs].trim(); + if var_name.is_empty() { + start = close_abs + 2; + continue; + } + + match env::var(var_name) { + Ok(value) => { + let full_match = &result[open_abs..close_abs + 2]; + let full_match_owned = full_match.to_string(); + result = result.replacen(&full_match_owned, &value, 1); + start = open_abs + value.len(); + } + Err(_) => { + anyhow::bail!("Environment variable {} not found", var_name); + } + } + } + + Ok(result) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_no_placeholders() { + let input = "This is a test string with no placeholders"; + assert_eq!( + VarStr::new(input.to_string()).sub_env_vars().unwrap(), + input + ); + } + + #[test] + fn test_with_placeholders() { + unsafe { env::set_var("FS_TEST_VAR", "environment variable") }; + let input = "This is a {{ FS_TEST_VAR }}"; + let expected = "This is a environment variable"; + assert_eq!( + VarStr::new(input.to_string()).sub_env_vars().unwrap(), + expected + ); + unsafe { env::remove_var("FS_TEST_VAR") }; + } +} diff --git a/src/streaming_planner/src/common/arrow_ext.rs b/src/streaming_planner/src/common/arrow_ext.rs new file mode 100644 index 00000000..f041ec6f --- /dev/null +++ b/src/streaming_planner/src/common/arrow_ext.rs @@ -0,0 +1,182 @@ +// 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::fmt::{Display, Formatter}; +use std::time::SystemTime; + +use datafusion::arrow::datatypes::{DataType, Field, TimeUnit}; + +pub struct DisplayAsSql<'a>(pub &'a DataType); + +impl Display for DisplayAsSql<'_> { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + match self.0 { + DataType::Boolean => write!(f, "BOOLEAN"), + DataType::Int8 | DataType::Int16 | DataType::Int32 => write!(f, "INT"), + DataType::Int64 => write!(f, "BIGINT"), + DataType::UInt8 | DataType::UInt16 | DataType::UInt32 => write!(f, "INT UNSIGNED"), + DataType::UInt64 => write!(f, "BIGINT UNSIGNED"), + DataType::Float16 | DataType::Float32 => write!(f, "FLOAT"), + DataType::Float64 => write!(f, "DOUBLE"), + DataType::Timestamp(_, _) => write!(f, "TIMESTAMP"), + DataType::Date32 => write!(f, "DATE"), + DataType::Date64 => write!(f, "DATETIME"), + DataType::Time32(_) => write!(f, "TIME"), + DataType::Time64(_) => write!(f, "TIME"), + DataType::Duration(_) => write!(f, "INTERVAL"), + DataType::Interval(_) => write!(f, "INTERVAL"), + DataType::Binary | DataType::FixedSizeBinary(_) | DataType::LargeBinary => { + write!(f, "BYTEA") + } + DataType::Utf8 | DataType::LargeUtf8 => write!(f, "TEXT"), + DataType::List(inner) => { + write!(f, "{}[]", DisplayAsSql(inner.data_type())) + } + dt => write!(f, "{dt}"), + } + } +} + +/// Arrow extension type markers for FunctionStream-specific semantics. +#[allow(clippy::upper_case_acronyms)] +#[derive(Copy, Clone, Eq, PartialEq, Debug)] +pub enum FsExtensionType { + JSON, +} + +impl FsExtensionType { + pub fn from_map(map: &HashMap) -> Option { + match map.get("ARROW:extension:name")?.as_str() { + "functionstream.json" => Some(Self::JSON), + _ => None, + } + } + + pub fn add_metadata(v: Option, field: Field) -> Field { + if let Some(v) = v { + let mut m = HashMap::new(); + match v { + FsExtensionType::JSON => { + m.insert( + "ARROW:extension:name".to_string(), + "functionstream.json".to_string(), + ); + } + } + field.with_metadata(m) + } else { + field + } + } +} + +pub trait GetArrowType { + fn arrow_type() -> DataType; +} + +pub trait GetArrowSchema { + fn arrow_schema() -> datafusion::arrow::datatypes::Schema; +} + +impl GetArrowType for T +where + T: GetArrowSchema, +{ + fn arrow_type() -> DataType { + DataType::Struct(Self::arrow_schema().fields.clone()) + } +} + +impl GetArrowType for bool { + fn arrow_type() -> DataType { + DataType::Boolean + } +} + +impl GetArrowType for i8 { + fn arrow_type() -> DataType { + DataType::Int8 + } +} + +impl GetArrowType for i16 { + fn arrow_type() -> DataType { + DataType::Int16 + } +} + +impl GetArrowType for i32 { + fn arrow_type() -> DataType { + DataType::Int32 + } +} + +impl GetArrowType for i64 { + fn arrow_type() -> DataType { + DataType::Int64 + } +} + +impl GetArrowType for u8 { + fn arrow_type() -> DataType { + DataType::UInt8 + } +} + +impl GetArrowType for u16 { + fn arrow_type() -> DataType { + DataType::UInt16 + } +} + +impl GetArrowType for u32 { + fn arrow_type() -> DataType { + DataType::UInt32 + } +} + +impl GetArrowType for u64 { + fn arrow_type() -> DataType { + DataType::UInt64 + } +} + +impl GetArrowType for f32 { + fn arrow_type() -> DataType { + DataType::Float32 + } +} + +impl GetArrowType for f64 { + fn arrow_type() -> DataType { + DataType::Float64 + } +} + +impl GetArrowType for String { + fn arrow_type() -> DataType { + DataType::Utf8 + } +} + +impl GetArrowType for Vec { + fn arrow_type() -> DataType { + DataType::Binary + } +} + +impl GetArrowType for SystemTime { + fn arrow_type() -> DataType { + DataType::Timestamp(TimeUnit::Nanosecond, None) + } +} diff --git a/src/streaming_planner/src/common/connector_options.rs b/src/streaming_planner/src/common/connector_options.rs new file mode 100644 index 00000000..0702d945 --- /dev/null +++ b/src/streaming_planner/src/common/connector_options.rs @@ -0,0 +1,449 @@ +// 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, HashMap}; +use std::num::{NonZero, NonZeroU64}; +use std::str::FromStr; +use std::time::Duration; + +use datafusion::common::{Result as DFResult, plan_datafusion_err}; +use datafusion::error::DataFusionError; +use datafusion::sql::sqlparser::ast::{Expr, Ident, SqlOption, Value as SqlValue, ValueWithSpan}; +use tracing::warn; + +use super::constants::{interval_duration_unit, with_opt_bool_str}; + +pub trait FromOpts: Sized { + fn from_opts(opts: &mut ConnectorOptions) -> DFResult; +} + +pub struct ConnectorOptions { + options: HashMap, + partitions: Vec, +} + +fn sql_expr_to_catalog_string(e: &Expr) -> String { + match e { + Expr::Value(ValueWithSpan { value, .. }) => match value { + SqlValue::SingleQuotedString(s) | SqlValue::DoubleQuotedString(s) => s.clone(), + SqlValue::NationalStringLiteral(s) => s.clone(), + SqlValue::HexStringLiteral(s) => s.clone(), + SqlValue::Number(n, _) => n.clone(), + SqlValue::Boolean(b) => b.to_string(), + SqlValue::Null => "NULL".to_string(), + other => other.to_string(), + }, + Expr::Identifier(ident) => ident.value.clone(), + other => other.to_string(), + } +} + +impl ConnectorOptions { + /// Build options from persisted catalog string maps (same semantics as SQL `WITH` literals). + pub fn from_flat_string_map(map: HashMap) -> DFResult { + let mut options = HashMap::with_capacity(map.len()); + for (k, v) in map { + options.insert( + k, + Expr::Value(SqlValue::SingleQuotedString(v).with_empty_span()), + ); + } + Ok(Self { + options, + partitions: Vec::new(), + }) + } + + pub fn new(sql_opts: &[SqlOption], partition_by: &Option>) -> DFResult { + let mut options = HashMap::new(); + + for option in sql_opts { + let SqlOption::KeyValue { key, value } = option else { + return Err(plan_datafusion_err!( + "invalid with option: '{}'; expected an `=` delimited key-value pair", + option + )); + }; + + options.insert(key.value.clone(), value.clone()); + } + + Ok(Self { + options, + partitions: partition_by.clone().unwrap_or_default(), + }) + } + + pub fn partitions(&self) -> &[Expr] { + &self.partitions + } + + pub fn pull_struct(&mut self) -> DFResult { + T::from_opts(self) + } + + pub fn pull_opt_str(&mut self, name: &str) -> DFResult> { + match self.options.remove(name) { + Some(Expr::Value(ValueWithSpan { + value: SqlValue::SingleQuotedString(s), + span: _, + })) => Ok(Some(s)), + Some(e) => Err(plan_datafusion_err!( + "expected with option '{}' to be a single-quoted string, but it was `{:?}`", + name, + e + )), + None => Ok(None), + } + } + + pub fn peek_opt_str(&self, name: &str) -> DFResult> { + match self.options.get(name) { + Some(Expr::Value(ValueWithSpan { + value: SqlValue::SingleQuotedString(s), + span: _, + })) => Ok(Some(s.clone())), + Some(e) => Err(plan_datafusion_err!( + "expected with option '{}' to be a single-quoted string, but it was `{:?}`", + name, + e + )), + None => Ok(None), + } + } + + pub fn pull_str(&mut self, name: &str) -> DFResult { + self.pull_opt_str(name)? + .ok_or_else(|| plan_datafusion_err!("required option '{}' not set", name)) + } + + pub fn pull_opt_bool(&mut self, name: &str) -> DFResult> { + match self.options.remove(name) { + Some(Expr::Value(ValueWithSpan { + value: SqlValue::Boolean(b), + span: _, + })) => Ok(Some(b)), + Some(Expr::Value(ValueWithSpan { + value: SqlValue::SingleQuotedString(s), + span: _, + })) => match s.as_str() { + with_opt_bool_str::TRUE | with_opt_bool_str::YES => Ok(Some(true)), + with_opt_bool_str::FALSE | with_opt_bool_str::NO => Ok(Some(false)), + _ => Err(plan_datafusion_err!( + "expected with option '{}' to be a boolean, but it was `'{}'`", + name, + s + )), + }, + Some(e) => Err(plan_datafusion_err!( + "expected with option '{}' to be a boolean, but it was `{:?}`", + name, + e + )), + None => Ok(None), + } + } + + pub fn pull_opt_u64(&mut self, name: &str) -> DFResult> { + match self.options.remove(name) { + Some(Expr::Value(ValueWithSpan { + value: SqlValue::Number(s, _), + span: _, + })) + | Some(Expr::Value(ValueWithSpan { + value: SqlValue::SingleQuotedString(s), + span: _, + })) => s.parse::().map(Some).map_err(|_| { + plan_datafusion_err!( + "expected with option '{}' to be an unsigned integer, but it was `{}`", + name, + s + ) + }), + Some(e) => Err(plan_datafusion_err!( + "expected with option '{}' to be an unsigned integer, but it was `{:?}`", + name, + e + )), + None => Ok(None), + } + } + + pub fn pull_opt_nonzero_u64(&mut self, name: &str) -> DFResult>> { + match self.pull_opt_u64(name)? { + Some(0) => Err(plan_datafusion_err!( + "expected with option '{name}' to be greater than 0, but it was 0" + )), + Some(i) => Ok(Some(NonZeroU64::new(i).unwrap())), + None => Ok(None), + } + } + + pub fn pull_opt_data_size_bytes(&mut self, name: &str) -> DFResult> { + self.pull_opt_str(name)? + .map(|s| { + s.parse::().map_err(|_| { + plan_datafusion_err!( + "expected with option '{}' to be a size in bytes (unsigned integer), but it was `{}`", + name, + s + ) + }) + }) + .transpose() + } + + pub fn pull_opt_i64(&mut self, name: &str) -> DFResult> { + match self.options.remove(name) { + Some(Expr::Value(ValueWithSpan { + value: SqlValue::Number(s, _), + span: _, + })) + | Some(Expr::Value(ValueWithSpan { + value: SqlValue::SingleQuotedString(s), + span: _, + })) => s.parse::().map(Some).map_err(|_| { + plan_datafusion_err!( + "expected with option '{}' to be an integer, but it was `{}`", + name, + s + ) + }), + Some(e) => Err(plan_datafusion_err!( + "expected with option '{}' to be an integer, but it was `{:?}`", + name, + e + )), + None => Ok(None), + } + } + + pub fn pull_i64(&mut self, name: &str) -> DFResult { + self.pull_opt_i64(name)? + .ok_or_else(|| plan_datafusion_err!("required option '{}' not set", name)) + } + + pub fn pull_u64(&mut self, name: &str) -> DFResult { + self.pull_opt_u64(name)? + .ok_or_else(|| plan_datafusion_err!("required option '{}' not set", name)) + } + + pub fn pull_opt_f64(&mut self, name: &str) -> DFResult> { + match self.options.remove(name) { + Some(Expr::Value(ValueWithSpan { + value: SqlValue::Number(s, _), + span: _, + })) + | Some(Expr::Value(ValueWithSpan { + value: SqlValue::SingleQuotedString(s), + span: _, + })) => s.parse::().map(Some).map_err(|_| { + plan_datafusion_err!( + "expected with option '{}' to be a double, but it was `{}`", + name, + s + ) + }), + Some(e) => Err(plan_datafusion_err!( + "expected with option '{}' to be a double, but it was `{:?}`", + name, + e + )), + None => Ok(None), + } + } + + pub fn pull_f64(&mut self, name: &str) -> DFResult { + self.pull_opt_f64(name)? + .ok_or_else(|| plan_datafusion_err!("required option '{}' not set", name)) + } + + pub fn pull_bool(&mut self, name: &str) -> DFResult { + self.pull_opt_bool(name)? + .ok_or_else(|| plan_datafusion_err!("required option '{}' not set", name)) + } + + pub fn pull_opt_duration(&mut self, name: &str) -> DFResult> { + match self.options.remove(name) { + Some(e) => Ok(Some(duration_from_sql_expr(&e).map_err(|e| { + plan_datafusion_err!("in with clause '{name}': {}", e) + })?)), + None => Ok(None), + } + } + + pub fn pull_opt_field(&mut self, name: &str) -> DFResult> { + match self.options.remove(name) { + Some(Expr::Value(ValueWithSpan { + value: SqlValue::SingleQuotedString(s), + span: _, + })) => { + warn!( + "Referred to a field in `{name}` with a string—this is deprecated and will be unsupported after Arroyo 0.14" + ); + Ok(Some(s)) + } + Some(Expr::Identifier(Ident { value, .. })) => Ok(Some(value)), + Some(e) => Err(plan_datafusion_err!( + "expected with option '{}' to be a field, but it was `{:?}`", + name, + e + )), + None => Ok(None), + } + } + + pub fn pull_opt_array(&mut self, name: &str) -> Option> { + Some(match self.options.remove(name)? { + Expr::Value(ValueWithSpan { + value: SqlValue::SingleQuotedString(s), + span, + }) => s + .split(',') + .map(|p| { + Expr::Value(ValueWithSpan { + value: SqlValue::SingleQuotedString(p.to_string()), + span, + }) + }) + .collect(), + Expr::Array(a) => a.elem, + e => vec![e], + }) + } + + pub fn pull_opt_parsed(&mut self, name: &str) -> DFResult> { + Ok(match self.pull_opt_str(name)? { + Some(s) => Some( + s.parse() + .map_err(|_| plan_datafusion_err!("invalid value '{s}' for {name}"))?, + ), + None => None, + }) + } + + pub fn keys(&self) -> impl Iterator { + self.options.keys() + } + + pub fn keys_with_prefix<'a, 'b>( + &'a self, + prefix: &'b str, + ) -> impl Iterator + 'b + where + 'a: 'b, + { + self.options.keys().filter(move |k| k.starts_with(prefix)) + } + + pub fn insert_str( + &mut self, + name: impl Into, + value: impl Into, + ) -> DFResult> { + let name = name.into(); + let value = value.into(); + let existing = self.pull_opt_str(&name)?; + self.options.insert( + name, + Expr::Value(SqlValue::SingleQuotedString(value).with_empty_span()), + ); + Ok(existing) + } + + pub fn is_empty(&self) -> bool { + self.options.is_empty() + } + + pub fn contains_key(&self, key: &str) -> bool { + self.options.contains_key(key) + } + + /// Drain all remaining options into string values (for connector runtime config). + pub fn drain_remaining_string_values(&mut self) -> DFResult> { + let taken = std::mem::take(&mut self.options); + let mut out = HashMap::with_capacity(taken.len()); + for (k, v) in taken { + out.insert(k, format!("{v}")); + } + Ok(out) + } + + /// Snapshot of all current `WITH` key/value pairs for catalog persistence (`SHOW CREATE TABLE`). + /// Call before any `pull_*` consumes options. + pub fn snapshot_for_catalog(&self) -> BTreeMap { + self.options + .iter() + .map(|(k, v)| (k.clone(), sql_expr_to_catalog_string(v))) + .collect() + } +} + +fn duration_from_sql_expr(expr: &Expr) -> Result { + match expr { + Expr::Interval(interval) => { + let s = match interval.value.as_ref() { + Expr::Value(ValueWithSpan { + value: SqlValue::SingleQuotedString(s), + .. + }) => s.clone(), + other => { + return Err(DataFusionError::Plan(format!( + "expected interval string literal, found {other}" + ))); + } + }; + parse_interval_to_duration(&s) + } + Expr::Value(ValueWithSpan { + value: SqlValue::SingleQuotedString(s), + .. + }) => parse_interval_to_duration(s), + other => Err(DataFusionError::Plan(format!( + "expected an interval expression, found {other}" + ))), + } +} + +fn parse_interval_to_duration(s: &str) -> Result { + let parts: Vec<&str> = s.split_whitespace().collect(); + if parts.len() != 2 { + return Err(DataFusionError::Plan(format!( + "invalid interval string '{s}'; expected ' '" + ))); + } + let value: u64 = parts[0] + .parse() + .map_err(|_| DataFusionError::Plan(format!("invalid interval number: {}", parts[0])))?; + let unit_lc = parts[1].to_lowercase(); + let unit = unit_lc.as_str(); + let duration = match unit { + interval_duration_unit::SECOND + | interval_duration_unit::SECONDS + | interval_duration_unit::S => Duration::from_secs(value), + interval_duration_unit::MINUTE + | interval_duration_unit::MINUTES + | interval_duration_unit::MIN => Duration::from_secs(value * 60), + interval_duration_unit::HOUR + | interval_duration_unit::HOURS + | interval_duration_unit::H => Duration::from_secs(value * 3600), + interval_duration_unit::DAY | interval_duration_unit::DAYS | interval_duration_unit::D => { + Duration::from_secs(value * 86400) + } + unit => { + return Err(DataFusionError::Plan(format!( + "unsupported interval unit '{unit}'" + ))); + } + }; + Ok(duration) +} diff --git a/src/streaming_planner/src/common/constants.rs b/src/streaming_planner/src/common/constants.rs new file mode 100644 index 00000000..8cdb68e3 --- /dev/null +++ b/src/streaming_planner/src/common/constants.rs @@ -0,0 +1,294 @@ +// 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. + +pub mod scalar_fn { + pub const GET_FIRST_JSON_OBJECT: &str = "get_first_json_object"; + pub const EXTRACT_JSON: &str = "extract_json"; + pub const EXTRACT_JSON_STRING: &str = "extract_json_string"; + pub const SERIALIZE_JSON_UNION: &str = "serialize_json_union"; + pub const MULTI_HASH: &str = "multi_hash"; +} + +pub mod window_fn { + pub const HOP: &str = "hop"; + pub const TUMBLE: &str = "tumble"; + pub const SESSION: &str = "session"; +} + +pub mod planning_placeholder_udf { + pub const UNNEST: &str = "unnest"; + pub const ROW_TIME: &str = "row_time"; + pub const LIST_ELEMENT_FIELD: &str = "field"; +} + +pub mod operator_feature { + pub const ASYNC_UDF: &str = "async-udf"; + pub const JOIN_WITH_EXPIRATION: &str = "join-with-expiration"; + pub const WINDOWED_JOIN: &str = "windowed-join"; + pub const SQL_WINDOW_FUNCTION: &str = "sql-window-function"; + pub const LOOKUP_JOIN: &str = "lookup-join"; + pub const SQL_TUMBLING_WINDOW_AGGREGATE: &str = "sql-tumbling-window-aggregate"; + pub const SQL_SLIDING_WINDOW_AGGREGATE: &str = "sql-sliding-window-aggregate"; + pub const SQL_SESSION_WINDOW_AGGREGATE: &str = "sql-session-window-aggregate"; + pub const SQL_UPDATING_AGGREGATE: &str = "sql-updating-aggregate"; + pub const KEY_BY_ROUTING: &str = "key-by-routing"; + pub const CONNECTOR_SOURCE: &str = "connector-source"; + pub const CONNECTOR_SINK: &str = "connector-sink"; +} + +pub mod extension_node { + pub const STREAM_WINDOW_AGGREGATE: &str = "StreamWindowAggregateNode"; + pub const STREAMING_WINDOW_FUNCTION: &str = "StreamingWindowFunctionNode"; + pub const EVENT_TIME_WATERMARK: &str = "EventTimeWatermarkNode"; + pub const CONTINUOUS_AGGREGATE: &str = "ContinuousAggregateNode"; + pub const SYSTEM_TIMESTAMP_INJECTOR: &str = "SystemTimestampInjectorNode"; + pub const STREAM_INGESTION: &str = "StreamIngestionNode"; + pub const STREAM_EGRESS: &str = "StreamEgressNode"; + pub const STREAM_PROJECTION: &str = "StreamProjectionNode"; + pub const REMOTE_TABLE_BOUNDARY: &str = "RemoteTableBoundaryNode"; + pub const REFERENCE_TABLE_SOURCE: &str = "ReferenceTableSource"; + pub const STREAM_REFERENCE_JOIN: &str = "StreamReferenceJoin"; + pub const KEY_EXTRACTION: &str = "KeyExtractionNode"; + pub const STREAMING_JOIN: &str = "StreamingJoinNode"; + pub const ASYNC_FUNCTION_EXECUTION: &str = "AsyncFunctionExecutionNode"; + pub const UNROLL_DEBEZIUM_PAYLOAD: &str = "UnrollDebeziumPayloadNode"; + pub const PACK_DEBEZIUM_ENVELOPE: &str = "PackDebeziumEnvelopeNode"; +} + +pub mod proto_operator_name { + pub const TUMBLING_WINDOW: &str = "TumblingWindow"; + pub const UPDATING_AGGREGATE: &str = "UpdatingAggregate"; + pub const WINDOW_FUNCTION: &str = "WindowFunction"; + pub const SLIDING_WINDOW_LABEL: &str = "sliding window"; + pub const INSTANT_WINDOW: &str = "InstantWindow"; + pub const INSTANT_WINDOW_LABEL: &str = "instant window"; +} + +pub mod runtime_operator_kind { + pub const STREAMING_JOIN: &str = "streaming_join"; + pub const WATERMARK_GENERATOR: &str = "watermark_generator"; + pub const STREAMING_WINDOW_EVALUATOR: &str = "streaming_window_evaluator"; +} + +pub mod factory_operator_name { + pub const CONNECTOR_SOURCE: &str = "ConnectorSource"; + pub const CONNECTOR_SINK: &str = "ConnectorSink"; + pub const KAFKA_SOURCE: &str = "KafkaSource"; + pub const KAFKA_SINK: &str = "KafkaSink"; +} + +pub mod cdc { + pub const BEFORE: &str = "before"; + pub const AFTER: &str = "after"; + pub const OP: &str = "op"; +} + +pub mod updating_state_field { + pub const IS_RETRACT: &str = "is_retract"; + pub const ID: &str = "id"; +} + +pub mod sql_field { + pub const ASYNC_RESULT: &str = "__async_result"; + pub const DEFAULT_KEY_LABEL: &str = "key"; + pub const DEFAULT_PROJECTION_LABEL: &str = "projection"; + pub const COMPUTED_WATERMARK: &str = "__watermark"; + pub const TIMESTAMP_FIELD: &str = "_timestamp"; + pub const UPDATING_META_FIELD: &str = "_updating_meta"; +} + +pub mod sql_planning_default { + pub const DEFAULT_PARALLELISM: usize = 1; + /// Default physical parallelism for `KeyBy` / key-extraction pipelines (configurable via YAML). + pub const DEFAULT_KEY_BY_PARALLELISM: usize = 1; + /// Parallelism for aggregations that run after `KeyBy` / shuffle on non-empty routing keys. + pub const KEYED_AGGREGATE_DEFAULT_PARALLELISM: usize = 8; + pub const PLANNING_TTL_SECS: u64 = 24 * 60 * 60; +} + +pub mod with_opt_bool_str { + pub const TRUE: &str = "true"; + pub const YES: &str = "yes"; + pub const FALSE: &str = "false"; + pub const NO: &str = "no"; +} + +pub mod interval_duration_unit { + pub const SECOND: &str = "second"; + pub const SECONDS: &str = "seconds"; + pub const S: &str = "s"; + pub const MINUTE: &str = "minute"; + pub const MINUTES: &str = "minutes"; + pub const MIN: &str = "min"; + pub const HOUR: &str = "hour"; + pub const HOURS: &str = "hours"; + pub const H: &str = "h"; + pub const DAY: &str = "day"; + pub const DAYS: &str = "days"; + pub const D: &str = "d"; +} + +pub mod connection_format_value { + pub const JSON: &str = "json"; + pub const CSV: &str = "csv"; + pub const JSONL: &str = "jsonl"; + pub const NDJSON: &str = "ndjson"; + pub const LANCE: &str = "lance"; + pub const ORC: &str = "orc"; + pub const DEBEZIUM_JSON: &str = "debezium_json"; + pub const AVRO: &str = "avro"; + pub const PARQUET: &str = "parquet"; + pub const PROTOBUF: &str = "protobuf"; + pub const RAW_STRING: &str = "raw_string"; + pub const RAW_BYTES: &str = "raw_bytes"; +} + +pub mod framing_method_value { + pub const NEWLINE: &str = "newline"; + pub const NEWLINE_DELIMITED: &str = "newline_delimited"; +} + +pub mod bad_data_value { + pub const FAIL: &str = "fail"; + pub const DROP: &str = "drop"; +} + +pub mod timestamp_format_value { + pub const RFC3339_SNAKE: &str = "rfc3339"; + pub const RFC3339_UPPER: &str = "RFC3339"; + pub const UNIX_MILLIS_SNAKE: &str = "unix_millis"; + pub const UNIX_MILLIS_PASCAL: &str = "UnixMillis"; +} + +pub mod decimal_encoding_value { + pub const NUMBER: &str = "number"; + pub const STRING: &str = "string"; + pub const BYTES: &str = "bytes"; +} + +pub mod json_compression_value { + pub const UNCOMPRESSED: &str = "uncompressed"; + pub const GZIP: &str = "gzip"; +} + +pub mod parquet_compression_value { + pub const UNCOMPRESSED: &str = "uncompressed"; + pub const SNAPPY: &str = "snappy"; + pub const GZIP: &str = "gzip"; + pub const ZSTD: &str = "zstd"; + pub const LZ4: &str = "lz4"; + pub const LZ4_RAW: &str = "lz4_raw"; +} + +pub mod date_part_keyword { + pub const YEAR: &str = "year"; + pub const MONTH: &str = "month"; + pub const WEEK: &str = "week"; + pub const DAY: &str = "day"; + pub const HOUR: &str = "hour"; + pub const MINUTE: &str = "minute"; + pub const SECOND: &str = "second"; + pub const MILLISECOND: &str = "millisecond"; + pub const MICROSECOND: &str = "microsecond"; + pub const NANOSECOND: &str = "nanosecond"; + pub const DOW: &str = "dow"; + pub const DOY: &str = "doy"; +} + +pub mod date_trunc_keyword { + pub const YEAR: &str = "year"; + pub const QUARTER: &str = "quarter"; + pub const MONTH: &str = "month"; + pub const WEEK: &str = "week"; + pub const DAY: &str = "day"; + pub const HOUR: &str = "hour"; + pub const MINUTE: &str = "minute"; + pub const SECOND: &str = "second"; +} + +pub mod mem_exec_join_side { + pub const LEFT: &str = "left"; + pub const RIGHT: &str = "right"; +} + +pub mod physical_plan_node_name { + pub const RW_LOCK_READER: &str = "rw_lock_reader"; + pub const UNBOUNDED_READER: &str = "unbounded_reader"; + pub const VEC_READER: &str = "vec_reader"; + pub const MEM_EXEC: &str = "mem_exec"; + pub const DEBEZIUM_UNROLLING_EXEC: &str = "debezium_unrolling_exec"; + pub const TO_DEBEZIUM_EXEC: &str = "to_debezium_exec"; +} + +pub mod window_function_udf { + pub const NAME: &str = "window"; +} + +pub mod window_interval_field { + pub const START: &str = "start"; + pub const END: &str = "end"; +} + +pub mod debezium_op_short { + pub const CREATE: &str = "c"; + pub const READ: &str = "r"; + pub const UPDATE: &str = "u"; + pub const DELETE: &str = "d"; +} + +pub mod connector_type { + pub const KAFKA: &str = "kafka"; + pub const KINESIS: &str = "kinesis"; + pub const FILESYSTEM: &str = "filesystem"; + pub const DELTA: &str = "delta"; + pub const ICEBERG: &str = "iceberg"; + pub const LANCE_DB: &str = "lanceDB"; + pub const S3: &str = "s3"; + pub const PULSAR: &str = "pulsar"; + pub const NATS: &str = "nats"; + pub const REDIS: &str = "redis"; + pub const MQTT: &str = "mqtt"; + pub const WEBSOCKET: &str = "websocket"; + pub const SSE: &str = "sse"; + pub const NEXMARK: &str = "nexmark"; + pub const BLACKHOLE: &str = "blackhole"; + pub const MEMORY: &str = "memory"; + pub const POSTGRES: &str = "postgres"; +} + +pub mod connection_table_role { + pub const SOURCE: &str = "source"; + pub const SINK: &str = "sink"; + pub const LOOKUP: &str = "lookup"; +} + +pub const SUPPORTED_CONNECTOR_ADAPTERS: &[&str] = &[ + connector_type::KAFKA, + connector_type::FILESYSTEM, + connector_type::S3, + connector_type::DELTA, + connector_type::ICEBERG, + connector_type::LANCE_DB, +]; + +pub mod kafka_with_value { + pub const SCAN_LATEST: &str = "latest"; + pub const SCAN_EARLIEST: &str = "earliest"; + pub const SCAN_GROUP_OFFSETS: &str = "group-offsets"; + pub const SCAN_GROUP: &str = "group"; + pub const ISOLATION_READ_COMMITTED: &str = "read_committed"; + pub const ISOLATION_READ_UNCOMMITTED: &str = "read_uncommitted"; + pub const SINK_COMMIT_EXACTLY_ONCE_HYPHEN: &str = "exactly-once"; + pub const SINK_COMMIT_EXACTLY_ONCE_UNDERSCORE: &str = "exactly_once"; + pub const SINK_COMMIT_AT_LEAST_ONCE_HYPHEN: &str = "at-least-once"; + pub const SINK_COMMIT_AT_LEAST_ONCE_UNDERSCORE: &str = "at_least_once"; +} diff --git a/src/streaming_planner/src/common/control.rs b/src/streaming_planner/src/common/control.rs new file mode 100644 index 00000000..59ef409e --- /dev/null +++ b/src/streaming_planner/src/common/control.rs @@ -0,0 +1,164 @@ +// 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::time::SystemTime; + +use function_stream_runtime_common::streaming_protocol::CheckpointBarrier; + +/// Control messages sent from the controller to worker tasks. +#[derive(Debug, Clone)] +pub enum ControlMessage { + Checkpoint(CheckpointBarrier), + Stop { + mode: StopMode, + }, + Commit { + epoch: u32, + commit_data: HashMap>>, + }, + LoadCompacted { + compacted: CompactionResult, + }, + NoOp, +} + +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub enum StopMode { + Graceful, + Immediate, +} + +#[derive(Debug, Clone)] +pub struct CompactionResult { + pub operator_id: String, + pub compacted_tables: HashMap, +} + +#[derive(Debug, Clone)] +pub struct TableCheckpointMetadata { + pub table_type: TableType, + pub data: Vec, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TableType { + GlobalKeyValue, + ExpiringKeyedTimeTable, +} + +/// Responses sent from worker tasks back to the controller. +#[derive(Debug, Clone)] +pub enum ControlResp { + CheckpointEvent(CheckpointEvent), + CheckpointCompleted(CheckpointCompleted), + TaskStarted { + node_id: u32, + task_index: usize, + start_time: SystemTime, + }, + TaskFinished { + node_id: u32, + task_index: usize, + }, + TaskFailed { + node_id: u32, + task_index: usize, + error: TaskError, + }, + Error { + node_id: u32, + operator_id: String, + task_index: usize, + message: String, + details: String, + }, +} + +#[derive(Debug, Clone)] +pub struct CheckpointCompleted { + pub checkpoint_epoch: u32, + pub node_id: u32, + pub operator_id: String, + pub subtask_metadata: SubtaskCheckpointMetadata, +} + +#[derive(Debug, Clone)] +pub struct SubtaskCheckpointMetadata { + pub subtask_index: u32, + pub start_time: u64, + pub finish_time: u64, + pub watermark: Option, + pub bytes: u64, + pub table_metadata: HashMap, + pub table_configs: HashMap, +} + +#[derive(Debug, Clone)] +pub struct TableSubtaskCheckpointMetadata { + pub subtask_index: u32, + pub table_type: TableType, + pub data: Vec, +} + +#[derive(Debug, Clone)] +pub struct TableConfig { + pub table_type: TableType, + pub config: Vec, + pub state_version: u32, +} + +#[derive(Debug, Clone)] +pub struct CheckpointEvent { + pub checkpoint_epoch: u32, + pub node_id: u32, + pub operator_id: String, + pub subtask_index: u32, + pub time: SystemTime, + pub event_type: TaskCheckpointEventType, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TaskCheckpointEventType { + StartedAlignment, + StartedCheckpointing, + FinishedOperatorSetup, + FinishedSync, + FinishedCommit, +} + +#[derive(Debug, Clone)] +pub struct TaskError { + pub job_id: String, + pub node_id: u32, + pub operator_id: String, + pub operator_subtask: u64, + pub error: String, + pub error_domain: ErrorDomain, + pub retry_hint: RetryHint, + pub details: String, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ErrorDomain { + User, + Internal, + External, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RetryHint { + NoRetry, + WithBackoff, +} diff --git a/src/streaming_planner/src/common/converter.rs b/src/streaming_planner/src/common/converter.rs new file mode 100644 index 00000000..a9023342 --- /dev/null +++ b/src/streaming_planner/src/common/converter.rs @@ -0,0 +1,95 @@ +// 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 arrow::row::{OwnedRow, RowConverter, RowParser, Rows, SortField}; +use arrow_array::{Array, ArrayRef, BooleanArray}; +use arrow_schema::{ArrowError, DataType}; +use std::sync::Arc; + +// need to handle the empty case as a row converter without sort fields emits empty Rows. +#[derive(Debug)] +pub enum Converter { + RowConverter(RowConverter), + Empty(RowConverter, Arc), +} + +impl Converter { + pub fn new(sort_fields: Vec) -> Result { + if sort_fields.is_empty() { + let array = Arc::new(BooleanArray::from(vec![false])); + Ok(Self::Empty( + RowConverter::new(vec![SortField::new(DataType::Boolean)])?, + array, + )) + } else { + Ok(Self::RowConverter(RowConverter::new(sort_fields)?)) + } + } + + pub fn convert_columns(&self, columns: &[Arc]) -> Result { + match self { + Converter::RowConverter(row_converter) => { + Ok(row_converter.convert_columns(columns)?.row(0).owned()) + } + Converter::Empty(row_converter, array) => Ok(row_converter + .convert_columns(std::slice::from_ref(array))? + .row(0) + .owned()), + } + } + + pub fn convert_all_columns( + &self, + columns: &[Arc], + num_rows: usize, + ) -> Result { + match self { + Converter::RowConverter(row_converter) => Ok(row_converter.convert_columns(columns)?), + Converter::Empty(row_converter, _array) => { + let array = Arc::new(BooleanArray::from(vec![false; num_rows])); + Ok(row_converter.convert_columns(&[array])?) + } + } + } + + pub fn convert_rows( + &self, + rows: Vec>, + ) -> Result, ArrowError> { + match self { + Converter::RowConverter(row_converter) => Ok(row_converter.convert_rows(rows)?), + Converter::Empty(_row_converter, _array) => Ok(vec![]), + } + } + + pub fn convert_raw_rows(&self, row_bytes: Vec<&[u8]>) -> Result, ArrowError> { + match self { + Converter::RowConverter(row_converter) => { + let parser = row_converter.parser(); + let mut row_list = vec![]; + for bytes in row_bytes { + let row = parser.parse(bytes); + row_list.push(row); + } + Ok(row_converter.convert_rows(row_list)?) + } + Converter::Empty(_row_converter, _array) => Ok(vec![]), + } + } + + pub fn parser(&self) -> Option { + match self { + Converter::RowConverter(r) => Some(r.parser()), + Converter::Empty(_, _) => None, + } + } +} diff --git a/src/streaming_planner/src/common/date.rs b/src/streaming_planner/src/common/date.rs new file mode 100644 index 00000000..ec310326 --- /dev/null +++ b/src/streaming_planner/src/common/date.rs @@ -0,0 +1,86 @@ +// 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 serde::Serialize; +use std::convert::TryFrom; + +use super::constants::{date_part_keyword, date_trunc_keyword}; + +#[derive(Debug, Clone, Copy, Eq, PartialEq, PartialOrd, Hash, Serialize)] +pub enum DatePart { + Year, + Month, + Week, + Day, + Hour, + Minute, + Second, + Millisecond, + Microsecond, + Nanosecond, + DayOfWeek, + DayOfYear, +} + +impl TryFrom<&str> for DatePart { + type Error = String; + + fn try_from(value: &str) -> Result { + let v = value.to_lowercase(); + match v.as_str() { + date_part_keyword::YEAR => Ok(DatePart::Year), + date_part_keyword::MONTH => Ok(DatePart::Month), + date_part_keyword::WEEK => Ok(DatePart::Week), + date_part_keyword::DAY => Ok(DatePart::Day), + date_part_keyword::HOUR => Ok(DatePart::Hour), + date_part_keyword::MINUTE => Ok(DatePart::Minute), + date_part_keyword::SECOND => Ok(DatePart::Second), + date_part_keyword::MILLISECOND => Ok(DatePart::Millisecond), + date_part_keyword::MICROSECOND => Ok(DatePart::Microsecond), + date_part_keyword::NANOSECOND => Ok(DatePart::Nanosecond), + date_part_keyword::DOW => Ok(DatePart::DayOfWeek), + date_part_keyword::DOY => Ok(DatePart::DayOfYear), + _ => Err(format!("'{value}' is not a valid DatePart")), + } + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, PartialOrd, Serialize)] +pub enum DateTruncPrecision { + Year, + Quarter, + Month, + Week, + Day, + Hour, + Minute, + Second, +} + +impl TryFrom<&str> for DateTruncPrecision { + type Error = String; + + fn try_from(value: &str) -> Result { + let v = value.to_lowercase(); + match v.as_str() { + date_trunc_keyword::YEAR => Ok(DateTruncPrecision::Year), + date_trunc_keyword::QUARTER => Ok(DateTruncPrecision::Quarter), + date_trunc_keyword::MONTH => Ok(DateTruncPrecision::Month), + date_trunc_keyword::WEEK => Ok(DateTruncPrecision::Week), + date_trunc_keyword::DAY => Ok(DateTruncPrecision::Day), + date_trunc_keyword::HOUR => Ok(DateTruncPrecision::Hour), + date_trunc_keyword::MINUTE => Ok(DateTruncPrecision::Minute), + date_trunc_keyword::SECOND => Ok(DateTruncPrecision::Second), + _ => Err(format!("'{value}' is not a valid DateTruncPrecision")), + } + } +} diff --git a/src/streaming_planner/src/common/debezium.rs b/src/streaming_planner/src/common/debezium.rs new file mode 100644 index 00000000..9dbc401f --- /dev/null +++ b/src/streaming_planner/src/common/debezium.rs @@ -0,0 +1,148 @@ +// 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 bincode::{Decode, Encode}; +use serde::{Deserialize, Serialize}; +use std::convert::TryFrom; +use std::fmt::Debug; + +pub trait Key: + Debug + Clone + Encode + Decode<()> + std::hash::Hash + PartialEq + Eq + Send + 'static +{ +} +impl + std::hash::Hash + PartialEq + Eq + Send + 'static> Key + for T +{ +} + +pub trait Data: Debug + Clone + Encode + Decode<()> + Send + 'static {} +impl + Send + 'static> Data for T {} + +#[derive(Debug, Clone, PartialEq, Encode, Decode, Serialize, Deserialize)] +pub enum UpdatingData { + Retract(T), + Update { old: T, new: T }, + Append(T), +} + +impl UpdatingData { + pub fn lower(&self) -> T { + match self { + UpdatingData::Retract(_) => panic!("cannot lower retractions"), + UpdatingData::Update { new, .. } => new.clone(), + UpdatingData::Append(t) => t.clone(), + } + } + + pub fn unwrap_append(&self) -> &T { + match self { + UpdatingData::Append(t) => t, + _ => panic!("UpdatingData is not an append"), + } + } +} + +#[derive(Clone, Encode, Decode, Debug, Serialize, Deserialize, PartialEq)] +#[serde(try_from = "DebeziumShadow")] +pub struct Debezium { + pub before: Option, + pub after: Option, + pub op: DebeziumOp, +} + +#[derive(Clone, Encode, Decode, Debug, Serialize, Deserialize, PartialEq)] +struct DebeziumShadow { + before: Option, + after: Option, + op: DebeziumOp, +} + +impl TryFrom> for Debezium { + type Error = &'static str; + + fn try_from(value: DebeziumShadow) -> Result { + match (value.op, &value.before, &value.after) { + (DebeziumOp::Create, _, None) => { + Err("`after` must be set for Debezium create messages") + } + (DebeziumOp::Update, None, _) => { + Err("`before` must be set for Debezium update messages") + } + (DebeziumOp::Update, _, None) => { + Err("`after` must be set for Debezium update messages") + } + (DebeziumOp::Delete, None, _) => { + Err("`before` must be set for Debezium delete messages") + } + _ => Ok(Debezium { + before: value.before, + after: value.after, + op: value.op, + }), + } + } +} + +#[derive(Copy, Clone, Encode, Decode, Debug, PartialEq)] +pub enum DebeziumOp { + Create, + Update, + Delete, +} + +#[allow(clippy::to_string_trait_impl)] +impl ToString for DebeziumOp { + fn to_string(&self) -> String { + match self { + DebeziumOp::Create => "c", + DebeziumOp::Update => "u", + DebeziumOp::Delete => "d", + } + .to_string() + } +} + +impl<'de> Deserialize<'de> for DebeziumOp { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + let s = String::deserialize(deserializer)?; + match s.as_str() { + "c" | "r" => Ok(DebeziumOp::Create), + "u" => Ok(DebeziumOp::Update), + "d" => Ok(DebeziumOp::Delete), + _ => Err(serde::de::Error::custom(format!("Invalid DebeziumOp {s}"))), + } + } +} + +impl Serialize for DebeziumOp { + fn serialize(&self, serializer: S) -> Result + where + S: serde::Serializer, + { + match self { + DebeziumOp::Create => serializer.serialize_str("c"), + DebeziumOp::Update => serializer.serialize_str("u"), + DebeziumOp::Delete => serializer.serialize_str("d"), + } + } +} + +#[derive(Copy, Clone, Encode, Decode, Debug, PartialEq, Serialize, Deserialize)] +pub enum JoinType { + Inner, + Left, + Right, + Full, +} diff --git a/src/streaming_planner/src/common/errors.rs b/src/streaming_planner/src/common/errors.rs new file mode 100644 index 00000000..b7e97400 --- /dev/null +++ b/src/streaming_planner/src/common/errors.rs @@ -0,0 +1,92 @@ +// 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::fmt; + +/// Result type for streaming operators and collectors. +pub type DataflowResult = std::result::Result; + +/// Unified error type for streaming dataflow operations. +#[derive(Debug)] +pub enum DataflowError { + Arrow(arrow_schema::ArrowError), + DataFusion(datafusion::error::DataFusionError), + Operator(String), + State(String), + Connector(String), + Internal(String), +} + +impl fmt::Display for DataflowError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + DataflowError::Arrow(e) => write!(f, "Arrow error: {e}"), + DataflowError::DataFusion(e) => write!(f, "DataFusion error: {e}"), + DataflowError::Operator(msg) => write!(f, "Operator error: {msg}"), + DataflowError::State(msg) => write!(f, "State error: {msg}"), + DataflowError::Connector(msg) => write!(f, "Connector error: {msg}"), + DataflowError::Internal(msg) => write!(f, "Internal error: {msg}"), + } + } +} + +impl std::error::Error for DataflowError {} + +impl DataflowError { + pub fn with_operator(self, operator_id: impl Into) -> Self { + let id = operator_id.into(); + match self { + DataflowError::Operator(m) => DataflowError::Operator(format!("{id}: {m}")), + other => DataflowError::Operator(format!("{id}: {other}")), + } + } +} + +impl From for DataflowError { + fn from(e: arrow_schema::ArrowError) -> Self { + DataflowError::Arrow(e) + } +} + +impl From for DataflowError { + fn from(e: datafusion::error::DataFusionError) -> Self { + DataflowError::DataFusion(e) + } +} + +/// Macro for creating connector errors. +#[macro_export] +macro_rules! connector_err { + ($($arg:tt)*) => { + $crate::common::errors::DataflowError::Connector(format!($($arg)*)) + }; +} + +/// State-related errors. +#[derive(Debug)] +pub enum StateError { + KeyNotFound(String), + SerializationError(String), + BackendError(String), +} + +impl fmt::Display for StateError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + StateError::KeyNotFound(key) => write!(f, "Key not found: {key}"), + StateError::SerializationError(msg) => write!(f, "Serialization error: {msg}"), + StateError::BackendError(msg) => write!(f, "State backend error: {msg}"), + } + } +} + +impl std::error::Error for StateError {} diff --git a/src/streaming_planner/src/common/format_from_opts.rs b/src/streaming_planner/src/common/format_from_opts.rs new file mode 100644 index 00000000..ffd29572 --- /dev/null +++ b/src/streaming_planner/src/common/format_from_opts.rs @@ -0,0 +1,182 @@ +// 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 `WITH` clause format / framing / bad-data options (Arroyo-compatible keys). + +use std::str::FromStr; + +use datafusion::common::{Result as DFResult, plan_datafusion_err, plan_err}; + +use super::connector_options::ConnectorOptions; +use super::constants::{bad_data_value, connection_format_value, framing_method_value}; +use super::formats::{ + AvroFormat, BadData, CsvFormat, DecimalEncoding, Format, Framing, JsonCompression, JsonFormat, + LanceFormat, NewlineDelimitedFraming, ParquetCompression, ParquetFormat, ProtobufFormat, + RawBytesFormat, RawStringFormat, TimestampFormat, +}; +use super::with_option_keys as opt; + +impl JsonFormat { + pub fn from_opts(opts: &mut ConnectorOptions) -> DFResult { + let mut j = JsonFormat::default(); + if let Some(v) = opts.pull_opt_bool(opt::JSON_CONFLUENT_SCHEMA_REGISTRY)? { + j.confluent_schema_registry = v; + } + if let Some(v) = opts.pull_opt_u64(opt::JSON_CONFLUENT_SCHEMA_VERSION)? { + j.schema_id = Some(v as u32); + } + if let Some(v) = opts.pull_opt_bool(opt::JSON_INCLUDE_SCHEMA)? { + j.include_schema = v; + } + if let Some(v) = opts.pull_opt_bool(opt::JSON_DEBEZIUM)? { + j.debezium = v; + } + if let Some(v) = opts.pull_opt_bool(opt::JSON_UNSTRUCTURED)? { + j.unstructured = v; + } + if let Some(s) = opts.pull_opt_str(opt::JSON_TIMESTAMP_FORMAT)? { + j.timestamp_format = TimestampFormat::try_from(s.as_str()) + .map_err(|_| plan_datafusion_err!("invalid json.timestamp_format '{}'", s))?; + } + if let Some(s) = opts.pull_opt_str(opt::JSON_DECIMAL_ENCODING)? { + j.decimal_encoding = DecimalEncoding::try_from(s.as_str()) + .map_err(|_| plan_datafusion_err!("invalid json.decimal_encoding '{s}'"))?; + } + if let Some(s) = opts.pull_opt_str(opt::JSON_COMPRESSION)? { + j.compression = JsonCompression::from_str(&s) + .map_err(|e| plan_datafusion_err!("invalid json.compression: {e}"))?; + } + Ok(j) + } +} + +impl Format { + pub fn from_opts(opts: &mut ConnectorOptions) -> DFResult> { + let Some(name) = opts.peek_opt_str(opt::FORMAT)? else { + return Ok(None); + }; + let n = name.to_lowercase(); + match n.as_str() { + connection_format_value::JSON => Ok(Some(Format::Json(JsonFormat::from_opts(opts)?))), + connection_format_value::CSV => Ok(Some(Format::Csv(CsvFormat {}))), + connection_format_value::LANCE => Ok(Some(Format::Lance(LanceFormat {}))), + connection_format_value::DEBEZIUM_JSON => { + let mut j = JsonFormat::from_opts(opts)?; + j.debezium = true; + Ok(Some(Format::Json(j))) + } + connection_format_value::AVRO => Ok(Some(Format::Avro(AvroFormat::from_opts(opts)?))), + connection_format_value::PARQUET => { + Ok(Some(Format::Parquet(ParquetFormat::from_opts(opts)?))) + } + connection_format_value::PROTOBUF => { + Ok(Some(Format::Protobuf(ProtobufFormat::from_opts(opts)?))) + } + connection_format_value::RAW_STRING => Ok(Some(Format::RawString(RawStringFormat {}))), + connection_format_value::RAW_BYTES => Ok(Some(Format::RawBytes(RawBytesFormat {}))), + _ => plan_err!("unknown format '{name}'"), + } + } +} + +impl AvroFormat { + fn from_opts(opts: &mut ConnectorOptions) -> DFResult { + let mut a = AvroFormat { + confluent_schema_registry: false, + raw_datums: false, + into_unstructured_json: false, + schema_id: None, + }; + if let Some(v) = opts.pull_opt_bool(opt::AVRO_CONFLUENT_SCHEMA_REGISTRY)? { + a.confluent_schema_registry = v; + } + if let Some(v) = opts.pull_opt_bool(opt::AVRO_RAW_DATUMS)? { + a.raw_datums = v; + } + if let Some(v) = opts.pull_opt_bool(opt::AVRO_INTO_UNSTRUCTURED_JSON)? { + a.into_unstructured_json = v; + } + if let Some(v) = opts.pull_opt_u64(opt::AVRO_SCHEMA_ID)? { + a.schema_id = Some(v as u32); + } + Ok(a) + } +} + +impl ParquetFormat { + fn from_opts(opts: &mut ConnectorOptions) -> DFResult { + let mut p = ParquetFormat::default(); + if let Some(s) = opts.pull_opt_str(opt::PARQUET_COMPRESSION)? { + p.compression = ParquetCompression::from_str(&s) + .map_err(|e| plan_datafusion_err!("invalid parquet.compression: {e}"))?; + } + if let Some(v) = opts.pull_opt_u64(opt::PARQUET_ROW_GROUP_BYTES)? { + p.row_group_bytes = Some(v); + } + Ok(p) + } +} + +impl ProtobufFormat { + fn from_opts(opts: &mut ConnectorOptions) -> DFResult { + let mut p = ProtobufFormat { + into_unstructured_json: false, + message_name: None, + compiled_schema: None, + confluent_schema_registry: false, + length_delimited: false, + }; + if let Some(v) = opts.pull_opt_bool(opt::PROTOBUF_INTO_UNSTRUCTURED_JSON)? { + p.into_unstructured_json = v; + } + if let Some(s) = opts.pull_opt_str(opt::PROTOBUF_MESSAGE_NAME)? { + p.message_name = Some(s); + } + if let Some(v) = opts.pull_opt_bool(opt::PROTOBUF_CONFLUENT_SCHEMA_REGISTRY)? { + p.confluent_schema_registry = v; + } + if let Some(v) = opts.pull_opt_bool(opt::PROTOBUF_LENGTH_DELIMITED)? { + p.length_delimited = v; + } + Ok(p) + } +} + +impl Framing { + pub fn from_opts(opts: &mut ConnectorOptions) -> DFResult> { + let method = opts.pull_opt_str(opt::FRAMING_METHOD)?; + match method.as_deref() { + None => Ok(None), + Some(framing_method_value::NEWLINE) | Some(framing_method_value::NEWLINE_DELIMITED) => { + let max = opts.pull_opt_u64(opt::FRAMING_MAX_LINE_LENGTH)?; + Ok(Some(Framing::Newline(NewlineDelimitedFraming { + max_line_length: max, + }))) + } + Some(other) => plan_err!("unknown framing.method '{other}'"), + } + } +} + +impl BadData { + pub fn from_opts(opts: &mut ConnectorOptions) -> DFResult { + let Some(s) = opts.pull_opt_str(opt::BAD_DATA)? else { + return Ok(BadData::Fail {}); + }; + let v = s.to_lowercase(); + match v.as_str() { + bad_data_value::FAIL => Ok(BadData::Fail {}), + bad_data_value::DROP => Ok(BadData::Drop {}), + _ => plan_err!("invalid bad_data '{s}'"), + } + } +} diff --git a/src/streaming_planner/src/common/formats.rs b/src/streaming_planner/src/common/formats.rs new file mode 100644 index 00000000..a47d93cf --- /dev/null +++ b/src/streaming_planner/src/common/formats.rs @@ -0,0 +1,267 @@ +// 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 serde::{Deserialize, Serialize}; +use std::convert::TryFrom; +use std::fmt::{Display, Formatter}; +use std::str::FromStr; + +use super::constants::{ + connection_format_value, decimal_encoding_value, json_compression_value, + parquet_compression_value, timestamp_format_value, +}; + +#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Default, Hash, PartialOrd)] +#[serde(rename_all = "snake_case")] +pub enum TimestampFormat { + #[default] + #[serde(rename = "rfc3339")] + RFC3339, + UnixMillis, +} + +impl TryFrom<&str> for TimestampFormat { + type Error = (); + + fn try_from(value: &str) -> Result { + match value { + timestamp_format_value::RFC3339_UPPER | timestamp_format_value::RFC3339_SNAKE => { + Ok(TimestampFormat::RFC3339) + } + timestamp_format_value::UNIX_MILLIS_PASCAL + | timestamp_format_value::UNIX_MILLIS_SNAKE => Ok(TimestampFormat::UnixMillis), + _ => Err(()), + } + } +} + +#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Default, Hash, PartialOrd)] +#[serde(rename_all = "snake_case")] +pub enum DecimalEncoding { + #[default] + Number, + String, + Bytes, +} + +impl TryFrom<&str> for DecimalEncoding { + type Error = (); + + fn try_from(s: &str) -> Result { + match s { + decimal_encoding_value::NUMBER => Ok(Self::Number), + decimal_encoding_value::STRING => Ok(Self::String), + decimal_encoding_value::BYTES => Ok(Self::Bytes), + _ => Err(()), + } + } +} + +#[derive(Serialize, Deserialize, Default, Copy, Clone, Debug, PartialEq, Eq, Hash, PartialOrd)] +#[serde(rename_all = "snake_case")] +pub enum JsonCompression { + #[default] + Uncompressed, + Gzip, +} + +impl FromStr for JsonCompression { + type Err = String; + + fn from_str(s: &str) -> Result { + match s { + json_compression_value::UNCOMPRESSED => Ok(JsonCompression::Uncompressed), + json_compression_value::GZIP => Ok(JsonCompression::Gzip), + _ => Err(format!("invalid json compression '{s}'")), + } + } +} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Default, Hash, PartialOrd)] +#[serde(rename_all = "snake_case")] +pub struct JsonFormat { + #[serde(default)] + pub confluent_schema_registry: bool, + #[serde(default, alias = "confluent_schema_version")] + pub schema_id: Option, + #[serde(default)] + pub include_schema: bool, + #[serde(default)] + pub debezium: bool, + #[serde(default)] + pub unstructured: bool, + #[serde(default)] + pub timestamp_format: TimestampFormat, + #[serde(default)] + pub decimal_encoding: DecimalEncoding, + #[serde(default)] + pub compression: JsonCompression, +} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash, PartialOrd)] +#[serde(rename_all = "snake_case")] +pub struct RawStringFormat {} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash, PartialOrd)] +#[serde(rename_all = "snake_case")] +pub struct RawBytesFormat {} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash, PartialOrd)] +#[serde(rename_all = "snake_case")] +pub struct CsvFormat {} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash, PartialOrd)] +#[serde(rename_all = "snake_case")] +pub struct LanceFormat {} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash, PartialOrd)] +#[serde(rename_all = "snake_case")] +pub struct AvroFormat { + #[serde(default)] + pub confluent_schema_registry: bool, + #[serde(default)] + pub raw_datums: bool, + #[serde(default)] + pub into_unstructured_json: bool, + #[serde(default)] + pub schema_id: Option, +} + +impl AvroFormat { + pub fn new( + confluent_schema_registry: bool, + raw_datums: bool, + into_unstructured_json: bool, + ) -> Self { + Self { + confluent_schema_registry, + raw_datums, + into_unstructured_json, + schema_id: None, + } + } +} + +#[derive(Serialize, Deserialize, Copy, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Default)] +#[serde(rename_all = "snake_case")] +pub enum ParquetCompression { + Uncompressed, + Snappy, + Gzip, + #[default] + Zstd, + Lz4, + Lz4Raw, +} + +impl FromStr for ParquetCompression { + type Err = String; + + fn from_str(s: &str) -> Result { + match s { + parquet_compression_value::UNCOMPRESSED => Ok(ParquetCompression::Uncompressed), + parquet_compression_value::SNAPPY => Ok(ParquetCompression::Snappy), + parquet_compression_value::GZIP => Ok(ParquetCompression::Gzip), + parquet_compression_value::ZSTD => Ok(ParquetCompression::Zstd), + parquet_compression_value::LZ4 => Ok(ParquetCompression::Lz4), + parquet_compression_value::LZ4_RAW => Ok(ParquetCompression::Lz4Raw), + _ => Err(format!("invalid parquet compression '{s}'")), + } + } +} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash, PartialOrd, Default)] +#[serde(rename_all = "snake_case")] +pub struct ParquetFormat { + #[serde(default)] + pub compression: ParquetCompression, + #[serde(default)] + pub row_group_bytes: Option, +} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash, PartialOrd)] +#[serde(rename_all = "snake_case")] +pub struct ProtobufFormat { + #[serde(default)] + pub into_unstructured_json: bool, + #[serde(default)] + pub message_name: Option, + #[serde(default)] + pub compiled_schema: Option>, + #[serde(default)] + pub confluent_schema_registry: bool, + #[serde(default)] + pub length_delimited: bool, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, Hash, PartialOrd)] +#[serde(rename_all = "snake_case", tag = "type")] +pub enum Format { + Json(JsonFormat), + Csv(CsvFormat), + Lance(LanceFormat), + Avro(AvroFormat), + Protobuf(ProtobufFormat), + Parquet(ParquetFormat), + RawString(RawStringFormat), + RawBytes(RawBytesFormat), +} + +impl Display for Format { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str(self.name()) + } +} + +impl Format { + pub fn name(&self) -> &'static str { + match self { + Format::Json(_) => connection_format_value::JSON, + Format::Csv(_) => connection_format_value::CSV, + Format::Lance(_) => connection_format_value::LANCE, + Format::Avro(_) => connection_format_value::AVRO, + Format::Protobuf(_) => connection_format_value::PROTOBUF, + Format::Parquet(_) => connection_format_value::PARQUET, + Format::RawString(_) => connection_format_value::RAW_STRING, + Format::RawBytes(_) => connection_format_value::RAW_BYTES, + } + } + + pub fn is_updating(&self) -> bool { + matches!(self, Format::Json(JsonFormat { debezium: true, .. })) + } +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, Hash, PartialOrd)] +#[serde(rename_all = "snake_case", tag = "behavior")] +pub enum BadData { + Fail {}, + Drop {}, +} + +impl Default for BadData { + fn default() -> Self { + BadData::Fail {} + } +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, Hash, PartialOrd)] +#[serde(rename_all = "snake_case", tag = "method")] +pub enum Framing { + Newline(NewlineDelimitedFraming), +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, Hash, PartialOrd)] +#[serde(rename_all = "snake_case")] +pub struct NewlineDelimitedFraming { + pub max_line_length: Option, +} diff --git a/src/streaming_planner/src/common/fs_schema.rs b/src/streaming_planner/src/common/fs_schema.rs new file mode 100644 index 00000000..9c548f69 --- /dev/null +++ b/src/streaming_planner/src/common/fs_schema.rs @@ -0,0 +1,470 @@ +// 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 super::{TIMESTAMP_FIELD, to_nanos}; +use crate::common::converter::Converter; +use arrow::compute::kernels::cmp::gt_eq; +use arrow::compute::kernels::numeric::div; +use arrow::compute::{SortColumn, filter_record_batch, lexsort_to_indices, partition, take}; +use arrow::row::SortField; +use arrow_array::types::UInt64Type; +use arrow_array::{PrimitiveArray, UInt64Array}; +use datafusion::arrow::array::builder::{ArrayBuilder, make_builder}; +use datafusion::arrow::array::{RecordBatch, TimestampNanosecondArray}; +use datafusion::arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaBuilder, TimeUnit}; +use datafusion::arrow::error::ArrowError; +use datafusion::common::{DataFusionError, Result as DFResult}; +use protocol::function_stream_graph; +use serde::{Deserialize, Serialize}; +use std::ops::Range; +use std::sync::Arc; +use std::time::SystemTime; + +#[derive(Debug, Copy, Clone)] +pub enum FieldValueType<'a> { + Int64(Option), + UInt64(Option), + Int32(Option), + String(Option<&'a str>), + Bytes(Option<&'a [u8]>), +} + +pub type FsSchemaRef = Arc; + +#[derive(Debug, Clone, Eq, PartialEq, Hash, Serialize, Deserialize)] +pub struct FsSchema { + pub schema: Arc, + pub timestamp_index: usize, + key_indices: Option>, + /// If defined, these indices are used for routing (i.e., which subtask gets which piece of data) + routing_key_indices: Option>, +} + +impl TryFrom for FsSchema { + type Error = DataFusionError; + fn try_from(schema_proto: function_stream_graph::FsSchema) -> Result { + let schema: Schema = serde_json::from_str(&schema_proto.arrow_schema) + .map_err(|e| DataFusionError::Plan(format!("Invalid arrow schema: {e}")))?; + let timestamp_index = schema_proto.timestamp_index as usize; + + let key_indices = schema_proto.has_keys.then(|| { + schema_proto + .key_indices + .into_iter() + .map(|index| index as usize) + .collect() + }); + + let routing_key_indices = schema_proto.has_routing_keys.then(|| { + schema_proto + .routing_key_indices + .into_iter() + .map(|index| index as usize) + .collect() + }); + + Ok(Self { + schema: Arc::new(schema), + timestamp_index, + key_indices, + routing_key_indices, + }) + } +} + +impl From for function_stream_graph::FsSchema { + fn from(schema: FsSchema) -> Self { + let arrow_schema = serde_json::to_string(schema.schema.as_ref()).unwrap(); + let timestamp_index = schema.timestamp_index as u32; + + let has_keys = schema.key_indices.is_some(); + let key_indices = schema + .key_indices + .map(|ks| ks.into_iter().map(|index| index as u32).collect()) + .unwrap_or_default(); + + let has_routing_keys = schema.routing_key_indices.is_some(); + let routing_key_indices = schema + .routing_key_indices + .map(|ks| ks.into_iter().map(|index| index as u32).collect()) + .unwrap_or_default(); + + Self { + arrow_schema, + timestamp_index, + key_indices, + has_keys, + routing_key_indices, + has_routing_keys, + } + } +} + +impl FsSchema { + pub fn new( + schema: Arc, + timestamp_index: usize, + key_indices: Option>, + routing_key_indices: Option>, + ) -> Self { + Self { + schema, + timestamp_index, + key_indices, + routing_key_indices, + } + } + pub fn new_unkeyed(schema: Arc, timestamp_index: usize) -> Self { + Self { + schema, + timestamp_index, + key_indices: None, + routing_key_indices: None, + } + } + pub fn new_keyed(schema: Arc, timestamp_index: usize, key_indices: Vec) -> Self { + Self { + schema, + timestamp_index, + key_indices: Some(key_indices), + routing_key_indices: None, + } + } + + pub fn from_fields(mut fields: Vec) -> Self { + if !fields.iter().any(|f| f.name() == TIMESTAMP_FIELD) { + fields.push(Field::new( + TIMESTAMP_FIELD, + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + )); + } + + Self::from_schema_keys(Arc::new(Schema::new(fields)), vec![]).unwrap() + } + + pub fn from_schema_unkeyed(schema: Arc) -> DFResult { + let timestamp_index = schema + .column_with_name(TIMESTAMP_FIELD) + .ok_or_else(|| { + DataFusionError::Plan(format!( + "no {TIMESTAMP_FIELD} field in schema, schema is {schema:?}" + )) + })? + .0; + + Ok(Self { + schema, + timestamp_index, + key_indices: None, + routing_key_indices: None, + }) + } + + pub fn from_schema_keys(schema: Arc, key_indices: Vec) -> DFResult { + let timestamp_index = schema + .column_with_name(TIMESTAMP_FIELD) + .ok_or_else(|| { + DataFusionError::Plan(format!( + "no {TIMESTAMP_FIELD} field in schema, schema is {schema:?}" + )) + })? + .0; + + Ok(Self { + schema, + timestamp_index, + key_indices: Some(key_indices), + routing_key_indices: None, + }) + } + + pub fn schema_without_timestamp(&self) -> Schema { + let mut builder = SchemaBuilder::from(self.schema.fields()); + builder.remove(self.timestamp_index); + builder.finish() + } + + pub fn remove_timestamp_column(&self, batch: &mut RecordBatch) { + batch.remove_column(self.timestamp_index); + } + + pub fn builders(&self) -> Vec> { + self.schema + .fields + .iter() + .map(|f| make_builder(f.data_type(), 8)) + .collect() + } + + pub fn timestamp_column<'a>(&self, batch: &'a RecordBatch) -> &'a TimestampNanosecondArray { + batch + .column(self.timestamp_index) + .as_any() + .downcast_ref::() + .unwrap() + } + + pub fn has_routing_keys(&self) -> bool { + self.routing_keys().map(|k| !k.is_empty()).unwrap_or(false) + } + + pub fn routing_keys(&self) -> Option<&Vec> { + self.routing_key_indices + .as_ref() + .or(self.key_indices.as_ref()) + } + + pub fn storage_keys(&self) -> Option<&Vec> { + self.key_indices.as_ref() + } + + pub fn clone_storage_key_indices(&self) -> Option> { + self.key_indices.clone() + } + + pub fn clone_routing_key_indices(&self) -> Option> { + self.routing_key_indices.clone() + } + + pub fn filter_by_time( + &self, + batch: RecordBatch, + cutoff: Option, + ) -> Result { + let Some(cutoff) = cutoff else { + // no watermark, so we just return the same batch. + return Ok(batch); + }; + // filter out late data + let timestamp_column = batch + .column(self.timestamp_index) + .as_any() + .downcast_ref::() + .ok_or_else(|| ArrowError::CastError( + format!("failed to downcast column {} of {:?} to timestamp. Schema is supposed to be {:?}", + self.timestamp_index, batch, self.schema)))?; + let cutoff_scalar = TimestampNanosecondArray::new_scalar(to_nanos(cutoff) as i64); + let on_time = gt_eq(timestamp_column, &cutoff_scalar)?; + filter_record_batch(&batch, &on_time) + } + + pub fn sort_columns(&self, batch: &RecordBatch, with_timestamp: bool) -> Vec { + let mut columns = vec![]; + if let Some(keys) = &self.key_indices { + columns.extend(keys.iter().map(|index| SortColumn { + values: batch.column(*index).clone(), + options: None, + })); + } + if with_timestamp { + columns.push(SortColumn { + values: batch.column(self.timestamp_index).clone(), + options: None, + }); + } + columns + } + + pub fn sort_fields(&self, with_timestamp: bool) -> Vec { + let mut sort_fields = vec![]; + if let Some(keys) = &self.key_indices { + sort_fields.extend(keys.iter()); + } + if with_timestamp { + sort_fields.push(self.timestamp_index); + } + self.sort_fields_by_indices(&sort_fields) + } + + fn sort_fields_by_indices(&self, indices: &[usize]) -> Vec { + indices + .iter() + .map(|index| SortField::new(self.schema.field(*index).data_type().clone())) + .collect() + } + + pub fn converter(&self, with_timestamp: bool) -> Result { + Converter::new(self.sort_fields(with_timestamp)) + } + + pub fn value_converter( + &self, + with_timestamp: bool, + generation_index: usize, + ) -> Result { + match &self.key_indices { + None => { + let mut indices = (0..self.schema.fields().len()).collect::>(); + indices.remove(generation_index); + if !with_timestamp { + indices.remove(self.timestamp_index); + } + Converter::new(self.sort_fields_by_indices(&indices)) + } + Some(keys) => { + let indices = (0..self.schema.fields().len()) + .filter(|index| { + !keys.contains(index) + && (with_timestamp || *index != self.timestamp_index) + && *index != generation_index + }) + .collect::>(); + Converter::new(self.sort_fields_by_indices(&indices)) + } + } + } + + pub fn value_indices(&self, with_timestamp: bool) -> Vec { + let field_count = self.schema.fields().len(); + match &self.key_indices { + None => { + let mut indices = (0..field_count).collect::>(); + + if !with_timestamp { + indices.remove(self.timestamp_index); + } + indices + } + Some(keys) => (0..field_count) + .filter(|index| { + !keys.contains(index) && (with_timestamp || *index != self.timestamp_index) + }) + .collect::>(), + } + } + + pub fn sort( + &self, + batch: RecordBatch, + with_timestamp: bool, + ) -> Result { + if self.key_indices.is_none() && !with_timestamp { + return Ok(batch); + } + let sort_columns = self.sort_columns(&batch, with_timestamp); + let sort_indices = lexsort_to_indices(&sort_columns, None).expect("should be able to sort"); + let columns = batch + .columns() + .iter() + .map(|c| take(c, &sort_indices, None).unwrap()) + .collect(); + + RecordBatch::try_new(batch.schema(), columns) + } + + pub fn partition( + &self, + batch: &RecordBatch, + with_timestamp: bool, + ) -> Result>, ArrowError> { + if self.key_indices.is_none() && !with_timestamp { + #[allow(clippy::single_range_in_vec_init)] + return Ok(vec![0..batch.num_rows()]); + } + + let mut partition_columns = vec![]; + + if let Some(keys) = &self.routing_keys() { + partition_columns.extend(keys.iter().map(|index| batch.column(*index).clone())); + } + if with_timestamp { + partition_columns.push(batch.column(self.timestamp_index).clone()); + } + + Ok(partition(&partition_columns)?.ranges()) + } + + pub fn unkeyed_batch(&self, batch: &RecordBatch) -> Result { + if self.key_indices.is_none() { + return Ok(batch.clone()); + } + let columns: Vec<_> = (0..batch.num_columns()) + .filter(|index| !self.key_indices.as_ref().unwrap().contains(index)) + .collect(); + batch.project(&columns) + } + + pub fn schema_without_keys(&self) -> Result { + if self.key_indices.is_none() { + return Ok(self.clone()); + } + let key_indices = self.key_indices.as_ref().unwrap(); + let unkeyed_schema = Schema::new( + self.schema + .fields() + .iter() + .enumerate() + .filter(|(index, _field)| !key_indices.contains(index)) + .map(|(_, field)| field.as_ref().clone()) + .collect::>(), + ); + let timestamp_index = unkeyed_schema.index_of(TIMESTAMP_FIELD)?; + Ok(Self { + schema: Arc::new(unkeyed_schema), + timestamp_index, + key_indices: None, + routing_key_indices: None, + }) + } + + pub fn with_fields(&self, fields: Vec) -> Result { + let schema = Arc::new(Schema::new_with_metadata( + fields, + self.schema.metadata.clone(), + )); + + let timestamp_index = schema.index_of(TIMESTAMP_FIELD)?; + let max_index = *[&self.key_indices, &self.routing_key_indices] + .iter() + .map(|indices| indices.as_ref().and_then(|k| k.iter().max())) + .max() + .flatten() + .unwrap_or(&0); + + if schema.fields.len() - 1 < max_index { + return Err(ArrowError::InvalidArgumentError(format!( + "expected at least {} fields, but were only {}", + max_index + 1, + schema.fields.len() + ))); + } + + Ok(Self { + schema, + timestamp_index, + key_indices: self.key_indices.clone(), + routing_key_indices: self.routing_key_indices.clone(), + }) + } + + pub fn with_additional_fields( + &self, + new_fields: impl Iterator, + ) -> Result { + let mut fields = self.schema.fields.to_vec(); + fields.extend(new_fields.map(Arc::new)); + + self.with_fields(fields) + } +} + +pub fn server_for_hash_array( + hash: &PrimitiveArray, + n: usize, +) -> Result, ArrowError> { + let range_size = u64::MAX / (n as u64) + 1; + let range_scalar = UInt64Array::new_scalar(range_size); + let division = div(hash, &range_scalar)?; + let result: &PrimitiveArray = division.as_any().downcast_ref().unwrap(); + Ok(result.clone()) +} diff --git a/src/streaming_planner/src/common/kafka_catalog.rs b/src/streaming_planner/src/common/kafka_catalog.rs new file mode 100644 index 00000000..51ceee67 --- /dev/null +++ b/src/streaming_planner/src/common/kafka_catalog.rs @@ -0,0 +1,116 @@ +// 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 serde::{Deserialize, Serialize}; +use std::collections::HashMap; + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct KafkaTable { + pub topic: String, + #[serde(flatten)] + pub kind: TableType, + #[serde(default)] + pub client_configs: HashMap, + pub value_subject: Option, +} + +impl KafkaTable { + pub fn subject(&self) -> String { + self.value_subject + .clone() + .unwrap_or_else(|| format!("{}-value", self.topic)) + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum TableType { + Source { + offset: KafkaTableSourceOffset, + read_mode: Option, + group_id: Option, + group_id_prefix: Option, + }, + Sink { + commit_mode: SinkCommitMode, + key_field: Option, + timestamp_field: Option, + }, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)] +#[serde(rename_all = "snake_case")] +pub enum KafkaTableSourceOffset { + Latest, + Earliest, + #[default] + Group, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum ReadMode { + ReadUncommitted, + ReadCommitted, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)] +#[serde(rename_all = "snake_case")] +pub enum SinkCommitMode { + #[default] + AtLeastOnce, + ExactlyOnce, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(rename_all = "camelCase")] +pub struct KafkaConfig { + pub bootstrap_servers: String, + #[serde(default)] + pub authentication: KafkaConfigAuthentication, + #[serde(default)] + pub schema_registry_enum: Option, + #[serde(default)] + pub connection_properties: HashMap, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)] +#[serde(tag = "type")] +pub enum KafkaConfigAuthentication { + #[default] + #[serde(rename = "None")] + None, + #[serde(rename = "AWS_MSK_IAM")] + AwsMskIam { region: String }, + #[serde(rename = "SASL")] + Sasl { + protocol: String, + mechanism: String, + username: String, + password: String, + }, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(tag = "type")] +pub enum SchemaRegistryConfig { + #[serde(rename = "None")] + None, + #[serde(rename = "Confluent Schema Registry")] + ConfluentSchemaRegistry { + endpoint: String, + #[serde(rename = "apiKey")] + api_key: Option, + #[serde(rename = "apiSecret")] + api_secret: Option, + }, +} diff --git a/src/streaming_planner/src/common/mod.rs b/src/streaming_planner/src/common/mod.rs new file mode 100644 index 00000000..6133294b --- /dev/null +++ b/src/streaming_planner/src/common/mod.rs @@ -0,0 +1,65 @@ +// 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. + +//! Shared core types and constants for FunctionStream streaming planner (`crate::common`). +//! +//! Used by the runtime, SQL planner, coordinator, and other subsystems — +//! analogous to `arroyo-types` + `arroyo-rpc` in Arroyo. + +pub mod arrow_ext; +pub mod connector_options; +pub mod constants; +pub mod control; +pub mod converter; +pub mod date; +pub mod debezium; +pub mod errors; +pub mod format_from_opts; +pub mod formats; +pub mod fs_schema; +pub mod kafka_catalog; +pub mod operator_config; +pub mod time_utils; +pub mod topology; +pub mod with_option_keys; + +// ── Re-exports from existing modules ── +pub use function_stream_runtime_common::streaming_protocol::{CheckpointBarrier, Watermark}; +pub use arrow_ext::FsExtensionType; +pub use time_utils::{from_nanos, to_micros, to_millis, to_nanos}; + +// ── Re-exports from new modules ── +pub use connector_options::ConnectorOptions; +pub use formats::{BadData, Format, Framing, JsonCompression, JsonFormat}; +pub use fs_schema::{FsSchema, FsSchemaRef}; +pub use operator_config::MetadataField; + +// ── Well-known column names ── +pub use constants::sql_field::{TIMESTAMP_FIELD, UPDATING_META_FIELD}; +pub use topology::render_program_topology; + +// ── Environment variables ── +pub const JOB_ID_ENV: &str = "JOB_ID"; +pub const RUN_ID_ENV: &str = "RUN_ID"; + +// ── Metric names ── +pub const MESSAGES_RECV: &str = "fs_worker_messages_recv"; +pub const MESSAGES_SENT: &str = "fs_worker_messages_sent"; +pub const BYTES_RECV: &str = "fs_worker_bytes_recv"; +pub const BYTES_SENT: &str = "fs_worker_bytes_sent"; +pub const BATCHES_RECV: &str = "fs_worker_batches_recv"; +pub const BATCHES_SENT: &str = "fs_worker_batches_sent"; +pub const TX_QUEUE_SIZE: &str = "fs_worker_tx_queue_size"; +pub const TX_QUEUE_REM: &str = "fs_worker_tx_queue_rem"; +pub const DESERIALIZATION_ERRORS: &str = "fs_worker_deserialization_errors"; + +pub const LOOKUP_KEY_INDEX_FIELD: &str = "__lookup_key_index"; diff --git a/src/streaming_planner/src/common/operator_config.rs b/src/streaming_planner/src/common/operator_config.rs new file mode 100644 index 00000000..209bee48 --- /dev/null +++ b/src/streaming_planner/src/common/operator_config.rs @@ -0,0 +1,21 @@ +// 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 serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MetadataField { + pub field_name: String, + pub key: String, + #[serde(default)] + pub data_type: Option, +} diff --git a/src/streaming_planner/src/common/time_utils.rs b/src/streaming_planner/src/common/time_utils.rs new file mode 100644 index 00000000..323445cd --- /dev/null +++ b/src/streaming_planner/src/common/time_utils.rs @@ -0,0 +1,74 @@ +// 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::hash::Hash; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +pub fn to_millis(time: SystemTime) -> u64 { + time.duration_since(UNIX_EPOCH).unwrap().as_millis() as u64 +} + +pub fn to_micros(time: SystemTime) -> u64 { + time.duration_since(UNIX_EPOCH).unwrap().as_micros() as u64 +} + +pub fn from_millis(ts: u64) -> SystemTime { + UNIX_EPOCH + Duration::from_millis(ts) +} + +pub fn from_micros(ts: u64) -> SystemTime { + UNIX_EPOCH + Duration::from_micros(ts) +} + +pub fn to_nanos(time: SystemTime) -> u128 { + time.duration_since(UNIX_EPOCH).unwrap().as_nanos() +} + +pub fn from_nanos(ts: u128) -> SystemTime { + UNIX_EPOCH + + Duration::from_secs((ts / 1_000_000_000) as u64) + + Duration::from_nanos((ts % 1_000_000_000) as u64) +} + +pub fn print_time(time: SystemTime) -> String { + chrono::DateTime::::from(time) + .format("%Y-%m-%d %H:%M:%S%.3f") + .to_string() +} + +/// Returns the number of days since the UNIX epoch (for Avro serialization). +pub fn days_since_epoch(time: SystemTime) -> i32 { + time.duration_since(UNIX_EPOCH) + .unwrap() + .as_secs() + .div_euclid(86400) as i32 +} + +pub fn single_item_hash_map, K: Hash + Eq, V>(key: I, value: V) -> HashMap { + let mut map = HashMap::new(); + map.insert(key.into(), value); + map +} + +pub fn string_to_map(s: &str, pair_delimiter: char) -> Option> { + if s.trim().is_empty() { + return Some(HashMap::new()); + } + + s.split(',') + .map(|s| { + let mut kv = s.trim().split(pair_delimiter); + Some((kv.next()?.trim().to_string(), kv.next()?.trim().to_string())) + }) + .collect() +} diff --git a/src/streaming_planner/src/common/topology.rs b/src/streaming_planner/src/common/topology.rs new file mode 100644 index 00000000..3b4f892f --- /dev/null +++ b/src/streaming_planner/src/common/topology.rs @@ -0,0 +1,295 @@ +// 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, VecDeque}; +use std::fmt::Write; + +use protocol::function_stream_graph::FsProgram; + +fn edge_type_label(edge_type: i32) -> &'static str { + match edge_type { + 1 => "Forward", + 2 => "Shuffle", + 3 => "LeftJoin", + 4 => "RightJoin", + _ => "Unknown", + } +} + +pub fn render_program_topology(program: &FsProgram) -> String { + if program.nodes.is_empty() { + return "(empty topology)".to_string(); + } + + struct EdgeInfo { + target: i32, + edge_type: i32, + } + struct InputInfo { + source: i32, + edge_type: i32, + } + + let node_map: BTreeMap = + program.nodes.iter().map(|n| (n.node_index, n)).collect(); + + let mut downstream: BTreeMap> = BTreeMap::new(); + let mut upstream: BTreeMap> = BTreeMap::new(); + let mut in_degree: BTreeMap = BTreeMap::new(); + + for idx in node_map.keys() { + in_degree.entry(*idx).or_insert(0); + } + for edge in &program.edges { + downstream.entry(edge.source).or_default().push(EdgeInfo { + target: edge.target, + edge_type: edge.edge_type, + }); + upstream.entry(edge.target).or_default().push(InputInfo { + source: edge.source, + edge_type: edge.edge_type, + }); + *in_degree.entry(edge.target).or_insert(0) += 1; + } + + // Kahn's topological sort + let mut queue: VecDeque = in_degree + .iter() + .filter(|(_, deg)| **deg == 0) + .map(|(idx, _)| *idx) + .collect(); + let mut topo_order: Vec = Vec::with_capacity(node_map.len()); + let mut remaining = in_degree.clone(); + while let Some(idx) = queue.pop_front() { + topo_order.push(idx); + if let Some(edges) = downstream.get(&idx) { + for e in edges { + if let Some(deg) = remaining.get_mut(&e.target) { + *deg -= 1; + if *deg == 0 { + queue.push_back(e.target); + } + } + } + } + } + for idx in node_map.keys() { + if !topo_order.contains(idx) { + topo_order.push(*idx); + } + } + + let is_source = |idx: &i32| upstream.get(idx).is_none_or(|v| v.is_empty()); + let is_sink = |idx: &i32| downstream.get(idx).is_none_or(|v| v.is_empty()); + + let mut out = String::new(); + let _ = writeln!( + out, + "Pipeline Topology ({} nodes, {} edges)", + program.nodes.len(), + program.edges.len(), + ); + let _ = writeln!(out, "{}", "=".repeat(50)); + + for (pos, &node_idx) in topo_order.iter().enumerate() { + let Some(node) = node_map.get(&node_idx) else { + continue; + }; + + let op_chain: String = node + .operators + .iter() + .map(|op| op.operator_name.as_str()) + .collect::>() + .join(" -> "); + + let role = if is_source(&node_idx) { + "Source" + } else if is_sink(&node_idx) { + "Sink" + } else { + "Operator" + }; + + let _ = writeln!(out); + let _ = writeln!( + out, + "[{role}] Node {node_idx} parallelism = {}", + node.parallelism, + ); + let _ = writeln!(out, " operators: {op_chain}"); + + if !node.description.is_empty() { + let _ = writeln!(out, " desc: {}", node.description); + } + + if let Some(inputs) = upstream.get(&node_idx) { + if inputs.len() == 1 { + let i = &inputs[0]; + let _ = writeln!( + out, + " input: <-- [{}] Node {}", + edge_type_label(i.edge_type), + i.source, + ); + } else if inputs.len() > 1 { + let _ = writeln!(out, " inputs:"); + for i in inputs { + let _ = writeln!( + out, + " <-- [{}] Node {}", + edge_type_label(i.edge_type), + i.source, + ); + } + } + } + + if let Some(outputs) = downstream.get(&node_idx) { + if outputs.len() == 1 { + let e = &outputs[0]; + let _ = writeln!( + out, + " output: --> [{}] Node {}", + edge_type_label(e.edge_type), + e.target, + ); + } else if outputs.len() > 1 { + let _ = writeln!(out, " outputs:"); + for e in outputs { + let _ = writeln!( + out, + " --> [{}] Node {}", + edge_type_label(e.edge_type), + e.target, + ); + } + } + } + + if pos < topo_order.len() - 1 { + let single_out = downstream.get(&node_idx).is_some_and(|v| v.len() == 1); + let next_idx = topo_order.get(pos + 1).copied(); + let is_direct = single_out + && next_idx + .is_some_and(|n| downstream.get(&node_idx).is_some_and(|v| v[0].target == n)); + let next_single_in = next_idx + .and_then(|n| upstream.get(&n)) + .is_some_and(|v| v.len() == 1); + + if is_direct && next_single_in { + let etype = downstream.get(&node_idx).unwrap()[0].edge_type; + let _ = writeln!(out, " |"); + let _ = writeln!(out, " | {}", edge_type_label(etype)); + let _ = writeln!(out, " v"); + } + } + } + + out.trim_end().to_string() +} + +#[cfg(test)] +mod tests { + use super::*; + use protocol::function_stream_graph::{ChainedOperator, FsEdge, FsNode, FsProgram}; + + fn make_node( + node_index: i32, + operators: Vec<(&str, &str)>, + desc: &str, + parallelism: u32, + ) -> FsNode { + FsNode { + node_index, + node_id: node_index as u32, + parallelism, + description: desc.to_string(), + operators: operators + .into_iter() + .map(|(id, name)| ChainedOperator { + operator_id: id.to_string(), + operator_name: name.to_string(), + operator_config: Vec::new(), + }) + .collect(), + edges: Vec::new(), + } + } + + fn make_edge(source: i32, target: i32, edge_type: i32) -> FsEdge { + FsEdge { + source, + target, + schema: None, + edge_type, + } + } + + #[test] + fn empty_program_renders_placeholder() { + let program = FsProgram { + nodes: vec![], + edges: vec![], + program_config: None, + }; + assert_eq!(render_program_topology(&program), "(empty topology)"); + } + + #[test] + fn linear_pipeline_renders_correctly() { + let program = FsProgram { + nodes: vec![ + make_node(0, vec![("src_0", "ConnectorSource")], "", 1), + make_node( + 1, + vec![("val_1", "Value"), ("wm_2", "ExpressionWatermark")], + "source -> watermark", + 1, + ), + make_node(2, vec![("sink_3", "ConnectorSink")], "sink (kafka)", 1), + ], + edges: vec![make_edge(0, 1, 1), make_edge(1, 2, 1)], + program_config: None, + }; + let result = render_program_topology(&program); + assert!(result.contains("[Source] Node 0")); + assert!(result.contains("[Operator] Node 1")); + assert!(result.contains("[Sink] Node 2")); + assert!(result.contains("ConnectorSource")); + assert!(result.contains("Value -> ExpressionWatermark")); + assert!(result.contains("Forward")); + } + + #[test] + fn join_topology_shows_multiple_inputs() { + let program = FsProgram { + nodes: vec![ + make_node(0, vec![("src_a", "ConnectorSource")], "source A", 1), + make_node(1, vec![("src_b", "ConnectorSource")], "source B", 1), + make_node(2, vec![("join_0", "WindowJoin")], "join node", 2), + make_node(3, vec![("sink_0", "ConnectorSink")], "sink", 1), + ], + edges: vec![ + make_edge(0, 2, 3), // LeftJoin + make_edge(1, 2, 4), // RightJoin + make_edge(2, 3, 1), // Forward + ], + program_config: None, + }; + let result = render_program_topology(&program); + assert!(result.contains("inputs:")); + assert!(result.contains("LeftJoin")); + assert!(result.contains("RightJoin")); + assert!(result.contains("[Operator] Node 2")); + } +} diff --git a/src/streaming_planner/src/common/with_option_keys.rs b/src/streaming_planner/src/common/with_option_keys.rs new file mode 100644 index 00000000..21bfa691 --- /dev/null +++ b/src/streaming_planner/src/common/with_option_keys.rs @@ -0,0 +1,105 @@ +// 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. + +pub const CONNECTOR: &str = "connector"; +pub const TYPE: &str = "type"; +pub const FORMAT: &str = "format"; +pub const DEFAULT_FORMAT_VALUE: &str = "json"; +pub const BAD_DATA: &str = "bad_data"; +pub const PARTITION_BY: &str = "partition_by"; +pub const PATH: &str = "path"; +pub const SINK_PATH: &str = "sink.path"; + +pub const EVENT_TIME_FIELD: &str = "event_time_field"; +pub const WATERMARK_FIELD: &str = "watermark_field"; + +pub const IDLE_MICROS: &str = "idle_micros"; +pub const IDLE_TIME: &str = "idle_time"; + +pub const LOOKUP_CACHE_MAX_BYTES: &str = "lookup.cache.max_bytes"; +pub const LOOKUP_CACHE_TTL: &str = "lookup.cache.ttl"; + +pub const CONNECTION_SCHEMA: &str = "connection_schema"; + +pub const ADAPTER: &str = "adapter"; + +// ── Kafka ───────────────────────────────────────────────────────────────── + +pub const KAFKA_BOOTSTRAP_SERVERS: &str = "bootstrap.servers"; +pub const KAFKA_BOOTSTRAP_SERVERS_LEGACY: &str = "bootstrap_servers"; +pub const KAFKA_TOPIC: &str = "topic"; +pub const KAFKA_RATE_LIMIT_MESSAGES_PER_SECOND: &str = "rate_limit.messages_per_second"; +pub const KAFKA_VALUE_SUBJECT: &str = "value.subject"; +pub const KAFKA_SCAN_STARTUP_MODE: &str = "scan.startup.mode"; +pub const KAFKA_ISOLATION_LEVEL: &str = "isolation.level"; +pub const KAFKA_GROUP_ID: &str = "group.id"; +pub const KAFKA_GROUP_ID_LEGACY: &str = "group_id"; +pub const KAFKA_GROUP_ID_PREFIX: &str = "group.id.prefix"; +pub const KAFKA_SINK_COMMIT_MODE: &str = "sink.commit.mode"; +pub const KAFKA_SINK_KEY_FIELD: &str = "sink.key.field"; +pub const KAFKA_KEY_FIELD_LEGACY: &str = "key.field"; +pub const KAFKA_SINK_TIMESTAMP_FIELD: &str = "sink.timestamp.field"; +pub const KAFKA_TIMESTAMP_FIELD_LEGACY: &str = "timestamp.field"; + +// ── JSON format ─────────────────────────────────────────────────────────── + +pub const JSON_CONFLUENT_SCHEMA_REGISTRY: &str = "json.confluent_schema_registry"; +pub const JSON_CONFLUENT_SCHEMA_VERSION: &str = "json.confluent_schema_version"; +pub const JSON_INCLUDE_SCHEMA: &str = "json.include_schema"; +pub const JSON_DEBEZIUM: &str = "json.debezium"; +pub const JSON_UNSTRUCTURED: &str = "json.unstructured"; +pub const JSON_TIMESTAMP_FORMAT: &str = "json.timestamp_format"; +pub const JSON_DECIMAL_ENCODING: &str = "json.decimal_encoding"; +pub const JSON_COMPRESSION: &str = "json.compression"; + +// ── Avro ────────────────────────────────────────────────────────────────── + +pub const AVRO_CONFLUENT_SCHEMA_REGISTRY: &str = "avro.confluent_schema_registry"; +pub const AVRO_RAW_DATUMS: &str = "avro.raw_datums"; +pub const AVRO_INTO_UNSTRUCTURED_JSON: &str = "avro.into_unstructured_json"; +pub const AVRO_SCHEMA_ID: &str = "avro.schema_id"; + +// ── Parquet ─────────────────────────────────────────────────────────────── + +pub const PARQUET_COMPRESSION: &str = "parquet.compression"; +pub const PARQUET_ROW_GROUP_BYTES: &str = "parquet.row_group_bytes"; + +// ── S3 ──────────────────────────────────────────────────────────────────── + +pub const S3_BUCKET: &str = "s3.bucket"; +pub const S3_REGION: &str = "s3.region"; +pub const S3_ENDPOINT: &str = "s3.endpoint"; +pub const S3_ACCESS_KEY_ID: &str = "s3.access_key_id"; +pub const S3_SECRET_ACCESS_KEY: &str = "s3.secret_access_key"; +pub const S3_SESSION_TOKEN: &str = "s3.session_token"; + +// ── Protobuf ──────────────────────────────────────────────────────────────── + +pub const PROTOBUF_INTO_UNSTRUCTURED_JSON: &str = "protobuf.into_unstructured_json"; +pub const PROTOBUF_MESSAGE_NAME: &str = "protobuf.message_name"; +pub const PROTOBUF_CONFLUENT_SCHEMA_REGISTRY: &str = "protobuf.confluent_schema_registry"; +pub const PROTOBUF_LENGTH_DELIMITED: &str = "protobuf.length_delimited"; + +// ── Framing ───────────────────────────────────────────────────────────────── + +pub const FRAMING_METHOD: &str = "framing.method"; +pub const FRAMING_MAX_LINE_LENGTH: &str = "framing.max_line_length"; + +pub const FORMAT_DEBEZIUM_FLAG: &str = "format.debezium"; + +// ── Streaming runtime common options ─────────────────────────────────────── + +pub const CHECKPOINT_INTERVAL_MS: &str = "checkpoint.interval.ms"; +pub const PIPELINE_PARALLELISM: &str = "pipeline.parallelism"; +pub const KEY_BY_PARALLELISM: &str = "key_by.parallelism"; +pub const OPERATOR_MEMORY_BYTES: &str = "operator.memory.bytes"; +pub const SINK_MEMORY_BYTES: &str = "sink.memory.bytes"; diff --git a/src/streaming_planner/src/connector/config.rs b/src/streaming_planner/src/connector/config.rs new file mode 100644 index 00000000..c3eaf00f --- /dev/null +++ b/src/streaming_planner/src/connector/config.rs @@ -0,0 +1,91 @@ +// 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 protocol::function_stream_graph::{ + DeltaSinkConfig, FilesystemSinkConfig, IcebergSinkConfig, KafkaSinkConfig, KafkaSourceConfig, + LanceDbSinkConfig, S3SinkConfig, connector_op, +}; + +#[derive(Debug, Clone)] +pub enum ConnectorConfig { + KafkaSource(KafkaSourceConfig), + KafkaSink(KafkaSinkConfig), + FilesystemSink(FilesystemSinkConfig), + DeltaSink(DeltaSinkConfig), + IcebergSink(IcebergSinkConfig), + S3Sink(S3SinkConfig), + LanceDbSink(LanceDbSinkConfig), +} + +impl ConnectorConfig { + pub fn to_proto_config(&self) -> connector_op::Config { + match self { + ConnectorConfig::KafkaSource(cfg) => connector_op::Config::KafkaSource(cfg.clone()), + ConnectorConfig::KafkaSink(cfg) => connector_op::Config::KafkaSink(cfg.clone()), + ConnectorConfig::FilesystemSink(cfg) => { + connector_op::Config::FilesystemSink(cfg.clone()) + } + ConnectorConfig::DeltaSink(cfg) => connector_op::Config::DeltaSink(cfg.clone()), + ConnectorConfig::IcebergSink(cfg) => connector_op::Config::IcebergSink(cfg.clone()), + ConnectorConfig::S3Sink(cfg) => connector_op::Config::S3Sink(cfg.clone()), + ConnectorConfig::LanceDbSink(cfg) => connector_op::Config::LancedbSink(cfg.clone()), + } + } +} + +impl PartialEq for ConnectorConfig { + fn eq(&self, other: &Self) -> bool { + use prost::Message; + match (self, other) { + (ConnectorConfig::KafkaSource(a), ConnectorConfig::KafkaSource(b)) => { + a.encode_to_vec() == b.encode_to_vec() + } + (ConnectorConfig::KafkaSink(a), ConnectorConfig::KafkaSink(b)) => { + a.encode_to_vec() == b.encode_to_vec() + } + (ConnectorConfig::FilesystemSink(a), ConnectorConfig::FilesystemSink(b)) => { + a.encode_to_vec() == b.encode_to_vec() + } + (ConnectorConfig::DeltaSink(a), ConnectorConfig::DeltaSink(b)) => { + a.encode_to_vec() == b.encode_to_vec() + } + (ConnectorConfig::IcebergSink(a), ConnectorConfig::IcebergSink(b)) => { + a.encode_to_vec() == b.encode_to_vec() + } + (ConnectorConfig::S3Sink(a), ConnectorConfig::S3Sink(b)) => { + a.encode_to_vec() == b.encode_to_vec() + } + (ConnectorConfig::LanceDbSink(a), ConnectorConfig::LanceDbSink(b)) => { + a.encode_to_vec() == b.encode_to_vec() + } + _ => false, + } + } +} + +impl Eq for ConnectorConfig {} + +impl std::hash::Hash for ConnectorConfig { + fn hash(&self, state: &mut H) { + use prost::Message; + std::mem::discriminant(self).hash(state); + match self { + ConnectorConfig::KafkaSource(cfg) => cfg.encode_to_vec().hash(state), + ConnectorConfig::KafkaSink(cfg) => cfg.encode_to_vec().hash(state), + ConnectorConfig::FilesystemSink(cfg) => cfg.encode_to_vec().hash(state), + ConnectorConfig::DeltaSink(cfg) => cfg.encode_to_vec().hash(state), + ConnectorConfig::IcebergSink(cfg) => cfg.encode_to_vec().hash(state), + ConnectorConfig::S3Sink(cfg) => cfg.encode_to_vec().hash(state), + ConnectorConfig::LanceDbSink(cfg) => cfg.encode_to_vec().hash(state), + } + } +} diff --git a/src/streaming_planner/src/connector/factory.rs b/src/streaming_planner/src/connector/factory.rs new file mode 100644 index 00000000..e89bafa1 --- /dev/null +++ b/src/streaming_planner/src/connector/factory.rs @@ -0,0 +1,67 @@ +// 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 datafusion::arrow::datatypes::Schema; +use datafusion::common::Result; + +use super::config::ConnectorConfig; +use super::registry::REGISTRY; +use super::sink::runtime_config::SinkRuntimeConfig; +use crate::common::connector_options::ConnectorOptions; +use crate::common::formats::{BadData, Format}; +use crate::schema::table_role::TableRole; + +pub fn build_connector_config( + connector_name: &str, + role: TableRole, + options: &mut ConnectorOptions, + format: &Option, + bad_data: BadData, +) -> Result { + let runtime_opts_map = options.snapshot_for_catalog().into_iter().collect(); + let runtime_props = + SinkRuntimeConfig::from_options_map(&runtime_opts_map)?.to_runtime_properties(); + match role { + TableRole::Ingestion | TableRole::Reference => REGISTRY + .get_source(connector_name)? + .build_source_config(options, format, bad_data), + TableRole::Egress => { + REGISTRY + .get_sink(connector_name)? + .build_sink_config(options, format, &runtime_props) + } + } +} + +pub fn build_connector_config_from_options( + connector_name: &str, + role: TableRole, + options: &mut ConnectorOptions, + format: &Option, + bad_data: BadData, +) -> Result { + build_connector_config(connector_name, role, options, format, bad_data) +} + +pub fn build_connector_config_from_catalog( + connector_name: &str, + role: TableRole, + opts: HashMap, + _physical_schema: &Schema, +) -> Result { + let mut options = ConnectorOptions::from_flat_string_map(opts)?; + let format = Format::from_opts(&mut options)?; + let bad_data = BadData::from_opts(&mut options)?; + build_connector_config(connector_name, role, &mut options, &format, bad_data) +} diff --git a/src/streaming_planner/src/connector/mod.rs b/src/streaming_planner/src/connector/mod.rs new file mode 100644 index 00000000..f477c976 --- /dev/null +++ b/src/streaming_planner/src/connector/mod.rs @@ -0,0 +1,18 @@ +// 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. + +pub mod config; +pub mod factory; +pub mod provider; +pub mod registry; +pub mod sink; +pub mod source; diff --git a/src/streaming_planner/src/connector/provider.rs b/src/streaming_planner/src/connector/provider.rs new file mode 100644 index 00000000..83e46aa7 --- /dev/null +++ b/src/streaming_planner/src/connector/provider.rs @@ -0,0 +1,52 @@ +// 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 datafusion::common::{DataFusionError, Result}; + +use super::config::ConnectorConfig; +use super::sink::runtime_config::SinkRuntimeProperties; +use crate::common::connector_options::ConnectorOptions; +use crate::common::formats::{BadData, Format}; + +pub trait SourceProvider: Send + Sync { + fn name(&self) -> &'static str; + + fn build_source_config( + &self, + options: &mut ConnectorOptions, + format: &Option, + bad_data: BadData, + ) -> Result; +} + +pub trait SinkProvider: Send + Sync { + fn name(&self) -> &'static str; + fn build_sink_config( + &self, + options: &mut ConnectorOptions, + format: &Option, + runtime_props: &SinkRuntimeProperties, + ) -> Result; +} + +pub fn require_option( + options: &mut ConnectorOptions, + key: &str, + connector_name: &str, +) -> Result { + options.pull_opt_str(key)?.ok_or_else(|| { + DataFusionError::Plan(format!( + "Connector '{}' requires option '{}' to be set", + connector_name, key + )) + }) +} diff --git a/src/streaming_planner/src/connector/registry.rs b/src/streaming_planner/src/connector/registry.rs new file mode 100644 index 00000000..4a8a8c1c --- /dev/null +++ b/src/streaming_planner/src/connector/registry.rs @@ -0,0 +1,86 @@ +// 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::{Arc, LazyLock}; + +use datafusion::common::{DataFusionError, Result}; + +use super::provider::{SinkProvider, SourceProvider}; +use super::sink::delta::DeltaSinkConnector; +use super::sink::filesystem::FilesystemSinkConnector; +use super::sink::iceberg::IcebergSinkConnector; +use super::sink::kafka::KafkaSinkConnector; +use super::sink::lancedb::LanceDbSinkConnector; +use super::sink::s3::S3SinkConnector; +use super::source::kafka::KafkaSourceConnector; + +pub struct ConnectorRegistry { + sources: HashMap>, + sinks: HashMap>, +} + +impl ConnectorRegistry { + fn new() -> Self { + let mut registry = Self { + sources: HashMap::new(), + sinks: HashMap::new(), + }; + + registry.register_source(Arc::new(KafkaSourceConnector)); + + registry.register_sink(Arc::new(KafkaSinkConnector)); + registry.register_sink(Arc::new(S3SinkConnector)); + registry.register_sink(Arc::new(FilesystemSinkConnector)); + registry.register_sink(Arc::new(DeltaSinkConnector)); + registry.register_sink(Arc::new(IcebergSinkConnector)); + registry.register_sink(Arc::new(LanceDbSinkConnector)); + + registry + } + + pub fn register_source(&mut self, provider: Arc) { + self.sources + .insert(provider.name().to_ascii_lowercase(), provider); + } + + pub fn register_sink(&mut self, provider: Arc) { + self.sinks + .insert(provider.name().to_ascii_lowercase(), provider); + } + + pub fn get_source(&self, connector_name: &str) -> Result> { + self.sources + .get(&connector_name.to_ascii_lowercase()) + .cloned() + .ok_or_else(|| { + DataFusionError::Plan(format!( + "Connector '{}' is not registered or does not support being used as a SOURCE", + connector_name + )) + }) + } + + pub fn get_sink(&self, connector_name: &str) -> Result> { + self.sinks + .get(&connector_name.to_ascii_lowercase()) + .cloned() + .ok_or_else(|| { + DataFusionError::Plan(format!( + "Connector '{}' is not registered or does not support being used as a SINK", + connector_name + )) + }) + } +} + +pub static REGISTRY: LazyLock = LazyLock::new(ConnectorRegistry::new); diff --git a/src/streaming_planner/src/connector/sink/delta.rs b/src/streaming_planner/src/connector/sink/delta.rs new file mode 100644 index 00000000..08d86d96 --- /dev/null +++ b/src/streaming_planner/src/connector/sink/delta.rs @@ -0,0 +1,60 @@ +// 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 datafusion::common::Result; +use protocol::function_stream_graph::{DeltaSinkConfig, SinkFormatProto}; + +use crate::common::Format; +use crate::common::connector_options::ConnectorOptions; +use crate::common::constants::connector_type; +use crate::connector::config::ConnectorConfig; +use crate::connector::provider::SinkProvider; +use crate::connector::sink::runtime_config::SinkRuntimeProperties; +use crate::connector::sink::utils::SinkUtils; + +pub struct DeltaSinkConnector; + +impl SinkProvider for DeltaSinkConnector { + fn name(&self) -> &'static str { + connector_type::DELTA + } + + fn build_sink_config( + &self, + options: &mut ConnectorOptions, + format: &Option, + runtime_props: &SinkRuntimeProperties, + ) -> Result { + let path = SinkUtils::require_path(options)?; + let parquet_compression = SinkUtils::extract_parquet_compression(options)?; + let format_proto = SinkUtils::resolve_sink_format( + format, + self.name(), + &[ + SinkFormatProto::SinkFormatCsv, + SinkFormatProto::SinkFormatJsonl, + SinkFormatProto::SinkFormatAvro, + SinkFormatProto::SinkFormatParquet, + SinkFormatProto::SinkFormatOrc, + ], + )?; + let extra_properties = options.drain_remaining_string_values()?; + + Ok(ConnectorConfig::DeltaSink(DeltaSinkConfig { + path, + format: format_proto, + parquet_compression, + extra_properties, + runtime_properties: runtime_props.clone(), + })) + } +} diff --git a/src/streaming_planner/src/connector/sink/filesystem.rs b/src/streaming_planner/src/connector/sink/filesystem.rs new file mode 100644 index 00000000..e529ed79 --- /dev/null +++ b/src/streaming_planner/src/connector/sink/filesystem.rs @@ -0,0 +1,60 @@ +// 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 datafusion::common::Result; +use protocol::function_stream_graph::{FilesystemSinkConfig, SinkFormatProto}; + +use crate::common::Format; +use crate::common::connector_options::ConnectorOptions; +use crate::common::constants::connector_type; +use crate::connector::config::ConnectorConfig; +use crate::connector::provider::SinkProvider; +use crate::connector::sink::runtime_config::SinkRuntimeProperties; +use crate::connector::sink::utils::SinkUtils; + +pub struct FilesystemSinkConnector; + +impl SinkProvider for FilesystemSinkConnector { + fn name(&self) -> &'static str { + connector_type::FILESYSTEM + } + + fn build_sink_config( + &self, + options: &mut ConnectorOptions, + format: &Option, + runtime_props: &SinkRuntimeProperties, + ) -> Result { + let path = SinkUtils::require_path(options)?; + let parquet_compression = SinkUtils::extract_parquet_compression(options)?; + let format_proto = SinkUtils::resolve_sink_format( + format, + self.name(), + &[ + SinkFormatProto::SinkFormatCsv, + SinkFormatProto::SinkFormatJsonl, + SinkFormatProto::SinkFormatAvro, + SinkFormatProto::SinkFormatParquet, + SinkFormatProto::SinkFormatOrc, + ], + )?; + let extra_properties = options.drain_remaining_string_values()?; + + Ok(ConnectorConfig::FilesystemSink(FilesystemSinkConfig { + path, + format: format_proto, + parquet_compression, + extra_properties, + runtime_properties: runtime_props.clone(), + })) + } +} diff --git a/src/streaming_planner/src/connector/sink/iceberg.rs b/src/streaming_planner/src/connector/sink/iceberg.rs new file mode 100644 index 00000000..4f4854a1 --- /dev/null +++ b/src/streaming_planner/src/connector/sink/iceberg.rs @@ -0,0 +1,57 @@ +// 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 datafusion::common::Result; +use protocol::function_stream_graph::{IcebergSinkConfig, SinkFormatProto}; + +use crate::common::Format; +use crate::common::connector_options::ConnectorOptions; +use crate::common::constants::connector_type; +use crate::connector::config::ConnectorConfig; +use crate::connector::provider::SinkProvider; +use crate::connector::sink::runtime_config::SinkRuntimeProperties; +use crate::connector::sink::utils::SinkUtils; + +pub struct IcebergSinkConnector; + +impl SinkProvider for IcebergSinkConnector { + fn name(&self) -> &'static str { + connector_type::ICEBERG + } + + fn build_sink_config( + &self, + options: &mut ConnectorOptions, + format: &Option, + runtime_props: &SinkRuntimeProperties, + ) -> Result { + let path = SinkUtils::require_path(options)?; + let parquet_compression = SinkUtils::extract_parquet_compression(options)?; + let format_proto = SinkUtils::resolve_sink_format( + format, + self.name(), + &[ + SinkFormatProto::SinkFormatCsv, + SinkFormatProto::SinkFormatParquet, + ], + )?; + let extra_properties = options.drain_remaining_string_values()?; + + Ok(ConnectorConfig::IcebergSink(IcebergSinkConfig { + path, + format: format_proto, + parquet_compression, + extra_properties, + runtime_properties: runtime_props.clone(), + })) + } +} diff --git a/src/streaming_planner/src/connector/sink/kafka.rs b/src/streaming_planner/src/connector/sink/kafka.rs new file mode 100644 index 00000000..79292929 --- /dev/null +++ b/src/streaming_planner/src/connector/sink/kafka.rs @@ -0,0 +1,159 @@ +// 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 datafusion::common::{Result, plan_datafusion_err, plan_err}; +use protocol::function_stream_graph::{ + DecimalEncodingProto, FormatConfig, JsonFormatConfig, KafkaAuthConfig, KafkaAuthNone, + KafkaSinkCommitMode, KafkaSinkConfig, RawBytesFormatConfig, RawStringFormatConfig, + TimestampFormatProto, format_config, kafka_auth_config, +}; + +use crate::common::connector_options::ConnectorOptions; +use crate::common::constants::{connector_type, kafka_with_value}; +use crate::common::formats::{ + DecimalEncoding as SqlDecimalEncoding, Format as SqlFormat, + TimestampFormat as SqlTimestampFormat, +}; +use crate::common::with_option_keys as opt; +use crate::connector::config::ConnectorConfig; +use crate::connector::provider::SinkProvider; +use crate::connector::sink::runtime_config::SinkRuntimeProperties; + +pub struct KafkaSinkConnector; + +impl KafkaSinkConnector { + fn sql_format_to_proto(fmt: &SqlFormat) -> Result { + match fmt { + SqlFormat::Json(j) => Ok(FormatConfig { + format: Some(format_config::Format::Json(JsonFormatConfig { + timestamp_format: match j.timestamp_format { + SqlTimestampFormat::RFC3339 => { + TimestampFormatProto::TimestampRfc3339 as i32 + } + SqlTimestampFormat::UnixMillis => { + TimestampFormatProto::TimestampUnixMillis as i32 + } + }, + decimal_encoding: match j.decimal_encoding { + SqlDecimalEncoding::Number => DecimalEncodingProto::DecimalNumber as i32, + SqlDecimalEncoding::String => DecimalEncodingProto::DecimalString as i32, + SqlDecimalEncoding::Bytes => DecimalEncodingProto::DecimalBytes as i32, + }, + include_schema: j.include_schema, + confluent_schema_registry: j.confluent_schema_registry, + schema_id: j.schema_id, + debezium: j.debezium, + unstructured: j.unstructured, + })), + }), + SqlFormat::RawString(_) => Ok(FormatConfig { + format: Some(format_config::Format::RawString(RawStringFormatConfig {})), + }), + SqlFormat::RawBytes(_) => Ok(FormatConfig { + format: Some(format_config::Format::RawBytes(RawBytesFormatConfig {})), + }), + other => plan_err!( + "Kafka sink connector: format '{}' is not supported", + other.name() + ), + } + } +} + +impl SinkProvider for KafkaSinkConnector { + fn name(&self) -> &'static str { + connector_type::KAFKA + } + + fn build_sink_config( + &self, + options: &mut ConnectorOptions, + format: &Option, + _runtime_props: &SinkRuntimeProperties, + ) -> Result { + let bootstrap_servers = match options.pull_opt_str(opt::KAFKA_BOOTSTRAP_SERVERS)? { + Some(s) => s, + None => options + .pull_opt_str(opt::KAFKA_BOOTSTRAP_SERVERS_LEGACY)? + .ok_or_else(|| { + plan_datafusion_err!( + "Kafka connector requires 'bootstrap.servers' in the WITH clause" + ) + })?, + }; + + let topic = options.pull_opt_str(opt::KAFKA_TOPIC)?.ok_or_else(|| { + plan_datafusion_err!("Kafka connector requires 'topic' in the WITH clause") + })?; + + let sql_format = format.as_ref().ok_or_else(|| { + plan_datafusion_err!( + "Kafka sink requires 'format' in the WITH clause (e.g. format = 'json')" + ) + })?; + let proto_format = Self::sql_format_to_proto(sql_format)?; + + let value_subject = options.pull_opt_str(opt::KAFKA_VALUE_SUBJECT)?; + + let commit_mode = match options + .pull_opt_str(opt::KAFKA_SINK_COMMIT_MODE)? + .as_deref() + { + Some(s) + if s == kafka_with_value::SINK_COMMIT_EXACTLY_ONCE_HYPHEN + || s == kafka_with_value::SINK_COMMIT_EXACTLY_ONCE_UNDERSCORE => + { + KafkaSinkCommitMode::KafkaSinkExactlyOnce as i32 + } + Some(s) + if s == kafka_with_value::SINK_COMMIT_AT_LEAST_ONCE_HYPHEN + || s == kafka_with_value::SINK_COMMIT_AT_LEAST_ONCE_UNDERSCORE => + { + KafkaSinkCommitMode::KafkaSinkAtLeastOnce as i32 + } + None => KafkaSinkCommitMode::KafkaSinkAtLeastOnce as i32, + Some(other) => return plan_err!("invalid sink.commit.mode '{other}'"), + }; + + let key_field = match options.pull_opt_str(opt::KAFKA_SINK_KEY_FIELD)? { + Some(s) => Some(s), + None => options.pull_opt_str(opt::KAFKA_KEY_FIELD_LEGACY)?, + }; + let timestamp_field = match options.pull_opt_str(opt::KAFKA_SINK_TIMESTAMP_FIELD)? { + Some(s) => Some(s), + None => options.pull_opt_str(opt::KAFKA_TIMESTAMP_FIELD_LEGACY)?, + }; + + let _ = options.pull_opt_str(opt::TYPE)?; + let _ = options.pull_opt_str(opt::CONNECTOR)?; + + let mut client_configs = options.drain_remaining_string_values()?; + client_configs.remove(opt::CHECKPOINT_INTERVAL_MS); + client_configs.remove(opt::PIPELINE_PARALLELISM); + client_configs.remove(opt::KEY_BY_PARALLELISM); + client_configs.remove(opt::FORMAT); + + Ok(ConnectorConfig::KafkaSink(KafkaSinkConfig { + topic, + bootstrap_servers, + commit_mode, + key_field, + timestamp_field, + auth: Some(KafkaAuthConfig { + auth: Some(kafka_auth_config::Auth::None(KafkaAuthNone {})), + }), + client_configs, + format: Some(proto_format), + value_subject, + })) + } +} diff --git a/src/streaming_planner/src/connector/sink/lancedb.rs b/src/streaming_planner/src/connector/sink/lancedb.rs new file mode 100644 index 00000000..87ce99be --- /dev/null +++ b/src/streaming_planner/src/connector/sink/lancedb.rs @@ -0,0 +1,61 @@ +// 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 datafusion::common::Result; +use protocol::function_stream_graph::{LanceDbSinkConfig, SinkFormatProto}; + +use crate::common::Format; +use crate::common::connector_options::ConnectorOptions; +use crate::common::with_option_keys as opt; +use crate::connector::config::ConnectorConfig; +use crate::connector::provider::SinkProvider; +use crate::connector::sink::runtime_config::SinkRuntimeProperties; +use crate::connector::sink::utils::SinkUtils; + +pub struct LanceDbSinkConnector; + +impl SinkProvider for LanceDbSinkConnector { + fn name(&self) -> &'static str { + "lancedb" + } + + fn build_sink_config( + &self, + options: &mut ConnectorOptions, + _format: &Option, + runtime_props: &SinkRuntimeProperties, + ) -> Result { + let path = SinkUtils::require_path(options)?; + + let s3_bucket = options.pull_opt_str(opt::S3_BUCKET)?; + let s3_region = options.pull_opt_str(opt::S3_REGION)?; + let s3_endpoint = options.pull_opt_str(opt::S3_ENDPOINT)?; + let s3_access_key_id = options.pull_opt_str(opt::S3_ACCESS_KEY_ID)?; + let s3_secret_access_key = options.pull_opt_str(opt::S3_SECRET_ACCESS_KEY)?; + let s3_session_token = options.pull_opt_str(opt::S3_SESSION_TOKEN)?; + + let extra_properties = options.drain_remaining_string_values()?; + + Ok(ConnectorConfig::LanceDbSink(LanceDbSinkConfig { + path, + format: SinkFormatProto::SinkFormatLance as i32, + s3_bucket, + s3_region, + s3_endpoint, + s3_access_key_id, + s3_secret_access_key, + s3_session_token, + extra_properties, + runtime_properties: runtime_props.clone(), + })) + } +} diff --git a/src/streaming_planner/src/connector/sink/mod.rs b/src/streaming_planner/src/connector/sink/mod.rs new file mode 100644 index 00000000..b7d645ca --- /dev/null +++ b/src/streaming_planner/src/connector/sink/mod.rs @@ -0,0 +1,20 @@ +// 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. + +pub mod delta; +pub mod filesystem; +pub mod iceberg; +pub mod kafka; +pub mod lancedb; +pub mod runtime_config; +pub mod s3; +pub mod utils; diff --git a/src/streaming_planner/src/connector/sink/runtime_config.rs b/src/streaming_planner/src/connector/sink/runtime_config.rs new file mode 100644 index 00000000..a383a7ef --- /dev/null +++ b/src/streaming_planner/src/connector/sink/runtime_config.rs @@ -0,0 +1,137 @@ +// 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 datafusion::common::{DataFusionError, Result, plan_err}; + +use function_stream_config::global_config::{ + DEFAULT_OPERATOR_STATE_STORE_MEMORY_BYTES, DEFAULT_SINK_BUFFER_MEMORY_BYTES, +}; +use function_stream_config::streaming_job::DEFAULT_CHECKPOINT_INTERVAL_MS; +use crate::common::connector_options::ConnectorOptions; +use crate::common::with_option_keys as opt; + +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub struct SinkRuntimeConfig { + pub pipeline_parallelism: Option, + pub key_by_parallelism: Option, + pub checkpoint_interval_ms: u64, + pub operator_memory_bytes: u64, + pub sink_memory_bytes: u64, +} + +pub type SinkRuntimeProperties = HashMap; + +impl SinkRuntimeConfig { + pub fn extract_from_options(options: &mut ConnectorOptions) -> Result { + let pipeline_parallelism = options + .pull_opt_u64(opt::PIPELINE_PARALLELISM)? + .map(|v| v as u32); + let key_by_parallelism = options + .pull_opt_u64(opt::KEY_BY_PARALLELISM)? + .map(|v| v as u32); + let checkpoint_interval_ms = options + .pull_opt_u64(opt::CHECKPOINT_INTERVAL_MS)? + .unwrap_or(DEFAULT_CHECKPOINT_INTERVAL_MS); + let operator_memory_bytes = options + .pull_opt_u64(opt::OPERATOR_MEMORY_BYTES)? + .unwrap_or(DEFAULT_OPERATOR_STATE_STORE_MEMORY_BYTES); + let sink_memory_bytes = options + .pull_opt_u64(opt::SINK_MEMORY_BYTES)? + .unwrap_or(DEFAULT_SINK_BUFFER_MEMORY_BYTES); + Ok(Self { + pipeline_parallelism, + key_by_parallelism, + checkpoint_interval_ms, + operator_memory_bytes, + sink_memory_bytes, + }) + } + + pub fn from_options_map(opts: &HashMap) -> Result { + let pipeline_parallelism = parse_opt_u32(opts, opt::PIPELINE_PARALLELISM)?; + let key_by_parallelism = parse_opt_u32(opts, opt::KEY_BY_PARALLELISM)?; + let checkpoint_interval_ms = parse_opt_u64(opts, opt::CHECKPOINT_INTERVAL_MS)? + .unwrap_or(DEFAULT_CHECKPOINT_INTERVAL_MS); + let operator_memory_bytes = parse_opt_u64(opts, opt::OPERATOR_MEMORY_BYTES)? + .unwrap_or(DEFAULT_OPERATOR_STATE_STORE_MEMORY_BYTES); + let sink_memory_bytes = parse_opt_u64(opts, opt::SINK_MEMORY_BYTES)? + .unwrap_or(DEFAULT_SINK_BUFFER_MEMORY_BYTES); + Ok(Self { + pipeline_parallelism, + key_by_parallelism, + checkpoint_interval_ms, + operator_memory_bytes, + sink_memory_bytes, + }) + } + + pub fn to_runtime_properties(&self) -> HashMap { + let mut out = HashMap::new(); + if let Some(v) = self.pipeline_parallelism { + out.insert(opt::PIPELINE_PARALLELISM.to_string(), v.to_string()); + } + if let Some(v) = self.key_by_parallelism { + out.insert(opt::KEY_BY_PARALLELISM.to_string(), v.to_string()); + } + out.insert( + opt::CHECKPOINT_INTERVAL_MS.to_string(), + self.checkpoint_interval_ms.to_string(), + ); + out.insert( + opt::OPERATOR_MEMORY_BYTES.to_string(), + self.operator_memory_bytes.to_string(), + ); + out.insert( + opt::SINK_MEMORY_BYTES.to_string(), + self.sink_memory_bytes.to_string(), + ); + out + } +} + +fn parse_opt_u32(opts: &HashMap, key: &str) -> Result> { + let Some(raw) = opts.get(key) else { + return Ok(None); + }; + let normalized = normalize_numeric_option(raw); + let parsed = normalized.parse::().map_err(|_| { + DataFusionError::Plan(format!( + "WITH option '{key}' expects unsigned integer, got '{raw}'" + )) + })?; + if parsed == 0 { + return plan_err!("WITH option '{key}' must be > 0"); + } + Ok(Some(parsed)) +} + +fn parse_opt_u64(opts: &HashMap, key: &str) -> Result> { + let Some(raw) = opts.get(key) else { + return Ok(None); + }; + let normalized = normalize_numeric_option(raw); + let parsed = normalized.parse::().map_err(|_| { + DataFusionError::Plan(format!( + "WITH option '{key}' expects unsigned integer, got '{raw}'" + )) + })?; + if parsed == 0 { + return plan_err!("WITH option '{key}' must be > 0"); + } + Ok(Some(parsed)) +} + +fn normalize_numeric_option(raw: &str) -> &str { + raw.trim().trim_matches('\'').trim_matches('"').trim() +} diff --git a/src/streaming_planner/src/connector/sink/s3.rs b/src/streaming_planner/src/connector/sink/s3.rs new file mode 100644 index 00000000..51cf85bf --- /dev/null +++ b/src/streaming_planner/src/connector/sink/s3.rs @@ -0,0 +1,75 @@ +// 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 datafusion::common::Result; +use protocol::function_stream_graph::{S3SinkConfig, SinkFormatProto}; + +use crate::common::Format; +use crate::common::connector_options::ConnectorOptions; +use crate::common::constants::connector_type; +use crate::common::with_option_keys as opt; +use crate::connector::config::ConnectorConfig; +use crate::connector::provider::SinkProvider; +use crate::connector::sink::runtime_config::SinkRuntimeProperties; +use crate::connector::sink::utils::SinkUtils; + +pub struct S3SinkConnector; + +impl SinkProvider for S3SinkConnector { + fn name(&self) -> &'static str { + connector_type::S3 + } + + fn build_sink_config( + &self, + options: &mut ConnectorOptions, + format: &Option, + runtime_props: &SinkRuntimeProperties, + ) -> Result { + let path = SinkUtils::require_path(options)?; + + let format_proto = SinkUtils::resolve_sink_format( + format, + self.name(), + &[ + SinkFormatProto::SinkFormatCsv, + SinkFormatProto::SinkFormatParquet, + ], + )?; + + let bucket = SinkUtils::require_str(options, opt::S3_BUCKET, self.name())?; + let region = options + .pull_opt_str(opt::S3_REGION)? + .unwrap_or_else(|| "us-east-1".to_string()); + let endpoint = options.pull_opt_str(opt::S3_ENDPOINT)?; + let access_key_id = options.pull_opt_str(opt::S3_ACCESS_KEY_ID)?; + let secret_access_key = options.pull_opt_str(opt::S3_SECRET_ACCESS_KEY)?; + let session_token = options.pull_opt_str(opt::S3_SESSION_TOKEN)?; + + let parquet_compression = SinkUtils::extract_parquet_compression(options)?; + let extra_properties = options.drain_remaining_string_values()?; + + Ok(ConnectorConfig::S3Sink(S3SinkConfig { + path, + format: format_proto, + bucket, + region, + endpoint, + access_key_id, + secret_access_key, + session_token, + parquet_compression, + extra_properties, + runtime_properties: runtime_props.clone(), + })) + } +} diff --git a/src/streaming_planner/src/connector/sink/utils.rs b/src/streaming_planner/src/connector/sink/utils.rs new file mode 100644 index 00000000..481ad3b2 --- /dev/null +++ b/src/streaming_planner/src/connector/sink/utils.rs @@ -0,0 +1,91 @@ +// 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 datafusion::common::{DataFusionError, Result, plan_err}; +use protocol::function_stream_graph::{ParquetCompressionProto, SinkFormatProto}; + +use crate::common::Format; +use crate::common::connector_options::ConnectorOptions; +use crate::common::constants::parquet_compression_value; +use crate::common::with_option_keys as opt; + +pub struct SinkUtils; + +impl SinkUtils { + pub fn require_path(options: &mut ConnectorOptions) -> Result { + if let Some(v) = options.pull_opt_str(opt::PATH)? { + return Ok(v); + } + if let Some(v) = options.pull_opt_str(opt::SINK_PATH)? { + return Ok(v); + } + plan_err!("Missing required WITH option 'path' (or 'sink.path')") + } + + pub fn extract_parquet_compression(options: &mut ConnectorOptions) -> Result> { + let Some(v) = options.pull_opt_str(opt::PARQUET_COMPRESSION)? else { + return Ok(None); + }; + let parsed = match v.to_ascii_lowercase().as_str() { + parquet_compression_value::UNCOMPRESSED => { + ParquetCompressionProto::ParquetCompressionUncompressed + } + parquet_compression_value::SNAPPY => ParquetCompressionProto::ParquetCompressionSnappy, + parquet_compression_value::GZIP => ParquetCompressionProto::ParquetCompressionGzip, + parquet_compression_value::ZSTD => ParquetCompressionProto::ParquetCompressionZstd, + parquet_compression_value::LZ4 => ParquetCompressionProto::ParquetCompressionLz4, + parquet_compression_value::LZ4_RAW => ParquetCompressionProto::ParquetCompressionLz4Raw, + other => return plan_err!("Unsupported parquet.compression '{other}'"), + }; + Ok(Some(parsed as i32)) + } + + pub fn require_str( + options: &mut ConnectorOptions, + key: &str, + connector: &str, + ) -> Result { + options.pull_opt_str(key)?.ok_or_else(|| { + DataFusionError::Plan(format!( + "Connector '{connector}' requires WITH option '{key}'" + )) + }) + } + + pub fn resolve_sink_format( + format: &Option, + connector_name: &str, + supported_formats: &[SinkFormatProto], + ) -> Result { + let proto_format = match format { + Some(Format::Csv(_)) => SinkFormatProto::SinkFormatCsv, + Some(Format::Json(_)) => SinkFormatProto::SinkFormatJsonl, + Some(Format::Avro(_)) => SinkFormatProto::SinkFormatAvro, + Some(Format::Parquet(_)) => SinkFormatProto::SinkFormatParquet, + Some(Format::Lance(_)) => SinkFormatProto::SinkFormatLance, + Some(f) => { + return plan_err!("Format '{f:?}' cannot be mapped to a sink format"); + } + None => { + return plan_err!("Connector '{connector_name}' requires a format to be specified"); + } + }; + + if !supported_formats.contains(&proto_format) { + return plan_err!( + "Format {proto_format:?} is not supported by connector '{connector_name}'" + ); + } + + Ok(proto_format as i32) + } +} diff --git a/src/streaming_planner/src/connector/source/kafka.rs b/src/streaming_planner/src/connector/source/kafka.rs new file mode 100644 index 00000000..71672279 --- /dev/null +++ b/src/streaming_planner/src/connector/source/kafka.rs @@ -0,0 +1,185 @@ +// 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 datafusion::common::{Result, plan_datafusion_err, plan_err}; +use protocol::function_stream_graph::{ + BadDataPolicy, DecimalEncodingProto, FormatConfig, JsonFormatConfig, KafkaAuthConfig, + KafkaAuthNone, KafkaOffsetMode, KafkaReadMode, KafkaSourceConfig, RawBytesFormatConfig, + RawStringFormatConfig, TimestampFormatProto, format_config, kafka_auth_config, +}; + +use crate::common::connector_options::ConnectorOptions; +use crate::common::constants::{connector_type, kafka_with_value}; +use crate::common::formats::{ + BadData, DecimalEncoding as SqlDecimalEncoding, Format as SqlFormat, + TimestampFormat as SqlTimestampFormat, +}; +use crate::common::with_option_keys as opt; +use crate::connector::config::ConnectorConfig; +use crate::connector::provider::SourceProvider; + +pub struct KafkaSourceConnector; + +impl KafkaSourceConnector { + fn sql_format_to_proto(fmt: &SqlFormat) -> Result { + match fmt { + SqlFormat::Json(j) => Ok(FormatConfig { + format: Some(format_config::Format::Json(JsonFormatConfig { + timestamp_format: match j.timestamp_format { + SqlTimestampFormat::RFC3339 => { + TimestampFormatProto::TimestampRfc3339 as i32 + } + SqlTimestampFormat::UnixMillis => { + TimestampFormatProto::TimestampUnixMillis as i32 + } + }, + decimal_encoding: match j.decimal_encoding { + SqlDecimalEncoding::Number => DecimalEncodingProto::DecimalNumber as i32, + SqlDecimalEncoding::String => DecimalEncodingProto::DecimalString as i32, + SqlDecimalEncoding::Bytes => DecimalEncodingProto::DecimalBytes as i32, + }, + include_schema: j.include_schema, + confluent_schema_registry: j.confluent_schema_registry, + schema_id: j.schema_id, + debezium: j.debezium, + unstructured: j.unstructured, + })), + }), + SqlFormat::RawString(_) => Ok(FormatConfig { + format: Some(format_config::Format::RawString(RawStringFormatConfig {})), + }), + SqlFormat::RawBytes(_) => Ok(FormatConfig { + format: Some(format_config::Format::RawBytes(RawBytesFormatConfig {})), + }), + other => plan_err!( + "Kafka source connector: format '{}' is not supported", + other.name() + ), + } + } + + fn bad_data_to_proto(bad: &BadData) -> i32 { + match bad { + BadData::Fail {} => BadDataPolicy::BadDataFail as i32, + BadData::Drop {} => BadDataPolicy::BadDataDrop as i32, + } + } +} + +impl SourceProvider for KafkaSourceConnector { + fn name(&self) -> &'static str { + connector_type::KAFKA + } + + fn build_source_config( + &self, + options: &mut ConnectorOptions, + format: &Option, + bad_data: BadData, + ) -> Result { + let bootstrap_servers = match options.pull_opt_str(opt::KAFKA_BOOTSTRAP_SERVERS)? { + Some(s) => s, + None => options + .pull_opt_str(opt::KAFKA_BOOTSTRAP_SERVERS_LEGACY)? + .ok_or_else(|| { + plan_datafusion_err!( + "Kafka connector requires 'bootstrap.servers' in the WITH clause" + ) + })?, + }; + + let topic = options.pull_opt_str(opt::KAFKA_TOPIC)?.ok_or_else(|| { + plan_datafusion_err!("Kafka connector requires 'topic' in the WITH clause") + })?; + + let sql_format = format.as_ref().ok_or_else(|| { + plan_datafusion_err!( + "Kafka source requires 'format' in the WITH clause (e.g. format = 'json')" + ) + })?; + let proto_format = Self::sql_format_to_proto(sql_format)?; + + let rate_limit = options + .pull_opt_u64(opt::KAFKA_RATE_LIMIT_MESSAGES_PER_SECOND)? + .map(|v| v.clamp(1, u32::MAX as u64) as u32) + .unwrap_or(0); + + let value_subject = options.pull_opt_str(opt::KAFKA_VALUE_SUBJECT)?; + + let offset_mode = match options + .pull_opt_str(opt::KAFKA_SCAN_STARTUP_MODE)? + .as_deref() + { + Some(s) if s == kafka_with_value::SCAN_LATEST => { + KafkaOffsetMode::KafkaOffsetLatest as i32 + } + Some(s) if s == kafka_with_value::SCAN_EARLIEST => { + KafkaOffsetMode::KafkaOffsetEarliest as i32 + } + Some(s) + if s == kafka_with_value::SCAN_GROUP_OFFSETS + || s == kafka_with_value::SCAN_GROUP => + { + KafkaOffsetMode::KafkaOffsetGroup as i32 + } + None => KafkaOffsetMode::KafkaOffsetGroup as i32, + Some(other) => { + return plan_err!( + "invalid scan.startup.mode '{other}'; expected latest, earliest, or group-offsets" + ); + } + }; + + let read_mode = match options.pull_opt_str(opt::KAFKA_ISOLATION_LEVEL)?.as_deref() { + Some(s) if s == kafka_with_value::ISOLATION_READ_COMMITTED => { + KafkaReadMode::KafkaReadCommitted as i32 + } + Some(s) if s == kafka_with_value::ISOLATION_READ_UNCOMMITTED => { + KafkaReadMode::KafkaReadUncommitted as i32 + } + None => KafkaReadMode::KafkaReadDefault as i32, + Some(other) => return plan_err!("invalid isolation.level '{other}'"), + }; + + let group_id = match options.pull_opt_str(opt::KAFKA_GROUP_ID)? { + Some(s) => Some(s), + None => options.pull_opt_str(opt::KAFKA_GROUP_ID_LEGACY)?, + }; + let group_id_prefix = options.pull_opt_str(opt::KAFKA_GROUP_ID_PREFIX)?; + + let _ = options.pull_opt_str(opt::TYPE)?; + let _ = options.pull_opt_str(opt::CONNECTOR)?; + + let mut client_configs = options.drain_remaining_string_values()?; + client_configs.remove(opt::CHECKPOINT_INTERVAL_MS); + client_configs.remove(opt::PIPELINE_PARALLELISM); + client_configs.remove(opt::KEY_BY_PARALLELISM); + client_configs.remove(opt::FORMAT); + + Ok(ConnectorConfig::KafkaSource(KafkaSourceConfig { + topic, + bootstrap_servers, + group_id, + group_id_prefix, + offset_mode, + read_mode, + auth: Some(KafkaAuthConfig { + auth: Some(kafka_auth_config::Auth::None(KafkaAuthNone {})), + }), + client_configs, + format: Some(proto_format), + bad_data_policy: Self::bad_data_to_proto(&bad_data), + rate_limit_msgs_per_sec: rate_limit, + value_subject, + })) + } +} diff --git a/src/streaming_planner/src/connector/source/mod.rs b/src/streaming_planner/src/connector/source/mod.rs new file mode 100644 index 00000000..b9574391 --- /dev/null +++ b/src/streaming_planner/src/connector/source/mod.rs @@ -0,0 +1,13 @@ +// 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. + +pub mod kafka; diff --git a/src/streaming_planner/src/functions/mod.rs b/src/streaming_planner/src/functions/mod.rs new file mode 100644 index 00000000..9ec1007f --- /dev/null +++ b/src/streaming_planner/src/functions/mod.rs @@ -0,0 +1,612 @@ +// 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::schema::StreamSchemaProvider; +use datafusion::arrow::array::{ + Array, ArrayRef, StringArray, UnionArray, + builder::{FixedSizeBinaryBuilder, ListBuilder, StringBuilder}, + cast::{AsArray, as_string_array}, + types::{Float64Type, Int64Type}, +}; +use datafusion::arrow::datatypes::{DataType, Field, UnionFields, UnionMode}; +use datafusion::arrow::row::{RowConverter, SortField}; +use datafusion::common::{DataFusionError, ScalarValue}; +use datafusion::common::{Result, TableReference}; +use datafusion::execution::FunctionRegistry; +use datafusion::logical_expr::expr::{Alias, ScalarFunction}; +use datafusion::logical_expr::{ + ColumnarValue, LogicalPlan, Projection, ScalarFunctionArgs, ScalarUDFImpl, Signature, + TypeSignature, Volatility, create_udf, +}; +use datafusion::prelude::{Expr, col}; +use serde_json_path::JsonPath; +use std::any::Any; +use std::collections::HashMap; +use std::fmt::{Debug, Write}; +use std::sync::{Arc, OnceLock}; + +use crate::common::constants::scalar_fn; + +/// Borrowed from DataFusion +/// +/// Creates a singleton `ScalarUDF` of the `$UDF` function named `$GNAME` and a +/// function named `$NAME` which returns that function named $NAME. +/// +/// This is used to ensure creating the list of `ScalarUDF` only happens once. +#[macro_export] +macro_rules! make_udf_function { + ($UDF:ty, $GNAME:ident, $NAME:ident) => { + /// Singleton instance of the function + static $GNAME: std::sync::OnceLock> = + std::sync::OnceLock::new(); + + /// Return a [`ScalarUDF`] for [`$UDF`] + /// + /// [`ScalarUDF`]: datafusion_expr::ScalarUDF + pub fn $NAME() -> std::sync::Arc { + $GNAME + .get_or_init(|| { + std::sync::Arc::new(datafusion::logical_expr::ScalarUDF::new_from_impl( + <$UDF>::default(), + )) + }) + .clone() + } + }; +} + +make_udf_function!(MultiHashFunction, MULTI_HASH, multi_hash); + +pub fn register_all(registry: &mut dyn FunctionRegistry) { + registry + .register_udf(Arc::new(create_udf( + scalar_fn::GET_FIRST_JSON_OBJECT, + vec![DataType::Utf8, DataType::Utf8], + DataType::Utf8, + Volatility::Immutable, + Arc::new(get_first_json_object), + ))) + .unwrap(); + + registry + .register_udf(Arc::new(create_udf( + scalar_fn::EXTRACT_JSON, + vec![DataType::Utf8, DataType::Utf8], + DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))), + Volatility::Immutable, + Arc::new(extract_json), + ))) + .unwrap(); + + registry + .register_udf(Arc::new(create_udf( + scalar_fn::EXTRACT_JSON_STRING, + vec![DataType::Utf8, DataType::Utf8], + DataType::Utf8, + Volatility::Immutable, + Arc::new(extract_json_string), + ))) + .unwrap(); + + registry + .register_udf(Arc::new(create_udf( + scalar_fn::SERIALIZE_JSON_UNION, + vec![DataType::Union(union_fields(), UnionMode::Sparse)], + DataType::Utf8, + Volatility::Immutable, + Arc::new(serialize_json_union), + ))) + .unwrap(); + + registry.register_udf(multi_hash()).unwrap(); +} + +fn parse_path(name: &str, path: &ScalarValue) -> Result> { + let path = match path { + ScalarValue::Utf8(Some(s)) => JsonPath::parse(s) + .map_err(|e| DataFusionError::Execution(format!("Invalid json path '{s}': {e:?}")))?, + ScalarValue::Utf8(None) => { + return Err(DataFusionError::Execution(format!( + "The path argument to {name} cannot be null" + ))); + } + _ => { + return Err(DataFusionError::Execution(format!( + "The path argument to {name} must be of type TEXT" + ))); + } + }; + + Ok(Arc::new(path)) +} + +// Hash function that can take any number of arguments and produces a fast (non-cryptographic) +// 128-bit hash from their string representations +#[derive(Debug)] +pub struct MultiHashFunction { + signature: Signature, +} + +impl MultiHashFunction { + pub fn invoke(&self, args: &[ColumnarValue]) -> Result { + let mut hasher = xxhash_rust::xxh3::Xxh3::new(); + + let all_scalar = args.iter().all(|a| matches!(a, ColumnarValue::Scalar(_))); + + let length = args + .iter() + .map(|t| match t { + ColumnarValue::Scalar(_) => 1, + ColumnarValue::Array(a) => a.len(), + }) + .max() + .ok_or_else(|| { + DataFusionError::Plan("multi_hash must have at least one argument".to_string()) + })?; + + let row_builder = RowConverter::new( + args.iter() + .map(|t| SortField::new(t.data_type().clone())) + .collect(), + )?; + + let arrays = args + .iter() + .map(|c| c.clone().into_array(length)) + .collect::>>()?; + let rows = row_builder.convert_columns(&arrays)?; + + if all_scalar { + hasher.update(rows.row(0).as_ref()); + let result = hasher.digest128().to_be_bytes().to_vec(); + hasher.reset(); + Ok(ColumnarValue::Scalar(ScalarValue::FixedSizeBinary( + size_of::() as i32, + Some(result), + ))) + } else { + let mut builder = + FixedSizeBinaryBuilder::with_capacity(length, size_of::() as i32); + + for row in rows.iter() { + hasher.update(row.as_ref()); + builder.append_value(hasher.digest128().to_be_bytes())?; + hasher.reset(); + } + + Ok(ColumnarValue::Array(Arc::new(builder.finish()))) + } + } +} + +impl Default for MultiHashFunction { + fn default() -> Self { + Self { + signature: Signature::new(TypeSignature::VariadicAny, Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for MultiHashFunction { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + scalar_fn::MULTI_HASH + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::FixedSizeBinary(size_of::() as i32)) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + self.invoke(&args.args) + } +} + +fn json_function( + name: &str, + f: F, + to_scalar: ToS, + args: &[ColumnarValue], +) -> Result +where + ArrayT: Array + FromIterator> + 'static, + F: Fn(serde_json::Value, &JsonPath) -> Option, + ToS: Fn(Option) -> ScalarValue, +{ + assert_eq!(args.len(), 2); + Ok(match (&args[0], &args[1]) { + (ColumnarValue::Array(values), ColumnarValue::Scalar(path)) => { + let path = parse_path(name, path)?; + let vs = as_string_array(values); + ColumnarValue::Array(Arc::new( + vs.iter() + .map(|s| s.and_then(|s| f(serde_json::from_str(s).ok()?, &path))) + .collect::(), + ) as ArrayRef) + } + (ColumnarValue::Scalar(value), ColumnarValue::Scalar(path)) => { + let path = parse_path(name, path)?; + let ScalarValue::Utf8(value) = value else { + return Err(DataFusionError::Execution(format!( + "The value argument to {name} must be of type TEXT" + ))); + }; + + let result = value + .as_ref() + .and_then(|v| f(serde_json::from_str(v).ok()?, &path)); + ColumnarValue::Scalar(to_scalar(result)) + } + _ => { + return Err(DataFusionError::Execution( + "The path argument to {name} must be a literal".to_string(), + )); + } + }) +} + +pub fn extract_json(args: &[ColumnarValue]) -> Result { + assert_eq!(args.len(), 2); + + let inner = |s, path: &JsonPath| { + Some( + path.query(&serde_json::from_str(s).ok()?) + .iter() + .map(|v| Some(v.to_string())) + .collect::>>(), + ) + }; + + Ok(match (&args[0], &args[1]) { + (ColumnarValue::Array(values), ColumnarValue::Scalar(path)) => { + let path = parse_path("extract_json", path)?; + let values = as_string_array(values); + + let mut builder = ListBuilder::with_capacity(StringBuilder::new(), values.len()); + + let queried = values.iter().map(|s| s.and_then(|s| inner(s, &path))); + + for v in queried { + builder.append_option(v); + } + + ColumnarValue::Array(Arc::new(builder.finish())) + } + (ColumnarValue::Scalar(value), ColumnarValue::Scalar(path)) => { + let path = parse_path("extract_json", path)?; + let ScalarValue::Utf8(v) = value else { + return Err(DataFusionError::Execution( + "The value argument to extract_json must be of type TEXT".to_string(), + )); + }; + + let mut builder = ListBuilder::with_capacity(StringBuilder::new(), 1); + let result = v.as_ref().and_then(|s| inner(s, &path)); + builder.append_option(result); + + ColumnarValue::Scalar(ScalarValue::List(Arc::new(builder.finish()))) + } + _ => { + return Err(DataFusionError::Execution( + "The path argument to extract_json must be a literal".to_string(), + )); + } + }) +} + +pub fn get_first_json_object(args: &[ColumnarValue]) -> Result { + json_function::( + "get_first_json_object", + |s, path| path.query(&s).first().map(|v| v.to_string()), + |s| s.as_deref().into(), + args, + ) +} + +pub fn extract_json_string(args: &[ColumnarValue]) -> Result { + json_function::( + "extract_json_string", + |s, path| { + path.query(&s) + .first() + .and_then(|v| v.as_str().map(|s| s.to_string())) + }, + |s| s.as_deref().into(), + args, + ) +} + +// This code is vendored from +// https://github.com/datafusion-contrib/datafusion-functions-json/blob/main/src/common_union.rs +// as the `is_json_union` function is not public. It should be kept in sync with that code so +// that we are able to detect JSON unions and rewrite them to serialized JSON for sinks. +pub fn is_json_union(data_type: &DataType) -> bool { + match data_type { + DataType::Union(fields, UnionMode::Sparse) => fields == &union_fields(), + _ => false, + } +} + +pub const TYPE_ID_NULL: i8 = 0; +const TYPE_ID_BOOL: i8 = 1; +const TYPE_ID_INT: i8 = 2; +const TYPE_ID_FLOAT: i8 = 3; +const TYPE_ID_STR: i8 = 4; +const TYPE_ID_ARRAY: i8 = 5; +const TYPE_ID_OBJECT: i8 = 6; + +fn union_fields() -> UnionFields { + static FIELDS: OnceLock = OnceLock::new(); + FIELDS + .get_or_init(|| { + let json_metadata: HashMap = + HashMap::from_iter(vec![("is_json".to_string(), "true".to_string())]); + UnionFields::from_iter([ + ( + TYPE_ID_NULL, + Arc::new(Field::new("null", DataType::Null, true)), + ), + ( + TYPE_ID_BOOL, + Arc::new(Field::new("bool", DataType::Boolean, false)), + ), + ( + TYPE_ID_INT, + Arc::new(Field::new("int", DataType::Int64, false)), + ), + ( + TYPE_ID_FLOAT, + Arc::new(Field::new("float", DataType::Float64, false)), + ), + ( + TYPE_ID_STR, + Arc::new(Field::new("str", DataType::Utf8, false)), + ), + ( + TYPE_ID_ARRAY, + Arc::new( + Field::new("array", DataType::Utf8, false) + .with_metadata(json_metadata.clone()), + ), + ), + ( + TYPE_ID_OBJECT, + Arc::new( + Field::new("object", DataType::Utf8, false) + .with_metadata(json_metadata.clone()), + ), + ), + ]) + }) + .clone() +} +// End vendored code + +pub fn serialize_json_union(args: &[ColumnarValue]) -> Result { + assert_eq!(args.len(), 1); + let array = match args.first().unwrap() { + ColumnarValue::Array(a) => a.clone(), + ColumnarValue::Scalar(s) => s.to_array_of_size(1)?, + }; + + let mut b = StringBuilder::with_capacity(array.len(), array.get_array_memory_size()); + + write_union(&mut b, &array)?; + + Ok(ColumnarValue::Array(Arc::new(b.finish()))) +} + +fn write_union(b: &mut StringBuilder, array: &ArrayRef) -> Result<(), std::fmt::Error> { + assert!( + is_json_union(array.data_type()), + "array item is not a valid JSON union" + ); + let json_union = array.as_any().downcast_ref::().unwrap(); + + for i in 0..json_union.len() { + if json_union.is_null(i) { + b.append_null(); + } else { + write_value(b, json_union.type_id(i), &json_union.value(i))?; + b.append_value(""); + } + } + + Ok(()) +} + +fn write_value(b: &mut StringBuilder, id: i8, a: &ArrayRef) -> Result<(), std::fmt::Error> { + match id { + TYPE_ID_NULL => write!(b, "null")?, + TYPE_ID_BOOL => write!(b, "{}", a.as_boolean().value(0))?, + TYPE_ID_INT => write!(b, "{}", a.as_primitive::().value(0))?, + TYPE_ID_FLOAT => write!(b, "{}", a.as_primitive::().value(0))?, + TYPE_ID_STR => { + // assumes that this is already a valid (escaped) json string as the only way to + // construct these values are by parsing (valid) JSON + b.write_char('"')?; + b.write_str(a.as_string::().value(0))?; + b.write_char('"')?; + } + TYPE_ID_ARRAY => { + b.write_str(a.as_string::().value(0))?; + } + TYPE_ID_OBJECT => { + b.write_str(a.as_string::().value(0))?; + } + _ => unreachable!("invalid union type in JSON union: {}", id), + } + + Ok(()) +} + +pub fn serialize_outgoing_json( + registry: &StreamSchemaProvider, + node: Arc, +) -> LogicalPlan { + let exprs = node + .schema() + .fields() + .iter() + .map(|f| { + if is_json_union(f.data_type()) { + Expr::Alias(Alias::new( + Expr::ScalarFunction(ScalarFunction::new_udf( + registry.udf(scalar_fn::SERIALIZE_JSON_UNION).unwrap(), + vec![col(f.name())], + )), + Option::::None, + f.name(), + )) + } else { + col(f.name()) + } + }) + .collect(); + + LogicalPlan::Projection(Projection::try_new(exprs, node).unwrap()) +} + +#[cfg(test)] +mod test { + use datafusion::arrow::array::StringArray; + use datafusion::arrow::array::builder::{ListBuilder, StringBuilder}; + use datafusion::common::ScalarValue; + use std::sync::Arc; + + #[test] + fn test_extract_json() { + let input = Arc::new(StringArray::from(vec![ + r#"{"a": 1, "b": 2, "c": { "d": "hello" }}"#, + r#"{"a": 3, "b": 4}"#, + r#"{"a": 5, "b": 6}"#, + ])); + + let path = "$.c.d"; + + let result = super::extract_json(&[ + super::ColumnarValue::Array(input), + super::ColumnarValue::Scalar(path.into()), + ]) + .unwrap(); + + let mut expected = ListBuilder::new(StringBuilder::new()); + expected.append_value(vec![Some("\"hello\"".to_string())]); + expected.append_value(Vec::>::new()); + expected.append_value(Vec::>::new()); + if let super::ColumnarValue::Array(result) = result { + assert_eq!(*result, expected.finish()); + } else { + panic!("Expected array, got scalar"); + } + + let result = super::extract_json(&[ + super::ColumnarValue::Scalar(r#"{"a": 1, "b": 2, "c": { "d": "hello" }}"#.into()), + super::ColumnarValue::Scalar(path.into()), + ]) + .unwrap(); + + let mut expected = ListBuilder::with_capacity(StringBuilder::new(), 1); + expected.append_value(vec![Some("\"hello\"".to_string())]); + + if let super::ColumnarValue::Scalar(ScalarValue::List(result)) = result { + assert_eq!(*result, expected.finish()); + } else { + panic!("Expected scalar list"); + } + } + + #[test] + fn test_get_first_json_object() { + let input = Arc::new(StringArray::from(vec![ + r#"{"a": 1, "b": 2}"#, + r#"{"a": 3}"#, + r#"{"a": 5, "b": 6}"#, + ])); + + let path = "$.b"; + + let result = super::get_first_json_object(&[ + super::ColumnarValue::Array(input), + super::ColumnarValue::Scalar(path.into()), + ]) + .unwrap(); + + let expected = StringArray::from(vec![Some("2"), None, Some("6")]); + + if let super::ColumnarValue::Array(result) = result { + assert_eq!(*result, expected); + } else { + panic!("Expected array, got scalar"); + } + + let result = super::get_first_json_object(&[ + super::ColumnarValue::Scalar(r#"{"a": 1, "b": 2, "c": { "d": "hello" }}"#.into()), + super::ColumnarValue::Scalar("$.c.d".into()), + ]) + .unwrap(); + + let expected = ScalarValue::Utf8(Some("\"hello\"".to_string())); + + if let super::ColumnarValue::Scalar(result) = result { + assert_eq!(result, expected); + } else { + panic!("Expected scalar"); + } + } + + #[test] + fn test_extract_json_string() { + let input = Arc::new(StringArray::from(vec![ + r#"{"a": 1, "b": 2, "c": { "d": "hello" }}"#, + r#"{"a": 3, "b": 4}"#, + r#"{"a": 5, "b": 6}"#, + ])); + + let path = "$.c.d"; + + let result = super::extract_json_string(&[ + super::ColumnarValue::Array(input), + super::ColumnarValue::Scalar(path.into()), + ]) + .unwrap(); + + let expected = StringArray::from(vec![Some("hello"), None, None]); + + if let super::ColumnarValue::Array(result) = result { + assert_eq!(*result, expected); + } else { + panic!("Expected array, got scalar"); + } + + let result = super::extract_json_string(&[ + super::ColumnarValue::Scalar(r#"{"a": 1, "b": 2, "c": { "d": "hello" }}"#.into()), + super::ColumnarValue::Scalar(path.into()), + ]) + .unwrap(); + + let expected = ScalarValue::Utf8(Some("hello".to_string())); + + if let super::ColumnarValue::Scalar(result) = result { + assert_eq!(result, expected); + } else { + panic!("Expected scalar"); + } + } +} diff --git a/src/streaming_planner/src/lib.rs b/src/streaming_planner/src/lib.rs new file mode 100644 index 00000000..040e0d42 --- /dev/null +++ b/src/streaming_planner/src/lib.rs @@ -0,0 +1,32 @@ +// 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. + +//! Streaming SQL planning (logical graph, connectors, schema, physical codec). +//! +//! The `function-stream` binary/library re-exports this crate as `function_stream::sql` for +//! stable `crate::sql::…` paths in in-tree modules. + +pub mod api; +pub mod common; + +pub mod analysis; +pub mod connector; +pub mod functions; +pub mod logical_node; +pub mod logical_planner; +pub mod parse; +pub mod physical; +pub mod planning_runtime; +pub mod schema; +pub mod types; + +pub use analysis::rewrite_plan; diff --git a/src/streaming_planner/src/logical_node/aggregate.rs b/src/streaming_planner/src/logical_node/aggregate.rs new file mode 100644 index 00000000..d8a96472 --- /dev/null +++ b/src/streaming_planner/src/logical_node/aggregate.rs @@ -0,0 +1,644 @@ +// 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::fmt::Formatter; +use std::sync::Arc; +use std::time::Duration; + +use arrow_array::types::IntervalMonthDayNanoType; +use datafusion::common::{Column, DFSchemaRef, Result, ScalarValue, internal_err}; +use datafusion::logical_expr::{ + self, BinaryExpr, Expr, Extension, LogicalPlan, UserDefinedLogicalNodeCore, + expr::ScalarFunction, +}; +use datafusion_common::{DFSchema, DataFusionError, plan_err}; +use datafusion_expr::Aggregate; +use datafusion_proto::physical_plan::to_proto::serialize_physical_expr; +use datafusion_proto::physical_plan::{AsExecutionPlan, DefaultPhysicalExtensionCodec}; +use datafusion_proto::protobuf::PhysicalPlanNode; +use prost::Message; +use protocol::function_stream_graph::{ + SessionWindowAggregateOperator, SlidingWindowAggregateOperator, TumblingWindowAggregateOperator, +}; + +use crate::multifield_partial_ord; +use crate::common::constants::{extension_node, proto_operator_name}; +use crate::common::{FsSchema, FsSchemaRef}; +use crate::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; +use crate::logical_node::{ + CompiledTopologyNode, StreamingOperatorBlueprint, SystemTimestampInjectorNode, +}; +use crate::logical_planner::planner::{NamedNode, Planner, SplitPlanOutput}; +use crate::physical::{StreamingExtensionCodec, window}; +use crate::types::{ + QualifiedField, TIMESTAMP_FIELD, WindowBehavior, WindowType, build_df_schema, + build_df_schema_with_metadata, extract_qualified_fields, +}; + +pub const STREAM_AGG_EXTENSION_NAME: &str = extension_node::STREAM_WINDOW_AGGREGATE; + +/// Represents a streaming windowed aggregation node in the logical plan. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct StreamWindowAggregateNode { + pub(crate) window_spec: WindowBehavior, + pub(crate) base_agg_plan: LogicalPlan, + pub(crate) output_schema: DFSchemaRef, + pub(crate) partition_keys: Vec, + pub(crate) post_aggregation_plan: LogicalPlan, +} + +multifield_partial_ord!( + StreamWindowAggregateNode, + base_agg_plan, + partition_keys, + post_aggregation_plan +); + +impl StreamWindowAggregateNode { + /// This node is only emitted after `KeyExtractionNode` in streaming rewrites; `partition_keys` + /// may be empty when GROUP BY is only a window call (window column stripped from key list), + /// but the pipeline still consumes a shuffle — use keyed aggregate parallelism. + fn parallelism_after_keyed_shuffle(&self, planner: &Planner) -> usize { + planner.keyed_aggregate_parallelism() + } + + /// Safely constructs a new node, computing the final projection without panicking. + pub fn try_new( + window_spec: WindowBehavior, + base_agg_plan: LogicalPlan, + partition_keys: Vec, + ) -> Result { + let post_aggregation_plan = + WindowBoundaryMath::build_post_aggregation(&base_agg_plan, window_spec.clone())?; + + Ok(Self { + window_spec, + base_agg_plan, + output_schema: post_aggregation_plan.schema().clone(), + partition_keys, + post_aggregation_plan, + }) + } + + fn build_tumbling_operator( + &self, + planner: &Planner, + node_id: usize, + input_schema: DFSchemaRef, + duration: Duration, + ) -> Result { + let binning_expr = planner.binning_function_proto(duration, input_schema.clone())?; + + let SplitPlanOutput { + partial_aggregation_plan, + partial_schema, + finish_plan, + } = planner.split_physical_plan(self.partition_keys.clone(), &self.base_agg_plan, true)?; + + let final_physical = planner.sync_plan(&self.post_aggregation_plan)?; + let final_physical_proto = PhysicalPlanNode::try_from_physical_plan( + final_physical, + &StreamingExtensionCodec::default(), + )?; + + let operator_config = TumblingWindowAggregateOperator { + name: proto_operator_name::TUMBLING_WINDOW.to_string(), + width_micros: duration.as_micros() as u64, + binning_function: binning_expr.encode_to_vec(), + input_schema: Some( + FsSchema::from_schema_keys( + Arc::new(input_schema.as_ref().into()), + self.partition_keys.clone(), + )? + .into(), + ), + partial_schema: Some(partial_schema.into()), + partial_aggregation_plan: partial_aggregation_plan.encode_to_vec(), + final_aggregation_plan: finish_plan.encode_to_vec(), + final_projection: Some(final_physical_proto.encode_to_vec()), + }; + + Ok(LogicalNode::single( + node_id as u32, + format!("tumbling_{node_id}"), + OperatorName::TumblingWindowAggregate, + operator_config.encode_to_vec(), + format!("TumblingWindow<{}>", operator_config.name), + self.parallelism_after_keyed_shuffle(planner), + )) + } + + fn build_sliding_operator( + &self, + planner: &Planner, + node_id: usize, + input_schema: DFSchemaRef, + duration: Duration, + slide_interval: Duration, + ) -> Result { + let binning_expr = planner.binning_function_proto(slide_interval, input_schema.clone())?; + + let SplitPlanOutput { + partial_aggregation_plan, + partial_schema, + finish_plan, + } = planner.split_physical_plan(self.partition_keys.clone(), &self.base_agg_plan, true)?; + + let final_physical = planner.sync_plan(&self.post_aggregation_plan)?; + let final_physical_proto = PhysicalPlanNode::try_from_physical_plan( + final_physical, + &StreamingExtensionCodec::default(), + )?; + + let operator_config = SlidingWindowAggregateOperator { + name: format!("SlidingWindow<{duration:?}>"), + width_micros: duration.as_micros() as u64, + slide_micros: slide_interval.as_micros() as u64, + binning_function: binning_expr.encode_to_vec(), + input_schema: Some( + FsSchema::from_schema_keys( + Arc::new(input_schema.as_ref().into()), + self.partition_keys.clone(), + )? + .into(), + ), + partial_schema: Some(partial_schema.into()), + partial_aggregation_plan: partial_aggregation_plan.encode_to_vec(), + final_aggregation_plan: finish_plan.encode_to_vec(), + final_projection: final_physical_proto.encode_to_vec(), + }; + + Ok(LogicalNode::single( + node_id as u32, + format!("sliding_window_{node_id}"), + OperatorName::SlidingWindowAggregate, + operator_config.encode_to_vec(), + proto_operator_name::SLIDING_WINDOW_LABEL.to_string(), + self.parallelism_after_keyed_shuffle(planner), + )) + } + + fn build_session_operator( + &self, + planner: &Planner, + node_id: usize, + input_schema: DFSchemaRef, + ) -> Result { + let WindowBehavior::FromOperator { + window: WindowType::Session { gap }, + window_index, + window_field, + is_nested: false, + } = &self.window_spec + else { + return plan_err!("Expected standard session window configuration"); + }; + + let output_fields = extract_qualified_fields(self.base_agg_plan.schema()); + let LogicalPlan::Aggregate(base_agg) = self.base_agg_plan.clone() else { + return plan_err!("Base plan must be an Aggregate node"); + }; + + let key_count = self.partition_keys.len(); + let unkeyed_schema = Arc::new(build_df_schema_with_metadata( + &output_fields[key_count..], + self.base_agg_plan.schema().metadata().clone(), + )?); + + let unkeyed_agg_node = Aggregate::try_new_with_schema( + base_agg.input.clone(), + vec![], + base_agg.aggr_expr.clone(), + unkeyed_schema, + )?; + + let physical_agg = planner.sync_plan(&LogicalPlan::Aggregate(unkeyed_agg_node))?; + let physical_agg_proto = PhysicalPlanNode::try_from_physical_plan( + physical_agg, + &StreamingExtensionCodec::default(), + )?; + + let operator_config = SessionWindowAggregateOperator { + name: format!("session_window_{node_id}"), + gap_micros: gap.as_micros() as u64, + window_field_name: window_field.name().to_string(), + window_index: *window_index as u64, + input_schema: Some( + FsSchema::from_schema_keys( + Arc::new(input_schema.as_ref().into()), + self.partition_keys.clone(), + )? + .into(), + ), + unkeyed_aggregate_schema: None, + partial_aggregation_plan: vec![], + final_aggregation_plan: physical_agg_proto.encode_to_vec(), + }; + + Ok(LogicalNode::single( + node_id as u32, + format!("SessionWindow<{gap:?}>"), + OperatorName::SessionWindowAggregate, + operator_config.encode_to_vec(), + operator_config.name.clone(), + self.parallelism_after_keyed_shuffle(planner), + )) + } + + fn build_instant_operator( + &self, + planner: &Planner, + node_id: usize, + input_schema: DFSchemaRef, + apply_final_projection: bool, + ) -> Result { + let ts_column_expr = Expr::Column(Column::new_unqualified(TIMESTAMP_FIELD.to_string())); + let binning_expr = planner.create_physical_expr(&ts_column_expr, &input_schema)?; + let binning_proto = + serialize_physical_expr(&binning_expr, &DefaultPhysicalExtensionCodec {})?; + + let final_projection_payload = if apply_final_projection { + let physical_plan = planner.sync_plan(&self.post_aggregation_plan)?; + let proto_node = PhysicalPlanNode::try_from_physical_plan( + physical_plan, + &StreamingExtensionCodec::default(), + )?; + Some(proto_node.encode_to_vec()) + } else { + None + }; + + let SplitPlanOutput { + partial_aggregation_plan, + partial_schema, + finish_plan, + } = planner.split_physical_plan(self.partition_keys.clone(), &self.base_agg_plan, true)?; + + let operator_config = TumblingWindowAggregateOperator { + name: proto_operator_name::INSTANT_WINDOW.to_string(), + width_micros: 0, + binning_function: binning_proto.encode_to_vec(), + input_schema: Some( + FsSchema::from_schema_keys( + Arc::new(input_schema.as_ref().into()), + self.partition_keys.clone(), + )? + .into(), + ), + partial_schema: Some(partial_schema.into()), + partial_aggregation_plan: partial_aggregation_plan.encode_to_vec(), + final_aggregation_plan: finish_plan.encode_to_vec(), + final_projection: final_projection_payload, + }; + + Ok(LogicalNode::single( + node_id as u32, + format!("instant_window_{node_id}"), + OperatorName::TumblingWindowAggregate, + operator_config.encode_to_vec(), + proto_operator_name::INSTANT_WINDOW_LABEL.to_string(), + self.parallelism_after_keyed_shuffle(planner), + )) + } +} + +impl StreamingOperatorBlueprint for StreamWindowAggregateNode { + fn operator_identity(&self) -> Option { + None + } + + fn compile_to_graph_node( + &self, + planner: &Planner, + node_id: usize, + mut input_schemas: Vec, + ) -> Result { + if input_schemas.len() != 1 { + return plan_err!("StreamWindowAggregateNode requires exactly one input schema"); + } + + let raw_schema = input_schemas.remove(0); + let df_schema = Arc::new(DFSchema::try_from(raw_schema.schema.as_ref().clone())?); + + let logical_operator = match &self.window_spec { + WindowBehavior::FromOperator { + window, is_nested, .. + } => { + if *is_nested { + self.build_instant_operator(planner, node_id, df_schema, true)? + } else { + match window { + WindowType::Tumbling { width } => { + self.build_tumbling_operator(planner, node_id, df_schema, *width)? + } + WindowType::Sliding { width, slide } => self + .build_sliding_operator(planner, node_id, df_schema, *width, *slide)?, + WindowType::Session { .. } => { + self.build_session_operator(planner, node_id, df_schema)? + } + WindowType::Instant => { + return plan_err!( + "Instant window is invalid within standard operator context" + ); + } + } + } + } + WindowBehavior::InData => self + .build_instant_operator(planner, node_id, df_schema, false) + .map_err(|e| e.context("Failed compiling instant window"))?, + }; + + let link = LogicalEdge::project_all(LogicalEdgeType::Shuffle, (*raw_schema).clone()); + Ok(CompiledTopologyNode { + execution_unit: logical_operator, + routing_edges: vec![link], + }) + } + + fn yielded_schema(&self) -> FsSchema { + let schema_ref = (*self.output_schema).clone().into(); + FsSchema::from_schema_unkeyed(Arc::new(schema_ref)) + .expect("StreamWindowAggregateNode output schema must contain timestamp column") + } +} + +impl UserDefinedLogicalNodeCore for StreamWindowAggregateNode { + fn name(&self) -> &str { + STREAM_AGG_EXTENSION_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.base_agg_plan] + } + + fn schema(&self) -> &DFSchemaRef { + &self.output_schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + let spec_desc = match &self.window_spec { + WindowBehavior::InData => "InData".to_string(), + WindowBehavior::FromOperator { window, .. } => format!("FromOperator({window:?})"), + }; + write!( + f, + "StreamWindowAggregate: {} | spec: {}", + self.schema(), + spec_desc + ) + } + + fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { + if inputs.len() != 1 { + return internal_err!("StreamWindowAggregateNode expects exactly 1 input"); + } + Self::try_new( + self.window_spec.clone(), + inputs[0].clone(), + self.partition_keys.clone(), + ) + } +} + +// ----------------------------------------------------------------------------- +// Dedicated boundary math for window bin / post-aggregation projection +// ----------------------------------------------------------------------------- + +struct WindowBoundaryMath; + +impl WindowBoundaryMath { + fn interval_nanos(nanos: i64) -> Expr { + Expr::Literal( + ScalarValue::IntervalMonthDayNano(Some(IntervalMonthDayNanoType::make_value( + 0, 0, nanos, + ))), + None, + ) + } + + fn build_post_aggregation( + agg_plan: &LogicalPlan, + window_spec: WindowBehavior, + ) -> Result { + let ts_field: QualifiedField = agg_plan + .inputs() + .first() + .ok_or_else(|| DataFusionError::Plan("Aggregate has no inputs".into()))? + .schema() + .qualified_field_with_unqualified_name(TIMESTAMP_FIELD)? + .into(); + + let plan_with_ts = LogicalPlan::Extension(Extension { + node: Arc::new(SystemTimestampInjectorNode::try_new( + agg_plan.clone(), + ts_field.qualifier().cloned(), + )?), + }); + + let (win_field, win_index, duration, is_nested) = match window_spec { + WindowBehavior::InData => return Ok(plan_with_ts), + WindowBehavior::FromOperator { + window, + window_field, + window_index, + is_nested, + } => match window { + WindowType::Tumbling { width } | WindowType::Sliding { width, .. } => { + (window_field, window_index, width, is_nested) + } + WindowType::Session { .. } => { + return Ok(LogicalPlan::Extension(Extension { + node: Arc::new(InjectWindowFieldNode::try_new( + plan_with_ts, + window_field, + window_index, + )?), + })); + } + WindowType::Instant => return Ok(plan_with_ts), + }, + }; + + if is_nested { + return Self::build_nested_projection(plan_with_ts, win_field, win_index, duration); + } + + let mut output_fields = extract_qualified_fields(agg_plan.schema()); + let mut projections: Vec<_> = output_fields + .iter() + .map(|f| Expr::Column(f.qualified_column())) + .collect(); + + let ts_col_expr = Expr::Column(Column::new(ts_field.qualifier().cloned(), ts_field.name())); + + output_fields.insert(win_index, win_field.clone()); + + let win_func_expr = Expr::ScalarFunction(ScalarFunction { + func: window(), + args: vec![ + ts_col_expr.clone(), + Expr::BinaryExpr(BinaryExpr { + left: Box::new(ts_col_expr.clone()), + op: logical_expr::Operator::Plus, + right: Box::new(Self::interval_nanos(duration.as_nanos() as i64)), + }), + ], + }); + + projections.insert( + win_index, + win_func_expr.alias_qualified(win_field.qualifier().cloned(), win_field.name()), + ); + + output_fields.push(ts_field); + + let bin_end_expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(ts_col_expr), + op: logical_expr::Operator::Plus, + right: Box::new(Self::interval_nanos((duration.as_nanos() - 1) as i64)), + }); + projections.push(bin_end_expr); + + Ok(LogicalPlan::Projection( + logical_expr::Projection::try_new_with_schema( + projections, + Arc::new(plan_with_ts), + Arc::new(build_df_schema(&output_fields)?), + )?, + )) + } + + fn build_nested_projection( + plan: LogicalPlan, + win_field: QualifiedField, + win_index: usize, + duration: Duration, + ) -> Result { + let ts_field: QualifiedField = plan + .schema() + .qualified_field_with_unqualified_name(TIMESTAMP_FIELD)? + .into(); + let ts_col_expr = Expr::Column(Column::new(ts_field.qualifier().cloned(), ts_field.name())); + + let mut output_fields = extract_qualified_fields(plan.schema()); + let mut projections: Vec<_> = output_fields + .iter() + .map(|f| Expr::Column(f.qualified_column())) + .collect(); + + output_fields.insert(win_index, win_field.clone()); + + let win_func_expr = Expr::ScalarFunction(ScalarFunction { + func: window(), + args: vec![ + Expr::BinaryExpr(BinaryExpr { + left: Box::new(ts_col_expr.clone()), + op: logical_expr::Operator::Minus, + right: Box::new(Self::interval_nanos(duration.as_nanos() as i64 - 1)), + }), + Expr::BinaryExpr(BinaryExpr { + left: Box::new(ts_col_expr), + op: logical_expr::Operator::Plus, + right: Box::new(Self::interval_nanos(1)), + }), + ], + }); + + projections.insert( + win_index, + win_func_expr.alias_qualified(win_field.qualifier().cloned(), win_field.name()), + ); + + Ok(LogicalPlan::Projection( + logical_expr::Projection::try_new_with_schema( + projections, + Arc::new(plan), + Arc::new(build_df_schema(&output_fields)?), + )?, + )) + } +} + +// ----------------------------------------------------------------------------- +// Field injection node (session window column placement) +// ----------------------------------------------------------------------------- + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct InjectWindowFieldNode { + pub(crate) upstream_plan: LogicalPlan, + pub(crate) target_field: QualifiedField, + pub(crate) insertion_index: usize, + pub(crate) new_schema: DFSchemaRef, +} + +multifield_partial_ord!(InjectWindowFieldNode, upstream_plan, insertion_index); + +impl InjectWindowFieldNode { + fn try_new( + upstream_plan: LogicalPlan, + target_field: QualifiedField, + insertion_index: usize, + ) -> Result { + let mut fields = extract_qualified_fields(upstream_plan.schema()); + fields.insert(insertion_index, target_field.clone()); + let meta = upstream_plan.schema().metadata().clone(); + + Ok(Self { + upstream_plan, + target_field, + insertion_index, + new_schema: Arc::new(build_df_schema_with_metadata(&fields, meta)?), + }) + } +} + +impl UserDefinedLogicalNodeCore for InjectWindowFieldNode { + fn name(&self) -> &str { + "InjectWindowFieldNode" + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.upstream_plan] + } + + fn schema(&self) -> &DFSchemaRef { + &self.new_schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!( + f, + "InjectWindowField: insert {:?} at offset {}", + self.target_field, self.insertion_index + ) + } + + fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { + if inputs.len() != 1 { + return internal_err!("InjectWindowFieldNode expects exactly 1 input"); + } + Self::try_new( + inputs[0].clone(), + self.target_field.clone(), + self.insertion_index, + ) + } +} diff --git a/src/streaming_planner/src/logical_node/async_udf.rs b/src/streaming_planner/src/logical_node/async_udf.rs new file mode 100644 index 00000000..3c70fc5e --- /dev/null +++ b/src/streaming_planner/src/logical_node/async_udf.rs @@ -0,0 +1,247 @@ +// 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::fmt::Formatter; +use std::sync::Arc; +use std::time::Duration; + +use datafusion::common::{DFSchemaRef, Result}; +use datafusion::logical_expr::{ + Expr, LogicalPlan, UserDefinedLogicalNode, UserDefinedLogicalNodeCore, +}; +use datafusion_common::{internal_err, plan_err}; +use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec; +use datafusion_proto::physical_plan::to_proto::serialize_physical_expr; +use prost::Message; +use protocol::function_stream_graph::{AsyncUdfOperator, AsyncUdfOrdering}; + +use crate::multifield_partial_ord; +use crate::common::constants::extension_node; +use crate::common::constants::sql_field; +use crate::common::{FsSchema, FsSchemaRef}; +use crate::logical_node::logical::{ + DylibUdfConfig, LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName, +}; +use crate::logical_node::streaming_operator_blueprint::{ + CompiledTopologyNode, StreamingOperatorBlueprint, +}; +use crate::logical_planner::planner::{NamedNode, Planner}; +use crate::types::{QualifiedField, build_df_schema, extract_qualified_fields}; + +pub const NODE_TYPE_NAME: &str = extension_node::ASYNC_FUNCTION_EXECUTION; + +/// Represents a logical node that executes an external asynchronous function (UDF) +/// and projects the final results into the streaming pipeline. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct AsyncFunctionExecutionNode { + pub(crate) upstream_plan: Arc, + pub(crate) operator_name: String, + pub(crate) function_config: DylibUdfConfig, + pub(crate) invocation_args: Vec, + pub(crate) result_projections: Vec, + pub(crate) preserve_ordering: bool, + pub(crate) concurrency_limit: usize, + pub(crate) execution_timeout: Duration, + pub(crate) resolved_schema: DFSchemaRef, +} + +multifield_partial_ord!( + AsyncFunctionExecutionNode, + upstream_plan, + operator_name, + function_config, + invocation_args, + result_projections, + preserve_ordering, + concurrency_limit, + execution_timeout +); + +impl AsyncFunctionExecutionNode { + /// Compiles logical expressions into serialized physical protobuf bytes. + fn compile_physical_expressions( + &self, + planner: &Planner, + expressions: &[Expr], + schema_context: &DFSchemaRef, + ) -> Result>> { + expressions + .iter() + .map(|logical_expr| { + let physical_expr = planner.create_physical_expr(logical_expr, schema_context)?; + let serialized = + serialize_physical_expr(&physical_expr, &DefaultPhysicalExtensionCodec {})?; + Ok(serialized.encode_to_vec()) + }) + .collect() + } + + /// Computes the intermediate schema which bridges the upstream output + /// and the raw asynchronous result injected by the UDF execution. + fn compute_intermediate_schema(&self) -> Result { + let mut fields = extract_qualified_fields(self.upstream_plan.schema()); + + let raw_result_field = QualifiedField::new( + None, + sql_field::ASYNC_RESULT, + self.function_config.return_type.clone(), + true, + ); + fields.push(raw_result_field); + + Ok(Arc::new(build_df_schema(&fields)?)) + } + + fn to_protobuf_config( + &self, + compiled_args: Vec>, + compiled_projections: Vec>, + ) -> AsyncUdfOperator { + let ordering_strategy = if self.preserve_ordering { + AsyncUdfOrdering::Ordered + } else { + AsyncUdfOrdering::Unordered + }; + + AsyncUdfOperator { + name: self.operator_name.clone(), + udf: Some(self.function_config.clone().into()), + arg_exprs: compiled_args, + final_exprs: compiled_projections, + ordering: ordering_strategy as i32, + max_concurrency: self.concurrency_limit as u32, + timeout_micros: self.execution_timeout.as_micros() as u64, + } + } +} + +impl StreamingOperatorBlueprint for AsyncFunctionExecutionNode { + fn operator_identity(&self) -> Option { + None + } + + fn compile_to_graph_node( + &self, + planner: &Planner, + node_index: usize, + mut input_schemas: Vec, + ) -> Result { + if input_schemas.len() != 1 { + return plan_err!("AsyncFunctionExecutionNode requires exactly one input schema"); + } + + let compiled_args = self.compile_physical_expressions( + planner, + &self.invocation_args, + self.upstream_plan.schema(), + )?; + + let intermediate_schema = self.compute_intermediate_schema()?; + let compiled_projections = self.compile_physical_expressions( + planner, + &self.result_projections, + &intermediate_schema, + )?; + + let operator_config = self.to_protobuf_config(compiled_args, compiled_projections); + + let logical_node = LogicalNode::single( + node_index as u32, + format!("async_udf_{node_index}"), + OperatorName::AsyncUdf, + operator_config.encode_to_vec(), + format!("AsyncUdf<{}>", self.operator_name), + planner.default_parallelism(), + ); + + let upstream_schema = input_schemas.remove(0); + let data_edge = + LogicalEdge::project_all(LogicalEdgeType::Forward, (*upstream_schema).clone()); + + Ok(CompiledTopologyNode { + execution_unit: logical_node, + routing_edges: vec![data_edge], + }) + } + + fn yielded_schema(&self) -> FsSchema { + let arrow_fields: Vec<_> = self + .resolved_schema + .fields() + .iter() + .map(|f| (**f).clone()) + .collect(); + + FsSchema::from_fields(arrow_fields) + } +} + +impl UserDefinedLogicalNodeCore for AsyncFunctionExecutionNode { + fn name(&self) -> &str { + NODE_TYPE_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.upstream_plan] + } + + fn schema(&self) -> &DFSchemaRef { + &self.resolved_schema + } + + fn expressions(&self) -> Vec { + self.invocation_args + .iter() + .chain(self.result_projections.iter()) + .cloned() + .collect() + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + write!( + f, + "AsyncFunctionExecution<{}>: Concurrency={}, Ordered={}", + self.operator_name, self.concurrency_limit, self.preserve_ordering + ) + } + + fn with_exprs_and_inputs( + &self, + exprs: Vec, + mut inputs: Vec, + ) -> Result { + if inputs.len() != 1 { + return internal_err!( + "AsyncFunctionExecutionNode expects exactly 1 input, but received {}", + inputs.len() + ); + } + + if UserDefinedLogicalNode::expressions(self) != exprs { + return internal_err!( + "Attempted to mutate async UDF expressions during logical planning, which is not supported." + ); + } + + Ok(Self { + upstream_plan: Arc::new(inputs.remove(0)), + operator_name: self.operator_name.clone(), + function_config: self.function_config.clone(), + invocation_args: self.invocation_args.clone(), + result_projections: self.result_projections.clone(), + preserve_ordering: self.preserve_ordering, + concurrency_limit: self.concurrency_limit, + execution_timeout: self.execution_timeout, + resolved_schema: self.resolved_schema.clone(), + }) + } +} diff --git a/src/streaming_planner/src/logical_node/debezium.rs b/src/streaming_planner/src/logical_node/debezium.rs new file mode 100644 index 00000000..9d558026 --- /dev/null +++ b/src/streaming_planner/src/logical_node/debezium.rs @@ -0,0 +1,393 @@ +// 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::sync::Arc; + +use arrow_schema::{DataType, Field, Schema}; +use datafusion::common::{ + DFSchema, DFSchemaRef, DataFusionError, Result, TableReference, internal_err, plan_err, +}; +use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion::physical_plan::DisplayAs; + +use crate::multifield_partial_ord; +use crate::common::constants::{cdc, extension_node}; +use crate::common::{FsSchema, FsSchemaRef, UPDATING_META_FIELD}; +use crate::logical_planner::planner::{NamedNode, Planner}; +use crate::physical::updating_meta_field; +use crate::types::TIMESTAMP_FIELD; + +use super::{CompiledTopologyNode, StreamingOperatorBlueprint}; + +// ----------------------------------------------------------------------------- +// Constants & Identifiers +// ----------------------------------------------------------------------------- + +pub const UNROLL_NODE_NAME: &str = extension_node::UNROLL_DEBEZIUM_PAYLOAD; +pub const PACK_NODE_NAME: &str = extension_node::PACK_DEBEZIUM_ENVELOPE; + +// ----------------------------------------------------------------------------- +// Core Schema Codec +// ----------------------------------------------------------------------------- + +/// Transforms between flat schemas and Debezium CDC envelopes. +pub struct DebeziumSchemaCodec; + +impl DebeziumSchemaCodec { + /// Wraps a flat physical schema into a Debezium CDC envelope structure. + pub fn wrap_into_envelope( + flat_schema: &DFSchemaRef, + qualifier_override: Option, + ) -> Result { + let ts_field = if flat_schema.has_column_with_unqualified_name(TIMESTAMP_FIELD) { + Some( + flat_schema + .field_with_unqualified_name(TIMESTAMP_FIELD)? + .clone(), + ) + } else { + None + }; + + let payload_fields: Vec<_> = flat_schema + .fields() + .iter() + .filter(|f| f.name() != TIMESTAMP_FIELD && f.name() != UPDATING_META_FIELD) + .cloned() + .collect(); + + let payload_struct_type = DataType::Struct(payload_fields.into()); + + let mut envelope_fields = vec![ + Arc::new(Field::new(cdc::BEFORE, payload_struct_type.clone(), true)), + Arc::new(Field::new(cdc::AFTER, payload_struct_type, true)), + Arc::new(Field::new(cdc::OP, DataType::Utf8, true)), + ]; + + if let Some(ts) = ts_field { + envelope_fields.push(Arc::new(ts)); + } + + let arrow_schema = Schema::new(envelope_fields); + let final_schema = match qualifier_override { + Some(qualifier) => DFSchema::try_from_qualified_schema(qualifier, &arrow_schema)?, + None => DFSchema::try_from(arrow_schema)?, + }; + + Ok(Arc::new(final_schema)) + } +} + +// ----------------------------------------------------------------------------- +// Logical Node: Unroll Debezium Payload +// ----------------------------------------------------------------------------- + +/// Decodes an incoming Debezium envelope into a flat, updating stream representation. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct UnrollDebeziumPayloadNode { + upstream_plan: LogicalPlan, + resolved_schema: DFSchemaRef, + pub pk_indices: Vec, + pk_names: Arc>, +} + +multifield_partial_ord!( + UnrollDebeziumPayloadNode, + upstream_plan, + pk_indices, + pk_names +); + +impl UnrollDebeziumPayloadNode { + pub fn try_new(upstream_plan: LogicalPlan, pk_names: Arc>) -> Result { + let input_schema = upstream_plan.schema(); + + let (before_idx, after_idx) = Self::validate_envelope_structure(input_schema)?; + + let payload_fields = Self::extract_payload_fields(input_schema, before_idx)?; + + let pk_indices = Self::map_primary_keys(payload_fields, &pk_names)?; + + let qualifier = Self::resolve_schema_qualifier(input_schema, before_idx, after_idx)?; + + let resolved_schema = + Self::compile_unrolled_schema(input_schema, payload_fields, qualifier)?; + + Ok(Self { + upstream_plan, + resolved_schema, + pk_indices, + pk_names, + }) + } + + fn validate_envelope_structure(schema: &DFSchemaRef) -> Result<(usize, usize)> { + let before_idx = schema + .index_of_column_by_name(None, cdc::BEFORE) + .ok_or_else(|| { + DataFusionError::Plan("Missing 'before' state column in CDC stream".into()) + })?; + + let after_idx = schema + .index_of_column_by_name(None, cdc::AFTER) + .ok_or_else(|| { + DataFusionError::Plan("Missing 'after' state column in CDC stream".into()) + })?; + + let op_idx = schema + .index_of_column_by_name(None, cdc::OP) + .ok_or_else(|| { + DataFusionError::Plan("Missing 'op' operation column in CDC stream".into()) + })?; + + let before_type = schema.field(before_idx).data_type(); + let after_type = schema.field(after_idx).data_type(); + + if before_type != after_type { + return plan_err!( + "State column type mismatch: 'before' is {before_type}, but 'after' is {after_type}" + ); + } + + if *schema.field(op_idx).data_type() != DataType::Utf8 { + return plan_err!("The '{}' column must be of type Utf8", cdc::OP); + } + + Ok((before_idx, after_idx)) + } + + fn extract_payload_fields( + schema: &DFSchemaRef, + state_idx: usize, + ) -> Result<&arrow_schema::Fields> { + match schema.field(state_idx).data_type() { + DataType::Struct(fields) => Ok(fields), + other => plan_err!("State columns must be of type Struct, found {other}"), + } + } + + fn map_primary_keys(fields: &arrow_schema::Fields, pk_names: &[String]) -> Result> { + pk_names + .iter() + .map(|pk| fields.find(pk).map(|(idx, _)| idx)) + .collect::>>() + .ok_or_else(|| { + DataFusionError::Plan("Specified primary key not found in payload schema".into()) + }) + } + + fn resolve_schema_qualifier( + schema: &DFSchemaRef, + before_idx: usize, + after_idx: usize, + ) -> Result> { + let before_qualifier = schema.qualified_field(before_idx).0; + let after_qualifier = schema.qualified_field(after_idx).0; + + match (before_qualifier, after_qualifier) { + (Some(bq), Some(aq)) if bq == aq => Ok(Some(bq.clone())), + (None, None) => Ok(None), + _ => plan_err!("'before' and 'after' columns must share the same namespace/qualifier"), + } + } + + fn compile_unrolled_schema( + original_schema: &DFSchemaRef, + payload_fields: &arrow_schema::Fields, + qualifier: Option, + ) -> Result { + let mut flat_fields = payload_fields.to_vec(); + + flat_fields.push(updating_meta_field()); + + let ts_idx = original_schema + .index_of_column_by_name(None, TIMESTAMP_FIELD) + .ok_or_else(|| { + DataFusionError::Plan(format!( + "Required event time field '{TIMESTAMP_FIELD}' is missing" + )) + })?; + + flat_fields.push(Arc::new(original_schema.field(ts_idx).clone())); + + let arrow_schema = Schema::new(flat_fields); + let compiled_schema = match qualifier { + Some(q) => DFSchema::try_from_qualified_schema(q, &arrow_schema)?, + None => DFSchema::try_from(arrow_schema)?, + }; + + Ok(Arc::new(compiled_schema)) + } +} + +impl UserDefinedLogicalNodeCore for UnrollDebeziumPayloadNode { + fn name(&self) -> &str { + UNROLL_NODE_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.upstream_plan] + } + + fn schema(&self) -> &DFSchemaRef { + &self.resolved_schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "UnrollDebeziumPayload") + } + + fn with_exprs_and_inputs( + &self, + _exprs: Vec, + mut inputs: Vec, + ) -> Result { + if inputs.len() != 1 { + return internal_err!( + "UnrollDebeziumPayloadNode expects exactly 1 input, got {}", + inputs.len() + ); + } + Self::try_new(inputs.remove(0), self.pk_names.clone()) + } +} + +impl StreamingOperatorBlueprint for UnrollDebeziumPayloadNode { + fn operator_identity(&self) -> Option { + None + } + + fn is_passthrough_boundary(&self) -> bool { + true + } + + fn compile_to_graph_node( + &self, + _: &Planner, + _: usize, + _: Vec, + ) -> Result { + plan_err!( + "UnrollDebeziumPayloadNode is a logical boundary and should not be physically planned" + ) + } + + fn yielded_schema(&self) -> FsSchema { + FsSchema::from_schema_unkeyed(Arc::new(self.resolved_schema.as_ref().into())) + .unwrap_or_else(|_| { + panic!("Failed to extract physical schema for {}", UNROLL_NODE_NAME) + }) + } +} + +// ----------------------------------------------------------------------------- +// Logical Node: Pack Debezium Envelope +// ----------------------------------------------------------------------------- + +/// Encodes a flat updating stream back into a Debezium CDC envelope representation. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct PackDebeziumEnvelopeNode { + upstream_plan: Arc, + envelope_schema: DFSchemaRef, +} + +multifield_partial_ord!(PackDebeziumEnvelopeNode, upstream_plan); + +impl PackDebeziumEnvelopeNode { + pub fn try_new(upstream_plan: LogicalPlan) -> Result { + let envelope_schema = DebeziumSchemaCodec::wrap_into_envelope(upstream_plan.schema(), None) + .map_err(|e| { + DataFusionError::Plan(format!("Failed to compile Debezium envelope schema: {e}")) + })?; + + Ok(Self { + upstream_plan: Arc::new(upstream_plan), + envelope_schema, + }) + } +} + +impl DisplayAs for PackDebeziumEnvelopeNode { + fn fmt_as( + &self, + _t: datafusion::physical_plan::DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "PackDebeziumEnvelope") + } +} + +impl UserDefinedLogicalNodeCore for PackDebeziumEnvelopeNode { + fn name(&self) -> &str { + PACK_NODE_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.upstream_plan] + } + + fn schema(&self) -> &DFSchemaRef { + &self.envelope_schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "PackDebeziumEnvelope") + } + + fn with_exprs_and_inputs( + &self, + _exprs: Vec, + mut inputs: Vec, + ) -> Result { + if inputs.len() != 1 { + return internal_err!( + "PackDebeziumEnvelopeNode expects exactly 1 input, got {}", + inputs.len() + ); + } + Self::try_new(inputs.remove(0)) + } +} + +impl StreamingOperatorBlueprint for PackDebeziumEnvelopeNode { + fn operator_identity(&self) -> Option { + None + } + + fn is_passthrough_boundary(&self) -> bool { + true + } + + fn compile_to_graph_node( + &self, + _: &Planner, + _: usize, + _: Vec, + ) -> Result { + internal_err!( + "PackDebeziumEnvelopeNode is a logical boundary and should not be physically planned" + ) + } + + fn yielded_schema(&self) -> FsSchema { + FsSchema::from_schema_unkeyed(Arc::new(self.envelope_schema.as_ref().into())) + .unwrap_or_else(|_| panic!("Failed to extract physical schema for {}", PACK_NODE_NAME)) + } +} diff --git a/src/streaming_planner/src/logical_node/extension_try_from.rs b/src/streaming_planner/src/logical_node/extension_try_from.rs new file mode 100644 index 00000000..c13532dc --- /dev/null +++ b/src/streaming_planner/src/logical_node/extension_try_from.rs @@ -0,0 +1,70 @@ +// 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::sync::Arc; + +use datafusion::common::{DataFusionError, Result}; +use datafusion::logical_expr::UserDefinedLogicalNode; + +use crate::logical_node::aggregate::StreamWindowAggregateNode; +use crate::logical_node::async_udf::AsyncFunctionExecutionNode; +use crate::logical_node::debezium::{PackDebeziumEnvelopeNode, UnrollDebeziumPayloadNode}; +use crate::logical_node::join::StreamingJoinNode; +use crate::logical_node::key_calculation::KeyExtractionNode; +use crate::logical_node::lookup::StreamReferenceJoinNode; +use crate::logical_node::projection::StreamProjectionNode; +use crate::logical_node::remote_table::RemoteTableBoundaryNode; +use crate::logical_node::sink::StreamEgressNode; +use crate::logical_node::streaming_operator_blueprint::StreamingOperatorBlueprint; +use crate::logical_node::table_source::StreamIngestionNode; +use crate::logical_node::updating_aggregate::ContinuousAggregateNode; +use crate::logical_node::watermark_node::EventTimeWatermarkNode; +use crate::logical_node::windows_function::StreamingWindowFunctionNode; + +fn try_from_t( + node: &dyn UserDefinedLogicalNode, +) -> std::result::Result<&dyn StreamingOperatorBlueprint, ()> { + node.as_any() + .downcast_ref::() + .map(|t| t as &dyn StreamingOperatorBlueprint) + .ok_or(()) +} + +impl<'a> TryFrom<&'a dyn UserDefinedLogicalNode> for &'a dyn StreamingOperatorBlueprint { + type Error = DataFusionError; + + fn try_from(node: &'a dyn UserDefinedLogicalNode) -> Result { + try_from_t::(node) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .map_err(|_| DataFusionError::Plan(format!("unexpected node: {}", node.name()))) + } +} + +impl<'a> TryFrom<&'a Arc> for &'a dyn StreamingOperatorBlueprint { + type Error = DataFusionError; + + fn try_from(node: &'a Arc) -> Result { + TryFrom::try_from(node.as_ref()) + } +} diff --git a/src/streaming_planner/src/logical_node/is_retract.rs b/src/streaming_planner/src/logical_node/is_retract.rs new file mode 100644 index 00000000..496edee6 --- /dev/null +++ b/src/streaming_planner/src/logical_node/is_retract.rs @@ -0,0 +1,82 @@ +// 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::sync::Arc; + +use datafusion::arrow::datatypes::{DataType, TimeUnit}; +use datafusion::common::{DFSchemaRef, Result, TableReference}; +use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; + +use crate::multifield_partial_ord; +use crate::physical::updating_meta_field; +use crate::types::{ + QualifiedField, TIMESTAMP_FIELD, build_df_schema, extract_qualified_fields, +}; + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct IsRetractExtension { + pub(crate) input: LogicalPlan, + pub(crate) schema: DFSchemaRef, + pub(crate) timestamp_qualifier: Option, +} + +multifield_partial_ord!(IsRetractExtension, input, timestamp_qualifier); + +impl IsRetractExtension { + pub fn new(input: LogicalPlan, timestamp_qualifier: Option) -> Self { + let mut output_fields = extract_qualified_fields(input.schema()); + + let timestamp_index = output_fields.len() - 1; + output_fields[timestamp_index] = QualifiedField::new( + timestamp_qualifier.clone(), + TIMESTAMP_FIELD, + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + ); + output_fields.push((timestamp_qualifier.clone(), updating_meta_field()).into()); + let schema = Arc::new(build_df_schema(&output_fields).unwrap()); + Self { + input, + schema, + timestamp_qualifier, + } + } +} + +impl UserDefinedLogicalNodeCore for IsRetractExtension { + fn name(&self) -> &str { + "IsRetractExtension" + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.input] + } + + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "IsRetractExtension") + } + + fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { + Ok(Self::new( + inputs[0].clone(), + self.timestamp_qualifier.clone(), + )) + } +} diff --git a/src/streaming_planner/src/logical_node/join.rs b/src/streaming_planner/src/logical_node/join.rs new file mode 100644 index 00000000..35d645f1 --- /dev/null +++ b/src/streaming_planner/src/logical_node/join.rs @@ -0,0 +1,211 @@ +// 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::fmt::Formatter; +use std::time::Duration; + +use datafusion::common::{DFSchemaRef, Result}; +use datafusion::logical_expr::expr::Expr; +use datafusion::logical_expr::{LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion_common::plan_err; +use datafusion_proto::physical_plan::AsExecutionPlan; +use datafusion_proto::protobuf::PhysicalPlanNode; +use prost::Message; +use protocol::function_stream_graph::JoinOperator; + +use crate::common::constants::{extension_node, runtime_operator_kind}; +use crate::common::{FsSchema, FsSchemaRef}; +use crate::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; +use crate::logical_node::{CompiledTopologyNode, StreamingOperatorBlueprint}; +use crate::logical_planner::planner::{NamedNode, Planner}; +use crate::physical::StreamingExtensionCodec; + +// ----------------------------------------------------------------------------- +// Constants +// ----------------------------------------------------------------------------- + +pub const STREAM_JOIN_NODE_TYPE: &str = extension_node::STREAMING_JOIN; + +// ----------------------------------------------------------------------------- +// Logical Node Definition +// ----------------------------------------------------------------------------- + +/// A logical plan node representing a streaming join operation. +/// It bridges the DataFusion logical plan with the physical streaming execution engine. +#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd)] +pub struct StreamingJoinNode { + pub(crate) underlying_plan: LogicalPlan, + pub(crate) instant_execution_mode: bool, + pub(crate) state_retention_ttl: Option, +} + +impl StreamingJoinNode { + /// Creates a new instance of the streaming join node. + pub fn new( + underlying_plan: LogicalPlan, + instant_execution_mode: bool, + state_retention_ttl: Option, + ) -> Self { + Self { + underlying_plan, + instant_execution_mode, + state_retention_ttl, + } + } + + /// Compiles the physical execution plan and serializes it into a Protobuf configuration payload. + fn compile_operator_config( + &self, + planner: &Planner, + node_identifier: &str, + left_schema: FsSchemaRef, + right_schema: FsSchemaRef, + ) -> Result { + let physical_plan = planner.sync_plan(&self.underlying_plan)?; + + let proto_node = PhysicalPlanNode::try_from_physical_plan( + physical_plan, + &StreamingExtensionCodec::default(), + )?; + + Ok(JoinOperator { + name: node_identifier.to_string(), + left_schema: Some(left_schema.as_ref().clone().into()), + right_schema: Some(right_schema.as_ref().clone().into()), + output_schema: Some(self.extract_fs_schema().into()), + join_plan: proto_node.encode_to_vec(), + ttl_micros: self.state_retention_ttl.map(|ttl| ttl.as_micros() as u64), + }) + } + + fn determine_operator_type(&self) -> OperatorName { + if self.instant_execution_mode { + OperatorName::InstantJoin + } else { + OperatorName::Join + } + } + + fn extract_fs_schema(&self) -> FsSchema { + FsSchema::from_schema_unkeyed(self.underlying_plan.schema().inner().clone()) + .expect("Fatal: Failed to convert internal join schema to FsSchema without keys") + } +} + +// ----------------------------------------------------------------------------- +// DataFusion Logical Node Core Implementation +// ----------------------------------------------------------------------------- + +impl UserDefinedLogicalNodeCore for StreamingJoinNode { + fn name(&self) -> &str { + STREAM_JOIN_NODE_TYPE + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.underlying_plan] + } + + fn schema(&self) -> &DFSchemaRef { + self.underlying_plan.schema() + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + write!( + f, + "StreamingJoinNode: Schema={}, InstantMode={}, TTL={:?}", + self.schema(), + self.instant_execution_mode, + self.state_retention_ttl + ) + } + + fn with_exprs_and_inputs( + &self, + _exprs: Vec, + mut inputs: Vec, + ) -> Result { + if inputs.len() != 1 { + return plan_err!( + "StreamingJoinNode expects exactly 1 underlying logical plan during recreation" + ); + } + + Ok(Self::new( + inputs.remove(0), + self.instant_execution_mode, + self.state_retention_ttl, + )) + } +} + +// ----------------------------------------------------------------------------- +// Streaming Graph Extension Implementation +// ----------------------------------------------------------------------------- + +impl StreamingOperatorBlueprint for StreamingJoinNode { + fn operator_identity(&self) -> Option { + None + } + + fn compile_to_graph_node( + &self, + planner: &Planner, + node_index: usize, + mut input_schemas: Vec, + ) -> Result { + if input_schemas.len() != 2 { + return plan_err!( + "Invalid topology: StreamingJoinNode requires exactly two upstream inputs, received {}", + input_schemas.len() + ); + } + + let right_schema = input_schemas.pop().unwrap(); + let left_schema = input_schemas.pop().unwrap(); + + let node_identifier = format!("stream_join_{node_index}"); + + let operator_config = self.compile_operator_config( + planner, + &node_identifier, + left_schema.clone(), + right_schema.clone(), + )?; + + let logical_node = LogicalNode::single( + node_index as u32, + node_identifier.clone(), + self.determine_operator_type(), + operator_config.encode_to_vec(), + runtime_operator_kind::STREAMING_JOIN.to_string(), + planner.default_parallelism(), + ); + + let left_edge = + LogicalEdge::project_all(LogicalEdgeType::LeftJoin, left_schema.as_ref().clone()); + let right_edge = + LogicalEdge::project_all(LogicalEdgeType::RightJoin, right_schema.as_ref().clone()); + + Ok(CompiledTopologyNode { + execution_unit: logical_node, + routing_edges: vec![left_edge, right_edge], + }) + } + + fn yielded_schema(&self) -> FsSchema { + self.extract_fs_schema() + } +} diff --git a/src/streaming_planner/src/logical_node/key_calculation.rs b/src/streaming_planner/src/logical_node/key_calculation.rs new file mode 100644 index 00000000..b27ed1ea --- /dev/null +++ b/src/streaming_planner/src/logical_node/key_calculation.rs @@ -0,0 +1,309 @@ +// 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::fmt::Formatter; +use std::sync::Arc; + +use datafusion::arrow::datatypes::{Field, Schema}; +use datafusion::common::{DFSchemaRef, Result, internal_err, plan_err}; +use datafusion::logical_expr::{Expr, ExprSchemable, LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion_common::DFSchema; +use datafusion_expr::col; +use datafusion_proto::physical_plan::to_proto::serialize_physical_expr; +use datafusion_proto::physical_plan::{AsExecutionPlan, DefaultPhysicalExtensionCodec}; +use datafusion_proto::protobuf::PhysicalPlanNode; +use itertools::Itertools; +use prost::Message; + +use protocol::function_stream_graph::{KeyPlanOperator, ProjectionOperator}; + +use crate::multifield_partial_ord; +use crate::common::constants::{extension_node, sql_field}; +use crate::common::{FsSchema, FsSchemaRef}; +use crate::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; +use crate::logical_node::{CompiledTopologyNode, StreamingOperatorBlueprint}; +use crate::logical_planner::planner::{NamedNode, Planner}; +use crate::physical::StreamingExtensionCodec; +use crate::types::{build_df_schema_with_metadata, extract_qualified_fields}; + +pub const EXTENSION_NODE_IDENTIFIER: &str = extension_node::KEY_EXTRACTION; + +/// Routing strategy for shuffling data across the stream topology. +#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd)] +pub enum KeyExtractionStrategy { + ColumnIndices(Vec), + CalculatedExpressions(Vec), +} + +/// Logical node that computes or extracts routing keys. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct KeyExtractionNode { + pub(crate) operator_label: Option, + pub(crate) upstream_plan: LogicalPlan, + pub(crate) extraction_strategy: KeyExtractionStrategy, + pub(crate) resolved_schema: DFSchemaRef, +} + +multifield_partial_ord!( + KeyExtractionNode, + operator_label, + upstream_plan, + extraction_strategy +); + +impl KeyExtractionNode { + /// Extracts keys and hides them from the downstream projection. + pub fn try_new_with_projection( + upstream_plan: LogicalPlan, + target_indices: Vec, + label: String, + ) -> Result { + let projected_fields: Vec<_> = extract_qualified_fields(upstream_plan.schema()) + .into_iter() + .enumerate() + .filter(|(idx, _)| !target_indices.contains(idx)) + .map(|(_, field)| field) + .collect(); + + let metadata = upstream_plan.schema().metadata().clone(); + let resolved_schema = build_df_schema_with_metadata(&projected_fields, metadata)?; + + Ok(Self { + operator_label: Some(label), + upstream_plan, + extraction_strategy: KeyExtractionStrategy::ColumnIndices(target_indices), + resolved_schema: Arc::new(resolved_schema), + }) + } + + /// Creates a node using an explicit strategy without changing the visible schema. + pub fn new(upstream_plan: LogicalPlan, strategy: KeyExtractionStrategy) -> Self { + let resolved_schema = upstream_plan.schema().clone(); + Self { + operator_label: None, + upstream_plan, + extraction_strategy: strategy, + resolved_schema, + } + } + + fn compile_index_router( + &self, + physical_plan_proto: PhysicalPlanNode, + indices: &[usize], + ) -> (Vec, OperatorName) { + let operator_config = KeyPlanOperator { + name: sql_field::DEFAULT_KEY_LABEL.into(), + physical_plan: physical_plan_proto.encode_to_vec(), + key_fields: indices.iter().map(|&idx| idx as u64).collect(), + }; + + (operator_config.encode_to_vec(), OperatorName::KeyBy) + } + + fn compile_expression_router( + &self, + planner: &Planner, + expressions: &[Expr], + input_schema_ref: &FsSchemaRef, + input_df_schema: &DFSchemaRef, + ) -> Result<(Vec, OperatorName)> { + let mut target_exprs = expressions.to_vec(); + + for field in input_schema_ref.schema.fields.iter() { + target_exprs.push(col(field.name())); + } + + let output_fs_schema = self.generate_fs_schema()?; + + for (compiled_expr, expected_field) in + target_exprs.iter().zip(output_fs_schema.schema.fields()) + { + let (expr_type, expr_nullable) = + compiled_expr.data_type_and_nullable(input_df_schema)?; + if expr_type != *expected_field.data_type() + || expr_nullable != expected_field.is_nullable() + { + return plan_err!( + "Type mismatch in key calculation: Expected {} (nullable: {}), got {} (nullable: {})", + expected_field.data_type(), + expected_field.is_nullable(), + expr_type, + expr_nullable + ); + } + } + + let mut physical_expr_payloads = Vec::with_capacity(target_exprs.len()); + for logical_expr in target_exprs { + let physical_expr = planner + .create_physical_expr(&logical_expr, input_df_schema) + .map_err(|e| e.context("Failed to physicalize PARTITION BY expression"))?; + + let serialized_expr = + serialize_physical_expr(&physical_expr, &DefaultPhysicalExtensionCodec {})?; + physical_expr_payloads.push(serialized_expr.encode_to_vec()); + } + + let operator_config = ProjectionOperator { + name: self + .operator_label + .as_deref() + .unwrap_or(sql_field::DEFAULT_KEY_LABEL) + .to_string(), + input_schema: Some(input_schema_ref.as_ref().clone().into()), + output_schema: Some(output_fs_schema.into()), + exprs: physical_expr_payloads, + }; + + Ok((operator_config.encode_to_vec(), OperatorName::Projection)) + } + + fn generate_fs_schema(&self) -> Result { + let base_arrow_schema = self.upstream_plan.schema().as_ref(); + + match &self.extraction_strategy { + KeyExtractionStrategy::ColumnIndices(indices) => { + FsSchema::from_schema_keys(Arc::new(base_arrow_schema.into()), indices.clone()) + } + KeyExtractionStrategy::CalculatedExpressions(expressions) => { + let mut composite_fields = + Vec::with_capacity(expressions.len() + base_arrow_schema.fields().len()); + + for (idx, expr) in expressions.iter().enumerate() { + let (data_type, nullable) = expr.data_type_and_nullable(base_arrow_schema)?; + composite_fields + .push(Field::new(format!("__key_{idx}"), data_type, nullable).into()); + } + + for field in base_arrow_schema.fields().iter() { + composite_fields.push(field.clone()); + } + + let final_schema = Arc::new(Schema::new(composite_fields)); + let key_mapping = (1..=expressions.len()).collect_vec(); + FsSchema::from_schema_keys(final_schema, key_mapping) + } + } + } +} + +impl StreamingOperatorBlueprint for KeyExtractionNode { + fn operator_identity(&self) -> Option { + None + } + + fn compile_to_graph_node( + &self, + planner: &Planner, + node_index: usize, + mut input_schemas: Vec, + ) -> Result { + if input_schemas.len() != 1 { + return plan_err!("KeyExtractionNode requires exactly one upstream input schema"); + } + + let input_schema_ref = input_schemas.remove(0); + let input_df_schema = Arc::new(DFSchema::try_from( + input_schema_ref.schema.as_ref().clone(), + )?); + + let physical_plan = planner.sync_plan(&self.upstream_plan)?; + let physical_plan_proto = PhysicalPlanNode::try_from_physical_plan( + physical_plan, + &StreamingExtensionCodec::default(), + )?; + + let (protobuf_payload, engine_operator_name) = match &self.extraction_strategy { + KeyExtractionStrategy::ColumnIndices(indices) => { + self.compile_index_router(physical_plan_proto, indices) + } + KeyExtractionStrategy::CalculatedExpressions(exprs) => { + self.compile_expression_router(planner, exprs, &input_schema_ref, &input_df_schema)? + } + }; + + let logical_node = LogicalNode::single( + node_index as u32, + format!("key_{node_index}"), + engine_operator_name, + protobuf_payload, + format!("Key<{}>", self.operator_label.as_deref().unwrap_or("_")), + planner.key_by_parallelism(), + ); + + let data_edge = + LogicalEdge::project_all(LogicalEdgeType::Forward, (*input_schema_ref).clone()); + + Ok(CompiledTopologyNode { + execution_unit: logical_node, + routing_edges: vec![data_edge], + }) + } + + fn yielded_schema(&self) -> FsSchema { + self.generate_fs_schema() + .expect("Fatal: Failed to generate output schema for KeyExtractionNode") + } +} + +impl UserDefinedLogicalNodeCore for KeyExtractionNode { + fn name(&self) -> &str { + EXTENSION_NODE_IDENTIFIER + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.upstream_plan] + } + + fn schema(&self) -> &DFSchemaRef { + &self.resolved_schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + write!( + f, + "KeyExtractionNode: Strategy={:?} | Schema={}", + self.extraction_strategy, self.resolved_schema + ) + } + + fn with_exprs_and_inputs( + &self, + exprs: Vec, + mut inputs: Vec, + ) -> Result { + if inputs.len() != 1 { + return internal_err!("KeyExtractionNode requires exactly 1 input logical plan"); + } + + let strategy = match &self.extraction_strategy { + KeyExtractionStrategy::ColumnIndices(indices) => { + KeyExtractionStrategy::ColumnIndices(indices.clone()) + } + KeyExtractionStrategy::CalculatedExpressions(_) => { + KeyExtractionStrategy::CalculatedExpressions(exprs) + } + }; + + Ok(Self { + operator_label: self.operator_label.clone(), + upstream_plan: inputs.remove(0), + extraction_strategy: strategy, + resolved_schema: self.resolved_schema.clone(), + }) + } +} diff --git a/src/streaming_planner/src/logical_node/logical/dylib_udf_config.rs b/src/streaming_planner/src/logical_node/logical/dylib_udf_config.rs new file mode 100644 index 00000000..9bf3368f --- /dev/null +++ b/src/streaming_planner/src/logical_node/logical/dylib_udf_config.rs @@ -0,0 +1,71 @@ +// 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 datafusion::arrow::datatypes::DataType; +use datafusion_proto::protobuf::ArrowType; +use prost::Message; +use protocol::function_stream_graph; + +#[derive(Clone, Debug, Eq, PartialEq, Hash, PartialOrd)] +pub struct DylibUdfConfig { + pub dylib_path: String, + pub arg_types: Vec, + pub return_type: DataType, + pub aggregate: bool, + pub is_async: bool, +} + +impl From for function_stream_graph::DylibUdfConfig { + fn from(from: DylibUdfConfig) -> Self { + function_stream_graph::DylibUdfConfig { + dylib_path: from.dylib_path, + arg_types: from + .arg_types + .iter() + .map(|t| { + ArrowType::try_from(t) + .expect("unsupported data type") + .encode_to_vec() + }) + .collect(), + return_type: ArrowType::try_from(&from.return_type) + .expect("unsupported data type") + .encode_to_vec(), + aggregate: from.aggregate, + is_async: from.is_async, + } + } +} + +impl From for DylibUdfConfig { + fn from(from: function_stream_graph::DylibUdfConfig) -> Self { + DylibUdfConfig { + dylib_path: from.dylib_path, + arg_types: from + .arg_types + .iter() + .map(|t| { + DataType::try_from( + &ArrowType::decode(&mut t.as_slice()).expect("invalid arrow type"), + ) + .expect("invalid arrow type") + }) + .collect(), + return_type: DataType::try_from( + &ArrowType::decode(&mut from.return_type.as_slice()).unwrap(), + ) + .expect("invalid arrow type"), + aggregate: from.aggregate, + is_async: from.is_async, + } + } +} diff --git a/src/streaming_planner/src/logical_node/logical/fs_program_convert.rs b/src/streaming_planner/src/logical_node/logical/fs_program_convert.rs new file mode 100644 index 00000000..53316bae --- /dev/null +++ b/src/streaming_planner/src/logical_node/logical/fs_program_convert.rs @@ -0,0 +1,200 @@ +// 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::str::FromStr; +use std::sync::Arc; + +use datafusion::common::{DataFusionError, Result as DFResult}; +use petgraph::graph::DiGraph; +use petgraph::prelude::EdgeRef; +use protocol::function_stream_graph::{ + ChainedOperator, EdgeType as ProtoEdgeType, FsEdge, FsNode, FsProgram, + FsSchema as ProtoFsSchema, +}; + +use crate::api::pipelines::{PipelineEdge, PipelineGraph, PipelineNode}; +use crate::common::FsSchema; + +use super::logical_edge::logical_edge_type_from_proto_i32; +use super::operator_chain::{ChainedLogicalOperator, OperatorChain}; +use super::operator_name::OperatorName; +use super::{LogicalEdge, LogicalNode, LogicalProgram, ProgramConfig}; + +impl TryFrom for LogicalProgram { + type Error = DataFusionError; + + fn try_from(value: FsProgram) -> DFResult { + let mut graph = DiGraph::new(); + let mut id_map = HashMap::with_capacity(value.nodes.len()); + + for node in value.nodes { + let operators = node + .operators + .into_iter() + .map(|op| { + let ChainedOperator { + operator_id, + operator_name: name_str, + operator_config, + } = op; + let operator_name = OperatorName::from_str(&name_str).map_err(|_| { + DataFusionError::Plan(format!("Invalid operator name: {name_str}")) + })?; + Ok(ChainedLogicalOperator { + operator_id, + operator_name, + operator_config, + }) + }) + .collect::>>()?; + + let edges = node + .edges + .into_iter() + .map(|e| { + let fs: FsSchema = e.try_into()?; + Ok(Arc::new(fs)) + }) + .collect::>>()?; + + let logical_node = LogicalNode { + node_id: node.node_id, + description: node.description, + operator_chain: OperatorChain { operators, edges }, + parallelism: node.parallelism as usize, + }; + + id_map.insert(node.node_index, graph.add_node(logical_node)); + } + + for edge in value.edges { + let source = *id_map.get(&edge.source).ok_or_else(|| { + DataFusionError::Plan("Graph integrity error: Missing source node".into()) + })?; + let target = *id_map.get(&edge.target).ok_or_else(|| { + DataFusionError::Plan("Graph integrity error: Missing target node".into()) + })?; + let schema = edge.schema.ok_or_else(|| { + DataFusionError::Plan("Graph integrity error: Missing edge schema".into()) + })?; + let edge_type = logical_edge_type_from_proto_i32(edge.edge_type)?; + + graph.add_edge( + source, + target, + LogicalEdge { + edge_type, + schema: Arc::new(FsSchema::try_from(schema)?), + }, + ); + } + + let program_config = value + .program_config + .map(ProgramConfig::from) + .unwrap_or_default(); + + Ok(LogicalProgram::new(graph, program_config)) + } +} + +impl From for FsProgram { + fn from(value: LogicalProgram) -> Self { + let nodes = value + .graph + .node_indices() + .filter_map(|idx| value.graph.node_weight(idx).map(|node| (idx, node))) + .map(|(idx, node)| FsNode { + node_index: idx.index() as i32, + node_id: node.node_id, + parallelism: node.parallelism as u32, + description: node.description.clone(), + operators: node + .operator_chain + .operators + .iter() + .map(|op| ChainedOperator { + operator_id: op.operator_id.clone(), + operator_name: op.operator_name.to_string(), + operator_config: op.operator_config.clone(), + }) + .collect(), + edges: node + .operator_chain + .edges + .iter() + .map(|edge| ProtoFsSchema::from((**edge).clone())) + .collect(), + }) + .collect(); + + let edges = value + .graph + .edge_indices() + .filter_map(|eidx| { + let edge = value.graph.edge_weight(eidx)?; + let (source, target) = value.graph.edge_endpoints(eidx)?; + Some(FsEdge { + source: source.index() as i32, + target: target.index() as i32, + schema: Some(ProtoFsSchema::from((*edge.schema).clone())), + edge_type: ProtoEdgeType::from(edge.edge_type) as i32, + }) + }) + .collect(); + + FsProgram { + nodes, + edges, + program_config: Some(value.program_config.into()), + } + } +} + +impl TryFrom for PipelineGraph { + type Error = DataFusionError; + + fn try_from(value: LogicalProgram) -> DFResult { + let nodes = value + .graph + .node_weights() + .map(|node| { + Ok(PipelineNode { + node_id: node.node_id, + operator: node.resolve_pipeline_operator_name()?, + description: node.description.clone(), + parallelism: node.parallelism as u32, + }) + }) + .collect::>>()?; + + let edges = value + .graph + .edge_references() + .filter_map(|edge| { + let src = value.graph.node_weight(edge.source())?; + let target = value.graph.node_weight(edge.target())?; + Some(PipelineEdge { + src_id: src.node_id, + dest_id: target.node_id, + key_type: "()".to_string(), + value_type: "()".to_string(), + edge_type: format!("{:?}", edge.weight().edge_type), + }) + }) + .collect(); + + Ok(PipelineGraph { nodes, edges }) + } +} diff --git a/src/streaming_planner/src/logical_node/logical/logical_edge.rs b/src/streaming_planner/src/logical_node/logical/logical_edge.rs new file mode 100644 index 00000000..d599f7f1 --- /dev/null +++ b/src/streaming_planner/src/logical_node/logical/logical_edge.rs @@ -0,0 +1,102 @@ +// 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::fmt::{Display, Formatter}; +use std::sync::Arc; + +use datafusion::common::{DataFusionError, Result}; +use protocol::function_stream_graph::EdgeType as ProtoEdgeType; +use serde::{Deserialize, Serialize}; + +use crate::common::FsSchema; + +#[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Ord, Serialize, Deserialize)] +pub enum LogicalEdgeType { + Forward, + Shuffle, + LeftJoin, + RightJoin, +} + +impl Display for LogicalEdgeType { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + let symbol = match self { + LogicalEdgeType::Forward => "→", + LogicalEdgeType::Shuffle => "⤨", + LogicalEdgeType::LeftJoin => "-[left]⤨", + LogicalEdgeType::RightJoin => "-[right]⤨", + }; + write!(f, "{symbol}") + } +} + +impl From for LogicalEdgeType { + fn from(value: ProtoEdgeType) -> Self { + match value { + ProtoEdgeType::Unused => { + panic!("Critical: Invalid EdgeType 'Unused' encountered") + } + ProtoEdgeType::Forward => Self::Forward, + ProtoEdgeType::Shuffle => Self::Shuffle, + ProtoEdgeType::LeftJoin => Self::LeftJoin, + ProtoEdgeType::RightJoin => Self::RightJoin, + } + } +} + +impl From for ProtoEdgeType { + fn from(value: LogicalEdgeType) -> Self { + match value { + LogicalEdgeType::Forward => Self::Forward, + LogicalEdgeType::Shuffle => Self::Shuffle, + LogicalEdgeType::LeftJoin => Self::LeftJoin, + LogicalEdgeType::RightJoin => Self::RightJoin, + } + } +} + +pub fn logical_edge_type_from_proto_i32(i: i32) -> Result { + let e = ProtoEdgeType::try_from(i).map_err(|_| { + DataFusionError::Plan(format!("invalid protobuf EdgeType discriminant {i}")) + })?; + match e { + ProtoEdgeType::Unused => Err(DataFusionError::Plan( + "Critical: Invalid EdgeType 'Unused' encountered".into(), + )), + ProtoEdgeType::Forward => Ok(LogicalEdgeType::Forward), + ProtoEdgeType::Shuffle => Ok(LogicalEdgeType::Shuffle), + ProtoEdgeType::LeftJoin => Ok(LogicalEdgeType::LeftJoin), + ProtoEdgeType::RightJoin => Ok(LogicalEdgeType::RightJoin), + } +} + +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] +pub struct LogicalEdge { + pub edge_type: LogicalEdgeType, + pub schema: Arc, +} + +impl LogicalEdge { + pub fn new(edge_type: LogicalEdgeType, schema: FsSchema) -> Self { + LogicalEdge { + edge_type, + schema: Arc::new(schema), + } + } + + pub fn project_all(edge_type: LogicalEdgeType, schema: FsSchema) -> Self { + LogicalEdge { + edge_type, + schema: Arc::new(schema), + } + } +} diff --git a/src/streaming_planner/src/logical_node/logical/logical_graph.rs b/src/streaming_planner/src/logical_node/logical/logical_graph.rs new file mode 100644 index 00000000..b877e2a0 --- /dev/null +++ b/src/streaming_planner/src/logical_node/logical/logical_graph.rs @@ -0,0 +1,30 @@ +// 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 petgraph::graph::DiGraph; + +use super::logical_edge::LogicalEdge; +use super::logical_node::LogicalNode; + +pub type LogicalGraph = DiGraph; + +pub trait Optimizer { + fn optimize_once(&self, plan: &mut LogicalGraph) -> bool; + + fn optimize(&self, plan: &mut LogicalGraph) { + loop { + if !self.optimize_once(plan) { + break; + } + } + } +} diff --git a/src/streaming_planner/src/logical_node/logical/logical_node.rs b/src/streaming_planner/src/logical_node/logical/logical_node.rs new file mode 100644 index 00000000..5f00dc4b --- /dev/null +++ b/src/streaming_planner/src/logical_node/logical/logical_node.rs @@ -0,0 +1,87 @@ +// 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::fmt::{Debug, Display, Formatter}; + +use datafusion::common::{DataFusionError, Result}; +use itertools::Itertools; +use serde::{Deserialize, Serialize}; + +use super::operator_chain::{ChainedLogicalOperator, OperatorChain}; +use super::operator_name::OperatorName; + +#[derive(Clone, Serialize, Deserialize)] +pub struct LogicalNode { + pub node_id: u32, + pub description: String, + pub operator_chain: OperatorChain, + pub parallelism: usize, +} + +impl LogicalNode { + pub fn single( + id: u32, + operator_id: String, + name: OperatorName, + config: Vec, + description: String, + parallelism: usize, + ) -> Self { + Self { + node_id: id, + description, + operator_chain: OperatorChain { + operators: vec![ChainedLogicalOperator { + operator_id, + operator_name: name, + operator_config: config, + }], + edges: vec![], + }, + parallelism, + } + } + + pub fn resolve_pipeline_operator_name(&self) -> Result { + let first_op = self.operator_chain.operators.first().ok_or_else(|| { + DataFusionError::Plan("Invalid LogicalNode: Operator chain is empty".into()) + })?; + + if let Some(connector_name) = first_op.extract_connector_name() { + return Ok(connector_name); + } + + if self.operator_chain.len() == 1 { + return Ok(first_op.operator_id.clone()); + } + + Ok("chained_op".to_string()) + } +} + +impl Display for LogicalNode { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.description) + } +} + +impl Debug for LogicalNode { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + let chain_path = self + .operator_chain + .operators + .iter() + .map(|op| op.operator_id.as_str()) + .join(" -> "); + write!(f, "{chain_path}[{}]", self.parallelism) + } +} diff --git a/src/streaming_planner/src/logical_node/logical/logical_program.rs b/src/streaming_planner/src/logical_node/logical/logical_program.rs new file mode 100644 index 00000000..119ac469 --- /dev/null +++ b/src/streaming_planner/src/logical_node/logical/logical_program.rs @@ -0,0 +1,153 @@ +// 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::hash_map::DefaultHasher; +use std::collections::{HashMap, HashSet}; +use std::hash::Hasher; +use std::sync::Arc; + +use datafusion::arrow::datatypes::Schema; +use datafusion::common::{DataFusionError, Result as DFResult}; +use petgraph::Direction; +use petgraph::dot::Dot; +use prost::Message; +use protocol::function_stream_graph::FsProgram; +use rand::distributions::Alphanumeric; +use rand::rngs::SmallRng; +use rand::{Rng, SeedableRng}; + +use super::logical_graph::{LogicalGraph, Optimizer}; +use super::operator_name::OperatorName; +use super::program_config::ProgramConfig; + +#[derive(Clone, Debug, Default)] +pub struct LogicalProgram { + pub graph: LogicalGraph, + pub program_config: ProgramConfig, +} + +impl LogicalProgram { + pub fn new(graph: LogicalGraph, program_config: ProgramConfig) -> Self { + Self { + graph, + program_config, + } + } + + pub fn optimize(&mut self, optimizer: &dyn Optimizer) { + optimizer.optimize(&mut self.graph); + } + + pub fn update_parallelism(&mut self, overrides: &HashMap) { + for node in self.graph.node_weights_mut() { + if let Some(&p) = overrides.get(&node.node_id) { + node.parallelism = p; + } + } + } + + pub fn dot(&self) -> String { + format!("{:?}", Dot::with_config(&self.graph, &[])) + } + + pub fn task_count(&self) -> usize { + self.graph.node_weights().map(|nw| nw.parallelism).sum() + } + + pub fn sources(&self) -> HashSet { + self.graph + .externals(Direction::Incoming) + .filter_map(|idx| self.graph.node_weight(idx)) + .map(|node| node.node_id) + .collect() + } + + pub fn get_hash(&self) -> String { + let mut hasher = DefaultHasher::new(); + let program_bytes = FsProgram::from(self.clone()).encode_to_vec(); + hasher.write(&program_bytes); + let rng = SmallRng::seed_from_u64(hasher.finish()); + rng.sample_iter(&Alphanumeric) + .take(16) + .map(|c| (c as char).to_ascii_lowercase()) + .collect() + } + + pub fn tasks_per_operator(&self) -> HashMap { + self.graph + .node_weights() + .flat_map(|node| { + node.operator_chain + .operators + .iter() + .map(move |op| (op.operator_id.clone(), node.parallelism)) + }) + .collect() + } + + pub fn operator_names_by_id(&self) -> HashMap { + self.graph + .node_weights() + .flat_map(|node| &node.operator_chain.operators) + .map(|op| { + let resolved_name = op + .extract_connector_name() + .unwrap_or_else(|| op.operator_name.to_string()); + (op.operator_id.clone(), resolved_name) + }) + .collect() + } + + pub fn tasks_per_node(&self) -> HashMap { + self.graph + .node_weights() + .map(|node| (node.node_id, node.parallelism)) + .collect() + } + + pub fn features(&self) -> HashSet { + self.graph + .node_weights() + .flat_map(|node| &node.operator_chain.operators) + .filter_map(|op| op.extract_feature()) + .collect() + } + + /// Arrow schema carried on edges into the connector-sink node, if present. + pub fn egress_arrow_schema(&self) -> Option> { + for idx in self.graph.node_indices() { + let node = self.graph.node_weight(idx)?; + if node + .operator_chain + .operators + .iter() + .any(|op| op.operator_name == OperatorName::ConnectorSink) + { + let e = self.graph.edges_directed(idx, Direction::Incoming).next()?; + return Some(Arc::clone(&e.weight().schema.schema)); + } + } + None + } + + pub fn encode_for_catalog(&self) -> DFResult> { + Ok(FsProgram::from(self.clone()).encode_to_vec()) + } + + pub fn decode_for_catalog(bytes: &[u8]) -> DFResult { + let proto = FsProgram::decode(bytes).map_err(|e| { + DataFusionError::Execution(format!("FsProgram catalog decode failed: {e}")) + })?; + LogicalProgram::try_from(proto) + } +} diff --git a/src/streaming_planner/src/logical_node/logical/mod.rs b/src/streaming_planner/src/logical_node/logical/mod.rs new file mode 100644 index 00000000..d2e9a327 --- /dev/null +++ b/src/streaming_planner/src/logical_node/logical/mod.rs @@ -0,0 +1,30 @@ +// 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. + +mod dylib_udf_config; +mod fs_program_convert; +mod logical_edge; +mod logical_graph; +mod logical_node; +mod logical_program; +mod operator_chain; +mod operator_name; +mod program_config; +mod python_udf_config; + +pub use dylib_udf_config::DylibUdfConfig; +pub use logical_edge::{LogicalEdge, LogicalEdgeType}; +pub use logical_graph::{LogicalGraph, Optimizer}; +pub use logical_node::LogicalNode; +pub use logical_program::LogicalProgram; +pub use operator_name::OperatorName; +pub use program_config::ProgramConfig; diff --git a/src/streaming_planner/src/logical_node/logical/operator_chain.rs b/src/streaming_planner/src/logical_node/logical/operator_chain.rs new file mode 100644 index 00000000..34be2f57 --- /dev/null +++ b/src/streaming_planner/src/logical_node/logical/operator_chain.rs @@ -0,0 +1,142 @@ +// 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::sync::Arc; + +use itertools::{EitherOrBoth, Itertools}; +use prost::Message; +use protocol::function_stream_graph::ConnectorOp; +use serde::{Deserialize, Serialize}; + +use super::operator_name::OperatorName; +use crate::common::FsSchema; + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct ChainedLogicalOperator { + pub operator_id: String, + pub operator_name: OperatorName, + pub operator_config: Vec, +} + +impl ChainedLogicalOperator { + pub fn extract_connector_name(&self) -> Option { + if matches!( + self.operator_name, + OperatorName::ConnectorSource | OperatorName::ConnectorSink + ) { + ConnectorOp::decode(self.operator_config.as_slice()) + .ok() + .map(|op| op.connector) + } else { + None + } + } + + pub fn extract_feature(&self) -> Option { + match self.operator_name { + OperatorName::AsyncUdf => Some("async-udf".to_string()), + OperatorName::Join => Some("join-with-expiration".to_string()), + OperatorName::InstantJoin => Some("windowed-join".to_string()), + OperatorName::WindowFunction => Some("sql-window-function".to_string()), + OperatorName::LookupJoin => Some("lookup-join".to_string()), + OperatorName::TumblingWindowAggregate => { + Some("sql-tumbling-window-aggregate".to_string()) + } + OperatorName::SlidingWindowAggregate => { + Some("sql-sliding-window-aggregate".to_string()) + } + OperatorName::SessionWindowAggregate => { + Some("sql-session-window-aggregate".to_string()) + } + OperatorName::UpdatingAggregate => Some("sql-updating-aggregate".to_string()), + OperatorName::ConnectorSource => { + self.extract_connector_name().map(|c| format!("{c}-source")) + } + OperatorName::ConnectorSink => { + self.extract_connector_name().map(|c| format!("{c}-sink")) + } + _ => None, + } + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct OperatorChain { + pub(crate) operators: Vec, + pub(crate) edges: Vec>, +} + +impl OperatorChain { + pub fn new(operator: ChainedLogicalOperator) -> Self { + Self { + operators: vec![operator], + edges: vec![], + } + } + + pub fn iter(&self) -> impl Iterator>)> { + self.operators + .iter() + .zip_longest(&self.edges) + .filter_map(|e| match e { + EitherOrBoth::Both(op, edge) => Some((op, Some(edge))), + EitherOrBoth::Left(op) => Some((op, None)), + EitherOrBoth::Right(_) => None, + }) + } + + pub fn iter_mut( + &mut self, + ) -> impl Iterator>)> { + self.operators + .iter_mut() + .zip_longest(&self.edges) + .filter_map(|e| match e { + EitherOrBoth::Both(op, edge) => Some((op, Some(edge))), + EitherOrBoth::Left(op) => Some((op, None)), + EitherOrBoth::Right(_) => None, + }) + } + + pub fn first(&self) -> &ChainedLogicalOperator { + self.operators + .first() + .expect("OperatorChain must contain at least one operator") + } + + pub fn len(&self) -> usize { + self.operators.len() + } + + pub fn is_empty(&self) -> bool { + self.operators.is_empty() + } + + pub fn is_source(&self) -> bool { + self.operators[0].operator_name == OperatorName::ConnectorSource + } + + pub fn is_sink(&self) -> bool { + self.operators[0].operator_name == OperatorName::ConnectorSink + } + + /// Operators safe to run at a higher upstream `TaskContext::parallelism` when fused after a + /// stateful node (e.g. window aggregate @ 8 → projection @ 1). + pub fn is_parallelism_upstream_expandable(&self) -> bool { + self.operators.iter().all(|op| { + matches!( + op.operator_name, + OperatorName::Projection | OperatorName::Value | OperatorName::ExpressionWatermark + ) + }) + } +} diff --git a/src/streaming_planner/src/logical_node/logical/operator_name.rs b/src/streaming_planner/src/logical_node/logical/operator_name.rs new file mode 100644 index 00000000..4377b0bd --- /dev/null +++ b/src/streaming_planner/src/logical_node/logical/operator_name.rs @@ -0,0 +1,82 @@ +// 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::str::FromStr; + +use serde::{Deserialize, Deserializer, Serialize, Serializer}; +use strum::{Display, EnumString, IntoStaticStr}; + +use crate::common::constants::operator_feature; + +#[derive(Clone, Copy, Debug, Eq, PartialEq, EnumString, Display, IntoStaticStr)] +pub enum OperatorName { + ExpressionWatermark, + Value, + KeyBy, + Projection, + AsyncUdf, + Join, + InstantJoin, + LookupJoin, + WindowFunction, + TumblingWindowAggregate, + SlidingWindowAggregate, + SessionWindowAggregate, + UpdatingAggregate, + ConnectorSource, + ConnectorSink, +} + +impl OperatorName { + /// Registry / worker lookup key; matches [`Display`] and protobuf operator names. + #[inline] + pub fn as_registry_key(self) -> &'static str { + self.into() + } + + pub fn feature_tag(self) -> Option<&'static str> { + match self { + Self::ExpressionWatermark | Self::Value | Self::Projection => None, + Self::AsyncUdf => Some(operator_feature::ASYNC_UDF), + Self::Join => Some(operator_feature::JOIN_WITH_EXPIRATION), + Self::InstantJoin => Some(operator_feature::WINDOWED_JOIN), + Self::WindowFunction => Some(operator_feature::SQL_WINDOW_FUNCTION), + Self::LookupJoin => Some(operator_feature::LOOKUP_JOIN), + Self::TumblingWindowAggregate => Some(operator_feature::SQL_TUMBLING_WINDOW_AGGREGATE), + Self::SlidingWindowAggregate => Some(operator_feature::SQL_SLIDING_WINDOW_AGGREGATE), + Self::SessionWindowAggregate => Some(operator_feature::SQL_SESSION_WINDOW_AGGREGATE), + Self::UpdatingAggregate => Some(operator_feature::SQL_UPDATING_AGGREGATE), + Self::KeyBy => Some(operator_feature::KEY_BY_ROUTING), + Self::ConnectorSource => Some(operator_feature::CONNECTOR_SOURCE), + Self::ConnectorSink => Some(operator_feature::CONNECTOR_SINK), + } + } +} + +impl Serialize for OperatorName { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + serializer.serialize_str(&self.to_string()) + } +} + +impl<'de> Deserialize<'de> for OperatorName { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let s = String::deserialize(deserializer)?; + Self::from_str(&s).map_err(serde::de::Error::custom) + } +} diff --git a/src/streaming_planner/src/logical_node/logical/program_config.rs b/src/streaming_planner/src/logical_node/logical/program_config.rs new file mode 100644 index 00000000..177326f4 --- /dev/null +++ b/src/streaming_planner/src/logical_node/logical/program_config.rs @@ -0,0 +1,33 @@ +// 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 protocol::function_stream_graph::FsProgramConfig; + +/// Placeholder program-level config (UDF tables live elsewhere; wire maps stay empty). +#[derive(Clone, Debug, Default)] +pub struct ProgramConfig {} + +impl From for FsProgramConfig { + fn from(_: ProgramConfig) -> Self { + Self { + udf_dylibs: Default::default(), + python_udfs: Default::default(), + } + } +} + +impl From for ProgramConfig { + fn from(_: FsProgramConfig) -> Self { + Self::default() + } +} diff --git a/src/streaming_planner/src/logical_node/logical/python_udf_config.rs b/src/streaming_planner/src/logical_node/logical/python_udf_config.rs new file mode 100644 index 00000000..6e7d5c66 --- /dev/null +++ b/src/streaming_planner/src/logical_node/logical/python_udf_config.rs @@ -0,0 +1,23 @@ +// 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::sync::Arc; + +use datafusion::arrow::datatypes::DataType; + +#[derive(Clone, Debug, Eq, PartialEq, Hash)] +pub struct PythonUdfConfig { + pub arg_types: Vec, + pub return_type: DataType, + pub name: Arc, + pub definition: Arc, +} diff --git a/src/streaming_planner/src/logical_node/lookup.rs b/src/streaming_planner/src/logical_node/lookup.rs new file mode 100644 index 00000000..9075da60 --- /dev/null +++ b/src/streaming_planner/src/logical_node/lookup.rs @@ -0,0 +1,256 @@ +// 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::fmt::Formatter; +use std::sync::Arc; + +use datafusion::common::{Column, DFSchemaRef, JoinType, Result, internal_err, plan_err}; +use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion::sql::TableReference; +use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec; +use datafusion_proto::physical_plan::to_proto::serialize_physical_expr; +use prost::Message; + +use protocol::function_stream_graph; +use protocol::function_stream_graph::{ConnectorOp, LookupJoinCondition, LookupJoinOperator}; + +use crate::multifield_partial_ord; +use crate::common::constants::extension_node; +use crate::common::{FsSchema, FsSchemaRef}; +use crate::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; +use crate::logical_node::{CompiledTopologyNode, StreamingOperatorBlueprint}; +use crate::logical_planner::planner::{NamedNode, Planner}; +use crate::schema::LookupTable; +use crate::schema::utils::add_timestamp_field_arrow; + +pub const DICTIONARY_SOURCE_NODE_NAME: &str = extension_node::REFERENCE_TABLE_SOURCE; +pub const STREAM_DICTIONARY_JOIN_NODE_NAME: &str = extension_node::STREAM_REFERENCE_JOIN; + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct ReferenceTableSourceNode { + pub(crate) source_definition: LookupTable, + pub(crate) resolved_schema: DFSchemaRef, +} + +multifield_partial_ord!(ReferenceTableSourceNode, source_definition); + +impl UserDefinedLogicalNodeCore for ReferenceTableSourceNode { + fn name(&self) -> &str { + DICTIONARY_SOURCE_NODE_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![] + } + + fn schema(&self) -> &DFSchemaRef { + &self.resolved_schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + write!(f, "ReferenceTableSource: Schema={}", self.resolved_schema) + } + + fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { + if !inputs.is_empty() { + return internal_err!( + "ReferenceTableSource is a leaf node and cannot accept upstream inputs" + ); + } + + Ok(Self { + source_definition: self.source_definition.clone(), + resolved_schema: self.resolved_schema.clone(), + }) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct StreamReferenceJoinNode { + pub(crate) upstream_stream_plan: LogicalPlan, + pub(crate) output_schema: DFSchemaRef, + pub(crate) external_dictionary: LookupTable, + pub(crate) equijoin_conditions: Vec<(Expr, Column)>, + pub(crate) post_join_filter: Option, + pub(crate) namespace_alias: Option, + pub(crate) join_semantics: JoinType, +} + +multifield_partial_ord!( + StreamReferenceJoinNode, + upstream_stream_plan, + external_dictionary, + equijoin_conditions, + post_join_filter, + namespace_alias +); + +impl StreamReferenceJoinNode { + fn compile_join_conditions(&self, planner: &Planner) -> Result> { + self.equijoin_conditions + .iter() + .map(|(logical_left_expr, right_column)| { + let physical_expr = + planner.create_physical_expr(logical_left_expr, &self.output_schema)?; + let serialized_expr = + serialize_physical_expr(&physical_expr, &DefaultPhysicalExtensionCodec {})?; + + Ok(LookupJoinCondition { + left_expr: serialized_expr.encode_to_vec(), + right_key: right_column.name.clone(), + }) + }) + .collect() + } + + fn map_api_join_type(&self) -> Result { + match self.join_semantics { + JoinType::Inner => Ok(function_stream_graph::JoinType::Inner as i32), + JoinType::Left => Ok(function_stream_graph::JoinType::Left as i32), + unsupported => plan_err!( + "Unsupported join type '{unsupported}' for dictionary lookups. Only INNER and LEFT joins are permitted." + ), + } + } + + fn build_engine_operator( + &self, + planner: &Planner, + _upstream_schema: &FsSchemaRef, + ) -> Result { + let internal_input_schema = + FsSchema::from_schema_unkeyed(Arc::new(self.output_schema.as_ref().into()))?; + let dictionary_physical_schema = self.external_dictionary.produce_physical_schema(); + let lookup_fs_schema = + FsSchema::from_schema_unkeyed(add_timestamp_field_arrow(dictionary_physical_schema))?; + + Ok(LookupJoinOperator { + input_schema: Some(internal_input_schema.into()), + lookup_schema: Some(lookup_fs_schema.clone().into()), + connector: Some(ConnectorOp { + connector: self.external_dictionary.adapter_type.clone(), + fs_schema: Some(lookup_fs_schema.into()), + name: self.external_dictionary.table_identifier.clone(), + description: self.external_dictionary.description.clone(), + config: Some(self.external_dictionary.connector_config.to_proto_config()), + }), + key_exprs: self.compile_join_conditions(planner)?, + join_type: self.map_api_join_type()?, + ttl_micros: self + .external_dictionary + .lookup_cache_ttl + .map(|t| t.as_micros() as u64), + max_capacity_bytes: self.external_dictionary.lookup_cache_max_bytes, + }) + } +} + +impl StreamingOperatorBlueprint for StreamReferenceJoinNode { + fn operator_identity(&self) -> Option { + None + } + + fn compile_to_graph_node( + &self, + planner: &Planner, + node_index: usize, + mut input_schemas: Vec, + ) -> Result { + if input_schemas.len() != 1 { + return plan_err!("StreamReferenceJoinNode requires exactly one upstream stream input"); + } + let upstream_schema = input_schemas.remove(0); + + let operator_config = self.build_engine_operator(planner, &upstream_schema)?; + + let logical_node = LogicalNode::single( + node_index as u32, + format!("lookup_join_{node_index}"), + OperatorName::LookupJoin, + operator_config.encode_to_vec(), + format!( + "DictionaryJoin<{}>", + self.external_dictionary.table_identifier + ), + planner.default_parallelism(), + ); + + let incoming_edge = + LogicalEdge::project_all(LogicalEdgeType::Shuffle, (*upstream_schema).clone()); + + Ok(CompiledTopologyNode { + execution_unit: logical_node, + routing_edges: vec![incoming_edge], + }) + } + + fn yielded_schema(&self) -> FsSchema { + FsSchema::from_schema_unkeyed(self.output_schema.inner().clone()) + .expect("Failed to convert lookup join output schema to FsSchema") + } +} + +impl UserDefinedLogicalNodeCore for StreamReferenceJoinNode { + fn name(&self) -> &str { + STREAM_DICTIONARY_JOIN_NODE_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.upstream_stream_plan] + } + + fn schema(&self) -> &DFSchemaRef { + &self.output_schema + } + + fn expressions(&self) -> Vec { + let mut exprs: Vec<_> = self + .equijoin_conditions + .iter() + .map(|(l, _)| l.clone()) + .collect(); + if let Some(filter) = &self.post_join_filter { + exprs.push(filter.clone()); + } + exprs + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + write!( + f, + "StreamReferenceJoin: join_type={:?} | {}", + self.join_semantics, self.output_schema + ) + } + + fn with_exprs_and_inputs(&self, _: Vec, inputs: Vec) -> Result { + if inputs.len() != 1 { + return internal_err!( + "StreamReferenceJoinNode expects exactly 1 upstream plan, got {}", + inputs.len() + ); + } + Ok(Self { + upstream_stream_plan: inputs[0].clone(), + output_schema: self.output_schema.clone(), + external_dictionary: self.external_dictionary.clone(), + equijoin_conditions: self.equijoin_conditions.clone(), + post_join_filter: self.post_join_filter.clone(), + namespace_alias: self.namespace_alias.clone(), + join_semantics: self.join_semantics, + }) + } +} diff --git a/src/streaming_planner/src/logical_node/macros.rs b/src/streaming_planner/src/logical_node/macros.rs new file mode 100644 index 00000000..4ce649c2 --- /dev/null +++ b/src/streaming_planner/src/logical_node/macros.rs @@ -0,0 +1,28 @@ +// 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. + +#[macro_export] +macro_rules! multifield_partial_ord { + ($ty:ty, $($field:tt), *) => { + impl PartialOrd for $ty { + fn partial_cmp(&self, other: &Self) -> Option { + $( + let cmp = self.$field.partial_cmp(&other.$field)?; + if cmp != std::cmp::Ordering::Equal { + return Some(cmp); + } + )* + Some(std::cmp::Ordering::Equal) + } + } + }; +} diff --git a/src/streaming_planner/src/logical_node/mod.rs b/src/streaming_planner/src/logical_node/mod.rs new file mode 100644 index 00000000..3ef02fc8 --- /dev/null +++ b/src/streaming_planner/src/logical_node/mod.rs @@ -0,0 +1,42 @@ +// 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. + +pub mod logical; + +mod macros; + +pub mod streaming_operator_blueprint; +pub use streaming_operator_blueprint::{CompiledTopologyNode, StreamingOperatorBlueprint}; + +pub mod aggregate; +pub mod debezium; +pub mod join; +pub mod key_calculation; +pub mod lookup; +pub mod projection; +pub mod remote_table; +pub mod sink; +pub mod table_source; +pub mod updating_aggregate; +pub mod watermark_node; +pub mod windows_function; + +pub mod timestamp_append; +pub use timestamp_append::SystemTimestampInjectorNode; + +pub mod async_udf; +pub use async_udf::AsyncFunctionExecutionNode; + +pub mod is_retract; +pub use is_retract::IsRetractExtension; + +mod extension_try_from; diff --git a/src/streaming_planner/src/logical_node/projection.rs b/src/streaming_planner/src/logical_node/projection.rs new file mode 100644 index 00000000..c7d06e29 --- /dev/null +++ b/src/streaming_planner/src/logical_node/projection.rs @@ -0,0 +1,239 @@ +// 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::fmt::Formatter; +use std::sync::Arc; + +use datafusion::common::{DFSchema, DFSchemaRef, Result, internal_err}; +use datafusion::logical_expr::{Expr, ExprSchemable, LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec; +use datafusion_proto::physical_plan::to_proto::serialize_physical_expr; +use prost::Message; + +use protocol::function_stream_graph::ProjectionOperator; + +use crate::multifield_partial_ord; +use crate::common::constants::{extension_node, sql_field}; +use crate::common::{FsSchema, FsSchemaRef}; +use crate::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; +use crate::logical_node::{CompiledTopologyNode, StreamingOperatorBlueprint}; +use crate::logical_planner::planner::{NamedNode, Planner}; +use crate::types::{QualifiedField, build_df_schema}; + +// ----------------------------------------------------------------------------- +// Constants & Identifiers +// ----------------------------------------------------------------------------- + +pub const STREAM_PROJECTION_NODE_NAME: &str = extension_node::STREAM_PROJECTION; +const DEFAULT_PROJECTION_LABEL: &str = sql_field::DEFAULT_PROJECTION_LABEL; + +// ----------------------------------------------------------------------------- +// Logical Node Definition +// ----------------------------------------------------------------------------- + +/// Projection within a streaming execution topology. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct StreamProjectionNode { + pub(crate) upstream_plans: Vec, + pub(crate) operator_label: Option, + pub(crate) projection_exprs: Vec, + pub(crate) resolved_schema: DFSchemaRef, + pub(crate) requires_shuffle: bool, +} + +multifield_partial_ord!(StreamProjectionNode, operator_label, projection_exprs); + +impl StreamProjectionNode { + pub fn try_new( + upstream_plans: Vec, + operator_label: Option, + projection_exprs: Vec, + ) -> Result { + if upstream_plans.is_empty() { + return internal_err!("StreamProjectionNode requires at least one upstream plan"); + } + let primary_input = &upstream_plans[0]; + let upstream_schema = primary_input.schema(); + + let mut projected_fields = Vec::with_capacity(projection_exprs.len()); + for logical_expr in &projection_exprs { + let arrow_field = logical_expr.to_field(upstream_schema)?; + projected_fields.push(QualifiedField::from(arrow_field)); + } + + let resolved_schema = Arc::new(build_df_schema(&projected_fields)?); + + Ok(Self { + upstream_plans, + operator_label, + projection_exprs, + resolved_schema, + requires_shuffle: false, + }) + } + + pub fn with_shuffle_routing(mut self) -> Self { + self.requires_shuffle = true; + self + } + + fn validate_uniform_schemas(input_schemas: &[FsSchemaRef]) -> Result { + if input_schemas.is_empty() { + return internal_err!("No input schemas provided to projection planner"); + } + let primary_schema = input_schemas[0].clone(); + + for schema in input_schemas.iter().skip(1) { + if **schema != *primary_schema { + return internal_err!( + "Schema mismatch: All upstream inputs to a projection node must share the identical schema topology." + ); + } + } + + Ok(primary_schema) + } + + fn compile_physical_expressions( + &self, + planner: &Planner, + input_df_schema: &DFSchemaRef, + ) -> Result>> { + self.projection_exprs + .iter() + .map(|logical_expr| { + let physical_expr = planner + .create_physical_expr(logical_expr, input_df_schema) + .map_err(|e| e.context("Failed to compile physical projection expression"))?; + + let serialized_expr = + serialize_physical_expr(&physical_expr, &DefaultPhysicalExtensionCodec {})?; + + Ok(serialized_expr.encode_to_vec()) + }) + .collect() + } +} + +// ----------------------------------------------------------------------------- +// Stream Extension Trait Implementation +// ----------------------------------------------------------------------------- + +impl StreamingOperatorBlueprint for StreamProjectionNode { + fn operator_identity(&self) -> Option { + None + } + + fn compile_to_graph_node( + &self, + planner: &Planner, + node_index: usize, + input_schemas: Vec, + ) -> Result { + let unified_input_schema = Self::validate_uniform_schemas(&input_schemas)?; + let input_df_schema = Arc::new(DFSchema::try_from( + unified_input_schema.schema.as_ref().clone(), + )?); + + let compiled_expr_payloads = + self.compile_physical_expressions(planner, &input_df_schema)?; + + let operator_config = ProjectionOperator { + name: self + .operator_label + .as_deref() + .unwrap_or(DEFAULT_PROJECTION_LABEL) + .to_string(), + input_schema: Some(unified_input_schema.as_ref().clone().into()), + output_schema: Some(self.yielded_schema().into()), + exprs: compiled_expr_payloads, + }; + + let node_identifier = format!("projection_{node_index}"); + let label = format!( + "ArrowProjection<{}>", + self.operator_label.as_deref().unwrap_or("_") + ); + + let logical_node = LogicalNode::single( + node_index as u32, + node_identifier, + OperatorName::Projection, + operator_config.encode_to_vec(), + label, + planner.default_parallelism(), + ); + + let routing_strategy = if self.requires_shuffle { + LogicalEdgeType::Shuffle + } else { + LogicalEdgeType::Forward + }; + + let outgoing_edge = + LogicalEdge::project_all(routing_strategy, (*unified_input_schema).clone()); + + Ok(CompiledTopologyNode { + execution_unit: logical_node, + routing_edges: vec![outgoing_edge], + }) + } + + fn yielded_schema(&self) -> FsSchema { + FsSchema::from_schema_unkeyed(Arc::new(self.resolved_schema.as_arrow().clone())) + .expect("Fatal: Failed to generate unkeyed output schema for projection") + } +} + +// ----------------------------------------------------------------------------- +// DataFusion Logical Node Hooks +// ----------------------------------------------------------------------------- + +impl UserDefinedLogicalNodeCore for StreamProjectionNode { + fn name(&self) -> &str { + STREAM_PROJECTION_NODE_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + self.upstream_plans.iter().collect() + } + + fn schema(&self) -> &DFSchemaRef { + &self.resolved_schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + write!( + f, + "StreamProjectionNode: RequiresShuffle={}, Schema={}", + self.requires_shuffle, self.resolved_schema + ) + } + + fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { + let mut new_node = Self::try_new( + inputs, + self.operator_label.clone(), + self.projection_exprs.clone(), + )?; + + if self.requires_shuffle { + new_node = new_node.with_shuffle_routing(); + } + + Ok(new_node) + } +} diff --git a/src/streaming_planner/src/logical_node/remote_table.rs b/src/streaming_planner/src/logical_node/remote_table.rs new file mode 100644 index 00000000..3b050c62 --- /dev/null +++ b/src/streaming_planner/src/logical_node/remote_table.rs @@ -0,0 +1,190 @@ +// 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::fmt::Formatter; +use std::sync::Arc; + +use datafusion::common::{DFSchemaRef, Result, TableReference, internal_err, plan_err}; +use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion_proto::physical_plan::AsExecutionPlan; +use datafusion_proto::protobuf::PhysicalPlanNode; +use prost::Message; + +use protocol::function_stream_graph::ValuePlanOperator; + +use crate::multifield_partial_ord; +use crate::common::constants::extension_node; +use crate::common::{FsSchema, FsSchemaRef}; +use crate::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; +use crate::logical_node::{CompiledTopologyNode, StreamingOperatorBlueprint}; +use crate::logical_planner::planner::{NamedNode, Planner}; +use crate::physical::StreamingExtensionCodec; + +// ----------------------------------------------------------------------------- +// Constants & Identifiers +// ----------------------------------------------------------------------------- + +pub const REMOTE_TABLE_NODE_NAME: &str = extension_node::REMOTE_TABLE_BOUNDARY; + +// ----------------------------------------------------------------------------- +// Logical Node Definition +// ----------------------------------------------------------------------------- + +/// Segments the execution graph and merges nodes sharing the same identifier; acts as a boundary. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct RemoteTableBoundaryNode { + pub(crate) upstream_plan: LogicalPlan, + pub(crate) table_identifier: TableReference, + pub(crate) resolved_schema: DFSchemaRef, + pub(crate) requires_materialization: bool, +} + +multifield_partial_ord!( + RemoteTableBoundaryNode, + upstream_plan, + table_identifier, + requires_materialization +); + +impl RemoteTableBoundaryNode { + fn compile_engine_operator(&self, planner: &Planner) -> Result> { + let physical_plan = planner.sync_plan(&self.upstream_plan)?; + + let physical_plan_proto = PhysicalPlanNode::try_from_physical_plan( + physical_plan, + &StreamingExtensionCodec::default(), + )?; + + let operator_config = ValuePlanOperator { + name: format!("value_calculation({})", self.table_identifier), + physical_plan: physical_plan_proto.encode_to_vec(), + }; + + Ok(operator_config.encode_to_vec()) + } + + fn validate_uniform_schemas(input_schemas: &[FsSchemaRef]) -> Result<()> { + if input_schemas.len() <= 1 { + return Ok(()); + } + + let primary_schema = &input_schemas[0]; + for schema in input_schemas.iter().skip(1) { + if *schema != *primary_schema { + return plan_err!( + "Topology error: Multiple input streams routed to the same remote table must share an identical schema structure." + ); + } + } + + Ok(()) + } +} + +// ----------------------------------------------------------------------------- +// Stream Extension Trait Implementation +// ----------------------------------------------------------------------------- + +impl StreamingOperatorBlueprint for RemoteTableBoundaryNode { + fn operator_identity(&self) -> Option { + if self.requires_materialization { + Some(NamedNode::RemoteTable(self.table_identifier.clone())) + } else { + None + } + } + + fn compile_to_graph_node( + &self, + planner: &Planner, + node_index: usize, + input_schemas: Vec, + ) -> Result { + Self::validate_uniform_schemas(&input_schemas)?; + + let operator_payload = self.compile_engine_operator(planner)?; + + let logical_node = LogicalNode::single( + node_index as u32, + format!("value_{node_index}"), + OperatorName::Value, + operator_payload, + self.table_identifier.to_string(), + planner.default_parallelism(), + ); + + let routing_edges: Vec = input_schemas + .into_iter() + .map(|schema| LogicalEdge::project_all(LogicalEdgeType::Forward, (*schema).clone())) + .collect(); + + Ok(CompiledTopologyNode { + execution_unit: logical_node, + routing_edges, + }) + } + + fn yielded_schema(&self) -> FsSchema { + FsSchema::from_schema_keys(Arc::new(self.resolved_schema.as_ref().into()), vec![]) + .expect("Fatal: Failed to generate output schema for remote table boundary") + } +} + +// ----------------------------------------------------------------------------- +// DataFusion Logical Node Hooks +// ----------------------------------------------------------------------------- + +impl UserDefinedLogicalNodeCore for RemoteTableBoundaryNode { + fn name(&self) -> &str { + REMOTE_TABLE_NODE_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.upstream_plan] + } + + fn schema(&self) -> &DFSchemaRef { + &self.resolved_schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + write!( + f, + "RemoteTableBoundaryNode: Identifier={}, Materialized={}, Schema={}", + self.table_identifier, self.requires_materialization, self.resolved_schema + ) + } + + fn with_exprs_and_inputs( + &self, + _exprs: Vec, + mut inputs: Vec, + ) -> Result { + if inputs.len() != 1 { + return internal_err!( + "RemoteTableBoundaryNode expects exactly 1 upstream logical plan, but received {}", + inputs.len() + ); + } + + Ok(Self { + upstream_plan: inputs.remove(0), + table_identifier: self.table_identifier.clone(), + resolved_schema: self.resolved_schema.clone(), + requires_materialization: self.requires_materialization, + }) + } +} diff --git a/src/streaming_planner/src/logical_node/sink.rs b/src/streaming_planner/src/logical_node/sink.rs new file mode 100644 index 00000000..e93ea02c --- /dev/null +++ b/src/streaming_planner/src/logical_node/sink.rs @@ -0,0 +1,247 @@ +// 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::fmt::Formatter; +use std::sync::Arc; + +use datafusion::common::{DFSchemaRef, Result, TableReference, plan_err}; +use datafusion::logical_expr::{Expr, Extension, LogicalPlan, UserDefinedLogicalNodeCore}; +use prost::Message; + +use crate::multifield_partial_ord; +use crate::common::constants::extension_node; +use crate::common::{FsSchema, FsSchemaRef, UPDATING_META_FIELD}; +use crate::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; +use crate::logical_node::{CompiledTopologyNode, StreamingOperatorBlueprint}; +use crate::logical_planner::planner::{NamedNode, Planner}; +use crate::schema::CatalogEntity; +use crate::schema::catalog::ExternalTable; + +use super::debezium::PackDebeziumEnvelopeNode; +use super::remote_table::RemoteTableBoundaryNode; + +// ----------------------------------------------------------------------------- +// Constants & Identifiers +// ----------------------------------------------------------------------------- + +pub const STREAM_EGRESS_NODE_NAME: &str = extension_node::STREAM_EGRESS; + +// ----------------------------------------------------------------------------- +// Logical Node Definition +// ----------------------------------------------------------------------------- + +/// Terminal node routing processed data into an external sink (e.g. Kafka, PostgreSQL). +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct StreamEgressNode { + pub(crate) target_identifier: TableReference, + pub(crate) destination_table: CatalogEntity, + pub(crate) egress_schema: DFSchemaRef, + upstream_plans: Arc>, +} + +multifield_partial_ord!(StreamEgressNode, target_identifier, upstream_plans); + +impl StreamEgressNode { + pub fn try_new( + target_identifier: TableReference, + destination_table: CatalogEntity, + initial_schema: DFSchemaRef, + upstream_plan: LogicalPlan, + ) -> Result { + let (mut processed_plan, mut resolved_schema) = + Self::apply_cdc_transformations(upstream_plan, initial_schema, &destination_table)?; + + Self::enforce_computational_boundary(&mut resolved_schema, &mut processed_plan); + + Ok(Self { + target_identifier, + destination_table, + egress_schema: resolved_schema, + upstream_plans: Arc::new(vec![processed_plan]), + }) + } + + fn apply_cdc_transformations( + plan: LogicalPlan, + schema: DFSchemaRef, + destination: &CatalogEntity, + ) -> Result<(LogicalPlan, DFSchemaRef)> { + let is_upstream_updating = plan + .schema() + .has_column_with_unqualified_name(UPDATING_META_FIELD); + + match destination { + CatalogEntity::ExternalConnector(b) => match b.as_ref() { + ExternalTable::Sink(sink) => { + let is_sink_updating = sink.is_updating(); + + match (is_upstream_updating, is_sink_updating) { + (_, true) => { + let debezium_encoder = PackDebeziumEnvelopeNode::try_new(plan)?; + let wrapped_plan = LogicalPlan::Extension(Extension { + node: Arc::new(debezium_encoder), + }); + let new_schema = wrapped_plan.schema().clone(); + + Ok((wrapped_plan, new_schema)) + } + (true, false) => { + plan_err!( + "Topology Mismatch: The upstream is producing an updating stream (CDC), \ + but the target sink '{}' is not configured to accept updates. \ + Hint: set `format = 'debezium_json'` in the WITH clause.", + sink.name() + ) + } + (false, false) => Ok((plan, schema)), + } + } + ExternalTable::Source(source) => { + let is_sink_updating = source.is_updating(); + match (is_upstream_updating, is_sink_updating) { + (_, true) => { + let debezium_encoder = PackDebeziumEnvelopeNode::try_new(plan)?; + let wrapped_plan = LogicalPlan::Extension(Extension { + node: Arc::new(debezium_encoder), + }); + let new_schema = wrapped_plan.schema().clone(); + Ok((wrapped_plan, new_schema)) + } + (true, false) => plan_err!( + "Topology Mismatch: upstream produces CDC but target '{}' is a non-updating source table", + source.name() + ), + (false, false) => Ok((plan, schema)), + } + } + ExternalTable::Lookup(_) => plan_err!( + "Topology Violation: A Lookup Table cannot be used as a streaming data sink." + ), + }, + CatalogEntity::ComputedTable { .. } => Ok((plan, schema)), + } + } + + fn enforce_computational_boundary(schema: &mut DFSchemaRef, plan: &mut LogicalPlan) { + let requires_boundary = if let LogicalPlan::Extension(extension) = plan { + let stream_ext: &dyn StreamingOperatorBlueprint = (&extension.node) + .try_into() + .expect("Fatal: Egress node encountered an extension that does not implement StreamingOperatorBlueprint"); + + stream_ext.is_passthrough_boundary() + } else { + true + }; + + if requires_boundary { + let boundary_node = RemoteTableBoundaryNode { + upstream_plan: plan.clone(), + table_identifier: TableReference::bare("sink projection"), + resolved_schema: schema.clone(), + requires_materialization: false, + }; + + *plan = LogicalPlan::Extension(Extension { + node: Arc::new(boundary_node), + }); + } + } +} + +// ----------------------------------------------------------------------------- +// Stream Extension Trait Implementation +// ----------------------------------------------------------------------------- + +impl StreamingOperatorBlueprint for StreamEgressNode { + fn operator_identity(&self) -> Option { + Some(NamedNode::Sink(self.target_identifier.clone())) + } + + fn compile_to_graph_node( + &self, + planner: &Planner, + node_index: usize, + input_schemas: Vec, + ) -> Result { + let connector_operator = self + .destination_table + .connector_op() + .map_err(|e| e.context("Failed to generate connector operation payload"))?; + + let operator_description = connector_operator.description.clone(); + let operator_payload = connector_operator.encode_to_vec(); + + let logical_node = LogicalNode::single( + node_index as u32, + format!("sink_{}_{node_index}", self.target_identifier), + OperatorName::ConnectorSink, + operator_payload, + operator_description, + planner.default_parallelism(), + ); + + let routing_edges: Vec = input_schemas + .into_iter() + .map(|input_schema| { + LogicalEdge::project_all(LogicalEdgeType::Forward, (*input_schema).clone()) + }) + .collect(); + + Ok(CompiledTopologyNode { + execution_unit: logical_node, + routing_edges, + }) + } + + fn yielded_schema(&self) -> FsSchema { + FsSchema::from_fields(vec![]) + } +} + +// ----------------------------------------------------------------------------- +// DataFusion Logical Node Hooks +// ----------------------------------------------------------------------------- + +impl UserDefinedLogicalNodeCore for StreamEgressNode { + fn name(&self) -> &str { + STREAM_EGRESS_NODE_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + self.upstream_plans.iter().collect() + } + + fn schema(&self) -> &DFSchemaRef { + &self.egress_schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + write!( + f, + "StreamEgressNode({:?}): Schema={}", + self.target_identifier, self.egress_schema + ) + } + + fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { + Ok(Self { + target_identifier: self.target_identifier.clone(), + destination_table: self.destination_table.clone(), + egress_schema: self.egress_schema.clone(), + upstream_plans: Arc::new(inputs), + }) + } +} diff --git a/src/streaming_planner/src/logical_node/streaming_operator_blueprint.rs b/src/streaming_planner/src/logical_node/streaming_operator_blueprint.rs new file mode 100644 index 00000000..d9afaeab --- /dev/null +++ b/src/streaming_planner/src/logical_node/streaming_operator_blueprint.rs @@ -0,0 +1,65 @@ +// 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::fmt::Debug; + +use datafusion::common::Result; + +use crate::common::{FsSchema, FsSchemaRef}; +use crate::logical_node::logical::{LogicalEdge, LogicalNode}; +use crate::logical_planner::planner::{NamedNode, Planner}; + +// ----------------------------------------------------------------------------- +// Core Execution Blueprint +// ----------------------------------------------------------------------------- + +/// Atomic unit within a streaming execution topology: translates streaming SQL into graph nodes. +pub trait StreamingOperatorBlueprint: Debug { + /// Canonical named identity for this operator, if any (sources, sinks, etc.). + fn operator_identity(&self) -> Option; + + /// Compiles this operator into a graph vertex and its incoming routing edges. + fn compile_to_graph_node( + &self, + compiler_context: &Planner, + node_id_sequence: usize, + upstream_schemas: Vec, + ) -> Result; + + /// Schema of records this operator yields downstream. + fn yielded_schema(&self) -> FsSchema; + + /// Logical passthrough boundary (no physical state change); default is stateful / materializing. + fn is_passthrough_boundary(&self) -> bool { + false + } +} + +// ----------------------------------------------------------------------------- +// Graph Topology Structures +// ----------------------------------------------------------------------------- + +/// Compiled vertex: execution unit plus upstream routing edges. +#[derive(Debug, Clone)] +pub struct CompiledTopologyNode { + pub execution_unit: LogicalNode, + pub routing_edges: Vec, +} + +impl CompiledTopologyNode { + pub fn new(execution_unit: LogicalNode, routing_edges: Vec) -> Self { + Self { + execution_unit, + routing_edges, + } + } +} diff --git a/src/streaming_planner/src/logical_node/table_source.rs b/src/streaming_planner/src/logical_node/table_source.rs new file mode 100644 index 00000000..c0067a7d --- /dev/null +++ b/src/streaming_planner/src/logical_node/table_source.rs @@ -0,0 +1,180 @@ +// 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::fmt::Formatter; +use std::sync::Arc; + +use datafusion::common::{DFSchemaRef, Result, TableReference, plan_err}; +use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; +use prost::Message; + +use crate::multifield_partial_ord; +use crate::common::constants::extension_node; +use crate::common::{FsSchema, FsSchemaRef}; +use crate::logical_node::debezium::DebeziumSchemaCodec; +use crate::logical_node::logical::{LogicalNode, OperatorName}; +use crate::logical_planner::planner::{NamedNode, Planner}; +use crate::schema::SourceTable; +use crate::schema::utils::add_timestamp_field; +use crate::types::build_df_schema; + +use super::{CompiledTopologyNode, StreamingOperatorBlueprint}; + +// ----------------------------------------------------------------------------- +// Constants & Identifiers +// ----------------------------------------------------------------------------- + +pub const STREAM_INGESTION_NODE_NAME: &str = extension_node::STREAM_INGESTION; + +// ----------------------------------------------------------------------------- +// Logical Node Definition +// ----------------------------------------------------------------------------- + +/// Foundational ingestion point: connects to external systems and injects raw or CDC data. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct StreamIngestionNode { + pub(crate) source_identifier: TableReference, + pub(crate) source_definition: SourceTable, + pub(crate) resolved_schema: DFSchemaRef, +} + +multifield_partial_ord!(StreamIngestionNode, source_identifier, source_definition); + +impl StreamIngestionNode { + pub fn try_new( + source_identifier: TableReference, + source_definition: SourceTable, + ) -> Result { + let resolved_schema = Self::build_ingestion_schema(&source_identifier, &source_definition)?; + + Ok(Self { + source_identifier, + source_definition, + resolved_schema, + }) + } + + fn build_ingestion_schema( + identifier: &TableReference, + definition: &SourceTable, + ) -> Result { + let physical_fields: Vec<_> = definition + .schema_specs + .iter() + .filter(|col| !col.is_computed()) + .map(|col| { + ( + Some(identifier.clone()), + Arc::new(col.arrow_field().clone()), + ) + .into() + }) + .collect(); + + let base_schema = Arc::new(build_df_schema(&physical_fields)?); + + let enveloped_schema = if definition.is_updating() { + DebeziumSchemaCodec::wrap_into_envelope(&base_schema, Some(identifier.clone()))? + } else { + base_schema + }; + + add_timestamp_field(enveloped_schema, Some(identifier.clone())) + } +} + +// ----------------------------------------------------------------------------- +// DataFusion Logical Node Hooks +// ----------------------------------------------------------------------------- + +impl UserDefinedLogicalNodeCore for StreamIngestionNode { + fn name(&self) -> &str { + STREAM_INGESTION_NODE_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![] + } + + fn schema(&self) -> &DFSchemaRef { + &self.resolved_schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + write!( + f, + "StreamIngestionNode({}): Schema={}", + self.source_identifier, self.resolved_schema + ) + } + + fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { + if !inputs.is_empty() { + return plan_err!( + "StreamIngestionNode acts as a leaf boundary and cannot accept upstream inputs." + ); + } + + Ok(Self { + source_identifier: self.source_identifier.clone(), + source_definition: self.source_definition.clone(), + resolved_schema: self.resolved_schema.clone(), + }) + } +} + +// ----------------------------------------------------------------------------- +// Core Execution Blueprint Implementation +// ----------------------------------------------------------------------------- + +impl StreamingOperatorBlueprint for StreamIngestionNode { + fn operator_identity(&self) -> Option { + Some(NamedNode::Source(self.source_identifier.clone())) + } + + fn compile_to_graph_node( + &self, + compiler_context: &Planner, + node_id_sequence: usize, + upstream_schemas: Vec, + ) -> Result { + if !upstream_schemas.is_empty() { + return plan_err!( + "Topology Violation: StreamIngestionNode is a source origin and cannot process upstream routing edges." + ); + } + + let sql_source = self.source_definition.as_sql_source()?; + let connector_payload = sql_source.source.config.encode_to_vec(); + let operator_description = sql_source.source.config.description.clone(); + + let execution_unit = LogicalNode::single( + node_id_sequence as u32, + format!("source_{}_{node_id_sequence}", self.source_identifier), + OperatorName::ConnectorSource, + connector_payload, + operator_description, + compiler_context.default_parallelism(), + ); + + Ok(CompiledTopologyNode::new(execution_unit, vec![])) + } + + fn yielded_schema(&self) -> FsSchema { + FsSchema::from_schema_keys(Arc::new(self.resolved_schema.as_ref().into()), vec![]) + .expect("Fatal: Failed to generate output schema for stream ingestion") + } +} diff --git a/src/streaming_planner/src/logical_node/timestamp_append.rs b/src/streaming_planner/src/logical_node/timestamp_append.rs new file mode 100644 index 00000000..14fb247d --- /dev/null +++ b/src/streaming_planner/src/logical_node/timestamp_append.rs @@ -0,0 +1,121 @@ +// 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::fmt::Formatter; + +use datafusion::common::{DFSchemaRef, Result, TableReference, internal_err}; +use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; + +use crate::multifield_partial_ord; +use crate::common::constants::extension_node; +use crate::schema::utils::{add_timestamp_field, has_timestamp_field}; + +// ----------------------------------------------------------------------------- +// Constants & Identifiers +// ----------------------------------------------------------------------------- + +pub const TIMESTAMP_INJECTOR_NODE_NAME: &str = extension_node::SYSTEM_TIMESTAMP_INJECTOR; + +// ----------------------------------------------------------------------------- +// Logical Node Definition +// ----------------------------------------------------------------------------- + +/// Injects the mandatory system `_timestamp` field into the upstream streaming schema. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct SystemTimestampInjectorNode { + pub(crate) upstream_plan: LogicalPlan, + pub(crate) target_qualifier: Option, + pub(crate) resolved_schema: DFSchemaRef, +} + +multifield_partial_ord!(SystemTimestampInjectorNode, upstream_plan, target_qualifier); + +impl SystemTimestampInjectorNode { + pub fn try_new( + upstream_plan: LogicalPlan, + target_qualifier: Option, + ) -> Result { + let upstream_schema = upstream_plan.schema(); + + if has_timestamp_field(upstream_schema) { + return internal_err!( + "Topology Violation: Attempted to inject a system timestamp into an upstream plan \ + that already contains one. \ + \nPlan:\n {:?} \nSchema:\n {:?}", + upstream_plan, + upstream_schema + ); + } + + let resolved_schema = + add_timestamp_field(upstream_schema.clone(), target_qualifier.clone())?; + + Ok(Self { + upstream_plan, + target_qualifier, + resolved_schema, + }) + } +} + +// ----------------------------------------------------------------------------- +// DataFusion Logical Node Hooks +// ----------------------------------------------------------------------------- + +impl UserDefinedLogicalNodeCore for SystemTimestampInjectorNode { + fn name(&self) -> &str { + TIMESTAMP_INJECTOR_NODE_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.upstream_plan] + } + + fn schema(&self) -> &DFSchemaRef { + &self.resolved_schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + let field_names = self + .resolved_schema + .fields() + .iter() + .map(|field| field.name().to_string()) + .collect::>() + .join(", "); + + write!( + f, + "SystemTimestampInjector(Qualifier={:?}): [{}]", + self.target_qualifier, field_names + ) + } + + fn with_exprs_and_inputs( + &self, + _exprs: Vec, + mut inputs: Vec, + ) -> Result { + if inputs.len() != 1 { + return internal_err!( + "SystemTimestampInjectorNode requires exactly 1 upstream logical plan, but received {}", + inputs.len() + ); + } + + Self::try_new(inputs.remove(0), self.target_qualifier.clone()) + } +} diff --git a/src/streaming_planner/src/logical_node/updating_aggregate.rs b/src/streaming_planner/src/logical_node/updating_aggregate.rs new file mode 100644 index 00000000..0343c370 --- /dev/null +++ b/src/streaming_planner/src/logical_node/updating_aggregate.rs @@ -0,0 +1,245 @@ +// 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::sync::Arc; +use std::time::Duration; + +use datafusion::common::{DFSchemaRef, Result, TableReference, ToDFSchema, internal_err, plan_err}; +use datafusion::logical_expr::expr::ScalarFunction; +use datafusion::logical_expr::{ + Expr, Extension, LogicalPlan, UserDefinedLogicalNodeCore, col, lit, +}; +use datafusion::prelude::named_struct; +use datafusion::scalar::ScalarValue; +use datafusion_proto::physical_plan::AsExecutionPlan; +use datafusion_proto::protobuf::PhysicalPlanNode; +use prost::Message; +use protocol::function_stream_graph::UpdatingAggregateOperator; + +use crate::common::constants::{extension_node, proto_operator_name, updating_state_field}; +use crate::common::{FsSchema, FsSchemaRef}; +use crate::functions::multi_hash; +use crate::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; +use crate::logical_node::{ + CompiledTopologyNode, IsRetractExtension, StreamingOperatorBlueprint, +}; +use crate::logical_planner::planner::{NamedNode, Planner}; +use crate::physical::StreamingExtensionCodec; + +// ----------------------------------------------------------------------------- +// Constants & Configuration +// ----------------------------------------------------------------------------- + +pub const CONTINUOUS_AGGREGATE_NODE_NAME: &str = extension_node::CONTINUOUS_AGGREGATE; + +const DEFAULT_FLUSH_INTERVAL_MICROS: u64 = 10_000_000; + +const STATIC_HASH_SIZE_BYTES: i32 = 16; + +// ----------------------------------------------------------------------------- +// Logical Node Definition +// ----------------------------------------------------------------------------- + +/// Stateful continuous aggregation: running aggregates with updating / retraction semantics. +#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd)] +pub struct ContinuousAggregateNode { + pub(crate) base_aggregate_plan: LogicalPlan, + pub(crate) partition_key_indices: Vec, + pub(crate) retract_injected_plan: LogicalPlan, + pub(crate) namespace_qualifier: Option, + pub(crate) state_retention_ttl: Duration, +} + +impl ContinuousAggregateNode { + pub fn try_new( + base_aggregate_plan: LogicalPlan, + partition_key_indices: Vec, + namespace_qualifier: Option, + state_retention_ttl: Duration, + ) -> Result { + let retract_injected_plan = LogicalPlan::Extension(Extension { + node: Arc::new(IsRetractExtension::new( + base_aggregate_plan.clone(), + namespace_qualifier.clone(), + )), + }); + + Ok(Self { + base_aggregate_plan, + partition_key_indices, + retract_injected_plan, + namespace_qualifier, + state_retention_ttl, + }) + } + + fn construct_state_metadata_expr(&self, upstream_schema: &FsSchemaRef) -> Expr { + let routing_keys: Vec = self + .partition_key_indices + .iter() + .map(|&idx| col(upstream_schema.schema.field(idx).name())) + .collect(); + + let state_id_hash = if routing_keys.is_empty() { + Expr::Literal( + ScalarValue::FixedSizeBinary( + STATIC_HASH_SIZE_BYTES, + Some(vec![0; STATIC_HASH_SIZE_BYTES as usize]), + ), + None, + ) + } else { + Expr::ScalarFunction(ScalarFunction { + func: multi_hash(), + args: routing_keys, + }) + }; + + named_struct(vec![ + lit(updating_state_field::IS_RETRACT), + lit(false), + lit(updating_state_field::ID), + state_id_hash, + ]) + } + + fn compile_operator_config( + &self, + planner: &Planner, + upstream_schema: &FsSchemaRef, + ) -> Result { + let upstream_df_schema = upstream_schema.schema.clone().to_dfschema()?; + + let physical_agg_plan = planner.sync_plan(&self.base_aggregate_plan)?; + let compiled_agg_payload = PhysicalPlanNode::try_from_physical_plan( + physical_agg_plan, + &StreamingExtensionCodec::default(), + )? + .encode_to_vec(); + + let meta_expr = self.construct_state_metadata_expr(upstream_schema); + let compiled_meta_expr = + planner.serialize_as_physical_expr(&meta_expr, &upstream_df_schema)?; + + Ok(UpdatingAggregateOperator { + name: proto_operator_name::UPDATING_AGGREGATE.to_string(), + input_schema: Some((**upstream_schema).clone().into()), + final_schema: Some(self.yielded_schema().into()), + aggregate_exec: compiled_agg_payload, + metadata_expr: compiled_meta_expr, + flush_interval_micros: DEFAULT_FLUSH_INTERVAL_MICROS, + ttl_micros: self.state_retention_ttl.as_micros() as u64, + }) + } +} + +// ----------------------------------------------------------------------------- +// DataFusion Logical Node Hooks +// ----------------------------------------------------------------------------- + +impl UserDefinedLogicalNodeCore for ContinuousAggregateNode { + fn name(&self) -> &str { + CONTINUOUS_AGGREGATE_NODE_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.base_aggregate_plan] + } + + fn schema(&self) -> &DFSchemaRef { + self.retract_injected_plan.schema() + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!( + f, + "ContinuousAggregateNode(TTL={:?})", + self.state_retention_ttl + ) + } + + fn with_exprs_and_inputs( + &self, + _exprs: Vec, + mut inputs: Vec, + ) -> Result { + if inputs.len() != 1 { + return internal_err!( + "ContinuousAggregateNode requires exactly 1 upstream input, got {}", + inputs.len() + ); + } + + Self::try_new( + inputs.remove(0), + self.partition_key_indices.clone(), + self.namespace_qualifier.clone(), + self.state_retention_ttl, + ) + } +} + +// ----------------------------------------------------------------------------- +// Core Execution Blueprint Implementation +// ----------------------------------------------------------------------------- + +impl StreamingOperatorBlueprint for ContinuousAggregateNode { + fn operator_identity(&self) -> Option { + None + } + + fn compile_to_graph_node( + &self, + planner: &Planner, + node_index: usize, + mut upstream_schemas: Vec, + ) -> Result { + if upstream_schemas.len() != 1 { + return plan_err!( + "Topology Violation: ContinuousAggregateNode requires exactly 1 upstream input, received {}", + upstream_schemas.len() + ); + } + + let upstream_schema = upstream_schemas.remove(0); + + let operator_config = self.compile_operator_config(planner, &upstream_schema)?; + + let parallelism = planner.keyed_aggregate_parallelism(); + + let logical_node = LogicalNode::single( + node_index as u32, + format!("updating_aggregate_{node_index}"), + OperatorName::UpdatingAggregate, + operator_config.encode_to_vec(), + proto_operator_name::UPDATING_AGGREGATE.to_string(), + parallelism, + ); + + let shuffle_edge = + LogicalEdge::project_all(LogicalEdgeType::Shuffle, (*upstream_schema).clone()); + + Ok(CompiledTopologyNode { + execution_unit: logical_node, + routing_edges: vec![shuffle_edge], + }) + } + + fn yielded_schema(&self) -> FsSchema { + FsSchema::from_schema_unkeyed(Arc::new(self.schema().as_ref().into())) + .expect("Fatal: Failed to generate unkeyed output schema for continuous aggregate") + } +} diff --git a/src/streaming_planner/src/logical_node/watermark_node.rs b/src/streaming_planner/src/logical_node/watermark_node.rs new file mode 100644 index 00000000..a8802a7b --- /dev/null +++ b/src/streaming_planner/src/logical_node/watermark_node.rs @@ -0,0 +1,229 @@ +// 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::fmt::Formatter; +use std::sync::Arc; + +use datafusion::common::{DFSchemaRef, Result, TableReference, internal_err, plan_err}; +use datafusion::error::DataFusionError; +use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec; +use datafusion_proto::physical_plan::to_proto::serialize_physical_expr; +use prost::Message; +use protocol::function_stream_graph::ExpressionWatermarkConfig; + +use crate::multifield_partial_ord; +use crate::common::constants::{extension_node, runtime_operator_kind}; +use crate::common::{FsSchema, FsSchemaRef}; +use crate::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; +use crate::logical_node::{CompiledTopologyNode, StreamingOperatorBlueprint}; +use crate::logical_planner::planner::{NamedNode, Planner}; +use crate::schema::utils::add_timestamp_field; +use crate::types::TIMESTAMP_FIELD; + +// ----------------------------------------------------------------------------- +// Constants & Identifiers +// ----------------------------------------------------------------------------- + +pub const EVENT_TIME_WATERMARK_NODE_NAME: &str = extension_node::EVENT_TIME_WATERMARK; + +const DEFAULT_WATERMARK_EMISSION_PERIOD_MICROS: u64 = 1_000_000; + +// ----------------------------------------------------------------------------- +// Logical Node Definition +// ----------------------------------------------------------------------------- + +/// Event-time watermark from a user strategy; drives time progress in stateful operators. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct EventTimeWatermarkNode { + pub(crate) upstream_plan: LogicalPlan, + pub(crate) namespace_qualifier: TableReference, + pub(crate) watermark_strategy_expr: Expr, + pub(crate) resolved_schema: DFSchemaRef, + pub(crate) internal_timestamp_offset: usize, +} + +multifield_partial_ord!( + EventTimeWatermarkNode, + upstream_plan, + namespace_qualifier, + watermark_strategy_expr, + internal_timestamp_offset +); + +impl EventTimeWatermarkNode { + pub fn try_new( + upstream_plan: LogicalPlan, + namespace_qualifier: TableReference, + watermark_strategy_expr: Expr, + ) -> Result { + let resolved_schema = add_timestamp_field( + upstream_plan.schema().clone(), + Some(namespace_qualifier.clone()), + )?; + + let internal_timestamp_offset = resolved_schema + .index_of_column_by_name(None, TIMESTAMP_FIELD) + .ok_or_else(|| { + DataFusionError::Plan(format!( + "Fatal: Failed to resolve mandatory temporal column '{}'", + TIMESTAMP_FIELD + )) + })?; + + Ok(Self { + upstream_plan, + namespace_qualifier, + watermark_strategy_expr, + resolved_schema, + internal_timestamp_offset, + }) + } + + pub fn generate_fs_schema(&self) -> FsSchema { + FsSchema::new_unkeyed( + Arc::new(self.resolved_schema.as_ref().into()), + self.internal_timestamp_offset, + ) + } + + fn compile_operator_config(&self, planner: &Planner) -> Result { + let physical_expr = + planner.create_physical_expr(&self.watermark_strategy_expr, &self.resolved_schema)?; + + let serialized_expr = + serialize_physical_expr(&physical_expr, &DefaultPhysicalExtensionCodec {})?; + + Ok(ExpressionWatermarkConfig { + period_micros: DEFAULT_WATERMARK_EMISSION_PERIOD_MICROS, + idle_time_micros: None, + expression: serialized_expr.encode_to_vec(), + input_schema: Some(self.generate_fs_schema().into()), + }) + } +} + +// ----------------------------------------------------------------------------- +// DataFusion Logical Node Hooks +// ----------------------------------------------------------------------------- + +impl UserDefinedLogicalNodeCore for EventTimeWatermarkNode { + fn name(&self) -> &str { + EVENT_TIME_WATERMARK_NODE_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.upstream_plan] + } + + fn schema(&self) -> &DFSchemaRef { + &self.resolved_schema + } + + fn expressions(&self) -> Vec { + vec![self.watermark_strategy_expr.clone()] + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + write!( + f, + "EventTimeWatermarkNode({}): Schema={}", + self.namespace_qualifier, self.resolved_schema + ) + } + + fn with_exprs_and_inputs( + &self, + mut exprs: Vec, + mut inputs: Vec, + ) -> Result { + if inputs.len() != 1 { + return internal_err!( + "EventTimeWatermarkNode requires exactly 1 upstream logical plan, but received {}", + inputs.len() + ); + } + if exprs.len() != 1 { + return internal_err!( + "EventTimeWatermarkNode requires exactly 1 watermark strategy expression, but received {}", + exprs.len() + ); + } + + let internal_timestamp_offset = self + .resolved_schema + .index_of_column_by_name(Some(&self.namespace_qualifier), TIMESTAMP_FIELD) + .ok_or_else(|| { + DataFusionError::Plan(format!( + "Optimizer Error: Lost tracking of temporal column '{}'", + TIMESTAMP_FIELD + )) + })?; + + Ok(Self { + upstream_plan: inputs.remove(0), + namespace_qualifier: self.namespace_qualifier.clone(), + watermark_strategy_expr: exprs.remove(0), + resolved_schema: self.resolved_schema.clone(), + internal_timestamp_offset, + }) + } +} + +// ----------------------------------------------------------------------------- +// Core Execution Blueprint Implementation +// ----------------------------------------------------------------------------- + +impl StreamingOperatorBlueprint for EventTimeWatermarkNode { + fn operator_identity(&self) -> Option { + Some(NamedNode::Watermark(self.namespace_qualifier.clone())) + } + + fn compile_to_graph_node( + &self, + planner: &Planner, + node_index: usize, + mut upstream_schemas: Vec, + ) -> Result { + if upstream_schemas.len() != 1 { + return plan_err!( + "Topology Violation: EventTimeWatermarkNode requires exactly 1 upstream input, received {}", + upstream_schemas.len() + ); + } + + let operator_config = self.compile_operator_config(planner)?; + + let execution_unit = LogicalNode::single( + node_index as u32, + format!("watermark_{node_index}"), + OperatorName::ExpressionWatermark, + operator_config.encode_to_vec(), + runtime_operator_kind::WATERMARK_GENERATOR.to_string(), + planner.default_parallelism(), + ); + + let incoming_edge = LogicalEdge::project_all( + LogicalEdgeType::Forward, + (*upstream_schemas.remove(0)).clone(), + ); + + Ok(CompiledTopologyNode { + execution_unit, + routing_edges: vec![incoming_edge], + }) + } + + fn yielded_schema(&self) -> FsSchema { + self.generate_fs_schema() + } +} diff --git a/src/streaming_planner/src/logical_node/windows_function.rs b/src/streaming_planner/src/logical_node/windows_function.rs new file mode 100644 index 00000000..5fdfb9ce --- /dev/null +++ b/src/streaming_planner/src/logical_node/windows_function.rs @@ -0,0 +1,191 @@ +// 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::fmt::Formatter; +use std::sync::Arc; + +use datafusion::common::{Column, DFSchema, DFSchemaRef, Result, internal_err, plan_err}; +use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec; +use datafusion_proto::physical_plan::to_proto::serialize_physical_expr; +use datafusion_proto::{physical_plan::AsExecutionPlan, protobuf::PhysicalPlanNode}; +use prost::Message; +use protocol::function_stream_graph::WindowFunctionOperator; + +use crate::common::constants::{extension_node, proto_operator_name, runtime_operator_kind}; +use crate::common::{FsSchema, FsSchemaRef}; +use crate::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; +use crate::logical_planner::planner::{NamedNode, Planner}; +use crate::physical::StreamingExtensionCodec; +use crate::types::TIMESTAMP_FIELD; + +use super::{CompiledTopologyNode, StreamingOperatorBlueprint}; + +// ----------------------------------------------------------------------------- +// Constants & Identifiers +// ----------------------------------------------------------------------------- + +pub const STREAMING_WINDOW_NODE_NAME: &str = extension_node::STREAMING_WINDOW_FUNCTION; + +// ----------------------------------------------------------------------------- +// Logical Node Definition +// ----------------------------------------------------------------------------- + +/// Stateful streaming window: temporal binning plus underlying window evaluation plan. +#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd)] +pub struct StreamingWindowFunctionNode { + pub(crate) underlying_evaluation_plan: LogicalPlan, + pub(crate) partition_key_indices: Vec, +} + +impl StreamingWindowFunctionNode { + pub fn new(underlying_evaluation_plan: LogicalPlan, partition_key_indices: Vec) -> Self { + Self { + underlying_evaluation_plan, + partition_key_indices, + } + } + + fn compile_temporal_binning_function( + &self, + planner: &Planner, + input_df_schema: &DFSchema, + ) -> Result> { + let timestamp_column = Expr::Column(Column::new_unqualified(TIMESTAMP_FIELD.to_string())); + + let physical_binning_expr = + planner.create_physical_expr(×tamp_column, input_df_schema)?; + + let serialized_expr = + serialize_physical_expr(&physical_binning_expr, &DefaultPhysicalExtensionCodec {})?; + + Ok(serialized_expr.encode_to_vec()) + } + + fn compile_physical_evaluation_plan(&self, planner: &Planner) -> Result> { + let physical_window_plan = planner.sync_plan(&self.underlying_evaluation_plan)?; + + let proto_plan_node = PhysicalPlanNode::try_from_physical_plan( + physical_window_plan, + &StreamingExtensionCodec::default(), + )?; + + Ok(proto_plan_node.encode_to_vec()) + } +} + +// ----------------------------------------------------------------------------- +// DataFusion Logical Node Hooks +// ----------------------------------------------------------------------------- + +impl UserDefinedLogicalNodeCore for StreamingWindowFunctionNode { + fn name(&self) -> &str { + STREAMING_WINDOW_NODE_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.underlying_evaluation_plan] + } + + fn schema(&self) -> &DFSchemaRef { + self.underlying_evaluation_plan.schema() + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + write!(f, "StreamingWindowFunction: Schema={}", self.schema()) + } + + fn with_exprs_and_inputs( + &self, + _exprs: Vec, + mut inputs: Vec, + ) -> Result { + if inputs.len() != 1 { + return internal_err!( + "StreamingWindowFunctionNode requires exactly 1 upstream input, got {}", + inputs.len() + ); + } + + Ok(Self::new( + inputs.remove(0), + self.partition_key_indices.clone(), + )) + } +} + +// ----------------------------------------------------------------------------- +// Core Execution Blueprint Implementation +// ----------------------------------------------------------------------------- + +impl StreamingOperatorBlueprint for StreamingWindowFunctionNode { + fn operator_identity(&self) -> Option { + None + } + + fn compile_to_graph_node( + &self, + planner: &Planner, + node_index: usize, + mut input_schemas: Vec, + ) -> Result { + if input_schemas.len() != 1 { + return plan_err!( + "Topology Violation: StreamingWindowFunctionNode requires exactly 1 upstream input schema, received {}", + input_schemas.len() + ); + } + + let input_schema = input_schemas.remove(0); + + let input_df_schema = DFSchema::try_from(input_schema.schema.as_ref().clone())?; + + let binning_payload = self.compile_temporal_binning_function(planner, &input_df_schema)?; + let evaluation_plan_payload = self.compile_physical_evaluation_plan(planner)?; + + let operator_config = WindowFunctionOperator { + name: proto_operator_name::WINDOW_FUNCTION.to_string(), + input_schema: Some(input_schema.as_ref().clone().into()), + binning_function: binning_payload, + window_function_plan: evaluation_plan_payload, + }; + + let parallelism = planner.keyed_aggregate_parallelism(); + + let logical_node = LogicalNode::single( + node_index as u32, + format!("window_function_{node_index}"), + OperatorName::WindowFunction, + operator_config.encode_to_vec(), + runtime_operator_kind::STREAMING_WINDOW_EVALUATOR.to_string(), + parallelism, + ); + + let routing_edge = + LogicalEdge::project_all(LogicalEdgeType::Shuffle, (*input_schema).clone()); + + Ok(CompiledTopologyNode { + execution_unit: logical_node, + routing_edges: vec![routing_edge], + }) + } + + fn yielded_schema(&self) -> FsSchema { + FsSchema::from_schema_unkeyed(Arc::new(self.schema().as_ref().clone().into())).expect( + "Fatal: Failed to generate unkeyed output schema for StreamingWindowFunctionNode", + ) + } +} diff --git a/src/streaming_planner/src/logical_planner/mod.rs b/src/streaming_planner/src/logical_planner/mod.rs new file mode 100644 index 00000000..d04ccce3 --- /dev/null +++ b/src/streaming_planner/src/logical_planner/mod.rs @@ -0,0 +1,16 @@ +// 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. + +pub mod optimizers; + +pub mod streaming_planner; +pub use streaming_planner as planner; diff --git a/src/streaming_planner/src/logical_planner/optimizers/chaining.rs b/src/streaming_planner/src/logical_planner/optimizers/chaining.rs new file mode 100644 index 00000000..60b36115 --- /dev/null +++ b/src/streaming_planner/src/logical_planner/optimizers/chaining.rs @@ -0,0 +1,200 @@ +// 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 petgraph::prelude::*; +use petgraph::visit::NodeIndexable; +use tracing::debug; + +use crate::logical_node::logical::{LogicalEdgeType, LogicalGraph, Optimizer}; + +pub struct ChainingOptimizer {} + +impl Optimizer for ChainingOptimizer { + fn optimize_once(&self, plan: &mut LogicalGraph) -> bool { + let mut match_found = None; + + for node_idx in plan.node_indices() { + let mut outgoing = plan.edges_directed(node_idx, Outgoing); + let first_out = outgoing.next(); + if first_out.is_none() || outgoing.next().is_some() { + continue; + } + let edge = first_out.unwrap(); + + if edge.weight().edge_type != LogicalEdgeType::Forward { + continue; + } + + let target_idx = edge.target(); + + let mut incoming = plan.edges_directed(target_idx, Incoming); + let first_in = incoming.next(); + if first_in.is_none() || incoming.next().is_some() { + continue; + } + + let source_node = plan.node_weight(node_idx).expect("Source node missing"); + let target_node = plan.node_weight(target_idx).expect("Target node missing"); + + let parallelism_ok = source_node.parallelism == target_node.parallelism + || target_node + .operator_chain + .is_parallelism_upstream_expandable(); + + if source_node.operator_chain.is_source() + || target_node.operator_chain.is_sink() + || !parallelism_ok + { + continue; + } + + match_found = Some((node_idx, target_idx, edge.id())); + break; + } + + if let Some((source_idx, target_idx, edge_id)) = match_found { + let edge_weight = plan.remove_edge(edge_id).expect("Edge should exist"); + + let target_outgoing: Vec<_> = plan + .edges_directed(target_idx, Outgoing) + .map(|e| (e.id(), e.target())) + .collect(); + + for (e_id, next_target_idx) in target_outgoing { + let weight = plan.remove_edge(e_id).expect("Outgoing edge missing"); + plan.add_edge(source_idx, next_target_idx, weight); + } + + let is_source_last = source_idx.index() == plan.node_bound() - 1; + + let target_node = plan + .remove_node(target_idx) + .expect("Target node should exist"); + + let actual_source_idx = if is_source_last { + target_idx + } else { + source_idx + }; + + let source_node = plan + .node_weight_mut(actual_source_idx) + .expect("Source node missing"); + + debug!( + "Chaining Optimizer: Fusing '{}' -> '{}'", + source_node.description, target_node.description + ); + + source_node.description = + format!("{} -> {}", source_node.description, target_node.description); + + source_node.parallelism = source_node.parallelism.max(target_node.parallelism); + + source_node + .operator_chain + .operators + .extend(target_node.operator_chain.operators); + source_node.operator_chain.edges.push(edge_weight.schema); + + return true; + } + + false + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit}; + + use crate::common::FsSchema; + use crate::logical_node::logical::{ + LogicalEdge, LogicalEdgeType, LogicalGraph, LogicalNode, OperatorName, Optimizer, + }; + + use super::ChainingOptimizer; + + fn forward_edge() -> LogicalEdge { + let s = Arc::new(Schema::new(vec![Field::new( + "_timestamp", + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + )])); + LogicalEdge::new(LogicalEdgeType::Forward, FsSchema::new_unkeyed(s, 0)) + } + + fn proj_node(id: u32, label: &str) -> LogicalNode { + LogicalNode::single( + id, + format!("op_{label}"), + OperatorName::Projection, + vec![], + label.to_string(), + 1, + ) + } + + fn source_node() -> LogicalNode { + LogicalNode::single( + 0, + "src".into(), + OperatorName::ConnectorSource, + vec![], + "source".into(), + 1, + ) + } + + /// Window aggregate at higher default parallelism may forward into projection @ 1: still fuse + /// so each branch does not reserve a separate global state-memory block for the same sub-chain. + #[test] + fn fusion_stateful_high_parallelism_into_expandable_low() { + let mut g = LogicalGraph::new(); + let n0 = g.add_node(source_node()); + let n1 = g.add_node(proj_node(1, "tumble")); + let n2 = g.add_node(proj_node(2, "proj")); + let n1w = g.node_weight_mut(n1).unwrap(); + n1w.parallelism = 8; + let e = forward_edge(); + g.add_edge(n0, n1, e.clone()); + g.add_edge(n1, n2, e); + + let changed = ChainingOptimizer {}.optimize_once(&mut g); + assert!(changed); + assert_eq!(g.node_count(), 2); + let fused = g + .node_weights() + .find(|n| n.description.contains("->")) + .unwrap(); + assert_eq!(fused.parallelism, 8); + assert_eq!(fused.operator_chain.len(), 2); + } + + /// Regression: upstream at last `NodeIndex` + remove non-last downstream swaps indices. + #[test] + fn fusion_remaps_when_upstream_was_last_node_index() { + let mut g = LogicalGraph::new(); + let n0 = g.add_node(source_node()); + let n1 = g.add_node(proj_node(1, "downstream")); + let n2 = g.add_node(proj_node(2, "upstream_last_index")); + let e = forward_edge(); + g.add_edge(n0, n2, e.clone()); + g.add_edge(n2, n1, e); + + let changed = ChainingOptimizer {}.optimize_once(&mut g); + assert!(changed); + assert_eq!(g.node_count(), 2); + } +} diff --git a/src/streaming_planner/src/logical_planner/optimizers/mod.rs b/src/streaming_planner/src/logical_planner/optimizers/mod.rs new file mode 100644 index 00000000..c7981313 --- /dev/null +++ b/src/streaming_planner/src/logical_planner/optimizers/mod.rs @@ -0,0 +1,20 @@ +// 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. + +//! Logical planner optimizers: graph-level chaining ([`ChainingOptimizer`]) and +//! DataFusion SQL logical-plan rules ([`produce_optimized_plan`]). + +mod chaining; +mod optimized_plan; + +pub use chaining::ChainingOptimizer; +pub use optimized_plan::produce_optimized_plan; diff --git a/src/streaming_planner/src/logical_planner/optimizers/optimized_plan.rs b/src/streaming_planner/src/logical_planner/optimizers/optimized_plan.rs new file mode 100644 index 00000000..df380fe1 --- /dev/null +++ b/src/streaming_planner/src/logical_planner/optimizers/optimized_plan.rs @@ -0,0 +1,95 @@ +// 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::sync::Arc; + +use datafusion::common::Result; +use datafusion::common::config::ConfigOptions; +use datafusion::logical_expr::LogicalPlan; +use datafusion::optimizer::OptimizerContext; +use datafusion::optimizer::OptimizerRule; +use datafusion::optimizer::common_subexpr_eliminate::CommonSubexprEliminate; +use datafusion::optimizer::decorrelate_lateral_join::DecorrelateLateralJoin; +use datafusion::optimizer::decorrelate_predicate_subquery::DecorrelatePredicateSubquery; +use datafusion::optimizer::eliminate_cross_join::EliminateCrossJoin; +use datafusion::optimizer::eliminate_duplicated_expr::EliminateDuplicatedExpr; +use datafusion::optimizer::eliminate_filter::EliminateFilter; +use datafusion::optimizer::eliminate_group_by_constant::EliminateGroupByConstant; +use datafusion::optimizer::eliminate_join::EliminateJoin; +use datafusion::optimizer::eliminate_limit::EliminateLimit; +use datafusion::optimizer::eliminate_nested_union::EliminateNestedUnion; +use datafusion::optimizer::eliminate_one_union::EliminateOneUnion; +use datafusion::optimizer::eliminate_outer_join::EliminateOuterJoin; +use datafusion::optimizer::extract_equijoin_predicate::ExtractEquijoinPredicate; +use datafusion::optimizer::filter_null_join_keys::FilterNullJoinKeys; +use datafusion::optimizer::optimizer::Optimizer; +use datafusion::optimizer::propagate_empty_relation::PropagateEmptyRelation; +use datafusion::optimizer::push_down_filter::PushDownFilter; +use datafusion::optimizer::push_down_limit::PushDownLimit; +use datafusion::optimizer::replace_distinct_aggregate::ReplaceDistinctWithAggregate; +use datafusion::optimizer::scalar_subquery_to_join::ScalarSubqueryToJoin; +use datafusion::optimizer::simplify_expressions::SimplifyExpressions; +use datafusion::sql::planner::SqlToRel; +use datafusion::sql::sqlparser::ast::Statement; + +use crate::schema::StreamSchemaProvider; + +/// Converts a SQL statement into an optimized DataFusion logical plan. +/// +/// Applies the DataFusion analyzer followed by a curated set of optimizer rules +/// suitable for streaming SQL (some rules like OptimizeProjections are excluded +/// because they can drop event-time calculation fields). +pub fn produce_optimized_plan( + statement: &Statement, + schema_provider: &StreamSchemaProvider, +) -> Result { + let sql_to_rel = SqlToRel::new(schema_provider); + let plan = sql_to_rel.sql_statement_to_plan(statement.clone())?; + + let analyzed_plan = schema_provider.analyzer.execute_and_check( + plan, + &ConfigOptions::default(), + |_plan, _rule| {}, + )?; + + let rules: Vec> = vec![ + Arc::new(EliminateNestedUnion::new()), + Arc::new(SimplifyExpressions::new()), + Arc::new(ReplaceDistinctWithAggregate::new()), + Arc::new(EliminateJoin::new()), + Arc::new(DecorrelatePredicateSubquery::new()), + Arc::new(ScalarSubqueryToJoin::new()), + Arc::new(DecorrelateLateralJoin::new()), + Arc::new(ExtractEquijoinPredicate::new()), + Arc::new(EliminateDuplicatedExpr::new()), + Arc::new(EliminateFilter::new()), + Arc::new(EliminateCrossJoin::new()), + Arc::new(EliminateLimit::new()), + Arc::new(PropagateEmptyRelation::new()), + Arc::new(EliminateOneUnion::new()), + Arc::new(FilterNullJoinKeys::default()), + Arc::new(EliminateOuterJoin::new()), + Arc::new(PushDownLimit::new()), + Arc::new(PushDownFilter::new()), + Arc::new(EliminateGroupByConstant::new()), + Arc::new(CommonSubexprEliminate::new()), + ]; + + let optimizer = Optimizer::with_rules(rules); + let optimized = optimizer.optimize( + analyzed_plan, + &OptimizerContext::default(), + |_plan, _rule| {}, + )?; + + Ok(optimized) +} diff --git a/src/streaming_planner/src/logical_planner/streaming_planner.rs b/src/streaming_planner/src/logical_planner/streaming_planner.rs new file mode 100644 index 00000000..ef7c7d38 --- /dev/null +++ b/src/streaming_planner/src/logical_planner/streaming_planner.rs @@ -0,0 +1,435 @@ +// 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::Arc; +use std::thread; +use std::time::Duration; + +use datafusion::arrow::datatypes::IntervalMonthDayNanoType; +use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor}; +use datafusion::common::{ + DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, Spans, plan_err, +}; +use datafusion::execution::context::SessionState; +use datafusion::execution::runtime_env::RuntimeEnvBuilder; +use datafusion::functions::datetime::date_bin; +use datafusion::logical_expr::{Expr, Extension, LogicalPlan, UserDefinedLogicalNode}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_planner::{DefaultPhysicalPlanner, ExtensionPlanner, PhysicalPlanner}; +use datafusion_proto::protobuf::{PhysicalExprNode, PhysicalPlanNode}; +use datafusion_proto::{ + physical_plan::AsExecutionPlan, + protobuf::{AggregateMode, physical_plan_node::PhysicalPlanType}, +}; +use petgraph::graph::{DiGraph, NodeIndex}; +use prost::Message; +use tokio::runtime::Builder; +use tokio::sync::oneshot; + +use async_trait::async_trait; +use datafusion_common::TableReference; +use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec; +use datafusion_proto::physical_plan::to_proto::serialize_physical_expr; + +use crate::common::constants::sql_planning_default; +use crate::common::{FsSchema, FsSchemaRef}; +use crate::logical_node::debezium::{ + PACK_NODE_NAME, UNROLL_NODE_NAME, UnrollDebeziumPayloadNode, +}; +use crate::logical_node::key_calculation::KeyExtractionNode; +use crate::logical_node::logical::{LogicalEdge, LogicalGraph, LogicalNode}; +use crate::logical_node::{CompiledTopologyNode, StreamingOperatorBlueprint}; +use crate::physical::{ + CdcDebeziumPackExec, CdcDebeziumUnrollExec, FsMemExec, StreamingDecodingContext, + StreamingExtensionCodec, +}; +use crate::schema::StreamSchemaProvider; +use crate::schema::utils::add_timestamp_field_arrow; + +pub struct SplitPlanOutput { + pub(crate) partial_aggregation_plan: PhysicalPlanNode, + pub(crate) partial_schema: FsSchema, + pub(crate) finish_plan: PhysicalPlanNode, +} +#[derive(Eq, Hash, PartialEq, Debug)] +pub enum NamedNode { + Source(TableReference), + Watermark(TableReference), + RemoteTable(TableReference), + Sink(TableReference), +} + +pub struct PlanToGraphVisitor<'a> { + graph: DiGraph, + output_schemas: HashMap, + named_nodes: HashMap, + traversal: Vec>, + planner: Planner<'a>, +} + +impl<'a> PlanToGraphVisitor<'a> { + pub fn new(schema_provider: &'a StreamSchemaProvider, session_state: &'a SessionState) -> Self { + Self { + graph: Default::default(), + output_schemas: Default::default(), + named_nodes: Default::default(), + traversal: vec![], + planner: Planner::new(schema_provider, session_state), + } + } +} + +pub struct Planner<'a> { + schema_provider: &'a StreamSchemaProvider, + planner: DefaultPhysicalPlanner, + session_state: &'a SessionState, +} + +impl<'a> Planner<'a> { + #[inline] + pub fn default_parallelism(&self) -> usize { + self.schema_provider.default_parallelism() + } + + #[inline] + pub fn key_by_parallelism(&self) -> usize { + self.schema_provider.key_by_parallelism() + } + + /// Parallelism for operators that consume a keyed shuffle (non-empty partition keys). + #[inline] + pub fn keyed_aggregate_parallelism(&self) -> usize { + sql_planning_default::KEYED_AGGREGATE_DEFAULT_PARALLELISM + } + + pub fn new( + schema_provider: &'a StreamSchemaProvider, + session_state: &'a SessionState, + ) -> Self { + let planner = + DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new(FsExtensionPlanner {})]); + Self { + schema_provider, + planner, + session_state, + } + } + + pub fn sync_plan(&self, plan: &LogicalPlan) -> Result> { + let fut = self.planner.create_physical_plan(plan, self.session_state); + let (tx, mut rx) = oneshot::channel(); + thread::scope(|s| { + let builder = thread::Builder::new(); + let builder = if cfg!(debug_assertions) { + builder.stack_size(10_000_000) + } else { + builder + }; + builder + .spawn_scoped(s, move || { + let rt = Builder::new_current_thread().enable_all().build().unwrap(); + rt.block_on(async { + let plan = fut.await; + tx.send(plan).unwrap(); + }); + }) + .unwrap(); + }); + + rx.try_recv().unwrap() + } + + pub fn create_physical_expr( + &self, + expr: &Expr, + input_dfschema: &DFSchema, + ) -> Result> { + self.planner + .create_physical_expr(expr, input_dfschema, self.session_state) + } + + pub fn serialize_as_physical_expr( + &self, + expr: &Expr, + schema: &DFSchema, + ) -> Result> { + let physical = self.create_physical_expr(expr, schema)?; + let proto = serialize_physical_expr(&physical, &DefaultPhysicalExtensionCodec {})?; + Ok(proto.encode_to_vec()) + } + + pub fn split_physical_plan( + &self, + key_indices: Vec, + aggregate: &LogicalPlan, + add_timestamp_field: bool, + ) -> Result { + let physical_plan = self.sync_plan(aggregate)?; + let codec = StreamingExtensionCodec { + context: StreamingDecodingContext::Planning, + }; + let mut physical_plan_node = + PhysicalPlanNode::try_from_physical_plan(physical_plan.clone(), &codec)?; + let PhysicalPlanType::Aggregate(mut final_aggregate_proto) = physical_plan_node + .physical_plan_type + .take() + .ok_or_else(|| DataFusionError::Plan("missing physical plan type".to_string()))? + else { + return plan_err!("unexpected physical plan type"); + }; + let AggregateMode::Final = final_aggregate_proto.mode() else { + return plan_err!("unexpected physical plan type"); + }; + + let partial_aggregation_plan = *final_aggregate_proto + .input + .take() + .ok_or_else(|| DataFusionError::Plan("missing input".to_string()))?; + + let partial_aggregation_exec_plan = partial_aggregation_plan.try_into_physical_plan( + self.schema_provider, + &RuntimeEnvBuilder::new().build().unwrap(), + &codec, + )?; + + let partial_schema = partial_aggregation_exec_plan.schema(); + let final_input_table_provider = FsMemExec::new("partial".into(), partial_schema.clone()); + + final_aggregate_proto.input = Some(Box::new(PhysicalPlanNode::try_from_physical_plan( + Arc::new(final_input_table_provider), + &codec, + )?)); + + let finish_plan = PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::Aggregate(final_aggregate_proto)), + }; + + let (partial_schema, timestamp_index) = if add_timestamp_field { + ( + add_timestamp_field_arrow((*partial_schema).clone()), + partial_schema.fields().len(), + ) + } else { + (partial_schema.clone(), partial_schema.fields().len() - 1) + }; + + let partial_schema = FsSchema::new_keyed(partial_schema, timestamp_index, key_indices); + + Ok(SplitPlanOutput { + partial_aggregation_plan, + partial_schema, + finish_plan, + }) + } + + pub fn binning_function_proto( + &self, + width: Duration, + input_schema: DFSchemaRef, + ) -> Result { + let date_bin = date_bin().call(vec![ + Expr::Literal( + ScalarValue::IntervalMonthDayNano(Some(IntervalMonthDayNanoType::make_value( + 0, + 0, + width.as_nanos() as i64, + ))), + None, + ), + Expr::Column(datafusion::common::Column { + relation: None, + name: "_timestamp".into(), + spans: Spans::new(), + }), + ]); + + let binning_function = self.create_physical_expr(&date_bin, &input_schema)?; + serialize_physical_expr(&binning_function, &DefaultPhysicalExtensionCodec {}) + } +} + +struct FsExtensionPlanner {} + +#[async_trait] +impl ExtensionPlanner for FsExtensionPlanner { + async fn plan_extension( + &self, + _planner: &dyn PhysicalPlanner, + node: &dyn UserDefinedLogicalNode, + _logical_inputs: &[&LogicalPlan], + physical_inputs: &[Arc], + _session_state: &SessionState, + ) -> Result>> { + let schema = node.schema().as_ref().into(); + if let Ok::<&dyn StreamingOperatorBlueprint, _>(stream_extension) = node.try_into() + && stream_extension.is_passthrough_boundary() + { + match node.name() { + UNROLL_NODE_NAME => { + let node = node + .as_any() + .downcast_ref::() + .unwrap(); + let input = physical_inputs[0].clone(); + return Ok(Some(Arc::new(CdcDebeziumUnrollExec::try_new( + input, + node.pk_indices.clone(), + )?))); + } + PACK_NODE_NAME => { + let input = physical_inputs[0].clone(); + return Ok(Some(Arc::new(CdcDebeziumPackExec::try_new(input)?))); + } + _ => return Ok(None), + } + } + let name = if let Some(key_extension) = node.as_any().downcast_ref::() { + key_extension.operator_label.clone() + } else { + None + }; + Ok(Some(Arc::new(FsMemExec::new( + name.unwrap_or("memory".to_string()), + Arc::new(schema), + )))) + } +} + +impl PlanToGraphVisitor<'_> { + fn add_index_to_traversal(&mut self, index: NodeIndex) { + if let Some(last) = self.traversal.last_mut() { + last.push(index); + } + } + + pub fn add_plan(&mut self, plan: LogicalPlan) -> Result<()> { + self.traversal.clear(); + plan.visit(self)?; + Ok(()) + } + + pub fn into_graph(self) -> LogicalGraph { + self.graph + } + + pub fn build_extension( + &mut self, + input_nodes: Vec, + extension: &dyn StreamingOperatorBlueprint, + ) -> Result<()> { + if let Some(node_name) = extension.operator_identity() + && self.named_nodes.contains_key(&node_name) + { + return plan_err!( + "extension {:?} has already been planned, shouldn't try again.", + node_name + ); + } + + let input_schemas = input_nodes + .iter() + .map(|index| { + Ok(self + .output_schemas + .get(index) + .ok_or_else(|| DataFusionError::Plan("missing input node".to_string()))? + .clone()) + }) + .collect::>>()?; + + let CompiledTopologyNode { + execution_unit, + routing_edges, + } = extension + .compile_to_graph_node(&self.planner, self.graph.node_count(), input_schemas) + .map_err(|e| e.context(format!("planning operator {extension:?}")))?; + + let node_index = self.graph.add_node(execution_unit); + self.add_index_to_traversal(node_index); + + for (source, edge) in input_nodes.into_iter().zip(routing_edges) { + self.graph.add_edge(source, node_index, edge); + } + + self.output_schemas + .insert(node_index, extension.yielded_schema().into()); + + if let Some(node_name) = extension.operator_identity() { + self.named_nodes.insert(node_name, node_index); + } + Ok(()) + } +} + +impl TreeNodeVisitor<'_> for PlanToGraphVisitor<'_> { + type Node = LogicalPlan; + + fn f_down(&mut self, node: &Self::Node) -> Result { + let LogicalPlan::Extension(Extension { node }) = node else { + return Ok(TreeNodeRecursion::Continue); + }; + + let stream_extension: &dyn StreamingOperatorBlueprint = node + .try_into() + .map_err(|e: DataFusionError| e.context("converting extension"))?; + if stream_extension.is_passthrough_boundary() { + return Ok(TreeNodeRecursion::Continue); + } + + if let Some(name) = stream_extension.operator_identity() + && let Some(node_index) = self.named_nodes.get(&name) + { + self.add_index_to_traversal(*node_index); + return Ok(TreeNodeRecursion::Jump); + } + + if !node.inputs().is_empty() { + self.traversal.push(vec![]); + } + + Ok(TreeNodeRecursion::Continue) + } + + fn f_up(&mut self, node: &Self::Node) -> Result { + let LogicalPlan::Extension(Extension { node }) = node else { + return Ok(TreeNodeRecursion::Continue); + }; + + let stream_extension: &dyn StreamingOperatorBlueprint = node + .try_into() + .map_err(|e: DataFusionError| e.context("planning extension"))?; + + if stream_extension.is_passthrough_boundary() { + return Ok(TreeNodeRecursion::Continue); + } + + if let Some(name) = stream_extension.operator_identity() + && self.named_nodes.contains_key(&name) + { + return Ok(TreeNodeRecursion::Continue); + } + + let input_nodes = if !node.inputs().is_empty() { + self.traversal.pop().unwrap_or_default() + } else { + vec![] + }; + let stream_extension: &dyn StreamingOperatorBlueprint = node + .try_into() + .map_err(|e: DataFusionError| e.context("converting extension"))?; + self.build_extension(input_nodes, stream_extension)?; + + Ok(TreeNodeRecursion::Continue) + } +} diff --git a/src/streaming_planner/src/parse.rs b/src/streaming_planner/src/parse.rs new file mode 100644 index 00000000..b3a90fbf --- /dev/null +++ b/src/streaming_planner/src/parse.rs @@ -0,0 +1,82 @@ +// 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. + +//! FunctionStream SQL parsing (`parse_sql`). +//! +//! This module only performs lexical/syntactic parsing into sqlparser +//! [`Statement`](datafusion::sql::sqlparser::ast::Statement) values using +//! [`FunctionStreamDialect`]. Mapping those AST nodes to coordinator `Statement` +//! values is done by `classify_statement` in the `function-stream` coordinator module. +//! +//! **Data-definition / pipeline shape (supported forms in the dialect)** +//! - **`CREATE TABLE ... (cols [, WATERMARK FOR ...]) WITH (...)`** — connector-backed source DDL +//! - **`CREATE TABLE ...`** other forms (including `AS SELECT` where the dialect accepts it) +//! - **`CREATE STREAMING TABLE ... WITH (...) AS SELECT ...`** +//! - **`DROP TABLE`** / **`DROP STREAMING TABLE`** +//! - **`SHOW TABLES`**, **`SHOW STREAMING TABLE(S)`**, **`SHOW CREATE TABLE`**, **`SHOW CREATE STREAMING TABLE`** +//! +//! **`INSERT` is not supported** at the coordinator layer — use `CREATE TABLE ... AS SELECT` or +//! `CREATE STREAMING TABLE ... AS SELECT` instead (see coordinator classification). + +use datafusion::common::{plan_err, Result}; +use datafusion::error::DataFusionError; +use datafusion::sql::sqlparser::ast::Statement as DFStatement; +use datafusion::sql::sqlparser::dialect::FunctionStreamDialect; +use datafusion::sql::sqlparser::parser::Parser; + +/// Parse SQL text into zero or more dialect [`Statement`](DFStatement) nodes. +pub fn parse_sql(query: &str) -> Result> { + let trimmed = query.trim(); + if trimmed.is_empty() { + return plan_err!("Query is empty"); + } + + let dialect = FunctionStreamDialect {}; + let statements = Parser::parse_sql(&dialect, trimmed) + .map_err(|e| DataFusionError::Plan(format!("SQL parse error: {e}")))?; + + if statements.is_empty() { + return plan_err!("No SQL statements found"); + } + + Ok(statements) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_multiple_statements_ast() { + let sql = concat!( + "CREATE TABLE t1 (id INT); ", + "CREATE STREAMING TABLE sk WITH ('connector' = 'kafka') AS SELECT id FROM t1", + ); + let stmts = parse_sql(sql).unwrap(); + assert_eq!(stmts.len(), 2); + assert!(matches!(stmts[0], DFStatement::CreateTable(_))); + assert!(matches!(stmts[1], DFStatement::CreateStreamingTable { .. })); + } + + #[test] + fn test_parse_empty() { + assert!(parse_sql("").is_err()); + assert!(parse_sql(" ").is_err()); + } + + #[test] + fn test_parse_select_yields_query_ast() { + let stmts = parse_sql("SELECT 1").unwrap(); + assert_eq!(stmts.len(), 1); + assert!(matches!(stmts[0], DFStatement::Query(_))); + } +} diff --git a/src/streaming_planner/src/physical/cdc/encode.rs b/src/streaming_planner/src/physical/cdc/encode.rs new file mode 100644 index 00000000..8457724c --- /dev/null +++ b/src/streaming_planner/src/physical/cdc/encode.rs @@ -0,0 +1,342 @@ +// 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::any::Any; +use std::collections::HashMap; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use datafusion::arrow::array::AsArray; +use datafusion::arrow::array::{ + Array, BooleanArray, FixedSizeBinaryArray, PrimitiveArray, RecordBatch, StringArray, + StructArray, TimestampNanosecondBuilder, UInt32Array, UInt32Builder, +}; +use datafusion::arrow::compute::take; +use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimestampNanosecondType}; +use datafusion::common::{DataFusionError, Result}; +use datafusion::execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext}; +use datafusion::physical_plan::{DisplayAs, ExecutionPlan, PlanProperties}; +use futures::{StreamExt, ready, stream::Stream}; + +use crate::common::constants::{cdc, debezium_op_short, physical_plan_node_name}; +use crate::common::{TIMESTAMP_FIELD, UPDATING_META_FIELD}; +use crate::physical::source_exec::make_stream_properties; + +// ============================================================================ +// CdcDebeziumPackExec (Execution Plan Node) +// ============================================================================ + +/// Packs internal flat changelog rows into Debezium-style `before` / `after` / `op` / timestamp. +/// +/// Intended as the last physical node before a sink that expects Debezium CDC envelopes. +#[derive(Debug)] +pub struct CdcDebeziumPackExec { + input: Arc, + schema: SchemaRef, + properties: PlanProperties, +} + +impl CdcDebeziumPackExec { + pub fn try_new(input: Arc) -> Result { + let input_schema = input.schema(); + let timestamp_index = input_schema.index_of(TIMESTAMP_FIELD)?; + + let struct_fields: Vec<_> = input_schema + .fields() + .iter() + .enumerate() + .filter_map(|(index, field)| { + if field.name() == UPDATING_META_FIELD || index == timestamp_index { + None + } else { + Some(field.clone()) + } + }) + .collect(); + + let payload_struct_type = DataType::Struct(struct_fields.into()); + + let before_field = Arc::new(Field::new(cdc::BEFORE, payload_struct_type.clone(), true)); + let after_field = Arc::new(Field::new(cdc::AFTER, payload_struct_type, true)); + let op_field = Arc::new(Field::new(cdc::OP, DataType::Utf8, false)); + let timestamp_field = Arc::new(input_schema.field(timestamp_index).clone()); + + let output_schema = Arc::new(Schema::new(vec![ + before_field, + after_field, + op_field, + timestamp_field, + ])); + + Ok(Self { + input, + schema: output_schema.clone(), + properties: make_stream_properties(output_schema), + }) + } + + pub fn from_decoded_parts(input: Arc, schema: SchemaRef) -> Self { + Self { + properties: make_stream_properties(schema.clone()), + input, + schema, + } + } +} + +impl DisplayAs for CdcDebeziumPackExec { + fn fmt_as( + &self, + _t: datafusion::physical_plan::DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "CdcDebeziumPackExec") + } +} + +impl ExecutionPlan for CdcDebeziumPackExec { + fn name(&self) -> &str { + physical_plan_node_name::TO_DEBEZIUM_EXEC + } + + fn as_any(&self) -> &dyn Any { + self + } + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + fn properties(&self) -> &PlanProperties { + &self.properties + } + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if children.len() != 1 { + return Err(DataFusionError::Internal( + "CdcDebeziumPackExec expects exactly 1 child".into(), + )); + } + Ok(Arc::new(Self::try_new(children[0].clone())?)) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let updating_meta_index = self.input.schema().index_of(UPDATING_META_FIELD).ok(); + let timestamp_index = self.input.schema().index_of(TIMESTAMP_FIELD)?; + + let struct_projection = (0..self.input.schema().fields().len()) + .filter(|index| (updating_meta_index != Some(*index)) && *index != timestamp_index) + .collect(); + + Ok(Box::pin(CdcDebeziumPackStream { + input: self.input.execute(partition, context)?, + schema: self.schema.clone(), + updating_meta_index, + timestamp_index, + struct_projection, + })) + } + + fn reset(&self) -> Result<()> { + self.input.reset() + } +} + +// ============================================================================ +// CdcDebeziumPackStream (Physical Stream Execution) +// ============================================================================ + +struct CdcDebeziumPackStream { + input: SendableRecordBatchStream, + schema: SchemaRef, + updating_meta_index: Option, + timestamp_index: usize, + struct_projection: Vec, +} + +#[derive(Debug)] +struct RowCompactionState { + first_idx: usize, + last_idx: usize, + first_is_create: bool, + last_is_create: bool, + max_timestamp: i64, +} + +impl CdcDebeziumPackStream { + fn compact_changelog<'a>( + num_rows: usize, + is_retract: &'a BooleanArray, + id_array: &'a FixedSizeBinaryArray, + timestamps: &'a PrimitiveArray, + ) -> (Vec<&'a [u8]>, HashMap<&'a [u8], RowCompactionState>) { + let mut state_map: HashMap<&[u8], RowCompactionState> = HashMap::new(); + let mut unique_order = vec![]; + + for i in 0..num_rows { + let row_id = id_array.value(i); + let is_create = !is_retract.value(i); + let timestamp = timestamps.value(i); + + state_map + .entry(row_id) + .and_modify(|state| { + state.last_idx = i; + state.last_is_create = is_create; + state.max_timestamp = state.max_timestamp.max(timestamp); + }) + .or_insert_with(|| { + unique_order.push(row_id); + RowCompactionState { + first_idx: i, + last_idx: i, + first_is_create: is_create, + last_is_create: is_create, + max_timestamp: timestamp, + } + }); + } + (unique_order, state_map) + } + + fn as_debezium_batch(&mut self, batch: &RecordBatch) -> Result { + let value_struct = batch.project(&self.struct_projection)?; + let timestamps = batch + .column(self.timestamp_index) + .as_primitive::(); + + let columns: Vec> = if let Some(meta_index) = self.updating_meta_index { + let metadata = batch.column(meta_index).as_struct(); + let is_retract = metadata.column(0).as_boolean(); + let row_ids = metadata.column(1).as_fixed_size_binary(); + + let (ordered_ids, state_map) = + Self::compact_changelog(batch.num_rows(), is_retract, row_ids, timestamps); + + let mut before_builder = UInt32Builder::with_capacity(state_map.len()); + let mut after_builder = UInt32Builder::with_capacity(state_map.len()); + let mut op_vec = Vec::with_capacity(state_map.len()); + let mut ts_builder = TimestampNanosecondBuilder::with_capacity(state_map.len()); + + for row_id in ordered_ids { + let state = state_map + .get(row_id) + .expect("row id from order must exist in map"); + + match (state.first_is_create, state.last_is_create) { + (true, true) => { + before_builder.append_null(); + after_builder.append_value(state.last_idx as u32); + op_vec.push(debezium_op_short::CREATE); + } + (false, false) => { + before_builder.append_value(state.first_idx as u32); + after_builder.append_null(); + op_vec.push(debezium_op_short::DELETE); + } + (false, true) => { + before_builder.append_value(state.first_idx as u32); + after_builder.append_value(state.last_idx as u32); + op_vec.push(debezium_op_short::UPDATE); + } + (true, false) => { + continue; + } + } + ts_builder.append_value(state.max_timestamp); + } + + let before_indices = before_builder.finish(); + let after_indices = after_builder.finish(); + + let before_array = Self::take_struct_columns(&value_struct, &before_indices)?; + let after_array = Self::take_struct_columns(&value_struct, &after_indices)?; + let op_array = StringArray::from(op_vec); + + vec![ + Arc::new(before_array), + Arc::new(after_array), + Arc::new(op_array), + Arc::new(ts_builder.finish()), + ] + } else { + let num_rows = value_struct.num_rows(); + + let after_array = StructArray::try_new( + value_struct.schema().fields().clone(), + value_struct.columns().to_vec(), + None, + )?; + let before_array = + StructArray::new_null(value_struct.schema().fields().clone(), num_rows); + + let op_array = StringArray::from_iter_values(std::iter::repeat_n( + debezium_op_short::CREATE, + num_rows, + )); + + vec![ + Arc::new(before_array), + Arc::new(after_array), + Arc::new(op_array), + batch.column(self.timestamp_index).clone(), + ] + }; + + Ok(RecordBatch::try_new(self.schema.clone(), columns)?) + } + + fn take_struct_columns( + value_struct: &RecordBatch, + indices: &UInt32Array, + ) -> Result { + let mut arrays: Vec> = Vec::with_capacity(value_struct.num_columns()); + + for col in value_struct.columns() { + arrays.push(take(col.as_ref(), indices, None)?); + } + + Ok(StructArray::try_new( + value_struct.schema().fields().clone(), + arrays, + indices.nulls().cloned(), + )?) + } +} + +impl Stream for CdcDebeziumPackStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.as_mut().get_mut(); + match ready!(this.input.poll_next_unpin(cx)) { + Some(Ok(batch)) => Poll::Ready(Some(this.as_debezium_batch(&batch))), + Some(Err(e)) => Poll::Ready(Some(Err(e))), + None => Poll::Ready(None), + } + } +} + +impl RecordBatchStream for CdcDebeziumPackStream { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} diff --git a/src/streaming_planner/src/physical/cdc/mod.rs b/src/streaming_planner/src/physical/cdc/mod.rs new file mode 100644 index 00000000..216dd4c1 --- /dev/null +++ b/src/streaming_planner/src/physical/cdc/mod.rs @@ -0,0 +1,17 @@ +// 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. + +mod encode; +mod unroll; + +pub use encode::CdcDebeziumPackExec; +pub use unroll::CdcDebeziumUnrollExec; diff --git a/src/streaming_planner/src/physical/cdc/unroll.rs b/src/streaming_planner/src/physical/cdc/unroll.rs new file mode 100644 index 00000000..963283cc --- /dev/null +++ b/src/streaming_planner/src/physical/cdc/unroll.rs @@ -0,0 +1,322 @@ +// 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::any::Any; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use datafusion::arrow::array::AsArray; +use datafusion::arrow::array::{ + BooleanBuilder, RecordBatch, StructArray, TimestampNanosecondBuilder, UInt32Builder, +}; +use datafusion::arrow::compute::{concat, take}; +use datafusion::arrow::datatypes::{ + DataType, Field, Schema, SchemaRef, TimeUnit, TimestampNanosecondType, +}; +use datafusion::common::{DataFusionError, Result, plan_err}; +use datafusion::execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext}; +use datafusion::logical_expr::ColumnarValue; +use datafusion::physical_plan::{DisplayAs, ExecutionPlan, PlanProperties}; +use futures::{StreamExt, ready, stream::Stream}; + +use crate::common::TIMESTAMP_FIELD; +use crate::common::constants::{cdc, debezium_op_short, physical_plan_node_name}; +use crate::functions::MultiHashFunction; +use crate::physical::meta::{updating_meta_field, updating_meta_fields}; +use crate::physical::source_exec::make_stream_properties; + +// ============================================================================ +// CdcDebeziumUnrollExec (Execution Plan Node) +// ============================================================================ + +/// Physical node that unrolls Debezium CDC payloads (`before` / `after` / `op`) into a flat +/// changelog stream with retract metadata. +/// +/// - `c` / `r` → emit `after` (`is_retract = false`) +/// - `d` → emit `before` (`is_retract = true`) +/// - `u` → emit `before` (retract) then `after` (insert) +#[derive(Debug)] +pub struct CdcDebeziumUnrollExec { + input: Arc, + schema: SchemaRef, + properties: PlanProperties, + primary_key_indices: Vec, +} + +impl CdcDebeziumUnrollExec { + /// Builds the node and validates Debezium payload schema constraints. + pub fn try_new(input: Arc, primary_key_indices: Vec) -> Result { + let input_schema = input.schema(); + + let before_index = input_schema.index_of(cdc::BEFORE)?; + let after_index = input_schema.index_of(cdc::AFTER)?; + let op_index = input_schema.index_of(cdc::OP)?; + let _timestamp_index = input_schema.index_of(TIMESTAMP_FIELD)?; + + let before_type = input_schema.field(before_index).data_type(); + let after_type = input_schema.field(after_index).data_type(); + + if before_type != after_type { + return Err(DataFusionError::Plan( + "CDC 'before' and 'after' columns must share the exact same DataType".to_string(), + )); + } + + if *input_schema.field(op_index).data_type() != DataType::Utf8 { + return Err(DataFusionError::Plan( + "CDC 'op' (operation) column must be of type Utf8 (String)".to_string(), + )); + } + + let DataType::Struct(fields) = before_type else { + return Err(DataFusionError::Plan( + "CDC 'before' and 'after' payload columns must be Structs".to_string(), + )); + }; + + let mut unrolled_fields = fields.to_vec(); + unrolled_fields.push(updating_meta_field()); + unrolled_fields.push(Arc::new(Field::new( + TIMESTAMP_FIELD, + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + ))); + + let schema = Arc::new(Schema::new(unrolled_fields)); + + Ok(Self { + input, + schema: schema.clone(), + properties: make_stream_properties(schema), + primary_key_indices, + }) + } + + /// Used when deserializing a plan with a pre-baked output schema (see [`StreamingExtensionCodec`]). + pub fn from_decoded_parts( + input: Arc, + schema: SchemaRef, + primary_key_indices: Vec, + ) -> Self { + Self { + properties: make_stream_properties(schema.clone()), + input, + schema, + primary_key_indices, + } + } + + pub fn primary_key_indices(&self) -> &[usize] { + &self.primary_key_indices + } +} + +impl DisplayAs for CdcDebeziumUnrollExec { + fn fmt_as( + &self, + _t: datafusion::physical_plan::DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "CdcDebeziumUnrollExec") + } +} + +impl ExecutionPlan for CdcDebeziumUnrollExec { + fn name(&self) -> &str { + physical_plan_node_name::DEBEZIUM_UNROLLING_EXEC + } + + fn as_any(&self) -> &dyn Any { + self + } + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + fn properties(&self) -> &PlanProperties { + &self.properties + } + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if children.len() != 1 { + return Err(DataFusionError::Internal( + "CdcDebeziumUnrollExec expects exactly one child".to_string(), + )); + } + Ok(Arc::new(Self { + input: children[0].clone(), + schema: self.schema.clone(), + properties: self.properties.clone(), + primary_key_indices: self.primary_key_indices.clone(), + })) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + Ok(Box::pin(CdcDebeziumUnrollStream::try_new( + self.input.execute(partition, context)?, + self.schema.clone(), + self.primary_key_indices.clone(), + )?)) + } + + fn reset(&self) -> Result<()> { + self.input.reset() + } +} + +// ============================================================================ +// CdcDebeziumUnrollStream (Physical Stream Execution) +// ============================================================================ + +struct CdcDebeziumUnrollStream { + input: SendableRecordBatchStream, + schema: SchemaRef, + before_index: usize, + after_index: usize, + op_index: usize, + timestamp_index: usize, + primary_key_indices: Vec, +} + +impl CdcDebeziumUnrollStream { + fn try_new( + input: SendableRecordBatchStream, + schema: SchemaRef, + primary_key_indices: Vec, + ) -> Result { + if primary_key_indices.is_empty() { + return plan_err!( + "A CDC source requires at least one primary key to maintain state correctly." + ); + } + + let input_schema = input.schema(); + Ok(Self { + input, + schema, + before_index: input_schema.index_of(cdc::BEFORE)?, + after_index: input_schema.index_of(cdc::AFTER)?, + op_index: input_schema.index_of(cdc::OP)?, + timestamp_index: input_schema.index_of(TIMESTAMP_FIELD)?, + primary_key_indices, + }) + } + + fn unroll_batch(&self, batch: &RecordBatch) -> Result { + let num_rows = batch.num_rows(); + if num_rows == 0 { + return Ok(RecordBatch::new_empty(self.schema.clone())); + } + + let before_col = batch.column(self.before_index); + let after_col = batch.column(self.after_index); + + let op_array = batch.column(self.op_index).as_string::(); + let timestamp_array = batch + .column(self.timestamp_index) + .as_primitive::(); + + let max_capacity = num_rows * 2; + let mut take_indices = UInt32Builder::with_capacity(max_capacity); + let mut is_retract_builder = BooleanBuilder::with_capacity(max_capacity); + let mut timestamp_builder = TimestampNanosecondBuilder::with_capacity(max_capacity); + + for i in 0..num_rows { + let op = op_array.value(i); + let ts = timestamp_array.value(i); + + match op { + debezium_op_short::CREATE | debezium_op_short::READ => { + take_indices.append_value((i + num_rows) as u32); + is_retract_builder.append_value(false); + timestamp_builder.append_value(ts); + } + debezium_op_short::DELETE => { + take_indices.append_value(i as u32); + is_retract_builder.append_value(true); + timestamp_builder.append_value(ts); + } + debezium_op_short::UPDATE => { + take_indices.append_value(i as u32); + is_retract_builder.append_value(true); + timestamp_builder.append_value(ts); + + take_indices.append_value((i + num_rows) as u32); + is_retract_builder.append_value(false); + timestamp_builder.append_value(ts); + } + _ => { + return Err(DataFusionError::Execution(format!( + "Encountered unexpected Debezium operation code: '{op}'" + ))); + } + } + } + + let take_indices = take_indices.finish(); + let unrolled_row_count = take_indices.len(); + + let combined_array = concat(&[before_col.as_ref(), after_col.as_ref()])?; + let unrolled_array = take(&combined_array, &take_indices, None)?; + + let mut final_columns = unrolled_array.as_struct().columns().to_vec(); + + let pk_columns: Vec = self + .primary_key_indices + .iter() + .map(|&idx| ColumnarValue::Array(Arc::clone(&final_columns[idx]))) + .collect(); + + let hash_column = MultiHashFunction::default().invoke(&pk_columns)?; + let ids_array = hash_column.into_array(unrolled_row_count)?; + + let meta_struct = StructArray::try_new( + updating_meta_fields(), + vec![Arc::new(is_retract_builder.finish()), ids_array], + None, + )?; + + final_columns.push(Arc::new(meta_struct)); + final_columns.push(Arc::new(timestamp_builder.finish())); + + Ok(RecordBatch::try_new(self.schema.clone(), final_columns)?) + } +} + +impl Stream for CdcDebeziumUnrollStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.as_mut().get_mut(); + match ready!(this.input.poll_next_unpin(cx)) { + Some(Ok(batch)) => Poll::Ready(Some(this.unroll_batch(&batch))), + Some(Err(e)) => Poll::Ready(Some(Err(e))), + None => Poll::Ready(None), + } + } +} + +impl RecordBatchStream for CdcDebeziumUnrollStream { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} diff --git a/src/streaming_planner/src/physical/codec.rs b/src/streaming_planner/src/physical/codec.rs new file mode 100644 index 00000000..2fb6fcfc --- /dev/null +++ b/src/streaming_planner/src/physical/codec.rs @@ -0,0 +1,307 @@ +// 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::sync::Arc; + +use datafusion::arrow::array::RecordBatch; +use datafusion::arrow::datatypes::Schema; +use datafusion::common::{DataFusionError, Result, UnnestOptions, not_impl_err}; +use datafusion::execution::FunctionRegistry; +use datafusion::logical_expr::ScalarUDF; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::unnest::{ListUnnest, UnnestExec}; +use datafusion_proto::physical_plan::PhysicalExtensionCodec; +use prost::Message; +use protocol::function_stream_graph::{ + DebeziumDecodeNode, DebeziumEncodeNode, FsExecNode, MemExecNode, UnnestExecNode, + fs_exec_node::Node, +}; +use tokio::sync::mpsc::UnboundedReceiver; + +use crate::analysis::UNNESTED_COL; +use crate::common::constants::{mem_exec_join_side, window_function_udf}; +use crate::physical::cdc::{CdcDebeziumPackExec, CdcDebeziumUnrollExec}; +use crate::physical::source_exec::{ + BufferedBatchesExec, InjectableSingleBatchExec, MpscReceiverStreamExec, PlanningPlaceholderExec, +}; +use crate::physical::udfs::window; + +// ============================================================================ +// StreamingExtensionCodec & StreamingDecodingContext +// ============================================================================ + +/// Worker-side context used when deserializing a physical plan from the coordinator. +/// +/// Planning uses [`PlanningPlaceholderExec`]; at runtime this selects the real source +/// implementation (locked batch, MPSC stream, join sides, etc.). +#[derive(Debug)] +pub enum StreamingDecodingContext { + None, + Planning, + SingleLockedBatch(Arc>>), + UnboundedBatchStream(Arc>>>), + LockedBatchVec(Arc>>), + LockedJoinPair { + left: Arc>>, + right: Arc>>, + }, + LockedJoinStream { + left: Arc>>>, + right: Arc>>>, + }, +} + +/// Codec for custom streaming physical extension nodes (`FsExecNode` protobuf). +#[derive(Debug)] +pub struct StreamingExtensionCodec { + pub context: StreamingDecodingContext, +} + +impl Default for StreamingExtensionCodec { + fn default() -> Self { + Self { + context: StreamingDecodingContext::None, + } + } +} + +impl PhysicalExtensionCodec for StreamingExtensionCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[Arc], + _registry: &dyn FunctionRegistry, + ) -> Result> { + let exec: FsExecNode = Message::decode(buf).map_err(|err| { + DataFusionError::Internal(format!("Failed to deserialize FsExecNode protobuf: {err}")) + })?; + + let node = exec.node.ok_or_else(|| { + DataFusionError::Internal("Decoded FsExecNode contains no inner node data".to_string()) + })?; + + match node { + Node::MemExec(mem) => self.decode_placeholder_exec(mem), + Node::UnnestExec(unnest) => decode_unnest_exec(unnest, inputs), + Node::DebeziumDecode(debezium) => decode_debezium_unroll(debezium, inputs), + Node::DebeziumEncode(debezium) => decode_debezium_pack(debezium, inputs), + } + } + + fn try_encode(&self, node: Arc, buf: &mut Vec) -> Result<()> { + let mut proto = None; + + if let Some(table) = node.as_any().downcast_ref::() { + let schema_json = serde_json::to_string(&table.schema).map_err(|e| { + DataFusionError::Internal(format!("Failed to serialize schema to JSON: {e}")) + })?; + + proto = Some(FsExecNode { + node: Some(Node::MemExec(MemExecNode { + table_name: table.table_name.clone(), + schema: schema_json, + })), + }); + } else if let Some(unnest) = node.as_any().downcast_ref::() { + let schema_json = serde_json::to_string(&unnest.schema()).map_err(|e| { + DataFusionError::Internal(format!("Failed to serialize unnest schema to JSON: {e}")) + })?; + + proto = Some(FsExecNode { + node: Some(Node::UnnestExec(UnnestExecNode { + schema: schema_json, + })), + }); + } else if let Some(decode) = node.as_any().downcast_ref::() { + let schema_json = serde_json::to_string(decode.schema().as_ref()).map_err(|e| { + DataFusionError::Internal(format!("Failed to serialize CDC unroll schema: {e}")) + })?; + + proto = Some(FsExecNode { + node: Some(Node::DebeziumDecode(DebeziumDecodeNode { + schema: schema_json, + primary_keys: decode + .primary_key_indices() + .iter() + .map(|&c| c as u64) + .collect(), + })), + }); + } else if let Some(encode) = node.as_any().downcast_ref::() { + let schema_json = serde_json::to_string(encode.schema().as_ref()).map_err(|e| { + DataFusionError::Internal(format!("Failed to serialize CDC pack schema: {e}")) + })?; + + proto = Some(FsExecNode { + node: Some(Node::DebeziumEncode(DebeziumEncodeNode { + schema: schema_json, + })), + }); + } + + if let Some(proto_node) = proto { + proto_node.encode(buf).map_err(|err| { + DataFusionError::Internal(format!("Failed to encode protobuf node: {err}")) + })?; + Ok(()) + } else { + Err(DataFusionError::Internal(format!( + "Cannot serialize unknown physical plan node: {node:?}" + ))) + } + } + + fn try_decode_udf(&self, name: &str, _buf: &[u8]) -> Result> { + if name == window_function_udf::NAME { + return Ok(window()); + } + not_impl_err!("PhysicalExtensionCodec does not support scalar function '{name}'") + } +} + +impl StreamingExtensionCodec { + fn decode_placeholder_exec(&self, mem_exec: MemExecNode) -> Result> { + let schema: Schema = serde_json::from_str(&mem_exec.schema).map_err(|e| { + DataFusionError::Internal(format!("Invalid schema JSON in exec codec: {e:?}")) + })?; + let schema = Arc::new(schema); + + match &self.context { + StreamingDecodingContext::SingleLockedBatch(single_batch) => Ok(Arc::new( + InjectableSingleBatchExec::new(schema, single_batch.clone()), + )), + StreamingDecodingContext::UnboundedBatchStream(unbounded_stream) => Ok(Arc::new( + MpscReceiverStreamExec::new(schema, unbounded_stream.clone()), + )), + StreamingDecodingContext::LockedBatchVec(locked_batches) => Ok(Arc::new( + BufferedBatchesExec::new(schema, locked_batches.clone()), + )), + StreamingDecodingContext::Planning => Ok(Arc::new(PlanningPlaceholderExec::new( + mem_exec.table_name, + schema, + ))), + StreamingDecodingContext::None => Err(DataFusionError::Internal( + "A valid StreamingDecodingContext is required to decode placeholders into execution streams.".into(), + )), + StreamingDecodingContext::LockedJoinPair { left, right } => { + match mem_exec.table_name.as_str() { + mem_exec_join_side::LEFT => Ok(Arc::new(InjectableSingleBatchExec::new( + schema, + left.clone(), + ))), + mem_exec_join_side::RIGHT => Ok(Arc::new(InjectableSingleBatchExec::new( + schema, + right.clone(), + ))), + _ => Err(DataFusionError::Internal(format!( + "Unknown join side table name: {}", + mem_exec.table_name + ))), + } + } + StreamingDecodingContext::LockedJoinStream { left, right } => { + match mem_exec.table_name.as_str() { + mem_exec_join_side::LEFT => Ok(Arc::new(MpscReceiverStreamExec::new( + schema, + left.clone(), + ))), + mem_exec_join_side::RIGHT => Ok(Arc::new(MpscReceiverStreamExec::new( + schema, + right.clone(), + ))), + _ => Err(DataFusionError::Internal(format!( + "Unknown join side table name: {}", + mem_exec.table_name + ))), + } + } + } + } +} + +fn decode_unnest_exec( + unnest: UnnestExecNode, + inputs: &[Arc], +) -> Result> { + let schema: Schema = serde_json::from_str(&unnest.schema) + .map_err(|e| DataFusionError::Internal(format!("Invalid unnest schema JSON: {e:?}")))?; + + let column = schema.index_of(UNNESTED_COL).map_err(|_| { + DataFusionError::Internal(format!( + "Unnest schema missing required column: {UNNESTED_COL}" + )) + })?; + + let input = inputs.first().ok_or_else(|| { + DataFusionError::Internal("UnnestExec requires exactly one input plan".to_string()) + })?; + + Ok(Arc::new(UnnestExec::new( + input.clone(), + vec![ListUnnest { + index_in_input_schema: column, + depth: 1, + }], + vec![], + Arc::new(schema), + UnnestOptions::default(), + ))) +} + +fn decode_debezium_unroll( + debezium: DebeziumDecodeNode, + inputs: &[Arc], +) -> Result> { + let schema = Arc::new( + serde_json::from_str::(&debezium.schema).map_err(|e| { + DataFusionError::Internal(format!("Invalid DebeziumDecode schema JSON: {e:?}")) + })?, + ); + + let input = inputs.first().ok_or_else(|| { + DataFusionError::Internal( + "CdcDebeziumUnrollExec requires exactly one input plan".to_string(), + ) + })?; + + let primary_keys = debezium + .primary_keys + .into_iter() + .map(|c| c as usize) + .collect(); + + Ok(Arc::new(CdcDebeziumUnrollExec::from_decoded_parts( + input.clone(), + schema, + primary_keys, + ))) +} + +fn decode_debezium_pack( + debezium: DebeziumEncodeNode, + inputs: &[Arc], +) -> Result> { + let schema = Arc::new( + serde_json::from_str::(&debezium.schema).map_err(|e| { + DataFusionError::Internal(format!("Invalid DebeziumEncode schema JSON: {e:?}")) + })?, + ); + + let input = inputs.first().ok_or_else(|| { + DataFusionError::Internal("CdcDebeziumPackExec requires exactly one input plan".to_string()) + })?; + + Ok(Arc::new(CdcDebeziumPackExec::from_decoded_parts( + input.clone(), + schema, + ))) +} diff --git a/src/streaming_planner/src/physical/meta.rs b/src/streaming_planner/src/physical/meta.rs new file mode 100644 index 00000000..ced37656 --- /dev/null +++ b/src/streaming_planner/src/physical/meta.rs @@ -0,0 +1,47 @@ +// 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::sync::{Arc, OnceLock}; + +use datafusion::arrow::datatypes::{DataType, Field, Fields}; + +use crate::common::UPDATING_META_FIELD; +use crate::common::constants::updating_state_field; + +pub fn updating_meta_fields() -> Fields { + static FIELDS: OnceLock = OnceLock::new(); + FIELDS + .get_or_init(|| { + Fields::from(vec![ + Field::new(updating_state_field::IS_RETRACT, DataType::Boolean, true), + Field::new( + updating_state_field::ID, + DataType::FixedSizeBinary(16), + true, + ), + ]) + }) + .clone() +} + +pub fn updating_meta_field() -> Arc { + static FIELD: OnceLock> = OnceLock::new(); + FIELD + .get_or_init(|| { + Arc::new(Field::new( + UPDATING_META_FIELD, + DataType::Struct(updating_meta_fields()), + false, + )) + }) + .clone() +} diff --git a/src/streaming_planner/src/physical/mod.rs b/src/streaming_planner/src/physical/mod.rs new file mode 100644 index 00000000..77f5c107 --- /dev/null +++ b/src/streaming_planner/src/physical/mod.rs @@ -0,0 +1,23 @@ +// 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. + +mod cdc; +mod codec; +mod meta; +mod source_exec; +mod udfs; + +pub use cdc::{CdcDebeziumPackExec, CdcDebeziumUnrollExec}; +pub use codec::{StreamingDecodingContext, StreamingExtensionCodec}; +pub use meta::{updating_meta_field, updating_meta_fields}; +pub use source_exec::FsMemExec; +pub use udfs::window; diff --git a/src/streaming_planner/src/physical/source_exec.rs b/src/streaming_planner/src/physical/source_exec.rs new file mode 100644 index 00000000..c94c1ca3 --- /dev/null +++ b/src/streaming_planner/src/physical/source_exec.rs @@ -0,0 +1,400 @@ +// 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::any::Any; +use std::mem; +use std::sync::Arc; + +use datafusion::arrow::array::RecordBatch; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::catalog::memory::MemorySourceConfig; +use datafusion::common::{DataFusionError, Result, Statistics, not_impl_err, plan_err}; +use datafusion::datasource::memory::DataSourceExec; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::physical_expr::EquivalenceProperties; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::memory::MemoryStream; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{DisplayAs, ExecutionPlan, Partitioning, PlanProperties}; +use futures::StreamExt; +use tokio::sync::mpsc::UnboundedReceiver; +use tokio_stream::wrappers::UnboundedReceiverStream; + +use crate::common::constants::physical_plan_node_name; + +/// Standard [`PlanProperties`] for a continuous, unbounded stream: incremental emission, +/// unknown partitioning, and unbounded boundedness (without requiring infinite memory). +pub fn create_unbounded_stream_properties(schema: SchemaRef) -> PlanProperties { + PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Unbounded { + requires_infinite_memory: false, + }, + ) +} + +/// Alias for call sites that still use the older name. +pub fn make_stream_properties(schema: SchemaRef) -> PlanProperties { + create_unbounded_stream_properties(schema) +} + +// ============================================================================ +// InjectableSingleBatchExec (formerly RwLockRecordBatchReader) +// ============================================================================ + +/// Yields exactly one [`RecordBatch`], injected via a lock before `execute` runs. +/// +/// For event-driven loops that receive a single batch from the network and run a DataFusion +/// plan over it, the batch is stored in the lock until execution starts. +#[derive(Debug)] +pub struct InjectableSingleBatchExec { + schema: SchemaRef, + injected_batch: Arc>>, + properties: PlanProperties, +} + +impl InjectableSingleBatchExec { + pub fn new( + schema: SchemaRef, + injected_batch: Arc>>, + ) -> Self { + Self { + schema: schema.clone(), + injected_batch, + properties: create_unbounded_stream_properties(schema), + } + } +} + +impl DisplayAs for InjectableSingleBatchExec { + fn fmt_as( + &self, + _t: datafusion::physical_plan::DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "InjectableSingleBatchExec") + } +} + +impl ExecutionPlan for InjectableSingleBatchExec { + fn as_any(&self) -> &dyn Any { + self + } + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + fn children(&self) -> Vec<&Arc> { + vec![] + } + fn properties(&self) -> &PlanProperties { + &self.properties + } + fn name(&self) -> &str { + physical_plan_node_name::RW_LOCK_READER + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + Err(DataFusionError::Internal( + "InjectableSingleBatchExec does not support children".into(), + )) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + let mut guard = self.injected_batch.write().map_err(|e| { + DataFusionError::Execution(format!("Failed to acquire write lock: {e}")) + })?; + + let batch = guard.take().ok_or_else(|| { + DataFusionError::Execution( + "Execution triggered, but no RecordBatch was injected into the node.".into(), + ) + })?; + + Ok(Box::pin(MemoryStream::try_new( + vec![batch], + self.schema.clone(), + None, + )?)) + } + + fn statistics(&self) -> Result { + Ok(Statistics::new_unknown(&self.schema)) + } + + fn reset(&self) -> Result<()> { + Ok(()) + } +} + +// ============================================================================ +// MpscReceiverStreamExec (formerly UnboundedRecordBatchReader) +// ============================================================================ + +/// Unbounded streaming source backed by a Tokio `mpsc` receiver. +/// +/// Bridges async producers (e.g. network threads) into a DataFusion pipeline. +#[derive(Debug)] +pub struct MpscReceiverStreamExec { + schema: SchemaRef, + channel_receiver: Arc>>>, + properties: PlanProperties, +} + +impl MpscReceiverStreamExec { + pub fn new( + schema: SchemaRef, + channel_receiver: Arc>>>, + ) -> Self { + Self { + schema: schema.clone(), + channel_receiver, + properties: create_unbounded_stream_properties(schema), + } + } +} + +impl DisplayAs for MpscReceiverStreamExec { + fn fmt_as( + &self, + _t: datafusion::physical_plan::DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "MpscReceiverStreamExec") + } +} + +impl ExecutionPlan for MpscReceiverStreamExec { + fn as_any(&self) -> &dyn Any { + self + } + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + fn children(&self) -> Vec<&Arc> { + vec![] + } + fn properties(&self) -> &PlanProperties { + &self.properties + } + fn name(&self) -> &str { + physical_plan_node_name::UNBOUNDED_READER + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + Err(DataFusionError::Internal( + "MpscReceiverStreamExec does not support children".into(), + )) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + let mut guard = self.channel_receiver.write().map_err(|e| { + DataFusionError::Execution(format!("Failed to acquire lock for MPSC receiver: {e}")) + })?; + + let receiver = guard.take().ok_or_else(|| { + DataFusionError::Execution( + "The MPSC receiver was already consumed by a previous execution.".into(), + ) + })?; + + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema.clone(), + UnboundedReceiverStream::new(receiver).map(Ok), + ))) + } + + fn statistics(&self) -> Result { + Ok(Statistics::new_unknown(&self.schema)) + } + + fn reset(&self) -> Result<()> { + Ok(()) + } +} + +// ============================================================================ +// BufferedBatchesExec (formerly RecordBatchVecReader) +// ============================================================================ + +/// Drains a growable, locked `Vec` when `execute` runs (micro-batching). +#[derive(Debug)] +pub struct BufferedBatchesExec { + schema: SchemaRef, + buffered_batches: Arc>>, + properties: PlanProperties, +} + +impl BufferedBatchesExec { + pub fn new( + schema: SchemaRef, + buffered_batches: Arc>>, + ) -> Self { + Self { + schema: schema.clone(), + buffered_batches, + properties: create_unbounded_stream_properties(schema), + } + } +} + +impl DisplayAs for BufferedBatchesExec { + fn fmt_as( + &self, + _t: datafusion::physical_plan::DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "BufferedBatchesExec") + } +} + +impl ExecutionPlan for BufferedBatchesExec { + fn as_any(&self) -> &dyn Any { + self + } + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + fn children(&self) -> Vec<&Arc> { + vec![] + } + fn properties(&self) -> &PlanProperties { + &self.properties + } + fn name(&self) -> &str { + physical_plan_node_name::VEC_READER + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + Err(DataFusionError::Internal( + "BufferedBatchesExec does not support children".into(), + )) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let mut guard = self.buffered_batches.write().map_err(|e| { + DataFusionError::Execution(format!("Failed to acquire lock for buffered batches: {e}")) + })?; + + let accumulated_batches = mem::take(&mut *guard); + + let memory_config = + MemorySourceConfig::try_new(&[accumulated_batches], self.schema.clone(), None)?; + + DataSourceExec::new(Arc::new(memory_config)).execute(partition, context) + } + + fn statistics(&self) -> Result { + Ok(Statistics::new_unknown(&self.schema)) + } + + fn reset(&self) -> Result<()> { + Ok(()) + } +} + +// ============================================================================ +#[derive(Debug, Clone)] +pub struct PlanningPlaceholderExec { + pub table_name: String, + pub schema: SchemaRef, + properties: PlanProperties, +} + +impl PlanningPlaceholderExec { + pub fn new(table_name: String, schema: SchemaRef) -> Self { + Self { + schema: schema.clone(), + table_name, + properties: create_unbounded_stream_properties(schema), + } + } +} + +impl DisplayAs for PlanningPlaceholderExec { + fn fmt_as( + &self, + _t: datafusion::physical_plan::DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "PlanningPlaceholderExec: schema={}", self.schema) + } +} + +impl ExecutionPlan for PlanningPlaceholderExec { + fn as_any(&self) -> &dyn Any { + self + } + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + fn children(&self) -> Vec<&Arc> { + vec![] + } + fn properties(&self) -> &PlanProperties { + &self.properties + } + fn name(&self) -> &str { + physical_plan_node_name::MEM_EXEC + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + not_impl_err!("PlanningPlaceholderExec does not accept children.") + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + plan_err!("PlanningPlaceholderExec cannot be executed; swap for a real source before run.") + } + + fn statistics(&self) -> Result { + Ok(Statistics::new_unknown(&self.schema)) + } + + fn reset(&self) -> Result<()> { + Ok(()) + } +} + +// Backward-compatible aliases +pub type FsMemExec = PlanningPlaceholderExec; +pub type RwLockRecordBatchReader = InjectableSingleBatchExec; +pub type UnboundedRecordBatchReader = MpscReceiverStreamExec; +pub type RecordBatchVecReader = BufferedBatchesExec; diff --git a/src/streaming_planner/src/physical/udfs.rs b/src/streaming_planner/src/physical/udfs.rs new file mode 100644 index 00000000..733fa79e --- /dev/null +++ b/src/streaming_planner/src/physical/udfs.rs @@ -0,0 +1,138 @@ +// 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::any::Any; +use std::sync::Arc; + +use datafusion::arrow::array::StructArray; +use datafusion::arrow::datatypes::{DataType, TimeUnit}; +use datafusion::common::{Result, ScalarValue, exec_err}; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility, +}; + +use crate::make_udf_function; +use crate::common::constants::window_function_udf; +use crate::schema::utils::window_arrow_struct; + +// ============================================================================ +// WindowFunctionUdf (User-Defined Scalar Function) +// ============================================================================ + +/// UDF that packs two nanosecond timestamps into the canonical window `Struct` type. +/// +/// Stream SQL uses a single struct column `[start, end)` for tumbling/hopping windows; +/// this keeps `GROUP BY` and physical codec alignment on one Arrow shape. +#[derive(Debug)] +pub struct WindowFunctionUdf { + signature: Signature, +} + +impl Default for WindowFunctionUdf { + fn default() -> Self { + Self { + signature: Signature::new( + TypeSignature::Exact(vec![ + DataType::Timestamp(TimeUnit::Nanosecond, None), + DataType::Timestamp(TimeUnit::Nanosecond, None), + ]), + Volatility::Immutable, + ), + } + } +} + +impl ScalarUDFImpl for WindowFunctionUdf { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + window_function_udf::NAME + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(window_arrow_struct()) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let columns = args.args; + + if columns.len() != 2 { + return exec_err!( + "Window UDF expected exactly 2 arguments, but received {}", + columns.len() + ); + } + + let DataType::Struct(fields) = window_arrow_struct() else { + return exec_err!( + "Internal Engine Error: window_arrow_struct() must return a Struct DataType" + ); + }; + + let start_val = &columns[0]; + let end_val = &columns[1]; + + if !matches!( + start_val.data_type(), + DataType::Timestamp(TimeUnit::Nanosecond, _) + ) { + return exec_err!("Window UDF expected first argument to be a Nanosecond Timestamp"); + } + if !matches!( + end_val.data_type(), + DataType::Timestamp(TimeUnit::Nanosecond, _) + ) { + return exec_err!("Window UDF expected second argument to be a Nanosecond Timestamp"); + } + + match (start_val, end_val) { + (ColumnarValue::Array(start_arr), ColumnarValue::Array(end_arr)) => { + let struct_array = + StructArray::try_new(fields, vec![start_arr.clone(), end_arr.clone()], None)?; + Ok(ColumnarValue::Array(Arc::new(struct_array))) + } + + (ColumnarValue::Array(start_arr), ColumnarValue::Scalar(end_scalar)) => { + let end_arr = end_scalar.to_array_of_size(start_arr.len())?; + let struct_array = + StructArray::try_new(fields, vec![start_arr.clone(), end_arr], None)?; + Ok(ColumnarValue::Array(Arc::new(struct_array))) + } + + (ColumnarValue::Scalar(start_scalar), ColumnarValue::Array(end_arr)) => { + let start_arr = start_scalar.to_array_of_size(end_arr.len())?; + let struct_array = + StructArray::try_new(fields, vec![start_arr, end_arr.clone()], None)?; + Ok(ColumnarValue::Array(Arc::new(struct_array))) + } + + (ColumnarValue::Scalar(start_scalar), ColumnarValue::Scalar(end_scalar)) => { + let struct_array = StructArray::try_new( + fields, + vec![start_scalar.to_array()?, end_scalar.to_array()?], + None, + )?; + Ok(ColumnarValue::Scalar(ScalarValue::Struct(Arc::new( + struct_array, + )))) + } + } + } +} + +make_udf_function!(WindowFunctionUdf, WINDOW_FUNCTION, window); diff --git a/src/streaming_planner/src/planning_runtime.rs b/src/streaming_planner/src/planning_runtime.rs new file mode 100644 index 00000000..8a03addc --- /dev/null +++ b/src/streaming_planner/src/planning_runtime.rs @@ -0,0 +1,35 @@ +// 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. + +//! Runtime-installed SQL planning defaults (from `GlobalConfig` / `conf/config.yaml`). + +use std::sync::OnceLock; + +use function_stream_config::streaming_job::ResolvedStreamingJobConfig; +use crate::common::constants::sql_planning_default; +use crate::types::SqlConfig; + +static SQL_PLANNING: OnceLock = OnceLock::new(); + +/// Installs [`SqlConfig`] derived from resolved streaming job YAML (KeyBy parallelism, etc.). +/// Safe to call once at bootstrap; later calls are ignored if already set. +pub fn install_sql_planning_from_streaming_job(job: &ResolvedStreamingJobConfig) { + let cfg = SqlConfig { + default_parallelism: sql_planning_default::DEFAULT_PARALLELISM, + key_by_parallelism: job.key_by_parallelism as usize, + }; + let _ = SQL_PLANNING.set(cfg).ok(); +} + +pub fn sql_planning_snapshot() -> SqlConfig { + SQL_PLANNING.get().cloned().unwrap_or_default() +} diff --git a/src/streaming_planner/src/schema/catalog.rs b/src/streaming_planner/src/schema/catalog.rs new file mode 100644 index 00000000..058c461d --- /dev/null +++ b/src/streaming_planner/src/schema/catalog.rs @@ -0,0 +1,609 @@ +// 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. + +//! External connector catalog: [`ExternalTable`] as [`SourceTable`] | [`SinkTable`] | [`LookupTable`]. + +use std::collections::BTreeMap; +use std::sync::Arc; +use std::time::Duration; + +use datafusion::arrow::datatypes::{DataType, Field, FieldRef, Schema}; +use datafusion::common::{Column, Result, plan_err}; +use datafusion::error::DataFusionError; +use datafusion::logical_expr::Expr; +use protocol::function_stream_graph::ConnectorOp; + +use super::column_descriptor::ColumnDescriptor; +use super::data_encoding_format::DataEncodingFormat; +use super::table::SqlSource; +use super::temporal_pipeline_config::TemporalPipelineConfig; +use crate::multifield_partial_ord; +use crate::common::constants::sql_field; +use crate::common::{Format, FsSchema}; +use crate::connector::config::ConnectorConfig; +use crate::types::ProcessingMode; + +#[derive(Debug, Clone)] +pub struct EngineDescriptor { + pub engine_type: String, + pub raw_payload: String, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum SyncMode { + AppendOnly, + Incremental, +} + +#[derive(Debug, Clone)] +pub struct TableExecutionUnit { + pub label: String, + pub engine_meta: EngineDescriptor, + pub sync_mode: SyncMode, + pub temporal_offset: TemporalPipelineConfig, +} + +/// The only legal shape an external-connector catalog row can take. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum ExternalTable { + Source(SourceTable), + Sink(SinkTable), + Lookup(LookupTable), +} + +impl ExternalTable { + #[inline] + pub fn name(&self) -> &str { + match self { + ExternalTable::Source(t) => t.table_identifier.as_str(), + ExternalTable::Sink(t) => t.table_identifier.as_str(), + ExternalTable::Lookup(t) => t.table_identifier.as_str(), + } + } + + #[inline] + pub fn adapter_type(&self) -> &str { + match self { + ExternalTable::Source(t) => t.adapter_type.as_str(), + ExternalTable::Sink(t) => t.adapter_type.as_str(), + ExternalTable::Lookup(t) => t.adapter_type.as_str(), + } + } + + #[inline] + pub fn description(&self) -> &str { + match self { + ExternalTable::Source(t) => t.description.as_str(), + ExternalTable::Sink(t) => t.description.as_str(), + ExternalTable::Lookup(t) => t.description.as_str(), + } + } + + #[inline] + pub fn schema_specs(&self) -> &[ColumnDescriptor] { + match self { + ExternalTable::Source(t) => &t.schema_specs, + ExternalTable::Sink(t) => &t.schema_specs, + ExternalTable::Lookup(t) => &t.schema_specs, + } + } + + #[inline] + pub fn connector_config(&self) -> &ConnectorConfig { + match self { + ExternalTable::Source(t) => &t.connector_config, + ExternalTable::Sink(t) => &t.connector_config, + ExternalTable::Lookup(t) => &t.connector_config, + } + } + + #[inline] + pub fn key_constraints(&self) -> &[String] { + match self { + ExternalTable::Source(t) => &t.key_constraints, + ExternalTable::Sink(t) => &t.key_constraints, + ExternalTable::Lookup(t) => &t.key_constraints, + } + } + + #[inline] + pub fn connection_format(&self) -> Option<&Format> { + match self { + ExternalTable::Source(t) => t.connection_format.as_ref(), + ExternalTable::Sink(t) => t.connection_format.as_ref(), + ExternalTable::Lookup(t) => t.connection_format.as_ref(), + } + } + + #[inline] + pub fn catalog_with_options(&self) -> &BTreeMap { + match self { + ExternalTable::Source(t) => &t.catalog_with_options, + ExternalTable::Sink(t) => &t.catalog_with_options, + ExternalTable::Lookup(t) => &t.catalog_with_options, + } + } + + pub fn produce_physical_schema(&self) -> Schema { + Schema::new( + self.schema_specs() + .iter() + .filter(|c| !c.is_computed()) + .map(|c| c.arrow_field().clone()) + .collect::>(), + ) + } + + pub fn connector_op(&self) -> ConnectorOp { + let physical = self.produce_physical_schema(); + let fields: Vec = physical + .fields() + .iter() + .map(|f| f.as_ref().clone()) + .collect(); + let fs_schema = FsSchema::from_fields(fields); + + ConnectorOp { + connector: self.adapter_type().to_string(), + fs_schema: Some(fs_schema.into()), + name: self.name().to_string(), + description: self.description().to_string(), + config: Some(self.connector_config().to_proto_config()), + } + } + + #[inline] + pub fn is_updating(&self) -> bool { + match self { + ExternalTable::Source(t) => t.is_updating(), + ExternalTable::Sink(t) => t + .connection_format + .as_ref() + .is_some_and(|f| f.is_updating()), + ExternalTable::Lookup(_) => false, + } + } + + /// Variant-agnostic view of "persisted Arrow fields post-planning". + /// Only Source / Lookup track inferred schema — Sinks derive theirs from the upstream plan. + pub fn effective_fields(&self) -> Vec { + match self { + ExternalTable::Source(t) => t.effective_fields(), + ExternalTable::Sink(t) => t.effective_fields(), + ExternalTable::Lookup(t) => t.effective_fields(), + } + } + + #[inline] + pub fn as_source(&self) -> Option<&SourceTable> { + match self { + ExternalTable::Source(t) => Some(t), + _ => None, + } + } + + #[inline] + pub fn as_sink(&self) -> Option<&SinkTable> { + match self { + ExternalTable::Sink(t) => Some(t), + _ => None, + } + } + + #[inline] + pub fn as_lookup(&self) -> Option<&LookupTable> { + match self { + ExternalTable::Lookup(t) => Some(t), + _ => None, + } + } +} + +/// Ingress external connector (`CREATE TABLE ... WITH (type='source', ...)`). +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct SourceTable { + pub table_identifier: String, + pub adapter_type: String, + pub schema_specs: Vec, + pub connector_config: ConnectorConfig, + pub temporal_config: TemporalPipelineConfig, + pub key_constraints: Vec, + pub payload_format: Option, + pub connection_format: Option, + pub description: String, + pub catalog_with_options: BTreeMap, + + // Planner / catalog; not in SQL text. + pub registry_id: Option, + pub inferred_fields: Option>, +} + +multifield_partial_ord!( + SourceTable, + registry_id, + adapter_type, + table_identifier, + description, + key_constraints, + connection_format, + catalog_with_options +); + +impl SourceTable { + #[inline] + pub fn name(&self) -> &str { + self.table_identifier.as_str() + } + + #[inline] + pub fn connector(&self) -> &str { + self.adapter_type.as_str() + } + + pub fn event_time_field(&self) -> Option<&str> { + self.temporal_config.event_column.as_deref() + } + + pub fn watermark_field(&self) -> Option<&str> { + self.temporal_config.watermark_strategy_column.as_deref() + } + + /// Watermark column safe to persist to the stream catalog. Omits the + /// generated `__watermark` column — that is only resolvable at compile + /// time, the catalog round-trip cannot reconstruct it. + pub fn stream_catalog_watermark_field(&self) -> Option { + self.temporal_config + .watermark_strategy_column + .as_deref() + .filter(|w| *w != sql_field::COMPUTED_WATERMARK) + .map(str::to_string) + } + + #[inline] + pub fn catalog_with_options(&self) -> &BTreeMap { + &self.catalog_with_options + } + + pub fn idle_time(&self) -> Option { + self.temporal_config.liveness_timeout + } + + pub fn produce_physical_schema(&self) -> Schema { + Schema::new( + self.schema_specs + .iter() + .filter(|c| !c.is_computed()) + .map(|c| c.arrow_field().clone()) + .collect::>(), + ) + } + + #[inline] + pub fn physical_schema(&self) -> Schema { + self.produce_physical_schema() + } + + pub fn effective_fields(&self) -> Vec { + self.inferred_fields.clone().unwrap_or_else(|| { + self.schema_specs + .iter() + .map(|c| Arc::new(c.arrow_field().clone())) + .collect() + }) + } + + pub fn convert_to_execution_unit(&self) -> Result { + if self.is_cdc_enabled() && self.schema_specs.iter().any(|c| c.is_computed()) { + return plan_err!("CDC cannot be mixed with computed columns natively"); + } + + let mode = if self.is_cdc_enabled() { + SyncMode::Incremental + } else { + SyncMode::AppendOnly + }; + + Ok(TableExecutionUnit { + label: self.table_identifier.clone(), + engine_meta: EngineDescriptor { + engine_type: self.adapter_type.clone(), + raw_payload: String::new(), + }, + sync_mode: mode, + temporal_offset: self.temporal_config.clone(), + }) + } + + #[inline] + pub fn to_execution_unit(&self) -> Result { + self.convert_to_execution_unit() + } + + fn is_cdc_enabled(&self) -> bool { + self.payload_format + .as_ref() + .is_some_and(|f| f.supports_delta_updates()) + } + + pub fn has_virtual_fields(&self) -> bool { + self.schema_specs.iter().any(|c| c.is_computed()) + } + + pub fn is_updating(&self) -> bool { + self.connection_format + .as_ref() + .is_some_and(|f| f.is_updating()) + || self.payload_format == Some(DataEncodingFormat::DebeziumJson) + } + + pub fn connector_op(&self) -> ConnectorOp { + let physical = self.produce_physical_schema(); + let fields: Vec = physical + .fields() + .iter() + .map(|f| f.as_ref().clone()) + .collect(); + let fs_schema = FsSchema::from_fields(fields); + + ConnectorOp { + connector: self.adapter_type.clone(), + fs_schema: Some(fs_schema.into()), + name: self.table_identifier.clone(), + description: self.description.clone(), + config: Some(self.connector_config.to_proto_config()), + } + } + + pub fn processing_mode(&self) -> ProcessingMode { + if self.is_updating() { + ProcessingMode::Update + } else { + ProcessingMode::Append + } + } + + pub fn timestamp_override(&self) -> Result> { + if let Some(field_name) = self.temporal_config.event_column.clone() { + if self.is_updating() { + return plan_err!("can't use event_time_field with update mode"); + } + let _field = self.get_time_column(&field_name)?; + Ok(Some(Expr::Column(Column::from_name(field_name.as_str())))) + } else { + Ok(None) + } + } + + fn get_time_column(&self, field_name: &str) -> Result<&ColumnDescriptor> { + self.schema_specs + .iter() + .find(|c| { + c.arrow_field().name() == field_name + && matches!(c.arrow_field().data_type(), DataType::Timestamp(..)) + }) + .ok_or_else(|| { + DataFusionError::Plan(format!("field {field_name} not found or not a timestamp")) + }) + } + + pub fn watermark_column(&self) -> Result> { + if let Some(field_name) = self.temporal_config.watermark_strategy_column.clone() { + let _field = self.get_time_column(&field_name)?; + Ok(Some(Expr::Column(Column::from_name(field_name.as_str())))) + } else { + Ok(None) + } + } + + pub fn as_sql_source(&self) -> Result { + if self.is_updating() && self.has_virtual_fields() { + return plan_err!("can't read from a source with virtual fields and update mode."); + } + + let timestamp_override = self.timestamp_override()?; + let watermark_column = self.watermark_column()?; + + let source = SqlSource { + id: self.registry_id, + struct_def: self + .schema_specs + .iter() + .filter(|c| !c.is_computed()) + .map(|c| Arc::new(c.arrow_field().clone())) + .collect(), + config: self.connector_op(), + processing_mode: self.processing_mode(), + idle_time: self.temporal_config.liveness_timeout, + }; + + Ok(SourceOperator { + name: self.table_identifier.clone(), + source, + timestamp_override, + watermark_column, + }) + } +} + +/// Egress external connector, or the sink of `CREATE STREAMING TABLE ... AS SELECT`. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct SinkTable { + pub table_identifier: String, + pub adapter_type: String, + pub schema_specs: Vec, + pub connector_config: ConnectorConfig, + pub partition_exprs: Arc>>, + pub key_constraints: Vec, + pub connection_format: Option, + pub description: String, + pub catalog_with_options: BTreeMap, +} + +multifield_partial_ord!( + SinkTable, + adapter_type, + table_identifier, + description, + key_constraints, + connection_format, + catalog_with_options +); + +impl SinkTable { + #[inline] + pub fn name(&self) -> &str { + self.table_identifier.as_str() + } + + #[inline] + pub fn connector(&self) -> &str { + self.adapter_type.as_str() + } + + #[inline] + pub fn catalog_with_options(&self) -> &BTreeMap { + &self.catalog_with_options + } + + pub fn produce_physical_schema(&self) -> Schema { + Schema::new( + self.schema_specs + .iter() + .filter(|c| !c.is_computed()) + .map(|c| c.arrow_field().clone()) + .collect::>(), + ) + } + + pub fn effective_fields(&self) -> Vec { + self.schema_specs + .iter() + .map(|c| Arc::new(c.arrow_field().clone())) + .collect() + } + + pub fn is_updating(&self) -> bool { + self.connection_format + .as_ref() + .is_some_and(|f| f.is_updating()) + } + + pub fn connector_op(&self) -> ConnectorOp { + let physical = self.produce_physical_schema(); + let fields: Vec = physical + .fields() + .iter() + .map(|f| f.as_ref().clone()) + .collect(); + let fs_schema = FsSchema::from_fields(fields); + + ConnectorOp { + connector: self.adapter_type.clone(), + fs_schema: Some(fs_schema.into()), + name: self.table_identifier.clone(), + description: self.description.clone(), + config: Some(self.connector_config.to_proto_config()), + } + } +} + +/// Lookup-join only; not a scan source. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct LookupTable { + pub table_identifier: String, + pub adapter_type: String, + pub schema_specs: Vec, + pub connector_config: ConnectorConfig, + pub key_constraints: Vec, + pub lookup_cache_max_bytes: Option, + pub lookup_cache_ttl: Option, + pub connection_format: Option, + pub description: String, + pub catalog_with_options: BTreeMap, + + pub registry_id: Option, + pub inferred_fields: Option>, +} + +multifield_partial_ord!( + LookupTable, + registry_id, + adapter_type, + table_identifier, + description, + key_constraints, + connection_format, + catalog_with_options +); + +impl LookupTable { + #[inline] + pub fn name(&self) -> &str { + self.table_identifier.as_str() + } + + #[inline] + pub fn connector(&self) -> &str { + self.adapter_type.as_str() + } + + #[inline] + pub fn catalog_with_options(&self) -> &BTreeMap { + &self.catalog_with_options + } + + pub fn produce_physical_schema(&self) -> Schema { + Schema::new( + self.schema_specs + .iter() + .filter(|c| !c.is_computed()) + .map(|c| c.arrow_field().clone()) + .collect::>(), + ) + } + + pub fn effective_fields(&self) -> Vec { + self.inferred_fields.clone().unwrap_or_else(|| { + self.schema_specs + .iter() + .map(|c| Arc::new(c.arrow_field().clone())) + .collect() + }) + } + + pub fn connector_op(&self) -> ConnectorOp { + let physical = self.produce_physical_schema(); + let fields: Vec = physical + .fields() + .iter() + .map(|f| f.as_ref().clone()) + .collect(); + let fs_schema = FsSchema::from_fields(fields); + + ConnectorOp { + connector: self.adapter_type.clone(), + fs_schema: Some(fs_schema.into()), + name: self.table_identifier.clone(), + description: self.description.clone(), + config: Some(self.connector_config.to_proto_config()), + } + } +} + +/// [`SourceTable`] as an ingestion logical node input. +#[derive(Debug, Clone)] +pub struct SourceOperator { + pub name: String, + pub source: SqlSource, + pub timestamp_override: Option, + pub watermark_column: Option, +} diff --git a/src/streaming_planner/src/schema/column_descriptor.rs b/src/streaming_planner/src/schema/column_descriptor.rs new file mode 100644 index 00000000..4228816f --- /dev/null +++ b/src/streaming_planner/src/schema/column_descriptor.rs @@ -0,0 +1,144 @@ +// 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 datafusion::arrow::datatypes::{DataType, Field, TimeUnit}; +use datafusion::logical_expr::Expr; + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum ColumnDescriptor { + Physical(Field), + SystemMeta { field: Field, meta_key: String }, + Computed { field: Field, logic: Box }, +} + +impl ColumnDescriptor { + #[inline] + pub fn new_physical(field: Field) -> Self { + Self::Physical(field) + } + + #[inline] + pub fn new_system_meta(field: Field, meta_key: impl Into) -> Self { + Self::SystemMeta { + field, + meta_key: meta_key.into(), + } + } + + #[inline] + pub fn new_computed(field: Field, logic: Expr) -> Self { + Self::Computed { + field, + logic: Box::new(logic), + } + } + + #[inline] + pub fn arrow_field(&self) -> &Field { + match self { + Self::Physical(f) => f, + Self::SystemMeta { field: f, .. } => f, + Self::Computed { field: f, .. } => f, + } + } + + #[inline] + pub fn into_arrow_field(self) -> Field { + match self { + Self::Physical(f) => f, + Self::SystemMeta { field: f, .. } => f, + Self::Computed { field: f, .. } => f, + } + } + + #[inline] + pub fn is_computed(&self) -> bool { + matches!(self, Self::Computed { .. }) + } + + #[inline] + pub fn is_physical(&self) -> bool { + matches!(self, Self::Physical(_)) + } + + #[inline] + pub fn system_meta_key(&self) -> Option<&str> { + if let Self::SystemMeta { meta_key, .. } = self { + Some(meta_key.as_str()) + } else { + None + } + } + + #[inline] + pub fn computation_logic(&self) -> Option<&Expr> { + if let Self::Computed { logic, .. } = self { + Some(logic) + } else { + None + } + } + + #[inline] + pub fn data_type(&self) -> &DataType { + self.arrow_field().data_type() + } + + pub fn set_nullable(&mut self, nullable: bool) { + let f = match self { + Self::Physical(f) => f, + Self::SystemMeta { field, .. } => field, + Self::Computed { field, .. } => field, + }; + *f = Field::new(f.name(), f.data_type().clone(), nullable) + .with_metadata(f.metadata().clone()); + } + + pub fn force_precision(&mut self, unit: TimeUnit) { + match self { + Self::Physical(f) => { + if let DataType::Timestamp(_, tz) = f.data_type() { + *f = Field::new( + f.name(), + DataType::Timestamp(unit, tz.clone()), + f.is_nullable(), + ); + } + } + Self::SystemMeta { field, .. } => { + if let DataType::Timestamp(_, tz) = field.data_type() { + *field = Field::new( + field.name(), + DataType::Timestamp(unit, tz.clone()), + field.is_nullable(), + ); + } + } + Self::Computed { field, .. } => { + if let DataType::Timestamp(_, tz) = field.data_type() { + *field = Field::new( + field.name(), + DataType::Timestamp(unit, tz.clone()), + field.is_nullable(), + ); + } + } + } + } +} + +impl From for ColumnDescriptor { + #[inline] + fn from(field: Field) -> Self { + Self::Physical(field) + } +} diff --git a/src/streaming_planner/src/schema/connection_type.rs b/src/streaming_planner/src/schema/connection_type.rs new file mode 100644 index 00000000..06a3df92 --- /dev/null +++ b/src/streaming_planner/src/schema/connection_type.rs @@ -0,0 +1,31 @@ +// 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::fmt; + +/// Describes the role of a connection in the streaming pipeline. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum ConnectionType { + Source, + Sink, + Lookup, +} + +impl fmt::Display for ConnectionType { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + ConnectionType::Source => write!(f, "source"), + ConnectionType::Sink => write!(f, "sink"), + ConnectionType::Lookup => write!(f, "lookup"), + } + } +} diff --git a/src/streaming_planner/src/schema/data_encoding_format.rs b/src/streaming_planner/src/schema/data_encoding_format.rs new file mode 100644 index 00000000..b589d683 --- /dev/null +++ b/src/streaming_planner/src/schema/data_encoding_format.rs @@ -0,0 +1,89 @@ +// 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 datafusion::arrow::datatypes::{DataType, Field}; +use datafusion::common::{Result, plan_err}; + +use super::column_descriptor::ColumnDescriptor; +use crate::common::Format; +use crate::common::constants::cdc; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] +pub enum DataEncodingFormat { + #[default] + Raw, + StandardJson, + DebeziumJson, + Avro, + Parquet, + Csv, + JsonL, + Orc, + Protobuf, +} + +impl DataEncodingFormat { + pub fn from_format(format: Option<&Format>) -> Self { + match format { + Some(Format::Json(j)) if j.debezium => Self::DebeziumJson, + Some(Format::Json(_)) => Self::StandardJson, + Some(Format::Avro(_)) => Self::Avro, + Some(Format::Parquet(_)) => Self::Parquet, + Some(Format::Csv(_)) => Self::Csv, + Some(Format::Protobuf(_)) => Self::Protobuf, + Some(Format::RawString(_)) | Some(Format::RawBytes(_)) | None => Self::Raw, + Some(_) => Self::Raw, + } + } + + pub fn is_cdc_format(&self) -> bool { + matches!(self, Self::DebeziumJson) + } + + #[inline] + pub fn supports_delta_updates(&self) -> bool { + self.is_cdc_format() + } + + pub fn apply_envelope( + &self, + logical_columns: Vec, + ) -> Result> { + if !self.is_cdc_format() { + return Ok(logical_columns); + } + + if logical_columns.is_empty() { + return Ok(logical_columns); + } + + if logical_columns.iter().any(|c| c.is_computed()) { + return plan_err!( + "Computed/Virtual columns are not supported directly inside a CDC source table; \ + define computed columns in a downstream VIEW or AS SELECT streaming query" + ); + } + + let inner_fields: Vec = logical_columns + .into_iter() + .map(|c| c.into_arrow_field()) + .collect(); + + let row_struct_type = DataType::Struct(inner_fields.into()); + + Ok(vec![ + ColumnDescriptor::new_physical(Field::new(cdc::BEFORE, row_struct_type.clone(), true)), + ColumnDescriptor::new_physical(Field::new(cdc::AFTER, row_struct_type, true)), + ColumnDescriptor::new_physical(Field::new(cdc::OP, DataType::Utf8, true)), + ]) + } +} diff --git a/src/streaming_planner/src/schema/introspection/ddl_formatter.rs b/src/streaming_planner/src/schema/introspection/ddl_formatter.rs new file mode 100644 index 00000000..15d8f707 --- /dev/null +++ b/src/streaming_planner/src/schema/introspection/ddl_formatter.rs @@ -0,0 +1,156 @@ +// 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 std::fmt::{self, Write}; + +use datafusion::arrow::datatypes::{DataType, Schema, TimeUnit}; + +use crate::common::constants::sql_field; + +pub struct DdlBuilder<'a> { + table_name: &'a str, + schema: &'a Schema, + watermark_column: Option<&'a str>, + primary_keys: &'a [String], + options: BTreeMap, +} + +impl<'a> DdlBuilder<'a> { + pub fn new(table_name: &'a str, schema: &'a Schema) -> Self { + Self { + table_name, + schema, + watermark_column: None, + primary_keys: &[], + options: BTreeMap::new(), + } + } + + pub fn with_watermark(mut self, watermark: Option<&'a str>) -> Self { + self.watermark_column = watermark; + self + } + + pub fn with_primary_keys(mut self, keys: &'a [String]) -> Self { + self.primary_keys = keys; + self + } + + pub fn with_options( + mut self, + opts: &BTreeMap, + role: &str, + connector: &str, + ) -> Self { + self.options = opts.clone(); + self.options + .entry("type".to_string()) + .or_insert_with(|| role.to_string()); + self.options + .entry("connector".to_string()) + .or_insert_with(|| connector.to_string()); + self + } +} + +impl<'a> fmt::Display for DdlBuilder<'a> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + writeln!(f, "CREATE TABLE {} (", self.table_name)?; + + let mut rows: Vec = Vec::new(); + for field in self.schema.fields() { + let null_constraint = if field.is_nullable() { "" } else { " NOT NULL" }; + rows.push(format!( + " {} {}{}", + field.name(), + format_data_type(field.data_type()), + null_constraint + )); + } + + if let Some(wm) = self.watermark_column + && wm != sql_field::COMPUTED_WATERMARK + { + rows.push(format!(" WATERMARK FOR {wm}")); + } + + if !self.primary_keys.is_empty() { + rows.push(format!(" PRIMARY KEY ({})", self.primary_keys.join(", "))); + } + + writeln!(f, "{}", rows.join(",\n"))?; + write!(f, ")")?; + + if !self.options.is_empty() { + writeln!(f)?; + writeln!(f, "WITH (")?; + let mut opt_lines: Vec = Vec::with_capacity(self.options.len()); + for (k, v) in &self.options { + let k_esc = k.replace('\'', "''"); + let v_esc = v.replace('\'', "''"); + opt_lines.push(format!(" '{k_esc}' = '{v_esc}'")); + } + write!(f, "{}\n);", opt_lines.join(",\n"))?; + } else { + write!(f, ";")?; + } + + Ok(()) + } +} + +pub fn format_data_type(dt: &DataType) -> String { + match dt { + DataType::Null => "NULL".to_string(), + DataType::Boolean => "BOOLEAN".to_string(), + DataType::Int8 => "TINYINT".to_string(), + DataType::Int16 => "SMALLINT".to_string(), + DataType::Int32 => "INT".to_string(), + DataType::Int64 => "BIGINT".to_string(), + DataType::UInt8 => "TINYINT UNSIGNED".to_string(), + DataType::UInt16 => "SMALLINT UNSIGNED".to_string(), + DataType::UInt32 => "INT UNSIGNED".to_string(), + DataType::UInt64 => "BIGINT UNSIGNED".to_string(), + DataType::Float16 => "FLOAT".to_string(), + DataType::Float32 => "REAL".to_string(), + DataType::Float64 => "DOUBLE".to_string(), + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => "VARCHAR".to_string(), + DataType::Binary | DataType::LargeBinary => "VARBINARY".to_string(), + DataType::Date32 | DataType::Date64 => "DATE".to_string(), + DataType::Timestamp(unit, tz) => match (unit, tz) { + (TimeUnit::Second, None) => "TIMESTAMP(0)".to_string(), + (TimeUnit::Millisecond, None) => "TIMESTAMP(3)".to_string(), + (TimeUnit::Microsecond, None) => "TIMESTAMP(6)".to_string(), + (TimeUnit::Nanosecond, None) => "TIMESTAMP(9)".to_string(), + (_, Some(_)) => "TIMESTAMP WITH TIME ZONE".to_string(), + }, + DataType::Decimal128(p, s) | DataType::Decimal256(p, s) => format!("DECIMAL({p}, {s})"), + _ => dt.to_string(), + } +} + +pub fn schema_columns_one_line(schema: &Schema) -> String { + let mut buf = String::new(); + for (idx, field) in schema.fields().iter().enumerate() { + if idx > 0 { + buf.push_str(", "); + } + let _ = write!( + buf, + "{}:{}", + field.name(), + format_data_type(field.data_type()) + ); + } + buf +} diff --git a/src/streaming_planner/src/schema/introspection/mod.rs b/src/streaming_planner/src/schema/introspection/mod.rs new file mode 100644 index 00000000..1ba9c816 --- /dev/null +++ b/src/streaming_planner/src/schema/introspection/mod.rs @@ -0,0 +1,21 @@ +// 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. + +pub mod ddl_formatter; +pub mod show_formatter; +pub mod stream_formatter; + +#[allow(unused_imports)] +pub use ddl_formatter::{DdlBuilder, format_data_type, schema_columns_one_line}; +pub use show_formatter::{catalog_table_row_detail, show_create_catalog_table}; +#[allow(unused_imports)] +pub use stream_formatter::{show_create_stream_table, stream_table_row_detail}; diff --git a/src/streaming_planner/src/schema/introspection/show_formatter.rs b/src/streaming_planner/src/schema/introspection/show_formatter.rs new file mode 100644 index 00000000..3857ba21 --- /dev/null +++ b/src/streaming_planner/src/schema/introspection/show_formatter.rs @@ -0,0 +1,100 @@ +// 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::common::constants::connection_table_role; +use crate::schema::catalog::ExternalTable; +use crate::schema::table::CatalogEntity; + +use super::ddl_formatter::DdlBuilder; + +impl ExternalTable { + pub fn to_ddl_string(&self) -> String { + match self { + ExternalTable::Source(source) => { + let schema = source.produce_physical_schema(); + DdlBuilder::new(&source.table_identifier, &schema) + .with_watermark(source.temporal_config.watermark_strategy_column.as_deref()) + .with_primary_keys(&source.key_constraints) + .with_options( + &source.catalog_with_options, + connection_table_role::SOURCE, + &source.adapter_type, + ) + .to_string() + } + ExternalTable::Sink(sink) => { + let schema = sink.produce_physical_schema(); + DdlBuilder::new(&sink.table_identifier, &schema) + .with_primary_keys(&sink.key_constraints) + .with_options( + &sink.catalog_with_options, + connection_table_role::SINK, + &sink.adapter_type, + ) + .to_string() + } + ExternalTable::Lookup(lookup) => { + let schema = lookup.produce_physical_schema(); + DdlBuilder::new(&lookup.table_identifier, &schema) + .with_primary_keys(&lookup.key_constraints) + .with_options( + &lookup.catalog_with_options, + connection_table_role::LOOKUP, + &lookup.adapter_type, + ) + .to_string() + } + } + } + + pub fn to_row_detail(&self) -> String { + match self { + ExternalTable::Source(s) => format!( + "{{ kind: 'source', connector: '{}', watermark: '{}', options_count: {} }}", + s.adapter_type, + s.temporal_config + .watermark_strategy_column + .as_deref() + .unwrap_or("none"), + s.catalog_with_options.len() + ), + ExternalTable::Sink(s) => format!( + "{{ kind: 'sink', connector: '{}', partitioned: {}, options_count: {} }}", + s.adapter_type, + s.partition_exprs.as_ref().is_some(), + s.catalog_with_options.len() + ), + ExternalTable::Lookup(s) => format!( + "{{ kind: 'lookup', connector: '{}', cache_ttl_secs: {}, options_count: {} }}", + s.adapter_type, + s.lookup_cache_ttl.map(|d| d.as_secs()).unwrap_or(0), + s.catalog_with_options.len() + ), + } + } +} + +pub fn show_create_catalog_table(table: &CatalogEntity) -> String { + match table { + CatalogEntity::ExternalConnector(ext) => ext.to_ddl_string(), + CatalogEntity::ComputedTable { name, .. } => { + format!("-- Logical query view\nCREATE VIEW {name} AS SELECT ...;") + } + } +} + +pub fn catalog_table_row_detail(table: &CatalogEntity) -> String { + match table { + CatalogEntity::ExternalConnector(ext) => ext.to_row_detail(), + CatalogEntity::ComputedTable { .. } => "{ kind: 'logical_view' }".to_string(), + } +} diff --git a/src/streaming_planner/src/schema/introspection/stream_formatter.rs b/src/streaming_planner/src/schema/introspection/stream_formatter.rs new file mode 100644 index 00000000..e3239ce8 --- /dev/null +++ b/src/streaming_planner/src/schema/introspection/stream_formatter.rs @@ -0,0 +1,120 @@ +// 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::sync::Arc; + +use datafusion::arrow::datatypes::Schema; + +use crate::common::constants::connection_table_role; +use crate::logical_node::logical::LogicalProgram; +use crate::schema::schema_provider::StreamTable; + +use super::ddl_formatter::DdlBuilder; + +impl StreamTable { + pub fn to_ddl_string(&self) -> String { + match self { + StreamTable::Source { + name, + connector, + schema, + event_time_field: _, + watermark_field, + with_options, + } => DdlBuilder::new(name, schema) + .with_watermark(watermark_field.as_deref()) + .with_options(with_options, connection_table_role::SOURCE, connector) + .to_string(), + StreamTable::Sink { name, program } => { + let schema: Arc = program + .egress_arrow_schema() + .unwrap_or_else(|| Arc::new(Schema::empty())); + + let mut ddl = format!("CREATE STREAMING TABLE {name} AS SELECT ...\n\n"); + ddl.push_str("/* === SINK SCHEMA === */\n"); + let schema_ddl = DdlBuilder::new(name, &schema).to_string(); + ddl.push_str(&schema_ddl); + ddl.push_str("\n\n/* === STREAMING TOPOLOGY === */\n"); + ddl.push_str(&format_pipeline(program)); + ddl + } + } + } + + pub fn to_row_detail(&self) -> String { + match self { + StreamTable::Source { + connector, + event_time_field, + watermark_field, + with_options, + .. + } => format!( + "{{ kind: 'stream_source', connector: '{}', event_time: '{}', watermark: '{}', options_count: {} }}", + connector, + event_time_field.as_deref().unwrap_or("none"), + watermark_field.as_deref().unwrap_or("none"), + with_options.len() + ), + StreamTable::Sink { program, .. } => format!( + "{{ kind: 'streaming_sink', tasks: {}, nodes: {} }}", + program.task_count(), + program.graph.node_count() + ), + } + } +} + +pub fn show_create_stream_table(table: &StreamTable) -> String { + table.to_ddl_string() +} + +pub fn stream_table_row_detail(table: &StreamTable) -> String { + table.to_row_detail() +} + +fn format_pipeline(program: &LogicalProgram) -> String { + let mut lines: Vec = Vec::new(); + lines.push(format!("Pipeline Hash : {}", program.get_hash())); + lines.push(format!("Total Tasks : {}", program.task_count())); + lines.push(format!("Node Count : {}", program.graph.node_count())); + lines.push(String::from("Operator Chains:")); + + for nw in program.graph.node_weights() { + let chain = nw + .operator_chain + .operators + .iter() + .map(|op| format!("{}[{}]", op.operator_name, op.operator_id)) + .collect::>() + .join(" -> "); + + lines.push(format!( + " Node {:<3} | Parallelism {:<3} | {}", + nw.node_id, nw.parallelism, chain + )); + } + + let dot = program.dot(); + const MAX_DOT: usize = 5_000; + if dot.len() > MAX_DOT { + lines.push(format!( + "\nGraphviz DOT (truncated, {} bytes omitted):\n{}...", + dot.len() - MAX_DOT, + &dot[..MAX_DOT] + )); + } else { + lines.push(format!("\nGraphviz DOT:\n{dot}")); + } + + lines.join("\n") +} diff --git a/src/streaming_planner/src/schema/mod.rs b/src/streaming_planner/src/schema/mod.rs new file mode 100644 index 00000000..e11e4808 --- /dev/null +++ b/src/streaming_planner/src/schema/mod.rs @@ -0,0 +1,30 @@ +// 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. + +pub mod catalog; +pub mod column_descriptor; +pub mod connection_type; +pub mod data_encoding_format; +pub mod introspection; +pub mod schema_provider; +pub mod table; +pub mod table_role; +pub mod temporal_pipeline_config; +pub mod utils; + +pub use catalog::{ExternalTable, LookupTable, SinkTable, SourceTable}; +pub use column_descriptor::ColumnDescriptor; +pub use introspection::{ + catalog_table_row_detail, schema_columns_one_line, show_create_catalog_table, +}; +pub use schema_provider::{ObjectName, StreamPlanningContext, StreamSchemaProvider, StreamTable}; +pub use table::CatalogEntity; diff --git a/src/streaming_planner/src/schema/schema_provider.rs b/src/streaming_planner/src/schema/schema_provider.rs new file mode 100644 index 00000000..75026c3f --- /dev/null +++ b/src/streaming_planner/src/schema/schema_provider.rs @@ -0,0 +1,469 @@ +// 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, HashMap, HashSet}; +use std::sync::Arc; + +use datafusion::arrow::datatypes::{self as datatypes, DataType, Field, Schema}; +use datafusion::common::{DataFusionError, Result as DataFusionResult}; +use datafusion::datasource::{DefaultTableSource, TableProvider, TableType}; +use datafusion::execution::{FunctionRegistry, SessionStateDefaults}; +use datafusion::logical_expr::expr_rewriter::FunctionRewrite; +use datafusion::logical_expr::planner::ExprPlanner; +use datafusion::logical_expr::{AggregateUDF, Expr, ScalarUDF, TableSource, WindowUDF}; +use datafusion::optimizer::Analyzer; +use datafusion::sql::TableReference; +use datafusion::sql::planner::ContextProvider; +use thiserror::Error; +use tracing::{debug, error, info}; +use unicase::UniCase; + +use crate::common::constants::{planning_placeholder_udf, window_fn}; +use crate::logical_node::logical::{DylibUdfConfig, LogicalProgram}; +use crate::schema::table::CatalogEntity; +use crate::schema::utils::window_arrow_struct; +use crate::types::{PlanningOptions, PlanningPlaceholderUdf, SqlConfig}; + +pub type ObjectName = UniCase; + +#[inline] +fn object_name(s: impl Into) -> ObjectName { + UniCase::new(s.into()) +} + +#[derive(Error, Debug)] +pub enum PlanningError { + #[error("Catalog table not found: {0}")] + TableNotFound(String), + #[error("Planning init failed: {0}")] + InitError(String), + #[error("Engine error: {0}")] + Engine(#[from] DataFusionError), +} + +impl From for DataFusionError { + fn from(err: PlanningError) -> Self { + match err { + PlanningError::Engine(inner) => inner, + other => DataFusionError::Plan(other.to_string()), + } + } +} + +#[derive(Clone, Debug)] +pub enum StreamTable { + Source { + name: String, + connector: String, + schema: Arc, + event_time_field: Option, + watermark_field: Option, + with_options: BTreeMap, + }, + Sink { + name: String, + program: LogicalProgram, + }, +} + +impl StreamTable { + pub fn name(&self) -> &str { + match self { + Self::Source { name, .. } | Self::Sink { name, .. } => name, + } + } + + pub fn schema(&self) -> Arc { + match self { + Self::Source { schema, .. } => Arc::clone(schema), + Self::Sink { program, .. } => program + .egress_arrow_schema() + .unwrap_or_else(|| Arc::new(Schema::empty())), + } + } +} + +#[derive(Debug, Clone)] +pub struct LogicalBatchInput { + pub table_name: String, + pub schema: Arc, +} + +#[async_trait::async_trait] +impl TableProvider for LogicalBatchInput { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn schema(&self) -> Arc { + Arc::clone(&self.schema) + } + + fn table_type(&self) -> TableType { + TableType::Temporary + } + + async fn scan( + &self, + _state: &dyn datafusion::catalog::Session, + _projection: Option<&Vec>, + _filters: &[Expr], + _limit: Option, + ) -> DataFusionResult> { + Ok(Arc::new(crate::physical::FsMemExec::new( + self.table_name.clone(), + Arc::clone(&self.schema), + ))) + } +} + +#[derive(Clone, Default)] +pub struct FunctionCatalog { + pub scalars: HashMap>, + pub aggregates: HashMap>, + pub windows: HashMap>, + pub planners: Vec>, +} + +#[derive(Clone, Default)] +pub struct TableCatalog { + pub streams: HashMap>, + pub catalogs: HashMap>, + pub source_defs: HashMap, +} + +#[derive(Clone, Default)] +pub struct StreamPlanningContext { + pub tables: TableCatalog, + pub functions: FunctionCatalog, + pub dylib_udfs: HashMap, + pub config_options: datafusion::config::ConfigOptions, + pub planning_options: PlanningOptions, + pub analyzer: Analyzer, + pub sql_config: SqlConfig, +} + +pub type StreamSchemaProvider = StreamPlanningContext; + +impl StreamPlanningContext { + pub fn builder() -> StreamPlanningContextBuilder { + StreamPlanningContextBuilder::default() + } + + #[inline] + pub fn default_parallelism(&self) -> usize { + self.sql_config.default_parallelism + } + + #[inline] + pub fn key_by_parallelism(&self) -> usize { + self.sql_config.key_by_parallelism + } + + pub fn try_new(config: SqlConfig) -> Result { + info!("Initializing StreamPlanningContext"); + let mut builder = StreamPlanningContextBuilder::default(); + builder + .with_streaming_extensions()? + .with_default_functions()?; + let mut ctx = builder.build(); + ctx.sql_config = config; + Ok(ctx) + } + + pub fn new() -> Self { + let config = crate::planning_runtime::sql_planning_snapshot(); + Self::try_new(config).expect("StreamPlanningContext bootstrap") + } + + pub fn register_stream_table(&mut self, table: StreamTable) { + let key = object_name(table.name().to_string()); + debug!(table = %key, "register stream table"); + self.tables.streams.insert(key, Arc::new(table)); + } + + pub fn get_stream_table(&self, name: &str) -> Option> { + self.tables + .streams + .get(&object_name(name.to_string())) + .cloned() + } + + pub fn register_catalog_table(&mut self, table: CatalogEntity) { + let key = object_name(table.name().to_string()); + debug!(table = %key, "register catalog table"); + self.tables.catalogs.insert(key, Arc::new(table)); + } + + pub fn get_catalog_table(&self, table_name: impl AsRef) -> Option<&CatalogEntity> { + self.tables + .catalogs + .get(&object_name(table_name.as_ref().to_string())) + .map(|t| t.as_ref()) + } + + pub fn get_catalog_table_mut( + &mut self, + table_name: impl AsRef, + ) -> Option<&mut CatalogEntity> { + self.tables + .catalogs + .get_mut(&object_name(table_name.as_ref().to_string())) + .map(Arc::make_mut) + } + + pub fn add_source_table( + &mut self, + name: String, + schema: Arc, + event_time_field: Option, + watermark_field: Option, + ) { + self.register_stream_table(StreamTable::Source { + name, + connector: "stream_catalog".to_string(), + schema, + event_time_field, + watermark_field, + with_options: BTreeMap::new(), + }); + } + + pub fn add_sink_table(&mut self, name: String, program: LogicalProgram) { + self.register_stream_table(StreamTable::Sink { name, program }); + } + + pub fn insert_table(&mut self, table: StreamTable) { + self.register_stream_table(table); + } + + pub fn insert_catalog_table(&mut self, table: CatalogEntity) { + self.register_catalog_table(table); + } + + pub fn get_table(&self, table_name: impl AsRef) -> Option<&StreamTable> { + self.tables + .streams + .get(&object_name(table_name.as_ref().to_string())) + .map(|a| a.as_ref()) + } + + pub fn get_table_mut(&mut self, table_name: impl AsRef) -> Option<&mut StreamTable> { + self.tables + .streams + .get_mut(&object_name(table_name.as_ref().to_string())) + .map(Arc::make_mut) + } + + pub fn get_async_udf_options(&self, _name: &str) -> Option { + None + } + + fn create_table_source(name: String, schema: Arc) -> Arc { + let provider = LogicalBatchInput { + table_name: name, + schema, + }; + Arc::new(DefaultTableSource::new(Arc::new(provider))) + } +} + +impl ContextProvider for StreamPlanningContext { + fn get_table_source(&self, name: TableReference) -> DataFusionResult> { + let name_str = name.table(); + match self.get_stream_table(name_str) { + Some(table) => Ok(Self::create_table_source(name.to_string(), table.schema())), + None => { + error!(table = %name_str, "stream table lookup failed"); + Err(DataFusionError::Plan(format!("Table {} not found", name))) + } + } + } + + fn get_function_meta(&self, name: &str) -> Option> { + self.functions.scalars.get(name).cloned() + } + + fn get_aggregate_meta(&self, name: &str) -> Option> { + self.functions.aggregates.get(name).cloned() + } + + fn get_window_meta(&self, name: &str) -> Option> { + self.functions.windows.get(name).cloned() + } + + fn get_variable_type(&self, _variable_names: &[String]) -> Option { + None + } + + fn options(&self) -> &datafusion::config::ConfigOptions { + &self.config_options + } + + fn udf_names(&self) -> Vec { + self.functions.scalars.keys().cloned().collect() + } + + fn udaf_names(&self) -> Vec { + self.functions.aggregates.keys().cloned().collect() + } + + fn udwf_names(&self) -> Vec { + self.functions.windows.keys().cloned().collect() + } + + fn get_expr_planners(&self) -> &[Arc] { + &self.functions.planners + } +} + +impl FunctionRegistry for StreamPlanningContext { + fn udfs(&self) -> HashSet { + self.functions.scalars.keys().cloned().collect() + } + + fn udf(&self, name: &str) -> DataFusionResult> { + self.functions + .scalars + .get(name) + .cloned() + .ok_or_else(|| DataFusionError::Plan(format!("No UDF with name {name}"))) + } + + fn udaf(&self, name: &str) -> DataFusionResult> { + self.functions + .aggregates + .get(name) + .cloned() + .ok_or_else(|| DataFusionError::Plan(format!("No UDAF with name {name}"))) + } + + fn udwf(&self, name: &str) -> DataFusionResult> { + self.functions + .windows + .get(name) + .cloned() + .ok_or_else(|| DataFusionError::Plan(format!("No UDWF with name {name}"))) + } + + fn register_function_rewrite( + &mut self, + rewrite: Arc, + ) -> DataFusionResult<()> { + self.analyzer.add_function_rewrite(rewrite); + Ok(()) + } + + fn register_udf(&mut self, udf: Arc) -> DataFusionResult>> { + Ok(self.functions.scalars.insert(udf.name().to_string(), udf)) + } + + fn register_udaf( + &mut self, + udaf: Arc, + ) -> DataFusionResult>> { + Ok(self + .functions + .aggregates + .insert(udaf.name().to_string(), udaf)) + } + + fn register_udwf(&mut self, udwf: Arc) -> DataFusionResult>> { + Ok(self.functions.windows.insert(udwf.name().to_string(), udwf)) + } + + fn register_expr_planner( + &mut self, + expr_planner: Arc, + ) -> DataFusionResult<()> { + self.functions.planners.push(expr_planner); + Ok(()) + } + + fn expr_planners(&self) -> Vec> { + self.functions.planners.clone() + } +} + +#[derive(Default)] +pub struct StreamPlanningContextBuilder { + context: StreamPlanningContext, +} + +impl StreamPlanningContextBuilder { + pub fn new() -> Self { + Self::default() + } + + pub fn with_default_functions(&mut self) -> Result<&mut Self, PlanningError> { + for p in SessionStateDefaults::default_scalar_functions() { + self.context.register_udf(p)?; + } + for p in SessionStateDefaults::default_aggregate_functions() { + self.context.register_udaf(p)?; + } + for p in SessionStateDefaults::default_window_functions() { + self.context.register_udwf(p)?; + } + for p in SessionStateDefaults::default_expr_planners() { + self.context.register_expr_planner(p)?; + } + Ok(self) + } + + pub fn with_streaming_extensions(&mut self) -> Result<&mut Self, PlanningError> { + let extensions = vec![ + PlanningPlaceholderUdf::new_with_return( + window_fn::HOP, + vec![ + DataType::Interval(datatypes::IntervalUnit::MonthDayNano), + DataType::Interval(datatypes::IntervalUnit::MonthDayNano), + ], + window_arrow_struct(), + ), + PlanningPlaceholderUdf::new_with_return( + window_fn::TUMBLE, + vec![DataType::Interval(datatypes::IntervalUnit::MonthDayNano)], + window_arrow_struct(), + ), + PlanningPlaceholderUdf::new_with_return( + window_fn::SESSION, + vec![DataType::Interval(datatypes::IntervalUnit::MonthDayNano)], + window_arrow_struct(), + ), + PlanningPlaceholderUdf::new_with_return( + planning_placeholder_udf::UNNEST, + vec![DataType::List(Arc::new(Field::new( + planning_placeholder_udf::LIST_ELEMENT_FIELD, + DataType::Utf8, + true, + )))], + DataType::Utf8, + ), + PlanningPlaceholderUdf::new_with_return( + planning_placeholder_udf::ROW_TIME, + vec![], + DataType::Timestamp(datatypes::TimeUnit::Nanosecond, None), + ), + ]; + + for ext in extensions { + self.context.register_udf(ext)?; + } + + Ok(self) + } + + pub fn build(self) -> StreamPlanningContext { + self.context + } +} diff --git a/src/streaming_planner/src/schema/table.rs b/src/streaming_planner/src/schema/table.rs new file mode 100644 index 00000000..912e9b00 --- /dev/null +++ b/src/streaming_planner/src/schema/table.rs @@ -0,0 +1,163 @@ +// 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::analysis::rewrite_plan; +use crate::logical_node::remote_table::RemoteTableBoundaryNode; +use crate::logical_planner::optimizers::produce_optimized_plan; +use crate::schema::StreamSchemaProvider; +use crate::schema::catalog::ExternalTable; +use crate::types::{ProcessingMode, QualifiedField}; +use datafusion::arrow::datatypes::FieldRef; +use datafusion::common::{Result, plan_err}; +use datafusion::logical_expr::{Extension, LogicalPlan}; +use datafusion::sql::sqlparser::ast::Statement; +use protocol::function_stream_graph::ConnectorOp; +use std::sync::Arc; +use std::time::Duration; + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum CatalogEntity { + /// Both payload variants are boxed so the enum is not padded to the largest field. + ExternalConnector(Box), + ComputedTable { + name: String, + logical_plan: Box, + }, +} + +impl CatalogEntity { + #[inline] + pub fn external(table: ExternalTable) -> Self { + Self::ExternalConnector(Box::new(table)) + } + + pub fn try_from_statement( + statement: &Statement, + schema_provider: &StreamSchemaProvider, + ) -> Result> { + use datafusion::logical_expr::{CreateMemoryTable, CreateView, DdlStatement}; + use datafusion::sql::sqlparser::ast::CreateTable; + + if let Statement::CreateTable(CreateTable { query: None, .. }) = statement { + return plan_err!( + "CREATE TABLE without AS SELECT is not supported; use CREATE TABLE ... AS SELECT or a connector table" + ); + } + + match produce_optimized_plan(statement, schema_provider) { + Ok(LogicalPlan::Ddl(DdlStatement::CreateView(CreateView { name, input, .. }))) + | Ok(LogicalPlan::Ddl(DdlStatement::CreateMemoryTable(CreateMemoryTable { + name, + input, + .. + }))) => { + let rewritten = rewrite_plan(input.as_ref().clone(), schema_provider)?; + let schema = rewritten.schema().clone(); + let remote = RemoteTableBoundaryNode { + upstream_plan: rewritten, + table_identifier: name.to_owned(), + resolved_schema: schema, + requires_materialization: true, + }; + Ok(Some(CatalogEntity::ComputedTable { + name: name.to_string(), + logical_plan: Box::new(LogicalPlan::Extension(Extension { + node: Arc::new(remote), + })), + })) + } + _ => Ok(None), + } + } + + pub fn name(&self) -> &str { + match self { + CatalogEntity::ComputedTable { name, .. } => name.as_str(), + CatalogEntity::ExternalConnector(e) => e.name(), + } + } + + pub fn get_fields(&self) -> Vec { + match self { + CatalogEntity::ExternalConnector(e) => e.effective_fields(), + CatalogEntity::ComputedTable { logical_plan, .. } => { + logical_plan.schema().fields().iter().cloned().collect() + } + } + } + + pub fn set_inferred_fields(&mut self, fields: Vec) -> Result<()> { + let CatalogEntity::ExternalConnector(ext) = self else { + return Ok(()); + }; + let ExternalTable::Source(t) = ext.as_mut() else { + return Ok(()); + }; + + if !t.schema_specs.is_empty() { + return Ok(()); + } + + if let Some(existing) = &t.inferred_fields { + let matches = existing.len() == fields.len() + && existing + .iter() + .zip(&fields) + .all(|(a, b)| a.name() == b.name() && a.data_type() == b.data_type()); + + if !matches { + return plan_err!("all inserts into a table must share the same schema"); + } + } + + let fields: Vec<_> = fields.into_iter().map(|f| f.field().clone()).collect(); + t.inferred_fields.replace(fields); + + Ok(()) + } + + pub fn connector_op(&self) -> Result { + match self { + CatalogEntity::ExternalConnector(e) => Ok(e.connector_op()), + CatalogEntity::ComputedTable { .. } => { + plan_err!("can't write to a query-defined table") + } + } + } + + pub fn partition_exprs(&self) -> Option<&Vec> { + let CatalogEntity::ExternalConnector(ext) = self else { + return None; + }; + let ExternalTable::Sink(s) = ext.as_ref() else { + return None; + }; + (*s.partition_exprs).as_ref() + } + + #[inline] + pub fn as_external(&self) -> Option<&ExternalTable> { + match self { + CatalogEntity::ExternalConnector(e) => Some(e.as_ref()), + _ => None, + } + } +} + +#[derive(Clone, Debug)] +pub struct SqlSource { + pub id: Option, + pub struct_def: Vec, + pub config: ConnectorOp, + pub processing_mode: ProcessingMode, + pub idle_time: Option, +} diff --git a/src/streaming_planner/src/schema/table_role.rs b/src/streaming_planner/src/schema/table_role.rs new file mode 100644 index 00000000..4574a6ea --- /dev/null +++ b/src/streaming_planner/src/schema/table_role.rs @@ -0,0 +1,104 @@ +// 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 datafusion::arrow::datatypes::{DataType, TimeUnit}; +use datafusion::common::{Result, plan_err}; +use datafusion::error::DataFusionError; + +use super::column_descriptor::ColumnDescriptor; +use super::connection_type::ConnectionType; +use crate::common::constants::{ + SUPPORTED_CONNECTOR_ADAPTERS, connection_table_role, connector_type, +}; +use crate::common::with_option_keys as opt; + +/// Role of a connector-backed table in the pipeline (ingest / egress / lookup). +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum TableRole { + Ingestion, + Egress, + Reference, +} + +impl From for ConnectionType { + fn from(r: TableRole) -> Self { + match r { + TableRole::Ingestion => ConnectionType::Source, + TableRole::Egress => ConnectionType::Sink, + TableRole::Reference => ConnectionType::Lookup, + } + } +} + +impl From for TableRole { + fn from(c: ConnectionType) -> Self { + match c { + ConnectionType::Source => TableRole::Ingestion, + ConnectionType::Sink => TableRole::Egress, + ConnectionType::Lookup => TableRole::Reference, + } + } +} + +pub fn validate_adapter_availability(adapter: &str) -> Result<()> { + if !SUPPORTED_CONNECTOR_ADAPTERS.contains(&adapter) { + return Err(DataFusionError::Plan(format!( + "Unknown adapter '{adapter}'" + ))); + } + Ok(()) +} + +pub fn apply_adapter_specific_rules( + adapter: &str, + mut cols: Vec, +) -> Vec { + match adapter { + a if a == connector_type::DELTA || a == connector_type::ICEBERG => { + for c in &mut cols { + if matches!(c.data_type(), DataType::Timestamp(_, _)) { + c.force_precision(TimeUnit::Microsecond); + } + } + cols + } + _ => cols, + } +} + +pub fn deduce_role(options: &HashMap) -> Result { + match options.get(opt::TYPE).map(|s| s.as_str()) { + None | Some(connection_table_role::SOURCE) => Ok(TableRole::Ingestion), + Some(connection_table_role::SINK) => Ok(TableRole::Egress), + Some(connection_table_role::LOOKUP) => Ok(TableRole::Reference), + Some(other) => plan_err!("Invalid role '{other}'"), + } +} + +pub fn serialize_backend_params( + adapter: &str, + options: &HashMap, +) -> Result { + let mut payload = serde_json::Map::new(); + payload.insert( + opt::ADAPTER.to_string(), + serde_json::Value::String(adapter.to_string()), + ); + + for (k, v) in options { + payload.insert(k.clone(), serde_json::Value::String(v.clone())); + } + + serde_json::to_string(&payload).map_err(|e| DataFusionError::Plan(e.to_string())) +} diff --git a/src/streaming_planner/src/schema/temporal_pipeline_config.rs b/src/streaming_planner/src/schema/temporal_pipeline_config.rs new file mode 100644 index 00000000..db751b1c --- /dev/null +++ b/src/streaming_planner/src/schema/temporal_pipeline_config.rs @@ -0,0 +1,58 @@ +// 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::time::Duration; + +use datafusion::common::{Result, plan_err}; +use datafusion::logical_expr::Expr; + +use super::column_descriptor::ColumnDescriptor; +use crate::common::constants::sql_field; + +/// Event-time and watermark configuration for streaming tables. +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct TemporalPipelineConfig { + pub event_column: Option, + pub watermark_strategy_column: Option, + pub liveness_timeout: Option, +} + +#[derive(Debug, Clone)] +pub struct TemporalSpec { + pub time_field: String, + pub watermark_expr: Option, +} + +pub fn resolve_temporal_logic( + columns: &[ColumnDescriptor], + time_meta: Option, +) -> Result { + let mut config = TemporalPipelineConfig::default(); + + if let Some(meta) = time_meta { + let field_exists = columns + .iter() + .any(|c| c.arrow_field().name() == meta.time_field.as_str()); + if !field_exists { + return plan_err!("Temporal field {} does not exist", meta.time_field); + } + config.event_column = Some(meta.time_field.clone()); + + if meta.watermark_expr.is_some() { + config.watermark_strategy_column = Some(sql_field::COMPUTED_WATERMARK.to_string()); + } else { + config.watermark_strategy_column = Some(meta.time_field); + } + } + + Ok(config) +} diff --git a/src/streaming_planner/src/schema/utils.rs b/src/streaming_planner/src/schema/utils.rs new file mode 100644 index 00000000..be5efbab --- /dev/null +++ b/src/streaming_planner/src/schema/utils.rs @@ -0,0 +1,79 @@ +// 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::Arc; + +use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; +use datafusion::common::{DFSchema, DFSchemaRef, Result as DFResult, TableReference}; + +use crate::common::constants::window_interval_field; +use crate::types::{QualifiedField, TIMESTAMP_FIELD}; + +/// Returns the Arrow struct type for a window (start, end) pair. +pub fn window_arrow_struct() -> DataType { + DataType::Struct( + vec![ + Arc::new(Field::new( + window_interval_field::START, + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + )), + Arc::new(Field::new( + window_interval_field::END, + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + )), + ] + .into(), + ) +} + +/// Adds a `_timestamp` field to a DFSchema if it doesn't already have one. +pub fn add_timestamp_field( + schema: DFSchemaRef, + qualifier: Option, +) -> DFResult { + if has_timestamp_field(&schema) { + return Ok(schema); + } + + let timestamp_field = QualifiedField::new( + qualifier, + TIMESTAMP_FIELD, + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + ); + Ok(Arc::new(schema.join(&DFSchema::new_with_metadata( + vec![timestamp_field.into()], + HashMap::new(), + )?)?)) +} + +/// Checks whether a DFSchema contains a `_timestamp` field. +pub fn has_timestamp_field(schema: &DFSchemaRef) -> bool { + schema + .fields() + .iter() + .any(|field| field.name() == TIMESTAMP_FIELD) +} + +/// Adds a `_timestamp` field to an Arrow Schema, returning a new SchemaRef. +pub fn add_timestamp_field_arrow(schema: Schema) -> SchemaRef { + let mut fields = schema.fields().to_vec(); + fields.push(Arc::new(Field::new( + TIMESTAMP_FIELD, + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + ))); + Arc::new(Schema::new(fields)) +} diff --git a/src/streaming_planner/src/types/data_type.rs b/src/streaming_planner/src/types/data_type.rs new file mode 100644 index 00000000..44bb8087 --- /dev/null +++ b/src/streaming_planner/src/types/data_type.rs @@ -0,0 +1,158 @@ +// 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::sync::Arc; + +use datafusion::arrow::datatypes::{ + DECIMAL_DEFAULT_SCALE, DECIMAL128_MAX_PRECISION, DataType, Field, IntervalUnit, TimeUnit, +}; +use datafusion::common::{Result, plan_datafusion_err, plan_err}; + +use crate::common::FsExtensionType; +use crate::common::constants::planning_placeholder_udf; + +pub fn convert_data_type( + sql_type: &datafusion::sql::sqlparser::ast::DataType, +) -> Result<(DataType, Option)> { + use datafusion::sql::sqlparser::ast::ArrayElemTypeDef; + use datafusion::sql::sqlparser::ast::DataType as SQLDataType; + + match sql_type { + SQLDataType::Array(ArrayElemTypeDef::AngleBracket(inner_sql_type)) + | SQLDataType::Array(ArrayElemTypeDef::SquareBracket(inner_sql_type, _)) => { + let (data_type, extension) = convert_simple_data_type(inner_sql_type)?; + + Ok(( + DataType::List(Arc::new(FsExtensionType::add_metadata( + extension, + Field::new( + planning_placeholder_udf::LIST_ELEMENT_FIELD, + data_type, + true, + ), + ))), + None, + )) + } + SQLDataType::Array(ArrayElemTypeDef::None) => { + plan_err!("Arrays with unspecified type is not supported") + } + other => convert_simple_data_type(other), + } +} + +fn convert_simple_data_type( + sql_type: &datafusion::sql::sqlparser::ast::DataType, +) -> Result<(DataType, Option)> { + use datafusion::sql::sqlparser::ast::DataType as SQLDataType; + use datafusion::sql::sqlparser::ast::{ExactNumberInfo, TimezoneInfo}; + + if matches!(sql_type, SQLDataType::JSON) { + return Ok((DataType::Utf8, Some(FsExtensionType::JSON))); + } + + let dt = match sql_type { + SQLDataType::Boolean | SQLDataType::Bool => Ok(DataType::Boolean), + SQLDataType::TinyInt(_) => Ok(DataType::Int8), + SQLDataType::SmallInt(_) | SQLDataType::Int2(_) => Ok(DataType::Int16), + SQLDataType::Int(_) | SQLDataType::Integer(_) | SQLDataType::Int4(_) => Ok(DataType::Int32), + SQLDataType::BigInt(_) | SQLDataType::Int8(_) => Ok(DataType::Int64), + SQLDataType::TinyIntUnsigned(_) => Ok(DataType::UInt8), + SQLDataType::SmallIntUnsigned(_) | SQLDataType::Int2Unsigned(_) => Ok(DataType::UInt16), + SQLDataType::IntUnsigned(_) + | SQLDataType::UnsignedInteger + | SQLDataType::Int4Unsigned(_) => Ok(DataType::UInt32), + SQLDataType::BigIntUnsigned(_) | SQLDataType::Int8Unsigned(_) => Ok(DataType::UInt64), + SQLDataType::Float(_) => Ok(DataType::Float32), + SQLDataType::Real | SQLDataType::Float4 => Ok(DataType::Float32), + SQLDataType::Double(_) | SQLDataType::DoublePrecision | SQLDataType::Float8 => { + Ok(DataType::Float64) + } + SQLDataType::Char(_) + | SQLDataType::Varchar(_) + | SQLDataType::Text + | SQLDataType::String(_) => Ok(DataType::Utf8), + SQLDataType::Timestamp(None, TimezoneInfo::None) | SQLDataType::Datetime(_) => { + Ok(DataType::Timestamp(TimeUnit::Nanosecond, None)) + } + SQLDataType::Timestamp(Some(precision), TimezoneInfo::None) => match *precision { + 0 => Ok(DataType::Timestamp(TimeUnit::Second, None)), + 3 => Ok(DataType::Timestamp(TimeUnit::Millisecond, None)), + 6 => Ok(DataType::Timestamp(TimeUnit::Microsecond, None)), + 9 => Ok(DataType::Timestamp(TimeUnit::Nanosecond, None)), + _ => { + return plan_err!( + "unsupported precision {} -- supported precisions are 0 (seconds), \ + 3 (milliseconds), 6 (microseconds), and 9 (nanoseconds)", + precision + ); + } + }, + SQLDataType::Date => Ok(DataType::Date32), + SQLDataType::Time(None, tz_info) + if matches!(tz_info, TimezoneInfo::None) + || matches!(tz_info, TimezoneInfo::WithoutTimeZone) => + { + Ok(DataType::Time64(TimeUnit::Nanosecond)) + } + SQLDataType::Numeric(exact_number_info) | SQLDataType::Decimal(exact_number_info) => { + let (precision, scale) = match *exact_number_info { + ExactNumberInfo::None => (None, None), + ExactNumberInfo::Precision(precision) => (Some(precision), None), + ExactNumberInfo::PrecisionAndScale(precision, scale) => { + (Some(precision), Some(scale)) + } + }; + make_decimal_type(precision, scale) + } + SQLDataType::Bytea => Ok(DataType::Binary), + SQLDataType::Interval => Ok(DataType::Interval(IntervalUnit::MonthDayNano)), + SQLDataType::Struct(fields, _) => { + let fields: Vec<_> = fields + .iter() + .map(|f| { + Ok::<_, datafusion::error::DataFusionError>(Arc::new(Field::new( + f.field_name + .as_ref() + .ok_or_else(|| { + plan_datafusion_err!("anonymous struct fields are not allowed") + })? + .to_string(), + convert_data_type(&f.field_type)?.0, + true, + ))) + }) + .collect::>()?; + Ok(DataType::Struct(fields.into())) + } + _ => return plan_err!("Unsupported SQL type {sql_type:?}"), + }; + + Ok((dt?, None)) +} + +fn make_decimal_type(precision: Option, scale: Option) -> Result { + let (precision, scale) = match (precision, scale) { + (Some(p), Some(s)) => (p as u8, s as i8), + (Some(p), None) => (p as u8, 0), + (None, Some(_)) => return plan_err!("Cannot specify only scale for decimal data type"), + (None, None) => (DECIMAL128_MAX_PRECISION, DECIMAL_DEFAULT_SCALE), + }; + + if precision == 0 || precision > DECIMAL128_MAX_PRECISION || scale.unsigned_abs() > precision { + plan_err!( + "Decimal(precision = {precision}, scale = {scale}) should satisfy `0 < precision <= 38`, and `scale <= precision`." + ) + } else { + Ok(DataType::Decimal128(precision, scale)) + } +} diff --git a/src/streaming_planner/src/types/df_field.rs b/src/streaming_planner/src/types/df_field.rs new file mode 100644 index 00000000..a32d7bc8 --- /dev/null +++ b/src/streaming_planner/src/types/df_field.rs @@ -0,0 +1,181 @@ +// 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::Arc; + +use datafusion::arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion::common::{Column, DFSchema, Result, TableReference}; + +// ============================================================================ +// QualifiedField (Strongly-typed Field Wrapper) +// ============================================================================ + +/// Arrow [`Field`] plus optional SQL [`TableReference`] qualifier. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct QualifiedField { + qualifier: Option, + field: FieldRef, +} + +// ============================================================================ +// Type Conversions (Interoperability with DataFusion) +// ============================================================================ + +impl From<(Option, FieldRef)> for QualifiedField { + fn from((qualifier, field): (Option, FieldRef)) -> Self { + Self { qualifier, field } + } +} + +impl From<(Option<&TableReference>, &Field)> for QualifiedField { + fn from((qualifier, field): (Option<&TableReference>, &Field)) -> Self { + Self { + qualifier: qualifier.cloned(), + field: Arc::new(field.clone()), + } + } +} + +impl From for (Option, FieldRef) { + fn from(value: QualifiedField) -> Self { + (value.qualifier, value.field) + } +} + +// ============================================================================ +// Core API +// ============================================================================ + +impl QualifiedField { + pub fn new( + qualifier: Option, + name: impl Into, + data_type: DataType, + nullable: bool, + ) -> Self { + Self { + qualifier, + field: Arc::new(Field::new(name, data_type, nullable)), + } + } + + pub fn new_unqualified(name: &str, data_type: DataType, nullable: bool) -> Self { + Self { + qualifier: None, + field: Arc::new(Field::new(name, data_type, nullable)), + } + } + + #[inline] + pub fn name(&self) -> &str { + self.field.name() + } + + #[inline] + pub fn data_type(&self) -> &DataType { + self.field.data_type() + } + + #[inline] + pub fn is_nullable(&self) -> bool { + self.field.is_nullable() + } + + #[inline] + pub fn metadata(&self) -> &HashMap { + self.field.metadata() + } + + #[inline] + pub fn qualifier(&self) -> Option<&TableReference> { + self.qualifier.as_ref() + } + + #[inline] + pub fn field(&self) -> &FieldRef { + &self.field + } + + pub fn qualified_name(&self) -> String { + match &self.qualifier { + Some(qualifier) => format!("{}.{}", qualifier, self.field.name()), + None => self.field.name().to_owned(), + } + } + + pub fn qualified_column(&self) -> Column { + Column { + relation: self.qualifier.clone(), + name: self.field.name().to_string(), + spans: Default::default(), + } + } + + pub fn unqualified_column(&self) -> Column { + Column { + relation: None, + name: self.field.name().to_string(), + spans: Default::default(), + } + } + + pub fn strip_qualifier(mut self) -> Self { + self.qualifier = None; + self + } + + pub fn with_nullable(mut self, nullable: bool) -> Self { + if self.field.is_nullable() == nullable { + return self; + } + let field = Arc::try_unwrap(self.field).unwrap_or_else(|arc| (*arc).clone()); + self.field = Arc::new(field.with_nullable(nullable)); + self + } + + pub fn with_metadata(mut self, metadata: HashMap) -> Self { + let field = Arc::try_unwrap(self.field).unwrap_or_else(|arc| (*arc).clone()); + self.field = Arc::new(field.with_metadata(metadata)); + self + } +} + +// ============================================================================ +// Schema Collection Helpers +// ============================================================================ + +pub fn extract_qualified_fields(schema: &DFSchema) -> Vec { + schema + .fields() + .iter() + .enumerate() + .map(|(i, field)| { + let (qualifier, _) = schema.qualified_field(i); + QualifiedField { + qualifier: qualifier.cloned(), + field: field.clone(), + } + }) + .collect() +} + +pub fn build_df_schema(fields: &[QualifiedField]) -> Result { + build_df_schema_with_metadata(fields, HashMap::new()) +} + +pub fn build_df_schema_with_metadata( + fields: &[QualifiedField], + metadata: HashMap, +) -> Result { + DFSchema::new_with_metadata(fields.iter().map(|f| f.clone().into()).collect(), metadata) +} diff --git a/src/streaming_planner/src/types/mod.rs b/src/streaming_planner/src/types/mod.rs new file mode 100644 index 00000000..a796bfea --- /dev/null +++ b/src/streaming_planner/src/types/mod.rs @@ -0,0 +1,65 @@ +// 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. + +mod data_type; +mod df_field; +pub mod placeholder_udf; +mod stream_schema; +mod window; + +use std::time::Duration; + +use crate::common::constants::sql_planning_default; + +pub use df_field::{ + QualifiedField, build_df_schema, build_df_schema_with_metadata, extract_qualified_fields, +}; +pub use placeholder_udf::PlanningPlaceholderUdf; +pub use window::WindowBehavior; +pub use window::{WindowType, extract_window_type}; + +pub use crate::common::constants::sql_field::TIMESTAMP_FIELD; + +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum ProcessingMode { + Append, + Update, +} + +#[derive(Clone, Debug)] +pub struct SqlConfig { + pub default_parallelism: usize, + /// Physical pipeline parallelism for [`KeyExtractionNode`](crate::logical_node::key_calculation::KeyExtractionNode) / KeyBy. + pub key_by_parallelism: usize, +} + +impl Default for SqlConfig { + fn default() -> Self { + Self { + default_parallelism: sql_planning_default::DEFAULT_PARALLELISM, + key_by_parallelism: sql_planning_default::DEFAULT_KEY_BY_PARALLELISM, + } + } +} + +#[derive(Clone)] +pub struct PlanningOptions { + pub ttl: Duration, +} + +impl Default for PlanningOptions { + fn default() -> Self { + Self { + ttl: Duration::from_secs(sql_planning_default::PLANNING_TTL_SECS), + } + } +} diff --git a/src/streaming_planner/src/types/placeholder_udf.rs b/src/streaming_planner/src/types/placeholder_udf.rs new file mode 100644 index 00000000..2f4f3f3c --- /dev/null +++ b/src/streaming_planner/src/types/placeholder_udf.rs @@ -0,0 +1,79 @@ +// 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::any::Any; +use std::fmt::{Debug, Formatter}; +use std::sync::Arc; + +use datafusion::arrow::datatypes::DataType; +use datafusion::common::{Result, internal_err}; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, +}; + +// ============================================================================ +// PlanningPlaceholderUdf +// ============================================================================ + +/// Logical-planning-only UDF: satisfies type checking until real functions are wired in. +pub struct PlanningPlaceholderUdf { + name: String, + signature: Signature, + return_type: DataType, +} + +impl Debug for PlanningPlaceholderUdf { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "PlanningPlaceholderUDF<{}>", self.name) + } +} + +impl ScalarUDFImpl for PlanningPlaceholderUdf { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + &self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(self.return_type.clone()) + } + + fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { + internal_err!( + "PlanningPlaceholderUDF '{}' was invoked during physical execution. \ + This indicates a bug in the stream query compiler: placeholders must be \ + swapped with actual physical UDFs before execution begins.", + self.name + ) + } +} + +impl PlanningPlaceholderUdf { + pub fn new_with_return( + name: impl Into, + args: Vec, + return_type: DataType, + ) -> Arc { + Arc::new(ScalarUDF::new_from_impl(Self { + name: name.into(), + signature: Signature::exact(args, Volatility::Volatile), + return_type, + })) + } +} diff --git a/src/streaming_planner/src/types/stream_schema.rs b/src/streaming_planner/src/types/stream_schema.rs new file mode 100644 index 00000000..c973386e --- /dev/null +++ b/src/streaming_planner/src/types/stream_schema.rs @@ -0,0 +1,133 @@ +// 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::sync::Arc; + +use datafusion::arrow::datatypes::{Field, Schema, SchemaRef}; +use datafusion::common::{DataFusionError, Result}; + +use super::TIMESTAMP_FIELD; + +// ============================================================================ +// StreamSchema +// ============================================================================ + +/// Schema wrapper for continuous streaming: requires event-time (`TIMESTAMP_FIELD`) for watermarks +/// and optionally tracks key column indices for partitioned state / shuffle. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StreamSchema { + schema: SchemaRef, + timestamp_index: usize, + key_indices: Option>, +} + +impl StreamSchema { + // ======================================================================== + // Raw Constructors (When indices are strictly known in advance) + // ======================================================================== + + /// Keyed stream when indices are already verified. + pub fn new_keyed(schema: SchemaRef, timestamp_index: usize, key_indices: Vec) -> Self { + Self { + schema, + timestamp_index, + key_indices: Some(key_indices), + } + } + + /// Unkeyed stream when `timestamp_index` is already verified. + pub fn new_unkeyed(schema: SchemaRef, timestamp_index: usize) -> Self { + Self { + schema, + timestamp_index, + key_indices: None, + } + } + + // ======================================================================== + // Safe Builders (Dynamically resolves and validates indices) + // ======================================================================== + + /// Unkeyed stream from a field list. Replaces the old `unwrap_or(0)` default when the timestamp + /// column was missing (silent wrong index / corruption). + pub fn try_from_fields(fields: impl Into>) -> Result { + let schema = Arc::new(Schema::new(fields.into())); + Self::try_from_schema_unkeyed(schema) + } + + /// Keyed stream from `SchemaRef`; resolves and validates the mandatory timestamp column. + pub fn try_from_schema_keyed(schema: SchemaRef, key_indices: Vec) -> Result { + let timestamp_index = schema + .column_with_name(TIMESTAMP_FIELD) + .ok_or_else(|| { + DataFusionError::Plan(format!( + "Streaming Topology Error: Mandatory event-time field '{}' is missing in the schema. \ + Current schema fields: {:?}", + TIMESTAMP_FIELD, + schema.fields() + )) + })? + .0; + + Ok(Self { + schema, + timestamp_index, + key_indices: Some(key_indices), + }) + } + + /// Unkeyed stream from `SchemaRef`; resolves and validates the mandatory timestamp column. + pub fn try_from_schema_unkeyed(schema: SchemaRef) -> Result { + let timestamp_index = schema + .column_with_name(TIMESTAMP_FIELD) + .ok_or_else(|| { + DataFusionError::Plan(format!( + "Streaming Topology Error: Mandatory event-time field '{}' is missing.", + TIMESTAMP_FIELD + )) + })? + .0; + + Ok(Self { + schema, + timestamp_index, + key_indices: None, + }) + } + + // ======================================================================== + // Zero-cost Getters + // ======================================================================== + + /// Underlying Arrow schema. + #[inline] + pub fn arrow_schema(&self) -> &SchemaRef { + &self.schema + } + + /// Physical column index used as event time / watermark driver. + #[inline] + pub fn timestamp_index(&self) -> usize { + self.timestamp_index + } + + /// Key column indices for shuffle / state, if keyed. + #[inline] + pub fn key_indices(&self) -> Option<&[usize]> { + self.key_indices.as_deref() + } + + #[inline] + pub fn is_keyed(&self) -> bool { + self.key_indices.is_some() + } +} diff --git a/src/streaming_planner/src/types/window.rs b/src/streaming_planner/src/types/window.rs new file mode 100644 index 00000000..69401f51 --- /dev/null +++ b/src/streaming_planner/src/types/window.rs @@ -0,0 +1,134 @@ +// 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::time::Duration; + +use datafusion::common::{Result, ScalarValue, not_impl_err, plan_err}; +use datafusion::logical_expr::Expr; +use datafusion::logical_expr::expr::{Alias, ScalarFunction}; + +use crate::common::constants::window_fn; + +use super::QualifiedField; + +// ============================================================================ +// Window Definitions +// ============================================================================ + +/// Temporal windowing semantics for streaming aggregations. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub enum WindowType { + Tumbling { width: Duration }, + Sliding { width: Duration, slide: Duration }, + Session { gap: Duration }, + Instant, +} + +/// How windowing is represented in the physical plan. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum WindowBehavior { + FromOperator { + window: WindowType, + window_field: QualifiedField, + window_index: usize, + is_nested: bool, + }, + InData, +} + +// ============================================================================ +// Logical Expression Parsers +// ============================================================================ + +pub fn extract_duration(expression: &Expr) -> Result { + match expression { + Expr::Literal(ScalarValue::IntervalDayTime(Some(val)), _) => { + let secs = (val.days as u64) * 24 * 60 * 60; + let millis = val.milliseconds as u64; + Ok(Duration::from_secs(secs) + Duration::from_millis(millis)) + } + Expr::Literal(ScalarValue::IntervalMonthDayNano(Some(val)), _) => { + if val.months != 0 { + return not_impl_err!( + "Streaming engine does not support window durations specified in months due to variable month lengths." + ); + } + let secs = (val.days as u64) * 24 * 60 * 60; + let nanos = val.nanoseconds as u64; + Ok(Duration::from_secs(secs) + Duration::from_nanos(nanos)) + } + _ => plan_err!( + "Unsupported window duration expression. Expected an interval literal (e.g., INTERVAL '1' MINUTE), got: {}", + expression + ), + } +} + +pub fn extract_window_type(expression: &Expr) -> Result> { + match expression { + Expr::ScalarFunction(ScalarFunction { func, args }) => match func.name() { + name if name == window_fn::HOP => { + if args.len() != 2 { + return plan_err!( + "hop() window function expects exactly 2 arguments (slide, width), got {}", + args.len() + ); + } + + let slide = extract_duration(&args[0])?; + let width = extract_duration(&args[1])?; + + if width.as_nanos() % slide.as_nanos() != 0 { + return plan_err!( + "Streaming Topology Error: hop() window width {:?} must be a perfect multiple of slide {:?}", + width, + slide + ); + } + + if slide == width { + Ok(Some(WindowType::Tumbling { width })) + } else { + Ok(Some(WindowType::Sliding { width, slide })) + } + } + + name if name == window_fn::TUMBLE => { + if args.len() != 1 { + return plan_err!( + "tumble() window function expects exactly 1 argument (width), got {}", + args.len() + ); + } + let width = extract_duration(&args[0])?; + Ok(Some(WindowType::Tumbling { width })) + } + + name if name == window_fn::SESSION => { + if args.len() != 1 { + return plan_err!( + "session() window function expects exactly 1 argument (gap), got {}", + args.len() + ); + } + let gap = extract_duration(&args[0])?; + Ok(Some(WindowType::Session { gap })) + } + + _ => Ok(None), + }, + + Expr::Alias(Alias { expr, .. }) => extract_window_type(expr), + + _ => Ok(None), + } +} From d80f700cfbc6a999b29d7eb72f0ea976edb78713 Mon Sep 17 00:00:00 2001 From: luoluoyuyu Date: Tue, 12 May 2026 22:15:16 +0800 Subject: [PATCH 5/7] =?UTF-8?q?refactor(workspace):=20=E4=B8=BB=E5=8C=85?= =?UTF-8?q?=E8=BF=81=E5=85=A5=20src/function-stream=20=E5=B9=B6=E9=9B=86?= =?UTF-8?q?=E4=B8=AD=20workspace=20=E5=85=83=E6=95=B0=E6=8D=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 根 Cargo.toml 改为纯 workspace，主程序与库移至 src/function-stream - 新增 [workspace.package] 统一 version / edition，成员 crate 使用 *.workspace = true - Makefile：VERSION 从根 Cargo.toml 读取，构建显式 -p function-stream，lite 使用 OPTIMIZE_FLAGS - wasm_host bindgen WIT 路径改为 ../../wit，以适配子目录中的 package manifest - 移除根 src 下旧入口与 runtime/storage 模块，streaming/wasm 等 crate 路径与导入同步调整 - examples-validator 与全仓 edition 对齐为 2024 Co-authored-by: Cursor --- Cargo.lock | 4 +- Cargo.toml | 91 +------------------ Makefile | 7 +- cli/cli/Cargo.toml | 4 +- examples/examples-validator/Cargo.toml | 2 +- protocol/Cargo.toml | 4 +- src/catalog/Cargo.toml | 4 +- src/catalog_storage/Cargo.toml | 4 +- src/catalog_storage/src/lib.rs | 2 +- .../src/stream_catalog/manager.rs | 4 +- .../src/stream_catalog/rocksdb_meta_store.rs | 2 +- src/catalog_storage/src/task/proto_codec.rs | 2 +- .../src/task/rocksdb_storage.rs | 2 +- src/catalog_storage/src/task/storage.rs | 2 +- src/common/Cargo.toml | 4 +- src/config/Cargo.toml | 4 +- src/coordinator/Cargo.toml | 4 +- .../legacy/dataset/show_functions_result.rs | 2 +- .../dataset/show_streaming_tables_result.rs | 2 +- .../src/legacy/execution/executor.rs | 8 +- src/coordinator/src/legacy/runtime_context.rs | 6 +- src/function-stream/Cargo.toml | 86 ++++++++++++++++++ .../src/lib.rs} | 33 ++++++- src/{ => function-stream/src}/main.rs | 66 ++++++++++++-- src/job_manager/Cargo.toml | 4 +- src/lib.rs | 25 ----- src/logger/Cargo.toml | 4 +- src/runtime.rs | 27 ------ src/runtime_common/Cargo.toml | 4 +- src/servicer/Cargo.toml | 4 +- src/servicer/src/legacy/initializer.rs | 28 +++--- src/servicer/src/legacy/memory_service.rs | 2 +- src/sqlparser/Cargo.toml | 4 +- src/streaming_planner/Cargo.toml | 4 +- src/streaming_runtime/Cargo.toml | 4 +- src/streaming_runtime/src/lib.rs | 2 +- .../src/streaming/api/context.rs | 10 +- .../src/streaming/api/operator.rs | 6 +- .../src/streaming/api/source.rs | 2 +- .../src/streaming/execution/operator_chain.rs | 8 +- .../src/streaming/execution/pipeline.rs | 14 +-- .../src/streaming/execution/source_driver.rs | 10 +- .../execution/tracker/watermark_tracker.rs | 2 +- .../src/streaming/factory/connector/delta.rs | 12 +-- .../factory/connector/dispatchers.rs | 6 +- .../streaming/factory/connector/filesystem.rs | 10 +- .../streaming/factory/connector/iceberg.rs | 12 +-- .../src/streaming/factory/connector/kafka.rs | 14 +-- .../streaming/factory/connector/lancedb.rs | 10 +- .../src/streaming/factory/connector/s3.rs | 10 +- .../streaming/factory/operator_constructor.rs | 4 +- .../src/streaming/factory/operator_factory.rs | 18 ++-- .../src/streaming/job/edge_manager.rs | 2 +- .../src/streaming/job/job_manager.rs | 32 +++---- .../src/streaming/job/models.rs | 2 +- .../src/streaming/network/endpoint.rs | 4 +- .../grouping/incremental_aggregate.rs | 14 +-- .../operators/joins/join_instance.rs | 10 +- .../operators/joins/join_with_expiration.rs | 10 +- .../src/streaming/operators/key_by.rs | 6 +- .../src/streaming/operators/key_operator.rs | 10 +- .../src/streaming/operators/projection.rs | 8 +- .../src/streaming/operators/sink/delta/mod.rs | 10 +- .../operators/sink/filesystem/mod.rs | 10 +- .../streaming/operators/sink/iceberg/mod.rs | 10 +- .../src/streaming/operators/sink/kafka/mod.rs | 8 +- .../streaming/operators/sink/lancedb/mod.rs | 6 +- .../src/streaming/operators/sink/s3/mod.rs | 6 +- .../streaming/operators/source/kafka/mod.rs | 6 +- .../operators/stateless_physical_executor.rs | 2 +- .../streaming/operators/value_execution.rs | 8 +- .../watermark/watermark_generator.rs | 8 +- .../windows/session_aggregating_window.rs | 10 +- .../windows/sliding_aggregating_window.rs | 10 +- .../windows/tumbling_aggregating_window.rs | 10 +- .../operators/windows/window_function.rs | 10 +- .../src/streaming/protocol/event.rs | 2 +- .../src/streaming/state/error.rs | 2 +- .../src/streaming/state/operator_state.rs | 6 +- src/wasm_runtime/Cargo.toml | 4 +- src/wasm_runtime/src/lib.rs | 6 +- src/wasm_runtime/src/state_backend/factory.rs | 10 +- .../src/state_backend/memory/factory.rs | 6 +- .../src/state_backend/memory/store.rs | 6 +- .../src/state_backend/rocksdb/factory.rs | 10 +- .../src/state_backend/rocksdb/store.rs | 6 +- src/wasm_runtime/src/state_backend/server.rs | 6 +- src/wasm_runtime/src/state_backend/store.rs | 16 ++-- .../src/wasm/input/input_protocol.rs | 2 +- .../src/wasm/input/input_provider.rs | 8 +- .../src/wasm/input/input_runner.rs | 18 ++-- src/wasm_runtime/src/wasm/input/interface.rs | 6 +- .../input/protocol/kafka/kafka_protocol.rs | 4 +- src/wasm_runtime/src/wasm/output/interface.rs | 4 +- .../src/wasm/output/output_protocol.rs | 2 +- .../src/wasm/output/output_provider.rs | 8 +- .../src/wasm/output/output_runner.rs | 18 ++-- .../output/protocol/kafka/kafka_protocol.rs | 4 +- .../src/wasm/processor/wasm/thread_pool.rs | 4 +- .../src/wasm/processor/wasm/wasm_host.rs | 28 +++--- .../src/wasm/processor/wasm/wasm_processor.rs | 6 +- .../processor/wasm/wasm_processor_trait.rs | 4 +- .../src/wasm/processor/wasm/wasm_task.rs | 24 ++--- .../src/wasm/task/builder/processor/mod.rs | 14 +-- .../src/wasm/task/builder/python/mod.rs | 18 ++-- .../src/wasm/task/builder/sink/mod.rs | 4 +- .../src/wasm/task/builder/source/mod.rs | 4 +- .../src/wasm/task/builder/task_builder.rs | 14 +-- .../src/wasm/task/control_mailbox.rs | 4 +- src/wasm_runtime/src/wasm/task/lifecycle.rs | 8 +- .../src/wasm/task/processor_config.rs | 2 +- .../src/wasm/taskexecutor/init_context.rs | 8 +- .../src/wasm/taskexecutor/task_manager.rs | 12 +-- 113 files changed, 581 insertions(+), 548 deletions(-) create mode 100644 src/function-stream/Cargo.toml rename src/{storage.rs => function-stream/src/lib.rs} (71%) rename src/{ => function-stream/src}/main.rs (80%) delete mode 100644 src/lib.rs delete mode 100644 src/runtime.rs diff --git a/Cargo.lock b/Cargo.lock index 97388459..e3e4b95a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3861,7 +3861,7 @@ dependencies = [ [[package]] name = "function-stream-cli" -version = "0.1.0" +version = "0.6.0" dependencies = [ "arrow-array 52.2.0", "arrow-ipc 52.2.0", @@ -6554,7 +6554,7 @@ dependencies = [ [[package]] name = "protocol" -version = "0.1.0" +version = "0.6.0" dependencies = [ "env_logger", "log", diff --git a/Cargo.toml b/Cargo.toml index 71c795b6..96adb4b3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [workspace] members = [ - ".", + "src/function-stream", "protocol", "cli/cli", "src/catalog", @@ -17,91 +17,10 @@ members = [ "src/streaming_planner", "src/wasm_runtime", ] +resolver = "2" +default-members = ["src/function-stream"] -[package] -name = "function-stream" +# Shared crate metadata for workspace members (`version.workspace = true`, etc.). +[workspace.package] version = "0.6.0" edition = "2024" - -[lib] -name = "function_stream" -path = "src/lib.rs" - -[[bin]] -name = "function-stream" -path = "src/main.rs" - - -[dependencies] -tokio = { version = "1.0", features = ["macros", "rt-multi-thread", "sync", "time", "net", "signal"] } -serde = { version = "1.0", features = ["derive"] } -serde_yaml = "0.9" -serde_json = "1.0" -uuid = { version = "1.0", features = ["v4", "v7"] } -log = "0.4" -tracing = "0.1" -tracing-subscriber = { version = "0.3", features = ["env-filter"] } -tracing-appender = "0.2" -anyhow = "1.0" -thiserror = "2" -tonic = { version = "0.12", features = ["default"] } -async-trait = "0.1" -num_cpus = "1.0" -protocol = { path = "./protocol" } -function-stream-config = { path = "src/config" } -function-stream-logger = { path = "src/logger" } -function-stream-runtime-common = { path = "src/runtime_common" } -function-stream-streaming-planner = { path = "src/streaming_planner" } -prost = "0.13" -rdkafka = { version = "0.38", features = ["cmake-build", "ssl", "gssapi", "curl"] } -crossbeam-channel = "0.5" -wasmtime = { version = "41.0.3", features = ["component-model", "async"] } -base64 = "0.22" -wasmtime-wasi = "41.0.3" -rocksdb = { version = "0.21", features = ["multi-threaded-cf", "lz4"] } -bincode = { version = "2", features = ["serde"] } -chrono = "0.4" -tokio-stream = "0.1.18" -lru = "0.12" -parking_lot = "0.12" -arrow = { version = "55", default-features = false } -arrow-array = "55" -arrow-ipc = "55" -arrow-schema = { version = "55", features = ["serde"] } -parquet = "55" -object_store = { version = "0.12.5", features = ["aws"] } -bytes = "1" -futures = "0.3" -serde_json_path = "0.7" -xxhash-rust = { version = "0.8", features = ["xxh3"] } -proctitle = "0.1" -unicase = "2.7" -petgraph = "0.7" -rand = { version = "0.8", features = ["small_rng"] } -itertools = "0.14" -strum = { version = "0.26", features = ["derive"] } - -arrow-json = {version = '55.2.0'} -apache-avro = "0.21" -datafusion = {git = 'https://github.com/FunctionStream/datafusion', branch = '48.0.1/fs'} -datafusion-common = {git = 'https://github.com/FunctionStream/datafusion', branch = '48.0.1/fs'} -datafusion-execution = {git = 'https://github.com/FunctionStream/datafusion', branch = '48.0.1/fs'} -datafusion-expr = {git = 'https://github.com/FunctionStream/datafusion', branch = '48.0.1/fs'} -datafusion-physical-expr = {git = 'https://github.com/FunctionStream/datafusion', branch = '48.0.1/fs'} -datafusion-proto = {git = 'https://github.com/FunctionStream/datafusion', branch = '48.0.1/fs'} - -sqlparser = { git = "https://github.com/FunctionStream/sqlparser-rs", branch = "0.58.0/fs" } - -ahash = "0.8" -governor = "0.8.0" -lance = { version = "4.0.0", default-features = false, features = ["aws"] } -arrow-array-lance = { package = "arrow-array", version = "57.3.0" } -arrow-ipc-lance = { package = "arrow-ipc", version = "57.3.0" } - -[features] -default = ["incremental-cache", "python"] -incremental-cache = ["wasmtime/incremental-cache"] -python = [] - -[dev-dependencies] -tempfile = "3.27.0" diff --git a/Makefile b/Makefile index 78138dae..d6340d9d 100644 --- a/Makefile +++ b/Makefile @@ -10,6 +10,7 @@ # See the License for the specific language governing permissions and # limitations under the License. APP_NAME := function-stream +# Version from root `[workspace.package]` (single source of truth). VERSION := $(shell grep '^version' Cargo.toml | head -1 | awk -F '"' '{print $$2}') DATE := $(shell date -u +"%Y-%m-%dT%H:%M:%SZ") @@ -106,6 +107,7 @@ build: .check-env .ensure-target .build-wasm @RUSTFLAGS="$(OPTIMIZE_FLAGS)" \ cargo build --release \ --target $(TRIPLE) \ + -p $(APP_NAME) \ --features python \ --quiet $(call log,BUILD,CLI) @@ -118,14 +120,15 @@ build: .check-env .ensure-target .build-wasm build-lite: .check-env .ensure-target $(call log,BUILD,Rust Lite [$(OS_NAME) / $(TRIPLE)]) - @RUSTFLAGS="$(INDUSTRIAL_RUSTFLAGS)" \ + @RUSTFLAGS="$(OPTIMIZE_FLAGS)" \ cargo build --release \ --target $(TRIPLE) \ + -p $(APP_NAME) \ --no-default-features \ --features incremental-cache \ --quiet $(call log,BUILD,CLI for dist) - @RUSTFLAGS="$(INDUSTRIAL_RUSTFLAGS)" \ + @RUSTFLAGS="$(OPTIMIZE_FLAGS)" \ cargo build --release \ --target $(TRIPLE) \ -p function-stream-cli \ diff --git a/cli/cli/Cargo.toml b/cli/cli/Cargo.toml index 49c0a881..75bd6cb4 100644 --- a/cli/cli/Cargo.toml +++ b/cli/cli/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "function-stream-cli" -version = "0.1.0" -edition = "2021" +version.workspace = true +edition.workspace = true [[bin]] name = "cli" diff --git a/examples/examples-validator/Cargo.toml b/examples/examples-validator/Cargo.toml index de87dbb3..9c204935 100644 --- a/examples/examples-validator/Cargo.toml +++ b/examples/examples-validator/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "examples-validator" version = "0.1.0" -edition = "2021" +edition = "2024" [workspace] diff --git a/protocol/Cargo.toml b/protocol/Cargo.toml index 51b1f3c1..96501c30 100644 --- a/protocol/Cargo.toml +++ b/protocol/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "protocol" -version = "0.1.0" -edition = "2024" +version.workspace = true +edition.workspace = true description = "Protocol Buffers protocol definitions for function stream" license = "MIT OR Apache-2.0" repository = "https://github.com/your-username/rust-function-stream" diff --git a/src/catalog/Cargo.toml b/src/catalog/Cargo.toml index 07437967..2a50735c 100644 --- a/src/catalog/Cargo.toml +++ b/src/catalog/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "function-stream-catalog" -version = "0.6.0" -edition = "2024" +version.workspace = true +edition.workspace = true [lib] name = "function_stream_catalog" diff --git a/src/catalog_storage/Cargo.toml b/src/catalog_storage/Cargo.toml index f76fb9d3..3bb9fed8 100644 --- a/src/catalog_storage/Cargo.toml +++ b/src/catalog_storage/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "function-stream-catalog-storage" -version = "0.6.0" -edition = "2024" +version.workspace = true +edition.workspace = true [lib] name = "function_stream_catalog_storage" diff --git a/src/catalog_storage/src/lib.rs b/src/catalog_storage/src/lib.rs index 4aa4cf8e..8fdca0bf 100644 --- a/src/catalog_storage/src/lib.rs +++ b/src/catalog_storage/src/lib.rs @@ -1,7 +1,7 @@ //! Persistent catalog storage implementations. //! //! The stream catalog manager and task persistence (`stream_catalog/`, `task/`) live in this -//! package and are compiled as part of `function-stream` via `#[path]` in `src/storage.rs`. +//! package and are compiled as part of `function-stream` via `#[path]` in `src/lib.rs` / `src/main.rs`. pub mod memory; pub mod rocksdb; diff --git a/src/catalog_storage/src/stream_catalog/manager.rs b/src/catalog_storage/src/stream_catalog/manager.rs index fa810e8d..cbdf30b3 100644 --- a/src/catalog_storage/src/stream_catalog/manager.rs +++ b/src/catalog_storage/src/stream_catalog/manager.rs @@ -627,7 +627,7 @@ pub fn restore_global_catalog_from_store() { } pub fn restore_streaming_jobs_from_store() { - use crate::runtime::streaming::job::JobManager; + use crate::streaming::job::JobManager; let Some(catalog) = CatalogManager::try_global() else { warn!("CatalogManager not available; skipping streaming job restore"); @@ -742,7 +742,7 @@ mod tests { use crate::sql::schema::column_descriptor::ColumnDescriptor; use crate::sql::schema::table::CatalogEntity; use crate::sql::schema::temporal_pipeline_config::TemporalPipelineConfig; - use crate::storage::stream_catalog::InMemoryMetaStore; + use crate::stream_catalog::InMemoryMetaStore; use super::CatalogManager; diff --git a/src/catalog_storage/src/stream_catalog/rocksdb_meta_store.rs b/src/catalog_storage/src/stream_catalog/rocksdb_meta_store.rs index 1537f278..58c0b35c 100644 --- a/src/catalog_storage/src/stream_catalog/rocksdb_meta_store.rs +++ b/src/catalog_storage/src/stream_catalog/rocksdb_meta_store.rs @@ -20,7 +20,7 @@ use rocksdb::{DB, Direction, IteratorMode, Options, WriteBatch}; use super::MetaStore; -/// Single-node durable KV used by [`crate::storage::stream_catalog::CatalogManager`]. +/// Single-node durable KV used by [`crate::stream_catalog::CatalogManager`]. pub struct RocksDbMetaStore { db: Arc, } diff --git a/src/catalog_storage/src/task/proto_codec.rs b/src/catalog_storage/src/task/proto_codec.rs index 6c1bc8df..78a0426d 100644 --- a/src/catalog_storage/src/task/proto_codec.rs +++ b/src/catalog_storage/src/task/proto_codec.rs @@ -20,7 +20,7 @@ use protocol::storage::{ }; use serde::{Deserialize, Serialize}; -use crate::runtime::common::ComponentState; +use crate::common::ComponentState; use super::storage::TaskModuleBytes; diff --git a/src/catalog_storage/src/task/rocksdb_storage.rs b/src/catalog_storage/src/task/rocksdb_storage.rs index cea0ceb9..62ffb2a0 100644 --- a/src/catalog_storage/src/task/rocksdb_storage.rs +++ b/src/catalog_storage/src/task/rocksdb_storage.rs @@ -20,7 +20,7 @@ use super::proto_codec::{ }; use super::storage::{StoredTaskInfo, TaskStorage}; use crate::config::storage::RocksDBStorageConfig; -use crate::runtime::common::ComponentState; +use crate::common::ComponentState; use anyhow::{Context, Result, anyhow}; use rocksdb::{ColumnFamilyDescriptor, DB, IteratorMode, Options, WriteBatch}; use std::path::Path; diff --git a/src/catalog_storage/src/task/storage.rs b/src/catalog_storage/src/task/storage.rs index 156ee5d8..6cbec9ea 100644 --- a/src/catalog_storage/src/task/storage.rs +++ b/src/catalog_storage/src/task/storage.rs @@ -10,7 +10,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::runtime::common::ComponentState; +use crate::common::ComponentState; use anyhow::Result; use serde::{Deserialize, Serialize}; diff --git a/src/common/Cargo.toml b/src/common/Cargo.toml index fb68a526..4c155c5e 100644 --- a/src/common/Cargo.toml +++ b/src/common/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "function-stream-common" -version = "0.6.0" -edition = "2024" +version.workspace = true +edition.workspace = true [lib] name = "function_stream_common" diff --git a/src/config/Cargo.toml b/src/config/Cargo.toml index edd2a8db..d1e437cd 100644 --- a/src/config/Cargo.toml +++ b/src/config/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "function-stream-config" -version = "0.6.0" -edition = "2024" +version.workspace = true +edition.workspace = true [lib] name = "function_stream_config" diff --git a/src/coordinator/Cargo.toml b/src/coordinator/Cargo.toml index 5322f8a2..f8a6e0b7 100644 --- a/src/coordinator/Cargo.toml +++ b/src/coordinator/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "function-stream-coordinator" -version = "0.6.0" -edition = "2024" +version.workspace = true +edition.workspace = true [lib] name = "function_stream_coordinator" diff --git a/src/coordinator/src/legacy/dataset/show_functions_result.rs b/src/coordinator/src/legacy/dataset/show_functions_result.rs index c16edf6d..18673f00 100644 --- a/src/coordinator/src/legacy/dataset/show_functions_result.rs +++ b/src/coordinator/src/legacy/dataset/show_functions_result.rs @@ -16,7 +16,7 @@ use arrow_array::{RecordBatch, StringArray}; use arrow_schema::{DataType, Field, Schema}; use super::DataSet; -use crate::storage::task::FunctionInfo; +use crate::task::FunctionInfo; #[derive(Clone, Debug)] pub struct ShowFunctionsResult { diff --git a/src/coordinator/src/legacy/dataset/show_streaming_tables_result.rs b/src/coordinator/src/legacy/dataset/show_streaming_tables_result.rs index cae597ac..5b5c8e84 100644 --- a/src/coordinator/src/legacy/dataset/show_streaming_tables_result.rs +++ b/src/coordinator/src/legacy/dataset/show_streaming_tables_result.rs @@ -16,7 +16,7 @@ use arrow_array::{Int32Array, StringArray}; use arrow_schema::{DataType, Field, Schema}; use super::DataSet; -use crate::runtime::streaming::job::StreamingJobSummary; +use crate::streaming::job::StreamingJobSummary; #[derive(Clone, Debug)] pub struct ShowStreamingTablesResult { diff --git a/src/coordinator/src/legacy/execution/executor.rs b/src/coordinator/src/legacy/execution/executor.rs index 22c5c1b7..7109a1d2 100644 --- a/src/coordinator/src/legacy/execution/executor.rs +++ b/src/coordinator/src/legacy/execution/executor.rs @@ -31,13 +31,13 @@ use crate::coordinator::statement::{ConfigSource, FunctionSource}; use crate::coordinator::streaming_table_options::{ parse_checkpoint_interval_ms, parse_pipeline_parallelism, }; -use crate::runtime::streaming::job::JobManager; -use crate::runtime::streaming::protocol::control::StopMode; -use crate::runtime::wasm::taskexecutor::TaskManager; +use crate::streaming::job::JobManager; +use crate::streaming::protocol::control::StopMode; +use crate::wasm::taskexecutor::TaskManager; use crate::sql::schema::catalog::ExternalTable; use crate::sql::schema::show_create_catalog_table; use crate::sql::schema::table::CatalogEntity; -use crate::storage::stream_catalog::CatalogManager; +use crate::stream_catalog::CatalogManager; #[derive(Error, Debug)] pub enum ExecuteError { diff --git a/src/coordinator/src/legacy/runtime_context.rs b/src/coordinator/src/legacy/runtime_context.rs index 21b9d876..513c703c 100644 --- a/src/coordinator/src/legacy/runtime_context.rs +++ b/src/coordinator/src/legacy/runtime_context.rs @@ -16,10 +16,10 @@ use std::sync::Arc; use anyhow::Result; -use crate::runtime::streaming::job::JobManager; -use crate::runtime::wasm::taskexecutor::TaskManager; +use crate::streaming::job::JobManager; +use crate::wasm::taskexecutor::TaskManager; use crate::sql::schema::StreamSchemaProvider; -use crate::storage::stream_catalog::CatalogManager; +use crate::stream_catalog::CatalogManager; /// Dependencies shared by analyze / plan / execute, analogous to installing globals in /// [`TaskManager`], [`CatalogManager`], and [`JobManager`]. diff --git a/src/function-stream/Cargo.toml b/src/function-stream/Cargo.toml new file mode 100644 index 00000000..baf9d5ea --- /dev/null +++ b/src/function-stream/Cargo.toml @@ -0,0 +1,86 @@ +[package] +name = "function-stream" +version.workspace = true +edition.workspace = true + +[lib] +name = "function_stream" +path = "src/lib.rs" + +[[bin]] +name = "function-stream" +path = "src/main.rs" + +[dependencies] +tokio = { version = "1.0", features = ["macros", "rt-multi-thread", "sync", "time", "net", "signal"] } +serde = { version = "1.0", features = ["derive"] } +serde_yaml = "0.9" +serde_json = "1.0" +uuid = { version = "1.0", features = ["v4", "v7"] } +log = "0.4" +tracing = "0.1" +tracing-subscriber = { version = "0.3", features = ["env-filter"] } +tracing-appender = "0.2" +anyhow = "1.0" +thiserror = "2" +tonic = { version = "0.12", features = ["default"] } +async-trait = "0.1" +num_cpus = "1.0" +protocol = { path = "../../protocol" } +function-stream-config = { path = "../config" } +function-stream-logger = { path = "../logger" } +function-stream-runtime-common = { path = "../runtime_common" } +function-stream-streaming-planner = { path = "../streaming_planner" } +prost = "0.13" +rdkafka = { version = "0.38", features = ["cmake-build", "ssl", "gssapi", "curl"] } +crossbeam-channel = "0.5" +wasmtime = { version = "41.0.3", features = ["component-model", "async"] } +base64 = "0.22" +wasmtime-wasi = "41.0.3" +rocksdb = { version = "0.21", features = ["multi-threaded-cf", "lz4"] } +bincode = { version = "2", features = ["serde"] } +chrono = "0.4" +tokio-stream = "0.1.18" +lru = "0.12" +parking_lot = "0.12" +arrow = { version = "55", default-features = false } +arrow-array = "55" +arrow-ipc = "55" +arrow-schema = { version = "55", features = ["serde"] } +parquet = "55" +object_store = { version = "0.12.5", features = ["aws"] } +bytes = "1" +futures = "0.3" +serde_json_path = "0.7" +xxhash-rust = { version = "0.8", features = ["xxh3"] } +proctitle = "0.1" +unicase = "2.7" +petgraph = "0.7" +rand = { version = "0.8", features = ["small_rng"] } +itertools = "0.14" +strum = { version = "0.26", features = ["derive"] } + +arrow-json = { version = "55.2.0" } +apache-avro = "0.21" +datafusion = { git = "https://github.com/FunctionStream/datafusion", branch = "48.0.1/fs" } +datafusion-common = { git = "https://github.com/FunctionStream/datafusion", branch = "48.0.1/fs" } +datafusion-execution = { git = "https://github.com/FunctionStream/datafusion", branch = "48.0.1/fs" } +datafusion-expr = { git = "https://github.com/FunctionStream/datafusion", branch = "48.0.1/fs" } +datafusion-physical-expr = { git = "https://github.com/FunctionStream/datafusion", branch = "48.0.1/fs" } +datafusion-proto = { git = "https://github.com/FunctionStream/datafusion", branch = "48.0.1/fs" } + +sqlparser = { git = "https://github.com/FunctionStream/sqlparser-rs", branch = "0.58.0/fs" } + +ahash = "0.8" +governor = "0.8.0" +lance = { version = "4.0.0", default-features = false, features = ["aws"] } +arrow-array-lance = { package = "arrow-array", version = "57.3.0" } +arrow-ipc-lance = { package = "arrow-ipc", version = "57.3.0" } + +[features] +default = ["incremental-cache", "python"] +incremental-cache = ["wasmtime/incremental-cache"] +python = [] + +[dev-dependencies] +tempfile = "3.27.0" diff --git a/src/storage.rs b/src/function-stream/src/lib.rs similarity index 71% rename from src/storage.rs rename to src/function-stream/src/lib.rs index 713b149e..fac7e8f1 100644 --- a/src/storage.rs +++ b/src/function-stream/src/lib.rs @@ -10,20 +10,39 @@ // See the License for the specific language governing permissions and // limitations under the License. -// Persistent / catalog-related modules compiled into the root crate. -// Paths are relative to `src/` (this file lives at `src/storage.rs`). +// Library crate for function-stream + +#![allow(dead_code)] use std::sync::Arc; use anyhow::Context; -#[path = "wasm_runtime/src/state_backend/mod.rs"] +pub use function_stream_config as config; +#[path = "../../coordinator/src/legacy/mod.rs"] +pub mod coordinator; +pub use function_stream_logger as logging; + +pub use function_stream_runtime_common::{common, memory}; + +#[path = "../../streaming_runtime/src/streaming/mod.rs"] +pub mod streaming; + +#[path = "../../streaming_runtime/src/util/mod.rs"] +pub mod util; + +#[path = "../../wasm_runtime/src/wasm/mod.rs"] +pub mod wasm; + +pub use wasm::{input, output, processor}; + +#[path = "../../wasm_runtime/src/state_backend/mod.rs"] pub mod state_backend; -#[path = "catalog_storage/src/stream_catalog/mod.rs"] +#[path = "../../catalog_storage/src/stream_catalog/mod.rs"] pub mod stream_catalog; -#[path = "catalog_storage/src/task/mod.rs"] +#[path = "../../catalog_storage/src/task/mod.rs"] pub mod task; /// Install the process-global [`stream_catalog::CatalogManager`] from configuration. @@ -59,3 +78,7 @@ pub fn initialize_stream_catalog(config: &crate::config::GlobalConfig) -> anyhow CatalogManager::init_global(store).context("Stream catalog (CatalogManager) global init failed") } + +#[path = "../../servicer/src/legacy/mod.rs"] +pub mod server; +pub use function_stream_streaming_planner as sql; diff --git a/src/main.rs b/src/function-stream/src/main.rs similarity index 80% rename from src/main.rs rename to src/function-stream/src/main.rs index eb5bcfec..33971370 100644 --- a/src/main.rs +++ b/src/function-stream/src/main.rs @@ -13,16 +13,70 @@ #![allow(dead_code)] pub use function_stream_config as config; -#[path = "coordinator/src/legacy/mod.rs"] +#[path = "../../coordinator/src/legacy/mod.rs"] mod coordinator; pub use function_stream_logger as logging; -mod runtime; -#[path = "servicer/src/legacy/mod.rs"] -mod server; -pub use function_stream_streaming_planner as sql; -mod storage; + +pub use function_stream_runtime_common::{common, memory}; + +use std::sync::Arc; use anyhow::{Context, Result}; + +#[path = "../../streaming_runtime/src/streaming/mod.rs"] +mod streaming; + +#[path = "../../streaming_runtime/src/util/mod.rs"] +mod util; + +#[path = "../../wasm_runtime/src/wasm/mod.rs"] +mod wasm; + +pub use wasm::{input, output, processor}; + +#[path = "../../wasm_runtime/src/state_backend/mod.rs"] +mod state_backend; + +#[path = "../../catalog_storage/src/stream_catalog/mod.rs"] +mod stream_catalog; + +#[path = "../../catalog_storage/src/task/mod.rs"] +mod task; + +pub fn initialize_stream_catalog(config: &crate::config::GlobalConfig) -> anyhow::Result<()> { + use stream_catalog::{CatalogManager, InMemoryMetaStore, MetaStore, RocksDbMetaStore}; + + let store: Arc = if !config.stream_catalog.persist { + Arc::new(InMemoryMetaStore::new()) + } else { + let path = config + .stream_catalog + .db_path + .as_ref() + .map(|p| crate::config::resolve_path(p)) + .unwrap_or_else(|| crate::config::get_data_dir().join("catalog.db")); + + std::fs::create_dir_all(&path).with_context(|| { + format!( + "Failed to create stream catalog RocksDB directory {}", + path.display() + ) + })?; + + Arc::new(RocksDbMetaStore::open(&path).with_context(|| { + format!( + "Failed to open stream catalog RocksDB at {}", + path.display() + ) + })?) + }; + + CatalogManager::init_global(store).context("Stream catalog (CatalogManager) global init failed") +} + +#[path = "../../servicer/src/legacy/mod.rs"] +mod server; +pub use function_stream_streaming_planner as sql; use std::thread; use tokio::sync::oneshot; diff --git a/src/job_manager/Cargo.toml b/src/job_manager/Cargo.toml index 1315a9e1..8edfb464 100644 --- a/src/job_manager/Cargo.toml +++ b/src/job_manager/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "function-stream-job-manager" -version = "0.6.0" -edition = "2024" +version.workspace = true +edition.workspace = true [lib] name = "function_stream_job_manager" diff --git a/src/lib.rs b/src/lib.rs deleted file mode 100644 index 177d8f72..00000000 --- a/src/lib.rs +++ /dev/null @@ -1,25 +0,0 @@ -// 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. - -// Library crate for function-stream - -#![allow(dead_code)] - -pub use function_stream_config as config; -#[path = "coordinator/src/legacy/mod.rs"] -pub mod coordinator; -pub use function_stream_logger as logging; -pub mod runtime; -#[path = "servicer/src/legacy/mod.rs"] -pub mod server; -pub use function_stream_streaming_planner as sql; -pub mod storage; diff --git a/src/logger/Cargo.toml b/src/logger/Cargo.toml index 8d584f74..273579a8 100644 --- a/src/logger/Cargo.toml +++ b/src/logger/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "function-stream-logger" -version = "0.6.0" -edition = "2024" +version.workspace = true +edition.workspace = true [lib] name = "function_stream_logger" diff --git a/src/runtime.rs b/src/runtime.rs deleted file mode 100644 index d08cb5c2..00000000 --- a/src/runtime.rs +++ /dev/null @@ -1,27 +0,0 @@ -// 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-tree runtime: streaming engine, util helpers, and WASM task runtime. -// Paths are relative to `src/` (this file lives at `src/runtime.rs`). - -pub use function_stream_runtime_common::{common, memory}; - -#[path = "streaming_runtime/src/streaming/mod.rs"] -pub mod streaming; - -#[path = "streaming_runtime/src/util/mod.rs"] -pub mod util; - -#[path = "wasm_runtime/src/wasm/mod.rs"] -pub mod wasm; - -pub use wasm::{input, output, processor}; diff --git a/src/runtime_common/Cargo.toml b/src/runtime_common/Cargo.toml index b7747b2a..f8e98ed6 100644 --- a/src/runtime_common/Cargo.toml +++ b/src/runtime_common/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "function-stream-runtime-common" -version = "0.6.0" -edition = "2024" +version.workspace = true +edition.workspace = true [lib] name = "function_stream_runtime_common" diff --git a/src/servicer/Cargo.toml b/src/servicer/Cargo.toml index 08860c22..16cd3109 100644 --- a/src/servicer/Cargo.toml +++ b/src/servicer/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "function-stream-servicer" -version = "0.6.0" -edition = "2024" +version.workspace = true +edition.workspace = true [lib] name = "function_stream_servicer" diff --git a/src/servicer/src/legacy/initializer.rs b/src/servicer/src/legacy/initializer.rs index 8a04608e..78b16c73 100644 --- a/src/servicer/src/legacy/initializer.rs +++ b/src/servicer/src/legacy/initializer.rs @@ -113,7 +113,7 @@ pub fn build_core_registry() -> ComponentRegistry { builder .register( "StreamCatalog", - crate::storage::stream_catalog::initialize_stream_catalog, + crate::stream_catalog::initialize_stream_catalog, ) .register("Coordinator", initialize_coordinator) .build() @@ -124,16 +124,16 @@ pub fn bootstrap_system(config: &GlobalConfig) -> Result<()> { registry.initialize_all(config)?; - crate::storage::stream_catalog::restore_global_catalog_from_store(); - crate::storage::stream_catalog::restore_streaming_jobs_from_store(); + crate::stream_catalog::restore_global_catalog_from_store(); + crate::stream_catalog::restore_streaming_jobs_from_store(); info!("System bootstrap finished. Node is ready to accept traffic."); Ok(()) } fn initialize_wasm_cache(config: &GlobalConfig) -> Result<()> { - crate::runtime::processor::wasm::wasm_cache::set_cache_config( - crate::runtime::processor::wasm::wasm_cache::WasmCacheConfig { + crate::processor::wasm::wasm_cache::set_cache_config( + crate::processor::wasm::wasm_cache::WasmCacheConfig { enabled: config.wasm.enable_cache, cache_dir: crate::config::paths::resolve_path(&config.wasm.cache_dir), max_size: config.wasm.max_cache_size, @@ -151,14 +151,14 @@ fn initialize_wasm_cache(config: &GlobalConfig) -> Result<()> { } fn initialize_task_manager(config: &GlobalConfig) -> Result<()> { - crate::runtime::wasm::taskexecutor::TaskManager::init(config) + crate::wasm::taskexecutor::TaskManager::init(config) .context("TaskManager service failed to start")?; Ok(()) } #[cfg(feature = "python")] fn initialize_python_service(config: &GlobalConfig) -> Result<()> { - crate::runtime::processor::python::PythonService::initialize(config) + crate::processor::python::PythonService::initialize(config) .context("Python Runtime initialization failed")?; Ok(()) } @@ -168,9 +168,9 @@ fn initialize_memory_service(config: &GlobalConfig) -> Result<()> { } fn initialize_job_manager(config: &GlobalConfig) -> Result<()> { - use crate::runtime::streaming::factory::OperatorFactory; - use crate::runtime::streaming::factory::Registry; - use crate::runtime::streaming::job::{JobManager, StateConfig}; + use crate::streaming::factory::OperatorFactory; + use crate::streaming::factory::Registry; + use crate::streaming::job::{JobManager, StateConfig}; use std::sync::Arc; let per_operator_memory_bytes = config @@ -199,16 +199,16 @@ fn initialize_job_manager(config: &GlobalConfig) -> Result<()> { } fn initialize_coordinator(_config: &GlobalConfig) -> Result<()> { - crate::runtime::wasm::taskexecutor::TaskManager::get() + crate::wasm::taskexecutor::TaskManager::get() .context("Dependency violation: Coordinator requires TaskManager")?; - crate::runtime::memory::try_global_memory_pool() + crate::memory::try_global_memory_pool() .context("Dependency violation: Coordinator requires MemoryService")?; - crate::storage::stream_catalog::CatalogManager::global() + crate::stream_catalog::CatalogManager::global() .context("Dependency violation: Coordinator requires StreamCatalog")?; - crate::runtime::streaming::job::JobManager::global() + crate::streaming::job::JobManager::global() .context("Dependency violation: Coordinator requires JobManager")?; Ok(()) diff --git a/src/servicer/src/legacy/memory_service.rs b/src/servicer/src/legacy/memory_service.rs index 2ba24eee..5c03350b 100644 --- a/src/servicer/src/legacy/memory_service.rs +++ b/src/servicer/src/legacy/memory_service.rs @@ -52,7 +52,7 @@ impl MemoryService { let total_pool_bytes = streaming_runtime_memory_bytes.saturating_add(operator_state_store_memory_bytes); - crate::runtime::memory::init_global_memory_pool(total_pool_bytes) + crate::memory::init_global_memory_pool(total_pool_bytes) .context("Global memory pool initialization failed")?; info!("MemoryService initialized"); diff --git a/src/sqlparser/Cargo.toml b/src/sqlparser/Cargo.toml index 6489639a..7d6b40a6 100644 --- a/src/sqlparser/Cargo.toml +++ b/src/sqlparser/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "function-stream-sqlparser" -version = "0.6.0" -edition = "2024" +version.workspace = true +edition.workspace = true [lib] name = "function_stream_sqlparser" diff --git a/src/streaming_planner/Cargo.toml b/src/streaming_planner/Cargo.toml index d1039f97..766aba1c 100644 --- a/src/streaming_planner/Cargo.toml +++ b/src/streaming_planner/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "function-stream-streaming-planner" -version = "0.6.0" -edition = "2024" +version.workspace = true +edition.workspace = true [lib] name = "function_stream_streaming_planner" diff --git a/src/streaming_runtime/Cargo.toml b/src/streaming_runtime/Cargo.toml index 732fd51a..026b50e1 100644 --- a/src/streaming_runtime/Cargo.toml +++ b/src/streaming_runtime/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "function-stream-streaming-runtime" -version = "0.6.0" -edition = "2024" +version.workspace = true +edition.workspace = true [lib] name = "function_stream_streaming_runtime" diff --git a/src/streaming_runtime/src/lib.rs b/src/streaming_runtime/src/lib.rs index c6c2c94d..aeb056d5 100644 --- a/src/streaming_runtime/src/lib.rs +++ b/src/streaming_runtime/src/lib.rs @@ -15,6 +15,6 @@ //! The streaming engine and shared runtime helpers (`streaming/`, `util/`) are //! implemented under [`src/streaming`] and [`src/util`] in this package. They are //! currently **compiled as part of the `function-stream` crate** via `#[path]` in -//! `src/runtime.rs`, sharing the root `crate::sql` name (re-exported streaming planner crate). +//! `src/lib.rs` / `src/main.rs`, sharing the root `crate::sql` name (re-exported streaming planner crate). pub const CRATE_NAME: &str = "function-stream-streaming-runtime"; diff --git a/src/streaming_runtime/src/streaming/api/context.rs b/src/streaming_runtime/src/streaming/api/context.rs index f2557e7a..a29d7ac2 100644 --- a/src/streaming_runtime/src/streaming/api/context.rs +++ b/src/streaming_runtime/src/streaming/api/context.rs @@ -19,11 +19,11 @@ use arrow_array::RecordBatch; use protocol::storage::SourceCheckpointInfo; use tokio::sync::mpsc; -use crate::runtime::memory::{MemoryBlock, MemoryPool, get_array_memory_size}; -use crate::runtime::streaming::network::endpoint::PhysicalSender; -use crate::runtime::streaming::protocol::control::JobMasterEvent; -use crate::runtime::streaming::protocol::event::{StreamEvent, TrackedEvent}; -use crate::runtime::streaming::state::IoManager; +use crate::memory::{MemoryBlock, MemoryPool, get_array_memory_size}; +use crate::streaming::network::endpoint::PhysicalSender; +use crate::streaming::protocol::control::JobMasterEvent; +use crate::streaming::protocol::event::{StreamEvent, TrackedEvent}; +use crate::streaming::state::IoManager; #[derive(Debug, Clone)] pub struct TaskContextConfig { diff --git a/src/streaming_runtime/src/streaming/api/operator.rs b/src/streaming_runtime/src/streaming/api/operator.rs index fc75e475..43535a8a 100644 --- a/src/streaming_runtime/src/streaming/api/operator.rs +++ b/src/streaming_runtime/src/streaming/api/operator.rs @@ -10,9 +10,9 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::runtime::streaming::api::context::TaskContext; -use crate::runtime::streaming::api::source::SourceOperator; -use crate::runtime::streaming::protocol::event::StreamOutput; +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::source::SourceOperator; +use crate::streaming::protocol::event::StreamOutput; use crate::sql::common::{CheckpointBarrier, Watermark}; use arrow_array::RecordBatch; use async_trait::async_trait; diff --git a/src/streaming_runtime/src/streaming/api/source.rs b/src/streaming_runtime/src/streaming/api/source.rs index 9c531f2c..e0ab6a92 100644 --- a/src/streaming_runtime/src/streaming/api/source.rs +++ b/src/streaming_runtime/src/streaming/api/source.rs @@ -10,7 +10,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::runtime::streaming::api::context::TaskContext; +use crate::streaming::api::context::TaskContext; use crate::sql::common::{CheckpointBarrier, Watermark}; use arrow_array::RecordBatch; use async_trait::async_trait; diff --git a/src/streaming_runtime/src/streaming/execution/operator_chain.rs b/src/streaming_runtime/src/streaming/execution/operator_chain.rs index 88e8f441..2cd392fc 100644 --- a/src/streaming_runtime/src/streaming/execution/operator_chain.rs +++ b/src/streaming_runtime/src/streaming/execution/operator_chain.rs @@ -13,10 +13,10 @@ use anyhow::anyhow; use async_trait::async_trait; -use crate::runtime::streaming::api::context::TaskContext; -use crate::runtime::streaming::api::operator::{Collector, Operator}; -use crate::runtime::streaming::error::RunError; -use crate::runtime::streaming::protocol::{ +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::operator::{Collector, Operator}; +use crate::streaming::error::RunError; +use crate::streaming::protocol::{ control::{ControlCommand, StopMode}, event::{StreamEvent, StreamOutput, TrackedEvent}, }; diff --git a/src/streaming_runtime/src/streaming/execution/pipeline.rs b/src/streaming_runtime/src/streaming/execution/pipeline.rs index 7c2ca17a..badc3601 100644 --- a/src/streaming_runtime/src/streaming/execution/pipeline.rs +++ b/src/streaming_runtime/src/streaming/execution/pipeline.rs @@ -14,16 +14,16 @@ use tokio::sync::mpsc::UnboundedReceiver; use tokio_stream::{StreamExt, StreamMap}; use tracing::{Instrument, info, info_span}; -use crate::runtime::streaming::api::context::TaskContext; -use crate::runtime::streaming::api::operator::Operator; -use crate::runtime::streaming::error::RunError; -use crate::runtime::streaming::execution::operator_chain::{ChainBuilder, OperatorDrive}; -use crate::runtime::streaming::execution::tracker::{ +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::operator::Operator; +use crate::streaming::error::RunError; +use crate::streaming::execution::operator_chain::{ChainBuilder, OperatorDrive}; +use crate::streaming::execution::tracker::{ barrier_aligner::{AlignmentStatus, BarrierAligner}, watermark_tracker::WatermarkTracker, }; -use crate::runtime::streaming::network::endpoint::BoxedEventStream; -use crate::runtime::streaming::protocol::{ +use crate::streaming::network::endpoint::BoxedEventStream; +use crate::streaming::protocol::{ control::ControlCommand, event::{StreamEvent, TrackedEvent}, }; diff --git a/src/streaming_runtime/src/streaming/execution/source_driver.rs b/src/streaming_runtime/src/streaming/execution/source_driver.rs index 0118c4ee..91058ab2 100644 --- a/src/streaming_runtime/src/streaming/execution/source_driver.rs +++ b/src/streaming_runtime/src/streaming/execution/source_driver.rs @@ -14,11 +14,11 @@ use tokio::sync::mpsc::UnboundedReceiver; use tokio::time::{Instant, sleep}; use tracing::{Instrument, info, info_span, warn}; -use crate::runtime::streaming::api::context::TaskContext; -use crate::runtime::streaming::api::source::{SourceCheckpointReport, SourceEvent, SourceOperator}; -use crate::runtime::streaming::error::RunError; -use crate::runtime::streaming::execution::OperatorDrive; -use crate::runtime::streaming::protocol::{ +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::source::{SourceCheckpointReport, SourceEvent, SourceOperator}; +use crate::streaming::error::RunError; +use crate::streaming::execution::OperatorDrive; +use crate::streaming::protocol::{ control::ControlCommand, event::{StreamEvent, TrackedEvent}, }; diff --git a/src/streaming_runtime/src/streaming/execution/tracker/watermark_tracker.rs b/src/streaming_runtime/src/streaming/execution/tracker/watermark_tracker.rs index af6fd0bc..87d9d89c 100644 --- a/src/streaming_runtime/src/streaming/execution/tracker/watermark_tracker.rs +++ b/src/streaming_runtime/src/streaming/execution/tracker/watermark_tracker.rs @@ -10,7 +10,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::runtime::streaming::protocol::event::{merge_watermarks, watermark_strictly_advances}; +use crate::streaming::protocol::event::{merge_watermarks, watermark_strictly_advances}; use crate::sql::common::Watermark; #[derive(Debug)] diff --git a/src/streaming_runtime/src/streaming/factory/connector/delta.rs b/src/streaming_runtime/src/streaming/factory/connector/delta.rs index 726f87ef..324147bd 100644 --- a/src/streaming_runtime/src/streaming/factory/connector/delta.rs +++ b/src/streaming_runtime/src/streaming/factory/connector/delta.rs @@ -18,14 +18,14 @@ use prost::Message; use protocol::function_stream_graph::ConnectorOp; use protocol::function_stream_graph::connector_op::Config; -use crate::runtime::streaming::api::operator::ConstructedOperator; -use crate::runtime::streaming::factory::connector::sink_props_codec::{ +use crate::streaming::api::operator::ConstructedOperator; +use crate::streaming::factory::connector::sink_props_codec::{ apply_common_sink_fields, normalized_props, parse_sink_memory_bytes, }; -use crate::runtime::streaming::factory::global::Registry; -use crate::runtime::streaming::factory::operator_constructor::OperatorConstructor; -use crate::runtime::streaming::operators::sink::delta::{DeltaFormat, DeltaSinkOperator}; -use crate::runtime::streaming::operators::sink::filesystem::compression_from_str; +use crate::streaming::factory::global::Registry; +use crate::streaming::factory::operator_constructor::OperatorConstructor; +use crate::streaming::operators::sink::delta::{DeltaFormat, DeltaSinkOperator}; +use crate::streaming::operators::sink::filesystem::compression_from_str; use crate::sql::common::constants::connection_format_value; use crate::sql::common::with_option_keys as opt; diff --git a/src/streaming_runtime/src/streaming/factory/connector/dispatchers.rs b/src/streaming_runtime/src/streaming/factory/connector/dispatchers.rs index 7d626600..e5997c38 100644 --- a/src/streaming_runtime/src/streaming/factory/connector/dispatchers.rs +++ b/src/streaming_runtime/src/streaming/factory/connector/dispatchers.rs @@ -16,9 +16,9 @@ use anyhow::{Context, Result, bail}; use prost::Message; use protocol::function_stream_graph::ConnectorOp; -use crate::runtime::streaming::api::operator::ConstructedOperator; -use crate::runtime::streaming::factory::global::Registry; -use crate::runtime::streaming::factory::operator_constructor::OperatorConstructor; +use crate::streaming::api::operator::ConstructedOperator; +use crate::streaming::factory::global::Registry; +use crate::streaming::factory::operator_constructor::OperatorConstructor; use crate::sql::common::constants::connector_type; use super::{ diff --git a/src/streaming_runtime/src/streaming/factory/connector/filesystem.rs b/src/streaming_runtime/src/streaming/factory/connector/filesystem.rs index 94101407..27ff9ec1 100644 --- a/src/streaming_runtime/src/streaming/factory/connector/filesystem.rs +++ b/src/streaming_runtime/src/streaming/factory/connector/filesystem.rs @@ -18,13 +18,13 @@ use prost::Message; use protocol::function_stream_graph::ConnectorOp; use protocol::function_stream_graph::connector_op::Config; -use crate::runtime::streaming::api::operator::ConstructedOperator; -use crate::runtime::streaming::factory::connector::sink_props_codec::{ +use crate::streaming::api::operator::ConstructedOperator; +use crate::streaming::factory::connector::sink_props_codec::{ apply_common_sink_fields, normalized_props, parse_sink_memory_bytes, }; -use crate::runtime::streaming::factory::global::Registry; -use crate::runtime::streaming::factory::operator_constructor::OperatorConstructor; -use crate::runtime::streaming::operators::sink::filesystem::{ +use crate::streaming::factory::global::Registry; +use crate::streaming::factory::operator_constructor::OperatorConstructor; +use crate::streaming::operators::sink::filesystem::{ FilesystemFormat, FilesystemSinkOperator, compression_from_str, }; use crate::sql::common::constants::connection_format_value; diff --git a/src/streaming_runtime/src/streaming/factory/connector/iceberg.rs b/src/streaming_runtime/src/streaming/factory/connector/iceberg.rs index 58e0809f..8c4ca56a 100644 --- a/src/streaming_runtime/src/streaming/factory/connector/iceberg.rs +++ b/src/streaming_runtime/src/streaming/factory/connector/iceberg.rs @@ -18,14 +18,14 @@ use prost::Message; use protocol::function_stream_graph::ConnectorOp; use protocol::function_stream_graph::connector_op::Config; -use crate::runtime::streaming::api::operator::ConstructedOperator; -use crate::runtime::streaming::factory::connector::sink_props_codec::{ +use crate::streaming::api::operator::ConstructedOperator; +use crate::streaming::factory::connector::sink_props_codec::{ apply_common_sink_fields, normalized_props, parse_sink_memory_bytes, }; -use crate::runtime::streaming::factory::global::Registry; -use crate::runtime::streaming::factory::operator_constructor::OperatorConstructor; -use crate::runtime::streaming::operators::sink::filesystem::compression_from_str; -use crate::runtime::streaming::operators::sink::iceberg::{IcebergFormat, IcebergSinkOperator}; +use crate::streaming::factory::global::Registry; +use crate::streaming::factory::operator_constructor::OperatorConstructor; +use crate::streaming::operators::sink::filesystem::compression_from_str; +use crate::streaming::operators::sink::iceberg::{IcebergFormat, IcebergSinkOperator}; use crate::sql::common::constants::connection_format_value; use crate::sql::common::with_option_keys as opt; diff --git a/src/streaming_runtime/src/streaming/factory/connector/kafka.rs b/src/streaming_runtime/src/streaming/factory/connector/kafka.rs index 17838e3e..7dcdafb9 100644 --- a/src/streaming_runtime/src/streaming/factory/connector/kafka.rs +++ b/src/streaming_runtime/src/streaming/factory/connector/kafka.rs @@ -24,16 +24,16 @@ use protocol::function_stream_graph::{ }; use tracing::info; -use crate::runtime::streaming::api::operator::ConstructedOperator; -use crate::runtime::streaming::api::source::SourceOffset; -use crate::runtime::streaming::factory::global::Registry; -use crate::runtime::streaming::factory::operator_constructor::OperatorConstructor; -use crate::runtime::streaming::format::{ +use crate::streaming::api::operator::ConstructedOperator; +use crate::streaming::api::source::SourceOffset; +use crate::streaming::factory::global::Registry; +use crate::streaming::factory::operator_constructor::OperatorConstructor; +use crate::streaming::format::{ BadDataPolicy as RtBadDataPolicy, DataSerializer, DecimalEncoding as RtDecimalEncoding, Format as RuntimeFormat, JsonFormat as RuntimeJsonFormat, TimestampFormat as RtTimestampFormat, }; -use crate::runtime::streaming::operators::sink::kafka::{ConsistencyMode, KafkaSinkOperator}; -use crate::runtime::streaming::operators::source::kafka::{ +use crate::streaming::operators::sink::kafka::{ConsistencyMode, KafkaSinkOperator}; +use crate::streaming::operators::source::kafka::{ BufferedDeserializer, KafkaSourceOperator, }; use crate::sql::common::FsSchema; diff --git a/src/streaming_runtime/src/streaming/factory/connector/lancedb.rs b/src/streaming_runtime/src/streaming/factory/connector/lancedb.rs index ad8bc246..61bfd927 100644 --- a/src/streaming_runtime/src/streaming/factory/connector/lancedb.rs +++ b/src/streaming_runtime/src/streaming/factory/connector/lancedb.rs @@ -18,13 +18,13 @@ use prost::Message; use protocol::function_stream_graph::ConnectorOp; use protocol::function_stream_graph::connector_op::Config; -use crate::runtime::streaming::api::operator::ConstructedOperator; -use crate::runtime::streaming::factory::connector::sink_props_codec::{ +use crate::streaming::api::operator::ConstructedOperator; +use crate::streaming::factory::connector::sink_props_codec::{ apply_common_sink_fields, normalized_props, }; -use crate::runtime::streaming::factory::global::Registry; -use crate::runtime::streaming::factory::operator_constructor::OperatorConstructor; -use crate::runtime::streaming::operators::sink::lancedb::LanceDbSinkOperator; +use crate::streaming::factory::global::Registry; +use crate::streaming::factory::operator_constructor::OperatorConstructor; +use crate::streaming::operators::sink::lancedb::LanceDbSinkOperator; use crate::sql::common::constants::connection_format_value; use crate::sql::common::with_option_keys as opt; diff --git a/src/streaming_runtime/src/streaming/factory/connector/s3.rs b/src/streaming_runtime/src/streaming/factory/connector/s3.rs index 4b67fb9e..729d1396 100644 --- a/src/streaming_runtime/src/streaming/factory/connector/s3.rs +++ b/src/streaming_runtime/src/streaming/factory/connector/s3.rs @@ -18,13 +18,13 @@ use prost::Message; use protocol::function_stream_graph::ConnectorOp; use protocol::function_stream_graph::connector_op::Config; -use crate::runtime::streaming::api::operator::ConstructedOperator; -use crate::runtime::streaming::factory::connector::sink_props_codec::{ +use crate::streaming::api::operator::ConstructedOperator; +use crate::streaming::factory::connector::sink_props_codec::{ apply_common_sink_fields, normalized_props, }; -use crate::runtime::streaming::factory::global::Registry; -use crate::runtime::streaming::factory::operator_constructor::OperatorConstructor; -use crate::runtime::streaming::operators::sink::s3::{ +use crate::streaming::factory::global::Registry; +use crate::streaming::factory::operator_constructor::OperatorConstructor; +use crate::streaming::operators::sink::s3::{ S3Format, S3SinkOperator, compression_from_str, }; use crate::sql::common::constants::connection_format_value; diff --git a/src/streaming_runtime/src/streaming/factory/operator_constructor.rs b/src/streaming_runtime/src/streaming/factory/operator_constructor.rs index 5d0ff7d7..2cd96e6d 100644 --- a/src/streaming_runtime/src/streaming/factory/operator_constructor.rs +++ b/src/streaming_runtime/src/streaming/factory/operator_constructor.rs @@ -13,8 +13,8 @@ use anyhow::Result; use std::sync::Arc; -use crate::runtime::streaming::api::operator::ConstructedOperator; -use crate::runtime::streaming::factory::global::Registry; +use crate::streaming::api::operator::ConstructedOperator; +use crate::streaming::factory::global::Registry; /// Builds a [`ConstructedOperator`] from serialized configuration and a [`Registry`]. pub trait OperatorConstructor: Send + Sync { diff --git a/src/streaming_runtime/src/streaming/factory/operator_factory.rs b/src/streaming_runtime/src/streaming/factory/operator_factory.rs index 1ce04eeb..17daa0f2 100644 --- a/src/streaming_runtime/src/streaming/factory/operator_factory.rs +++ b/src/streaming_runtime/src/streaming/factory/operator_factory.rs @@ -11,13 +11,13 @@ // limitations under the License. use super::operator_constructor::OperatorConstructor; -use crate::runtime::streaming::api::operator::ConstructedOperator; -use crate::runtime::streaming::factory::connector::{ +use crate::streaming::api::operator::ConstructedOperator; +use crate::streaming::factory::connector::{ ConnectorSinkDispatcher, ConnectorSourceDispatcher, }; -use crate::runtime::streaming::factory::global::Registry; -use crate::runtime::streaming::operators::grouping::IncrementalAggregatingConstructor; -use crate::runtime::streaming::operators::joins::{ +use crate::streaming::factory::global::Registry; +use crate::streaming::operators::grouping::IncrementalAggregatingConstructor; +use crate::streaming::operators::joins::{ InstantJoinConstructor, JoinWithExpirationConstructor, }; use anyhow::{Result, anyhow}; @@ -26,12 +26,12 @@ use protocol::function_stream_graph::ProjectionOperator as ProjectionOperatorPro use std::collections::HashMap; use std::sync::Arc; -use crate::runtime::streaming::operators::watermark::WatermarkGeneratorConstructor; -use crate::runtime::streaming::operators::windows::{ +use crate::streaming::operators::watermark::WatermarkGeneratorConstructor; +use crate::streaming::operators::windows::{ SessionAggregatingWindowConstructor, SlidingAggregatingWindowConstructor, TumblingAggregateWindowConstructor, WindowFunctionConstructor, }; -use crate::runtime::streaming::operators::{ +use crate::streaming::operators::{ KeyExecutionOperator, ProjectionOperator, StatelessPhysicalExecutor, ValueExecutionOperator, }; use protocol::function_stream_graph::{ @@ -125,7 +125,7 @@ impl OperatorFactory { ); self.register_named(OperatorName::ConnectorSink, Box::new(ConnectorSinkBridge)); - crate::runtime::streaming::factory::register_kafka_connector_plugins(self); + crate::streaming::factory::register_kafka_connector_plugins(self); } } diff --git a/src/streaming_runtime/src/streaming/job/edge_manager.rs b/src/streaming_runtime/src/streaming/job/edge_manager.rs index 00c94485..32bb3722 100644 --- a/src/streaming_runtime/src/streaming/job/edge_manager.rs +++ b/src/streaming_runtime/src/streaming/job/edge_manager.rs @@ -16,7 +16,7 @@ use anyhow::{Result, anyhow}; use tokio::sync::mpsc; use tracing::{debug, info, warn}; -use crate::runtime::streaming::protocol::event::TrackedEvent; +use crate::streaming::protocol::event::TrackedEvent; use protocol::function_stream_graph::{FsEdge, FsNode}; const DEFAULT_CHANNEL_CAPACITY: usize = 2048; diff --git a/src/streaming_runtime/src/streaming/job/job_manager.rs b/src/streaming_runtime/src/streaming/job/job_manager.rs index e4b9916b..d6b268c4 100644 --- a/src/streaming_runtime/src/streaming/job/job_manager.rs +++ b/src/streaming_runtime/src/streaming/job/job_manager.rs @@ -30,22 +30,22 @@ use crate::config::{ DEFAULT_CHECKPOINT_INTERVAL_MS, DEFAULT_OPERATOR_STATE_STORE_MEMORY_BYTES, DEFAULT_PIPELINE_PARALLELISM, }; -use crate::runtime::memory::global_memory_pool; -use crate::runtime::streaming::api::context::TaskContext; -use crate::runtime::streaming::api::operator::{ConstructedOperator, Operator}; -use crate::runtime::streaming::api::source::SourceOperator; -use crate::runtime::streaming::execution::{ChainBuilder, Pipeline, SourceDriver}; -use crate::runtime::streaming::factory::OperatorFactory; -use crate::runtime::streaming::job::edge_manager::EdgeManager; -use crate::runtime::streaming::job::models::{ +use crate::memory::global_memory_pool; +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::operator::{ConstructedOperator, Operator}; +use crate::streaming::api::source::SourceOperator; +use crate::streaming::execution::{ChainBuilder, Pipeline, SourceDriver}; +use crate::streaming::factory::OperatorFactory; +use crate::streaming::job::edge_manager::EdgeManager; +use crate::streaming::job::models::{ PhysicalExecutionGraph, PhysicalPipeline, PipelineStatus, StreamingJobRollupStatus, }; -use crate::runtime::streaming::network::endpoint::{BoxedEventStream, PhysicalSender}; -use crate::runtime::streaming::protocol::control::{ControlCommand, JobMasterEvent, StopMode}; -use crate::runtime::streaming::protocol::event::CheckpointBarrier; -use crate::runtime::streaming::state::{IoManager, IoPool, NoopMetricsCollector}; +use crate::streaming::network::endpoint::{BoxedEventStream, PhysicalSender}; +use crate::streaming::protocol::control::{ControlCommand, JobMasterEvent, StopMode}; +use crate::streaming::protocol::event::CheckpointBarrier; +use crate::streaming::state::{IoManager, IoPool, NoopMetricsCollector}; use crate::sql::logical_node::logical::OperatorName; -use crate::storage::stream_catalog::CatalogManager; +use crate::stream_catalog::CatalogManager; #[derive(Debug, Clone)] pub struct StreamingJobSummary { @@ -80,7 +80,7 @@ pub struct StateConfig { pub pipeline_parallelism: u32, pub job_manager_control_plane_threads: u32, pub job_manager_data_plane_threads: u32, - /// Total bytes shared by all [`crate::runtime::streaming::state::OperatorStateStore`] (global pool). + /// Total bytes shared by all [`crate::streaming::state::OperatorStateStore`] (global pool). pub per_operator_memory_bytes: u64, } @@ -103,7 +103,7 @@ impl Default for StateConfig { static GLOBAL_JOB_MANAGER: OnceLock> = OnceLock::new(); -/// Operators that create an [`crate::runtime::streaming::state::OperatorStateStore`] at runtime. +/// Operators that create an [`crate::streaming::state::OperatorStateStore`] at runtime. fn pipeline_state_store_operator_count(operators: &[ChainedOperator]) -> usize { operators .iter() @@ -159,7 +159,7 @@ struct CheckpointCoordinatorConfig { } impl PipelineRunner { - async fn run(self) -> Result<(), crate::runtime::streaming::error::RunError> { + async fn run(self) -> Result<(), crate::streaming::error::RunError> { match self { PipelineRunner::Source(driver) => driver.run().await, PipelineRunner::Standard(pipeline) => pipeline.run().await, diff --git a/src/streaming_runtime/src/streaming/job/models.rs b/src/streaming_runtime/src/streaming/job/models.rs index e81649f2..a6d226bd 100644 --- a/src/streaming_runtime/src/streaming/job/models.rs +++ b/src/streaming_runtime/src/streaming/job/models.rs @@ -19,7 +19,7 @@ use protocol::function_stream_graph::FsProgram; use tokio::sync::mpsc; use tokio::task::JoinHandle; -use crate::runtime::streaming::protocol::control::ControlCommand; +use crate::streaming::protocol::control::ControlCommand; #[derive(Debug, Clone, PartialEq)] pub enum PipelineStatus { diff --git a/src/streaming_runtime/src/streaming/network/endpoint.rs b/src/streaming_runtime/src/streaming/network/endpoint.rs index ae75e6fc..12550977 100644 --- a/src/streaming_runtime/src/streaming/network/endpoint.rs +++ b/src/streaming_runtime/src/streaming/network/endpoint.rs @@ -10,8 +10,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::runtime::streaming::protocol::event::StreamEvent; -use crate::runtime::streaming::protocol::event::TrackedEvent; +use crate::streaming::protocol::event::StreamEvent; +use crate::streaming::protocol::event::TrackedEvent; use anyhow::{Result, anyhow}; use std::pin::Pin; use tokio::sync::mpsc; diff --git a/src/streaming_runtime/src/streaming/operators/grouping/incremental_aggregate.rs b/src/streaming_runtime/src/streaming/operators/grouping/incremental_aggregate.rs index ff997c11..eaaedddf 100644 --- a/src/streaming_runtime/src/streaming/operators/grouping/incremental_aggregate.rs +++ b/src/streaming_runtime/src/streaming/operators/grouping/incremental_aggregate.rs @@ -41,13 +41,13 @@ use std::{collections::HashMap, mem, sync::Arc}; use tracing::{debug, info, warn}; // ========================================================================= // ========================================================================= -use crate::runtime::streaming::StreamOutput; -use crate::runtime::streaming::api::context::TaskContext; -use crate::runtime::streaming::api::operator::{Collector, Operator}; -use crate::runtime::streaming::factory::Registry; -use crate::runtime::streaming::operators::{Key, UpdatingCache}; -use crate::runtime::streaming::state::OperatorStateStore; -use crate::runtime::util::decode_aggregate; +use crate::streaming::StreamOutput; +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::operator::{Collector, Operator}; +use crate::streaming::factory::Registry; +use crate::streaming::operators::{Key, UpdatingCache}; +use crate::streaming::state::OperatorStateStore; +use crate::util::decode_aggregate; use crate::sql::common::{ CheckpointBarrier, FsSchema, TIMESTAMP_FIELD, UPDATING_META_FIELD, Watermark, to_nanos, }; diff --git a/src/streaming_runtime/src/streaming/operators/joins/join_instance.rs b/src/streaming_runtime/src/streaming/operators/joins/join_instance.rs index 5c04c1fb..e1a0e71b 100644 --- a/src/streaming_runtime/src/streaming/operators/joins/join_instance.rs +++ b/src/streaming_runtime/src/streaming/operators/joins/join_instance.rs @@ -25,11 +25,11 @@ use std::sync::{Arc, RwLock}; use std::time::UNIX_EPOCH; use tracing::{info, warn}; -use crate::runtime::streaming::StreamOutput; -use crate::runtime::streaming::api::context::TaskContext; -use crate::runtime::streaming::api::operator::{Collector, Operator}; -use crate::runtime::streaming::factory::Registry; -use crate::runtime::streaming::state::OperatorStateStore; +use crate::streaming::StreamOutput; +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::operator::{Collector, Operator}; +use crate::streaming::factory::Registry; +use crate::streaming::state::OperatorStateStore; use crate::sql::common::{CheckpointBarrier, FsSchema, FsSchemaRef, Watermark}; use crate::sql::physical::{StreamingDecodingContext, StreamingExtensionCodec}; use async_trait::async_trait; diff --git a/src/streaming_runtime/src/streaming/operators/joins/join_with_expiration.rs b/src/streaming_runtime/src/streaming/operators/joins/join_with_expiration.rs index 5ed8dfa3..ea7e8477 100644 --- a/src/streaming_runtime/src/streaming/operators/joins/join_with_expiration.rs +++ b/src/streaming_runtime/src/streaming/operators/joins/join_with_expiration.rs @@ -24,11 +24,11 @@ use std::sync::{Arc, RwLock}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use tracing::{info, warn}; -use crate::runtime::streaming::StreamOutput; -use crate::runtime::streaming::api::context::TaskContext; -use crate::runtime::streaming::api::operator::{Collector, Operator}; -use crate::runtime::streaming::factory::Registry; -use crate::runtime::streaming::state::OperatorStateStore; +use crate::streaming::StreamOutput; +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::operator::{Collector, Operator}; +use crate::streaming::factory::Registry; +use crate::streaming::state::OperatorStateStore; use crate::sql::common::{CheckpointBarrier, FsSchema, Watermark}; use crate::sql::physical::{StreamingDecodingContext, StreamingExtensionCodec}; use async_trait::async_trait; diff --git a/src/streaming_runtime/src/streaming/operators/key_by.rs b/src/streaming_runtime/src/streaming/operators/key_by.rs index 90c55d08..7e17a768 100644 --- a/src/streaming_runtime/src/streaming/operators/key_by.rs +++ b/src/streaming_runtime/src/streaming/operators/key_by.rs @@ -19,9 +19,9 @@ use datafusion_common::hash_utils::create_hashes; use datafusion_physical_expr::expressions::Column; use std::sync::Arc; -use crate::runtime::streaming::StreamOutput; -use crate::runtime::streaming::api::context::TaskContext; -use crate::runtime::streaming::api::operator::{Collector, Operator}; +use crate::streaming::StreamOutput; +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::operator::{Collector, Operator}; use crate::sql::common::{CheckpointBarrier, Watermark}; use protocol::function_stream_graph::KeyPlanOperator; diff --git a/src/streaming_runtime/src/streaming/operators/key_operator.rs b/src/streaming_runtime/src/streaming/operators/key_operator.rs index 7a89d2f2..7ac364bc 100644 --- a/src/streaming_runtime/src/streaming/operators/key_operator.rs +++ b/src/streaming_runtime/src/streaming/operators/key_operator.rs @@ -13,12 +13,12 @@ //! Key-by over the physical plan output: key column(s) are **values** projected by the plan //! (e.g. `_key_user_id`); **shuffle / `StreamOutput::Keyed` uses `u64` hashes** computed by //! [`datafusion_common::hash_utils::create_hashes`] on those columns — same mechanism as -//! [`crate::runtime::streaming::operators::key_by::KeyByOperator`]. +//! [`crate::streaming::operators::key_by::KeyByOperator`]. -use crate::runtime::streaming::StreamOutput; -use crate::runtime::streaming::api::context::TaskContext; -use crate::runtime::streaming::api::operator::{Collector, Operator}; -use crate::runtime::streaming::operators::StatelessPhysicalExecutor; +use crate::streaming::StreamOutput; +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::operator::{Collector, Operator}; +use crate::streaming::operators::StatelessPhysicalExecutor; use crate::sql::common::{CheckpointBarrier, Watermark}; use ahash::RandomState; use anyhow::{Result, anyhow}; diff --git a/src/streaming_runtime/src/streaming/operators/projection.rs b/src/streaming_runtime/src/streaming/operators/projection.rs index b84d74aa..cff269f1 100644 --- a/src/streaming_runtime/src/streaming/operators/projection.rs +++ b/src/streaming_runtime/src/streaming/operators/projection.rs @@ -22,10 +22,10 @@ use std::sync::Arc; use protocol::function_stream_graph::ProjectionOperator as ProjectionOperatorProto; -use crate::runtime::streaming::StreamOutput; -use crate::runtime::streaming::api::context::TaskContext; -use crate::runtime::streaming::api::operator::{Collector, Operator}; -use crate::runtime::streaming::factory::global::Registry; +use crate::streaming::StreamOutput; +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::operator::{Collector, Operator}; +use crate::streaming::factory::global::Registry; use crate::sql::common::{CheckpointBarrier, FsSchema, FsSchemaRef, Watermark}; use crate::sql::logical_node::logical::OperatorName; diff --git a/src/streaming_runtime/src/streaming/operators/sink/delta/mod.rs b/src/streaming_runtime/src/streaming/operators/sink/delta/mod.rs index 4df6b3b5..1bce0cd5 100644 --- a/src/streaming_runtime/src/streaming/operators/sink/delta/mod.rs +++ b/src/streaming_runtime/src/streaming/operators/sink/delta/mod.rs @@ -26,11 +26,11 @@ use parquet::basic::Compression; use tokio::io::AsyncWriteExt; use tracing::{debug, info, warn}; -use crate::runtime::memory::{MemoryBlock, try_global_memory_pool}; -use crate::runtime::streaming::StreamOutput; -use crate::runtime::streaming::api::context::TaskContext; -use crate::runtime::streaming::api::operator::{Collector, Operator}; -use crate::runtime::streaming::format::encoder::FormatEncoder; +use crate::memory::{MemoryBlock, try_global_memory_pool}; +use crate::streaming::StreamOutput; +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::operator::{Collector, Operator}; +use crate::streaming::format::encoder::FormatEncoder; use crate::sql::common::constants::factory_operator_name; use crate::sql::common::with_option_keys as opt; use crate::sql::common::{CheckpointBarrier, Watermark}; diff --git a/src/streaming_runtime/src/streaming/operators/sink/filesystem/mod.rs b/src/streaming_runtime/src/streaming/operators/sink/filesystem/mod.rs index a865a752..250ff0c3 100644 --- a/src/streaming_runtime/src/streaming/operators/sink/filesystem/mod.rs +++ b/src/streaming_runtime/src/streaming/operators/sink/filesystem/mod.rs @@ -20,11 +20,11 @@ use parquet::basic::Compression; use tokio::io::AsyncWriteExt; use tracing::{debug, info, warn}; -use crate::runtime::memory::{MemoryBlock, try_global_memory_pool}; -use crate::runtime::streaming::StreamOutput; -use crate::runtime::streaming::api::context::TaskContext; -use crate::runtime::streaming::api::operator::{Collector, Operator}; -use crate::runtime::streaming::format::encoder::FormatEncoder; +use crate::memory::{MemoryBlock, try_global_memory_pool}; +use crate::streaming::StreamOutput; +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::operator::{Collector, Operator}; +use crate::streaming::format::encoder::FormatEncoder; use crate::sql::common::constants::factory_operator_name; use crate::sql::common::{CheckpointBarrier, Watermark}; diff --git a/src/streaming_runtime/src/streaming/operators/sink/iceberg/mod.rs b/src/streaming_runtime/src/streaming/operators/sink/iceberg/mod.rs index b6c17414..801e34b0 100644 --- a/src/streaming_runtime/src/streaming/operators/sink/iceberg/mod.rs +++ b/src/streaming_runtime/src/streaming/operators/sink/iceberg/mod.rs @@ -26,11 +26,11 @@ use parquet::basic::Compression; use tokio::io::AsyncWriteExt; use tracing::{debug, info, warn}; -use crate::runtime::memory::{MemoryBlock, try_global_memory_pool}; -use crate::runtime::streaming::StreamOutput; -use crate::runtime::streaming::api::context::TaskContext; -use crate::runtime::streaming::api::operator::{Collector, Operator}; -use crate::runtime::streaming::format::encoder::FormatEncoder; +use crate::memory::{MemoryBlock, try_global_memory_pool}; +use crate::streaming::StreamOutput; +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::operator::{Collector, Operator}; +use crate::streaming::format::encoder::FormatEncoder; use crate::sql::common::constants::factory_operator_name; use crate::sql::common::with_option_keys as opt; use crate::sql::common::{CheckpointBarrier, Watermark}; diff --git a/src/streaming_runtime/src/streaming/operators/sink/kafka/mod.rs b/src/streaming_runtime/src/streaming/operators/sink/kafka/mod.rs index b30bc572..e87326e8 100644 --- a/src/streaming_runtime/src/streaming/operators/sink/kafka/mod.rs +++ b/src/streaming_runtime/src/streaming/operators/sink/kafka/mod.rs @@ -36,10 +36,10 @@ use std::time::Duration; use tokio::time::sleep; use tracing::{info, warn}; -use crate::runtime::streaming::StreamOutput; -use crate::runtime::streaming::api::context::TaskContext; -use crate::runtime::streaming::api::operator::{Collector, Operator}; -use crate::runtime::streaming::format::DataSerializer; +use crate::streaming::StreamOutput; +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::operator::{Collector, Operator}; +use crate::streaming::format::DataSerializer; use crate::sql::common::constants::factory_operator_name; use crate::sql::common::{CheckpointBarrier, FsSchema, Watermark}; // ============================================================================ diff --git a/src/streaming_runtime/src/streaming/operators/sink/lancedb/mod.rs b/src/streaming_runtime/src/streaming/operators/sink/lancedb/mod.rs index 4e8d8309..fa5b6760 100644 --- a/src/streaming_runtime/src/streaming/operators/sink/lancedb/mod.rs +++ b/src/streaming_runtime/src/streaming/operators/sink/lancedb/mod.rs @@ -23,9 +23,9 @@ use lance::dataset::{WriteMode, WriteParams}; use lance::io::{ObjectStoreParams, StorageOptionsAccessor}; use tracing::{info, warn}; -use crate::runtime::streaming::StreamOutput; -use crate::runtime::streaming::api::context::TaskContext; -use crate::runtime::streaming::api::operator::{Collector, Operator}; +use crate::streaming::StreamOutput; +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::operator::{Collector, Operator}; use crate::sql::common::constants::factory_operator_name; use crate::sql::common::{CheckpointBarrier, Watermark}; diff --git a/src/streaming_runtime/src/streaming/operators/sink/s3/mod.rs b/src/streaming_runtime/src/streaming/operators/sink/s3/mod.rs index 715b5b86..77dcca97 100644 --- a/src/streaming_runtime/src/streaming/operators/sink/s3/mod.rs +++ b/src/streaming_runtime/src/streaming/operators/sink/s3/mod.rs @@ -26,9 +26,9 @@ use parquet::basic::Compression; use parquet::file::properties::WriterProperties; use tracing::{info, warn}; -use crate::runtime::streaming::StreamOutput; -use crate::runtime::streaming::api::context::TaskContext; -use crate::runtime::streaming::api::operator::{Collector, Operator}; +use crate::streaming::StreamOutput; +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::operator::{Collector, Operator}; use crate::sql::common::constants::factory_operator_name; use crate::sql::common::with_option_keys as opt; use crate::sql::common::{CheckpointBarrier, Watermark}; diff --git a/src/streaming_runtime/src/streaming/operators/source/kafka/mod.rs b/src/streaming_runtime/src/streaming/operators/source/kafka/mod.rs index 75edb968..3e7bffc9 100644 --- a/src/streaming_runtime/src/streaming/operators/source/kafka/mod.rs +++ b/src/streaming_runtime/src/streaming/operators/source/kafka/mod.rs @@ -31,11 +31,11 @@ use std::num::NonZeroU32; use std::time::{Duration, Instant}; use tracing::{debug, error, info, warn}; -use crate::runtime::streaming::api::context::TaskContext; -use crate::runtime::streaming::api::source::{ +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::source::{ SourceCheckpointReport, SourceEvent, SourceOffset, SourceOperator, }; -use crate::runtime::streaming::format::{BadDataPolicy, DataDeserializer, Format}; +use crate::streaming::format::{BadDataPolicy, DataDeserializer, Format}; use crate::sql::common::fs_schema::FieldValueType; use crate::sql::common::{CheckpointBarrier, MetadataField}; diff --git a/src/streaming_runtime/src/streaming/operators/stateless_physical_executor.rs b/src/streaming_runtime/src/streaming/operators/stateless_physical_executor.rs index eb595d31..69f9dc45 100644 --- a/src/streaming_runtime/src/streaming/operators/stateless_physical_executor.rs +++ b/src/streaming_runtime/src/streaming/operators/stateless_physical_executor.rs @@ -24,7 +24,7 @@ use datafusion_proto::protobuf::PhysicalPlanNode; use futures::StreamExt; use prost::Message; -use crate::runtime::streaming::factory::Registry; +use crate::streaming::factory::Registry; use crate::sql::physical::{StreamingDecodingContext, StreamingExtensionCodec}; pub struct StatelessPhysicalExecutor { diff --git a/src/streaming_runtime/src/streaming/operators/value_execution.rs b/src/streaming_runtime/src/streaming/operators/value_execution.rs index b93cd78b..69059317 100644 --- a/src/streaming_runtime/src/streaming/operators/value_execution.rs +++ b/src/streaming_runtime/src/streaming/operators/value_execution.rs @@ -15,10 +15,10 @@ use arrow_array::RecordBatch; use async_trait::async_trait; use futures::StreamExt; -use crate::runtime::streaming::StreamOutput; -use crate::runtime::streaming::api::context::TaskContext; -use crate::runtime::streaming::api::operator::{Collector, Operator}; -use crate::runtime::streaming::operators::StatelessPhysicalExecutor; +use crate::streaming::StreamOutput; +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::operator::{Collector, Operator}; +use crate::streaming::operators::StatelessPhysicalExecutor; use crate::sql::common::{CheckpointBarrier, Watermark}; pub struct ValueExecutionOperator { diff --git a/src/streaming_runtime/src/streaming/operators/watermark/watermark_generator.rs b/src/streaming_runtime/src/streaming/operators/watermark/watermark_generator.rs index 497553eb..2a38f8f0 100644 --- a/src/streaming_runtime/src/streaming/operators/watermark/watermark_generator.rs +++ b/src/streaming_runtime/src/streaming/operators/watermark/watermark_generator.rs @@ -25,10 +25,10 @@ use std::sync::Arc; use std::time::{Duration, SystemTime}; use tracing::debug; -use crate::runtime::streaming::StreamOutput; -use crate::runtime::streaming::api::context::TaskContext; -use crate::runtime::streaming::api::operator::{Collector, Operator}; -use crate::runtime::streaming::factory::Registry; +use crate::streaming::StreamOutput; +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::operator::{Collector, Operator}; +use crate::streaming::factory::Registry; use crate::sql::common::{CheckpointBarrier, FsSchema, Watermark, from_nanos, to_millis}; use async_trait::async_trait; use protocol::function_stream_graph::ExpressionWatermarkConfig; diff --git a/src/streaming_runtime/src/streaming/operators/windows/session_aggregating_window.rs b/src/streaming_runtime/src/streaming/operators/windows/session_aggregating_window.rs index 2056cdd9..c48bd7a9 100644 --- a/src/streaming_runtime/src/streaming/operators/windows/session_aggregating_window.rs +++ b/src/streaming_runtime/src/streaming/operators/windows/session_aggregating_window.rs @@ -36,11 +36,11 @@ use std::time::{Duration, SystemTime}; use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel}; use tracing::info; -use crate::runtime::streaming::StreamOutput; -use crate::runtime::streaming::api::context::TaskContext; -use crate::runtime::streaming::api::operator::{Collector, Operator}; -use crate::runtime::streaming::factory::Registry; -use crate::runtime::streaming::state::OperatorStateStore; +use crate::streaming::StreamOutput; +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::operator::{Collector, Operator}; +use crate::streaming::factory::Registry; +use crate::streaming::state::OperatorStateStore; use crate::sql::common::converter::Converter; use crate::sql::common::{ CheckpointBarrier, FsSchema, FsSchemaRef, Watermark, from_nanos, to_nanos, diff --git a/src/streaming_runtime/src/streaming/operators/windows/sliding_aggregating_window.rs b/src/streaming_runtime/src/streaming/operators/windows/sliding_aggregating_window.rs index f18b3b14..267f289f 100644 --- a/src/streaming_runtime/src/streaming/operators/windows/sliding_aggregating_window.rs +++ b/src/streaming_runtime/src/streaming/operators/windows/sliding_aggregating_window.rs @@ -33,11 +33,11 @@ use std::time::{Duration, SystemTime}; use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel}; use tracing::info; -use crate::runtime::streaming::StreamOutput; -use crate::runtime::streaming::api::context::TaskContext; -use crate::runtime::streaming::api::operator::{Collector, Operator}; -use crate::runtime::streaming::factory::Registry; -use crate::runtime::streaming::state::OperatorStateStore; +use crate::streaming::StreamOutput; +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::operator::{Collector, Operator}; +use crate::streaming::factory::Registry; +use crate::streaming::state::OperatorStateStore; use crate::sql::common::{CheckpointBarrier, FsSchema, Watermark, from_nanos, to_nanos}; use crate::sql::physical::{StreamingDecodingContext, StreamingExtensionCodec}; use async_trait::async_trait; diff --git a/src/streaming_runtime/src/streaming/operators/windows/tumbling_aggregating_window.rs b/src/streaming_runtime/src/streaming/operators/windows/tumbling_aggregating_window.rs index 5c805625..9588d800 100644 --- a/src/streaming_runtime/src/streaming/operators/windows/tumbling_aggregating_window.rs +++ b/src/streaming_runtime/src/streaming/operators/windows/tumbling_aggregating_window.rs @@ -34,11 +34,11 @@ use std::time::{Duration, SystemTime}; use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel}; use tracing::{info, warn}; -use crate::runtime::streaming::StreamOutput; -use crate::runtime::streaming::api::context::TaskContext; -use crate::runtime::streaming::api::operator::{Collector, Operator}; -use crate::runtime::streaming::factory::Registry; -use crate::runtime::streaming::state::OperatorStateStore; +use crate::streaming::StreamOutput; +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::operator::{Collector, Operator}; +use crate::streaming::factory::Registry; +use crate::streaming::state::OperatorStateStore; use crate::sql::common::time_utils::print_time; use crate::sql::common::{CheckpointBarrier, FsSchema, Watermark, from_nanos, to_nanos}; use crate::sql::physical::{StreamingDecodingContext, StreamingExtensionCodec}; diff --git a/src/streaming_runtime/src/streaming/operators/windows/window_function.rs b/src/streaming_runtime/src/streaming/operators/windows/window_function.rs index 37815b78..e9c83b97 100644 --- a/src/streaming_runtime/src/streaming/operators/windows/window_function.rs +++ b/src/streaming_runtime/src/streaming/operators/windows/window_function.rs @@ -26,11 +26,11 @@ use std::time::SystemTime; use tokio::sync::mpsc::{UnboundedReceiver, unbounded_channel}; use tracing::{info, warn}; -use crate::runtime::streaming::StreamOutput; -use crate::runtime::streaming::api::context::TaskContext; -use crate::runtime::streaming::api::operator::{Collector, Operator}; -use crate::runtime::streaming::factory::Registry; -use crate::runtime::streaming::state::OperatorStateStore; +use crate::streaming::StreamOutput; +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::operator::{Collector, Operator}; +use crate::streaming::factory::Registry; +use crate::streaming::state::OperatorStateStore; use crate::sql::common::time_utils::print_time; use crate::sql::common::{ CheckpointBarrier, FsSchema, FsSchemaRef, Watermark, from_nanos, to_nanos, diff --git a/src/streaming_runtime/src/streaming/protocol/event.rs b/src/streaming_runtime/src/streaming/protocol/event.rs index 8c0a4989..0db2135d 100644 --- a/src/streaming_runtime/src/streaming/protocol/event.rs +++ b/src/streaming_runtime/src/streaming/protocol/event.rs @@ -14,7 +14,7 @@ use std::sync::Arc; use arrow_array::RecordBatch; -use crate::runtime::memory::MemoryTicket; +use crate::memory::MemoryTicket; pub use function_stream_runtime_common::streaming_protocol::{ CheckpointBarrier, Watermark, merge_watermarks, watermark_strictly_advances, diff --git a/src/streaming_runtime/src/streaming/state/error.rs b/src/streaming_runtime/src/streaming/state/error.rs index 37bc6481..cab033db 100644 --- a/src/streaming_runtime/src/streaming/state/error.rs +++ b/src/streaming_runtime/src/streaming/state/error.rs @@ -13,7 +13,7 @@ use crossbeam_channel::TrySendError; use thiserror::Error; -use crate::runtime::memory::MemoryAllocationError; +use crate::memory::MemoryAllocationError; #[derive(Error, Debug)] pub enum StateEngineError { diff --git a/src/streaming_runtime/src/streaming/state/operator_state.rs b/src/streaming_runtime/src/streaming/state/operator_state.rs index a3514461..224671f2 100644 --- a/src/streaming_runtime/src/streaming/state/operator_state.rs +++ b/src/streaming_runtime/src/streaming/state/operator_state.rs @@ -13,7 +13,7 @@ use super::error::{Result, StateEngineError}; use super::io_manager::{CompactJob, IoManager, SpillJob}; use super::metrics::StateMetricsCollector; -use crate::runtime::memory::{MemoryBlock, MemoryTicket}; +use crate::memory::{MemoryBlock, MemoryTicket}; use arrow_array::builder::{BinaryBuilder, BooleanBuilder, UInt64Builder}; use arrow_array::{Array, BinaryArray, RecordBatch, UInt64Array}; use arrow_schema::{DataType, Field, Schema}; @@ -813,7 +813,7 @@ mod tests { use super::super::io_manager::IoPool; use super::super::metrics::NoopMetricsCollector; use super::*; - use crate::runtime::memory::{MemoryBlock, MemoryPool, global_memory_pool}; + use crate::memory::{MemoryBlock, MemoryPool, global_memory_pool}; use arrow_array::Int64Array; use tempfile::TempDir; @@ -836,7 +836,7 @@ mod tests { const TEST_OPERATOR_MEMORY: u64 = 2 * 1024 * 1024; fn ensure_global_memory_pool() { - use crate::runtime::memory::{init_global_memory_pool, try_global_memory_pool}; + use crate::memory::{init_global_memory_pool, try_global_memory_pool}; use std::sync::Once; static INIT: Once = Once::new(); INIT.call_once(|| { diff --git a/src/wasm_runtime/Cargo.toml b/src/wasm_runtime/Cargo.toml index 4cb1eae2..2f385240 100644 --- a/src/wasm_runtime/Cargo.toml +++ b/src/wasm_runtime/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "function-stream-wasm-runtime" -version = "0.6.0" -edition = "2024" +version.workspace = true +edition.workspace = true [lib] name = "function_stream_wasm_runtime" diff --git a/src/wasm_runtime/src/lib.rs b/src/wasm_runtime/src/lib.rs index dc036f7d..12b3ebee 100644 --- a/src/wasm_runtime/src/lib.rs +++ b/src/wasm_runtime/src/lib.rs @@ -13,10 +13,10 @@ //! WebAssembly execution runtime. //! //! Implementation lives under `src/wasm/` in this package. It is currently **compiled as -//! part of the `function-stream` crate** via `#[path]` in `src/runtime.rs`, so paths -//! like `crate::sql` (streaming planner dependency) and `crate::runtime::memory` keep resolving. +//! part of the `function-stream` crate** via `#[path]` in `src/lib.rs` / `src/main.rs`, so paths +//! like `crate::sql` (streaming planner dependency) and `crate::memory` keep resolving. //! //! Operator state storage (`state_backend/`) also lives in this package and is compiled via -//! `#[path]` from `src/storage.rs` as `crate::storage::state_backend`. +//! `#[path]` from the root crate as `crate::state_backend`. pub const CRATE_NAME: &str = "function-stream-wasm-runtime"; diff --git a/src/wasm_runtime/src/state_backend/factory.rs b/src/wasm_runtime/src/state_backend/factory.rs index 5dc38632..08cb06ca 100644 --- a/src/wasm_runtime/src/state_backend/factory.rs +++ b/src/wasm_runtime/src/state_backend/factory.rs @@ -10,8 +10,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::storage::state_backend::error::BackendError; -use crate::storage::state_backend::store::StateStore; +use crate::state_backend::error::BackendError; +use crate::state_backend::store::StateStore; use std::path::Path; use std::sync::Arc; @@ -47,11 +47,11 @@ pub fn get_factory_for_task>( task_name: String, created_at: u64, base_dir: Option

, - rocksdb_config: Option, + rocksdb_config: Option, ) -> Result, BackendError> { match factory_type { FactoryType::Memory => { - Ok(crate::storage::state_backend::memory::MemoryStateStoreFactory::default_factory()) + Ok(crate::state_backend::memory::MemoryStateStoreFactory::default_factory()) } FactoryType::RocksDB => { let base_dir = base_dir.ok_or_else(|| { @@ -63,7 +63,7 @@ pub fn get_factory_for_task>( .join(format!("{}-{}", task_name, created_at)); let config = rocksdb_config.unwrap_or_default(); - let factory = crate::storage::state_backend::rocksdb::RocksDBStateStoreFactory::new( + let factory = crate::state_backend::rocksdb::RocksDBStateStoreFactory::new( db_path, config, )?; diff --git a/src/wasm_runtime/src/state_backend/memory/factory.rs b/src/wasm_runtime/src/state_backend/memory/factory.rs index b62bd444..0526ecbd 100644 --- a/src/wasm_runtime/src/state_backend/memory/factory.rs +++ b/src/wasm_runtime/src/state_backend/memory/factory.rs @@ -11,8 +11,8 @@ // limitations under the License. use super::store::MemoryStateStore; -use crate::storage::state_backend::error::BackendError; -use crate::storage::state_backend::factory::StateStoreFactory; +use crate::state_backend::error::BackendError; +use crate::state_backend::factory::StateStoreFactory; use std::sync::{Arc, Mutex}; pub struct MemoryStateStoreFactory {} @@ -43,7 +43,7 @@ impl StateStoreFactory for MemoryStateStoreFactory { fn new_state_store( &self, _column_family: Option, - ) -> Result, BackendError> { + ) -> Result, BackendError> { Ok(Box::new(MemoryStateStore::new())) } } diff --git a/src/wasm_runtime/src/state_backend/memory/store.rs b/src/wasm_runtime/src/state_backend/memory/store.rs index e65e839d..e95bf51f 100644 --- a/src/wasm_runtime/src/state_backend/memory/store.rs +++ b/src/wasm_runtime/src/state_backend/memory/store.rs @@ -10,8 +10,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::storage::state_backend::error::BackendError; -use crate::storage::state_backend::store::{StateIterator, StateStore}; +use crate::state_backend::error::BackendError; +use crate::state_backend::store::{StateIterator, StateStore}; use std::collections::HashMap; use std::sync::{Arc, Mutex}; @@ -94,7 +94,7 @@ impl StateStore for MemoryStateStore { user_key: Vec, value: Vec, ) -> Result<(), BackendError> { - let key_bytes = crate::storage::state_backend::key_builder::build_key( + let key_bytes = crate::state_backend::key_builder::build_key( &key_group, &key, &namespace, &user_key, ); diff --git a/src/wasm_runtime/src/state_backend/rocksdb/factory.rs b/src/wasm_runtime/src/state_backend/rocksdb/factory.rs index 11554e13..74759b91 100644 --- a/src/wasm_runtime/src/state_backend/rocksdb/factory.rs +++ b/src/wasm_runtime/src/state_backend/rocksdb/factory.rs @@ -10,8 +10,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::storage::state_backend::error::BackendError; -use crate::storage::state_backend::factory::StateStoreFactory; +use crate::state_backend::error::BackendError; +use crate::state_backend::factory::StateStoreFactory; use rocksdb::{ColumnFamilyDescriptor, DB, Options}; use std::path::Path; use std::sync::{Arc, Mutex}; @@ -56,7 +56,7 @@ impl StateStoreFactory for RocksDBStateStoreFactory { fn new_state_store( &self, column_family: Option, - ) -> Result, BackendError> { + ) -> Result, BackendError> { self.new_state_store(column_family) } } @@ -137,7 +137,7 @@ impl RocksDBStateStoreFactory { pub fn new_state_store( &self, column_family: Option, - ) -> Result, BackendError> { + ) -> Result, BackendError> { if let Some(ref cf_name) = column_family && cf_name != "default" && self.db.cf_handle(cf_name).is_none() @@ -158,7 +158,7 @@ impl RocksDBStateStoreFactory { } } - crate::storage::state_backend::rocksdb::store::RocksDBStateStore::new_with_factory( + crate::state_backend::rocksdb::store::RocksDBStateStore::new_with_factory( self.db.clone(), column_family, ) diff --git a/src/wasm_runtime/src/state_backend/rocksdb/store.rs b/src/wasm_runtime/src/state_backend/rocksdb/store.rs index e8a2ad13..b7d98e4c 100644 --- a/src/wasm_runtime/src/state_backend/rocksdb/store.rs +++ b/src/wasm_runtime/src/state_backend/rocksdb/store.rs @@ -10,9 +10,9 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::storage::state_backend::error::BackendError; -use crate::storage::state_backend::key_builder::{build_key, increment_key, is_all_0xff}; -use crate::storage::state_backend::store::{StateIterator, StateStore}; +use crate::state_backend::error::BackendError; +use crate::state_backend::key_builder::{build_key, increment_key, is_all_0xff}; +use crate::state_backend::store::{StateIterator, StateStore}; use rocksdb::{ BlockBasedOptions, Cache, ColumnFamilyDescriptor, DB, DBCompressionType, Direction, IteratorMode, Options, ReadOptions, WriteBatch, WriteOptions, diff --git a/src/wasm_runtime/src/state_backend/server.rs b/src/wasm_runtime/src/state_backend/server.rs index efeceaa4..85c3e8c3 100644 --- a/src/wasm_runtime/src/state_backend/server.rs +++ b/src/wasm_runtime/src/state_backend/server.rs @@ -12,11 +12,11 @@ use crate::config::storage::{StateStorageConfig, StateStorageType}; use crate::config::{get_state_dir, get_state_dir_for_base}; -use crate::storage::state_backend::error::BackendError; -use crate::storage::state_backend::factory::{ +use crate::state_backend::error::BackendError; +use crate::state_backend::factory::{ FactoryType, StateStoreFactory, get_factory_for_task, }; -use crate::storage::state_backend::rocksdb::RocksDBConfig; +use crate::state_backend::rocksdb::RocksDBConfig; use std::fs; use std::path::PathBuf; use std::sync::Arc; diff --git a/src/wasm_runtime/src/state_backend/store.rs b/src/wasm_runtime/src/state_backend/store.rs index aaa02988..8e61d08a 100644 --- a/src/wasm_runtime/src/state_backend/store.rs +++ b/src/wasm_runtime/src/state_backend/store.rs @@ -10,7 +10,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::storage::state_backend::error::BackendError; +use crate::state_backend::error::BackendError; pub type StateIteratorItem = Result, Vec)>, BackendError>; @@ -107,7 +107,7 @@ pub trait StateStore: Send + Sync { user_key: Vec, value: Vec, ) -> Result<(), BackendError> { - let key_bytes = crate::storage::state_backend::key_builder::build_key( + let key_bytes = crate::state_backend::key_builder::build_key( &key_group, &key, &namespace, &user_key, ); self.put_state(key_bytes, value) @@ -132,7 +132,7 @@ pub trait StateStore: Send + Sync { namespace: Vec, user_key: Vec, ) -> Result>, BackendError> { - let key_bytes = crate::storage::state_backend::key_builder::build_key( + let key_bytes = crate::state_backend::key_builder::build_key( &key_group, &key, &namespace, &user_key, ); self.get_state(key_bytes) @@ -156,7 +156,7 @@ pub trait StateStore: Send + Sync { namespace: Vec, user_key: Vec, ) -> Result<(), BackendError> { - let key_bytes = crate::storage::state_backend::key_builder::build_key( + let key_bytes = crate::state_backend::key_builder::build_key( &key_group, &key, &namespace, &user_key, ); self.delete_state(key_bytes) @@ -199,7 +199,7 @@ pub trait StateStore: Send + Sync { key: Vec, namespace: Vec, ) -> Result { - let prefix_bytes = crate::storage::state_backend::key_builder::build_key( + let prefix_bytes = crate::state_backend::key_builder::build_key( &key_group, &key, &namespace, @@ -228,13 +228,13 @@ pub trait StateStore: Send + Sync { start_inclusive: Vec, end_exclusive: Vec, ) -> Result>, BackendError> { - let start_key = crate::storage::state_backend::key_builder::build_key( + let start_key = crate::state_backend::key_builder::build_key( &key_group, &key, &namespace, &start_inclusive, ); - let end_key = crate::storage::state_backend::key_builder::build_key( + let end_key = crate::state_backend::key_builder::build_key( &key_group, &key, &namespace, @@ -279,7 +279,7 @@ pub trait StateStore: Send + Sync { key: Vec, namespace: Vec, ) -> Result, BackendError> { - let prefix = crate::storage::state_backend::key_builder::build_key( + let prefix = crate::state_backend::key_builder::build_key( &key_group, &key, &namespace, diff --git a/src/wasm_runtime/src/wasm/input/input_protocol.rs b/src/wasm_runtime/src/wasm/input/input_protocol.rs index 50294201..782a7f4f 100644 --- a/src/wasm_runtime/src/wasm/input/input_protocol.rs +++ b/src/wasm_runtime/src/wasm/input/input_protocol.rs @@ -10,7 +10,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::runtime::wasm::buffer_and_event::BufferOrEvent; +use crate::wasm::buffer_and_event::BufferOrEvent; use std::time::Duration; pub trait InputProtocol: Send + Sync + 'static { diff --git a/src/wasm_runtime/src/wasm/input/input_provider.rs b/src/wasm_runtime/src/wasm/input/input_provider.rs index 8eee649d..d7ed64b9 100644 --- a/src/wasm_runtime/src/wasm/input/input_provider.rs +++ b/src/wasm_runtime/src/wasm/input/input_provider.rs @@ -10,8 +10,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::runtime::input::Input; -use crate::runtime::wasm::task::InputConfig; +use crate::input::Input; +use crate::wasm::task::InputConfig; pub struct InputProvider; @@ -66,8 +66,8 @@ impl InputProvider { extra, runtime: _, } => { - use crate::runtime::input::InputRunner; - use crate::runtime::input::protocol::kafka::{KafkaConfig, KafkaProtocol}; + use crate::input::InputRunner; + use crate::input::protocol::kafka::{KafkaConfig, KafkaProtocol}; let servers: Vec = bootstrap_servers .split(',') diff --git a/src/wasm_runtime/src/wasm/input/input_runner.rs b/src/wasm_runtime/src/wasm/input/input_runner.rs index ece85e3d..24b96dae 100644 --- a/src/wasm_runtime/src/wasm/input/input_runner.rs +++ b/src/wasm_runtime/src/wasm/input/input_runner.rs @@ -10,13 +10,13 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::runtime::common::TaskCompletionFlag; -use crate::runtime::input::input_protocol::InputProtocol; -use crate::runtime::input::{Input, InputState}; -use crate::runtime::processor::function_error::FunctionErrorReport; -use crate::runtime::wasm::buffer_and_event::BufferOrEvent; -use crate::runtime::wasm::task::ControlMailBox; -use crate::runtime::wasm::task::InputRuntimeConfig; +use crate::common::TaskCompletionFlag; +use crate::input::input_protocol::InputProtocol; +use crate::input::{Input, InputState}; +use crate::processor::function_error::FunctionErrorReport; +use crate::wasm::buffer_and_event::BufferOrEvent; +use crate::wasm::task::ControlMailBox; +use crate::wasm::task::InputRuntimeConfig; use crossbeam_channel::{Receiver, Sender, bounded, unbounded}; use std::sync::{Arc, Mutex}; use std::thread; @@ -250,7 +250,7 @@ impl InputRunner

{ impl Input for InputRunner

{ fn init_with_context( &mut self, - init_context: &crate::runtime::wasm::taskexecutor::InitContext, + init_context: &crate::wasm::taskexecutor::InitContext, ) -> Result<(), Box> { if !matches!(*self.state.lock().unwrap(), InputState::Uninitialized) { return Ok(()); @@ -298,7 +298,7 @@ impl Input for InputRunner

{ }) .map_err(|e| Box::new(std::io::Error::other(e)) as Box)?; - use crate::runtime::processor::wasm::thread_pool::{ThreadGroup, ThreadGroupType}; + use crate::processor::wasm::thread_pool::{ThreadGroup, ThreadGroupType}; let mut group = ThreadGroup::new( ThreadGroupType::Input(self.group_id), format!("Input-g{}", self.group_id), diff --git a/src/wasm_runtime/src/wasm/input/interface.rs b/src/wasm_runtime/src/wasm/input/interface.rs index 06da4923..cbfc61cc 100644 --- a/src/wasm_runtime/src/wasm/input/interface.rs +++ b/src/wasm_runtime/src/wasm/input/interface.rs @@ -10,10 +10,10 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::runtime::wasm::buffer_and_event::BufferOrEvent; -use crate::runtime::wasm::taskexecutor::InitContext; +use crate::wasm::buffer_and_event::BufferOrEvent; +use crate::wasm::taskexecutor::InitContext; -pub use crate::runtime::common::ComponentState as InputState; +pub use crate::common::ComponentState as InputState; pub trait Input: Send + Sync { fn init_with_context( diff --git a/src/wasm_runtime/src/wasm/input/protocol/kafka/kafka_protocol.rs b/src/wasm_runtime/src/wasm/input/protocol/kafka/kafka_protocol.rs index 1fb487a6..928c9050 100644 --- a/src/wasm_runtime/src/wasm/input/protocol/kafka/kafka_protocol.rs +++ b/src/wasm_runtime/src/wasm/input/protocol/kafka/kafka_protocol.rs @@ -11,8 +11,8 @@ // limitations under the License. use super::config::KafkaConfig; -use crate::runtime::input::input_protocol::InputProtocol; -use crate::runtime::wasm::buffer_and_event::BufferOrEvent; +use crate::input::input_protocol::InputProtocol; +use crate::wasm::buffer_and_event::BufferOrEvent; use rdkafka::Message; use rdkafka::TopicPartitionList; use rdkafka::config::ClientConfig; diff --git a/src/wasm_runtime/src/wasm/output/interface.rs b/src/wasm_runtime/src/wasm/output/interface.rs index 21c3055d..e38217a1 100644 --- a/src/wasm_runtime/src/wasm/output/interface.rs +++ b/src/wasm_runtime/src/wasm/output/interface.rs @@ -10,8 +10,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::runtime::wasm::buffer_and_event::BufferOrEvent; -use crate::runtime::wasm::taskexecutor::InitContext; +use crate::wasm::buffer_and_event::BufferOrEvent; +use crate::wasm::taskexecutor::InitContext; pub trait Output: Send + Sync { fn init_with_context( diff --git a/src/wasm_runtime/src/wasm/output/output_protocol.rs b/src/wasm_runtime/src/wasm/output/output_protocol.rs index 6140d3eb..8efc5f7d 100644 --- a/src/wasm_runtime/src/wasm/output/output_protocol.rs +++ b/src/wasm_runtime/src/wasm/output/output_protocol.rs @@ -10,7 +10,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::runtime::wasm::buffer_and_event::BufferOrEvent; +use crate::wasm::buffer_and_event::BufferOrEvent; pub trait OutputProtocol: Send + Sync + 'static { fn name(&self) -> String; diff --git a/src/wasm_runtime/src/wasm/output/output_provider.rs b/src/wasm_runtime/src/wasm/output/output_provider.rs index 25ca8431..52d13544 100644 --- a/src/wasm_runtime/src/wasm/output/output_provider.rs +++ b/src/wasm_runtime/src/wasm/output/output_provider.rs @@ -10,8 +10,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::runtime::output::Output; -use crate::runtime::wasm::task::OutputConfig; +use crate::output::Output; +use crate::wasm::task::OutputConfig; pub struct OutputProvider; @@ -59,8 +59,8 @@ impl OutputProvider { extra, runtime: _, } => { - use crate::runtime::output::output_runner::OutputRunner; - use crate::runtime::output::protocol::kafka::{ + use crate::output::output_runner::OutputRunner; + use crate::output::protocol::kafka::{ KafkaOutputProtocol, KafkaProducerConfig, }; diff --git a/src/wasm_runtime/src/wasm/output/output_runner.rs b/src/wasm_runtime/src/wasm/output/output_runner.rs index ca6d780c..b732693a 100644 --- a/src/wasm_runtime/src/wasm/output/output_runner.rs +++ b/src/wasm_runtime/src/wasm/output/output_runner.rs @@ -10,13 +10,13 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::runtime::common::{ComponentState, TaskCompletionFlag}; -use crate::runtime::output::Output; -use crate::runtime::output::output_protocol::OutputProtocol; -use crate::runtime::processor::function_error::FunctionErrorReport; -use crate::runtime::wasm::buffer_and_event::BufferOrEvent; -use crate::runtime::wasm::task::ControlMailBox; -use crate::runtime::wasm::task::OutputRuntimeConfig; +use crate::common::{ComponentState, TaskCompletionFlag}; +use crate::output::Output; +use crate::output::output_protocol::OutputProtocol; +use crate::processor::function_error::FunctionErrorReport; +use crate::wasm::buffer_and_event::BufferOrEvent; +use crate::wasm::task::ControlMailBox; +use crate::wasm::task::OutputRuntimeConfig; use crossbeam_channel::{Receiver, Sender, bounded, unbounded}; use std::sync::{Arc, Mutex}; use std::thread; @@ -288,7 +288,7 @@ impl OutputRunner

{ impl Output for OutputRunner

{ fn init_with_context( &mut self, - ctx: &crate::runtime::wasm::taskexecutor::InitContext, + ctx: &crate::wasm::taskexecutor::InitContext, ) -> Result<(), Box> { if !matches!(*self.state.lock().unwrap(), ComponentState::Uninitialized) { return Ok(()); @@ -333,7 +333,7 @@ impl Output for OutputRunner

{ }) .map_err(|e| Box::new(std::io::Error::other(e)) as Box)?; - use crate::runtime::processor::wasm::thread_pool::{ThreadGroup, ThreadGroupType}; + use crate::processor::wasm::thread_pool::{ThreadGroup, ThreadGroupType}; let mut group = ThreadGroup::new( ThreadGroupType::Output(self.output_id), format!("Output-{}", self.output_id), diff --git a/src/wasm_runtime/src/wasm/output/protocol/kafka/kafka_protocol.rs b/src/wasm_runtime/src/wasm/output/protocol/kafka/kafka_protocol.rs index d9e6db4d..787804a2 100644 --- a/src/wasm_runtime/src/wasm/output/protocol/kafka/kafka_protocol.rs +++ b/src/wasm_runtime/src/wasm/output/protocol/kafka/kafka_protocol.rs @@ -11,8 +11,8 @@ // limitations under the License. use super::producer_config::KafkaProducerConfig; -use crate::runtime::output::output_protocol::OutputProtocol; -use crate::runtime::wasm::buffer_and_event::BufferOrEvent; +use crate::output::output_protocol::OutputProtocol; +use crate::wasm::buffer_and_event::BufferOrEvent; use rdkafka::producer::{BaseRecord, DefaultProducerContext, Producer, ThreadedProducer}; use std::sync::Mutex; use std::time::Duration; diff --git a/src/wasm_runtime/src/wasm/processor/wasm/thread_pool.rs b/src/wasm_runtime/src/wasm/processor/wasm/thread_pool.rs index c78eef79..d625d86a 100644 --- a/src/wasm_runtime/src/wasm/processor/wasm/thread_pool.rs +++ b/src/wasm_runtime/src/wasm/processor/wasm/thread_pool.rs @@ -10,8 +10,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::runtime::common::ComponentState; -use crate::runtime::processor::wasm::wasm_task::WasmTask; +use crate::common::ComponentState; +use crate::processor::wasm::wasm_task::WasmTask; use std::collections::HashMap; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; diff --git a/src/wasm_runtime/src/wasm/processor/wasm/wasm_host.rs b/src/wasm_runtime/src/wasm/processor/wasm/wasm_host.rs index 2bf7d4f0..9990e260 100644 --- a/src/wasm_runtime/src/wasm/processor/wasm/wasm_host.rs +++ b/src/wasm_runtime/src/wasm/processor/wasm/wasm_host.rs @@ -10,10 +10,10 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::runtime::output::Output; -use crate::runtime::processor::wasm::wasm_cache; -use crate::runtime::wasm::buffer_and_event::BufferOrEvent; -use crate::storage::state_backend::{StateStore, StateStoreFactory}; +use crate::output::Output; +use crate::processor::wasm::wasm_cache; +use crate::wasm::buffer_and_event::BufferOrEvent; +use crate::state_backend::{StateStore, StateStoreFactory}; use std::sync::{Arc, OnceLock}; use wasmtime::component::{Component, HasData, Linker, Resource, bindgen}; use wasmtime::{Config, Engine, Store}; @@ -67,7 +67,7 @@ fn get_global_engine(_wasm_size: usize) -> anyhow::Result> { bindgen!({ world: "processor", - path: "wit", + path: "../../wit", with: { "functionstream:core/kv.store": FunctionStreamStoreHandle, "functionstream:core/kv.iterator": FunctionStreamIteratorHandle, @@ -86,7 +86,7 @@ impl Drop for FunctionStreamStoreHandle { } pub struct FunctionStreamIteratorHandle { - pub state_iterator: Box, + pub state_iterator: Box, } pub struct HostState { @@ -205,7 +205,7 @@ impl HostStore for HostState { .get(&self_) .map_err(|e| Error::Other(format!("Failed to get store resource: {}", e)))?; - let real_key = crate::storage::state_backend::key_builder::build_key( + let real_key = crate::state_backend::key_builder::build_key( &key.key_group, &key.key, &key.namespace, @@ -228,7 +228,7 @@ impl HostStore for HostState { .get(&self_) .map_err(|e| Error::Other(format!("Failed to get store resource: {}", e)))?; - let real_key = crate::storage::state_backend::key_builder::build_key( + let real_key = crate::state_backend::key_builder::build_key( &key.key_group, &key.key, &key.namespace, @@ -251,7 +251,7 @@ impl HostStore for HostState { .get(&self_) .map_err(|e| Error::Other(format!("Failed to get store resource: {}", e)))?; - let real_key = crate::storage::state_backend::key_builder::build_key( + let real_key = crate::state_backend::key_builder::build_key( &key.key_group, &key.key, &key.namespace, @@ -284,7 +284,7 @@ impl HostStore for HostState { .get(&self_) .map_err(|e| Error::Other(format!("Failed to get store resource: {}", e)))?; - let prefix_key = crate::storage::state_backend::key_builder::build_key( + let prefix_key = crate::state_backend::key_builder::build_key( &key.key_group, &key.key, &key.namespace, @@ -323,13 +323,13 @@ impl HostStore for HostState { .get(&self_) .map_err(|e| Error::Other(format!("Failed to get store resource: {}", e)))?; - let start_key = crate::storage::state_backend::key_builder::build_key( + let start_key = crate::state_backend::key_builder::build_key( &key_group, &key, &namespace, &start_inclusive, ); - let end_key = crate::storage::state_backend::key_builder::build_key( + let end_key = crate::state_backend::key_builder::build_key( &key_group, &key, &namespace, @@ -449,7 +449,7 @@ pub fn create_wasm_host_with_component( engine: &Engine, component: &Component, outputs: Vec>, - init_context: &crate::runtime::wasm::taskexecutor::InitContext, + init_context: &crate::wasm::taskexecutor::InitContext, task_name: String, create_time: u64, ) -> anyhow::Result<(Processor, Store)> { @@ -495,7 +495,7 @@ pub fn create_wasm_host_with_component( pub fn create_wasm_host( wasm_bytes: &[u8], outputs: Vec>, - init_context: &crate::runtime::wasm::taskexecutor::InitContext, + init_context: &crate::wasm::taskexecutor::InitContext, task_name: String, create_time: u64, ) -> anyhow::Result<(Processor, Store)> { diff --git a/src/wasm_runtime/src/wasm/processor/wasm/wasm_processor.rs b/src/wasm_runtime/src/wasm/processor/wasm/wasm_processor.rs index 52234bfe..741280bb 100644 --- a/src/wasm_runtime/src/wasm/processor/wasm/wasm_processor.rs +++ b/src/wasm_runtime/src/wasm/processor/wasm/wasm_processor.rs @@ -17,7 +17,7 @@ use super::wasm_host::{HostState, Processor}; use super::wasm_processor_trait::WasmProcessor; -use crate::runtime::output::Output; +use crate::output::Output; use std::cell::RefCell; use std::error::Error; use std::fmt; @@ -134,7 +134,7 @@ impl WasmProcessorImpl { impl WasmProcessor for WasmProcessorImpl { fn init_with_context( &mut self, - _init_context: &crate::runtime::wasm::taskexecutor::InitContext, + _init_context: &crate::wasm::taskexecutor::InitContext, ) -> Result<(), Box> { if self.initialized { log::warn!("WasmProcessor '{}' already initialized", self.name); @@ -405,7 +405,7 @@ impl WasmProcessor for WasmProcessorImpl { fn init_wasm_host( &mut self, outputs: Vec>, - init_context: &crate::runtime::wasm::taskexecutor::InitContext, + init_context: &crate::wasm::taskexecutor::InitContext, task_name: String, create_time: u64, ) -> Result<(), Box> { diff --git a/src/wasm_runtime/src/wasm/processor/wasm/wasm_processor_trait.rs b/src/wasm_runtime/src/wasm/processor/wasm/wasm_processor_trait.rs index fb2c17fb..c699fe8b 100644 --- a/src/wasm_runtime/src/wasm/processor/wasm/wasm_processor_trait.rs +++ b/src/wasm_runtime/src/wasm/processor/wasm/wasm_processor_trait.rs @@ -10,8 +10,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::runtime::output::Output; -use crate::runtime::wasm::taskexecutor::InitContext; +use crate::output::Output; +use crate::wasm::taskexecutor::InitContext; pub trait WasmProcessor: Send + Sync { fn process( diff --git a/src/wasm_runtime/src/wasm/processor/wasm/wasm_task.rs b/src/wasm_runtime/src/wasm/processor/wasm/wasm_task.rs index 4330aaaf..45659dd9 100644 --- a/src/wasm_runtime/src/wasm/processor/wasm/wasm_task.rs +++ b/src/wasm_runtime/src/wasm/processor/wasm/wasm_task.rs @@ -13,14 +13,14 @@ use super::input_strategy::{InputStrategy, RoundRobinStrategy, from_selector_name}; use super::thread_pool::ThreadGroup; use super::wasm_processor_trait::WasmProcessor; -use crate::runtime::common::{ComponentState, TaskCompletionFlag}; -use crate::runtime::input::Input; -use crate::runtime::output::Output; -use crate::runtime::processor::function_error::FunctionErrorReport; -use crate::runtime::wasm::buffer_and_event::BufferOrEvent; -use crate::runtime::wasm::task::ProcessorRuntimeConfig; -use crate::runtime::wasm::task::{ControlMailBox, TaskControlSignal, TaskLifecycle}; -use crate::storage::task::FunctionInfo; +use crate::common::{ComponentState, TaskCompletionFlag}; +use crate::input::Input; +use crate::output::Output; +use crate::processor::function_error::FunctionErrorReport; +use crate::wasm::buffer_and_event::BufferOrEvent; +use crate::wasm::task::ProcessorRuntimeConfig; +use crate::wasm::task::{ControlMailBox, TaskControlSignal, TaskLifecycle}; +use crate::task::FunctionInfo; use crossbeam_channel::{Receiver, after, select, unbounded}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::mpsc; @@ -120,7 +120,7 @@ impl WasmTask { pub fn init_with_context( &mut self, - init_context: &crate::runtime::wasm::taskexecutor::InitContext, + init_context: &crate::wasm::taskexecutor::InitContext, ) -> Result<(), Box> { let mut inputs = self.inputs.take().ok_or_else(|| { Box::new(std::io::Error::other("inputs already moved to thread")) @@ -235,7 +235,7 @@ impl WasmTask { ))) as Box })?; - use crate::runtime::processor::wasm::thread_pool::{ThreadGroup, ThreadGroupType}; + use crate::processor::wasm::thread_pool::{ThreadGroup, ThreadGroupType}; let mut main_runloop_group = ThreadGroup::new( ThreadGroupType::MainRunloop, format!("MainRunloop-{}", self.task_name), @@ -262,7 +262,7 @@ impl WasmTask { shared_state: Arc>, failure_cause: Arc>>, execution_state: Arc>, - _init_context: crate::runtime::wasm::taskexecutor::InitContext, + _init_context: crate::wasm::taskexecutor::InitContext, ) { let mut state = TaskState::Initialized; let mut last_idx: usize = 0; @@ -729,7 +729,7 @@ impl WasmTask { impl TaskLifecycle for WasmTask { fn init_with_context( &mut self, - init_context: &crate::runtime::wasm::taskexecutor::InitContext, + init_context: &crate::wasm::taskexecutor::InitContext, ) -> Result<(), Box> { ::init_with_context(self, init_context) } diff --git a/src/wasm_runtime/src/wasm/task/builder/processor/mod.rs b/src/wasm_runtime/src/wasm/task/builder/processor/mod.rs index c1306924..1ff39a7d 100644 --- a/src/wasm_runtime/src/wasm/task/builder/processor/mod.rs +++ b/src/wasm_runtime/src/wasm/task/builder/processor/mod.rs @@ -14,13 +14,13 @@ // // Specifically handles building logic for Processor type configuration -use crate::runtime::input::{Input, InputProvider}; -use crate::runtime::output::{Output, OutputProvider}; -use crate::runtime::processor::wasm::wasm_processor::WasmProcessorImpl; -use crate::runtime::processor::wasm::wasm_processor_trait::WasmProcessor; -use crate::runtime::processor::wasm::wasm_task::WasmTask; -use crate::runtime::wasm::task::yaml_keys::{TYPE, type_values}; -use crate::runtime::wasm::task::{InputConfig, OutputConfig, ProcessorConfig, WasmTaskConfig}; +use crate::input::{Input, InputProvider}; +use crate::output::{Output, OutputProvider}; +use crate::processor::wasm::wasm_processor::WasmProcessorImpl; +use crate::processor::wasm::wasm_processor_trait::WasmProcessor; +use crate::processor::wasm::wasm_task::WasmTask; +use crate::wasm::task::yaml_keys::{TYPE, type_values}; +use crate::wasm::task::{InputConfig, OutputConfig, ProcessorConfig, WasmTaskConfig}; use serde_yaml::Value; use std::sync::Arc; diff --git a/src/wasm_runtime/src/wasm/task/builder/python/mod.rs b/src/wasm_runtime/src/wasm/task/builder/python/mod.rs index 1b31d2e5..cb48eec2 100644 --- a/src/wasm_runtime/src/wasm/task/builder/python/mod.rs +++ b/src/wasm_runtime/src/wasm/task/builder/python/mod.rs @@ -14,14 +14,14 @@ // // Specifically handles building logic for python runtime configuration -use crate::runtime::input::{Input, InputProvider}; -use crate::runtime::output::{Output, OutputProvider}; -use crate::runtime::processor::python::get_python_engine_and_component; -use crate::runtime::processor::wasm::wasm_processor::WasmProcessorImpl; -use crate::runtime::processor::wasm::wasm_processor_trait::WasmProcessor; -use crate::runtime::processor::wasm::wasm_task::WasmTask; -use crate::runtime::wasm::task::yaml_keys::{TYPE, type_values}; -use crate::runtime::wasm::task::{InputConfig, OutputConfig, ProcessorConfig, WasmTaskConfig}; +use crate::input::{Input, InputProvider}; +use crate::output::{Output, OutputProvider}; +use crate::processor::python::get_python_engine_and_component; +use crate::processor::wasm::wasm_processor::WasmProcessorImpl; +use crate::processor::wasm::wasm_processor_trait::WasmProcessor; +use crate::processor::wasm::wasm_task::WasmTask; +use crate::wasm::task::yaml_keys::{TYPE, type_values}; +use crate::wasm::task::{InputConfig, OutputConfig, ProcessorConfig, WasmTaskConfig}; use serde_yaml::Value; use std::sync::Arc; @@ -33,7 +33,7 @@ impl PythonBuilder { yaml_value: &Value, modules: &[(String, Vec)], create_time: u64, - ) -> Result, Box> + ) -> Result, Box> { let config_type = yaml_value .get(TYPE) diff --git a/src/wasm_runtime/src/wasm/task/builder/sink/mod.rs b/src/wasm_runtime/src/wasm/task/builder/sink/mod.rs index 65e8bc95..844b0aef 100644 --- a/src/wasm_runtime/src/wasm/task/builder/sink/mod.rs +++ b/src/wasm_runtime/src/wasm/task/builder/sink/mod.rs @@ -14,8 +14,8 @@ // // Specifically handles building logic for Sink type configuration (future support) -use crate::runtime::processor::wasm::wasm_task::WasmTask; -use crate::runtime::wasm::task::yaml_keys::{TYPE, type_values}; +use crate::processor::wasm::wasm_task::WasmTask; +use crate::wasm::task::yaml_keys::{TYPE, type_values}; use serde_yaml::Value; use std::sync::Arc; diff --git a/src/wasm_runtime/src/wasm/task/builder/source/mod.rs b/src/wasm_runtime/src/wasm/task/builder/source/mod.rs index fc81bea9..938a5046 100644 --- a/src/wasm_runtime/src/wasm/task/builder/source/mod.rs +++ b/src/wasm_runtime/src/wasm/task/builder/source/mod.rs @@ -14,8 +14,8 @@ // // Specifically handles building logic for Source type configuration (future support) -use crate::runtime::processor::wasm::wasm_task::WasmTask; -use crate::runtime::wasm::task::yaml_keys::{TYPE, type_values}; +use crate::processor::wasm::wasm_task::WasmTask; +use crate::wasm::task::yaml_keys::{TYPE, type_values}; use serde_yaml::Value; use std::sync::Arc; diff --git a/src/wasm_runtime/src/wasm/task/builder/task_builder.rs b/src/wasm_runtime/src/wasm/task/builder/task_builder.rs index 2246d6d8..9ca7236d 100644 --- a/src/wasm_runtime/src/wasm/task/builder/task_builder.rs +++ b/src/wasm_runtime/src/wasm/task/builder/task_builder.rs @@ -15,13 +15,13 @@ //! Provides unified factory methods to create TaskLifecycle instances from YAML config. //! Dispatches to specific builders (Processor, Source, Sink, Python) based on task type. -use crate::runtime::wasm::task::TaskLifecycle; -use crate::runtime::wasm::task::builder::processor::ProcessorBuilder; +use crate::wasm::task::TaskLifecycle; +use crate::wasm::task::builder::processor::ProcessorBuilder; #[cfg(feature = "python")] -use crate::runtime::wasm::task::builder::python::PythonBuilder; -use crate::runtime::wasm::task::builder::sink::SinkBuilder; -use crate::runtime::wasm::task::builder::source::SourceBuilder; -use crate::runtime::wasm::task::yaml_keys::{NAME, TYPE, type_values}; +use crate::wasm::task::builder::python::PythonBuilder; +use crate::wasm::task::builder::sink::SinkBuilder; +use crate::wasm::task::builder::source::SourceBuilder; +use crate::wasm::task::yaml_keys::{NAME, TYPE, type_values}; use serde_yaml::Value; use std::sync::Arc; @@ -169,7 +169,7 @@ impl TaskBuilder { /// Build and unwrap WASM task from Arc fn build_wasm_task( result: Result< - Arc, + Arc, Box, >, task_name: &str, diff --git a/src/wasm_runtime/src/wasm/task/control_mailbox.rs b/src/wasm_runtime/src/wasm/task/control_mailbox.rs index 8aaf2de3..07397ac6 100644 --- a/src/wasm_runtime/src/wasm/task/control_mailbox.rs +++ b/src/wasm_runtime/src/wasm/task/control_mailbox.rs @@ -10,8 +10,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::runtime::common::TaskCompletionFlag; -use crate::runtime::processor::function_error::FunctionErrorReport; +use crate::common::TaskCompletionFlag; +use crate::processor::function_error::FunctionErrorReport; use crossbeam_channel::Sender; #[derive(Clone)] diff --git a/src/wasm_runtime/src/wasm/task/lifecycle.rs b/src/wasm_runtime/src/wasm/task/lifecycle.rs index ea00f7c2..827cddb5 100644 --- a/src/wasm_runtime/src/wasm/task/lifecycle.rs +++ b/src/wasm_runtime/src/wasm/task/lifecycle.rs @@ -14,10 +14,10 @@ // // Defines the complete lifecycle management interface for Task, including initialization, start, stop, checkpoint and close -use crate::runtime::common::ComponentState; -use crate::runtime::wasm::task::control_mailbox::ControlMailBox; -use crate::runtime::wasm::taskexecutor::InitContext; -use crate::storage::task::FunctionInfo; +use crate::common::ComponentState; +use crate::wasm::task::control_mailbox::ControlMailBox; +use crate::wasm::taskexecutor::InitContext; +use crate::task::FunctionInfo; use std::sync::Arc; /// Task lifecycle management interface diff --git a/src/wasm_runtime/src/wasm/task/processor_config.rs b/src/wasm_runtime/src/wasm/task/processor_config.rs index a3069adc..8e642ed6 100644 --- a/src/wasm_runtime/src/wasm/task/processor_config.rs +++ b/src/wasm_runtime/src/wasm/task/processor_config.rs @@ -608,7 +608,7 @@ impl WasmTaskConfig { task_name: String, value: &Value, ) -> Result> { - use crate::runtime::wasm::task::yaml_keys::{INPUT_GROUPS, INPUTS, NAME, OUTPUTS}; + use crate::wasm::task::yaml_keys::{INPUT_GROUPS, INPUTS, NAME, OUTPUTS}; // 1. Get name from config (if exists), otherwise use the passed task_name let config_name = value diff --git a/src/wasm_runtime/src/wasm/taskexecutor/init_context.rs b/src/wasm_runtime/src/wasm/taskexecutor/init_context.rs index fca44a32..982c00e7 100644 --- a/src/wasm_runtime/src/wasm/taskexecutor/init_context.rs +++ b/src/wasm_runtime/src/wasm/taskexecutor/init_context.rs @@ -14,10 +14,10 @@ // // Provides various resources needed for task initialization, including state storage, task storage, thread pool, etc. -use crate::runtime::processor::wasm::thread_pool::{TaskThreadPool, ThreadGroup}; -use crate::runtime::wasm::task::ControlMailBox; -use crate::storage::state_backend::StateStorageServer; -use crate::storage::task::TaskStorage; +use crate::processor::wasm::thread_pool::{TaskThreadPool, ThreadGroup}; +use crate::wasm::task::ControlMailBox; +use crate::state_backend::StateStorageServer; +use crate::task::TaskStorage; use std::sync::{Arc, Mutex}; #[derive(Clone)] diff --git a/src/wasm_runtime/src/wasm/taskexecutor/task_manager.rs b/src/wasm_runtime/src/wasm/taskexecutor/task_manager.rs index 897e0a3d..0c45931c 100644 --- a/src/wasm_runtime/src/wasm/taskexecutor/task_manager.rs +++ b/src/wasm_runtime/src/wasm/taskexecutor/task_manager.rs @@ -11,12 +11,12 @@ // limitations under the License. use crate::config::GlobalConfig; -use crate::runtime::common::ComponentState; -use crate::runtime::processor::wasm::thread_pool::{GlobalTaskThreadPool, TaskThreadPool}; -use crate::runtime::wasm::task::{TaskBuilder, TaskLifecycle}; -use crate::runtime::wasm::taskexecutor::init_context::InitContext; -use crate::storage::state_backend::StateStorageServer; -use crate::storage::task::{ +use crate::common::ComponentState; +use crate::processor::wasm::thread_pool::{GlobalTaskThreadPool, TaskThreadPool}; +use crate::wasm::task::{TaskBuilder, TaskLifecycle}; +use crate::wasm::taskexecutor::init_context::InitContext; +use crate::state_backend::StateStorageServer; +use crate::task::{ FunctionInfo, StoredTaskInfo, TaskModuleBytes, TaskStorage, TaskStorageFactory, }; From 3271589bb35bb92c48c4f34c048be92663450913 Mon Sep 17 00:00:00 2001 From: luoluoyuyu Date: Wed, 13 May 2026 22:01:40 +0800 Subject: [PATCH 6/7] refactor: flatten legacy modules, fix clippy, ignore Python bytecode - Move coordinator sources from legacy/ to src/coordinator/src; embed via coordinator_body.rs; update function-stream #[path]. - Move servicer sources from legacy/; embed via servicer_body.rs; update function-stream #[path]. - Remove unused config legacy shim (src/config/src/legacy/mod.rs). - Drop empty function-stream-job-manager workspace member; refresh Cargo.lock. - CLI: collapse nested ifs in repl.rs for clippy::collapsible_if. - streaming_planner: add crate-level #![allow(dead_code)] for reserved APIs under -D warnings. - .gitignore: **/__pycache__/ and *.py[cod] so bytecode is never committed. --- .gitignore | 4 ++ Cargo.lock | 4 -- Cargo.toml | 1 - cli/cli/src/repl.rs | 36 +++++++++--------- src/config/src/legacy/mod.rs | 38 ------------------- .../src/{legacy => }/analyze/analysis.rs | 0 .../src/{legacy => }/analyze/analyzer.rs | 0 .../src/{legacy => }/analyze/mod.rs | 0 .../src/{legacy => }/coordinator.rs | 0 .../{legacy/mod.rs => coordinator_body.rs} | 0 .../src/{legacy => }/dataset/data_set.rs | 0 .../{legacy => }/dataset/execute_result.rs | 0 .../src/{legacy => }/dataset/mod.rs | 0 .../dataset/show_catalog_tables_result.rs | 0 .../show_create_streaming_table_result.rs | 0 .../dataset/show_create_table_result.rs | 0 .../dataset/show_functions_result.rs | 0 .../dataset/show_streaming_tables_result.rs | 0 .../src/{legacy => }/execution/executor.rs | 0 .../src/{legacy => }/execution/mod.rs | 0 .../src/{legacy => }/execution_context.rs | 0 .../src/{legacy => }/plan/ast_utils.rs | 0 .../{legacy => }/plan/compile_error_plan.rs | 0 .../{legacy => }/plan/create_function_plan.rs | 0 .../plan/create_python_function_plan.rs | 0 .../{legacy => }/plan/create_table_plan.rs | 0 .../src/{legacy => }/plan/ddl_compiler.rs | 0 .../{legacy => }/plan/drop_function_plan.rs | 0 .../plan/drop_streaming_table_plan.rs | 0 .../src/{legacy => }/plan/drop_table_plan.rs | 0 .../{legacy => }/plan/logical_plan_visitor.rs | 0 .../{legacy => }/plan/lookup_table_plan.rs | 0 src/coordinator/src/{legacy => }/plan/mod.rs | 0 .../src/{legacy => }/plan/optimizer.rs | 0 .../plan/show_catalog_tables_plan.rs | 0 .../plan/show_create_streaming_table_plan.rs | 0 .../plan/show_create_table_plan.rs | 0 .../{legacy => }/plan/show_functions_plan.rs | 0 .../plan/show_streaming_tables_plan.rs | 0 .../{legacy => }/plan/start_function_plan.rs | 0 .../{legacy => }/plan/stop_function_plan.rs | 0 .../{legacy => }/plan/streaming_compiler.rs | 0 .../plan/streaming_table_connector_plan.rs | 0 .../{legacy => }/plan/streaming_table_plan.rs | 0 .../src/{legacy => }/plan/visitor.rs | 0 .../src/{legacy => }/runtime_context.rs | 0 .../src/{legacy => }/sql_classify.rs | 0 .../{legacy => }/statement/create_function.rs | 0 .../statement/create_python_function.rs | 0 .../{legacy => }/statement/create_table.rs | 0 .../{legacy => }/statement/drop_function.rs | 0 .../statement/drop_streaming_table.rs | 0 .../src/{legacy => }/statement/drop_table.rs | 0 .../src/{legacy => }/statement/mod.rs | 0 .../statement/show_catalog_tables.rs | 0 .../statement/show_create_streaming_table.rs | 0 .../statement/show_create_table.rs | 0 .../{legacy => }/statement/show_functions.rs | 0 .../statement/show_streaming_tables.rs | 0 .../{legacy => }/statement/start_function.rs | 0 .../{legacy => }/statement/stop_function.rs | 0 .../{legacy => }/statement/streaming_table.rs | 0 .../src/{legacy => }/statement/visitor.rs | 0 .../{legacy => }/streaming_table_options.rs | 0 src/coordinator/src/{legacy => }/tool/mod.rs | 0 src/function-stream/src/lib.rs | 4 +- src/function-stream/src/main.rs | 4 +- src/job_manager/Cargo.toml | 8 ---- src/job_manager/src/lib.rs | 3 -- src/servicer/src/{legacy => }/handler.rs | 0 src/servicer/src/{legacy => }/initializer.rs | 0 .../src/{legacy => }/memory_service.rs | 0 src/servicer/src/{legacy => }/service.rs | 0 .../src/{legacy/mod.rs => servicer_body.rs} | 0 src/streaming_planner/src/lib.rs | 2 + 75 files changed, 28 insertions(+), 76 deletions(-) delete mode 100644 src/config/src/legacy/mod.rs rename src/coordinator/src/{legacy => }/analyze/analysis.rs (100%) rename src/coordinator/src/{legacy => }/analyze/analyzer.rs (100%) rename src/coordinator/src/{legacy => }/analyze/mod.rs (100%) rename src/coordinator/src/{legacy => }/coordinator.rs (100%) rename src/coordinator/src/{legacy/mod.rs => coordinator_body.rs} (100%) rename src/coordinator/src/{legacy => }/dataset/data_set.rs (100%) rename src/coordinator/src/{legacy => }/dataset/execute_result.rs (100%) rename src/coordinator/src/{legacy => }/dataset/mod.rs (100%) rename src/coordinator/src/{legacy => }/dataset/show_catalog_tables_result.rs (100%) rename src/coordinator/src/{legacy => }/dataset/show_create_streaming_table_result.rs (100%) rename src/coordinator/src/{legacy => }/dataset/show_create_table_result.rs (100%) rename src/coordinator/src/{legacy => }/dataset/show_functions_result.rs (100%) rename src/coordinator/src/{legacy => }/dataset/show_streaming_tables_result.rs (100%) rename src/coordinator/src/{legacy => }/execution/executor.rs (100%) rename src/coordinator/src/{legacy => }/execution/mod.rs (100%) rename src/coordinator/src/{legacy => }/execution_context.rs (100%) rename src/coordinator/src/{legacy => }/plan/ast_utils.rs (100%) rename src/coordinator/src/{legacy => }/plan/compile_error_plan.rs (100%) rename src/coordinator/src/{legacy => }/plan/create_function_plan.rs (100%) rename src/coordinator/src/{legacy => }/plan/create_python_function_plan.rs (100%) rename src/coordinator/src/{legacy => }/plan/create_table_plan.rs (100%) rename src/coordinator/src/{legacy => }/plan/ddl_compiler.rs (100%) rename src/coordinator/src/{legacy => }/plan/drop_function_plan.rs (100%) rename src/coordinator/src/{legacy => }/plan/drop_streaming_table_plan.rs (100%) rename src/coordinator/src/{legacy => }/plan/drop_table_plan.rs (100%) rename src/coordinator/src/{legacy => }/plan/logical_plan_visitor.rs (100%) rename src/coordinator/src/{legacy => }/plan/lookup_table_plan.rs (100%) rename src/coordinator/src/{legacy => }/plan/mod.rs (100%) rename src/coordinator/src/{legacy => }/plan/optimizer.rs (100%) rename src/coordinator/src/{legacy => }/plan/show_catalog_tables_plan.rs (100%) rename src/coordinator/src/{legacy => }/plan/show_create_streaming_table_plan.rs (100%) rename src/coordinator/src/{legacy => }/plan/show_create_table_plan.rs (100%) rename src/coordinator/src/{legacy => }/plan/show_functions_plan.rs (100%) rename src/coordinator/src/{legacy => }/plan/show_streaming_tables_plan.rs (100%) rename src/coordinator/src/{legacy => }/plan/start_function_plan.rs (100%) rename src/coordinator/src/{legacy => }/plan/stop_function_plan.rs (100%) rename src/coordinator/src/{legacy => }/plan/streaming_compiler.rs (100%) rename src/coordinator/src/{legacy => }/plan/streaming_table_connector_plan.rs (100%) rename src/coordinator/src/{legacy => }/plan/streaming_table_plan.rs (100%) rename src/coordinator/src/{legacy => }/plan/visitor.rs (100%) rename src/coordinator/src/{legacy => }/runtime_context.rs (100%) rename src/coordinator/src/{legacy => }/sql_classify.rs (100%) rename src/coordinator/src/{legacy => }/statement/create_function.rs (100%) rename src/coordinator/src/{legacy => }/statement/create_python_function.rs (100%) rename src/coordinator/src/{legacy => }/statement/create_table.rs (100%) rename src/coordinator/src/{legacy => }/statement/drop_function.rs (100%) rename src/coordinator/src/{legacy => }/statement/drop_streaming_table.rs (100%) rename src/coordinator/src/{legacy => }/statement/drop_table.rs (100%) rename src/coordinator/src/{legacy => }/statement/mod.rs (100%) rename src/coordinator/src/{legacy => }/statement/show_catalog_tables.rs (100%) rename src/coordinator/src/{legacy => }/statement/show_create_streaming_table.rs (100%) rename src/coordinator/src/{legacy => }/statement/show_create_table.rs (100%) rename src/coordinator/src/{legacy => }/statement/show_functions.rs (100%) rename src/coordinator/src/{legacy => }/statement/show_streaming_tables.rs (100%) rename src/coordinator/src/{legacy => }/statement/start_function.rs (100%) rename src/coordinator/src/{legacy => }/statement/stop_function.rs (100%) rename src/coordinator/src/{legacy => }/statement/streaming_table.rs (100%) rename src/coordinator/src/{legacy => }/statement/visitor.rs (100%) rename src/coordinator/src/{legacy => }/streaming_table_options.rs (100%) rename src/coordinator/src/{legacy => }/tool/mod.rs (100%) delete mode 100644 src/job_manager/Cargo.toml delete mode 100644 src/job_manager/src/lib.rs rename src/servicer/src/{legacy => }/handler.rs (100%) rename src/servicer/src/{legacy => }/initializer.rs (100%) rename src/servicer/src/{legacy => }/memory_service.rs (100%) rename src/servicer/src/{legacy => }/service.rs (100%) rename src/servicer/src/{legacy/mod.rs => servicer_body.rs} (100%) diff --git a/.gitignore b/.gitignore index 5dffeb0f..62f1dfb1 100644 --- a/.gitignore +++ b/.gitignore @@ -38,3 +38,7 @@ python/**/target/ tests/integration/target/ tests/integration/.venv/ tests/integration/install + +# Python bytecode (do not commit) +**/__pycache__/ +*.py[cod] diff --git a/Cargo.lock b/Cargo.lock index e3e4b95a..a1176666 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3891,10 +3891,6 @@ dependencies = [ name = "function-stream-coordinator" version = "0.6.0" -[[package]] -name = "function-stream-job-manager" -version = "0.6.0" - [[package]] name = "function-stream-logger" version = "0.6.0" diff --git a/Cargo.toml b/Cargo.toml index 96adb4b3..324c882b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,7 +8,6 @@ members = [ "src/common", "src/config", "src/coordinator", - "src/job_manager", "src/logger", "src/runtime_common", "src/servicer", diff --git a/cli/cli/src/repl.rs b/cli/cli/src/repl.rs index b442bd07..7761a2a7 100644 --- a/cli/cli/src/repl.rs +++ b/cli/cli/src/repl.rs @@ -158,17 +158,17 @@ impl Repl { } // 3. Strict Data Check: Only proceed if data is explicitly present and non-empty - if let Some(bytes) = response.data { - if !bytes.is_empty() { - // format_arrow_data returns Ok(Some(Table)) ONLY if row_count > 0 - match self.format_arrow_data(&bytes) { - Ok(Some(table)) => println!("{}", table), - Ok(None) => { - // Data was present but contained 0 rows (e.g., empty result set) - // We print nothing here to keep output clean as requested - } - Err(e) => eprintln!("Failed to parse result data: {}", e), + if let Some(bytes) = response.data + && !bytes.is_empty() + { + // format_arrow_data returns Ok(Some(Table)) ONLY if row_count > 0 + match self.format_arrow_data(&bytes) { + Ok(Some(table)) => println!("{}", table), + Ok(None) => { + // Data was present but contained 0 rows (e.g., empty result set) + // We print nothing here to keep output clean as requested } + Err(e) => eprintln!("Failed to parse result data: {}", e), } } @@ -403,10 +403,10 @@ impl Repl { println!(); } - if !skip_save_history { - if let Some(ref mut ed) = repl.lock().await.editor { - let _ = ed.save_history(".function-stream-cli-history"); - } + if !skip_save_history + && let Some(ref mut ed) = repl.lock().await.editor + { + let _ = ed.save_history(".function-stream-cli-history"); } Ok(()) } @@ -448,10 +448,10 @@ impl Repl { } fn add_history_entry(&mut self, entry: &str) { - if let Some(ed) = self.editor.as_mut() { - if !entry.trim().is_empty() { - let _ = ed.add_history_entry(entry.trim()); - } + if let Some(ed) = self.editor.as_mut() + && !entry.trim().is_empty() + { + let _ = ed.add_history_entry(entry.trim()); } } diff --git a/src/config/src/legacy/mod.rs b/src/config/src/legacy/mod.rs deleted file mode 100644 index e60dcfde..00000000 --- a/src/config/src/legacy/mod.rs +++ /dev/null @@ -1,38 +0,0 @@ -// 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. - -pub mod global_config; -pub mod loader; -pub mod log_config; -pub mod paths; -pub mod python_config; -pub mod service_config; -pub mod storage; -pub mod streaming_job; -pub mod system; -pub mod wasm_config; - -pub use global_config::{ - DEFAULT_OPERATOR_STATE_STORE_MEMORY_BYTES, DEFAULT_STREAMING_RUNTIME_MEMORY_BYTES, GlobalConfig, -}; -pub use loader::load_global_config; -pub use log_config::LogConfig; -#[allow(unused_imports)] -pub use paths::{ - ENV_CONF, ENV_HOME, find_config_file, get_app_log_path, get_conf_dir, get_data_dir, - get_log_path, get_logs_dir, get_project_root, get_python_cache_dir, get_python_cwasm_path, - get_python_wasm_path, get_state_dir, get_state_dir_for_base, get_task_dir, get_wasm_cache_dir, - resolve_path, -}; -#[cfg(feature = "python")] -pub use python_config::PythonConfig; -pub use streaming_job::{DEFAULT_CHECKPOINT_INTERVAL_MS, DEFAULT_PIPELINE_PARALLELISM}; diff --git a/src/coordinator/src/legacy/analyze/analysis.rs b/src/coordinator/src/analyze/analysis.rs similarity index 100% rename from src/coordinator/src/legacy/analyze/analysis.rs rename to src/coordinator/src/analyze/analysis.rs diff --git a/src/coordinator/src/legacy/analyze/analyzer.rs b/src/coordinator/src/analyze/analyzer.rs similarity index 100% rename from src/coordinator/src/legacy/analyze/analyzer.rs rename to src/coordinator/src/analyze/analyzer.rs diff --git a/src/coordinator/src/legacy/analyze/mod.rs b/src/coordinator/src/analyze/mod.rs similarity index 100% rename from src/coordinator/src/legacy/analyze/mod.rs rename to src/coordinator/src/analyze/mod.rs diff --git a/src/coordinator/src/legacy/coordinator.rs b/src/coordinator/src/coordinator.rs similarity index 100% rename from src/coordinator/src/legacy/coordinator.rs rename to src/coordinator/src/coordinator.rs diff --git a/src/coordinator/src/legacy/mod.rs b/src/coordinator/src/coordinator_body.rs similarity index 100% rename from src/coordinator/src/legacy/mod.rs rename to src/coordinator/src/coordinator_body.rs diff --git a/src/coordinator/src/legacy/dataset/data_set.rs b/src/coordinator/src/dataset/data_set.rs similarity index 100% rename from src/coordinator/src/legacy/dataset/data_set.rs rename to src/coordinator/src/dataset/data_set.rs diff --git a/src/coordinator/src/legacy/dataset/execute_result.rs b/src/coordinator/src/dataset/execute_result.rs similarity index 100% rename from src/coordinator/src/legacy/dataset/execute_result.rs rename to src/coordinator/src/dataset/execute_result.rs diff --git a/src/coordinator/src/legacy/dataset/mod.rs b/src/coordinator/src/dataset/mod.rs similarity index 100% rename from src/coordinator/src/legacy/dataset/mod.rs rename to src/coordinator/src/dataset/mod.rs diff --git a/src/coordinator/src/legacy/dataset/show_catalog_tables_result.rs b/src/coordinator/src/dataset/show_catalog_tables_result.rs similarity index 100% rename from src/coordinator/src/legacy/dataset/show_catalog_tables_result.rs rename to src/coordinator/src/dataset/show_catalog_tables_result.rs diff --git a/src/coordinator/src/legacy/dataset/show_create_streaming_table_result.rs b/src/coordinator/src/dataset/show_create_streaming_table_result.rs similarity index 100% rename from src/coordinator/src/legacy/dataset/show_create_streaming_table_result.rs rename to src/coordinator/src/dataset/show_create_streaming_table_result.rs diff --git a/src/coordinator/src/legacy/dataset/show_create_table_result.rs b/src/coordinator/src/dataset/show_create_table_result.rs similarity index 100% rename from src/coordinator/src/legacy/dataset/show_create_table_result.rs rename to src/coordinator/src/dataset/show_create_table_result.rs diff --git a/src/coordinator/src/legacy/dataset/show_functions_result.rs b/src/coordinator/src/dataset/show_functions_result.rs similarity index 100% rename from src/coordinator/src/legacy/dataset/show_functions_result.rs rename to src/coordinator/src/dataset/show_functions_result.rs diff --git a/src/coordinator/src/legacy/dataset/show_streaming_tables_result.rs b/src/coordinator/src/dataset/show_streaming_tables_result.rs similarity index 100% rename from src/coordinator/src/legacy/dataset/show_streaming_tables_result.rs rename to src/coordinator/src/dataset/show_streaming_tables_result.rs diff --git a/src/coordinator/src/legacy/execution/executor.rs b/src/coordinator/src/execution/executor.rs similarity index 100% rename from src/coordinator/src/legacy/execution/executor.rs rename to src/coordinator/src/execution/executor.rs diff --git a/src/coordinator/src/legacy/execution/mod.rs b/src/coordinator/src/execution/mod.rs similarity index 100% rename from src/coordinator/src/legacy/execution/mod.rs rename to src/coordinator/src/execution/mod.rs diff --git a/src/coordinator/src/legacy/execution_context.rs b/src/coordinator/src/execution_context.rs similarity index 100% rename from src/coordinator/src/legacy/execution_context.rs rename to src/coordinator/src/execution_context.rs diff --git a/src/coordinator/src/legacy/plan/ast_utils.rs b/src/coordinator/src/plan/ast_utils.rs similarity index 100% rename from src/coordinator/src/legacy/plan/ast_utils.rs rename to src/coordinator/src/plan/ast_utils.rs diff --git a/src/coordinator/src/legacy/plan/compile_error_plan.rs b/src/coordinator/src/plan/compile_error_plan.rs similarity index 100% rename from src/coordinator/src/legacy/plan/compile_error_plan.rs rename to src/coordinator/src/plan/compile_error_plan.rs diff --git a/src/coordinator/src/legacy/plan/create_function_plan.rs b/src/coordinator/src/plan/create_function_plan.rs similarity index 100% rename from src/coordinator/src/legacy/plan/create_function_plan.rs rename to src/coordinator/src/plan/create_function_plan.rs diff --git a/src/coordinator/src/legacy/plan/create_python_function_plan.rs b/src/coordinator/src/plan/create_python_function_plan.rs similarity index 100% rename from src/coordinator/src/legacy/plan/create_python_function_plan.rs rename to src/coordinator/src/plan/create_python_function_plan.rs diff --git a/src/coordinator/src/legacy/plan/create_table_plan.rs b/src/coordinator/src/plan/create_table_plan.rs similarity index 100% rename from src/coordinator/src/legacy/plan/create_table_plan.rs rename to src/coordinator/src/plan/create_table_plan.rs diff --git a/src/coordinator/src/legacy/plan/ddl_compiler.rs b/src/coordinator/src/plan/ddl_compiler.rs similarity index 100% rename from src/coordinator/src/legacy/plan/ddl_compiler.rs rename to src/coordinator/src/plan/ddl_compiler.rs diff --git a/src/coordinator/src/legacy/plan/drop_function_plan.rs b/src/coordinator/src/plan/drop_function_plan.rs similarity index 100% rename from src/coordinator/src/legacy/plan/drop_function_plan.rs rename to src/coordinator/src/plan/drop_function_plan.rs diff --git a/src/coordinator/src/legacy/plan/drop_streaming_table_plan.rs b/src/coordinator/src/plan/drop_streaming_table_plan.rs similarity index 100% rename from src/coordinator/src/legacy/plan/drop_streaming_table_plan.rs rename to src/coordinator/src/plan/drop_streaming_table_plan.rs diff --git a/src/coordinator/src/legacy/plan/drop_table_plan.rs b/src/coordinator/src/plan/drop_table_plan.rs similarity index 100% rename from src/coordinator/src/legacy/plan/drop_table_plan.rs rename to src/coordinator/src/plan/drop_table_plan.rs diff --git a/src/coordinator/src/legacy/plan/logical_plan_visitor.rs b/src/coordinator/src/plan/logical_plan_visitor.rs similarity index 100% rename from src/coordinator/src/legacy/plan/logical_plan_visitor.rs rename to src/coordinator/src/plan/logical_plan_visitor.rs diff --git a/src/coordinator/src/legacy/plan/lookup_table_plan.rs b/src/coordinator/src/plan/lookup_table_plan.rs similarity index 100% rename from src/coordinator/src/legacy/plan/lookup_table_plan.rs rename to src/coordinator/src/plan/lookup_table_plan.rs diff --git a/src/coordinator/src/legacy/plan/mod.rs b/src/coordinator/src/plan/mod.rs similarity index 100% rename from src/coordinator/src/legacy/plan/mod.rs rename to src/coordinator/src/plan/mod.rs diff --git a/src/coordinator/src/legacy/plan/optimizer.rs b/src/coordinator/src/plan/optimizer.rs similarity index 100% rename from src/coordinator/src/legacy/plan/optimizer.rs rename to src/coordinator/src/plan/optimizer.rs diff --git a/src/coordinator/src/legacy/plan/show_catalog_tables_plan.rs b/src/coordinator/src/plan/show_catalog_tables_plan.rs similarity index 100% rename from src/coordinator/src/legacy/plan/show_catalog_tables_plan.rs rename to src/coordinator/src/plan/show_catalog_tables_plan.rs diff --git a/src/coordinator/src/legacy/plan/show_create_streaming_table_plan.rs b/src/coordinator/src/plan/show_create_streaming_table_plan.rs similarity index 100% rename from src/coordinator/src/legacy/plan/show_create_streaming_table_plan.rs rename to src/coordinator/src/plan/show_create_streaming_table_plan.rs diff --git a/src/coordinator/src/legacy/plan/show_create_table_plan.rs b/src/coordinator/src/plan/show_create_table_plan.rs similarity index 100% rename from src/coordinator/src/legacy/plan/show_create_table_plan.rs rename to src/coordinator/src/plan/show_create_table_plan.rs diff --git a/src/coordinator/src/legacy/plan/show_functions_plan.rs b/src/coordinator/src/plan/show_functions_plan.rs similarity index 100% rename from src/coordinator/src/legacy/plan/show_functions_plan.rs rename to src/coordinator/src/plan/show_functions_plan.rs diff --git a/src/coordinator/src/legacy/plan/show_streaming_tables_plan.rs b/src/coordinator/src/plan/show_streaming_tables_plan.rs similarity index 100% rename from src/coordinator/src/legacy/plan/show_streaming_tables_plan.rs rename to src/coordinator/src/plan/show_streaming_tables_plan.rs diff --git a/src/coordinator/src/legacy/plan/start_function_plan.rs b/src/coordinator/src/plan/start_function_plan.rs similarity index 100% rename from src/coordinator/src/legacy/plan/start_function_plan.rs rename to src/coordinator/src/plan/start_function_plan.rs diff --git a/src/coordinator/src/legacy/plan/stop_function_plan.rs b/src/coordinator/src/plan/stop_function_plan.rs similarity index 100% rename from src/coordinator/src/legacy/plan/stop_function_plan.rs rename to src/coordinator/src/plan/stop_function_plan.rs diff --git a/src/coordinator/src/legacy/plan/streaming_compiler.rs b/src/coordinator/src/plan/streaming_compiler.rs similarity index 100% rename from src/coordinator/src/legacy/plan/streaming_compiler.rs rename to src/coordinator/src/plan/streaming_compiler.rs diff --git a/src/coordinator/src/legacy/plan/streaming_table_connector_plan.rs b/src/coordinator/src/plan/streaming_table_connector_plan.rs similarity index 100% rename from src/coordinator/src/legacy/plan/streaming_table_connector_plan.rs rename to src/coordinator/src/plan/streaming_table_connector_plan.rs diff --git a/src/coordinator/src/legacy/plan/streaming_table_plan.rs b/src/coordinator/src/plan/streaming_table_plan.rs similarity index 100% rename from src/coordinator/src/legacy/plan/streaming_table_plan.rs rename to src/coordinator/src/plan/streaming_table_plan.rs diff --git a/src/coordinator/src/legacy/plan/visitor.rs b/src/coordinator/src/plan/visitor.rs similarity index 100% rename from src/coordinator/src/legacy/plan/visitor.rs rename to src/coordinator/src/plan/visitor.rs diff --git a/src/coordinator/src/legacy/runtime_context.rs b/src/coordinator/src/runtime_context.rs similarity index 100% rename from src/coordinator/src/legacy/runtime_context.rs rename to src/coordinator/src/runtime_context.rs diff --git a/src/coordinator/src/legacy/sql_classify.rs b/src/coordinator/src/sql_classify.rs similarity index 100% rename from src/coordinator/src/legacy/sql_classify.rs rename to src/coordinator/src/sql_classify.rs diff --git a/src/coordinator/src/legacy/statement/create_function.rs b/src/coordinator/src/statement/create_function.rs similarity index 100% rename from src/coordinator/src/legacy/statement/create_function.rs rename to src/coordinator/src/statement/create_function.rs diff --git a/src/coordinator/src/legacy/statement/create_python_function.rs b/src/coordinator/src/statement/create_python_function.rs similarity index 100% rename from src/coordinator/src/legacy/statement/create_python_function.rs rename to src/coordinator/src/statement/create_python_function.rs diff --git a/src/coordinator/src/legacy/statement/create_table.rs b/src/coordinator/src/statement/create_table.rs similarity index 100% rename from src/coordinator/src/legacy/statement/create_table.rs rename to src/coordinator/src/statement/create_table.rs diff --git a/src/coordinator/src/legacy/statement/drop_function.rs b/src/coordinator/src/statement/drop_function.rs similarity index 100% rename from src/coordinator/src/legacy/statement/drop_function.rs rename to src/coordinator/src/statement/drop_function.rs diff --git a/src/coordinator/src/legacy/statement/drop_streaming_table.rs b/src/coordinator/src/statement/drop_streaming_table.rs similarity index 100% rename from src/coordinator/src/legacy/statement/drop_streaming_table.rs rename to src/coordinator/src/statement/drop_streaming_table.rs diff --git a/src/coordinator/src/legacy/statement/drop_table.rs b/src/coordinator/src/statement/drop_table.rs similarity index 100% rename from src/coordinator/src/legacy/statement/drop_table.rs rename to src/coordinator/src/statement/drop_table.rs diff --git a/src/coordinator/src/legacy/statement/mod.rs b/src/coordinator/src/statement/mod.rs similarity index 100% rename from src/coordinator/src/legacy/statement/mod.rs rename to src/coordinator/src/statement/mod.rs diff --git a/src/coordinator/src/legacy/statement/show_catalog_tables.rs b/src/coordinator/src/statement/show_catalog_tables.rs similarity index 100% rename from src/coordinator/src/legacy/statement/show_catalog_tables.rs rename to src/coordinator/src/statement/show_catalog_tables.rs diff --git a/src/coordinator/src/legacy/statement/show_create_streaming_table.rs b/src/coordinator/src/statement/show_create_streaming_table.rs similarity index 100% rename from src/coordinator/src/legacy/statement/show_create_streaming_table.rs rename to src/coordinator/src/statement/show_create_streaming_table.rs diff --git a/src/coordinator/src/legacy/statement/show_create_table.rs b/src/coordinator/src/statement/show_create_table.rs similarity index 100% rename from src/coordinator/src/legacy/statement/show_create_table.rs rename to src/coordinator/src/statement/show_create_table.rs diff --git a/src/coordinator/src/legacy/statement/show_functions.rs b/src/coordinator/src/statement/show_functions.rs similarity index 100% rename from src/coordinator/src/legacy/statement/show_functions.rs rename to src/coordinator/src/statement/show_functions.rs diff --git a/src/coordinator/src/legacy/statement/show_streaming_tables.rs b/src/coordinator/src/statement/show_streaming_tables.rs similarity index 100% rename from src/coordinator/src/legacy/statement/show_streaming_tables.rs rename to src/coordinator/src/statement/show_streaming_tables.rs diff --git a/src/coordinator/src/legacy/statement/start_function.rs b/src/coordinator/src/statement/start_function.rs similarity index 100% rename from src/coordinator/src/legacy/statement/start_function.rs rename to src/coordinator/src/statement/start_function.rs diff --git a/src/coordinator/src/legacy/statement/stop_function.rs b/src/coordinator/src/statement/stop_function.rs similarity index 100% rename from src/coordinator/src/legacy/statement/stop_function.rs rename to src/coordinator/src/statement/stop_function.rs diff --git a/src/coordinator/src/legacy/statement/streaming_table.rs b/src/coordinator/src/statement/streaming_table.rs similarity index 100% rename from src/coordinator/src/legacy/statement/streaming_table.rs rename to src/coordinator/src/statement/streaming_table.rs diff --git a/src/coordinator/src/legacy/statement/visitor.rs b/src/coordinator/src/statement/visitor.rs similarity index 100% rename from src/coordinator/src/legacy/statement/visitor.rs rename to src/coordinator/src/statement/visitor.rs diff --git a/src/coordinator/src/legacy/streaming_table_options.rs b/src/coordinator/src/streaming_table_options.rs similarity index 100% rename from src/coordinator/src/legacy/streaming_table_options.rs rename to src/coordinator/src/streaming_table_options.rs diff --git a/src/coordinator/src/legacy/tool/mod.rs b/src/coordinator/src/tool/mod.rs similarity index 100% rename from src/coordinator/src/legacy/tool/mod.rs rename to src/coordinator/src/tool/mod.rs diff --git a/src/function-stream/src/lib.rs b/src/function-stream/src/lib.rs index fac7e8f1..a3fc9234 100644 --- a/src/function-stream/src/lib.rs +++ b/src/function-stream/src/lib.rs @@ -19,7 +19,7 @@ use std::sync::Arc; use anyhow::Context; pub use function_stream_config as config; -#[path = "../../coordinator/src/legacy/mod.rs"] +#[path = "../../coordinator/src/coordinator_body.rs"] pub mod coordinator; pub use function_stream_logger as logging; @@ -79,6 +79,6 @@ pub fn initialize_stream_catalog(config: &crate::config::GlobalConfig) -> anyhow CatalogManager::init_global(store).context("Stream catalog (CatalogManager) global init failed") } -#[path = "../../servicer/src/legacy/mod.rs"] +#[path = "../../servicer/src/servicer_body.rs"] pub mod server; pub use function_stream_streaming_planner as sql; diff --git a/src/function-stream/src/main.rs b/src/function-stream/src/main.rs index 33971370..b7e1093f 100644 --- a/src/function-stream/src/main.rs +++ b/src/function-stream/src/main.rs @@ -13,7 +13,7 @@ #![allow(dead_code)] pub use function_stream_config as config; -#[path = "../../coordinator/src/legacy/mod.rs"] +#[path = "../../coordinator/src/coordinator_body.rs"] mod coordinator; pub use function_stream_logger as logging; @@ -74,7 +74,7 @@ pub fn initialize_stream_catalog(config: &crate::config::GlobalConfig) -> anyhow CatalogManager::init_global(store).context("Stream catalog (CatalogManager) global init failed") } -#[path = "../../servicer/src/legacy/mod.rs"] +#[path = "../../servicer/src/servicer_body.rs"] mod server; pub use function_stream_streaming_planner as sql; use std::thread; diff --git a/src/job_manager/Cargo.toml b/src/job_manager/Cargo.toml deleted file mode 100644 index 8edfb464..00000000 --- a/src/job_manager/Cargo.toml +++ /dev/null @@ -1,8 +0,0 @@ -[package] -name = "function-stream-job-manager" -version.workspace = true -edition.workspace = true - -[lib] -name = "function_stream_job_manager" -path = "src/lib.rs" diff --git a/src/job_manager/src/lib.rs b/src/job_manager/src/lib.rs deleted file mode 100644 index 6877187d..00000000 --- a/src/job_manager/src/lib.rs +++ /dev/null @@ -1,3 +0,0 @@ -//! Job lifecycle management and scheduling. - -pub const CRATE_NAME: &str = "function-stream-job-manager"; diff --git a/src/servicer/src/legacy/handler.rs b/src/servicer/src/handler.rs similarity index 100% rename from src/servicer/src/legacy/handler.rs rename to src/servicer/src/handler.rs diff --git a/src/servicer/src/legacy/initializer.rs b/src/servicer/src/initializer.rs similarity index 100% rename from src/servicer/src/legacy/initializer.rs rename to src/servicer/src/initializer.rs diff --git a/src/servicer/src/legacy/memory_service.rs b/src/servicer/src/memory_service.rs similarity index 100% rename from src/servicer/src/legacy/memory_service.rs rename to src/servicer/src/memory_service.rs diff --git a/src/servicer/src/legacy/service.rs b/src/servicer/src/service.rs similarity index 100% rename from src/servicer/src/legacy/service.rs rename to src/servicer/src/service.rs diff --git a/src/servicer/src/legacy/mod.rs b/src/servicer/src/servicer_body.rs similarity index 100% rename from src/servicer/src/legacy/mod.rs rename to src/servicer/src/servicer_body.rs diff --git a/src/streaming_planner/src/lib.rs b/src/streaming_planner/src/lib.rs index 040e0d42..2dbc2b47 100644 --- a/src/streaming_planner/src/lib.rs +++ b/src/streaming_planner/src/lib.rs @@ -10,6 +10,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +#![allow(dead_code)] // Planner keeps helpers/types for upcoming features; strict -D dead_code is noisy. + //! Streaming SQL planning (logical graph, connectors, schema, physical codec). //! //! The `function-stream` binary/library re-exports this crate as `function_stream::sql` for From a9ba1c10908c8d8ed9bcce92d0512544df70589a Mon Sep 17 00:00:00 2001 From: luoluoyuyu Date: Wed, 13 May 2026 22:13:53 +0800 Subject: [PATCH 7/7] style: apply rustfmt across workspace crates Run rustfmt on CLI, coordinator, catalog_storage, streaming_planner, streaming_runtime, and wasm_runtime so cargo fmt --check passes for the main binary and CLI packages. --- cli/cli/src/repl.rs | 14 +++------ .../src/task/rocksdb_storage.rs | 2 +- src/coordinator/src/coordinator_body.rs | 2 +- src/coordinator/src/execution/executor.rs | 6 ++-- src/coordinator/src/runtime_context.rs | 4 +-- src/coordinator/src/sql_classify.rs | 15 ++++----- src/streaming_planner/src/common/mod.rs | 2 +- .../src/connector/sink/runtime_config.rs | 4 +-- .../src/logical_node/aggregate.rs | 2 +- .../src/logical_node/async_udf.rs | 2 +- .../src/logical_node/debezium.rs | 2 +- .../src/logical_node/is_retract.rs | 4 +-- .../src/logical_node/key_calculation.rs | 2 +- .../src/logical_node/lookup.rs | 2 +- .../src/logical_node/projection.rs | 2 +- .../src/logical_node/remote_table.rs | 2 +- .../src/logical_node/sink.rs | 2 +- .../src/logical_node/table_source.rs | 2 +- .../src/logical_node/timestamp_append.rs | 2 +- .../src/logical_node/updating_aggregate.rs | 4 +-- .../src/logical_node/watermark_node.rs | 2 +- .../src/logical_planner/streaming_planner.rs | 15 ++------- src/streaming_planner/src/parse.rs | 2 +- src/streaming_planner/src/physical/udfs.rs | 2 +- src/streaming_planner/src/planning_runtime.rs | 2 +- src/streaming_planner/src/schema/catalog.rs | 2 +- .../src/streaming/api/operator.rs | 2 +- .../src/streaming/api/source.rs | 2 +- .../src/streaming/execution/operator_chain.rs | 2 +- .../src/streaming/execution/pipeline.rs | 2 +- .../src/streaming/execution/source_driver.rs | 2 +- .../execution/tracker/watermark_tracker.rs | 2 +- .../src/streaming/factory/connector/delta.rs | 4 +-- .../factory/connector/dispatchers.rs | 2 +- .../streaming/factory/connector/filesystem.rs | 4 +-- .../streaming/factory/connector/iceberg.rs | 4 +-- .../src/streaming/factory/connector/kafka.rs | 6 ++-- .../streaming/factory/connector/lancedb.rs | 4 +-- .../src/streaming/factory/connector/s3.rs | 8 ++--- .../src/streaming/factory/operator_factory.rs | 8 ++--- .../src/streaming/job/job_manager.rs | 4 +-- .../grouping/incremental_aggregate.rs | 8 ++--- .../operators/joins/join_instance.rs | 4 +-- .../operators/joins/join_with_expiration.rs | 4 +-- .../src/streaming/operators/key_by.rs | 2 +- .../src/streaming/operators/key_operator.rs | 2 +- .../src/streaming/operators/projection.rs | 4 +-- .../src/streaming/operators/sink/delta/mod.rs | 6 ++-- .../operators/sink/filesystem/mod.rs | 4 +-- .../streaming/operators/sink/iceberg/mod.rs | 6 ++-- .../src/streaming/operators/sink/kafka/mod.rs | 4 +-- .../streaming/operators/sink/lancedb/mod.rs | 4 +-- .../src/streaming/operators/sink/s3/mod.rs | 6 ++-- .../streaming/operators/source/kafka/mod.rs | 4 +-- .../operators/stateless_physical_executor.rs | 2 +- .../streaming/operators/value_execution.rs | 2 +- .../watermark/watermark_generator.rs | 2 +- .../windows/session_aggregating_window.rs | 10 +++--- .../windows/sliding_aggregating_window.rs | 4 +-- .../windows/tumbling_aggregating_window.rs | 8 ++--- .../operators/windows/window_function.rs | 10 +++--- src/wasm_runtime/src/state_backend/factory.rs | 5 ++- .../src/state_backend/memory/store.rs | 5 ++- src/wasm_runtime/src/state_backend/server.rs | 4 +-- src/wasm_runtime/src/state_backend/store.rs | 31 ++++++------------- .../src/wasm/output/output_provider.rs | 4 +-- .../src/wasm/processor/wasm/wasm_host.rs | 2 +- .../src/wasm/processor/wasm/wasm_task.rs | 2 +- .../src/wasm/task/builder/python/mod.rs | 3 +- src/wasm_runtime/src/wasm/task/lifecycle.rs | 2 +- .../src/wasm/taskexecutor/init_context.rs | 2 +- .../src/wasm/taskexecutor/task_manager.rs | 8 ++--- 72 files changed, 138 insertions(+), 184 deletions(-) diff --git a/cli/cli/src/repl.rs b/cli/cli/src/repl.rs index 7761a2a7..01209ba1 100644 --- a/cli/cli/src/repl.rs +++ b/cli/cli/src/repl.rs @@ -17,7 +17,7 @@ use arrow_ipc::reader::StreamReader; use arrow_schema::DataType; use comfy_table::presets::UTF8_FULL; use comfy_table::{Attribute, Cell, Color, ContentArrangement, Table, TableComponent}; -use protocol::cli::{function_stream_service_client::FunctionStreamServiceClient, SqlRequest}; +use protocol::cli::{SqlRequest, function_stream_service_client::FunctionStreamServiceClient}; use rustyline::error::ReadlineError; use rustyline::{Config, DefaultEditor, EditMode}; use std::fmt; @@ -243,11 +243,7 @@ impl Repl { } } - if has_rows { - Ok(Some(table)) - } else { - Ok(None) - } + if has_rows { Ok(Some(table)) } else { Ok(None) } } fn extract_value(&self, column: &dyn Array, row: usize) -> String { @@ -317,7 +313,7 @@ impl Repl { #[cfg(unix)] let mut sigterm = { - use tokio::signal::unix::{signal, SignalKind}; + use tokio::signal::unix::{SignalKind, signal}; signal(SignalKind::terminate()).expect("failed to register SIGTERM handler") }; @@ -403,9 +399,7 @@ impl Repl { println!(); } - if !skip_save_history - && let Some(ref mut ed) = repl.lock().await.editor - { + if !skip_save_history && let Some(ref mut ed) = repl.lock().await.editor { let _ = ed.save_history(".function-stream-cli-history"); } Ok(()) diff --git a/src/catalog_storage/src/task/rocksdb_storage.rs b/src/catalog_storage/src/task/rocksdb_storage.rs index 62ffb2a0..a3927515 100644 --- a/src/catalog_storage/src/task/rocksdb_storage.rs +++ b/src/catalog_storage/src/task/rocksdb_storage.rs @@ -19,8 +19,8 @@ use super::proto_codec::{ encode_task_module_bytes, }; use super::storage::{StoredTaskInfo, TaskStorage}; -use crate::config::storage::RocksDBStorageConfig; use crate::common::ComponentState; +use crate::config::storage::RocksDBStorageConfig; use anyhow::{Context, Result, anyhow}; use rocksdb::{ColumnFamilyDescriptor, DB, IteratorMode, Options, WriteBatch}; use std::path::Path; diff --git a/src/coordinator/src/coordinator_body.rs b/src/coordinator/src/coordinator_body.rs index 640d7692..0b5b7135 100644 --- a/src/coordinator/src/coordinator_body.rs +++ b/src/coordinator/src/coordinator_body.rs @@ -18,8 +18,8 @@ mod execution; mod execution_context; mod plan; mod runtime_context; -mod statement; mod sql_classify; +mod statement; mod streaming_table_options; mod tool; diff --git a/src/coordinator/src/execution/executor.rs b/src/coordinator/src/execution/executor.rs index 7109a1d2..a58dc6c0 100644 --- a/src/coordinator/src/execution/executor.rs +++ b/src/coordinator/src/execution/executor.rs @@ -31,13 +31,13 @@ use crate::coordinator::statement::{ConfigSource, FunctionSource}; use crate::coordinator::streaming_table_options::{ parse_checkpoint_interval_ms, parse_pipeline_parallelism, }; -use crate::streaming::job::JobManager; -use crate::streaming::protocol::control::StopMode; -use crate::wasm::taskexecutor::TaskManager; use crate::sql::schema::catalog::ExternalTable; use crate::sql::schema::show_create_catalog_table; use crate::sql::schema::table::CatalogEntity; use crate::stream_catalog::CatalogManager; +use crate::streaming::job::JobManager; +use crate::streaming::protocol::control::StopMode; +use crate::wasm::taskexecutor::TaskManager; #[derive(Error, Debug)] pub enum ExecuteError { diff --git a/src/coordinator/src/runtime_context.rs b/src/coordinator/src/runtime_context.rs index 513c703c..2f8dc2a7 100644 --- a/src/coordinator/src/runtime_context.rs +++ b/src/coordinator/src/runtime_context.rs @@ -16,10 +16,10 @@ use std::sync::Arc; use anyhow::Result; -use crate::streaming::job::JobManager; -use crate::wasm::taskexecutor::TaskManager; use crate::sql::schema::StreamSchemaProvider; use crate::stream_catalog::CatalogManager; +use crate::streaming::job::JobManager; +use crate::wasm::taskexecutor::TaskManager; /// Dependencies shared by analyze / plan / execute, analogous to installing globals in /// [`TaskManager`], [`CatalogManager`], and [`JobManager`]. diff --git a/src/coordinator/src/sql_classify.rs b/src/coordinator/src/sql_classify.rs index fa43a347..0c1417d4 100644 --- a/src/coordinator/src/sql_classify.rs +++ b/src/coordinator/src/sql_classify.rs @@ -15,7 +15,7 @@ use std::collections::HashMap; -use datafusion::common::{plan_err, Result}; +use datafusion::common::{Result, plan_err}; use datafusion::error::DataFusionError; use datafusion::sql::sqlparser::ast::{ ObjectType, ShowCreateObject, SqlOption, Statement as DFStatement, @@ -50,9 +50,9 @@ pub fn classify_statement(stmt: DFStatement) -> Result> { DFStatement::ShowStreamingTable => Ok(Box::new(ShowStreamingTables::new())), DFStatement::ShowCreate { obj_type, obj_name } => match obj_type { ShowCreateObject::Table => Ok(Box::new(ShowCreateTable::new(obj_name.to_string()))), - ShowCreateObject::StreamingTable => { - Ok(Box::new(ShowCreateStreamingTable::new(obj_name.to_string()))) - } + ShowCreateObject::StreamingTable => Ok(Box::new(ShowCreateStreamingTable::new( + obj_name.to_string(), + ))), _ => plan_err!( "SHOW CREATE {obj_type} is not supported; use SHOW CREATE TABLE or SHOW CREATE STREAMING TABLE " ), @@ -88,11 +88,12 @@ pub fn classify_statement(stmt: DFStatement) -> Result> { } let table_name = names[0].to_string(); Ok(Box::new(DropStreamingTableStatement::new( - table_name, - *if_exists, + table_name, *if_exists, ))) } - _ => plan_err!("Only DROP TABLE and DROP STREAMING TABLE are supported in this SQL frontend"), + _ => plan_err!( + "Only DROP TABLE and DROP STREAMING TABLE are supported in this SQL frontend" + ), } } DFStatement::Insert { .. } => plan_err!( diff --git a/src/streaming_planner/src/common/mod.rs b/src/streaming_planner/src/common/mod.rs index 6133294b..40c11c6c 100644 --- a/src/streaming_planner/src/common/mod.rs +++ b/src/streaming_planner/src/common/mod.rs @@ -33,8 +33,8 @@ pub mod topology; pub mod with_option_keys; // ── Re-exports from existing modules ── -pub use function_stream_runtime_common::streaming_protocol::{CheckpointBarrier, Watermark}; pub use arrow_ext::FsExtensionType; +pub use function_stream_runtime_common::streaming_protocol::{CheckpointBarrier, Watermark}; pub use time_utils::{from_nanos, to_micros, to_millis, to_nanos}; // ── Re-exports from new modules ── diff --git a/src/streaming_planner/src/connector/sink/runtime_config.rs b/src/streaming_planner/src/connector/sink/runtime_config.rs index a383a7ef..b68f7aa9 100644 --- a/src/streaming_planner/src/connector/sink/runtime_config.rs +++ b/src/streaming_planner/src/connector/sink/runtime_config.rs @@ -14,12 +14,12 @@ use std::collections::HashMap; use datafusion::common::{DataFusionError, Result, plan_err}; +use crate::common::connector_options::ConnectorOptions; +use crate::common::with_option_keys as opt; use function_stream_config::global_config::{ DEFAULT_OPERATOR_STATE_STORE_MEMORY_BYTES, DEFAULT_SINK_BUFFER_MEMORY_BYTES, }; use function_stream_config::streaming_job::DEFAULT_CHECKPOINT_INTERVAL_MS; -use crate::common::connector_options::ConnectorOptions; -use crate::common::with_option_keys as opt; #[derive(Debug, Clone, PartialEq, Eq, Default)] pub struct SinkRuntimeConfig { diff --git a/src/streaming_planner/src/logical_node/aggregate.rs b/src/streaming_planner/src/logical_node/aggregate.rs index d8a96472..cef337c0 100644 --- a/src/streaming_planner/src/logical_node/aggregate.rs +++ b/src/streaming_planner/src/logical_node/aggregate.rs @@ -30,7 +30,6 @@ use protocol::function_stream_graph::{ SessionWindowAggregateOperator, SlidingWindowAggregateOperator, TumblingWindowAggregateOperator, }; -use crate::multifield_partial_ord; use crate::common::constants::{extension_node, proto_operator_name}; use crate::common::{FsSchema, FsSchemaRef}; use crate::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; @@ -38,6 +37,7 @@ use crate::logical_node::{ CompiledTopologyNode, StreamingOperatorBlueprint, SystemTimestampInjectorNode, }; use crate::logical_planner::planner::{NamedNode, Planner, SplitPlanOutput}; +use crate::multifield_partial_ord; use crate::physical::{StreamingExtensionCodec, window}; use crate::types::{ QualifiedField, TIMESTAMP_FIELD, WindowBehavior, WindowType, build_df_schema, diff --git a/src/streaming_planner/src/logical_node/async_udf.rs b/src/streaming_planner/src/logical_node/async_udf.rs index 3c70fc5e..a3d1f68c 100644 --- a/src/streaming_planner/src/logical_node/async_udf.rs +++ b/src/streaming_planner/src/logical_node/async_udf.rs @@ -24,7 +24,6 @@ use datafusion_proto::physical_plan::to_proto::serialize_physical_expr; use prost::Message; use protocol::function_stream_graph::{AsyncUdfOperator, AsyncUdfOrdering}; -use crate::multifield_partial_ord; use crate::common::constants::extension_node; use crate::common::constants::sql_field; use crate::common::{FsSchema, FsSchemaRef}; @@ -35,6 +34,7 @@ use crate::logical_node::streaming_operator_blueprint::{ CompiledTopologyNode, StreamingOperatorBlueprint, }; use crate::logical_planner::planner::{NamedNode, Planner}; +use crate::multifield_partial_ord; use crate::types::{QualifiedField, build_df_schema, extract_qualified_fields}; pub const NODE_TYPE_NAME: &str = extension_node::ASYNC_FUNCTION_EXECUTION; diff --git a/src/streaming_planner/src/logical_node/debezium.rs b/src/streaming_planner/src/logical_node/debezium.rs index 9d558026..2b13f090 100644 --- a/src/streaming_planner/src/logical_node/debezium.rs +++ b/src/streaming_planner/src/logical_node/debezium.rs @@ -19,10 +19,10 @@ use datafusion::common::{ use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; use datafusion::physical_plan::DisplayAs; -use crate::multifield_partial_ord; use crate::common::constants::{cdc, extension_node}; use crate::common::{FsSchema, FsSchemaRef, UPDATING_META_FIELD}; use crate::logical_planner::planner::{NamedNode, Planner}; +use crate::multifield_partial_ord; use crate::physical::updating_meta_field; use crate::types::TIMESTAMP_FIELD; diff --git a/src/streaming_planner/src/logical_node/is_retract.rs b/src/streaming_planner/src/logical_node/is_retract.rs index 496edee6..17fad0db 100644 --- a/src/streaming_planner/src/logical_node/is_retract.rs +++ b/src/streaming_planner/src/logical_node/is_retract.rs @@ -18,9 +18,7 @@ use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; use crate::multifield_partial_ord; use crate::physical::updating_meta_field; -use crate::types::{ - QualifiedField, TIMESTAMP_FIELD, build_df_schema, extract_qualified_fields, -}; +use crate::types::{QualifiedField, TIMESTAMP_FIELD, build_df_schema, extract_qualified_fields}; #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct IsRetractExtension { diff --git a/src/streaming_planner/src/logical_node/key_calculation.rs b/src/streaming_planner/src/logical_node/key_calculation.rs index b27ed1ea..8bd03b6e 100644 --- a/src/streaming_planner/src/logical_node/key_calculation.rs +++ b/src/streaming_planner/src/logical_node/key_calculation.rs @@ -26,12 +26,12 @@ use prost::Message; use protocol::function_stream_graph::{KeyPlanOperator, ProjectionOperator}; -use crate::multifield_partial_ord; use crate::common::constants::{extension_node, sql_field}; use crate::common::{FsSchema, FsSchemaRef}; use crate::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; use crate::logical_node::{CompiledTopologyNode, StreamingOperatorBlueprint}; use crate::logical_planner::planner::{NamedNode, Planner}; +use crate::multifield_partial_ord; use crate::physical::StreamingExtensionCodec; use crate::types::{build_df_schema_with_metadata, extract_qualified_fields}; diff --git a/src/streaming_planner/src/logical_node/lookup.rs b/src/streaming_planner/src/logical_node/lookup.rs index 9075da60..ca95c4a6 100644 --- a/src/streaming_planner/src/logical_node/lookup.rs +++ b/src/streaming_planner/src/logical_node/lookup.rs @@ -23,12 +23,12 @@ use prost::Message; use protocol::function_stream_graph; use protocol::function_stream_graph::{ConnectorOp, LookupJoinCondition, LookupJoinOperator}; -use crate::multifield_partial_ord; use crate::common::constants::extension_node; use crate::common::{FsSchema, FsSchemaRef}; use crate::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; use crate::logical_node::{CompiledTopologyNode, StreamingOperatorBlueprint}; use crate::logical_planner::planner::{NamedNode, Planner}; +use crate::multifield_partial_ord; use crate::schema::LookupTable; use crate::schema::utils::add_timestamp_field_arrow; diff --git a/src/streaming_planner/src/logical_node/projection.rs b/src/streaming_planner/src/logical_node/projection.rs index c7d06e29..f56b109f 100644 --- a/src/streaming_planner/src/logical_node/projection.rs +++ b/src/streaming_planner/src/logical_node/projection.rs @@ -21,12 +21,12 @@ use prost::Message; use protocol::function_stream_graph::ProjectionOperator; -use crate::multifield_partial_ord; use crate::common::constants::{extension_node, sql_field}; use crate::common::{FsSchema, FsSchemaRef}; use crate::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; use crate::logical_node::{CompiledTopologyNode, StreamingOperatorBlueprint}; use crate::logical_planner::planner::{NamedNode, Planner}; +use crate::multifield_partial_ord; use crate::types::{QualifiedField, build_df_schema}; // ----------------------------------------------------------------------------- diff --git a/src/streaming_planner/src/logical_node/remote_table.rs b/src/streaming_planner/src/logical_node/remote_table.rs index 3b050c62..302228e7 100644 --- a/src/streaming_planner/src/logical_node/remote_table.rs +++ b/src/streaming_planner/src/logical_node/remote_table.rs @@ -21,12 +21,12 @@ use prost::Message; use protocol::function_stream_graph::ValuePlanOperator; -use crate::multifield_partial_ord; use crate::common::constants::extension_node; use crate::common::{FsSchema, FsSchemaRef}; use crate::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; use crate::logical_node::{CompiledTopologyNode, StreamingOperatorBlueprint}; use crate::logical_planner::planner::{NamedNode, Planner}; +use crate::multifield_partial_ord; use crate::physical::StreamingExtensionCodec; // ----------------------------------------------------------------------------- diff --git a/src/streaming_planner/src/logical_node/sink.rs b/src/streaming_planner/src/logical_node/sink.rs index e93ea02c..f0b66187 100644 --- a/src/streaming_planner/src/logical_node/sink.rs +++ b/src/streaming_planner/src/logical_node/sink.rs @@ -17,12 +17,12 @@ use datafusion::common::{DFSchemaRef, Result, TableReference, plan_err}; use datafusion::logical_expr::{Expr, Extension, LogicalPlan, UserDefinedLogicalNodeCore}; use prost::Message; -use crate::multifield_partial_ord; use crate::common::constants::extension_node; use crate::common::{FsSchema, FsSchemaRef, UPDATING_META_FIELD}; use crate::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; use crate::logical_node::{CompiledTopologyNode, StreamingOperatorBlueprint}; use crate::logical_planner::planner::{NamedNode, Planner}; +use crate::multifield_partial_ord; use crate::schema::CatalogEntity; use crate::schema::catalog::ExternalTable; diff --git a/src/streaming_planner/src/logical_node/table_source.rs b/src/streaming_planner/src/logical_node/table_source.rs index c0067a7d..dad8f921 100644 --- a/src/streaming_planner/src/logical_node/table_source.rs +++ b/src/streaming_planner/src/logical_node/table_source.rs @@ -17,12 +17,12 @@ use datafusion::common::{DFSchemaRef, Result, TableReference, plan_err}; use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; use prost::Message; -use crate::multifield_partial_ord; use crate::common::constants::extension_node; use crate::common::{FsSchema, FsSchemaRef}; use crate::logical_node::debezium::DebeziumSchemaCodec; use crate::logical_node::logical::{LogicalNode, OperatorName}; use crate::logical_planner::planner::{NamedNode, Planner}; +use crate::multifield_partial_ord; use crate::schema::SourceTable; use crate::schema::utils::add_timestamp_field; use crate::types::build_df_schema; diff --git a/src/streaming_planner/src/logical_node/timestamp_append.rs b/src/streaming_planner/src/logical_node/timestamp_append.rs index 14fb247d..a490243a 100644 --- a/src/streaming_planner/src/logical_node/timestamp_append.rs +++ b/src/streaming_planner/src/logical_node/timestamp_append.rs @@ -15,8 +15,8 @@ use std::fmt::Formatter; use datafusion::common::{DFSchemaRef, Result, TableReference, internal_err}; use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; -use crate::multifield_partial_ord; use crate::common::constants::extension_node; +use crate::multifield_partial_ord; use crate::schema::utils::{add_timestamp_field, has_timestamp_field}; // ----------------------------------------------------------------------------- diff --git a/src/streaming_planner/src/logical_node/updating_aggregate.rs b/src/streaming_planner/src/logical_node/updating_aggregate.rs index 0343c370..14f52ea7 100644 --- a/src/streaming_planner/src/logical_node/updating_aggregate.rs +++ b/src/streaming_planner/src/logical_node/updating_aggregate.rs @@ -29,9 +29,7 @@ use crate::common::constants::{extension_node, proto_operator_name, updating_sta use crate::common::{FsSchema, FsSchemaRef}; use crate::functions::multi_hash; use crate::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; -use crate::logical_node::{ - CompiledTopologyNode, IsRetractExtension, StreamingOperatorBlueprint, -}; +use crate::logical_node::{CompiledTopologyNode, IsRetractExtension, StreamingOperatorBlueprint}; use crate::logical_planner::planner::{NamedNode, Planner}; use crate::physical::StreamingExtensionCodec; diff --git a/src/streaming_planner/src/logical_node/watermark_node.rs b/src/streaming_planner/src/logical_node/watermark_node.rs index a8802a7b..8d51fc47 100644 --- a/src/streaming_planner/src/logical_node/watermark_node.rs +++ b/src/streaming_planner/src/logical_node/watermark_node.rs @@ -21,12 +21,12 @@ use datafusion_proto::physical_plan::to_proto::serialize_physical_expr; use prost::Message; use protocol::function_stream_graph::ExpressionWatermarkConfig; -use crate::multifield_partial_ord; use crate::common::constants::{extension_node, runtime_operator_kind}; use crate::common::{FsSchema, FsSchemaRef}; use crate::logical_node::logical::{LogicalEdge, LogicalEdgeType, LogicalNode, OperatorName}; use crate::logical_node::{CompiledTopologyNode, StreamingOperatorBlueprint}; use crate::logical_planner::planner::{NamedNode, Planner}; +use crate::multifield_partial_ord; use crate::schema::utils::add_timestamp_field; use crate::types::TIMESTAMP_FIELD; diff --git a/src/streaming_planner/src/logical_planner/streaming_planner.rs b/src/streaming_planner/src/logical_planner/streaming_planner.rs index ef7c7d38..a5f3d5f6 100644 --- a/src/streaming_planner/src/logical_planner/streaming_planner.rs +++ b/src/streaming_planner/src/logical_planner/streaming_planner.rs @@ -44,9 +44,7 @@ use datafusion_proto::physical_plan::to_proto::serialize_physical_expr; use crate::common::constants::sql_planning_default; use crate::common::{FsSchema, FsSchemaRef}; -use crate::logical_node::debezium::{ - PACK_NODE_NAME, UNROLL_NODE_NAME, UnrollDebeziumPayloadNode, -}; +use crate::logical_node::debezium::{PACK_NODE_NAME, UNROLL_NODE_NAME, UnrollDebeziumPayloadNode}; use crate::logical_node::key_calculation::KeyExtractionNode; use crate::logical_node::logical::{LogicalEdge, LogicalGraph, LogicalNode}; use crate::logical_node::{CompiledTopologyNode, StreamingOperatorBlueprint}; @@ -113,10 +111,7 @@ impl<'a> Planner<'a> { sql_planning_default::KEYED_AGGREGATE_DEFAULT_PARALLELISM } - pub fn new( - schema_provider: &'a StreamSchemaProvider, - session_state: &'a SessionState, - ) -> Self { + pub fn new(schema_provider: &'a StreamSchemaProvider, session_state: &'a SessionState) -> Self { let planner = DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new(FsExtensionPlanner {})]); Self { @@ -159,11 +154,7 @@ impl<'a> Planner<'a> { .create_physical_expr(expr, input_dfschema, self.session_state) } - pub fn serialize_as_physical_expr( - &self, - expr: &Expr, - schema: &DFSchema, - ) -> Result> { + pub fn serialize_as_physical_expr(&self, expr: &Expr, schema: &DFSchema) -> Result> { let physical = self.create_physical_expr(expr, schema)?; let proto = serialize_physical_expr(&physical, &DefaultPhysicalExtensionCodec {})?; Ok(proto.encode_to_vec()) diff --git a/src/streaming_planner/src/parse.rs b/src/streaming_planner/src/parse.rs index b3a90fbf..563214aa 100644 --- a/src/streaming_planner/src/parse.rs +++ b/src/streaming_planner/src/parse.rs @@ -27,7 +27,7 @@ //! **`INSERT` is not supported** at the coordinator layer — use `CREATE TABLE ... AS SELECT` or //! `CREATE STREAMING TABLE ... AS SELECT` instead (see coordinator classification). -use datafusion::common::{plan_err, Result}; +use datafusion::common::{Result, plan_err}; use datafusion::error::DataFusionError; use datafusion::sql::sqlparser::ast::Statement as DFStatement; use datafusion::sql::sqlparser::dialect::FunctionStreamDialect; diff --git a/src/streaming_planner/src/physical/udfs.rs b/src/streaming_planner/src/physical/udfs.rs index 733fa79e..0a11ba4f 100644 --- a/src/streaming_planner/src/physical/udfs.rs +++ b/src/streaming_planner/src/physical/udfs.rs @@ -20,8 +20,8 @@ use datafusion::logical_expr::{ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility, }; -use crate::make_udf_function; use crate::common::constants::window_function_udf; +use crate::make_udf_function; use crate::schema::utils::window_arrow_struct; // ============================================================================ diff --git a/src/streaming_planner/src/planning_runtime.rs b/src/streaming_planner/src/planning_runtime.rs index 8a03addc..cb2ab93e 100644 --- a/src/streaming_planner/src/planning_runtime.rs +++ b/src/streaming_planner/src/planning_runtime.rs @@ -14,9 +14,9 @@ use std::sync::OnceLock; -use function_stream_config::streaming_job::ResolvedStreamingJobConfig; use crate::common::constants::sql_planning_default; use crate::types::SqlConfig; +use function_stream_config::streaming_job::ResolvedStreamingJobConfig; static SQL_PLANNING: OnceLock = OnceLock::new(); diff --git a/src/streaming_planner/src/schema/catalog.rs b/src/streaming_planner/src/schema/catalog.rs index 058c461d..35966289 100644 --- a/src/streaming_planner/src/schema/catalog.rs +++ b/src/streaming_planner/src/schema/catalog.rs @@ -26,10 +26,10 @@ use super::column_descriptor::ColumnDescriptor; use super::data_encoding_format::DataEncodingFormat; use super::table::SqlSource; use super::temporal_pipeline_config::TemporalPipelineConfig; -use crate::multifield_partial_ord; use crate::common::constants::sql_field; use crate::common::{Format, FsSchema}; use crate::connector::config::ConnectorConfig; +use crate::multifield_partial_ord; use crate::types::ProcessingMode; #[derive(Debug, Clone)] diff --git a/src/streaming_runtime/src/streaming/api/operator.rs b/src/streaming_runtime/src/streaming/api/operator.rs index 43535a8a..56cefcee 100644 --- a/src/streaming_runtime/src/streaming/api/operator.rs +++ b/src/streaming_runtime/src/streaming/api/operator.rs @@ -10,10 +10,10 @@ // See the License for the specific language governing permissions and // limitations under the License. +use crate::sql::common::{CheckpointBarrier, Watermark}; use crate::streaming::api::context::TaskContext; use crate::streaming::api::source::SourceOperator; use crate::streaming::protocol::event::StreamOutput; -use crate::sql::common::{CheckpointBarrier, Watermark}; use arrow_array::RecordBatch; use async_trait::async_trait; diff --git a/src/streaming_runtime/src/streaming/api/source.rs b/src/streaming_runtime/src/streaming/api/source.rs index e0ab6a92..4bed912f 100644 --- a/src/streaming_runtime/src/streaming/api/source.rs +++ b/src/streaming_runtime/src/streaming/api/source.rs @@ -10,8 +10,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::streaming::api::context::TaskContext; use crate::sql::common::{CheckpointBarrier, Watermark}; +use crate::streaming::api::context::TaskContext; use arrow_array::RecordBatch; use async_trait::async_trait; use protocol::storage::{ diff --git a/src/streaming_runtime/src/streaming/execution/operator_chain.rs b/src/streaming_runtime/src/streaming/execution/operator_chain.rs index 2cd392fc..38584c5d 100644 --- a/src/streaming_runtime/src/streaming/execution/operator_chain.rs +++ b/src/streaming_runtime/src/streaming/execution/operator_chain.rs @@ -13,6 +13,7 @@ use anyhow::anyhow; use async_trait::async_trait; +use crate::sql::common::CheckpointBarrier; use crate::streaming::api::context::TaskContext; use crate::streaming::api::operator::{Collector, Operator}; use crate::streaming::error::RunError; @@ -20,7 +21,6 @@ use crate::streaming::protocol::{ control::{ControlCommand, StopMode}, event::{StreamEvent, StreamOutput, TrackedEvent}, }; -use crate::sql::common::CheckpointBarrier; // ============================================================================ // Core Traits diff --git a/src/streaming_runtime/src/streaming/execution/pipeline.rs b/src/streaming_runtime/src/streaming/execution/pipeline.rs index badc3601..6b6b1878 100644 --- a/src/streaming_runtime/src/streaming/execution/pipeline.rs +++ b/src/streaming_runtime/src/streaming/execution/pipeline.rs @@ -14,6 +14,7 @@ use tokio::sync::mpsc::UnboundedReceiver; use tokio_stream::{StreamExt, StreamMap}; use tracing::{Instrument, info, info_span}; +use crate::sql::common::Watermark; use crate::streaming::api::context::TaskContext; use crate::streaming::api::operator::Operator; use crate::streaming::error::RunError; @@ -27,7 +28,6 @@ use crate::streaming::protocol::{ control::ControlCommand, event::{StreamEvent, TrackedEvent}, }; -use crate::sql::common::Watermark; pub struct Pipeline { chain_head: Box, diff --git a/src/streaming_runtime/src/streaming/execution/source_driver.rs b/src/streaming_runtime/src/streaming/execution/source_driver.rs index 91058ab2..72cc160e 100644 --- a/src/streaming_runtime/src/streaming/execution/source_driver.rs +++ b/src/streaming_runtime/src/streaming/execution/source_driver.rs @@ -14,6 +14,7 @@ use tokio::sync::mpsc::UnboundedReceiver; use tokio::time::{Instant, sleep}; use tracing::{Instrument, info, info_span, warn}; +use crate::sql::common::CheckpointBarrier; use crate::streaming::api::context::TaskContext; use crate::streaming::api::source::{SourceCheckpointReport, SourceEvent, SourceOperator}; use crate::streaming::error::RunError; @@ -22,7 +23,6 @@ use crate::streaming::protocol::{ control::ControlCommand, event::{StreamEvent, TrackedEvent}, }; -use crate::sql::common::CheckpointBarrier; pub struct SourceDriver { operator: Box, diff --git a/src/streaming_runtime/src/streaming/execution/tracker/watermark_tracker.rs b/src/streaming_runtime/src/streaming/execution/tracker/watermark_tracker.rs index 87d9d89c..ec44726d 100644 --- a/src/streaming_runtime/src/streaming/execution/tracker/watermark_tracker.rs +++ b/src/streaming_runtime/src/streaming/execution/tracker/watermark_tracker.rs @@ -10,8 +10,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::streaming::protocol::event::{merge_watermarks, watermark_strictly_advances}; use crate::sql::common::Watermark; +use crate::streaming::protocol::event::{merge_watermarks, watermark_strictly_advances}; #[derive(Debug)] pub struct WatermarkTracker { diff --git a/src/streaming_runtime/src/streaming/factory/connector/delta.rs b/src/streaming_runtime/src/streaming/factory/connector/delta.rs index 324147bd..0c684e94 100644 --- a/src/streaming_runtime/src/streaming/factory/connector/delta.rs +++ b/src/streaming_runtime/src/streaming/factory/connector/delta.rs @@ -18,6 +18,8 @@ use prost::Message; use protocol::function_stream_graph::ConnectorOp; use protocol::function_stream_graph::connector_op::Config; +use crate::sql::common::constants::connection_format_value; +use crate::sql::common::with_option_keys as opt; use crate::streaming::api::operator::ConstructedOperator; use crate::streaming::factory::connector::sink_props_codec::{ apply_common_sink_fields, normalized_props, parse_sink_memory_bytes, @@ -26,8 +28,6 @@ use crate::streaming::factory::global::Registry; use crate::streaming::factory::operator_constructor::OperatorConstructor; use crate::streaming::operators::sink::delta::{DeltaFormat, DeltaSinkOperator}; use crate::streaming::operators::sink::filesystem::compression_from_str; -use crate::sql::common::constants::connection_format_value; -use crate::sql::common::with_option_keys as opt; pub struct DeltaSinkDispatcher; diff --git a/src/streaming_runtime/src/streaming/factory/connector/dispatchers.rs b/src/streaming_runtime/src/streaming/factory/connector/dispatchers.rs index e5997c38..bcd540a4 100644 --- a/src/streaming_runtime/src/streaming/factory/connector/dispatchers.rs +++ b/src/streaming_runtime/src/streaming/factory/connector/dispatchers.rs @@ -16,10 +16,10 @@ use anyhow::{Context, Result, bail}; use prost::Message; use protocol::function_stream_graph::ConnectorOp; +use crate::sql::common::constants::connector_type; use crate::streaming::api::operator::ConstructedOperator; use crate::streaming::factory::global::Registry; use crate::streaming::factory::operator_constructor::OperatorConstructor; -use crate::sql::common::constants::connector_type; use super::{ DeltaSinkDispatcher, FilesystemSinkDispatcher, IcebergSinkDispatcher, LanceDbSinkDispatcher, diff --git a/src/streaming_runtime/src/streaming/factory/connector/filesystem.rs b/src/streaming_runtime/src/streaming/factory/connector/filesystem.rs index 27ff9ec1..b680590c 100644 --- a/src/streaming_runtime/src/streaming/factory/connector/filesystem.rs +++ b/src/streaming_runtime/src/streaming/factory/connector/filesystem.rs @@ -18,6 +18,8 @@ use prost::Message; use protocol::function_stream_graph::ConnectorOp; use protocol::function_stream_graph::connector_op::Config; +use crate::sql::common::constants::connection_format_value; +use crate::sql::common::with_option_keys as opt; use crate::streaming::api::operator::ConstructedOperator; use crate::streaming::factory::connector::sink_props_codec::{ apply_common_sink_fields, normalized_props, parse_sink_memory_bytes, @@ -27,8 +29,6 @@ use crate::streaming::factory::operator_constructor::OperatorConstructor; use crate::streaming::operators::sink::filesystem::{ FilesystemFormat, FilesystemSinkOperator, compression_from_str, }; -use crate::sql::common::constants::connection_format_value; -use crate::sql::common::with_option_keys as opt; pub struct FilesystemSinkDispatcher; diff --git a/src/streaming_runtime/src/streaming/factory/connector/iceberg.rs b/src/streaming_runtime/src/streaming/factory/connector/iceberg.rs index 8c4ca56a..a6285cbf 100644 --- a/src/streaming_runtime/src/streaming/factory/connector/iceberg.rs +++ b/src/streaming_runtime/src/streaming/factory/connector/iceberg.rs @@ -18,6 +18,8 @@ use prost::Message; use protocol::function_stream_graph::ConnectorOp; use protocol::function_stream_graph::connector_op::Config; +use crate::sql::common::constants::connection_format_value; +use crate::sql::common::with_option_keys as opt; use crate::streaming::api::operator::ConstructedOperator; use crate::streaming::factory::connector::sink_props_codec::{ apply_common_sink_fields, normalized_props, parse_sink_memory_bytes, @@ -26,8 +28,6 @@ use crate::streaming::factory::global::Registry; use crate::streaming::factory::operator_constructor::OperatorConstructor; use crate::streaming::operators::sink::filesystem::compression_from_str; use crate::streaming::operators::sink::iceberg::{IcebergFormat, IcebergSinkOperator}; -use crate::sql::common::constants::connection_format_value; -use crate::sql::common::with_option_keys as opt; pub struct IcebergSinkDispatcher; diff --git a/src/streaming_runtime/src/streaming/factory/connector/kafka.rs b/src/streaming_runtime/src/streaming/factory/connector/kafka.rs index 7dcdafb9..ad77af06 100644 --- a/src/streaming_runtime/src/streaming/factory/connector/kafka.rs +++ b/src/streaming_runtime/src/streaming/factory/connector/kafka.rs @@ -24,6 +24,7 @@ use protocol::function_stream_graph::{ }; use tracing::info; +use crate::sql::common::FsSchema; use crate::streaming::api::operator::ConstructedOperator; use crate::streaming::api::source::SourceOffset; use crate::streaming::factory::global::Registry; @@ -33,10 +34,7 @@ use crate::streaming::format::{ Format as RuntimeFormat, JsonFormat as RuntimeJsonFormat, TimestampFormat as RtTimestampFormat, }; use crate::streaming::operators::sink::kafka::{ConsistencyMode, KafkaSinkOperator}; -use crate::streaming::operators::source::kafka::{ - BufferedDeserializer, KafkaSourceOperator, -}; -use crate::sql::common::FsSchema; +use crate::streaming::operators::source::kafka::{BufferedDeserializer, KafkaSourceOperator}; const DEFAULT_SOURCE_BATCH_SIZE: usize = 1024; diff --git a/src/streaming_runtime/src/streaming/factory/connector/lancedb.rs b/src/streaming_runtime/src/streaming/factory/connector/lancedb.rs index 61bfd927..3deb055d 100644 --- a/src/streaming_runtime/src/streaming/factory/connector/lancedb.rs +++ b/src/streaming_runtime/src/streaming/factory/connector/lancedb.rs @@ -18,6 +18,8 @@ use prost::Message; use protocol::function_stream_graph::ConnectorOp; use protocol::function_stream_graph::connector_op::Config; +use crate::sql::common::constants::connection_format_value; +use crate::sql::common::with_option_keys as opt; use crate::streaming::api::operator::ConstructedOperator; use crate::streaming::factory::connector::sink_props_codec::{ apply_common_sink_fields, normalized_props, @@ -25,8 +27,6 @@ use crate::streaming::factory::connector::sink_props_codec::{ use crate::streaming::factory::global::Registry; use crate::streaming::factory::operator_constructor::OperatorConstructor; use crate::streaming::operators::sink::lancedb::LanceDbSinkOperator; -use crate::sql::common::constants::connection_format_value; -use crate::sql::common::with_option_keys as opt; pub struct LanceDbSinkDispatcher; diff --git a/src/streaming_runtime/src/streaming/factory/connector/s3.rs b/src/streaming_runtime/src/streaming/factory/connector/s3.rs index 729d1396..462b25d9 100644 --- a/src/streaming_runtime/src/streaming/factory/connector/s3.rs +++ b/src/streaming_runtime/src/streaming/factory/connector/s3.rs @@ -18,17 +18,15 @@ use prost::Message; use protocol::function_stream_graph::ConnectorOp; use protocol::function_stream_graph::connector_op::Config; +use crate::sql::common::constants::connection_format_value; +use crate::sql::common::with_option_keys as opt; use crate::streaming::api::operator::ConstructedOperator; use crate::streaming::factory::connector::sink_props_codec::{ apply_common_sink_fields, normalized_props, }; use crate::streaming::factory::global::Registry; use crate::streaming::factory::operator_constructor::OperatorConstructor; -use crate::streaming::operators::sink::s3::{ - S3Format, S3SinkOperator, compression_from_str, -}; -use crate::sql::common::constants::connection_format_value; -use crate::sql::common::with_option_keys as opt; +use crate::streaming::operators::sink::s3::{S3Format, S3SinkOperator, compression_from_str}; pub struct S3SinkDispatcher; diff --git a/src/streaming_runtime/src/streaming/factory/operator_factory.rs b/src/streaming_runtime/src/streaming/factory/operator_factory.rs index 17daa0f2..10b9d3bf 100644 --- a/src/streaming_runtime/src/streaming/factory/operator_factory.rs +++ b/src/streaming_runtime/src/streaming/factory/operator_factory.rs @@ -12,14 +12,10 @@ use super::operator_constructor::OperatorConstructor; use crate::streaming::api::operator::ConstructedOperator; -use crate::streaming::factory::connector::{ - ConnectorSinkDispatcher, ConnectorSourceDispatcher, -}; +use crate::streaming::factory::connector::{ConnectorSinkDispatcher, ConnectorSourceDispatcher}; use crate::streaming::factory::global::Registry; use crate::streaming::operators::grouping::IncrementalAggregatingConstructor; -use crate::streaming::operators::joins::{ - InstantJoinConstructor, JoinWithExpirationConstructor, -}; +use crate::streaming::operators::joins::{InstantJoinConstructor, JoinWithExpirationConstructor}; use anyhow::{Result, anyhow}; use prost::Message; use protocol::function_stream_graph::ProjectionOperator as ProjectionOperatorProto; diff --git a/src/streaming_runtime/src/streaming/job/job_manager.rs b/src/streaming_runtime/src/streaming/job/job_manager.rs index d6b268c4..c4c0a826 100644 --- a/src/streaming_runtime/src/streaming/job/job_manager.rs +++ b/src/streaming_runtime/src/streaming/job/job_manager.rs @@ -31,6 +31,8 @@ use crate::config::{ DEFAULT_PIPELINE_PARALLELISM, }; use crate::memory::global_memory_pool; +use crate::sql::logical_node::logical::OperatorName; +use crate::stream_catalog::CatalogManager; use crate::streaming::api::context::TaskContext; use crate::streaming::api::operator::{ConstructedOperator, Operator}; use crate::streaming::api::source::SourceOperator; @@ -44,8 +46,6 @@ use crate::streaming::network::endpoint::{BoxedEventStream, PhysicalSender}; use crate::streaming::protocol::control::{ControlCommand, JobMasterEvent, StopMode}; use crate::streaming::protocol::event::CheckpointBarrier; use crate::streaming::state::{IoManager, IoPool, NoopMetricsCollector}; -use crate::sql::logical_node::logical::OperatorName; -use crate::stream_catalog::CatalogManager; #[derive(Debug, Clone)] pub struct StreamingJobSummary { diff --git a/src/streaming_runtime/src/streaming/operators/grouping/incremental_aggregate.rs b/src/streaming_runtime/src/streaming/operators/grouping/incremental_aggregate.rs index eaaedddf..55ec2006 100644 --- a/src/streaming_runtime/src/streaming/operators/grouping/incremental_aggregate.rs +++ b/src/streaming_runtime/src/streaming/operators/grouping/incremental_aggregate.rs @@ -41,6 +41,10 @@ use std::{collections::HashMap, mem, sync::Arc}; use tracing::{debug, info, warn}; // ========================================================================= // ========================================================================= +use crate::sql::common::{ + CheckpointBarrier, FsSchema, TIMESTAMP_FIELD, UPDATING_META_FIELD, Watermark, to_nanos, +}; +use crate::sql::physical::updating_meta_fields; use crate::streaming::StreamOutput; use crate::streaming::api::context::TaskContext; use crate::streaming::api::operator::{Collector, Operator}; @@ -48,10 +52,6 @@ use crate::streaming::factory::Registry; use crate::streaming::operators::{Key, UpdatingCache}; use crate::streaming::state::OperatorStateStore; use crate::util::decode_aggregate; -use crate::sql::common::{ - CheckpointBarrier, FsSchema, TIMESTAMP_FIELD, UPDATING_META_FIELD, Watermark, to_nanos, -}; -use crate::sql::physical::updating_meta_fields; #[derive(Debug, Copy, Clone)] struct BatchData { diff --git a/src/streaming_runtime/src/streaming/operators/joins/join_instance.rs b/src/streaming_runtime/src/streaming/operators/joins/join_instance.rs index e1a0e71b..232c2985 100644 --- a/src/streaming_runtime/src/streaming/operators/joins/join_instance.rs +++ b/src/streaming_runtime/src/streaming/operators/joins/join_instance.rs @@ -25,13 +25,13 @@ use std::sync::{Arc, RwLock}; use std::time::UNIX_EPOCH; use tracing::{info, warn}; +use crate::sql::common::{CheckpointBarrier, FsSchema, FsSchemaRef, Watermark}; +use crate::sql::physical::{StreamingDecodingContext, StreamingExtensionCodec}; use crate::streaming::StreamOutput; use crate::streaming::api::context::TaskContext; use crate::streaming::api::operator::{Collector, Operator}; use crate::streaming::factory::Registry; use crate::streaming::state::OperatorStateStore; -use crate::sql::common::{CheckpointBarrier, FsSchema, FsSchemaRef, Watermark}; -use crate::sql::physical::{StreamingDecodingContext, StreamingExtensionCodec}; use async_trait::async_trait; use protocol::function_stream_graph::JoinOperator; diff --git a/src/streaming_runtime/src/streaming/operators/joins/join_with_expiration.rs b/src/streaming_runtime/src/streaming/operators/joins/join_with_expiration.rs index ea7e8477..139e88cd 100644 --- a/src/streaming_runtime/src/streaming/operators/joins/join_with_expiration.rs +++ b/src/streaming_runtime/src/streaming/operators/joins/join_with_expiration.rs @@ -24,13 +24,13 @@ use std::sync::{Arc, RwLock}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use tracing::{info, warn}; +use crate::sql::common::{CheckpointBarrier, FsSchema, Watermark}; +use crate::sql::physical::{StreamingDecodingContext, StreamingExtensionCodec}; use crate::streaming::StreamOutput; use crate::streaming::api::context::TaskContext; use crate::streaming::api::operator::{Collector, Operator}; use crate::streaming::factory::Registry; use crate::streaming::state::OperatorStateStore; -use crate::sql::common::{CheckpointBarrier, FsSchema, Watermark}; -use crate::sql::physical::{StreamingDecodingContext, StreamingExtensionCodec}; use async_trait::async_trait; use protocol::function_stream_graph::JoinOperator; diff --git a/src/streaming_runtime/src/streaming/operators/key_by.rs b/src/streaming_runtime/src/streaming/operators/key_by.rs index 7e17a768..f1e2831d 100644 --- a/src/streaming_runtime/src/streaming/operators/key_by.rs +++ b/src/streaming_runtime/src/streaming/operators/key_by.rs @@ -19,10 +19,10 @@ use datafusion_common::hash_utils::create_hashes; use datafusion_physical_expr::expressions::Column; use std::sync::Arc; +use crate::sql::common::{CheckpointBarrier, Watermark}; use crate::streaming::StreamOutput; use crate::streaming::api::context::TaskContext; use crate::streaming::api::operator::{Collector, Operator}; -use crate::sql::common::{CheckpointBarrier, Watermark}; use protocol::function_stream_graph::KeyPlanOperator; diff --git a/src/streaming_runtime/src/streaming/operators/key_operator.rs b/src/streaming_runtime/src/streaming/operators/key_operator.rs index 7ac364bc..a8064d2e 100644 --- a/src/streaming_runtime/src/streaming/operators/key_operator.rs +++ b/src/streaming_runtime/src/streaming/operators/key_operator.rs @@ -15,11 +15,11 @@ //! [`datafusion_common::hash_utils::create_hashes`] on those columns — same mechanism as //! [`crate::streaming::operators::key_by::KeyByOperator`]. +use crate::sql::common::{CheckpointBarrier, Watermark}; use crate::streaming::StreamOutput; use crate::streaming::api::context::TaskContext; use crate::streaming::api::operator::{Collector, Operator}; use crate::streaming::operators::StatelessPhysicalExecutor; -use crate::sql::common::{CheckpointBarrier, Watermark}; use ahash::RandomState; use anyhow::{Result, anyhow}; use arrow::compute::{sort_to_indices, take}; diff --git a/src/streaming_runtime/src/streaming/operators/projection.rs b/src/streaming_runtime/src/streaming/operators/projection.rs index cff269f1..54882547 100644 --- a/src/streaming_runtime/src/streaming/operators/projection.rs +++ b/src/streaming_runtime/src/streaming/operators/projection.rs @@ -22,12 +22,12 @@ use std::sync::Arc; use protocol::function_stream_graph::ProjectionOperator as ProjectionOperatorProto; +use crate::sql::common::{CheckpointBarrier, FsSchema, FsSchemaRef, Watermark}; +use crate::sql::logical_node::logical::OperatorName; use crate::streaming::StreamOutput; use crate::streaming::api::context::TaskContext; use crate::streaming::api::operator::{Collector, Operator}; use crate::streaming::factory::global::Registry; -use crate::sql::common::{CheckpointBarrier, FsSchema, FsSchemaRef, Watermark}; -use crate::sql::logical_node::logical::OperatorName; pub struct ProjectionOperator { name: String, diff --git a/src/streaming_runtime/src/streaming/operators/sink/delta/mod.rs b/src/streaming_runtime/src/streaming/operators/sink/delta/mod.rs index 1bce0cd5..b2d16c60 100644 --- a/src/streaming_runtime/src/streaming/operators/sink/delta/mod.rs +++ b/src/streaming_runtime/src/streaming/operators/sink/delta/mod.rs @@ -27,13 +27,13 @@ use tokio::io::AsyncWriteExt; use tracing::{debug, info, warn}; use crate::memory::{MemoryBlock, try_global_memory_pool}; +use crate::sql::common::constants::factory_operator_name; +use crate::sql::common::with_option_keys as opt; +use crate::sql::common::{CheckpointBarrier, Watermark}; use crate::streaming::StreamOutput; use crate::streaming::api::context::TaskContext; use crate::streaming::api::operator::{Collector, Operator}; use crate::streaming::format::encoder::FormatEncoder; -use crate::sql::common::constants::factory_operator_name; -use crate::sql::common::with_option_keys as opt; -use crate::sql::common::{CheckpointBarrier, Watermark}; /// Flush early when buffered batches exceed this size. const DEFAULT_MAX_BUFFER_BYTES: usize = 64 * 1024 * 1024; diff --git a/src/streaming_runtime/src/streaming/operators/sink/filesystem/mod.rs b/src/streaming_runtime/src/streaming/operators/sink/filesystem/mod.rs index 250ff0c3..3461c4a4 100644 --- a/src/streaming_runtime/src/streaming/operators/sink/filesystem/mod.rs +++ b/src/streaming_runtime/src/streaming/operators/sink/filesystem/mod.rs @@ -21,12 +21,12 @@ use tokio::io::AsyncWriteExt; use tracing::{debug, info, warn}; use crate::memory::{MemoryBlock, try_global_memory_pool}; +use crate::sql::common::constants::factory_operator_name; +use crate::sql::common::{CheckpointBarrier, Watermark}; use crate::streaming::StreamOutput; use crate::streaming::api::context::TaskContext; use crate::streaming::api::operator::{Collector, Operator}; use crate::streaming::format::encoder::FormatEncoder; -use crate::sql::common::constants::factory_operator_name; -use crate::sql::common::{CheckpointBarrier, Watermark}; const DEFAULT_MAX_BUFFER_BYTES: usize = 64 * 1024 * 1024; diff --git a/src/streaming_runtime/src/streaming/operators/sink/iceberg/mod.rs b/src/streaming_runtime/src/streaming/operators/sink/iceberg/mod.rs index 801e34b0..51715aa2 100644 --- a/src/streaming_runtime/src/streaming/operators/sink/iceberg/mod.rs +++ b/src/streaming_runtime/src/streaming/operators/sink/iceberg/mod.rs @@ -27,13 +27,13 @@ use tokio::io::AsyncWriteExt; use tracing::{debug, info, warn}; use crate::memory::{MemoryBlock, try_global_memory_pool}; +use crate::sql::common::constants::factory_operator_name; +use crate::sql::common::with_option_keys as opt; +use crate::sql::common::{CheckpointBarrier, Watermark}; use crate::streaming::StreamOutput; use crate::streaming::api::context::TaskContext; use crate::streaming::api::operator::{Collector, Operator}; use crate::streaming::format::encoder::FormatEncoder; -use crate::sql::common::constants::factory_operator_name; -use crate::sql::common::with_option_keys as opt; -use crate::sql::common::{CheckpointBarrier, Watermark}; const DEFAULT_MAX_BUFFER_BYTES: usize = 64 * 1024 * 1024; diff --git a/src/streaming_runtime/src/streaming/operators/sink/kafka/mod.rs b/src/streaming_runtime/src/streaming/operators/sink/kafka/mod.rs index e87326e8..1266ebdd 100644 --- a/src/streaming_runtime/src/streaming/operators/sink/kafka/mod.rs +++ b/src/streaming_runtime/src/streaming/operators/sink/kafka/mod.rs @@ -36,12 +36,12 @@ use std::time::Duration; use tokio::time::sleep; use tracing::{info, warn}; +use crate::sql::common::constants::factory_operator_name; +use crate::sql::common::{CheckpointBarrier, FsSchema, Watermark}; use crate::streaming::StreamOutput; use crate::streaming::api::context::TaskContext; use crate::streaming::api::operator::{Collector, Operator}; use crate::streaming::format::DataSerializer; -use crate::sql::common::constants::factory_operator_name; -use crate::sql::common::{CheckpointBarrier, FsSchema, Watermark}; // ============================================================================ // ============================================================================ diff --git a/src/streaming_runtime/src/streaming/operators/sink/lancedb/mod.rs b/src/streaming_runtime/src/streaming/operators/sink/lancedb/mod.rs index fa5b6760..e69a59ee 100644 --- a/src/streaming_runtime/src/streaming/operators/sink/lancedb/mod.rs +++ b/src/streaming_runtime/src/streaming/operators/sink/lancedb/mod.rs @@ -23,11 +23,11 @@ use lance::dataset::{WriteMode, WriteParams}; use lance::io::{ObjectStoreParams, StorageOptionsAccessor}; use tracing::{info, warn}; +use crate::sql::common::constants::factory_operator_name; +use crate::sql::common::{CheckpointBarrier, Watermark}; use crate::streaming::StreamOutput; use crate::streaming::api::context::TaskContext; use crate::streaming::api::operator::{Collector, Operator}; -use crate::sql::common::constants::factory_operator_name; -use crate::sql::common::{CheckpointBarrier, Watermark}; pub struct LanceDbSinkOperator { table_name: String, diff --git a/src/streaming_runtime/src/streaming/operators/sink/s3/mod.rs b/src/streaming_runtime/src/streaming/operators/sink/s3/mod.rs index 77dcca97..f374b95c 100644 --- a/src/streaming_runtime/src/streaming/operators/sink/s3/mod.rs +++ b/src/streaming_runtime/src/streaming/operators/sink/s3/mod.rs @@ -26,12 +26,12 @@ use parquet::basic::Compression; use parquet::file::properties::WriterProperties; use tracing::{info, warn}; -use crate::streaming::StreamOutput; -use crate::streaming::api::context::TaskContext; -use crate::streaming::api::operator::{Collector, Operator}; use crate::sql::common::constants::factory_operator_name; use crate::sql::common::with_option_keys as opt; use crate::sql::common::{CheckpointBarrier, Watermark}; +use crate::streaming::StreamOutput; +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::operator::{Collector, Operator}; #[derive(Debug, Clone, Copy)] pub enum S3Format { diff --git a/src/streaming_runtime/src/streaming/operators/source/kafka/mod.rs b/src/streaming_runtime/src/streaming/operators/source/kafka/mod.rs index 3e7bffc9..7f5e09d4 100644 --- a/src/streaming_runtime/src/streaming/operators/source/kafka/mod.rs +++ b/src/streaming_runtime/src/streaming/operators/source/kafka/mod.rs @@ -31,13 +31,13 @@ use std::num::NonZeroU32; use std::time::{Duration, Instant}; use tracing::{debug, error, info, warn}; +use crate::sql::common::fs_schema::FieldValueType; +use crate::sql::common::{CheckpointBarrier, MetadataField}; use crate::streaming::api::context::TaskContext; use crate::streaming::api::source::{ SourceCheckpointReport, SourceEvent, SourceOffset, SourceOperator, }; use crate::streaming::format::{BadDataPolicy, DataDeserializer, Format}; -use crate::sql::common::fs_schema::FieldValueType; -use crate::sql::common::{CheckpointBarrier, MetadataField}; pub trait BatchDeserializer: Send + 'static { fn deserialize_slice( diff --git a/src/streaming_runtime/src/streaming/operators/stateless_physical_executor.rs b/src/streaming_runtime/src/streaming/operators/stateless_physical_executor.rs index 69f9dc45..560ae8c5 100644 --- a/src/streaming_runtime/src/streaming/operators/stateless_physical_executor.rs +++ b/src/streaming_runtime/src/streaming/operators/stateless_physical_executor.rs @@ -24,8 +24,8 @@ use datafusion_proto::protobuf::PhysicalPlanNode; use futures::StreamExt; use prost::Message; -use crate::streaming::factory::Registry; use crate::sql::physical::{StreamingDecodingContext, StreamingExtensionCodec}; +use crate::streaming::factory::Registry; pub struct StatelessPhysicalExecutor { batch: Arc>>, diff --git a/src/streaming_runtime/src/streaming/operators/value_execution.rs b/src/streaming_runtime/src/streaming/operators/value_execution.rs index 69059317..d971668f 100644 --- a/src/streaming_runtime/src/streaming/operators/value_execution.rs +++ b/src/streaming_runtime/src/streaming/operators/value_execution.rs @@ -15,11 +15,11 @@ use arrow_array::RecordBatch; use async_trait::async_trait; use futures::StreamExt; +use crate::sql::common::{CheckpointBarrier, Watermark}; use crate::streaming::StreamOutput; use crate::streaming::api::context::TaskContext; use crate::streaming::api::operator::{Collector, Operator}; use crate::streaming::operators::StatelessPhysicalExecutor; -use crate::sql::common::{CheckpointBarrier, Watermark}; pub struct ValueExecutionOperator { name: String, diff --git a/src/streaming_runtime/src/streaming/operators/watermark/watermark_generator.rs b/src/streaming_runtime/src/streaming/operators/watermark/watermark_generator.rs index 2a38f8f0..ea70b41b 100644 --- a/src/streaming_runtime/src/streaming/operators/watermark/watermark_generator.rs +++ b/src/streaming_runtime/src/streaming/operators/watermark/watermark_generator.rs @@ -25,11 +25,11 @@ use std::sync::Arc; use std::time::{Duration, SystemTime}; use tracing::debug; +use crate::sql::common::{CheckpointBarrier, FsSchema, Watermark, from_nanos, to_millis}; use crate::streaming::StreamOutput; use crate::streaming::api::context::TaskContext; use crate::streaming::api::operator::{Collector, Operator}; use crate::streaming::factory::Registry; -use crate::sql::common::{CheckpointBarrier, FsSchema, Watermark, from_nanos, to_millis}; use async_trait::async_trait; use protocol::function_stream_graph::ExpressionWatermarkConfig; diff --git a/src/streaming_runtime/src/streaming/operators/windows/session_aggregating_window.rs b/src/streaming_runtime/src/streaming/operators/windows/session_aggregating_window.rs index c48bd7a9..b42fcc13 100644 --- a/src/streaming_runtime/src/streaming/operators/windows/session_aggregating_window.rs +++ b/src/streaming_runtime/src/streaming/operators/windows/session_aggregating_window.rs @@ -36,17 +36,17 @@ use std::time::{Duration, SystemTime}; use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel}; use tracing::info; -use crate::streaming::StreamOutput; -use crate::streaming::api::context::TaskContext; -use crate::streaming::api::operator::{Collector, Operator}; -use crate::streaming::factory::Registry; -use crate::streaming::state::OperatorStateStore; use crate::sql::common::converter::Converter; use crate::sql::common::{ CheckpointBarrier, FsSchema, FsSchemaRef, Watermark, from_nanos, to_nanos, }; use crate::sql::physical::{StreamingDecodingContext, StreamingExtensionCodec}; use crate::sql::schema::utils::window_arrow_struct; +use crate::streaming::StreamOutput; +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::operator::{Collector, Operator}; +use crate::streaming::factory::Registry; +use crate::streaming::state::OperatorStateStore; use async_trait::async_trait; use protocol::function_stream_graph::SessionWindowAggregateOperator; // ============================================================================ diff --git a/src/streaming_runtime/src/streaming/operators/windows/sliding_aggregating_window.rs b/src/streaming_runtime/src/streaming/operators/windows/sliding_aggregating_window.rs index 267f289f..b9c57835 100644 --- a/src/streaming_runtime/src/streaming/operators/windows/sliding_aggregating_window.rs +++ b/src/streaming_runtime/src/streaming/operators/windows/sliding_aggregating_window.rs @@ -33,13 +33,13 @@ use std::time::{Duration, SystemTime}; use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel}; use tracing::info; +use crate::sql::common::{CheckpointBarrier, FsSchema, Watermark, from_nanos, to_nanos}; +use crate::sql::physical::{StreamingDecodingContext, StreamingExtensionCodec}; use crate::streaming::StreamOutput; use crate::streaming::api::context::TaskContext; use crate::streaming::api::operator::{Collector, Operator}; use crate::streaming::factory::Registry; use crate::streaming::state::OperatorStateStore; -use crate::sql::common::{CheckpointBarrier, FsSchema, Watermark, from_nanos, to_nanos}; -use crate::sql::physical::{StreamingDecodingContext, StreamingExtensionCodec}; use async_trait::async_trait; use protocol::function_stream_graph::SlidingWindowAggregateOperator; // ============================================================================ diff --git a/src/streaming_runtime/src/streaming/operators/windows/tumbling_aggregating_window.rs b/src/streaming_runtime/src/streaming/operators/windows/tumbling_aggregating_window.rs index 9588d800..7d45a684 100644 --- a/src/streaming_runtime/src/streaming/operators/windows/tumbling_aggregating_window.rs +++ b/src/streaming_runtime/src/streaming/operators/windows/tumbling_aggregating_window.rs @@ -34,15 +34,15 @@ use std::time::{Duration, SystemTime}; use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel}; use tracing::{info, warn}; +use crate::sql::common::time_utils::print_time; +use crate::sql::common::{CheckpointBarrier, FsSchema, Watermark, from_nanos, to_nanos}; +use crate::sql::physical::{StreamingDecodingContext, StreamingExtensionCodec}; +use crate::sql::schema::utils::add_timestamp_field_arrow; use crate::streaming::StreamOutput; use crate::streaming::api::context::TaskContext; use crate::streaming::api::operator::{Collector, Operator}; use crate::streaming::factory::Registry; use crate::streaming::state::OperatorStateStore; -use crate::sql::common::time_utils::print_time; -use crate::sql::common::{CheckpointBarrier, FsSchema, Watermark, from_nanos, to_nanos}; -use crate::sql::physical::{StreamingDecodingContext, StreamingExtensionCodec}; -use crate::sql::schema::utils::add_timestamp_field_arrow; use async_trait::async_trait; use protocol::function_stream_graph::TumblingWindowAggregateOperator; diff --git a/src/streaming_runtime/src/streaming/operators/windows/window_function.rs b/src/streaming_runtime/src/streaming/operators/windows/window_function.rs index e9c83b97..83f14551 100644 --- a/src/streaming_runtime/src/streaming/operators/windows/window_function.rs +++ b/src/streaming_runtime/src/streaming/operators/windows/window_function.rs @@ -26,16 +26,16 @@ use std::time::SystemTime; use tokio::sync::mpsc::{UnboundedReceiver, unbounded_channel}; use tracing::{info, warn}; -use crate::streaming::StreamOutput; -use crate::streaming::api::context::TaskContext; -use crate::streaming::api::operator::{Collector, Operator}; -use crate::streaming::factory::Registry; -use crate::streaming::state::OperatorStateStore; use crate::sql::common::time_utils::print_time; use crate::sql::common::{ CheckpointBarrier, FsSchema, FsSchemaRef, Watermark, from_nanos, to_nanos, }; use crate::sql::physical::{StreamingDecodingContext, StreamingExtensionCodec}; +use crate::streaming::StreamOutput; +use crate::streaming::api::context::TaskContext; +use crate::streaming::api::operator::{Collector, Operator}; +use crate::streaming::factory::Registry; +use crate::streaming::state::OperatorStateStore; use async_trait::async_trait; // ============================================================================ diff --git a/src/wasm_runtime/src/state_backend/factory.rs b/src/wasm_runtime/src/state_backend/factory.rs index 08cb06ca..2bc0b55c 100644 --- a/src/wasm_runtime/src/state_backend/factory.rs +++ b/src/wasm_runtime/src/state_backend/factory.rs @@ -63,9 +63,8 @@ pub fn get_factory_for_task>( .join(format!("{}-{}", task_name, created_at)); let config = rocksdb_config.unwrap_or_default(); - let factory = crate::state_backend::rocksdb::RocksDBStateStoreFactory::new( - db_path, config, - )?; + let factory = + crate::state_backend::rocksdb::RocksDBStateStoreFactory::new(db_path, config)?; Ok(Arc::new(factory)) } diff --git a/src/wasm_runtime/src/state_backend/memory/store.rs b/src/wasm_runtime/src/state_backend/memory/store.rs index e95bf51f..5d07ad99 100644 --- a/src/wasm_runtime/src/state_backend/memory/store.rs +++ b/src/wasm_runtime/src/state_backend/memory/store.rs @@ -94,9 +94,8 @@ impl StateStore for MemoryStateStore { user_key: Vec, value: Vec, ) -> Result<(), BackendError> { - let key_bytes = crate::state_backend::key_builder::build_key( - &key_group, &key, &namespace, &user_key, - ); + let key_bytes = + crate::state_backend::key_builder::build_key(&key_group, &key, &namespace, &user_key); let existing = self.get_state(key_bytes.clone())?; diff --git a/src/wasm_runtime/src/state_backend/server.rs b/src/wasm_runtime/src/state_backend/server.rs index 85c3e8c3..1fb24c38 100644 --- a/src/wasm_runtime/src/state_backend/server.rs +++ b/src/wasm_runtime/src/state_backend/server.rs @@ -13,9 +13,7 @@ use crate::config::storage::{StateStorageConfig, StateStorageType}; use crate::config::{get_state_dir, get_state_dir_for_base}; use crate::state_backend::error::BackendError; -use crate::state_backend::factory::{ - FactoryType, StateStoreFactory, get_factory_for_task, -}; +use crate::state_backend::factory::{FactoryType, StateStoreFactory, get_factory_for_task}; use crate::state_backend::rocksdb::RocksDBConfig; use std::fs; use std::path::PathBuf; diff --git a/src/wasm_runtime/src/state_backend/store.rs b/src/wasm_runtime/src/state_backend/store.rs index 8e61d08a..cb897f6b 100644 --- a/src/wasm_runtime/src/state_backend/store.rs +++ b/src/wasm_runtime/src/state_backend/store.rs @@ -107,9 +107,8 @@ pub trait StateStore: Send + Sync { user_key: Vec, value: Vec, ) -> Result<(), BackendError> { - let key_bytes = crate::state_backend::key_builder::build_key( - &key_group, &key, &namespace, &user_key, - ); + let key_bytes = + crate::state_backend::key_builder::build_key(&key_group, &key, &namespace, &user_key); self.put_state(key_bytes, value) } @@ -132,9 +131,8 @@ pub trait StateStore: Send + Sync { namespace: Vec, user_key: Vec, ) -> Result>, BackendError> { - let key_bytes = crate::state_backend::key_builder::build_key( - &key_group, &key, &namespace, &user_key, - ); + let key_bytes = + crate::state_backend::key_builder::build_key(&key_group, &key, &namespace, &user_key); self.get_state(key_bytes) } @@ -156,9 +154,8 @@ pub trait StateStore: Send + Sync { namespace: Vec, user_key: Vec, ) -> Result<(), BackendError> { - let key_bytes = crate::state_backend::key_builder::build_key( - &key_group, &key, &namespace, &user_key, - ); + let key_bytes = + crate::state_backend::key_builder::build_key(&key_group, &key, &namespace, &user_key); self.delete_state(key_bytes) } @@ -199,12 +196,8 @@ pub trait StateStore: Send + Sync { key: Vec, namespace: Vec, ) -> Result { - let prefix_bytes = crate::state_backend::key_builder::build_key( - &key_group, - &key, - &namespace, - &[], - ); + let prefix_bytes = + crate::state_backend::key_builder::build_key(&key_group, &key, &namespace, &[]); self.delete_prefix_bytes(prefix_bytes) } @@ -279,12 +272,8 @@ pub trait StateStore: Send + Sync { key: Vec, namespace: Vec, ) -> Result, BackendError> { - let prefix = crate::state_backend::key_builder::build_key( - &key_group, - &key, - &namespace, - &[], - ); + let prefix = + crate::state_backend::key_builder::build_key(&key_group, &key, &namespace, &[]); self.scan(prefix) } } diff --git a/src/wasm_runtime/src/wasm/output/output_provider.rs b/src/wasm_runtime/src/wasm/output/output_provider.rs index 52d13544..33b470d1 100644 --- a/src/wasm_runtime/src/wasm/output/output_provider.rs +++ b/src/wasm_runtime/src/wasm/output/output_provider.rs @@ -60,9 +60,7 @@ impl OutputProvider { runtime: _, } => { use crate::output::output_runner::OutputRunner; - use crate::output::protocol::kafka::{ - KafkaOutputProtocol, KafkaProducerConfig, - }; + use crate::output::protocol::kafka::{KafkaOutputProtocol, KafkaProducerConfig}; let servers: Vec = bootstrap_servers .split(',') diff --git a/src/wasm_runtime/src/wasm/processor/wasm/wasm_host.rs b/src/wasm_runtime/src/wasm/processor/wasm/wasm_host.rs index 9990e260..348f3063 100644 --- a/src/wasm_runtime/src/wasm/processor/wasm/wasm_host.rs +++ b/src/wasm_runtime/src/wasm/processor/wasm/wasm_host.rs @@ -12,8 +12,8 @@ use crate::output::Output; use crate::processor::wasm::wasm_cache; -use crate::wasm::buffer_and_event::BufferOrEvent; use crate::state_backend::{StateStore, StateStoreFactory}; +use crate::wasm::buffer_and_event::BufferOrEvent; use std::sync::{Arc, OnceLock}; use wasmtime::component::{Component, HasData, Linker, Resource, bindgen}; use wasmtime::{Config, Engine, Store}; diff --git a/src/wasm_runtime/src/wasm/processor/wasm/wasm_task.rs b/src/wasm_runtime/src/wasm/processor/wasm/wasm_task.rs index 45659dd9..ea0f5dc1 100644 --- a/src/wasm_runtime/src/wasm/processor/wasm/wasm_task.rs +++ b/src/wasm_runtime/src/wasm/processor/wasm/wasm_task.rs @@ -17,10 +17,10 @@ use crate::common::{ComponentState, TaskCompletionFlag}; use crate::input::Input; use crate::output::Output; use crate::processor::function_error::FunctionErrorReport; +use crate::task::FunctionInfo; use crate::wasm::buffer_and_event::BufferOrEvent; use crate::wasm::task::ProcessorRuntimeConfig; use crate::wasm::task::{ControlMailBox, TaskControlSignal, TaskLifecycle}; -use crate::task::FunctionInfo; use crossbeam_channel::{Receiver, after, select, unbounded}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::mpsc; diff --git a/src/wasm_runtime/src/wasm/task/builder/python/mod.rs b/src/wasm_runtime/src/wasm/task/builder/python/mod.rs index cb48eec2..8bc9aa2d 100644 --- a/src/wasm_runtime/src/wasm/task/builder/python/mod.rs +++ b/src/wasm_runtime/src/wasm/task/builder/python/mod.rs @@ -33,8 +33,7 @@ impl PythonBuilder { yaml_value: &Value, modules: &[(String, Vec)], create_time: u64, - ) -> Result, Box> - { + ) -> Result, Box> { let config_type = yaml_value .get(TYPE) .and_then(|v| v.as_str()) diff --git a/src/wasm_runtime/src/wasm/task/lifecycle.rs b/src/wasm_runtime/src/wasm/task/lifecycle.rs index 827cddb5..cbe7955d 100644 --- a/src/wasm_runtime/src/wasm/task/lifecycle.rs +++ b/src/wasm_runtime/src/wasm/task/lifecycle.rs @@ -15,9 +15,9 @@ // Defines the complete lifecycle management interface for Task, including initialization, start, stop, checkpoint and close use crate::common::ComponentState; +use crate::task::FunctionInfo; use crate::wasm::task::control_mailbox::ControlMailBox; use crate::wasm::taskexecutor::InitContext; -use crate::task::FunctionInfo; use std::sync::Arc; /// Task lifecycle management interface diff --git a/src/wasm_runtime/src/wasm/taskexecutor/init_context.rs b/src/wasm_runtime/src/wasm/taskexecutor/init_context.rs index 982c00e7..edd8837d 100644 --- a/src/wasm_runtime/src/wasm/taskexecutor/init_context.rs +++ b/src/wasm_runtime/src/wasm/taskexecutor/init_context.rs @@ -15,9 +15,9 @@ // Provides various resources needed for task initialization, including state storage, task storage, thread pool, etc. use crate::processor::wasm::thread_pool::{TaskThreadPool, ThreadGroup}; -use crate::wasm::task::ControlMailBox; use crate::state_backend::StateStorageServer; use crate::task::TaskStorage; +use crate::wasm::task::ControlMailBox; use std::sync::{Arc, Mutex}; #[derive(Clone)] diff --git a/src/wasm_runtime/src/wasm/taskexecutor/task_manager.rs b/src/wasm_runtime/src/wasm/taskexecutor/task_manager.rs index 0c45931c..d88c3df4 100644 --- a/src/wasm_runtime/src/wasm/taskexecutor/task_manager.rs +++ b/src/wasm_runtime/src/wasm/taskexecutor/task_manager.rs @@ -10,15 +10,13 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::config::GlobalConfig; use crate::common::ComponentState; +use crate::config::GlobalConfig; use crate::processor::wasm::thread_pool::{GlobalTaskThreadPool, TaskThreadPool}; +use crate::state_backend::StateStorageServer; +use crate::task::{FunctionInfo, StoredTaskInfo, TaskModuleBytes, TaskStorage, TaskStorageFactory}; use crate::wasm::task::{TaskBuilder, TaskLifecycle}; use crate::wasm::taskexecutor::init_context::InitContext; -use crate::state_backend::StateStorageServer; -use crate::task::{ - FunctionInfo, StoredTaskInfo, TaskModuleBytes, TaskStorage, TaskStorageFactory, -}; use anyhow::{Context, Result, anyhow}; use parking_lot::RwLock;