diff --git a/Cargo.lock b/Cargo.lock index d3c2156c043..d00b4b1eac3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1906,7 +1906,7 @@ dependencies = [ "bitflags 2.10.0", "cexpr", "clang-sys", - "itertools 0.11.0", + "itertools 0.13.0", "proc-macro2", "quote", "regex", @@ -1924,7 +1924,7 @@ dependencies = [ "bitflags 2.10.0", "cexpr", "clang-sys", - "itertools 0.11.0", + "itertools 0.13.0", "proc-macro2", "quote", "regex", @@ -1942,7 +1942,7 @@ dependencies = [ "bitflags 2.10.0", "cexpr", "clang-sys", - "itertools 0.11.0", + "itertools 0.13.0", "log", "prettyplease", "proc-macro2", @@ -3678,6 +3678,18 @@ dependencies = [ "datafusion-physical-expr-common", ] +[[package]] +name = "datafusion-functions-json" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13ff70cb2c1960f03ba647aa2813fb1efba4c33bc221d973dd4a462a6376359a" +dependencies = [ + "datafusion", + "jiter", + "log", + "paste", +] + [[package]] name = "datafusion-functions-nested" version = "53.1.0" @@ -3976,7 +3988,7 @@ dependencies = [ [[package]] name = "dbsp" -version = "0.322.0" +version = "0.323.0" dependencies = [ "anyhow", "arc-swap", @@ -4066,7 +4078,7 @@ dependencies = [ [[package]] name = "dbsp_adapters" -version = "0.322.0" +version = "0.323.0" dependencies = [ "actix", "actix-codec", @@ -4210,7 +4222,7 @@ dependencies = [ [[package]] name = "dbsp_nexmark" -version = "0.322.0" +version = "0.323.0" dependencies = [ "anyhow", "ascii_table", @@ -5194,7 +5206,7 @@ dependencies = [ [[package]] name = "fda" -version = "0.322.0" +version = "0.323.0" dependencies = [ "anyhow", "arrow", @@ -5246,7 +5258,7 @@ dependencies = [ [[package]] name = "feldera-adapterlib" -version = "0.322.0" +version = "0.323.0" dependencies = [ "actix-web", "anyhow", @@ -5256,6 +5268,7 @@ dependencies = [ "bytemuck", "chrono", "datafusion", + "datafusion-functions-json", "dbsp", "dyn-clone", "erased-serde 0.3.31", @@ -5280,7 +5293,7 @@ dependencies = [ [[package]] name = "feldera-buffer-cache" -version = "0.322.0" +version = "0.323.0" dependencies = [ "crossbeam-utils", "enum-map", @@ -5308,7 +5321,7 @@ dependencies = [ [[package]] name = "feldera-datagen" -version = "0.322.0" +version = "0.323.0" dependencies = [ "anyhow", "async-channel 2.5.0", @@ -5334,7 +5347,7 @@ dependencies = [ [[package]] name = "feldera-fxp" -version = "0.322.0" +version = "0.323.0" dependencies = [ "bytecheck", "dbsp", @@ -5354,7 +5367,7 @@ dependencies = [ [[package]] name = "feldera-iceberg" -version = "0.322.0" +version = "0.323.0" dependencies = [ "anyhow", "chrono", @@ -5377,7 +5390,7 @@ dependencies = [ [[package]] name = "feldera-ir" -version = "0.322.0" +version = "0.323.0" dependencies = [ "proptest", "proptest-derive", @@ -5389,7 +5402,7 @@ dependencies = [ [[package]] name = "feldera-macros" -version = "0.322.0" +version = "0.323.0" dependencies = [ "prettyplease", "proc-macro2", @@ -5399,7 +5412,7 @@ dependencies = [ [[package]] name = "feldera-observability" -version = "0.322.0" +version = "0.323.0" dependencies = [ "actix-http", "awc", @@ -5414,7 +5427,7 @@ dependencies = [ [[package]] name = "feldera-rest-api" -version = "0.322.0" +version = "0.323.0" dependencies = [ "chrono", "feldera-observability", @@ -5433,7 +5446,7 @@ dependencies = [ [[package]] name = "feldera-samply" -version = "0.322.0" +version = "0.323.0" dependencies = [ "crossbeam", "feldera-size-of", @@ -5469,7 +5482,7 @@ dependencies = [ [[package]] name = "feldera-sqllib" -version = "0.322.0" +version = "0.323.0" dependencies = [ "arcstr", "base58", @@ -5513,7 +5526,7 @@ dependencies = [ [[package]] name = "feldera-storage" -version = "0.322.0" +version = "0.323.0" dependencies = [ "anyhow", "crossbeam", @@ -5536,7 +5549,7 @@ dependencies = [ [[package]] name = "feldera-types" -version = "0.322.0" +version = "0.323.0" dependencies = [ "actix-web", "anyhow", @@ -6567,7 +6580,7 @@ dependencies = [ "js-sys", "log", "wasm-bindgen", - "windows-core 0.57.0", + "windows-core 0.61.2", ] [[package]] @@ -7084,6 +7097,21 @@ dependencies = [ "jiff-tzdb", ] +[[package]] +name = "jiter" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "020ba671987d7444d251d3ee5340be1bf4606cd6c0b53e6f4066b5a1ee376b22" +dependencies = [ + "ahash 0.8.12", + "bitvec", + "lexical-parse-float", + "num-bigint", + "num-traits", + "pyo3", + "smallvec", +] + [[package]] name = "jni" version = "0.22.4" @@ -7512,9 +7540,8 @@ dependencies = [ [[package]] name = "marrow" -version = "0.2.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f5240d6977234968ff9ad254bfa73aa397fb51e41dcb22b1eb85835e9295485b" +version = "0.2.7" +source = "git+https://github.com/ryzhyk/serde_arrow.git?rev=9555b34fb85d4323e4381ad5a09872c1696a0d34#9555b34fb85d4323e4381ad5a09872c1696a0d34" dependencies = [ "arrow-array", "arrow-buffer", @@ -8813,7 +8840,7 @@ checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" [[package]] name = "pipeline-manager" -version = "0.322.0" +version = "0.323.0" dependencies = [ "actix-cors", "actix-files", @@ -9590,7 +9617,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.11.0", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.117", @@ -9603,7 +9630,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b" dependencies = [ "anyhow", - "itertools 0.11.0", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.117", @@ -9667,6 +9694,66 @@ dependencies = [ "syn 1.0.109", ] +[[package]] +name = "pyo3" +version = "0.28.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91fd8e38a3b50ed1167fb981cd6fd60147e091784c427b8f7183a7ee32c31c12" +dependencies = [ + "libc", + "num-bigint", + "num-traits", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", +] + +[[package]] +name = "pyo3-build-config" +version = "0.28.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e368e7ddfdeb98c9bca7f8383be1648fd84ab466bf2bc015e94008db6d35611e" +dependencies = [ + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.28.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f29e10af80b1f7ccaf7f69eace800a03ecd13e883acfacc1e5d0988605f651e" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.28.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df6e520eff47c45997d2fc7dd8214b25dd1310918bbb2642156ef66a67f29813" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.28.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4cdc218d835738f81c2338f822078af45b4afdf8b2e33cbb5916f108b813acb" +dependencies = [ + "heck 0.5.0", + "proc-macro2", + "pyo3-build-config", + "quote", + "syn 2.0.117", +] + [[package]] name = "quad-rand" version = "0.2.3" @@ -10023,7 +10110,7 @@ dependencies = [ [[package]] name = "readers" -version = "0.322.0" +version = "0.323.0" dependencies = [ "async-std", "csv", @@ -11304,9 +11391,8 @@ dependencies = [ [[package]] name = "serde_arrow" -version = "0.14.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "26e4ac1bef72720318e2c67bd19b972d17084840f3188a585021828122c43c2c" +version = "0.14.2" +source = "git+https://github.com/ryzhyk/serde_arrow.git?rev=9555b34fb85d4323e4381ad5a09872c1696a0d34#9555b34fb85d4323e4381ad5a09872c1696a0d34" dependencies = [ "arrow-array", "arrow-schema", @@ -11722,7 +11808,7 @@ checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" [[package]] name = "sltsqlvalue" -version = "0.322.0" +version = "0.323.0" dependencies = [ "dbsp", "feldera-sqllib", @@ -12193,7 +12279,7 @@ checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" [[package]] name = "storage-test-compat" -version = "0.322.0" +version = "0.323.0" dependencies = [ "dbsp", "derive_more 1.0.0", @@ -12451,6 +12537,12 @@ dependencies = [ "xattr", ] +[[package]] +name = "target-lexicon" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" + [[package]] name = "target-triple" version = "0.1.4" diff --git a/Cargo.toml b/Cargo.toml index 04dfef6851d..30b36a51762 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [workspace.package] authors = ["Feldera Team "] -version = "0.322.0" +version = "0.323.0" license = "MIT OR Apache-2.0" homepage = "https://github.com/feldera/feldera" repository = "https://github.com/feldera/feldera" @@ -108,7 +108,9 @@ csv = "1.2.2" csv-core = "0.1.10" dashmap = "6.1.0" datafusion = "53.1" -dbsp = { path = "crates/dbsp", version = "0.322.0" } +dbsp = { path = "crates/dbsp", version = "0.323.0" } +# JSON functions applied to strings. +datafusion-functions-json = "0.53.1" dbsp_nexmark = { path = "crates/nexmark" } deadpool-postgres = "0.14.1" # Feldera fork of delta-rs: upstream tag `rust-v0.32.3` + 3 patches, on branch @@ -136,20 +138,20 @@ erased-serde = "0.3.31" fake = "2.10" fastbloom = "0.14.0" fdlimit = "0.3.0" -feldera-buffer-cache = { version = "0.322.0", path = "crates/buffer-cache" } +feldera-buffer-cache = { version = "0.323.0", path = "crates/buffer-cache" } feldera-cloud1-client = "0.1.2" feldera-datagen = { path = "crates/datagen" } -feldera-fxp = { version = "0.322.0", path = "crates/fxp", features = ["dbsp"] } +feldera-fxp = { version = "0.323.0", path = "crates/fxp", features = ["dbsp"] } feldera-iceberg = { path = "crates/iceberg" } -feldera-observability = { version = "0.322.0", path = "crates/feldera-observability" } -feldera-macros = { version = "0.322.0", path = "crates/feldera-macros" } -feldera-sqllib = { version = "0.322.0", path = "crates/sqllib" } -feldera-storage = { version = "0.322.0", path = "crates/storage" } -feldera-types = { version = "0.322.0", path = "crates/feldera-types" } -feldera-rest-api = { version = "0.322.0", path = "crates/rest-api" } -feldera-ir = { version = "0.322.0", path = "crates/ir" } -feldera-adapterlib = { version = "0.322.0", path = "crates/adapterlib" } -feldera-samply = { version = "0.322.0", path = "crates/samply" } +feldera-observability = { version = "0.323.0", path = "crates/feldera-observability" } +feldera-macros = { version = "0.323.0", path = "crates/feldera-macros" } +feldera-sqllib = { version = "0.323.0", path = "crates/sqllib" } +feldera-storage = { version = "0.323.0", path = "crates/storage" } +feldera-types = { version = "0.323.0", path = "crates/feldera-types" } +feldera-rest-api = { version = "0.323.0", path = "crates/rest-api" } +feldera-ir = { version = "0.323.0", path = "crates/ir" } +feldera-adapterlib = { version = "0.323.0", path = "crates/adapterlib" } +feldera-samply = { version = "0.323.0", path = "crates/samply" } flate2 = "1.1.0" form_urlencoded = "1.2.0" fs_extra = "1.3.0" @@ -253,7 +255,15 @@ schema_registry_converter = "4.2.0" semver = "1.0.27" serde = "1.0.213" serde_dynamo = { version = "4.3.0", features = ["aws-sdk-dynamodb+1"] } -serde_arrow = { version = "0.14.1", features = ["arrow-58"] } +# serde_arrow fork that recognizes the "+00:00" timezone Iceberg uses for +# `timestamptz` columns (upstream 0.14.2 accepts only the name "UTC"). +# Remove the git source when https://github.com/chmp/serde_arrow/pull/320 +# is released. `version` keeps the crates that depend on serde_arrow +# publishable to crates.io; cargo drops the git source when publishing. +# The git source goes on the dependency rather than in a `[patch]` section +# because patches do not apply to generated SQL pipeline workspaces, which +# import this dependency through `crates/adapters`. +serde_arrow = { version = "0.14.2", git = "https://github.com/ryzhyk/serde_arrow.git", rev = "9555b34fb85d4323e4381ad5a09872c1696a0d34", features = ["arrow-58"] } serde_bytes = "0.11.15" serde_json = { version = "1.0.132", features = ["arbitrary_precision"] } serde_json_path_to_error = "0.1.5" diff --git a/crates/adapterlib/Cargo.toml b/crates/adapterlib/Cargo.toml index 8b06684aba8..031a23575c0 100644 --- a/crates/adapterlib/Cargo.toml +++ b/crates/adapterlib/Cargo.toml @@ -42,6 +42,7 @@ num-derive = { workspace = true } tokio = { workspace = true, features = ["sync"] } xxhash-rust = { workspace = true, features = ["xxh3"] } datafusion = { workspace = true } +datafusion-functions-json = { workspace = true } thiserror = { workspace = true } serde_json_path_to_error = { workspace = true } chrono = { workspace = true } diff --git a/crates/adapterlib/src/utils/datafusion.rs b/crates/adapterlib/src/utils/datafusion.rs index c828838d8d3..70fb60f14ca 100644 --- a/crates/adapterlib/src/utils/datafusion.rs +++ b/crates/adapterlib/src/utils/datafusion.rs @@ -219,11 +219,16 @@ where ); let session_config = customize_config(session_config); - let state = SessionStateBuilder::new() + let mut state = SessionStateBuilder::new() .with_config(session_config) .with_runtime_env(runtime_env) .with_default_features() .build(); + // JSON functions for querying VARIANT columns, which reach DataFusion + // as JSON-encoded strings. Note: the crate's `->` and `?` operators do + // not parse in the default (generic) SQL dialect; `->>` works. + datafusion_functions_json::register_all(&mut state) + .expect("registering JSON functions on a fresh session state cannot fail"); SessionContext::from(state) } @@ -364,6 +369,35 @@ pub fn columns_referenced_by_order_by(order_by: &str) -> Result Ok(columns) } +/// Takes a column name from an external table schema and returns a quoted +/// string that can be used in datafusion queries like `select "foo""bar" from my_table`. +pub fn quote_sql_identifier>(ident: S) -> String { + format!("\"{}\"", ident.as_ref().replace("\"", "\"\"")) +} + +/// A set of column names compared case-insensitively. External table schemas +/// (Delta, Iceberg) carry no case-sensitivity information, so names are stored +/// and probed in lowercased form. +/// +/// SQL is case-sensitive for quoted column names, but an external table cannot +/// hold two columns with the same lowercase form, so collapsing to a single +/// canonical form is safe here. +#[derive(Default)] +pub struct ColumnNameSet { + lowercase: BTreeSet, +} + +impl ColumnNameSet { + pub fn from_names(names: impl IntoIterator) -> Self { + let lowercase = names.into_iter().map(|c| c.to_lowercase()).collect(); + Self { lowercase } + } + + pub fn contains(&self, name: &str) -> bool { + self.lowercase.contains(&name.to_lowercase()) + } +} + /// Convert a value of the timestamp column returned by a SQL query into a valid /// SQL expression. pub fn timestamp_to_sql_expression(column_type: &ColumnType, expr: &str) -> String { @@ -704,6 +738,40 @@ mod tests { assert_eq!(ctx.copied_config().target_partitions(), 99); } + /// Every context built here must provide the JSON function family. + #[test] + fn create_session_context_registers_json_functions() { + let storage = TempStorage::new("feldera-datafusion-create-session-context-json-test"); + let cfg = pipeline_config( + RuntimeConfig { + workers: 1, + ..Default::default() + }, + Some(storage.path()), + ); + let env = create_runtime_env(&cfg).unwrap(); + let ctx = create_session_context(&cfg, env); + let state = ctx.state(); + for function in [ + "json_get", + "json_get_str", + "json_get_int", + "json_get_float", + "json_get_bool", + "json_get_json", + "json_get_array", + "json_as_text", + "json_contains", + "json_length", + "json_object_keys", + ] { + assert!( + state.scalar_functions().contains_key(function), + "JSON function '{function}' is not registered" + ); + } + } + /// Tripwire: `clean_stale_scratch_entries` refuses to walk a directory /// whose final component isn't `DATAFUSION_TEMP_DIR`, so a misuse can't /// recursively wipe an arbitrary path. diff --git a/crates/adapters/src/adhoc.rs b/crates/adapters/src/adhoc.rs index 44ead9b7d1c..893677715da 100644 --- a/crates/adapters/src/adhoc.rs +++ b/crates/adapters/src/adhoc.rs @@ -705,6 +705,137 @@ mod tests { assert_eq!(total_rows, 1); } + /// Ad-hoc queries can take apart VARIANT columns with the JSON function + /// family from `datafusion-functions-json`. VARIANT values reach + /// DataFusion as JSON-encoded Utf8 strings; the in-memory table below + /// stands in for a materialized table with a VARIANT column. + /// + /// The state must come from `create_session_context`, which registers + /// the functions; the plain `test_state()` used elsewhere in this + /// module does not have them. + #[tokio::test] + async fn json_functions_work_in_adhoc_queries() { + use datafusion::arrow::array::{Array, Int64Array, StringArray}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::datasource::MemTable; + use feldera_adapterlib::utils::datafusion::{create_runtime_env, create_session_context}; + use feldera_types::config::{PipelineConfig, RuntimeConfig}; + + // Tripwire: stock DataFusion does not know these functions; they + // must come from the registration in `create_session_context`. If + // this starts failing, DataFusion gained native JSON functions and + // the `datafusion-functions-json` dependency may be redundant. + let err = execute_sql_with_state(test_state(), "SELECT json_get_str('{}', 'a')") + .await + .expect_err("stock DataFusion should not provide json_get_str"); + assert!( + format!("{err}").contains("json_get_str"), + "unexpected error: {err}" + ); + + let config = PipelineConfig { + global: RuntimeConfig { + workers: 1, + ..Default::default() + }, + multihost: None, + name: None, + given_name: None, + storage_config: None, + secrets_dir: None, + inputs: Default::default(), + outputs: Default::default(), + program_ir: None, + }; + let runtime_env = create_runtime_env(&config).unwrap(); + let ctx = create_session_context(&config, runtime_env); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("val", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec![ + Some(r#"{"name":"Bob","scores":[8,10],"active":true}"#), + Some(r#"{"name":"Ann","scores":[3],"extra":{"x":5}}"#), + None, + ])), + ], + ) + .unwrap(); + let mem = MemTable::try_new(schema, vec![vec![batch]]).unwrap(); + ctx.register_table("t", Arc::new(mem)).unwrap(); + + // Typed getters in the projection. A missing key ('scores' has one + // element in row 2; indexes are 0-based) and a NULL document both + // yield NULL. + let batches = collect_rows( + ctx.state(), + "SELECT json_get_str(val, 'name') AS name, \ + json_get_int(val, 'scores', 1) AS second_score \ + FROM t ORDER BY id", + ) + .await; + let names = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .expect("string column"); + assert_eq!(names.value(0), "Bob"); + assert_eq!(names.value(1), "Ann"); + assert!(names.is_null(2)); + let scores = batches[0] + .column(1) + .as_any() + .downcast_ref::() + .expect("int64 column"); + assert_eq!(scores.value(0), 10); + assert!(scores.is_null(1)); + assert!(scores.is_null(2)); + + // JSON predicates filter on document contents. + for (query, expected_id) in [ + ("SELECT id FROM t WHERE json_contains(val, 'extra')", 2), + ( + "SELECT id FROM t WHERE json_get_bool(val, 'active') = TRUE", + 1, + ), + ] { + let batches = collect_rows(ctx.state(), query).await; + let ids = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .expect("int64 column"); + assert_eq!(ids.len(), 1, "query: {query}"); + assert_eq!(ids.value(0), expected_id, "query: {query}"); + } + + // The `->>` operator and casting `json_get` to a concrete type. + let batches = collect_rows( + ctx.state(), + "SELECT val->>'name' AS name, \ + CAST(json_get(val, 'scores', 0) AS BIGINT) AS first_score \ + FROM t WHERE json_get_str(val, 'name') = 'Bob'", + ) + .await; + let names = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .expect("string column"); + assert_eq!(names.value(0), "Bob"); + let scores = batches[0] + .column(1) + .as_any() + .downcast_ref::() + .expect("int64 column"); + assert_eq!(scores.value(0), 8); + } + /// Selecting from a non-materialized source fails during *physical /// planning* — when the scan node is built — not during scan execution. /// `execute_adhoc_stream` must surface that as an `Err`, letting the HTTP diff --git a/crates/adapters/src/integrated/delta_table/input.rs b/crates/adapters/src/integrated/delta_table/input.rs index 75fe65db27b..ce054f1ac65 100644 --- a/crates/adapters/src/integrated/delta_table/input.rs +++ b/crates/adapters/src/integrated/delta_table/input.rs @@ -39,10 +39,10 @@ use feldera_adapterlib::format::{ParseError, StagedInputBuffer}; use feldera_adapterlib::metrics::{ConnectorMetrics, ValueType}; use feldera_adapterlib::transport::{InputQueueEntry, Resume, Watermark, parse_resume_info}; use feldera_adapterlib::utils::datafusion::{ - array_to_string, columns_referenced_by_expression, columns_referenced_by_order_by, - create_session_context_with, execute_query_collect, execute_singleton_query, - timestamp_to_sql_expression, validate_sql_expression, validate_sql_order_by, - validate_timestamp_column, + ColumnNameSet, array_to_string, columns_referenced_by_expression, + columns_referenced_by_order_by, create_session_context_with, execute_query_collect, + execute_singleton_query, quote_sql_identifier, timestamp_to_sql_expression, + validate_sql_expression, validate_sql_order_by, validate_timestamp_column, }; use feldera_storage::tokio::TOKIO_DEDICATED_IO; use feldera_types::adapter_stats::ConnectorHealth; @@ -112,12 +112,6 @@ const DEFAULT_MAX_CONCURRENT_READERS: usize = 6; /// connector initialization. static MAX_CONCURRENT_READERS: AtomicUsize = AtomicUsize::new(0); -/// Takes a column name from a DeltaLake schema and returns a quoted string -/// that can be used in datafusion queries like `select "foo""bar" from my_table`. -fn quote_sql_identifier>(ident: S) -> String { - format!("\"{}\"", ident.as_ref().replace("\"", "\"\"")) -} - /// Format a DataFusion error, appending actionable guidance when the /// underlying variant is `ResourcesExhausted` (the shared memory pool ran /// out). `find_root` walks past `Context(...)` / `ArrowError(...)` wrappers @@ -883,29 +877,6 @@ struct CatchupFollowState { transaction: Option>, } -/// A set of column names compared case-insensitively. Delta schemas carry no -/// case-sensitivity information, so names are stored and probed in lowercased -/// form (mirrors the long-standing matching in [`used_columns`](DeltaTableInputEndpointInner::used_columns)). -/// -/// SQL is case-sensitive for quoted column names, but an external table cannot -/// hold two columns with the same lowercase form, so collapsing to a single -/// canonical form is safe here. -#[derive(Default)] -struct ColumnNameSet { - lowercase: BTreeSet, -} - -impl ColumnNameSet { - fn from_names(names: impl IntoIterator) -> Self { - let lowercase = names.into_iter().map(|c| c.to_lowercase()).collect(); - Self { lowercase } - } - - fn contains(&self, name: &str) -> bool { - self.lowercase.contains(&name.to_lowercase()) - } -} - struct DeltaTableInputEndpointInner { endpoint_name: String, schema: Relation, @@ -1090,8 +1061,7 @@ impl DeltaTableInputEndpointInner { fn skip_unused_columns(&self) -> bool { // Old-style: property in the connector configuration. // New-style: property in the table definition. - self.config.skip_unused_columns - || self.schema.get_property("skip_unused_columns") == Some("true") + self.config.skip_unused_columns || self.schema.skip_unused_columns() } fn new_follow_transaction_label(&self) -> Option> { diff --git a/crates/adapters/src/static_compile/catalog.rs b/crates/adapters/src/static_compile/catalog.rs index c7a2e4585a2..2189233d93e 100644 --- a/crates/adapters/src/static_compile/catalog.rs +++ b/crates/adapters/src/static_compile/catalog.rs @@ -101,6 +101,24 @@ impl Catalog { .shard_workers_accumulate(hosts[ordinal].workers.clone()) .into_parts(); + // The output gather is collaborative: every host shards its slice of + // the view to the owning host's workers, which accumulate and emit + // it. The accumulator's `enable_count` is per-host and is enabled + // only when an output connector attaches, which in a multihost + // pipeline happens only on the owning host (connectors, including + // dynamic HTTP `listen`, are assigned to one host). A non-owning + // host would therefore leave its accumulator disabled and send an + // empty batch instead of its slice, silently dropping every view + // record computed on that host. + // + // Enable the accumulator on every host so the gather is always + // complete. Materialized views already get this via the integral's + // `into_enabled_stream`; do the same for the delta gather. This + // gathers a view even when it has no connector; a future + // optimization could propagate the owning host's connector state to + // the other hosts to skip that work. + enabled_count.enable(); + let integral = if integrate { Some( stream.shard_workers_accumulate_integrate_trace(hosts[ordinal].workers.clone()), diff --git a/crates/adapters/src/test.rs b/crates/adapters/src/test.rs index 93d28f6b47b..4e64e0c3f46 100644 --- a/crates/adapters/src/test.rs +++ b/crates/adapters/src/test.rs @@ -59,14 +59,16 @@ use crate::catalog::InputCollectionHandle; use crate::format::get_input_format; use crate::transport::input_transport_config_to_endpoint; pub use data::{ - DatabricksPeople, DeltaTestStruct, EmbeddedStruct, IcebergTestStruct, KeyStruct, - S3TablesTestStruct, TestStruct, TestStruct2, generate_test_batch, generate_test_batches, - generate_test_batches_with_weights, + DatabricksPeople, DeltaTestStruct, EmbeddedStruct, IcebergSubsetTestStruct, IcebergTestStruct, + KeyStruct, S3TablesTestStruct, TestStruct, TestStruct2, generate_test_batch, + generate_test_batches, generate_test_batches_with_weights, }; use dbsp::circuit::{CircuitConfig, NodeId}; use dbsp::utils::Tup2; use feldera_types::format::json::{JsonFlavor, JsonLines, JsonParserConfig, JsonUpdateFormat}; -use feldera_types::program_schema::{Field, Relation, SqlIdentifier}; +use feldera_types::program_schema::{ + Field, PropertyValue, Relation, SourcePosition, SqlIdentifier, +}; pub use mock_dezset::{ MockDeZSet, MockUpdate, wait_for_output_count, wait_for_output_ordered, wait_for_output_unordered, @@ -242,6 +244,44 @@ where + for<'de> DeserializeWithContext<'de, SqlSerdeConfig, Variant> + Sync, { + test_circuit_with_properties::(config, schema, &[], persistent_output_ids) +} + +/// Like [`test_circuit`], but sets the given properties on the input relations, +/// the way the SQL compiler records table-level `WITH` properties (e.g., +/// `skip_unused_columns`) in the program schema. +pub fn test_circuit_with_properties( + config: CircuitConfig, + schema: &[Field], + input_properties: &[(&str, &str)], + persistent_output_ids: &[Option<&str>], +) -> (DBSPHandle, Box) +where + T: DBData + + SerializeWithContext + + for<'de> DeserializeWithContext<'de, SqlSerdeConfig, Variant> + + Sync, +{ + let zero_position = SourcePosition { + start_line_number: 0, + start_column: 0, + end_line_number: 0, + end_column: 0, + }; + let input_properties: BTreeMap = input_properties + .iter() + .map(|(key, value)| { + ( + key.to_string(), + PropertyValue { + value: value.to_string(), + key_position: zero_position, + value_position: zero_position, + }, + ) + }) + .collect(); + let schema = schema.to_vec(); let persistent_output_ids = persistent_output_ids .iter() @@ -264,7 +304,7 @@ where format!("test_input{i}").into(), schema.clone(), false, - BTreeMap::new(), + input_properties.clone(), )) .unwrap(); diff --git a/crates/adapters/src/test/data.rs b/crates/adapters/src/test/data.rs index 86c7c9789cd..691e3574cf2 100644 --- a/crates/adapters/src/test/data.rs +++ b/crates/adapters/src/test/data.rs @@ -6,7 +6,7 @@ use arrow::datatypes::{DataType, Schema, TimeUnit}; use dbsp::utils::Tup2; use feldera_macros::IsNone; use feldera_sqllib::{ - ByteArray, Date, F32, F64, SqlDecimal, SqlString, Time, Timestamp, Uuid, Variant, + ByteArray, Date, F32, F64, SqlDecimal, SqlString, Time, Timestamp, TimestampTz, Uuid, Variant, }; use feldera_types::program_schema::{ColumnType, Field, Relation, SqlIdentifier}; use feldera_types::{ @@ -784,6 +784,7 @@ pub struct IcebergTestStruct { //pub uuid: ByteArray, pub fixed: ByteArray, pub varbin: ByteArray, + pub tstz: TimestampTz, } impl Arbitrary for IcebergTestStruct { @@ -808,6 +809,8 @@ impl Arbitrary for IcebergTestStruct { 0u64..24 * 3600 * 1_000_000_000, // Generate timestamps within a 1-year range 1704070800i64..1735693200, + // Generate timestamptz within a 1-year range (microseconds since epoch) + 1704070800_000_000i64..1735693200_000_000, ), // String in the range "0".."1000" ( @@ -822,10 +825,10 @@ impl Arbitrary for IcebergTestStruct { ) .prop_map( |( - (b, i, l, r, d, dec_num, dec_scale, dt, tm, ts), + (b, i, l, r, d, dec_num, dec_scale, dt, tm, ts, tstz), (s, /*uuid,*/ fixed, varbin), ): ( - (bool, i32, i64, f32, f64, i128, u32, i32, u64, i64), + (bool, i32, i64, f32, f64, i128, u32, i32, u64, i64, i64), (i32, /*Vec,*/ Vec, Vec), )| { IcebergTestStruct { @@ -842,6 +845,7 @@ impl Arbitrary for IcebergTestStruct { // uuid: ByteArray::from_vec(uuid), fixed: ByteArray::from_vec(fixed), varbin: ByteArray::new(&varbin), + tstz: TimestampTz::from_microseconds(tstz), } }, ) @@ -869,6 +873,11 @@ impl IcebergTestStruct { // arrow::datatypes::Field::new("uuid", DataType::FixedSizeBinary(16), false), arrow::datatypes::Field::new("fixed", DataType::FixedSizeBinary(5), false), arrow::datatypes::Field::new("varbin", DataType::Binary, false), + arrow::datatypes::Field::new( + "tstz", + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + false, + ), ])) } @@ -887,6 +896,7 @@ impl IcebergTestStruct { // Field::new("uuid".into(), ColumnType::fixed(16, false)), Field::new("fixed".into(), ColumnType::fixed(5, false)), Field::new("varbin".into(), ColumnType::varbinary(false)), + Field::new("tstz".into(), ColumnType::timestamp_tz(false)), ] } @@ -906,6 +916,7 @@ impl IcebergTestStruct { // Field::new("uuid".into(), ColumnType::fixed(16, false)), Field::new("fixed".into(), ColumnType::fixed(5, false)), Field::new("varbin".into(), ColumnType::varbinary(false)), + Field::new("tstz".into(), ColumnType::timestamp_tz(false)), ]; fields @@ -922,7 +933,7 @@ impl IcebergTestStruct { } } -serialize_table_record!(IcebergTestStruct[12]{ +serialize_table_record!(IcebergTestStruct[13]{ b["b"]: bool, i["i"]: i32, l["l"]: i64, @@ -935,10 +946,11 @@ serialize_table_record!(IcebergTestStruct[12]{ s["s"]: String, // uuid["uuid"]: ByteArray, fixed["fixed"]: ByteArray, - varbin["varbin"]: ByteArray + varbin["varbin"]: ByteArray, + tstz["tstz"]: TimestampTz }); -deserialize_table_record!(IcebergTestStruct["IcebergTestStruct", Variant, 12] { +deserialize_table_record!(IcebergTestStruct["IcebergTestStruct", Variant, 13] { (b, "b", false, bool, |_| None), (i, "i", false, i32, |_| None), (l, "l", false, i64, |_| None), @@ -951,7 +963,8 @@ deserialize_table_record!(IcebergTestStruct["IcebergTestStruct", Variant, 12] { (s, "s", false, String, |_| None), // (uuid, "uuid", false, ByteArray, None), (fixed, "fixed", false, ByteArray, |_| None), - (varbin, "varbin", false, ByteArray, |_| None) + (varbin, "varbin", false, ByteArray, |_| None), + (tstz, "tstz", false, TimestampTz, |_| None) }); /// Records in the Amazon S3 Tables integration-test table (`dev.test_table`). @@ -1003,6 +1016,57 @@ deserialize_table_record!(S3TablesTestStruct["S3TablesTestStruct", Variant, 3] { (created_at, "created_at", true, Option, |_| Some(None)) }); +/// SQL table that declares a subset of [`IcebergTestStruct`]'s columns, for +/// testing that the Iceberg connector only reads declared columns. +/// +/// `l` is nullable and marked unused, so the connector skips it (and the +/// deserializer NULL-fills it) when the `skip_unused_columns` table property +/// is set. +#[derive( + Debug, + Default, + PartialEq, + Eq, + PartialOrd, + Ord, + serde::Serialize, + Clone, + Hash, + SizeOf, + rkyv::Archive, + rkyv::Serialize, + rkyv::Deserialize, + IsNone, +)] +#[archive_attr(derive(Ord, Eq, PartialEq, PartialOrd))] +pub struct IcebergSubsetTestStruct { + pub i: i32, + pub s: String, + pub l: Option, +} + +impl IcebergSubsetTestStruct { + pub fn schema() -> Vec { + vec![ + Field::new("i".into(), ColumnType::int(false)), + Field::new("s".into(), ColumnType::varchar(false)), + Field::new("l".into(), ColumnType::bigint(true)).with_unused(true), + ] + } +} + +serialize_table_record!(IcebergSubsetTestStruct[3]{ + i["i"]: i32, + s["s"]: String, + l["l"]: Option +}); + +deserialize_table_record!(IcebergSubsetTestStruct["IcebergSubsetTestStruct", Variant, 3] { + (i, "i", false, i32, |_| None), + (s, "s", false, String, |_| None), + (l, "l", false, Option, |_| Some(None)) +}); + /// Struct will all types supported by the DeltaLake connector. #[derive( Debug, diff --git a/crates/adapters/src/test/iceberg.rs b/crates/adapters/src/test/iceberg.rs index 990056c55c5..f596071d9fb 100644 --- a/crates/adapters/src/test/iceberg.rs +++ b/crates/adapters/src/test/iceberg.rs @@ -8,7 +8,7 @@ use crossbeam::channel::Receiver; use dbsp::DBData; use feldera_sqllib::Variant; #[cfg(feature = "iceberg-tests-fs")] -use feldera_sqllib::{ByteArray, F32, F64}; +use feldera_sqllib::{ByteArray, F32, F64, Timestamp, TimestampTz}; use feldera_types::{ program_schema::Field, serde_with_context::{DeserializeWithContext, SerializeWithContext, SqlSerdeConfig}, @@ -23,6 +23,8 @@ use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitEx #[cfg(feature = "iceberg-tests-fs")] use std::io::Write; +#[cfg(feature = "iceberg-tests-fs")] +use super::IcebergSubsetTestStruct; #[cfg(any( feature = "iceberg-tests-fs", feature = "iceberg-tests-glue", @@ -31,7 +33,7 @@ use std::io::Write; use super::IcebergTestStruct; #[cfg(feature = "iceberg-tests-s3tables")] use super::S3TablesTestStruct; -use super::test_circuit; +use super::test_circuit_with_properties; fn init_logging() { let _ = tracing_subscriber::registry() @@ -64,7 +66,14 @@ fn data_to_ndjson(data: Vec) -> NamedTempFile { } /// Read a snapshot of an Iceberg table with records of type `T` to a temporary JSON file. -fn iceberg_snapshot_to_json(schema: &[Field], config: &HashMap) -> NamedTempFile +/// +/// `table_properties` are set on the input relation, the way table-level SQL +/// `WITH` properties (e.g., `skip_unused_columns`) reach the connector. +fn iceberg_snapshot_to_json( + schema: &[Field], + table_properties: &[(&str, &str)], + config: &HashMap, +) -> NamedTempFile where T: DBData + SerializeWithContext @@ -81,8 +90,12 @@ where let mut config = config.clone(); config.insert("mode".to_string(), "snapshot".to_string()); - let (input_pipeline, err_receiver) = - iceberg_input_pipeline::(schema, &config, &json_file.path().display().to_string()); + let (input_pipeline, err_receiver) = iceberg_input_pipeline::( + schema, + table_properties, + &config, + &json_file.path().display().to_string(), + ); input_pipeline.start(); wait( || input_pipeline.status().pipeline_complete() || err_receiver.len() > 0, @@ -102,6 +115,7 @@ where /// Build a pipeline that reads from an Iceberg table and writes to a JSON file. fn iceberg_input_pipeline( schema: &[Field], + table_properties: &[(&str, &str)], config: &HashMap, output_file_path: &str, ) -> (Controller, Receiver) @@ -147,11 +161,26 @@ where .unwrap(); let schema = schema.to_vec(); + let table_properties: Vec<(String, String)> = table_properties + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect(); let (err_sender, err_receiver) = crossbeam::channel::unbounded(); let controller = Controller::with_test_config( - move |workers| Ok(test_circuit::(workers, &schema, &[None])), + move |workers| { + let table_properties: Vec<(&str, &str)> = table_properties + .iter() + .map(|(k, v)| (k.as_str(), v.as_str())) + .collect(); + Ok(test_circuit_with_properties::( + workers, + &schema, + &table_properties, + &[None], + )) + }, &config, Box::new(move |e, _| { let msg = format!("iceberg_input_test: error: {e}"); @@ -187,6 +216,7 @@ fn data(n_records: usize) -> Vec { // uuid: ByteArray::new([0u8; 16].as_slice()), fixed: ByteArray::new([0u8; 5].as_slice()), varbin: ByteArray::new([0u8; 5].as_slice()), + tstz: TimestampTz::from(Timestamp::from_naiveDateTime(time)), }); time += std::time::Duration::from_secs(1); @@ -222,17 +252,15 @@ fn iceberg_localfs_input_test_ordered_with_filter() { ); } +/// Create a local Iceberg table populated with `data` and return its metadata +/// location. With `extra_columns`, the table gets columns that no test SQL +/// schema declares (see `--extra-columns` in `create_test_table_s3.py`). #[cfg(feature = "iceberg-tests-fs")] -fn iceberg_localfs_input_test( - extra_config: &[(String, String)], - filter: &dyn Fn(&IcebergTestStruct) -> bool, -) { - let data = data(1_000_000); - +fn create_localfs_table(data: &[IcebergTestStruct], extra_columns: bool) -> String { let table_dir = tempfile::TempDir::new().unwrap(); let table_path = table_dir.path().display().to_string(); - let ndjson_file = data_to_ndjson(data.clone()); + let ndjson_file = data_to_ndjson(data.to_vec()); println!("wrote test data to {}", ndjson_file.path().display()); // Uncomment to inspect output parquet files produced by the test. @@ -241,11 +269,16 @@ fn iceberg_localfs_input_test( let script_path = "../iceberg/src/test/create_test_table_s3.py"; // Run the Python script using the Python interpreter - let output = std::process::Command::new("python3") + let mut command = std::process::Command::new("python3"); + command .arg(script_path) .arg("--catalog=sql") .arg(format!("--warehouse-path={table_path}")) - .arg(format!("--json-file={}", ndjson_file.path().display())) + .arg(format!("--json-file={}", ndjson_file.path().display())); + if extra_columns { + command.arg("--extra-columns"); + } + let output = command .output() .map_err(|e| { format!("Error running '{script_path}' script to generate an Iceberg table: {e}") @@ -262,15 +295,26 @@ fn iceberg_localfs_input_test( } // The script should print table metadata location on the last line. - let metadata_path = String::from_utf8(output.stdout.clone()) + String::from_utf8(output.stdout.clone()) .unwrap() .lines() .last() .unwrap() - .to_string(); + .to_string() +} + +#[cfg(feature = "iceberg-tests-fs")] +fn iceberg_localfs_input_test( + extra_config: &[(String, String)], + filter: &dyn Fn(&IcebergTestStruct) -> bool, +) { + let data = data(1_000_000); + + let metadata_path = create_localfs_table(&data, false); let mut json_file = iceberg_snapshot_to_json::( &IcebergTestStruct::schema_with_lateness(), + &[], &[("metadata_location".to_string(), metadata_path.to_string())] .into_iter() .chain(extra_config.into_iter().cloned()) @@ -290,6 +334,66 @@ fn iceberg_localfs_input_test( assert_eq!(zset, expected_zset); } +/// Read a table through a SQL declaration that names only a few of its +/// columns, while the table also holds columns (including a `uuid` one, a +/// type no test struct models) that the connector must ignore because it +/// selects the declared columns instead of `*`. +/// +/// With `skip_unused` (the `skip_unused_columns` table property), the +/// connector must additionally not read the nullable `l` column, which the +/// SQL schema marks unused, so `l` comes out NULL. This variant fails if the +/// connector falls back to reading all columns. +#[cfg(feature = "iceberg-tests-fs")] +fn iceberg_localfs_input_subset_test(skip_unused: bool) { + let data = data(100_000); + + let metadata_path = create_localfs_table(&data, true); + + let table_properties: &[(&str, &str)] = if skip_unused { + &[("skip_unused_columns", "true")] + } else { + &[] + }; + + let mut json_file = iceberg_snapshot_to_json::( + &IcebergSubsetTestStruct::schema(), + table_properties, + &[("metadata_location".to_string(), metadata_path)] + .into_iter() + .collect::>(), + ); + + let expected_zset = dbsp::OrdZSet::from_tuples( + (), + data.into_iter() + .map(|x| IcebergSubsetTestStruct { + i: x.i, + s: x.s, + l: if skip_unused { None } else { Some(x.l) }, + }) + .map(|x| dbsp::utils::Tup2(dbsp::utils::Tup2(x, ()), 1)) + .collect(), + ); + let zset = file_to_zset::(json_file.as_file_mut()); + + assert_eq!(zset, expected_zset); +} + +/// The connector reads only the columns the SQL table declares. +#[test] +#[cfg(feature = "iceberg-tests-fs")] +fn iceberg_localfs_input_test_subset_schema() { + iceberg_localfs_input_subset_test(false); +} + +/// The `skip_unused_columns` table property also drops declared-but-unused +/// columns from the read. +#[test] +#[cfg(feature = "iceberg-tests-fs")] +fn iceberg_localfs_input_test_skip_unused_columns() { + iceberg_localfs_input_subset_test(true); +} + #[test] #[cfg(feature = "iceberg-tests-glue")] fn iceberg_glue_s3_input_test() { @@ -297,6 +401,7 @@ fn iceberg_glue_s3_input_test() { // Read delta table unordered. let mut json_file = iceberg_snapshot_to_json::( &IcebergTestStruct::schema_with_lateness(), + &[], &[ ("catalog_type".to_string(), "glue".to_string()), ( @@ -305,7 +410,7 @@ fn iceberg_glue_s3_input_test() { ), ( "table_name".to_string(), - "iceberg_test.test_table".to_string(), + "iceberg_test.test_table_v2".to_string(), ), ( "glue.access-key-id".to_string(), @@ -354,6 +459,7 @@ fn iceberg_s3tables_input_test() { // `AWS_ACCESS_KEY_ID`/`AWS_SECRET_ACCESS_KEY`(/`AWS_SESSION_TOKEN`) exported. let mut json_file = iceberg_snapshot_to_json::( &S3TablesTestStruct::schema(), + &[], &[ ("catalog_type".to_string(), "s3tables".to_string()), ( @@ -381,6 +487,7 @@ fn iceberg_rest_s3_input_test() { // Read delta table unordered. let mut json_file = iceberg_snapshot_to_json::( &IcebergTestStruct::schema_with_lateness(), + &[], &[ ("catalog_type".to_string(), "rest".to_string()), ("rest.uri".to_string(), "http://localhost:8181".to_string()), @@ -390,7 +497,7 @@ fn iceberg_rest_s3_input_test() { ), ( "table_name".to_string(), - "iceberg_test.test_table".to_string(), + "iceberg_test.test_table_v2".to_string(), ), ( "s3.access-key-id".to_string(), diff --git a/crates/dbsp/src/circuit/circuit_builder.rs b/crates/dbsp/src/circuit/circuit_builder.rs index 91f59bccf5a..e07d24b8860 100644 --- a/crates/dbsp/src/circuit/circuit_builder.rs +++ b/crates/dbsp/src/circuit/circuit_builder.rs @@ -3939,31 +3939,12 @@ where SndOp: SinkOperator, RcvOp: SourceOperator, { - let sender_id = self.add_node(|id| { - self.log_circuit_event(&CircuitEvent::operator( - GlobalNodeId::child_of(self, id), - sender.name(), - sender.location(), - )); - - let node = SinkNode::new(sender, input_stream.clone(), self.clone(), id); - self.connect_stream(input_stream, id, input_preference); - (node, id) - }); - - let output_stream = self.add_node(|id| { - self.log_circuit_event(&CircuitEvent::operator( - GlobalNodeId::child_of(self, id), - receiver.name(), - receiver.location(), - )); - - let node = SourceNode::new(receiver, self.clone(), id); - let output_stream = node.output_stream(); - (node, output_stream) - }); - - self.add_dependency(sender_id, output_stream.local_node_id()); + let sender_id = self.add_sink_with_preference(sender, input_stream, input_preference); + let output_stream = self.add_source(receiver); + self.add_dependency( + sender_id.local_node_id().unwrap(), + output_stream.local_node_id(), + ); output_stream } diff --git a/crates/dbsp/src/circuit/runtime.rs b/crates/dbsp/src/circuit/runtime.rs index 8e24494cb96..9720ffe1dff 100644 --- a/crates/dbsp/src/circuit/runtime.rs +++ b/crates/dbsp/src/circuit/runtime.rs @@ -7,7 +7,7 @@ use crate::SchedulerError; use crate::circuit::checkpointer::Checkpointer; use crate::circuit::dbsp_handle::StepSize; use crate::error::Error as DbspError; -use crate::operator::communication::Exchange; +use crate::operator::communication::{Exchange, ExchangeActivity}; use crate::storage::backend::StorageBackend; use crate::storage::file::format::Compression; use crate::storage::file::writer::Parameters; @@ -1383,7 +1383,8 @@ where match Runtime::runtime() { Some(runtime) if Runtime::num_workers() > 1 => { let exchange_id = runtime.sequence_next().try_into().unwrap(); - let exchange = Exchange::with_runtime(&runtime, exchange_id); + let exchange = + Exchange::with_runtime(&runtime, exchange_id, ExchangeActivity::AllSteps); let identifier = Arc::new(format!("broadcast {name} (exchange {exchange_id})")); Self::MultiThreaded { @@ -1426,8 +1427,9 @@ where .await; exchange - .receive_all(|data| serde_json::from_slice(&data).unwrap()) + .receive_all(|data| serde_json::from_slice(&data).unwrap(), None) .await + .0 }) .await .ok_or(SchedulerError::Killed) diff --git a/crates/dbsp/src/operator/communication.rs b/crates/dbsp/src/operator/communication.rs index ec1eb6b591a..73ab4ef33cc 100644 --- a/crates/dbsp/src/operator/communication.rs +++ b/crates/dbsp/src/operator/communication.rs @@ -3,6 +3,9 @@ mod gather; mod shard; pub(crate) use exchange::{ - Exchange, ExchangeClients, ExchangeDelivery, ExchangeDirectory, ExchangeId, pop_flushed, + Exchange, ExchangeClients, ExchangeDelivery, ExchangeDirectory, ExchangeId, MessageType, + pop_flushed, +}; +pub use exchange::{ + ExchangeActivity, ExchangeReceiver, ExchangeSender, Mailbox, new_exchange_operators, }; -pub use exchange::{ExchangeReceiver, ExchangeSender, Mailbox, new_exchange_operators}; diff --git a/crates/dbsp/src/operator/communication/exchange.rs b/crates/dbsp/src/operator/communication/exchange.rs index b6c2e6d289f..71f22513b01 100644 --- a/crates/dbsp/src/operator/communication/exchange.rs +++ b/crates/dbsp/src/operator/communication/exchange.rs @@ -23,6 +23,7 @@ use crate::{ }; use binrw::{BinRead, BinWrite}; use crossbeam_utils::CachePadded; +use enum_map::{Enum, EnumMap}; use feldera_samply::Span; use feldera_storage::fbuf::FBuf; use itertools::Itertools; @@ -239,6 +240,45 @@ struct ExchangeMessage { data: Vec, } +/// Distinguishes messages by size. +/// +/// We segregate big and small messages into separate queues so that simple +/// broadcast and consensus messages don't get delayed behind bigger messages. +#[derive(Copy, Clone, Debug, PartialEq, Eq, Enum)] +pub(crate) enum MessageSize { + /// A big message. + Big, + + /// A small message. + Small, +} + +impl MessageSize { + /// Constructs `MessageSize` from a count of `bytes`. + pub(crate) fn from_bytes(bytes: usize) -> Self { + if bytes <= 4096 { + Self::Small + } else { + Self::Big + } + } +} + +/// Categorizes an [ExchangeMessage]. +/// +/// We segregate messages in different categories into different queues so that +/// messages in one category do not delay those in other categories. +#[derive(Copy, Clone, Debug, PartialEq, Eq, Enum)] +pub(crate) enum MessageType { + /// Messages sent via synchronous exchange. + Synchronous( + /// Message size. + MessageSize, + ), + /// Messages sent via streaming exchange. + Streaming, +} + pub struct ExchangeClient { tx: ByteBoundedSender, } @@ -342,6 +382,10 @@ impl ExchangeClient { }) .expect("remote exchange failed") } + + pub async fn wait(&self) { + self.tx.bound.wait().await; + } } /// Uniquely identifies an `Exchange` or `ShardedAccumulator` within a circuit. @@ -484,9 +528,11 @@ pub struct ExchangeClients { /// tries to send data to one. listener: OnceCell>, - /// Maps from a range of worker IDs to the RPC client used to contact those + /// Maps from a range of worker IDs to the RPC clients used to contact those /// workers. Only worker IDs for remote workers appear in the map. - clients: Vec<(Host, OnceCell)>, + /// + /// We use one RPC client per [MessageType] per [Host]. + clients: Vec<(Host, EnumMap>)>, } impl ExchangeClients { @@ -509,14 +555,14 @@ impl ExchangeClients { clients: runtime .layout() .other_hosts() - .map(|host| (host.clone(), OnceCell::new())) + .map(|host| (host.clone(), Default::default())) .collect(), } } /// Returns a client for `worker`, which must be a remote worker ID, first /// establishing a connection if there isn't one yet. - pub async fn connect(&self, worker: usize) -> &ExchangeClient { + pub async fn connect(&self, worker: usize, message_type: MessageType) -> &ExchangeClient { self.listener .get_or_init(|| async { if let Some(runtime) = self.runtime.upgrade() @@ -540,9 +586,20 @@ impl ExchangeClients { .iter() .find(|(host, _client)| host.workers.contains(&worker)) .unwrap(); - cell.get_or_init(|| ExchangeClient::new(host.address, &host.workers)) + cell[message_type] + .get_or_init(|| ExchangeClient::new(host.address, &host.workers)) .await } + + pub async fn wait(&self) { + for (_, clients) in &self.clients { + for client in clients.values() { + if let Some(client) = client.get() { + client.wait().await; + } + } + } + } } struct CallbackInner { @@ -703,6 +760,9 @@ pub(crate) struct Exchange { /// The number of bytes serialized. deserialized_bytes: AtomicUsize, + + /// When the exchange is active. + activity: ExchangeActivity, } // Stop Rust from complaining about unused field. @@ -763,6 +823,7 @@ where clients: Arc, exchange_id: ExchangeId, directory: &ExchangeDirectory, + activity: ExchangeActivity, ) -> Arc { let npeers = Runtime::num_workers(); let mailboxes: Vec>>> = @@ -788,6 +849,7 @@ where mailboxes, deserialization_usecs: AtomicU64::new(0), deserialized_bytes: AtomicUsize::new(0), + activity, }); directory.insert(exchange_id, exchange.clone()); @@ -806,39 +868,6 @@ where sender * self.npeers + receiver } - /// True if all `receiver`'s incoming mailboxes contain data. - /// - /// Once this function returns true, a subsequent `try_receive_all` - /// operation is guaranteed for `receiver`. - fn ready_to_receive(&self, receiver: usize) -> bool { - debug_assert!(receiver < self.npeers); - self.receiver_counters[receiver].load(Ordering::Acquire) == self.npeers - } - - /// Register callback to be invoked whenever the `ready_to_receive` - /// condition becomes true. - /// - /// The callback can be setup at most once (e.g., when a scheduler attaches - /// to the circuit) and cannot be unregistered. Notifications delivered - /// before the callback is registered are lost. The client should call - /// `ready_to_receive` after installing the callback to check - /// the status. - /// - /// After the callback has been registered, notifications are delivered with - /// at-least-once semantics: a notification is generated whenever the - /// status changes from not ready to ready, but spurious notifications - /// can occur occasionally. The user must check the status explicitly - /// by calling `ready_to_receive` or be prepared that `receive_all` - /// can fail. - pub(crate) fn register_receiver_callback(&self, receiver: usize, cb: F) - where - F: Fn() + Send + Sync + 'static, - { - debug_assert!(receiver < self.npeers); - - self.receiver_callbacks[receiver].set_callback(cb); - } - /// Locks and returns the mailbox for the sender/receiver pair. fn mailbox(&self, sender: usize, receiver: usize) -> MutexGuard<'_, Option>> { self.mailboxes[self.mailbox_index(sender, receiver)] @@ -849,7 +878,11 @@ where /// Create a new `Exchange` instance if an instance with the same id /// (created by another thread) does not yet exist within `runtime`. /// The number of peers will be set to `runtime.num_workers()`. - pub(crate) fn with_runtime(runtime: &Runtime, exchange_id: ExchangeId) -> Arc { + pub(crate) fn with_runtime( + runtime: &Runtime, + exchange_id: ExchangeId, + activity: ExchangeActivity, + ) -> Arc { // It's tempting to move the following calls to create the // `ExchangeDirectory` and `ExchangeClients` into `Exchange::new`, but // don't do it: all three of these access `runtime.local_store` and @@ -859,7 +892,7 @@ where runtime .local_store() .entry(ExchangeCacheId::new(exchange_id)) - .or_insert_with(|| Exchange::new(runtime, clients, exchange_id, &directory)) + .or_insert_with(|| Exchange::new(runtime, clients, exchange_id, &directory, activity)) .value() .clone() } @@ -987,16 +1020,19 @@ where }) .collect_vec(); + let message_type = MessageType::Synchronous(MessageSize::from_bytes( + items.iter().map(|fbuf| fbuf.len()).sum(), + )); + // We discard the return value that could allow us to wait // for the channel tx buffer to drain, because exchange is // synchronous, meaning that it will drain before we send // the next message. - let _ = self.clients.connect(receivers.start).await.send( - global_node_id.clone(), - self.exchange_id, - sender, - items, - ); + let _ = self + .clients + .connect(receivers.start, message_type) + .await + .send(global_node_id.clone(), self.exchange_id, sender, items); } } } @@ -1004,7 +1040,16 @@ where /// Read all incoming messages for this worker, waiting for data to arrive /// as needed. - pub(crate) async fn receive_all(&self, deserialize: D) -> Vec + /// + /// When the data is ready, but before reading it, this method swaps + /// `start_wait_usecs` with 0 and returns the old value along with the data. + /// This allows the caller to obtain the waiting time incurred just after it + /// became ready. + pub(crate) async fn receive_all( + &self, + deserialize: D, + start_wait_usecs: Option<&AtomicU64>, + ) -> (Vec, Option) where D: Fn(AlignedVec) -> T, { @@ -1024,6 +1069,11 @@ where } } + let start_wait_usecs = start_wait_usecs.and_then(|v| { + let start_wait_usecs = v.swap(0, Ordering::Acquire); + (start_wait_usecs != 0).then_some(start_wait_usecs) + }); + let mut data = Vec::with_capacity(self.npeers); for sender in 0..self.npeers { let mailbox = self.mailbox(sender, receiver).take().unwrap(); @@ -1035,7 +1085,7 @@ where } } - data + (data, start_wait_usecs) } } @@ -1062,6 +1112,19 @@ where } } +#[derive(Copy, Clone, Debug, PartialEq, Eq)] +enum Phase { + Active, + Flush, + Commit, +} + +impl Phase { + fn is_inactive(&self, activity: ExchangeActivity) -> bool { + *self == Phase::Commit && activity == ExchangeActivity::InputOnly + } +} + /// Operator that partitions incoming data across all workers. /// /// This operator works in tandem with [`ExchangeReceiver`], which reassembles @@ -1100,13 +1163,6 @@ where /// ExchangeSender ExchangeReceiver /// ``` /// -/// `ExchangeSender` is an asynchronous operator., i.e., -/// [`ExchangeSender::is_async`] returns `true`. It becomes schedulable -/// ([`ExchangeSender::ready`] returns `true`) once all peers have retrieved -/// values written by the operator in the previous clock cycle. The scheduler -/// should use [`ExchangeSender::register_ready_callback`] to get notified when -/// the operator becomes schedulable. -/// /// `ExchangeSender` doesn't have a public constructor and must be instantiated /// using the [`new_exchange_operators`] function, which creates an /// [`ExchangeSender`]/[`ExchangeReceiver`] pair of operators and connects them @@ -1182,7 +1238,7 @@ where /// use dbsp::{ /// operator::{communication::new_exchange_operators, Generator}, /// circuit::{WorkerLocation, WorkerLocations}, -/// operator::communication::Mailbox, +/// operator::communication::{ExchangeActivity, Mailbox}, /// Circuit, RootCircuit, Runtime, /// storage::file::to_bytes_dyn, /// trace::aligned_deserialize, @@ -1218,6 +1274,7 @@ where /// }, /// |data| aligned_deserialize(&data[..]),/// // Reassemble received values into a vector. /// |v: &mut Vec, n| v.push(n), +/// ExchangeActivity::AllSteps, /// ).unwrap(); /// /// // Add exchange operators to the circuit. @@ -1259,7 +1316,7 @@ where // Input batch sizes. input_batch_stats: BatchSizeStats, - flushed: bool, + phase: Phase, // The instant when the sender produced its outputs, and the // receiver starts waiting for all other workers to produce their @@ -1286,7 +1343,7 @@ where outputs: Vec::with_capacity(Runtime::num_workers()), exchange, input_batch_stats: BatchSizeStats::new(), - flushed: false, + phase: Phase::Active, start_wait_usecs, phantom: PhantomData, } @@ -1324,8 +1381,12 @@ where true } + fn start_transaction(&mut self) { + self.phase = Phase::Active; + } + fn flush(&mut self) { - self.flushed = true; + self.phase = Phase::Flush; } } @@ -1340,6 +1401,16 @@ where } async fn eval_owned(&mut self, input: D) { + if self.phase.is_inactive(self.exchange.activity) { + return; + }; + let flushed = if self.phase == Phase::Flush { + self.phase = Phase::Commit; + true + } else { + false + }; + self.input_batch_stats.add_batch(input.num_entries_deep()); debug_assert!(self.ready()); @@ -1350,16 +1421,14 @@ where let data = self.outputs.drain(..).map(|mailbox| match mailbox { Mailbox::Tx(mut data) => { - data.push(self.flushed as u8); + data.push(flushed as u8); Mailbox::Tx(data) } Mailbox::Rx(_) => unreachable!(), - Mailbox::Plain(item) => Mailbox::Plain((item, self.flushed)), + Mailbox::Plain(item) => Mailbox::Plain((item, flushed)), }); self.exchange.send_all(&self.global_node_id, data).await; - - self.flushed = false; } fn input_preference(&self) -> OwnershipPreference { @@ -1375,18 +1444,10 @@ where /// peer. /// /// See [`ExchangeSender`] documentation for details. -/// -/// `ExchangeReceiver` is an asynchronous operator., i.e., -/// [`ExchangeReceiver::is_async`] returns `true`. It becomes schedulable -/// ([`ExchangeReceiver::ready`] returns `true`) once all peers have sent values -/// for this worker in the current clock cycle. The scheduler should use -/// [`ExchangeReceiver::register_ready_callback`] to get notified when the -/// operator becomes schedulable. pub struct ExchangeReceiver where T: Send + 'static + Clone, { - worker_index: usize, location: OperatorLocation, init: IF, deserialize: D, @@ -1396,6 +1457,7 @@ where flush_complete: bool, start_wait_usecs: Arc, total_wait_time: Arc, + phase: Phase, // Output batch sizes. output_batch_stats: BatchSizeStats, @@ -1406,7 +1468,6 @@ where T: Send + 'static + Clone + Debug, { pub(crate) fn new( - worker_index: usize, location: OperatorLocation, exchange: Arc>, init: IF, @@ -1414,10 +1475,7 @@ where deserialize: D, combine: L, ) -> Self { - debug_assert!(worker_index < Runtime::num_workers()); - Self { - worker_index, location, init, combine, @@ -1428,12 +1486,9 @@ where output_batch_stats: BatchSizeStats::new(), start_wait_usecs, total_wait_time: Arc::new(AtomicU64::new(0)), + phase: Phase::Active, } } - - fn is_ready(&self) -> bool { - self.exchange.ready_to_receive(self.worker_index) - } } impl Operator for ExchangeReceiver @@ -1464,50 +1519,12 @@ where }); } - fn is_async(&self) -> bool { + fn fixedpoint(&self, _scope: Scope) -> bool { true } - fn register_ready_callback(&mut self, cb: F) - where - F: Fn() + Send + Sync + 'static, - { - let start_wait_usecs = self.start_wait_usecs.clone(); - let total_wait_time = self.total_wait_time.clone(); - let exchange = self.exchange.clone(); - let worker_index = self.worker_index; - - let cb = move || { - if exchange.ready_to_receive(worker_index) { - // The callback can be invoked multiple times per step. - // Reset start_wait_usecs to 0 to make sure we don't double-count. - let start = start_wait_usecs.swap(0, Ordering::Acquire); - if start != 0 { - let end = current_time_usecs(); - if end > start { - let wait_time_usecs = end - start; - // if worker_index == 0 { - // info!( - // "{worker_index}: {} +{wait_time_usecs}", - // exchange.exchange_id() - // ); - // } - total_wait_time.fetch_add(wait_time_usecs, Ordering::AcqRel); - } - } - } - cb() - }; - self.exchange - .register_receiver_callback(self.worker_index, cb) - } - - fn ready(&self) -> bool { - self.is_ready() - } - - fn fixedpoint(&self, _scope: Scope) -> bool { - true + fn start_transaction(&mut self) { + self.phase = Phase::Active; } fn flush(&mut self) { @@ -1542,7 +1559,10 @@ where D: Fn(AlignedVec) -> T + Send + Sync + 'static, { async fn eval(&mut self) -> O { - debug_assert!(self.ready()); + if self.phase.is_inactive(self.exchange.activity) { + return (self.init)(); + } + let deserialize = |mut vec: AlignedVec| { let flushed = pop_flushed(&mut vec); let value = (self.deserialize)(vec); @@ -1550,7 +1570,16 @@ where }; let mut combined = (self.init)(); - let res = self.exchange.receive_all(deserialize).await; + let (res, start_wait_usecs) = self + .exchange + .receive_all(deserialize, Some(&self.start_wait_usecs)) + .await; + if let Some(start_wait_usecs) = start_wait_usecs { + self.total_wait_time.fetch_add( + current_time_usecs().saturating_sub(start_wait_usecs), + Ordering::Release, + ); + } for (data, flushed) in res { if flushed { self.flush_count += 1; @@ -1566,6 +1595,7 @@ where self.flush_complete = true; self.flush_count = 0; + self.phase = Phase::Commit; } self.output_batch_stats @@ -1588,6 +1618,28 @@ impl TypedMapKey for DirectoryId { type Value = ExchangeDirectory; } +/// The microsteps during which an exchange is active. +#[derive(Copy, Clone, Debug, PartialEq, Eq)] +pub enum ExchangeActivity { + /// The exchange is active in every microstep. + /// + /// This includes pre-commit and commit microstep. + AllSteps, + + /// The exchange is active only during pre-commit microsteps. + /// + /// This allows for optimizations for exchanges used for sharding data from + /// input operators, which don't exchange any data during commit. + /// + /// # Limitation + /// + /// The current implementation only works for operators that flush in the + /// same (micro)step in every worker. This is true for input operators, + /// which flush as soon as the transaction starts committing, but it is not + /// necessarily true for other operators. + InputOnly, +} + /// Create an [`ExchangeSender`]/[`ExchangeReceiver`] operator pair. /// /// See [`ExchangeSender`] documentation for details and example usage. @@ -1619,6 +1671,7 @@ pub fn new_exchange_operators( partition: PL, deserialize: D, combine: CL, + activity: ExchangeActivity, ) -> Option<(ExchangeSender, ExchangeReceiver)> where TO: Clone, @@ -1632,11 +1685,10 @@ where return None; } let runtime = Runtime::runtime().unwrap(); - let worker_index = Runtime::worker_index(); let exchange_id = runtime.sequence_next().try_into().unwrap(); let start_wait_usecs = Arc::new(AtomicU64::new(0)); - let exchange = Exchange::with_runtime(&runtime, exchange_id); + let exchange = Exchange::with_runtime(&runtime, exchange_id, activity); let sender = ExchangeSender::new( location, exchange.clone(), @@ -1644,7 +1696,6 @@ where partition, ); let receiver = ExchangeReceiver::new( - worker_index, location, exchange, init, @@ -1670,7 +1721,7 @@ mod tests { }, operator::{ Generator, - communication::{Mailbox, new_exchange_operators}, + communication::{ExchangeActivity, Mailbox, new_exchange_operators}, }, storage::file::{to_bytes, to_bytes_dyn}, trace::aligned_deserialize, @@ -1691,7 +1742,8 @@ mod tests { // value `(sender, n)` to each receiver, where `sender` is the sender's // worker number in round `n`. fn circuit() { - let exchange = Exchange::with_runtime(&Runtime::runtime().unwrap(), 0); + let exchange = + Exchange::with_runtime(&Runtime::runtime().unwrap(), 0, ExchangeActivity::AllSteps); TOKIO.block_on(async { let sender = Runtime::worker_index(); let n_workers = Runtime::num_workers(); @@ -1703,8 +1755,8 @@ mod tests { }) .await; - let received = exchange - .receive_all(|data| aligned_deserialize(&data[..])) + let (received, _) = exchange + .receive_all(|data| aligned_deserialize(&data[..]), None) .await; let expected = (0..n_workers).map(|worker| (worker, round)).collect_vec(); @@ -1821,6 +1873,7 @@ mod tests { }, |data| aligned_deserialize(&data[..]), |v: &mut Vec, n| v.push(n), + ExchangeActivity::AllSteps, ) .unwrap(); diff --git a/crates/dbsp/src/operator/dynamic.rs b/crates/dbsp/src/operator/dynamic.rs index 5bf472b5f1e..f35b4db6108 100644 --- a/crates/dbsp/src/operator/dynamic.rs +++ b/crates/dbsp/src/operator/dynamic.rs @@ -34,6 +34,7 @@ pub mod time_series; pub mod trace; pub(crate) mod upsert; +pub use communication::shard::Sharder; pub(crate) use communication::shard::shard_batch; /// The "standard" indexed Z-set type used by monomorphic diff --git a/crates/dbsp/src/operator/dynamic/balance/accumulate_trace_balanced.rs b/crates/dbsp/src/operator/dynamic/balance/accumulate_trace_balanced.rs index 06b87e61ab6..ec2ca8fb3a4 100644 --- a/crates/dbsp/src/operator/dynamic/balance/accumulate_trace_balanced.rs +++ b/crates/dbsp/src/operator/dynamic/balance/accumulate_trace_balanced.rs @@ -20,7 +20,7 @@ use crate::{ operator::{ Z1, async_stream_operators::{StreamingTernarySinkOperator, StreamingTernarySinkWrapper}, - communication::{Exchange, ExchangeReceiver}, + communication::{Exchange, ExchangeActivity, ExchangeReceiver}, dynamic::{ accumulate_trace::{AccumulateBoundsId, AccumulateTraceAppend, AccumulateZ1Trace}, balance::{ @@ -145,13 +145,12 @@ where // during the current transaction. It will be set to `true` exactly once per transaction. // Once all senders are flushed, the receiver can report itself as flushed too. let exchange: Arc> = - Exchange::with_runtime(&runtime, exchange_id); + Exchange::with_runtime(&runtime, exchange_id, ExchangeActivity::AllSteps); // Exchange receiver let batch_factories_clone = batch_factories.clone(); let receiver = circuit.add_source(ExchangeReceiver::new( - worker_index, Some(location), exchange.clone(), || Vec::new(), diff --git a/crates/dbsp/src/operator/dynamic/communication/gather.rs b/crates/dbsp/src/operator/dynamic/communication/gather.rs index ce2c962ae82..40eb252178b 100644 --- a/crates/dbsp/src/operator/dynamic/communication/gather.rs +++ b/crates/dbsp/src/operator/dynamic/communication/gather.rs @@ -8,7 +8,7 @@ use crate::{ runtime::{WorkerLocation, WorkerLocations}, }, circuit_cache_key, - operator::communication::{Mailbox, new_exchange_operators}, + operator::communication::{ExchangeActivity, Mailbox, new_exchange_operators}, trace::{Batch, deserialize_indexed_wset, merge_batches, serialize_indexed_wset}, }; use arc_swap::ArcSwap; @@ -159,6 +159,7 @@ where }, move |data| deserialize_indexed_wset(&factories_clone, &data), |batches: &mut Vec, batch: B| batches.push(batch), + ExchangeActivity::AllSteps, ) .unwrap(); diff --git a/crates/dbsp/src/operator/dynamic/communication/shard.rs b/crates/dbsp/src/operator/dynamic/communication/shard.rs index 44f43d59c54..f1e212d6f71 100644 --- a/crates/dbsp/src/operator/dynamic/communication/shard.rs +++ b/crates/dbsp/src/operator/dynamic/communication/shard.rs @@ -16,7 +16,7 @@ use crate::{ }, circuit_cache_key, dynamic::{Data, DataTrait, DynPair, DynPairs, Factory}, - operator::communication::{Mailbox, new_exchange_operators}, + operator::communication::{ExchangeActivity, Mailbox, new_exchange_operators}, storage::file::SerializerInner, trace::{ Batch, BatchReader, Builder, IndexedWSetSerializer, deserialize_indexed_wset, merge_batches, @@ -35,68 +35,59 @@ fn all_workers() -> Range { 0..Runtime::num_workers() } -impl Stream +/// Builder for sharding a batch or pairs. +pub struct Sharder<'a, C, B> { + stream: &'a Stream, + workers: Range, + activity: ExchangeActivity, +} + +impl<'a, C, B> Sharder<'a, C, B> where C: Circuit, - IB: BatchReader