-
-
Save igiloh-pinecone/a350cddbcf532df1f72b618a491e2f45 to your computer and use it in GitHub Desktop.
Simplified code for demonstrating issue with `pyo3-asyncio` for struct method
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
// ------------------------ pinecone-core lib------------------------ | |
// This lib builds to a pure Rust lib (not `cdylib`) | |
// This is a single data unit. My python function takes a list of these and passes them to a `tonic`-generated GRPC client | |
#[derive(Debug, Default, Clone)] | |
pub struct Vector{ | |
pub id: String, | |
pub values: Vec<f32>, | |
pub sparse_values: HashMap<u32, f32>, | |
} | |
// Low-level class that transforms input data to GRPC messages and Prost structs | |
#[derive(Debug, Clone)] | |
pub struct DataplaneGrpcClient{ | |
inner: VectorServiceClient<Channel> // This struct is auto-generated by [tonic](https://github.com/hyperium/tonic). | |
// Tonic requires that all methods will use `&mut self` | |
} | |
impl DataplaneGrpcClient{ | |
pub async fn upsert(&mut self, namespace: &str, vectors: &[Vector]) -> Result<u32, tonic::Status> { | |
// Convert input to GRPC message | |
let grpc_vectos: Vec<GrpcVector> = vectors.iter().map(|v| v.into()).collect(); | |
// Call auto-generated tonic client code | |
let res = self.inner.upsert(UpsertRequest{namespace: namespace.into(), vectors: grpc_vectos}).await?; | |
Ok(res.into_inner().upserted_count) | |
} | |
} | |
// Public sturct which exposes functionality to users | |
pub struct Index { | |
dataplane_client: DataplaneGrpcClient, | |
} | |
impl Index { | |
pub async fn upsert(&mut self, | |
namespace: &str, | |
vectors: &Vec<Vector>) -> PineconeResult<Vec<String>> { | |
let ids: Vec<String> = vectors.iter().map(|v| v.id.clone()).collect(); | |
let res = self.dataplane_client.upsert(&namespace, &vectors).await?; | |
Ok(ids) | |
} | |
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
// ------------------------ pinecone-client lib----------------------- | |
// This lib wraps pinecone-core with `pyo3` and builds a `cdylib`. | |
#[pyclass] | |
pub struct PyVector { | |
pub inner: pinecone_core::Vector | |
} | |
impl From<Vector> for pinecone_core::Vector { | |
fn from(vector: Vector) -> pinecone_core::Vector { | |
vector.inner | |
} | |
} | |
#[pyclass] | |
pub struct Index{ | |
inner: Arc<Mutex<core_index::Index>>, | |
} | |
impl Index{ | |
pub fn new(inner: core_index::Index) -> Self { | |
Self{inner: Arc::new(Mutex::new(inner))} | |
} | |
} | |
#[pymethods] | |
impl Index{ | |
pub fn upsert_async<'a>(&mut self, py: Python<'a>, vectors: Vec<PyRef<Vector>>, namespace: &'a str) -> PyResult<&'a PyAny>{ | |
let mut inner_index = self.inner.clone(); | |
let core_vectors: Vec<pinecone_core::Vector> = vectors.into_iter().map(|x| x.borrow().inner.clone()).collect(); | |
pyo3_asyncio::tokio::future_into_py(py, async move { | |
let res = inner_index.lock().await.upsert(&namespace, &core_vectors, None).await?; | |
Ok(res) | |
// Also tried: | |
// Ok(Python::with_gil(|py| res.into_py(py))) | |
}) | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment