fix(connect): count returns the number of rows in the table

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)
 }

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"