Dataset not supporting RecordBatchReader

[Arttii]

Hi, looking at the code and the docs, the dataset should support using a RecordBatchReader, but it does not seem to be possible at the moment. Is this not implemented, or am I missing something?

[wjones127]

Given what you linked, this is in PyArrow, right? It should be working, at least in 7.0.0. We have unit tests for writing with RecordBatchReaders.

What code did you try? And what version of PyArrow?

[Arttii]

@wjones127 thanks for the reply, and sorry for the lack of detail. I am using 7.0.0, but I wanted to create a dataset to read. Basically, I was trying to get a dataset to read from a Flight response. Basically:

reader = client.do_get(fl.Ticket('nyc/2015')).to_reader()
dataset = ds.dataset(reader)

Is something like this planned, or I am just misunderstanding it?

[wjones127]

That should be working just like you wrote, and is definitely an intended us case. Could you give more detail on the error you are getting?

[Arttii]

Running this:

pyarrow.__version__
df = pd.DataFrame({"a": [1, 2, 3]})
table = pyarrow.Table.from_pandas(df)
userdata_parquet_dataset= pyarrow.dataset.dataset(table)
batches=[r for r in userdata_parquet_dataset.to_batches()]
reader=pyarrow.dataset.Scanner.from_batches(batches,userdata_parquet_dataset.schema).to_reader()
dataset=pyarrow.dataset.dataset(reader)

gets me this:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Input In [15], in <cell line: 7>()
      5 batches=[r for r in userdata_parquet_dataset.to_batches()]
      6 reader=pyarrow.dataset.Scanner.from_batches(batches,userdata_parquet_dataset.schema).to_reader()
----> 7 dataset=pyarrow.dataset.dataset(reader)

File ~/workdir/duckdb/.venv/lib/python3.9/site-packages/pyarrow/dataset.py:687, in dataset(source, schema, format, filesystem, partitioning, partition_base_dir, exclude_invalid_files, ignore_prefixes)
    685     return _in_memory_dataset(source, **kwargs)
    686 else:
--> 687     raise TypeError(
    688         'Expected a path-like, list of path-likes or a list of Datasets '
    689         'instead of the given type: {}'.format(type(source).__name__)
    690     )

TypeError: Expected a path-like, list of path-likes or a list of Datasets instead of the given type: RecordBatchReader

[wjones127]

Sorry, I read your question wrong initially; I thought you were referring to write_dataset, which definitely does accept a RBR.

It sounds like the lines that handled RBR (introduced in this PR) went missing. 😕 I'll file a Jira issue to get those readded.

[wjones127]

Actually I think is change in API is intentional, given #10070. We just forgot to update the docstring (which I'll create a Jira for).

Instead, you should convert the RecordBatchReader directly into a scanner using Scanner.from_batches(reader).

[Arttii]

I see, thanks for clarifying. Is there no way to have a rereadable one-shot source? I was working on reading a returned reader from a Flight server in DuckDB and I there was no other good way from the API to support this at the moment. Or perhaps the data needs to be reread from the server.

[wjones127]

Is there no way to have a rereadable one-shot source?

If you want to make a one-shot source re-readable, you'll have to collect it into a table. That might work for you if it's not larger than memory.

But I think the idea of a (re-readable) dataset that wraps a flight data source has been something we've thought of, but haven't yet implemented. Here is the relevant Jira issue: https://issues.apache.org/jira/browse/ARROW-10524.

[Arttii]

Oke thanks for the help! I will look into the issue and see if i could try implementing it.