Skip to content

Instantly share code, notes, and snippets.

@igiloh-pinecone
Last active March 12, 2023 16:51
Show Gist options
  • Save igiloh-pinecone/a350cddbcf532df1f72b618a491e2f45 to your computer and use it in GitHub Desktop.
Save igiloh-pinecone/a350cddbcf532df1f72b618a491e2f45 to your computer and use it in GitHub Desktop.
Simplified code for demonstrating issue with `pyo3-asyncio` for struct method
// ------------------------ 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)
}
}
// ------------------------ 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