diff --git a/src/service/search/sql/rewriter/track_total_hits.rs b/src/service/search/sql/rewriter/track_total_hits.rs index cae773d015..99fc9c2f33 100644 --- a/src/service/search/sql/rewriter/track_total_hits.rs +++ b/src/service/search/sql/rewriter/track_total_hits.rs @@ -16,10 +16,9 @@ use std::ops::ControlFlow; use sqlparser::ast::{ - DuplicateTreatment, Expr, Function, FunctionArg, FunctionArgExpr, FunctionArgumentList, - FunctionArguments, GroupByExpr, Ident, ObjectName, ObjectNamePart, Query, Select, SelectFlavor, - SelectItem, SetExpr, TableFactor, TableWithJoins, VisitorMut, - helpers::attached_token::AttachedToken, + Expr, Function, FunctionArg, FunctionArgExpr, FunctionArgumentList, FunctionArguments, + GroupByExpr, Ident, ObjectName, ObjectNamePart, Query, Select, SelectFlavor, SelectItem, + SetExpr, TableFactor, TableWithJoins, VisitorMut, helpers::attached_token::AttachedToken, }; pub struct TrackTotalHitsVisitor {} @@ -36,44 +35,96 @@ impl VisitorMut for TrackTotalHitsVisitor { fn pre_visit_query(&mut self, query: &mut Query) -> ControlFlow { match query.body.as_mut() { SetExpr::Select(select) => { - let (field_expr, duplicate_treatment) = if select.distinct.is_some() { - match select.projection.first() { - Some(SelectItem::UnnamedExpr(expr)) => ( - FunctionArgExpr::Expr(expr.clone()), - Some(DuplicateTreatment::Distinct), - ), - Some(SelectItem::ExprWithAlias { expr, alias: _ }) => ( - FunctionArgExpr::Expr(expr.clone()), - Some(DuplicateTreatment::Distinct), - ), - _ => (FunctionArgExpr::Wildcard, None), - } + if select.distinct.is_some() { + // For DISTINCT queries, we need to wrap the query in a subquery + // and count the results, since DataFusion doesn't support COUNT DISTINCT with + // multiple arguments + let original_query = query.clone(); + let subquery = Box::new(SetExpr::Select(Box::new(Select { + select_token: AttachedToken::empty(), + distinct: None, + top: None, + top_before_distinct: false, + projection: vec![SelectItem::ExprWithAlias { + expr: Expr::Function(Function { + name: ObjectName(vec![ObjectNamePart::Identifier(Ident::new( + "count", + ))]), + parameters: FunctionArguments::None, + args: FunctionArguments::List(FunctionArgumentList { + args: vec![FunctionArg::Unnamed(FunctionArgExpr::Wildcard)], + duplicate_treatment: None, + clauses: vec![], + }), + filter: None, + null_treatment: None, + over: None, + within_group: vec![], + uses_odbc_syntax: false, + }), + alias: Ident::new("zo_sql_num"), + }], + into: None, + from: vec![TableWithJoins { + relation: TableFactor::Derived { + lateral: false, + subquery: Box::new(original_query), + alias: None, + }, + joins: vec![], + }], + lateral_views: vec![], + selection: None, + group_by: GroupByExpr::Expressions(vec![], vec![]), + having: None, + prewhere: None, + sort_by: vec![], + cluster_by: vec![], + distribute_by: vec![], + named_window: vec![], + qualify: None, + window_before_qualify: false, + connect_by: None, + value_table_mode: None, + flavor: SelectFlavor::Standard, + }))); + *query = Query { + with: None, + body: subquery, + order_by: None, + limit: None, + offset: None, + fetch: None, + limit_by: vec![], + for_clause: None, + locks: vec![], + settings: None, + format_clause: None, + }; } else { - (FunctionArgExpr::Wildcard, None) - }; - - select.group_by = GroupByExpr::Expressions(vec![], vec![]); - select.having = None; - select.sort_by = vec![]; - select.projection = vec![SelectItem::ExprWithAlias { - expr: Expr::Function(Function { - name: ObjectName(vec![ObjectNamePart::Identifier(Ident::new("count"))]), - parameters: FunctionArguments::None, - args: FunctionArguments::List(FunctionArgumentList { - args: vec![FunctionArg::Unnamed(field_expr)], - duplicate_treatment, - clauses: vec![], + // For non-DISTINCT queries, use the original approach + select.group_by = GroupByExpr::Expressions(vec![], vec![]); + select.having = None; + select.sort_by = vec![]; + select.projection = vec![SelectItem::ExprWithAlias { + expr: Expr::Function(Function { + name: ObjectName(vec![ObjectNamePart::Identifier(Ident::new("count"))]), + parameters: FunctionArguments::None, + args: FunctionArguments::List(FunctionArgumentList { + args: vec![FunctionArg::Unnamed(FunctionArgExpr::Wildcard)], + duplicate_treatment: None, + clauses: vec![], + }), + filter: None, + null_treatment: None, + over: None, + within_group: vec![], + uses_odbc_syntax: false, }), - filter: None, - null_treatment: None, - over: None, - within_group: vec![], - uses_odbc_syntax: false, - }), - alias: Ident::new("zo_sql_num"), - }]; - select.distinct = None; - query.order_by = None; + alias: Ident::new("zo_sql_num"), + }]; + query.order_by = None; + } } SetExpr::SetOperation { .. } => { let select = Box::new(SetExpr::Select(Box::new(Select { @@ -240,4 +291,48 @@ mod tests { let expected_sql = "SELECT count(*) AS zo_sql_num FROM (SELECT name FROM t1 UNION SELECT name FROM t2 UNION SELECT name FROM t3)"; assert_eq!(statement.to_string(), expected_sql); } + + #[test] + fn test_track_total_hits_distinct_single_column() { + let sql = "SELECT DISTINCT name FROM t WHERE name = 'a'"; + let mut statement = sqlparser::parser::Parser::parse_sql(&GenericDialect {}, &sql) + .unwrap() + .pop() + .unwrap(); + let mut track_total_hits_visitor = TrackTotalHitsVisitor::new(); + let _ = statement.visit(&mut track_total_hits_visitor); + // For DISTINCT queries, we wrap in a subquery to count results + let expected_sql = + "SELECT count(*) AS zo_sql_num FROM (SELECT DISTINCT name FROM t WHERE name = 'a')"; + assert_eq!(statement.to_string(), expected_sql); + } + + #[test] + fn test_track_total_hits_distinct_multiple_columns() { + let sql = "SELECT DISTINCT unique_id, continent FROM oly WHERE continent = 'ASI'"; + let mut statement = sqlparser::parser::Parser::parse_sql(&GenericDialect {}, &sql) + .unwrap() + .pop() + .unwrap(); + let mut track_total_hits_visitor = TrackTotalHitsVisitor::new(); + let _ = statement.visit(&mut track_total_hits_visitor); + // For DISTINCT queries with multiple columns, we wrap in a subquery to count results + let expected_sql = "SELECT count(*) AS zo_sql_num FROM (SELECT DISTINCT unique_id, continent FROM oly WHERE continent = 'ASI')"; + assert_eq!(statement.to_string(), expected_sql); + } + + #[test] + fn test_track_total_hits_distinct_three_columns() { + let sql = + "SELECT DISTINCT unique_id, continent, bronze_medals FROM oly WHERE continent = 'ASI'"; + let mut statement = sqlparser::parser::Parser::parse_sql(&GenericDialect {}, &sql) + .unwrap() + .pop() + .unwrap(); + let mut track_total_hits_visitor = TrackTotalHitsVisitor::new(); + let _ = statement.visit(&mut track_total_hits_visitor); + // For DISTINCT queries with multiple columns, we wrap in a subquery to count results + let expected_sql = "SELECT count(*) AS zo_sql_num FROM (SELECT DISTINCT unique_id, continent, bronze_medals FROM oly WHERE continent = 'ASI')"; + assert_eq!(statement.to_string(), expected_sql); + } }