Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 29 additions & 0 deletions crates/adapterlib/src/utils/datafusion.rs
Original file line number Diff line number Diff line change
Expand Up @@ -364,6 +364,35 @@ pub fn columns_referenced_by_order_by(order_by: &str) -> Result<BTreeSet<String>
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<S: AsRef<str>>(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<String>,
}

impl ColumnNameSet {
pub fn from_names(names: impl IntoIterator<Item = String>) -> 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 {
Expand Down
40 changes: 5 additions & 35 deletions crates/adapters/src/integrated/delta_table/input.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<S: AsRef<str>>(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
Expand Down Expand Up @@ -883,29 +877,6 @@ struct CatchupFollowState {
transaction: Option<Option<String>>,
}

/// 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<String>,
}

impl ColumnNameSet {
fn from_names(names: impl IntoIterator<Item = String>) -> 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,
Expand Down Expand Up @@ -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<Option<String>> {
Expand Down
50 changes: 45 additions & 5 deletions crates/adapters/src/test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -242,6 +244,44 @@ where
+ for<'de> DeserializeWithContext<'de, SqlSerdeConfig, Variant>
+ Sync,
{
test_circuit_with_properties::<T>(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<T>(
config: CircuitConfig,
schema: &[Field],
input_properties: &[(&str, &str)],
persistent_output_ids: &[Option<&str>],
) -> (DBSPHandle, Box<dyn CircuitCatalog>)
where
T: DBData
+ SerializeWithContext<SqlSerdeConfig>
+ 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<String, PropertyValue> = 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()
Expand All @@ -264,7 +304,7 @@ where
format!("test_input{i}").into(),
schema.clone(),
false,
BTreeMap::new(),
input_properties.clone(),
))
.unwrap();

Expand Down
51 changes: 51 additions & 0 deletions crates/adapters/src/test/data.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1003,6 +1003,57 @@ deserialize_table_record!(S3TablesTestStruct["S3TablesTestStruct", Variant, 3] {
(created_at, "created_at", true, Option<Timestamp>, |_| 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<i64>,
}

impl IcebergSubsetTestStruct {
pub fn schema() -> Vec<Field> {
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<i64>
});

deserialize_table_record!(IcebergSubsetTestStruct["IcebergSubsetTestStruct", Variant, 3] {
(i, "i", false, i32, |_| None),
(s, "s", false, String, |_| None),
(l, "l", false, Option<i64>, |_| Some(None))
});

/// Struct will all types supported by the DeltaLake connector.
#[derive(
Debug,
Expand Down
Loading
Loading