diff --git a/src/config/src/utils/schema.rs b/src/config/src/utils/schema.rs index f1541d7a5f..f88d00e9a0 100644 --- a/src/config/src/utils/schema.rs +++ b/src/config/src/utils/schema.rs @@ -131,9 +131,9 @@ fn infer_json_schema_from_object( } else if v.is_f64() { convert_data_type(fields, key, DataType::Float64)?; } else { - return Err(ArrowError::SchemaError( - "Cannot infer schema from non-basic-number type value".to_string(), - )); + return Err(ArrowError::SchemaError(format!( + "Cannot infer schema from non-basic-number type value: {v:?}", + ))); } } Value::Bool(_) => { diff --git a/src/service/search/datafusion/optimizer/physical_optimizer/rewrite_match.rs b/src/service/search/datafusion/optimizer/physical_optimizer/rewrite_match.rs index 54a9cc6ff4..0b68fb1992 100644 --- a/src/service/search/datafusion/optimizer/physical_optimizer/rewrite_match.rs +++ b/src/service/search/datafusion/optimizer/physical_optimizer/rewrite_match.rs @@ -118,12 +118,21 @@ impl TreeNodeRewriter for PlanRewriter { .rewrite(&mut add_fst_fields_to_projection)? .data; - // 3. Rewrite match_all/fuzzy_match_all - let mut expr_rewriter = - MatchAllRewriter::new(input.schema().clone(), self.fields.clone()); - let rewritten_predicate = predicate.clone().rewrite(&mut expr_rewriter).data()?; + // 3. Update column indices in the predicate to match the new schema + let new_schema = input.schema(); + let mut column_index_rewriter = ColumnIndexRewriter::new(new_schema.clone()); + let predicate_with_updated_indices = predicate + .clone() + .rewrite(&mut column_index_rewriter) + .data()?; - // 4. Create new filter with rewritten predicate + // 4. Rewrite match_all/fuzzy_match_all + let mut expr_rewriter = MatchAllRewriter::new(new_schema.clone(), self.fields.clone()); + let rewritten_predicate = predicate_with_updated_indices + .rewrite(&mut expr_rewriter) + .data()?; + + // 5. Create new filter with rewritten predicate let new_filter = FilterExec::try_new(rewritten_predicate, input)? .with_projection(add_fst_fields_to_projection.filter_projection.clone())?; return Ok(Transformed::yes(Arc::new(new_filter))); @@ -132,6 +141,38 @@ impl TreeNodeRewriter for PlanRewriter { } } +// Rewriter for updating column indices when schema changes +#[derive(Debug, Clone)] +struct ColumnIndexRewriter { + new_schema: SchemaRef, +} + +impl ColumnIndexRewriter { + fn new(new_schema: SchemaRef) -> Self { + Self { new_schema } + } +} + +impl TreeNodeRewriter for ColumnIndexRewriter { + type Node = Arc; + + fn f_up(&mut self, expr: Self::Node) -> Result> { + // Check if this is a Column expression + if let Some(column) = expr.as_any().downcast_ref::() { + let field_name = column.name(); + // Find the new index in the new schema + if let Ok(new_index) = self.new_schema.index_of(field_name) { + // Only update if the index changed + if new_index != column.index() { + let new_column = Arc::new(Column::new(field_name, new_index)); + return Ok(Transformed::yes(new_column as Arc)); + } + } + } + Ok(Transformed::no(expr)) + } +} + // Rewriter for match_all() to LIKE expressions in physical plans #[derive(Debug, Clone)] pub struct MatchAllRewriter { @@ -454,7 +495,9 @@ mod tests { }; use super::*; - use crate::service::search::datafusion::udf::match_all_udf; + use crate::service::search::datafusion::{ + table_provider::empty_table::NewEmptyTable, udf::match_all_udf, + }; #[tokio::test] async fn test_rewrite_match_physical() { @@ -555,4 +598,69 @@ mod tests { assert_batches_eq!(item.1, &data); } } + + #[tokio::test] + async fn test_column_index_update_with_match_all_and_str_match() { + let schema = Arc::new(Schema::new(vec![ + Field::new("_timestamp", DataType::Int64, false), + Field::new("msg", DataType::Utf8, false), + Field::new("error", DataType::Utf8, false), + Field::new("message", DataType::Utf8, false), + ])); + + let fields = vec![ + ("error".to_string(), DataType::Utf8), + ("msg".to_string(), DataType::Utf8), + ("message".to_string(), DataType::Utf8), + ]; + + let state = SessionStateBuilder::new() + .with_config(SessionConfig::new()) + .with_runtime_env(Arc::new(RuntimeEnvBuilder::new().build().unwrap())) + .with_default_features() + .with_physical_optimizer_rules(vec![Arc::new(RewriteMatchPhysical::new( + fields.clone(), + ))]) + .build(); + let ctx = SessionContext::new_with_state(state); + let provider = NewEmptyTable::new("t", schema).with_partitions(8); + ctx.register_table("t", Arc::new(provider)).unwrap(); + ctx.register_udf(match_all_udf::MATCH_ALL_UDF.clone()); + ctx.register_udf( + crate::service::search::datafusion::udf::str_match_udf::STR_MATCH_UDF.clone(), + ); + + let sql = + "select count(*) from t where match_all('test') and str_match(message, 'success')"; + let plan = ctx.state().create_logical_plan(sql).await.unwrap(); + let physical_plan = ctx.state().create_physical_plan(&plan).await.unwrap(); + + // Verify that all column indices in FilterExec match the input schema + let _ = physical_plan.apply(&mut |node: &Arc| -> Result { + if let Some(filter) = node.as_any().downcast_ref::() { + let input_schema = filter.input().schema(); + let predicate = filter.predicate(); + + // Traverse the predicate expression tree to find all Column references + let _ = predicate.apply(&mut |expr: &Arc| -> Result { + if let Some(column) = expr.as_any().downcast_ref::() { + let column_name = column.name(); + let column_index = column.index(); + + // Find the expected index in the input schema + match input_schema.index_of(column_name) { + Ok(expected_index) => { + if column_index != expected_index { + panic!("Column '{column_name}' has index {column_index} in filter predicate but should be {expected_index} according to input schema"); + } + } + Err(_) => panic!("Column '{column_name}' with index {column_index} not found in input schema"), + } + } + Ok(TreeNodeRecursion::Continue) + }); + } + Ok(TreeNodeRecursion::Continue) + }); + } }