mirror of
https://github.com/lancedb/lancedb.git
synced 2026-01-09 05:12:58 +00:00
feat: refactor the query API and add query support to the python async API (#1113)
In addition, there are also a number of changes in nodejs to the docstrings of existing methods because this PR adds a jsdoc linter.
This commit is contained in:
51
python/src/arrow.rs
Normal file
51
python/src/arrow.rs
Normal file
@@ -0,0 +1,51 @@
|
||||
// use arrow::datatypes::SchemaRef;
|
||||
// use lancedb::arrow::SendableRecordBatchStream;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow::{
|
||||
datatypes::SchemaRef,
|
||||
pyarrow::{IntoPyArrow, ToPyArrow},
|
||||
};
|
||||
use futures::stream::StreamExt;
|
||||
use lancedb::arrow::SendableRecordBatchStream;
|
||||
use pyo3::{pyclass, pymethods, PyAny, PyObject, PyRef, PyResult, Python};
|
||||
use pyo3_asyncio::tokio::future_into_py;
|
||||
|
||||
use crate::error::PythonErrorExt;
|
||||
|
||||
#[pyclass]
|
||||
pub struct RecordBatchStream {
|
||||
schema: SchemaRef,
|
||||
inner: Arc<tokio::sync::Mutex<SendableRecordBatchStream>>,
|
||||
}
|
||||
|
||||
impl RecordBatchStream {
|
||||
pub fn new(inner: SendableRecordBatchStream) -> Self {
|
||||
let schema = inner.schema().clone();
|
||||
Self {
|
||||
schema,
|
||||
inner: Arc::new(tokio::sync::Mutex::new(inner)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl RecordBatchStream {
|
||||
pub fn schema(&self, py: Python) -> PyResult<PyObject> {
|
||||
(*self.schema).clone().into_pyarrow(py)
|
||||
}
|
||||
|
||||
pub fn next(self_: PyRef<'_, Self>) -> PyResult<&PyAny> {
|
||||
let inner = self_.inner.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let inner_next = inner.lock().await.next().await;
|
||||
inner_next
|
||||
.map(|item| {
|
||||
let item = item.infer_error()?;
|
||||
Python::with_gil(|py| item.to_pyarrow(py))
|
||||
})
|
||||
.transpose()
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user