From 68553904c76f47ac35ac477e0a8a4dfd4f42cb99 Mon Sep 17 00:00:00 2001
From: Krishna Vishal <krishnavishal2008@gmail.com>
Date: Sat, 18 Jan 2025 11:18:12 +0530
Subject: [PATCH] Converted the unconditional unwrap to a match which handles
 the case when the function is COUNT and args are None and replaces the args.
 Solves https://github.com/tursodatabase/limbo/issues/725

---
 core/translate/select.rs | 24 +++++++++++++++++++++---
 1 file changed, 21 insertions(+), 3 deletions(-)

diff --git a/core/translate/select.rs b/core/translate/select.rs
index 157f55d9d..e75201d43 100644
--- a/core/translate/select.rs
+++ b/core/translate/select.rs
@@ -1,3 +1,5 @@
+use std::ops::Deref;
+
 use super::emitter::emit_program;
 use super::expr::get_name;
 use super::plan::SelectQueryType;
@@ -9,10 +11,10 @@ use crate::translate::planner::{
     parse_where, resolve_aggregates, OperatorIdCounter,
 };
 use crate::util::normalize_ident;
-use crate::SymbolTable;
 use crate::{schema::Schema, vdbe::builder::ProgramBuilder, Result};
-use sqlite3_parser::ast;
+use crate::{LimboError, SymbolTable};
 use sqlite3_parser::ast::ResultColumn;
+use sqlite3_parser::ast::{self, Expr};
 
 pub fn translate_select(
     program: &mut ProgramBuilder,
@@ -116,9 +118,25 @@ pub fn prepare_select_plan(
                                     args_count,
                                 ) {
                                     Ok(Func::Agg(f)) => {
+                                        let count_args = vec![ast::Expr::Literal(
+                                            ast::Literal::Numeric("1".to_string()),
+                                        )];
+                                        let agg_args: Result<Vec<Expr>, LimboError> = match args {
+                                            // if args is None and its COUNT
+                                            None if name.0.to_uppercase() == "COUNT" => {
+                                                Ok(count_args)
+                                            }
+                                            // if args is None and the function is not COUNT
+                                            None => crate::bail_parse_error!(
+                                                "Aggregate function {} requires arguments",
+                                                name.0
+                                            ),
+                                            Some(args) => Ok(args.clone()),
+                                        };
+
                                         let agg = Aggregate {
                                             func: f,
-                                            args: args.as_ref().unwrap().clone(),
+                                            args: agg_args?.clone(),
                                             original_expr: expr.clone(),
                                         };
                                         aggregate_expressions.push(agg.clone());