From 2f498bb426cc33d0be2960d78dfe570d524f5ffc Mon Sep 17 00:00:00 2001 From: Mihai Budiu Date: Tue, 14 Jul 2026 18:11:48 -0700 Subject: [PATCH 01/23] [SLT] Skip test that produces integer overflow (fails on Postgres as well) Signed-off-by: Mihai Budiu --- sql-to-dbsp-compiler/slt/skip.txt | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/sql-to-dbsp-compiler/slt/skip.txt b/sql-to-dbsp-compiler/slt/skip.txt index 5ea8a3c96fe..f670804cf01 100644 --- a/sql-to-dbsp-compiler/slt/skip.txt +++ b/sql-to-dbsp-compiler/slt/skip.txt @@ -33,4 +33,9 @@ SELECT ALL + NULLIF ( - COUNT ( * ), + 47 * - 93 * MAX ( 45 ) * - 63 * - 56 / - SELECT ALL - col0 AS col2, - 22 / + - 41 * + col0 AS col2 FROM tab2 AS cor0 GROUP BY col0 HAVING + SUM ( ALL - - col2 ) IS NOT NULL // test/random/groupby/slt_good_12.test: ambiguous query; Postgres rejects this query as well SELECT + COUNT ( * ) AS col1, - 0 * + - 69 * + ( + 53 ) AS col1 FROM tab0 WHERE NOT NULL IS NOT NULL GROUP BY col2 HAVING ( COUNT ( col1 ) ) IS NULL - +// test/index/random/1000/slt_good_6.test: test 8506-8510 (identical!): arithmetic overflow; reproduced on Postgres +SELECT ALL + - CAST ( NULL AS INTEGER ) + + col3 FROM tab0 AS cor0 WHERE col3 * + col0 * + - CAST ( col1 AS INTEGER ) * col4 + - + col4 IS NULL +SELECT ALL + - CAST ( NULL AS INTEGER ) + + col3 FROM tab1 AS cor0 WHERE col3 * + col0 * + - CAST ( col1 AS INTEGER ) * col4 + - + col4 IS NULL +SELECT ALL + - CAST ( NULL AS INTEGER ) + + col3 FROM tab2 AS cor0 WHERE col3 * + col0 * + - CAST ( col1 AS INTEGER ) * col4 + - + col4 IS NULL +SELECT ALL + - CAST ( NULL AS INTEGER ) + + col3 FROM tab3 AS cor0 WHERE col3 * + col0 * + - CAST ( col1 AS INTEGER ) * col4 + - + col4 IS NULL +SELECT ALL + - CAST ( NULL AS INTEGER ) + + col3 FROM tab4 AS cor0 WHERE col3 * + col0 * + - CAST ( col1 AS INTEGER ) * col4 + - + col4 IS NULL From 3c8a2b708a09df8b3db97947ed0672ef4f74e1f0 Mon Sep 17 00:00:00 2001 From: Leonid Ryzhyk Date: Tue, 14 Jul 2026 23:19:31 -0700 Subject: [PATCH 02/23] [Connectors] Iceberg connector reads only columns declared in the SQL table Replace the `select *` snapshot queries with an explicit column list: the intersection of the Iceberg snapshot schema and the SQL table declaration, mirroring the Delta connector's `used_column_list`. Undeclared table columns previously wasted IO and could fail deserialization for column types Feldera does not support. The connector also honors the table-level `skip_unused_columns` property: a declared-but-unused column is skipped when it is nullable or has a default and `snapshot_filter` does not reference it. The property check lives in `Relation::skip_unused_columns`, used by both the Delta and Iceberg connectors. Share `ColumnNameSet` and `quote_sql_identifier` between the Delta and Iceberg connectors by moving them to feldera-adapterlib. Signed-off-by: Leonid Ryzhyk --- crates/adapterlib/src/utils/datafusion.rs | 29 ++ .../src/integrated/delta_table/input.rs | 40 +-- crates/adapters/src/test.rs | 50 ++- crates/adapters/src/test/data.rs | 51 +++ crates/adapters/src/test/iceberg.rs | 138 +++++++- crates/feldera-types/src/program_schema.rs | 46 +++ crates/iceberg/src/input.rs | 309 ++++++++++++++++-- .../iceberg/src/test/create_test_table_s3.py | 62 ++-- .../docs/connectors/sources/iceberg.md | 7 + docs.feldera.com/docs/sql/grammar.md | 4 +- 10 files changed, 639 insertions(+), 97 deletions(-) diff --git a/crates/adapterlib/src/utils/datafusion.rs b/crates/adapterlib/src/utils/datafusion.rs index c828838d8d3..ab3643d5546 100644 --- a/crates/adapterlib/src/utils/datafusion.rs +++ b/crates/adapterlib/src/utils/datafusion.rs @@ -364,6 +364,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 { 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/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..b86e0f5d632 100644 --- a/crates/adapters/src/test/data.rs +++ b/crates/adapters/src/test/data.rs @@ -1003,6 +1003,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..c93734a5d00 100644 --- a/crates/adapters/src/test/iceberg.rs +++ b/crates/adapters/src/test/iceberg.rs @@ -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}"); @@ -222,17 +251,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 +268,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 +294,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 +333,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 +400,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()), ( @@ -354,6 +458,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 +486,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()), diff --git a/crates/feldera-types/src/program_schema.rs b/crates/feldera-types/src/program_schema.rs index e3b0a676134..8789ee1bbae 100644 --- a/crates/feldera-types/src/program_schema.rs +++ b/crates/feldera-types/src/program_schema.rs @@ -258,6 +258,13 @@ impl Relation { self.properties.get(name).map(|p| p.value.as_str()) } + /// `true` if the table declaration sets the `skip_unused_columns` + /// property, which asks connectors not to read declared-but-unused + /// columns. + pub fn skip_unused_columns(&self) -> bool { + self.get_property("skip_unused_columns") == Some("true") + } + pub fn with_primary_key<'a>( mut self, primary_key: impl IntoIterator, @@ -1218,6 +1225,45 @@ mod tests { } } + /// Only the exact property value "true" enables `skip_unused_columns`; + /// connectors have always compared the value verbatim. + #[test] + fn relation_skip_unused_columns_property() { + use super::{PropertyValue, Relation, SourcePosition}; + use std::collections::BTreeMap; + + let zero = SourcePosition { + start_line_number: 0, + start_column: 0, + end_line_number: 0, + end_column: 0, + }; + let relation = |value: Option<&str>| { + let mut properties = BTreeMap::new(); + if let Some(value) = value { + properties.insert( + "skip_unused_columns".to_string(), + PropertyValue { + value: value.to_string(), + key_position: zero, + value_position: zero, + }, + ); + } + Relation::new( + SqlIdentifier::new("t", false), + Vec::new(), + false, + properties, + ) + }; + + assert!(relation(Some("true")).skip_unused_columns()); + assert!(!relation(Some("false")).skip_unused_columns()); + assert!(!relation(Some("TRUE")).skip_unused_columns()); + assert!(!relation(None).skip_unused_columns()); + } + #[test] fn serde_sql_type() { for (sql_str_base, expected_value) in [ diff --git a/crates/iceberg/src/input.rs b/crates/iceberg/src/input.rs index fbfab988835..b2bfe46f292 100644 --- a/crates/iceberg/src/input.rs +++ b/crates/iceberg/src/input.rs @@ -1,6 +1,8 @@ use crate::iceberg_input_serde_config; use anyhow::{anyhow, bail, Error as AnyError, Result as AnyResult}; use chrono::{DateTime, Utc}; +use datafusion::arrow::datatypes::Schema as ArrowSchema; +use datafusion::catalog::TableProvider; use datafusion::prelude::{DataFrame, SQLOptions, SessionContext}; use dbsp::circuit::tokio::TOKIO; use feldera_adapterlib::{ @@ -12,14 +14,16 @@ use feldera_adapterlib::{ IntegratedInputEndpoint, NonFtInputReaderCommand, }, utils::datafusion::{ - array_to_string, create_session_context, execute_query_collect, execute_singleton_query, + array_to_string, columns_referenced_by_expression, create_session_context, + execute_query_collect, execute_singleton_query, quote_sql_identifier, timestamp_to_sql_expression, validate_sql_expression, validate_timestamp_column, + ColumnNameSet, }, PipelineState, }; use feldera_types::{ config::{FtModel, PipelineConfig}, - program_schema::Relation, + program_schema::{Field, Relation}, transport::iceberg::{IcebergCatalogType, IcebergReaderConfig}, }; use futures_util::StreamExt; @@ -44,7 +48,7 @@ use iceberg_catalog_s3tables::{ use iceberg_datafusion::IcebergStaticTableProvider; use iceberg_storage_opendal::OpenDalResolvingStorageFactory; use log::{debug, info, trace}; -use std::{sync::Arc, thread}; +use std::{collections::BTreeSet, sync::Arc, thread}; use tokio::{ select, sync::{ @@ -70,6 +74,89 @@ const S3TABLES_PROP_SESSION_TOKEN: &str = "aws_session_token"; const S3TABLES_PROP_PROFILE_NAME: &str = "profile_name"; const S3TABLES_PROP_REGION_NAME: &str = "region_name"; +/// SQL columns named by the connector's own `snapshot_filter` expression, +/// case-folded for matching. These are kept even when marked unused, so the +/// expression is guaranteed to resolve them (mirrors the Delta connector; for +/// the snapshot queries issued here a `where` clause could also reference +/// unprojected columns). +/// +/// The filter is validated before columns are computed, so it parses here; a +/// parse error is therefore unreachable and contributes no columns rather +/// than failing the connector a second time. +fn config_referenced_columns(config: &IcebergReaderConfig) -> ColumnNameSet { + let mut columns = BTreeSet::new(); + if let Some(filter) = &config.snapshot_filter { + columns.extend(columns_referenced_by_expression(filter).unwrap_or_default()); + } + ColumnNameSet::from_names(columns) +} + +/// True if a column's *shape* allows omitting it: no user-visible result +/// depends on it (`unused`), and omitting it lets us substitute NULL or its +/// default value (it is nullable or has a default). +/// +/// This is the shape-only rule; [`can_skip_column`] adds the +/// filter-reference check before a column is actually skipped. +fn is_unused_and_omittable(field: &Field) -> bool { + field.unused && (field.columntype.nullable || field.default.is_some()) +} + +/// True if a column may actually be skipped: its shape permits omitting it +/// ([`is_unused_and_omittable`]) *and* the `snapshot_filter` expression does +/// not reference it. +fn can_skip_column(field: &Field, config_referenced: &ColumnNameSet) -> bool { + is_unused_and_omittable(field) && !config_referenced.contains(&field.name.name()) +} + +/// SQL columns the connector reads, matched case-insensitively against +/// Iceberg column names. When the SQL table declaration sets the +/// `skip_unused_columns` property, skippable unused columns are removed. +fn used_sql_columns(schema: &Relation, config: &IcebergReaderConfig) -> ColumnNameSet { + let skip = schema.skip_unused_columns(); + let config_referenced = config_referenced_columns(config); + ColumnNameSet::from_names( + schema + .fields + .iter() + .filter(|f| !skip || !can_skip_column(f, &config_referenced)) + .map(|f| f.name.name()), + ) +} + +/// Compute the subset of columns in the Iceberg table snapshot schema that +/// occur in the SQL table declaration (minus columns removed by +/// [`used_sql_columns`]), preserving the table schema's column order. +/// +/// Snapshot queries select this subset instead of `*`, so the connector never +/// reads table columns the pipeline does not ingest. Besides being wasteful, +/// reading such columns can fail when their Arrow type is not supported by +/// the Feldera deserializer. +/// +/// Iceberg schemas carry no case-sensitivity information, so column names are +/// matched in lowercased form (see [`ColumnNameSet`]). +fn used_columns( + table_schema: &ArrowSchema, + schema: &Relation, + config: &IcebergReaderConfig, +) -> Vec { + let used = used_sql_columns(schema, config); + table_schema + .fields() + .iter() + .filter(|f| used.contains(f.name())) + .map(|f| f.name().to_string()) + .collect() +} + +/// Quoted, comma-separated column list for `select {} from snapshot` queries. +fn used_column_list(columns: &[String]) -> String { + columns + .iter() + .map(quote_sql_identifier) + .collect::>() + .join(", ") +} + enum SnapshotDescr { /// Open the latest snapshot (default) Latest, @@ -304,20 +391,27 @@ impl IcebergInputEndpointInner { } } - /// Load the entire table snapshot as a single "select * where " query. + /// Load the entire table snapshot as a single + /// "select where " query. async fn read_unordered_snapshot( &self, + used_columns: &[String], input_stream: &mut dyn ArrowStream, receiver: &mut Receiver, ) { - // Execute the snapshot query; push snapshot data to the circuit. - info!("iceberg {}: reading initial snapshot", &self.endpoint_name,); + let column_names = used_column_list(used_columns); - let mut snapshot_query = "select * from snapshot".to_string(); + let mut snapshot_query = format!("select {column_names} from snapshot"); if let Some(filter) = &self.config.snapshot_filter { snapshot_query = format!("{snapshot_query} where {filter}"); } + // Execute the snapshot query; push snapshot data to the circuit. + info!( + "iceberg {}: reading initial snapshot: {snapshot_query}", + &self.endpoint_name, + ); + self.execute_snapshot_query(&snapshot_query, "initial snapshot", input_stream, receiver) .await; @@ -330,17 +424,19 @@ impl IcebergInputEndpointInner { async fn read_ordered_snapshot( &self, + used_columns: &[String], input_stream: &mut dyn ArrowStream, schema: &Relation, receiver: &mut Receiver, ) { - self.read_ordered_snapshot_inner(input_stream, schema, receiver) + self.read_ordered_snapshot_inner(used_columns, input_stream, schema, receiver) .await .unwrap_or_else(|e| self.consumer.error(true, e, None)); } async fn read_ordered_snapshot_inner( &self, + used_columns: &[String], input_stream: &mut dyn ArrowStream, schema: &Relation, receiver: &mut Receiver, @@ -409,6 +505,8 @@ impl IcebergInputEndpointInner { let min = timestamp_to_sql_expression(×tamp_field.columntype, &min); let max = timestamp_to_sql_expression(×tamp_field.columntype, &max); + let column_names = used_column_list(used_columns); + let mut start = min.clone(); let mut done = "false".to_string(); @@ -423,7 +521,7 @@ impl IcebergInputEndpointInner { // Query the table for the range. let mut range_query = - format!("select * from snapshot where {timestamp_column} >= {start} and {timestamp_column} < {end}"); + format!("select {column_names} from snapshot where {timestamp_column} >= {start} and {timestamp_column} < {end}"); if let Some(filter) = &self.config.snapshot_filter { range_query = format!("{range_query} and {filter}"); } @@ -460,9 +558,12 @@ impl IcebergInputEndpointInner { let table = Arc::new(table); - if let Err(e) = self.prepare_snapshot_query(&table, &schema).await { - let _ = init_status_sender.send(Err(e)).await; - return; + let used_columns = match self.prepare_snapshot_query(&table, &schema).await { + Err(e) => { + let _ = init_status_sender.send(Err(e)).await; + return; + } + Ok(used_columns) => used_columns, }; // Code before this point is part of endpoint initialization. @@ -472,12 +573,17 @@ impl IcebergInputEndpointInner { if self.config.snapshot() && self.config.timestamp_column.is_none() { // Read snapshot chunk-by-chunk. - self.read_unordered_snapshot(input_stream.as_mut(), &mut receiver) + self.read_unordered_snapshot(&used_columns, input_stream.as_mut(), &mut receiver) .await; } else if self.config.snapshot() { // Read the entire snapshot in one query. - self.read_ordered_snapshot(input_stream.as_mut(), &schema, &mut receiver) - .await; + self.read_ordered_snapshot( + &used_columns, + input_stream.as_mut(), + &schema, + &mut receiver, + ) + .await; }; self.consumer.eoi(); @@ -819,15 +925,23 @@ impl IcebergInputEndpointInner { /// /// * register snapshot as a datafusion table /// * validate snapshot config: filter condition and timestamp column + /// + /// Returns the columns snapshot queries must select (see [`used_columns`]); + /// empty when the configuration requires no snapshot. async fn prepare_snapshot_query( &self, table: &IcebergTable, schema: &Relation, - ) -> Result<(), ControllerError> { + ) -> Result, ControllerError> { if !self.config.snapshot() { - return Ok(()); + return Ok(Vec::new()); } + // Validate the filter before `config_referenced_columns` extracts + // column names from it, so an invalid filter fails with a parse error + // rather than being silently ignored during column selection. + self.validate_snapshot_filter()?; + trace!( "iceberg {}: registering table with Datafusion", &self.endpoint_name, @@ -870,6 +984,14 @@ impl IcebergInputEndpointInner { ) })?; + let used_columns = used_columns(provider.schema().as_ref(), schema, &self.config); + if used_columns.is_empty() { + return Err(ControllerError::invalid_transport_configuration( + &self.endpoint_name, + "the connector would read no columns: none of the columns declared in the SQL table exist in the Iceberg table (columns skipped via the 'skip_unused_columns' table property don't count); check that the connector points to the correct table", + )); + } + self.datafusion .register_table("snapshot", Arc::new(provider)) .map_err(|e| { @@ -880,8 +1002,6 @@ impl IcebergInputEndpointInner { ) })?; - self.validate_snapshot_filter()?; - if let Some(timestamp_column) = &self.config.timestamp_column { validate_timestamp_column( &self.endpoint_name, @@ -893,7 +1013,7 @@ impl IcebergInputEndpointInner { .await?; }; - Ok(()) + Ok(used_columns) } /// Execute a SQL query to load a complete or partial snapshot of the table. @@ -1013,10 +1133,159 @@ async fn wait_running(receiver: &mut Receiver) { #[cfg(test)] mod tests { use super::*; + use datafusion::arrow::datatypes::{DataType, Field as ArrowField}; + use feldera_types::program_schema::{ColumnType, PropertyValue, SourcePosition, SqlIdentifier}; + use serde_json::json; + use std::collections::BTreeMap; #[test] fn storage_factory_constructs() { // Smoke test; scheme dispatch is covered upstream in iceberg-rust. let _factory = storage_factory(); } + + fn config(value: serde_json::Value) -> IcebergReaderConfig { + serde_json::from_value(value).unwrap() + } + + fn snapshot_config() -> IcebergReaderConfig { + config(json!({"mode": "snapshot", "metadata_location": "/tmp/metadata.json"})) + } + + /// SQL field named `name`; `nullable`, `unused`, and `default` control + /// whether `skip_unused_columns` may skip it. + fn field(name: &str, nullable: bool, unused: bool, default: Option<&str>) -> Field { + let mut field = Field::new(SqlIdentifier::new(name, false), ColumnType::int(nullable)) + .with_unused(unused); + field.default = default.map(String::from); + field + } + + fn relation(fields: Vec, skip_unused_columns: bool) -> Relation { + let zero = SourcePosition { + start_line_number: 0, + start_column: 0, + end_line_number: 0, + end_column: 0, + }; + let mut properties = BTreeMap::new(); + if skip_unused_columns { + properties.insert( + "skip_unused_columns".to_string(), + PropertyValue { + value: "true".to_string(), + key_position: zero, + value_position: zero, + }, + ); + } + Relation::new( + SqlIdentifier::new("test_table", false), + fields, + false, + properties, + ) + } + + fn arrow_schema(names: &[&str]) -> ArrowSchema { + ArrowSchema::new( + names + .iter() + .map(|name| ArrowField::new(*name, DataType::Int32, true)) + .collect::>(), + ) + } + + /// The projection is the intersection of the Iceberg schema and the SQL + /// declaration: table-only columns (`b`, `uuid`) are never read, and + /// SQL-only columns (`missing`) are never selected. Output follows table + /// schema order, not declaration order. + #[test] + fn used_columns_selects_sql_declared_table_columns_in_table_order() { + let schema = relation( + vec![ + field("c", true, false, None), + field("a", true, false, None), + field("missing", true, false, None), + ], + false, + ); + let table_schema = arrow_schema(&["a", "b", "c", "uuid"]); + + assert_eq!( + used_columns(&table_schema, &schema, &snapshot_config()), + vec!["a", "c"] + ); + } + + /// Names match case-insensitively, and the query uses the table's own + /// spelling (quoted identifiers resolve case-sensitively in datafusion). + #[test] + fn used_columns_matches_case_insensitively() { + let schema = relation(vec![field("tstz", true, false, None)], false); + let table_schema = arrow_schema(&["TsTz"]); + + assert_eq!( + used_columns(&table_schema, &schema, &snapshot_config()), + vec!["TsTz"] + ); + } + + #[test] + fn used_columns_empty_when_nothing_matches() { + let schema = relation(vec![field("x", true, false, None)], false); + let table_schema = arrow_schema(&["a", "b"]); + + assert!(used_columns(&table_schema, &schema, &snapshot_config()).is_empty()); + } + + /// With the `skip_unused_columns` table property, a column is dropped only + /// when it is unused *and* omittable (nullable or defaulted) *and* not + /// referenced by `snapshot_filter`. Without the property, everything the + /// SQL table declares is read. + #[test] + fn skip_unused_columns_drops_only_omittable_unreferenced_columns() { + let fields = vec![ + field("used", false, false, None), // read by a view -> kept + field("unused_nullable", true, true, None), // skippable -> dropped + field("unused_nonnull", false, true, None), // not omittable -> kept + field("unused_default", false, true, Some("0")), // defaulted -> dropped + field("unused_filtered", true, true, None), // filter needs it -> kept + ]; + let table_schema = arrow_schema(&[ + "used", + "unused_nullable", + "unused_nonnull", + "unused_default", + "unused_filtered", + ]); + let config = config(json!({ + "mode": "snapshot", + "metadata_location": "/tmp/metadata.json", + "snapshot_filter": "unused_filtered > 10", + })); + + assert_eq!( + used_columns(&table_schema, &relation(fields.clone(), false), &config).len(), + 5 + ); + assert_eq!( + used_columns(&table_schema, &relation(fields, true), &config), + vec!["used", "unused_nonnull", "unused_filtered"] + ); + } + + #[test] + fn used_column_list_quotes_identifiers() { + let columns = vec![ + "simple".to_string(), + "with\"quote".to_string(), + "With Space".to_string(), + ]; + + assert_eq!( + used_column_list(&columns), + r#""simple", "with""quote", "With Space""# + ); + } } diff --git a/crates/iceberg/src/test/create_test_table_s3.py b/crates/iceberg/src/test/create_test_table_s3.py index fabb75d53ec..df216bebd30 100644 --- a/crates/iceberg/src/test/create_test_table_s3.py +++ b/crates/iceberg/src/test/create_test_table_s3.py @@ -20,6 +20,7 @@ DecimalType, BinaryType, FixedType, + UUIDType, ) from pyiceberg.partitioning import PartitionSpec, PartitionField from pyiceberg.transforms import DayTransform @@ -28,6 +29,7 @@ import datetime import os import sys +import uuid import pyarrow as pa import pandas as pd import numpy as np @@ -55,6 +57,12 @@ help="Number of rows to generate (default: 1000000)", ) parser.add_argument("--json-file", help="JSON file to load data from") +parser.add_argument( + "--extra-columns", + action="store_true", + help="Add columns that the Feldera SQL test schemas do not declare " + "(a 'uuid' and a string column), which the connector must not read", +) args = parser.parse_args() @@ -110,7 +118,7 @@ pass # Iceberg schema (matches `IcebergTestStruct`) -schema = Schema( +schema_fields = [ NestedField(1, "b", BooleanType(), required=True), NestedField(2, "i", IntegerType(), required=True), NestedField(3, "l", LongType(), required=True), @@ -121,29 +129,40 @@ NestedField(8, "tm", TimeType(), required=True), NestedField(9, "ts", TimestampType(), required=True), NestedField(10, "s", StringType(), required=True), - # NestedField(11, "uuid", UUIDType(), required=True), NestedField(11, "fixed", FixedType(5), required=True), NestedField(12, "varbin", BinaryType(), required=True), -) +] # Equivalent arrow schema -arrow_schema = pa.schema( - [ - pa.field("b", pa.bool_(), nullable=False), - pa.field("i", pa.int32(), nullable=False), - pa.field("l", pa.int64(), nullable=False), - pa.field("r", pa.float32(), nullable=False), - pa.field("d", pa.float64(), nullable=False), - pa.field("dec", pa.decimal128(10, 3), nullable=False), - pa.field("dt", pa.date32(), nullable=False), - pa.field("tm", pa.time64("us"), nullable=False), - pa.field("ts", pa.timestamp("us"), nullable=False), - pa.field("s", pa.string(), nullable=False), - # pa.field("uuid", pa.binary(16), nullable=False), - pa.field("fixed", pa.binary(5), nullable=False), - pa.field("varbin", pa.binary(), nullable=False), +arrow_fields = [ + pa.field("b", pa.bool_(), nullable=False), + pa.field("i", pa.int32(), nullable=False), + pa.field("l", pa.int64(), nullable=False), + pa.field("r", pa.float32(), nullable=False), + pa.field("d", pa.float64(), nullable=False), + pa.field("dec", pa.decimal128(10, 3), nullable=False), + pa.field("dt", pa.date32(), nullable=False), + pa.field("tm", pa.time64("us"), nullable=False), + pa.field("ts", pa.timestamp("us"), nullable=False), + pa.field("s", pa.string(), nullable=False), + pa.field("fixed", pa.binary(5), nullable=False), + pa.field("varbin", pa.binary(), nullable=False), +] + +# Columns that the Feldera SQL test schemas do not declare; the connector +# must never select them. `uuid` exercises an extension-typed column. +if args.extra_columns: + schema_fields += [ + NestedField(13, "uuid", UUIDType(), required=False), + NestedField(14, "extra_s", StringType(), required=False), ] -) + arrow_fields += [ + pa.field("uuid", pa.uuid(), nullable=True), + pa.field("extra_s", pa.string(), nullable=True), + ] + +schema = Schema(*schema_fields) +arrow_schema = pa.schema(arrow_fields) partition_spec = PartitionSpec( PartitionField(source_id=9, field_id=1000, transform=DayTransform(), name="date") @@ -234,6 +253,11 @@ # print(pandas_df.head()) +if args.extra_columns: + print("Adding extra columns not declared in the Feldera SQL test schemas") + pandas_df["uuid"] = [uuid.uuid4().bytes for _ in range(len(pandas_df))] + pandas_df["extra_s"] = [f"extra_{i}" for i in range(len(pandas_df))] + print("Generating Pandas dataframe") # pyiceberg does not support nanosecond timestamps diff --git a/docs.feldera.com/docs/connectors/sources/iceberg.md b/docs.feldera.com/docs/connectors/sources/iceberg.md index b101c154db8..7f48f9280d2 100644 --- a/docs.feldera.com/docs/connectors/sources/iceberg.md +++ b/docs.feldera.com/docs/connectors/sources/iceberg.md @@ -156,6 +156,13 @@ The following table lists supported Iceberg data types and corresponding Feldera Types that are currently not supported include Iceberg's nested data types (`struct`s, `list`s and `map`s), `uuid`, and timestamps with time zone. +## Column selection + +The connector reads only the columns that appear in the Feldera SQL table +declaration. Other columns of the Iceberg table are never read. In addition, +when the table declaration sets the [`skip_unused_columns` property](/sql/grammar#ignoring-unused-columns), the connector skips declared columns that no view uses, provided they are +nullable or have default values. + [some other operator] let stream = self .add_source(input) - .dyn_shard(&factories.indexed_zset_factories); + .dyn_sharder() + .with_activity(ExchangeActivity::InputOnly) + .shard(&factories.indexed_zset_factories); let zset_handle = >>::new( factories.pair_factory, diff --git a/crates/dbsp/src/operator/dynamic/time_series/waterline.rs b/crates/dbsp/src/operator/dynamic/time_series/waterline.rs index 10a2bd40769..1b919acb79f 100644 --- a/crates/dbsp/src/operator/dynamic/time_series/waterline.rs +++ b/crates/dbsp/src/operator/dynamic/time_series/waterline.rs @@ -5,7 +5,7 @@ use size_of::SizeOf; use crate::circuit::checkpointer::Checkpoint; use crate::circuit::runtime::{WorkerLocation, WorkerLocations}; use crate::dynamic::{DynData, DynUnit}; -use crate::operator::communication::Mailbox; +use crate::operator::communication::{ExchangeActivity, Mailbox}; use crate::operator::dynamic::{MonoIndexedZSet, MonoZSet}; use crate::{ Circuit, NumEntries, RootCircuit, Stream, @@ -111,6 +111,7 @@ where *result = waterline; } }, + ExchangeActivity::AllSteps, ); match exchange { @@ -191,6 +192,7 @@ where let old_result = clone_box(result); least_upper_bound(&old_result, &waterline, result.as_mut()); }, + ExchangeActivity::AllSteps, ); match exchange { diff --git a/crates/sqllib/src/string_interner.rs b/crates/sqllib/src/string_interner.rs index b4c2f79b4d9..60747c0fefe 100644 --- a/crates/sqllib/src/string_interner.rs +++ b/crates/sqllib/src/string_interner.rs @@ -10,7 +10,7 @@ use dbsp::{ Circuit, DynZWeight, OrdZSet, RootCircuit, Runtime, Stream, ZWeight, circuit::{LocalStoreMarker, WorkerLocation, WorkerLocations}, dynamic::{DowncastTrait, DynData}, - operator::communication::{Mailbox, new_exchange_operators}, + operator::communication::{ExchangeActivity, Mailbox, new_exchange_operators}, storage::file::to_bytes, trace::{ BatchReader, BatchReaderFactories, Cursor, OrdIndexedWSet as DynOrdIndexedWSet, @@ -315,6 +315,7 @@ pub fn build_string_interner( snapshot.extend(remote_snapshot); } }, + ExchangeActivity::AllSteps, ); let by_id = match exchange { Some((sender, receiver)) => interned_strings From 00765df4c056b8ac1d40f9cf941478ad4bfdbbc5 Mon Sep 17 00:00:00 2001 From: Ben Pfaff Date: Fri, 12 Jun 2026 13:30:20 -0700 Subject: [PATCH 08/23] [dbsp] Segregate inter-host exchange channels by message size and type. Until now, when one host in a DBSP pipeline sends another one a message as part of synchronous or streaming exchange, it sent it over a single host-to-host TCP connection. This means that synchronous exchange has to wait for streaming exchange messages, which increases latency. This commit uses different TCP connections for streaming and synchronous exchange, which helps. In addition, it uses separate TCP connections for small and large synchronous exchange; we speculate that giving small synchronous messages priority can help with broadcast and consensus messages (this is unproven). Signed-off-by: Ben Pfaff --- crates/dbsp/src/operator/communication.rs | 3 +- .../src/operator/communication/exchange.rs | 68 ++++++++++++++++--- .../operator/dynamic/sharded_accumulator.rs | 15 ++-- 3 files changed, 67 insertions(+), 19 deletions(-) diff --git a/crates/dbsp/src/operator/communication.rs b/crates/dbsp/src/operator/communication.rs index a8ed35132ca..73ab4ef33cc 100644 --- a/crates/dbsp/src/operator/communication.rs +++ b/crates/dbsp/src/operator/communication.rs @@ -3,7 +3,8 @@ 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, diff --git a/crates/dbsp/src/operator/communication/exchange.rs b/crates/dbsp/src/operator/communication/exchange.rs index 511c32780f3..23197f55d50 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, } @@ -484,9 +524,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 +551,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,7 +582,8 @@ 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 } } @@ -963,16 +1006,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); } } } diff --git a/crates/dbsp/src/operator/dynamic/sharded_accumulator.rs b/crates/dbsp/src/operator/dynamic/sharded_accumulator.rs index 5f6f1ee62e1..51ba069782b 100644 --- a/crates/dbsp/src/operator/dynamic/sharded_accumulator.rs +++ b/crates/dbsp/src/operator/dynamic/sharded_accumulator.rs @@ -28,7 +28,8 @@ use crate::{ circuit_cache_key, operator::{ communication::{ - ExchangeClients, ExchangeDelivery, ExchangeDirectory, ExchangeId, pop_flushed, + ExchangeClients, ExchangeDelivery, ExchangeDirectory, ExchangeId, MessageType, + pop_flushed, }, dynamic::{ accumulator::{Accumulation, EnableCount}, @@ -312,12 +313,12 @@ where item.push(flush as u8); } let this = self.clone(); - if let Some(waiter) = this.clients.connect(receivers.start).await.send( - name.clone(), - this.exchange_id, - sender, - items, - ) { + if let Some(waiter) = this + .clients + .connect(receivers.start, MessageType::Streaming) + .await + .send(name.clone(), this.exchange_id, sender, items) + { remote_waiters.push(waiter); } } From 2c705bf2b3ff95a89c9a4eeb70ff1d26d6e4b308 Mon Sep 17 00:00:00 2001 From: Ben Pfaff Date: Wed, 17 Jun 2026 11:18:25 -0700 Subject: [PATCH 09/23] [dbsp] Simplify `ChildCircuit::add_exchange_with_preference`. This should be equivalent code, just shorter. Just noticed this while reading code, there is nothing important here. Signed-off-by: Ben Pfaff --- crates/dbsp/src/circuit/circuit_builder.rs | 31 +++++----------------- 1 file changed, 6 insertions(+), 25 deletions(-) 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 } From b4ef0389b3879fb39bdfca3bec4a2c4ad3810a86 Mon Sep 17 00:00:00 2001 From: Ben Pfaff Date: Wed, 17 Jun 2026 12:59:20 -0700 Subject: [PATCH 10/23] [dbsp] Break sharded accumulator waiting for merging into separate operator. Before this commit, the sharded accumulator waited for spines to merge in the ShardedAccumulatorSender operator before it returned a value from its eval() method. This prevented any other work from happening in the receiver and any other part of the circuit downstream from the receiver, because there is an artificial dependency between the sender and receiver created by add_exchange_operator(). One way to solve the problem is to remove that dependency (it is not essential for streaming exchange) but it still makes some sense to have it; it does make sense to receive after sending. Another reason is that ExchangeActivity::InputOnly currently expects all of the workers to flush in the same step, but without the dependency I am not confident without further thought that that would still be the case, since it would become nondeterministic which operator (sender or receiver) would run first. Instead, this commit breaks off the waiting part of ShardedAccumulatorSender into a new operator ShardedAccumulatorWaiter, which just does the wait. It adds a dependency on ShardedAccumulatorReceiver for that operator, so that it runs after everything has been sent and received in a step. The new waiter operator does no work, it only keeps the step from finishing before the spine used for accumulating data has merged to an acceptable number of batches. This way, the sender runs quickly and the wait does not block the receiver. Remote waiters are a different problem. This commit handles remote waiters by adding a single operator to wait for all of them at the very end of a step (by making the operator depend on all the sending operators). This isn't a problem in synchronous exchange (or perhaps it is a different problem) because there it is the receiver, not the sender, that adds the batch to the spine. Signed-off-by: Ben Pfaff --- .../src/operator/communication/exchange.rs | 14 ++ .../operator/dynamic/sharded_accumulator.rs | 183 ++++++++++++++---- 2 files changed, 164 insertions(+), 33 deletions(-) diff --git a/crates/dbsp/src/operator/communication/exchange.rs b/crates/dbsp/src/operator/communication/exchange.rs index 23197f55d50..08fa320fc54 100644 --- a/crates/dbsp/src/operator/communication/exchange.rs +++ b/crates/dbsp/src/operator/communication/exchange.rs @@ -382,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. @@ -586,6 +590,16 @@ impl ExchangeClients { .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 { diff --git a/crates/dbsp/src/operator/dynamic/sharded_accumulator.rs b/crates/dbsp/src/operator/dynamic/sharded_accumulator.rs index 51ba069782b..38b066325de 100644 --- a/crates/dbsp/src/operator/dynamic/sharded_accumulator.rs +++ b/crates/dbsp/src/operator/dynamic/sharded_accumulator.rs @@ -6,6 +6,7 @@ use std::{ panic::Location, pin::Pin, sync::{Arc, Mutex, MutexGuard}, + time::Instant, }; use feldera_samply::Span; @@ -16,7 +17,7 @@ use size_of::{HumanBytes, SizeOf, TotalSize}; use crate::{ Circuit, NumEntries, Runtime, Scope, Stream, circuit::{ - GlobalNodeId, OwnershipPreference, StepSize, WorkerLocation, WorkerLocations, + GlobalNodeId, NodeId, OwnershipPreference, StepSize, WorkerLocation, WorkerLocations, circuit_builder::StreamId, metadata::{ ALLOCATED_MEMORY_BYTES, BatchSizeStats, INPUT_BATCHES_STATS, MEMORY_ALLOCATIONS_COUNT, @@ -43,6 +44,8 @@ circuit_cache_key!(local StreamingExchangeCacheId(ExchangeId => Arc((StreamId, Range) => Accumulation>>>)); +circuit_cache_key!(ShardedAccumulatorRemoteWaiterId(() => NodeId)); + impl Stream where C: Circuit, @@ -76,6 +79,21 @@ where && runtime.layout().n_workers() > 1 && runtime.get_step_size() == StepSize::Microsteps { + let remote_waiter_node_id = if runtime.layout().is_multihost() { + let clients = ExchangeClients::for_runtime(&runtime); + Some(*self.circuit().cache_get_or_insert_with( + ShardedAccumulatorRemoteWaiterId::new(()), + move || { + let waiter = self + .circuit() + .add_source(ShardedAccumulatorRemoteWaiter::new(clients.clone())); + waiter.local_node_id() + }, + )) + } else { + None + }; + self.circuit() .cache_get_or_insert_with( ShardedAccumulatorId::new((self.stream_id(), workers.clone())), @@ -89,7 +107,13 @@ where factories, ); let enable_count = exchange.enable_count.clone(); - let stream = self + let local_waiter = + self.circuit() + .add_source(ShardedAccumulatorLocalWaiter::new( + Some(Location::caller()), + exchange.clone(), + )); + let receiver = self .circuit() .add_exchange( ShardedAccumulatorSender::new( @@ -100,8 +124,14 @@ where self, ) .mark_sharded_workers(workers.clone()); + self.circuit() + .add_dependency(receiver.local_node_id(), local_waiter.local_node_id()); + if let Some(remote_waiter_node_id) = remote_waiter_node_id { + self.circuit() + .add_dependency(receiver.local_node_id(), remote_waiter_node_id); + } Accumulation { - stream, + stream: receiver, enable_count, } }, @@ -135,6 +165,7 @@ where /// The RPC clients to contact remote hosts. clients: Arc, + /// One [Rxq] for each of `local_workers`. rxq: Vec>>, enable_count: EnableCount, @@ -261,7 +292,6 @@ where let worker_locations = WorkerLocations::for_layout(layout); let mut data = batches.into_iter(); let mut remote_waiters = Vec::new(); - let mut local_waiters = Vec::new(); let mut serialized_bytes = 0; for receivers in layout.all_hosts() { match worker_locations[receivers.start] { @@ -272,18 +302,7 @@ where .expect("data should include one item per peer") .into_plain() .expect("local data should not be serialized"); - if self.deliver(&self.factories, sender, receiver, item, flush) - && !flush - && let Some(waiter) = self - .rxq(receiver) - .spines - .back() - .unwrap() - .spine - .backpressure_waiter() - { - local_waiters.push((receiver, waiter)); - } + self.deliver(&self.factories, sender, receiver, item, flush); } } WorkerLocation::Remote => { @@ -326,23 +345,6 @@ where } } - if !local_waiters.is_empty() { - let _span = Span::new("local send wait") - .with_category("Exchange") - .with_tooltip(|| { - format!( - "{name} wait for batches to merge in {} receive queues (for workers {})", - local_waiters.len(), - local_waiters - .iter() - .map(|(receiver, _waiter)| receiver) - .format(", ") - ) - }); - for (_receiver, waiter) in local_waiters { - waiter.await; - } - } if !remote_waiters.is_empty() { let _span = Span::new("remote send wait") .with_category("Exchange") @@ -372,6 +374,43 @@ where } } +impl ShardedAccumulator +where + B: Batch, +{ + async fn wait(&self, name: Arc) { + let start = Instant::now(); + let mut local_waiters = Vec::new(); + for (rxq, worker) in self.rxq.iter().zip(self.local_workers.clone()) { + // This is intentionally two separate statements to avoid holding + // the lock while waiting. + let waiter = rxq + .lock() + .unwrap() + .spines + .front() + .and_then(|entry| entry.spine.backpressure_waiter()); + if let Some(waiter) = waiter { + local_waiters.push(worker); + waiter.await; + } + } + if !local_waiters.is_empty() { + Span::new("local send wait") + .with_start(start) + .with_category("Exchange") + .with_tooltip(|| { + format!( + "{name} wait for batches to merge in {} receive queues (for workers {})", + local_waiters.len(), + local_waiters.iter().format(", ") + ) + }) + .record(); + } + } +} + impl ExchangeDelivery for ShardedAccumulator where B: Batch