From fdc183041bab6b6aec8b76d9658240953f444ef4 Mon Sep 17 00:00:00 2001
From: Andrew Gazelka <andrew.gazelka@gmail.com>
Date: Wed, 11 Dec 2024 14:22:46 -0800
Subject: [PATCH] fix(connect): `count` returns the number of rows in the table

---
 .../src/translation/expr/unresolved_function.rs     | 11 ++++++++++-
 tests/connect/test_count.py                         | 13 +++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)
 create mode 100644 tests/connect/test_count.py

diff --git a/src/daft-connect/src/translation/expr/unresolved_function.rs b/src/daft-connect/src/translation/expr/unresolved_function.rs
index af70a1e8f2..959307e779 100644
--- a/src/daft-connect/src/translation/expr/unresolved_function.rs
+++ b/src/daft-connect/src/translation/expr/unresolved_function.rs
@@ -1,4 +1,5 @@
 use daft_core::count_mode::CountMode;
+use daft_schema::dtype::DataType;
 use eyre::{bail, Context};
 use spark_connect::expression::UnresolvedFunction;
 
@@ -79,7 +80,15 @@ pub fn handle_count(arguments: Vec<daft_dsl::ExprRef>) -> eyre::Result<daft_dsl:
 
     let [arg] = arguments;
 
-    let count = arg.count(CountMode::All);
+    let count = if arg == daft_dsl::lit(1_i32) {
+        // Count(Literal(1)) is handled differently by Daft. Generally speaking though, what
+        // this means is we are counting the number of rows in the table.
+        daft_dsl::col("*")
+            .count(CountMode::All)
+            .cast(&DataType::Int64)
+    } else {
+        arg.count(CountMode::All).cast(&DataType::Int64)
+    };
 
     Ok(count)
 }
diff --git a/tests/connect/test_count.py b/tests/connect/test_count.py
new file mode 100644
index 0000000000..1cb48dd47b
--- /dev/null
+++ b/tests/connect/test_count.py
@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+
+def test_count(spark_session):
+    # Create a range using Spark
+    # For example, creating a range from 0 to 9
+    spark_range = spark_session.range(10)  # Creates DataFrame with numbers 0 to 9
+
+    # Convert to Pandas DataFrame
+    count = spark_range.count()
+
+    # Verify the DataFrame has expected values
+    assert count == 10, "DataFrame should have 10 rows"