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. +