[Python][Rust] Create extension point in python for Dataset/Scanner

[changhiskhan]

Describe the enhancement requested

As the Arrow ecosystem grows ever richer, desire paths emerge :)

Integrating Arrow based projects written in Rust works great across the C data interface. But it doesn't allow lazy execution or pushdowns in the same way that pyarrow Dataset/Scanner's do.

My proposal here is to expose Dataset/Scanner python abc's with s.t. rust libraries can extend via pyo3+python so higher level tooling (like duckdb for example, can query these without having to transfer the whole Table into memory first).

In keeping with the same principles as the C data interface, I think it would be sufficient for this python interface to be very minimal: Dataset with schema and scanner methods, Scanner with projected_schema and to_reader methods. The to_reader should return a RecordBatchReader which would then link pyo3 datasets into the Arrow C data interface.

Thanks for your consideration!

Component(s)

Python

[changhiskhan]

to be precise, this is already doable today as a workaround, what i'm proposing here is:

1. formalize this extension point
2. if possible, make the top level abstraction pure python, so subclasses doesn't need to deal with cython etc if coming from Rust

[lidavidm]

From a quick read, is this about formalizing an extension point per se or actually just an interface? Some way to say, "given a filter/column selection, please try to push these down and give me a C Data Interface stream"

[wjones127]

