diff --git a/Cargo.lock b/Cargo.lock index d8dcf1fd67a..1f406a05533 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4631,6 +4631,7 @@ dependencies = [ "prost", "prost-build", "protobuf-src", + "rstest", "tokio", "tracing", ] diff --git a/rust/lance-datafusion/Cargo.toml b/rust/lance-datafusion/Cargo.toml index 7f93ab619cd..bb48cae0db7 100644 --- a/rust/lance-datafusion/Cargo.toml +++ b/rust/lance-datafusion/Cargo.toml @@ -42,6 +42,7 @@ protobuf-src = {version = "2.1", optional = true} [dev-dependencies] lance-datagen.workspace = true +rstest.workspace = true [features] geo = ["dep:lance-geo"] diff --git a/rust/lance-datafusion/src/planner.rs b/rust/lance-datafusion/src/planner.rs index 31e3c2d9e4c..add2d276973 100644 --- a/rust/lance-datafusion/src/planner.rs +++ b/rust/lance-datafusion/src/planner.rs @@ -931,12 +931,21 @@ impl Planner { /// Create Logical [Expr] from a SQL expression. /// + /// Simple column references are resolved directly against the schema, + /// including backtick-quoted names (e.g. `` `id` `` resolves to field `id`) + /// and case-insensitive matching. Compound expressions fall through to + /// the full SQL parser. + /// /// Note: the returned expression must be passed through `optimize_filter()` /// before being passed to `create_physical_expr()`. pub fn parse_expr(&self, expr: &str) -> Result { // First check if it's a simple column reference (no operators, functions, etc.) // resolve_column_name tries exact match first, then falls back to case-insensitive - let resolved_name = self.resolve_column_name(expr); + let bare = expr + .strip_prefix('`') + .and_then(|s| s.strip_suffix('`')) + .unwrap_or(expr); + let resolved_name = self.resolve_column_name(bare); if self.schema.field_with_name(&resolved_name).is_ok() { return Ok(Expr::Column(Column::from_name(resolved_name))); } @@ -1098,6 +1107,7 @@ mod tests { prelude::{array_element, get_field}, }; use datafusion_functions::core::expr_ext::FieldAccessor; + use rstest::rstest; #[test] fn test_parse_filter_simple() { @@ -1170,6 +1180,34 @@ mod tests { ); } + #[rstest] + #[case::bare("id", "id")] + #[case::backtick_quoted("`id`", "id")] + #[case::case_insensitive("ID", "id")] + #[case::backtick_case_insensitive("`ID`", "id")] + #[case::reserved_word("`order`", "order")] + fn test_parse_expr_column_resolution(#[case] input: &str, #[case] expected_name: &str) { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("order", DataType::Utf8, true), + ])); + let planner = Planner::new(schema); + + let expr = planner.parse_expr(input).unwrap(); + assert_eq!(expr, Expr::Column(Column::from_name(expected_name))); + } + + #[test] + fn test_parse_expr_sql_fallback() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + ])); + let planner = Planner::new(schema); + + let expr = planner.parse_expr("id + 1").unwrap(); + assert_eq!(expr, col("id") + lit(1_i32)); + } + #[test] fn test_parse_deep_logical_filter() { let planner = Planner::new(Arc::new(Schema::empty()));