Hi there, I'm making a contribution for <https://...
# general
n
Hi there, I'm making a contribution for https://github.com/Eventual-Inc/Daft/issues/3786 I think I've done it, but I was wondering if someone can help me explain how testing works. I already see function signatures in
tests/connect/test_io.py
Copy code
@pytest.mark.skip(reason="<https://github.com/Eventual-Inc/Daft/issues/3786>")
def test_write_csv_with_delimiter(make_df, make_spark_df, spark_session, tmp_path):
    pass
But the thing is, I don't know what these arguments mean, and I don't see that function being invoked anywhere else so I don't know where it's being provided those arguments. But the other thing is, its function signature is different from the ones above:
Copy code
def test_csv_basic_roundtrip(make_spark_df, assert_spark_equals, spark_session, tmp_path):
And I think I might need the
assert_spark_equals
function in my scenario, I was wondering if someone is willing to guide me in the right direction.
r
The tests uses pytest which supports fixtures from the conftest (1) and contest (2). So just add
assert_spark_equals
to your test function's signature.
m
Hi @Nathan Cai! Did this help? Or do you want any more assistance here? We ❤️ PRs!
n
Hi there, thank you so much for helping, I'll make the PR soon
Hi @Malcolm Greaves, I just made a PR https://github.com/Eventual-Inc/Daft/pull/5430
m
Wonderful! We'll get someone to review it!
n
Wait, don't
I just realized it only works for empty csv files
m
oh! well please fix that then 🙂 I also reviewed it (https://github.com/Eventual-Inc/Daft/pull/5430#pullrequestreview-3362711289) -- overall very good stuff! Just requesting some changes.
I just realized it only works for empty csv files
What kind of test can you add to this PR to check this? 🤔
n
Hi there Malcolm, thank you so much for your review! I realized that to set a custom delimiter for a non-empty tabular, I need to dig into the Rust side of things. The problem is, I can't find where it actually writes into a CSV file. I assume it's here: In
DataFrame.write_csv
Copy code
builder = self._builder.write_tabular(
            root_dir=root_dir,
            partition_cols=cols,
            write_mode=WriteMode.from_str(write_mode),
            file_format=FileFormat.Csv,
            io_config=io_config,
)
In the
LogicPlanBuilder
class, I looked deeper into the rust bindings, and it calls the table_write method of the LogicPlanBuilder, But in
table_write
method, I don't see anywhere that handles the logic of writing to files,
Copy code
pub fn table_write(
        &self,
        root_dir: &str,
        write_mode: WriteMode,
        file_format: FileFormat,
        partition_cols: Option<Vec<ExprRef>>,
        compression: Option<String>,
        io_config: Option<IOConfig>,
    ) -> DaftResult<Self> {
        let expr_resolver = ExprResolver::default();

        let partition_cols = partition_cols
            .map(|cols| expr_resolver.resolve(cols, self.plan.clone()))
            .transpose()?;

        let sink_info = SinkInfo::OutputFileInfo(OutputFileInfo::new(
            root_dir.into(),
            write_mode,
            file_format,
            partition_cols,
            compression,
            io_config,
        ));

        let logical_plan: LogicalPlan =
            ops::Sink::try_new(self.plan.clone(), sink_info.into())?.into();
        Ok(self.with_new_plan(logical_plan))
    }
The next step I'd assume is under
logical_plan
but to me it seems to simply return the fields when I look into the
try_new
function?
m
Daft design crash course -- Daft is lazy. All of those operations are actually building a plan (logical plan) which will describe what needs to be done. Nothing is actually ran until one adds a materializing operation -- the most common ones are
collect
(bring everything into memory on a single machine),
show
(get the first 8 records, print them, then stop), and
write_<XYZ>
(write the results of the dataframe as format
<XYZ>
). So you're not finding these things when you goto definition because you're just seeing the logical plan being constructed.
(I'll get you pointers to the right files in a moment...)
The actual CSV writer implementation is here: https://github.com/Eventual-Inc/Daft/blob/main/daft/io/writer.py#L179 There's a long and winding road to get here 😅 : The Rust code for writing CSVs in
daft.DataFrame.write_csv
follows this path: ### 1. Python Entry Point -
daft/dataframe/dataframe.py:846
-
DataFrame.write_csv()
method ### 2. Rust Physical Plan Node -
src/daft-physical-plan/src/ops/csv.rs
- Defines
TabularWriteCsv
struct that represents the CSV write operation in the physical plan ### 3. Rust Writer Factory -
src/daft-writers/src/lib.rs
-
make_physical_writer_factory()
creates writer factories based on file format -
src/daft-writers/src/physical.rs
-
PhysicalWriterFactory
routes to the appropriate writer ### 4. Actual CSV Writing Implementation For CSV writes, the implementation uses PyArrow (not native Rust): - Rust side:
src/daft-writers/src/pyarrow.rs:48-75
-
PyArrowWriter::new_csv_writer()
creates a Python-based CSV writer - Python side:
daft/io/writer.py:213-246
-
CSVFileWriter
class uses PyArrow's
pacsv.CSVWriter
to write CSV files The key implementation is in the
CSVFileWriter
class at
daft/io/writer.py:213
, which: - Uses PyArrow's
pacsv.CSVWriter
(Arrow CSV writer) - Calls
write_table()
on PyArrow tables - Handles file creation, partitioning, and metadata Note: Unlike Parquet and JSON which have native Rust implementations, CSV writing currently delegates to PyArrow through the Python/Rust bridge (PyO3).
n
Hi there, thank you so much for the details, but a couple of questions. There's a lot of layers in the stack trace I would need to pass the delimiter to, PyArrow's CSVWriter I need to know every part of the stack calls to pass the delimiter argument, you said that Rust doesn't have native CSV support but there is a crate for that? https://docs.rs/csv/latest/csv/ Do you mean CSV support isn't implemented in Rust's STL? I do need
m
you said that Rust doesn't have native CSV support but there is a crate for that?
Ah sorry, wasn't clear enough I suppose -- it's not that Rust can't do it. It's just that we implemented a different way to do this in Daft.
Addressed some more in your DM -- let me know there (or here) if you have remaining questions or comments!