Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Read Arrow C Stream from Arrow PyCapsule Interface #501

Closed
wants to merge 6 commits into from
Closed

Read Arrow C Stream from Arrow PyCapsule Interface #501

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
0c39d5d
Read from Arrow C Data interface
kylebarron Jul 23, 2024
3db45e0
add comment
kylebarron Jul 23, 2024
9014157
Merge remote-tracking branch 'origin/main' into fork/kylebarron/kyle/…
jonmmease Aug 14, 2024
98aae9c
Finish merge
jonmmease Aug 15, 2024
d690caa
update Cargo.toml
jonmmease Aug 15, 2024
26340fa
fmt
jonmmease Aug 15, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 71 additions & 2 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ chrono-tz = {version = "0.9.0", features=["case-insensitive", "filter-by-regex"]
reqwest = { version = "0.11.22", default-features = false }
tokio = { version = "1.36.0" }
pyo3 = { version = "0.21.1" }
pyo3-arrow = { version = "0.2.0" }
pythonize = { version = "0.21.1" }
prost = { version = "0.12.3" }
prost-types = { version = "0.12.3" }
Expand Down
2 changes: 2 additions & 0 deletions python/vegafusion/vegafusion/runtime.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -215,6 +215,8 @@ def _import_or_register_inline_datasets(self, inline_datasets=None):
pass

imported_inline_datasets[name] = PandasDatasource(value)
elif hasattr(value, "__arrow_c_stream__"):
imported_inline_datasets[name] = value
elif hasattr(value, "__dataframe__"):
# Let polars convert to pyarrow since it has broader support than the raw dataframe interchange
# protocol, and "This operation is mostly zero copy."
Expand Down
6 changes: 5 additions & 1 deletion vegafusion-common/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ description = "Common components required by multiple VegaFusion crates"
license = "BSD-3-Clause"

[features]
pyarrow = [ "pyo3", "arrow/pyarrow",]
pyarrow = [ "pyo3", "arrow/pyarrow", "pyo3-arrow"]
json = [ "serde_json/preserve_order", "arrow/json", "chrono",]
prettyprint = [ "arrow/prettyprint",]
proto = ["datafusion-proto", "datafusion-proto-common"]
Expand Down Expand Up @@ -52,6 +52,10 @@ optional = true
workspace = true
optional = true

[dependencies.pyo3-arrow]
workspace = true
optional = true

[dependencies.jni]
version = "0.21.1"
optional = true
Expand Down
10 changes: 9 additions & 1 deletion vegafusion-common/src/data/table.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,9 +34,10 @@ use {
use {
arrow::pyarrow::{FromPyArrow, ToPyArrow},
pyo3::{
conversion::FromPyObjectBound,
prelude::*,
types::{PyList, PyTuple},
Bound, PyAny, PyErr, PyObject, Python,
Bound, PyAny, PyErr, PyObject, PyResult, Python,
},
};

Expand Down Expand Up @@ -271,6 +272,13 @@ impl VegaFusionTable {
}
}

#[cfg(feature = "pyarrow")]
pub fn from_arrow_c_stream(table: &Bound<PyAny>) -> PyResult<Self> {
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this a public API or is this an internal API called by your Python code? If it's a public API, it might be preferred to overload the existing from_pyarrow method.

You can check for both the pycapsule interface and fall back to pyarrow in a single method. That's a recommended approach here.

You can do this by first trying to extract a PyTable; this will work for any recent pyarrow table object (or pyarrow record batch reader). And then you can fallback to your existing from_pyarrow code.

    pub fn from_arrow(table: &Bound<PyAny>) -> PyResult<Self> {
        if let Ok(table) = data.extract::<PyTable>() {
            let (batches, schema) = table.into_inner();
            Ok(VegaFusionTable::try_new(schema, batches)?)
        } else {
            Self::from_pyarrow(table)
        }
    }

let pytable = pyo3_arrow::PyTable::from_py_object_bound(table.as_borrowed())?;
Comment on lines +276 to +277
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If you keep this method the same to only support PyTable, PyTable implements FromPyObject, so you can make it the parameter directly.

Suggested change
pub fn from_arrow_c_stream(table: &Bound<PyAny>) -> PyResult<Self> {
let pytable = pyo3_arrow::PyTable::from_py_object_bound(table.as_borrowed())?;
pub fn from_arrow_c_stream(table: pyo3_arrow::PyTable) -> PyResult<Self> {

let (batches, schema) = pytable.into_inner();
Ok(VegaFusionTable::try_new(schema, batches)?)
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I never know whether batches, schema or schema, batches is a better ordering 🥲

}

#[cfg(feature = "pyarrow")]
pub fn from_pyarrow(pyarrow_table: &Bound<PyAny>) -> std::result::Result<Self, PyErr> {
// Extract table.schema as a Rust Schema
Expand Down
4 changes: 4 additions & 0 deletions vegafusion-python-embed/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -190,6 +190,10 @@ impl PyVegaFusionRuntime {
.scan_py_datasource(inline_dataset.to_object(py)),
)?;
VegaFusionDataset::DataFrame(df)
} else if inline_dataset.hasattr("__arrow_c_stream__")? {
// Import via Arrow PyCapsule Interface
let table = VegaFusionTable::from_arrow_c_stream(inline_dataset)?;
VegaFusionDataset::from_table_ipc_bytes(&table.to_ipc_bytes()?)?
} else {
// Assume PyArrow Table
// We convert to ipc bytes for two reasons:
Expand Down
Loading