Skip to content

Commit

Permalink
feat: Use FFI to extract Series from different Polars binaries (#18964)
Browse files Browse the repository at this point in the history
  • Loading branch information
ritchie46 authored Sep 27, 2024
1 parent cafc163 commit 74b5307
Showing 1 changed file with 41 additions and 4 deletions.
45 changes: 41 additions & 4 deletions crates/polars-python/src/map/lazy.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
use polars::prelude::*;
use pyo3::ffi::Py_uintptr_t;
use pyo3::prelude::*;
use pyo3::types::PyList;
use pyo3::types::{PyDict, PyList};

use crate::py_modules::POLARS;
use crate::series::PySeries;
Expand Down Expand Up @@ -42,9 +43,45 @@ impl ToSeries for PyObject {
}
},
};
let pyseries = py_pyseries.extract::<PySeries>(py).unwrap();
// Finally get the actual Series
Ok(pyseries.series)
let s = match py_pyseries.extract::<PySeries>(py) {
Ok(pyseries) => pyseries.series,
// This happens if the executed Polars is not from this source.
// Currently only happens in PC-workers
// For now use arrow to convert
// Eventually we must use Polars' Series Export as that can deal with
// multiple chunks
Err(_) => {
use polars::export::arrow::ffi;
let kwargs = PyDict::new_bound(py);
kwargs.set_item("in_place", true).unwrap();
py_pyseries
.call_method_bound(py, "rechunk", (), Some(&kwargs))
.map_err(|e| polars_err!(ComputeError: "could not rechunk: {e}"))?;

// Prepare a pointer to receive the Array struct.
let array = Box::new(ffi::ArrowArray::empty());
let schema = Box::new(ffi::ArrowSchema::empty());

let array_ptr = &*array as *const ffi::ArrowArray;
let schema_ptr = &*schema as *const ffi::ArrowSchema;
// SAFETY:
// this is unsafe as it write to the pointers we just prepared
py_pyseries
.call_method1(
py,
"_export_arrow_to_c",
(array_ptr as Py_uintptr_t, schema_ptr as Py_uintptr_t),
)
.map_err(|e| polars_err!(ComputeError: "{e}"))?;

unsafe {
let field = ffi::import_field_from_c(schema.as_ref())?;
let array = ffi::import_array_from_c(*array, field.dtype)?;
Series::from_arrow(field.name, array)?
}
},
};
Ok(s)
}
}

Expand Down

0 comments on commit 74b5307

Please sign in to comment.