diff --git a/rust/lancedb/src/table/query.rs b/rust/lancedb/src/table/query.rs index bce8fd5a4..a955ac171 100644 --- a/rust/lancedb/src/table/query.rs +++ b/rust/lancedb/src/table/query.rs @@ -327,7 +327,7 @@ pub async fn create_plan( scanner .create_plan() .await - .map_err(|error| enrich_field_not_found(error, &arrow_schema)) + .map_err(|error| enrich_lance_field_not_found(error, &arrow_schema)) } /// Replace DataFusion's top-level field candidates with qualified leaf paths. @@ -336,19 +336,24 @@ pub async fn create_plan( /// top-level Arrow fields. This makes a missing leaf look unavailable even when it /// exists below a struct. Keep every other Lance/DataFusion error unchanged and /// enrich only this one schema error at the LanceDB query boundary. -fn enrich_field_not_found(error: lance::Error, schema: &Schema) -> Error { - let Some(field) = find_missing_field(&error) else { - return error.into(); - }; +fn enrich_lance_field_not_found(error: lance::Error, schema: &Schema) -> Error { + field_not_found_diagnostic(&error, schema).unwrap_or_else(|| error.into()) +} - let schema_error = SchemaError::FieldNotFound { - field: Box::new(field.clone()), - valid_fields: leaf_field_columns(schema), - }; - let error = DataFusionError::SchemaError(Box::new(schema_error), Box::new(None)); - Error::InvalidInput { - message: error.to_string(), +fn field_not_found_diagnostic( + error: &(dyn std::error::Error + 'static), + schema: &Schema, +) -> Option { + let field = find_missing_field(error)?; + let valid_fields = leaf_field_paths(schema); + let mut message = format!("Schema error: No field named {}", field.quoted_flat_name()); + if !valid_fields.is_empty() { + message.push_str(". Valid fields are "); + message.push_str(&valid_fields.join(", ")); } + message.push('.'); + + Some(Error::InvalidInput { message }) } fn find_missing_field<'a>(error: &'a (dyn std::error::Error + 'static)) -> Option<&'a Column> { @@ -362,32 +367,26 @@ fn find_missing_field<'a>(error: &'a (dyn std::error::Error + 'static)) -> Optio error.source().and_then(find_missing_field) } -fn leaf_field_columns(schema: &Schema) -> Vec { - fn visit(fields: &arrow_schema::Fields, path: &mut Vec, columns: &mut Vec) { +fn leaf_field_paths(schema: &Schema) -> Vec { + fn visit(fields: &arrow_schema::Fields, path: &mut Vec, paths: &mut Vec) { for field in fields { path.push(field.name().clone()); match field.data_type() { DataType::Struct(children) if !children.is_empty() => { - visit(children, path, columns); + visit(children, path, paths); } _ => { - let qualified_name = path - .iter() - .map(|part| { - datafusion_common::utils::quote_identifier(part.as_str()).into_owned() - }) - .collect::>() - .join("."); - columns.push(Column::from_qualified_name_ignore_case(qualified_name)); + let segments = path.iter().map(String::as_str).collect::>(); + paths.push(lance_core::datatypes::format_field_path(&segments)); } } path.pop(); } } - let mut columns = Vec::new(); - visit(schema.fields(), &mut Vec::new(), &mut columns); - columns + let mut paths = Vec::new(); + visit(schema.fields(), &mut Vec::new(), &mut paths); + paths } //Helper functions below @@ -939,7 +938,7 @@ mod tests { } #[tokio::test] - async fn test_missing_filter_field_lists_nested_fields() { + async fn test_missing_filter_field_lists_nested_fields_in_local_planners() { use crate::connect; use arrow_schema::{DataType, Field, Schema}; @@ -984,11 +983,71 @@ mod tests { .await .err() .expect("query should reject the unqualified nested field"); + let expected = "No field named year. Valid fields are id, vector, content, metadata.year, metadata.genre."; assert!( - error.to_string().contains( - "No field named year. Valid fields are id, vector, content, metadata.year, metadata.genre." - ), + error.to_string().contains(expected), + "unexpected error: {error}" + ); + + table.set_unenforced_primary_key(["id"]).await.unwrap(); + table + .set_lsm_write_spec(crate::table::LsmWriteSpec::unsharded()) + .await + .unwrap(); + let lsm_error = table + .query() + .only_if("year = 2024") + .execute() + .await + .err() + .expect("LSM query should reject the unqualified nested field"); + + assert!( + lsm_error.to_string().contains(expected), + "unexpected LSM error: {lsm_error}" + ); + } + + #[test] + fn test_leaf_field_paths_preserve_arbitrary_depth() { + use arrow_schema::{DataType, Field, Schema}; + + fn nested_field(path: &[&str]) -> Field { + let mut segments = path.iter().rev(); + let mut field = Field::new( + *segments.next().expect("path must have a leaf"), + DataType::Int32, + false, + ); + for segment in segments { + field = Field::new(*segment, DataType::Struct(vec![field].into()), false); + } + field + } + + let schema = Schema::new(vec![ + nested_field(&["a", "b", "c", "d", "e"]), + nested_field(&["metadata", "child.with.dot"]), + ]); + + assert_eq!( + leaf_field_paths(&schema), + vec!["a.b.c.d.e", "metadata.`child.with.dot`"] + ); + + let source = DataFusionError::SchemaError( + Box::new(SchemaError::FieldNotFound { + field: Box::new(Column::from_name("missing")), + valid_fields: Vec::new(), + }), + Box::new(None), + ); + let error = field_not_found_diagnostic(&source, &schema).unwrap(); + assert!( + error + .to_string() + .contains("Valid fields are a.b.c.d.e, metadata.`child.with.dot`"), "unexpected error: {error}" ); } diff --git a/rust/lancedb/src/table/query/lsm.rs b/rust/lancedb/src/table/query/lsm.rs index 074d13476..085f9a10a 100644 --- a/rust/lancedb/src/table/query/lsm.rs +++ b/rust/lancedb/src/table/query/lsm.rs @@ -27,6 +27,8 @@ use std::sync::Arc; use arrow_array::Array; use arrow_schema::{DataType, Schema as ArrowSchema}; +use datafusion::common::{DataFusionError, ToDFSchema}; +use datafusion::prelude::SessionContext; use datafusion_physical_plan::expressions::Column; use datafusion_physical_plan::projection::ProjectionExec; use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr}; @@ -370,6 +372,7 @@ fn base_scanner( snapshots: Vec, in_memory: HashMap, ) -> Result { + let schema = ArrowSchema::from(dataset.schema()); let mut scanner = LsmScanner::new(Arc::new(dataset.clone()), snapshots, pk_columns); for (shard_id, memtables) in in_memory { scanner = scanner.with_in_memory_memtables(shard_id, memtables); @@ -380,7 +383,20 @@ fn base_scanner( } if let Some(filter) = &query.base.filter { scanner = match filter { - QueryFilter::Sql(sql) => scanner.filter(sql)?, + QueryFilter::Sql(sql) => { + // Parse here instead of inside `LsmScanner::filter` so the typed + // DataFusion `FieldNotFound` error is still available for the + // same nested-field enrichment used by the ordinary scanner. + let df_schema = schema.clone().to_dfschema().map_err(|error| { + enrich_filter_error(error, &schema, "Failed to create DFSchema") + })?; + let expr = SessionContext::new() + .parse_sql_expr(sql, &df_schema) + .map_err(|error| { + enrich_filter_error(error, &schema, "Failed to parse filter expression") + })?; + scanner.filter_expr(expr) + } QueryFilter::Datafusion(expr) => scanner.filter_expr(expr.clone()), QueryFilter::Substrait(_) => { return Err(Error::NotSupported { @@ -392,6 +408,12 @@ fn base_scanner( Ok(scanner) } +fn enrich_filter_error(error: DataFusionError, schema: &ArrowSchema, context: &str) -> Error { + super::field_not_found_diagnostic(&error, schema).unwrap_or_else(|| Error::InvalidInput { + message: format!("{context}: {error}"), + }) +} + /// Plain scan: filter / projection / limit over base ∪ SSTables ∪ in-memory. /// The plain scan applies limit and offset inside the planner. async fn plain_plan(