I've had similar challenges with supporting datasets in delta-rs. Another aspect you'll need to think about is supporting Filesystems. In Rust, that means calling into Python functions, which I fear can be sub-optimal because of the GIL. I don't think there's a practical way to directly access the underlying C++ implemented FS unless we made the ABI stable (which I don't see us doing in the foreseeable future).

A route I'm exploring right now is using the ADBC as a stable ABI for pushing down scan queries to storage formats and systems. It probably make more sense for table formats like Delta Lake, which have database-like semantics, than file formats like Lance (which I assume is the projection with the use case you are discussing).

if possible, make the top level abstraction pure python, so subclasses doesn't need to deal with cython etc if coming from Rust

It's harder, but part of me would prefer a stable C ABI, because it would mean the extension could be used in any language, not just Python.

[changhiskhan]

From a quick read, is this about formalizing an extension point per se or actually just an interface? Some way to say, "given a filter/column selection, please try to push these down and give me a C Data Interface stream"

You're right. Good question. I think just the interface actually, so this issue is incorrectly named. Ie Rust packages could implement the same interface that higher level tooling will call, without necessarily needing to inherit from the cython classes defining Dataset/Scanner today.

[changhiskhan]

I've had similar challenges with supporting datasets in delta-rs. Another aspect you'll need to think about is supporting Filesystems. In Rust, that means calling into Python functions, which I fear can be sub-optimal because of the GIL. I don't think there's a practical way to directly access the underlying C++ implemented FS unless we made the ABI stable (which I don't see us doing in the foreseeable future).

Does object store rs work for this?

A route I'm exploring right now is using the ADBC as a stable ABI for pushing down scan queries to storage formats and systems. It probably make more sense for table formats like Delta Lake, which have database-like semantics, than file formats like Lance (which I assume is the projection with the use case you are discussing).

if possible, make the top level abstraction pure python, so subclasses doesn't need to deal with cython etc if coming from Rust

Yeah most def pure python

It's harder, but part of me would prefer a stable C ABI, because it would mean the extension could be used in any language, not just Python.

[lidavidm]

It sounds like what you basically want is a

struct DataFrameProducer {
    int read(struct DataFrameProducer* self, ??? filters, ??? selection, struct ArrowArrayStream* out);
};

with corresponding wrappers/carriers in Python, Rust, Go, etc. Then this can be fed into DuckDB, Ballista, Acero, etc. and can be produced by ADBC, Acero, DuckDB, etc.

You can have nearly both "pure Python" and "C ABI", I think. A Python-level interface could be transformed by PyArrow into the C ABI and vice versa. Where possible, the Python-level interface should let you 'extract' the underlying C ABI, if it exists, but otherwise we can push the responsibility of the GIL and such into PyArrow (or something like that). (So basically, shove all the non-Python code into PyArrow.)

The question is what the filter/selection format should be; ideally it would be language agnostic and implementation agnostic and so Dataset's expressions aren't great there.

[wjones127]

The question is what the filter/selection format should be; ideally it would be language agnostic and implementation agnostic and so Dataset's expressions aren't great there.

Substrait recently merged a message type for standalone expression lists. :)

substrait-io/substrait#405

[westonpace]

What if we created a read type in substrait for datasets? It would require encoding file system options in Substrait as well. However, this seems like a valuable feature.

Then the C API / interface is just the typical substrait consumer API of "substrait plan in, record batch reader out"

[jorisvandenbossche]

I have been recently thinking about this as well, triggered by our work to support the DataFrame Interchange protocol in pyarrow (#14804, https://data-apis.org/dataframe-protocol/latest/purpose_and_scope.html).
For some context, this "interchange protocol" is a python specific one to allow interchanging dataframe-data, typically between different dataframe libraries (eg convert a dataframe of library1 into a dataframe of library2 without both libraries having to know about each other). Essentially, this is very similar to our C Data Interface (it specifies how to share the memory, and in the end also has pointers to buffers of a certain size, just like the C struct), but because it is a python API, it gives you some functionality on top of those raw buffers (or, before you get to the buffers).

As illustration, using latest pyarrow (using a pyarrow.Table, so all data is already in memory, but we could add the same to something backed by a not-yet materialized stream or data on disk):

>>> table = pa.table({'a': [1, 2, 3], 'b': [4, 5, 6]})
# accessing this interchange object
>>> interchange_object = table.__dataframe__()
>>> interchange_object
<pyarrow.interchange.dataframe._PyArrowDataFrame at 0x7ffb64f45420>
# this doesn't yet need to materialize all the buffers, but you can inspect metadata
>>> interchange_object.num_columns()
2
>>> interchange_object.column_names()
['a', 'b']
# you can select a subset (i.e. simple projection)
>>> subset = interchange_object.select_columns_by_name(['a'])
>>> subset.num_columns()
1
# only when actually asking for the buffers of one chunk of a column, the data needs
# to be in memory (to pass a pointer to the buffers)
>>> subset.get_column_by_name("a").get_buffers()
{'data': (PyArrowBuffer({'bufsize': 24, 'ptr': 140717840175104, 'device': 'CPU'}),
  (<DtypeKind.INT: 0>, 64, 'l', '=')),
 'validity': None,
 'offsets': None}

To be clear, I don't propose we do something exactly like that (and personally, I think it's also a missed opportunity for the DataFrame Interchange protocol to not use Arrow for the memory layout specification, and not give access to data as arrow data or using the Arrow C Interface).
But it has some nice properties compared to our "raw" C Data Interface. It allows to first inspect the data (columns, data types) before you actually start materializing/consuming the data chunk by chunk, and allows you to get the data of only a subset of the columns. While when using the C Stream for a Table or generically RecordBatchReader, you directly get the struct with all buffers for all columns. So if you don't need all columns, you need to ensure to first subset the data before accessing the C interface data, but this means you need to know the API of the object you are getting the data from (partly removing the benefit of the generality of the C Interface).

So having some standard Python interface that would give you delayed and queryable (filter/simple projection) access to Arrow C Interface data sounds really interesting, or having this as an additional C ABI (for which we can still provide such a Python interface as well) like David sketched above.

[westonpace]

From an implementation perspective I suspect we can satisfy any of these proposed APIs. If we need to come up with a new API then my preference is Substrait, but if the consensus heads in some other direction I'm fine with that too.

@jorisvandenbossche , your proposal seems fine, but I don't see anything in there for filesystems. I think this is for on-disk data moreso than purely in-memory data. Though I believe your approach could be adapted to include filesystems.

Does object store rs work for this?

Yes, I would assume that object store rs would be able to satisfy this but I'm not familiar with the capabilities. For example, my idea of how this would work in Substrait would be:

# This would be usable as ReadRel::read_type
message Dataset {

  # This is already definedin ReadRel and is basically a list of files
  # and a format object which defines things like delimiter (for CSV)
  LocalFiles files = 0;
  oneof filesystem {
    LocalFilesystem = 1;
    S3Filesystem = 2;
    ExtensionFilesystem = 3;
  }

  message LocalFilesystem {}
  message S3Filesystem {
    string region;
    string client_id;
    string client_secret; // could be omitted if credentials negotiated elsewhere
  }
  message ExtensionFilesystem {
    google.protobuf.Any details;
  }

}

The equivalent C interface would just be structifying those messages.

[jorisvandenbossche]

What's the reason that the potential interface would need to be aware of filesystems?

[westonpace]

I think the goal is to create a C API capable of representing datasets? I think that would require filesystems. But I may be mistaken on either of these points.