diff --git a/src/daft-connect/src/op/execute/write.rs b/src/daft-connect/src/op/execute/write.rs index 44696f8164..aaf63fae6e 100644 --- a/src/daft-connect/src/op/execute/write.rs +++ b/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()?; diff --git a/src/daft-connect/src/translation/logical_plan/read/data_source.rs b/src/daft-connect/src/translation/logical_plan/read/data_source.rs index 90164dd0bd..863b5e8f1d 100644 --- a/src/daft-connect/src/translation/logical_plan/read/data_source.rs +++ b/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) } diff --git a/tests/connect/test_csv.py b/tests/connect/test_csv.py new file mode 100644 index 0000000000..18e986ccdf --- /dev/null +++ b/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)