Add blocksize to DocumentDataset.read_* that uses dask_cudf.read_*

[praateekmahajan]

This PR introduces a new codepath that leverages dask_cudf.read_json / read_parquet directly rather than our existing from_map implementation.

Important

The newer implementation supports fewer use-cases compared to original, however it supports the most frequently used implementations (i.e jsonl and parquet files without add_filename)

New code gotchya's

1. blocksize > filesize doesn't work when backend='cpu' and filetype='jsonl' since pandas doesn't support reading multiple files together.
2. If underlying data has inconsistent schema / metadata the read might fail. See Supporting inconsistent schemas in read_json dask/dask#11595

Existing (fpp implementtion) unexpected behavior

As we added more tests for existing code, uncovered a few cases where it doesn't work as expected

1. filename column can't be selected when backend='pandas', filetype='parquet' and add_columns=True
2. Specifying input_meta with pandas jsonl didn't select only subset of columns. That behavior has been fixed.

Differences between Old vs New

Discussion Points	New Implementation	Existing Implementation
Underlying API	dask_cudf.read_*	dd.from_map(read_single_partition)
backend	Works with pandas / cudf	Works with pandas / cudf
filetype	Only supports jsonl and parquet	Supports json, along with jsonl and parquet
add_filename	Only when filetype is jsonl	Supports all filetypes
input_meta	Returns only columns in input_meta.keys()	Returns only columns in input_meta.keys()
meta as **kwarg	Not required as the first file is used to parse the schema	Is required otherwise can result in OOM (see benchmark row 2 below)
Inconsitent Schema	Doesn't always work (see dask/dask#11595)	Is supported

Follow up Issues

1. filename column inaccessible with pandas backend and parquet #427
2. Supporting inconsistent schemas in read_json dask/dask#11595
3. Add include_path_column to read_parquet dask/dask#6575
4. https://github.com/rapidsai/cudf/pull/17554/files

Benchmarking

Reading 6000 files of ~25mb each, i.e ~145gb over 8GPUs

add_filename	partition_size	input_meta	Using dask.read_json #285	Providing meta in dask.from_map #291
False	2gb	Specified	24.9 s ± 330 ms	25.9 s ± 520 ms
False	2gb	None	24.9 s ± 470 ms	OOM
True	2gb	Specified	55 s ± 177 ms	53.2 s ± 350 ms per loop
True	2gb	None	54.8 s ± 248 ms	64s ± 289 ms per loop

Using dask.read_json #285	Providing meta in dask.from_map #291
First two are add_filename=False, latter two are True where we see a lower utilization	The first one is add_filename=False, and the latter are True where we see a lower utilization

Usage

# Add snippet demonstrating usage

Checklist

• I am familiar with the Contributing Guide.
• New or Existing tests cover these changes.
• The documentation is up to date with these changes.

[praateekmahajan]

Observed an interesting thing here with pandas when "id" : f"{file_id}_{record_id}" then pd.read_json(.., lines=True).dtypes["id"] returns int

[praateekmahajan]

Marked it as xfail but this is something that we should look at. I would imagine that the errors in test_read_data_select_columns are related, but I could be wrong (or just tired looking at the same errors)

[sarahyurick]

This is demonstrating the note in the README:
"filename column cannot be selected when backend="pandas", filetype="parquet" and add_columns=True"
?

This is saying that the "filename" column does not get added when backend="pandas" and file_type="parquet"? Or are you just saying it gets dropped by the select_columns function?

[praateekmahajan]

df.columns doesn't return filename as a column, and I believe that's because filename is a reserved field in the pyarrow schema. However if we were to change our filename field to file_name / path things would work as expected. So if it sounds good I'll create a followup issue and do the renaming in the next PR (and also remove the xfail)

[sarahyurick]

Yes, that works for me. Thanks for the explanation.

[praateekmahajan]

The sorting also isn't respected under the read_files_fpp path. Removing the sorting still results in test_read_data_select_columns failures but far fewer than when we sort.

[VibhuJawa]

I like this a lot, we can have underlying data that has columns with order not respected. Doing this makes a lot of sense especially for formats like jsonl

[praateekmahajan]

These are the tests that are failing right now

[praateekmahajan]

Please TAL here as the behavior of when we pass input_meta results in different outputs depending on fpp vs blocksize

[sarahyurick]

Okay, I think I understand what is being demonstrated here.

Even though Pandas read_data_fpp and Pandas/cuDF read_data_blocksize cannot prune columns while reading, can't we still have it prune the columns after reading? I know that doesn't save on I/O but then at least the user is getting the resulting DataFrame that they would expect.

[praateekmahajan]

That's a good point. Would you want to do a select_columns([input_meta.keys()]) when input_type=='jsonl' (since input_meta in it's current setting only supports json)? We can do that, though I worry that it'll become convolved along with columns (and maybe even add_filename)

[sarahyurick]

I think as long as we are thoroughly checking all the conditions we expect to have to do this, it should be okay. Happy to discuss this offline if anything is unclear here.

[praateekmahajan]

@sarahyurick I made the changes in the last commit however the test_io now fails due to since the test does

            dataset = DocumentDataset.read_json(
                temp_file.name, input_meta='{"id": "float"}'
            )

        output_meta = str({col: str(dtype) for col, dtype in dataset.df.dtypes.items()})

        expected_meta = (
            "{'date': 'datetime64[ns, UTC]', 'id': 'float64', 'text': 'object'}"
        )

I was going to change it but just wanted to confirm if it's okay. FWIW, that test only passes in pandas backend case, but would've failed if backend would've been cudf, i.e the new tests are more thorough and caught the behavior difference.

[sarahyurick]

Added a couple more nits, then should be good on my end. Thanks!

[sarahyurick]

Suggested change

	match="add_filename and blocksize cannot be set at the same time for parquet files",
	match="add_filename and blocksize cannot be set at the same time for Parquet files",

[sarahyurick]

#434 should fix the gpuCI build errors, then we can make sure all GPU tests pass, then merge!