[FEAT] connect: read/write -> csv, write -> json

src/daft-connect/src/op/execute/write.rs

@@ -54,9 +54,12 @@ impl Session {
                 bail!("Source is required");
             };
 
-            if source != "parquet" {
-                bail!("Unsupported source: {source}; only parquet is supported");
-            }
+            let file_format = match &*source {
+                "parquet" => FileFormat::Parquet,
+                "csv" => FileFormat::Csv,
+                "json" => FileFormat::Json,
+                _ => bail!("Unsupported source: {source}; only parquet and csv are supported"),
+            };
 
             let Ok(mode) = SaveMode::try_from(mode) else {
                 bail!("Invalid save mode: {mode}");
@@ -113,7 +116,7 @@ impl Session {
 
             plan.builder = plan
                 .builder
-                .table_write(&path, FileFormat::Parquet, None, None, None)
+                .table_write(&path, file_format, None, None, None)
                 .wrap_err("Failed to create table write plan")?;
 
             let optimized_plan = plan.builder.optimize()?;

src/daft-connect/src/translation/logical_plan/read/data_source.rs

@@ -1,5 +1,5 @@
 use daft_logical_plan::LogicalPlanBuilder;
-use daft_scan::builder::ParquetScanBuilder;
+use daft_scan::builder::{CsvScanBuilder, ParquetScanBuilder};
 use eyre::{bail, ensure, WrapErr};
 use tracing::warn;
 
@@ -18,10 +18,6 @@ pub async fn data_source(
         bail!("Format is required");
     };
 
-    if format != "parquet" {
-        bail!("Unsupported format: {format}; only parquet is supported");
-    }
-
     ensure!(!paths.is_empty(), "Paths are required");
 
     if let Some(schema) = schema {
@@ -36,10 +32,23 @@ pub async fn data_source(
         warn!("Ignoring predicates: {predicates:?}; not yet implemented");
     }
 
-    let builder = ParquetScanBuilder::new(paths)
-        .finish()
-        .await
-        .wrap_err("Failed to create parquet scan builder")?;
+    let plan = match &*format {
+        "parquet" => ParquetScanBuilder::new(paths)
+            .finish()
+            .await
+            .wrap_err("Failed to create parquet scan builder")?,
+        "csv" => CsvScanBuilder::new(paths)
+            .finish()
+            .await
+            .wrap_err("Failed to create csv scan builder")?,
+        "json" => {
+            // todo(completeness): implement json reading
+            bail!("json reading is not yet implemented");
+        }
+        other => {
+            bail!("Unsupported format: {other}; only parquet and csv are supported");
+        }
+    };
 
-    Ok(builder)
+    Ok(plan)
 }

tests/connect/test_csv.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+import tempfile
+import shutil
+import os
+
+
+def test_write_csv(spark_session):
+    # Create a temporary directory
+    temp_dir = tempfile.mkdtemp()
+    try:
+        # Create DataFrame from range(10)
+        df = spark_session.range(10)
+
+        # Write DataFrame to CSV directory
+        csv_dir = os.path.join(temp_dir, "test.csv")
+        df.write.csv(csv_dir)
+
+        # List all files in the CSV directory
+        csv_files = [f for f in os.listdir(csv_dir) if f.endswith('.csv')]
+        print(f"CSV files in directory: {csv_files}")
+
+        # Assert there is at least one CSV file
+        assert len(csv_files) > 0, "Expected at least one CSV file to be written"
+
+        # Read back from the CSV directory (not specific file)
+        df_read = spark_session.read.csv(csv_dir)
+
+        # Verify the data is unchanged
+        df_pandas = df.toPandas()
+        df_read_pandas = df_read.toPandas()
+        assert df_pandas["id"].equals(df_read_pandas["id"]), "Data should be unchanged after write/read"
+
+    finally:
+        # Clean up temp directory
+        shutil.rmtree(temp_dir)