[FIX] .count() ... WIP

src/daft-connect/src/lib.rs

@@ -284,6 +284,8 @@ impl SparkConnectService for DaftSparkConnectService {
     ) -> Result<Response<AnalyzePlanResponse>, Status> {
         use spark_connect::analyze_plan_request::*;
         let request = request.into_inner();
+
+        println!("got analyze plan request: {request:#?}");
 
         let AnalyzePlanRequest {
             session_id,

tests/connect/test_alias.py

@@ -12,6 +12,7 @@ def test_alias(spark_session):
 
     # Verify the alias was set correctly
     assert df_renamed.schema != df.schema, "Schema should be changed after alias"
+    assert df_renamed.count() == df.count(), "Row count should be unchanged after alias"
 
     # Verify the data is unchanged but column name is different
     df_rows = df.collect()

tests/connect/test_count.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+
+def test_count(spark_session):
+    # Create DataFrame from range(10)
+    df = spark_session.range(10)
+
+    # Verify count is correct
+    assert df.count() == 10, "DataFrame should have 10 rows"