From b4cdbddad0bef3444eb3e0cdc87dbd73788082cd Mon Sep 17 00:00:00 2001 From: Loakesh Indiran <66478092+Loaki07@users.noreply.github.com> Date: Fri, 8 Aug 2025 08:43:19 +0530 Subject: [PATCH] fix: track_total_hits to consider multiple distinct cols (#7894) ### **User description** ## Problem Description When using `SELECT DISTINCT` queries with multiple columns in OpenObserve, the track total hits functionality was incorrectly counting only the first column instead of all columns, and the response field name was incorrect. ### Example Problem Query ```sql SELECT DISTINCT unique_id, continent FROM oly WHERE continent = 'ASI' ``` ### Issues Encountered 1. **Incorrect Count**: The track total hits was only counting distinct values of the first column (`unique_id`) instead of distinct combinations of all columns (`unique_id, continent`). 2. **DataFusion Error**: The original approach tried to use `count(DISTINCT unique_id, continent)` which caused DataFusion to throw an error: ``` "This feature is not implemented: COUNT DISTINCT with multiple arguments" ``` ## Solution ### Approach Instead of trying to use `COUNT DISTINCT` with multiple arguments (which DataFusion doesn't support), we now use a subquery approach: **Before (causing errors):** ```sql SELECT count(DISTINCT unique_id, continent) AS zo_sql_num FROM oly WHERE continent = 'ASI' ``` **After (DataFusion compatible):** ```sql SELECT count(*) AS zo_sql_num FROM (SELECT DISTINCT unique_id, continent FROM oly WHERE continent = 'ASI') ``` ## Tests image image (18) ## Testing 1. **Single column DISTINCT**: ```sql SELECT DISTINCT name FROM t WHERE name = 'a' ``` Expected: `SELECT count(*) AS zo_sql_num FROM (SELECT DISTINCT name FROM t WHERE name = 'a')` 2. **Multiple columns DISTINCT**: ```sql SELECT DISTINCT unique_id, continent FROM oly WHERE continent = 'ASI' ``` Expected: `SELECT count(*) AS zo_sql_num FROM (SELECT DISTINCT unique_id, continent FROM oly WHERE continent = 'ASI')` 3. **Three columns DISTINCT**: ```sql SELECT DISTINCT unique_id, continent, bronze_medals FROM oly WHERE continent = 'ASI' ``` Expected: `SELECT count(*) AS zo_sql_num FROM (SELECT DISTINCT unique_id, continent, bronze_medals FROM oly WHERE continent = 'ASI')` ___ ### **PR Type** Bug fix, Tests ___ ### **Description** - Wrap DISTINCT queries in subquery for counting - Simplify non-DISTINCT count to wildcard - Remove unsupported COUNT DISTINCT arguments - Add unit tests for DISTINCT cases ___ ### Diagram Walkthrough ```mermaid flowchart LR A["Original DISTINCT query"] B["Wrapped subquery"] C["Count(*) AS zo_sql_num"] A -- "wrap" --> B B -- "count" --> C ```

File Walkthrough
Relevant files
Bug fix
track_total_hits.rs
Support multi-column DISTINCT and tests                                   
src/service/search/sql/rewriter/track_total_hits.rs
  • Use subquery wrapper for multi-column DISTINCT
  • Simplify count logic for non-DISTINCT queries
  • Remove DuplicateTreatment import
  • Add unit tests for DISTINCT scenarios
 +135/-40
___ --- .../search/sql/rewriter/track_total_hits.rs | 175 ++++++++++++++---- 1 file changed, 135 insertions(+), 40 deletions(-) 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); + } }