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 <img width="1802" height="1009" alt="image" src="https://github.com/user-attachments/assets/ec302f65-9115-4584-a9de-a2a51cbad502" /> <img width="1804" height="1009" alt="image (18)" src="https://github.com/user-attachments/assets/e38d9b51-afa8-4bfe-a61a-5f6fa1683eb7" /> ## 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 ``` <details> <summary><h3> File Walkthrough</h3></summary> <table><thead><tr><th></th><th align="left">Relevant files</th></tr></thead><tbody><tr><td><strong>Bug fix</strong></td><td><table> <tr> <td> <details> <summary><strong>track_total_hits.rs</strong><dd><code>Support multi-column DISTINCT and tests</code>                                    </dd></summary> <hr> src/service/search/sql/rewriter/track_total_hits.rs <ul><li>Use subquery wrapper for multi-column DISTINCT<br> <li> Simplify count logic for non-DISTINCT queries<br> <li> Remove DuplicateTreatment import<br> <li> Add unit tests for DISTINCT scenarios</ul> </details> </td> <td><a href="https://github.com/openobserve/openobserve/pull/7894/files#diff-723facaa72e99da8a0736ca449529fe2395059d1196c54943c1622c7f99f29d5">+135/-40</a></td> </tr> </table></td></tr></tr></tbody></table> </details> ___
2026-11-06 00:09:22 +00:00 · 2025-08-08 08:43:19 +05:30
parent 646c6c681f
commit b4cdbddad0
1 changed files with 135 additions and 40 deletions
@@ -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<Self::Break> {
        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);
+    }
 }