fresh-borzoni commented on code in PR #613:
URL: https://github.com/apache/fluss-rust/pull/613#discussion_r3391957957
##########
bindings/python/src/table.rs:
##########
@@ -2780,6 +2842,116 @@ impl LogScanner {
}
}
+/// One-shot bounded scanner over a single bucket.
+///
+/// Obtained via
`table.new_scan().limit(n).create_bucket_batch_scanner(bucket)`.
+/// The scan runs on the first `next_batch()` and yields its single batch once,
+/// then is spent. Honors the configured limit and any projection.
+#[pyclass]
+pub struct BatchScanner {
+ inner: Arc<Mutex<LimitBatchScanner>>,
+ bucket: TableBucket,
+ projected_schema: SchemaRef,
+}
+
+#[pymethods]
+impl BatchScanner {
+ /// The bucket scanned by this batch scanner.
+ #[getter]
+ fn bucket(&self) -> TableBucket {
+ self.bucket.clone()
+ }
+
+ /// Run the scan and return its batch, or `None` once the scanner is spent.
+ ///
+ /// The scan runs on the first call and is not retried; on error, create a
+ /// new scanner.
+ fn next_batch<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+ let inner = Arc::clone(&self.inner);
+ future_into_py(py, async move {
+ let mut scanner = inner.lock().await;
+ let batch = scanner
+ .next_batch()
+ .await
+ .map_err(|e| FlussError::from_core_error(&e))?;
+ Python::attach(|py| match batch {
+ Some(sb) => Ok(Some(Py::new(py,
RecordBatch::from_scan_batch(sb))?)),
+ None => Ok(None),
+ })
+ })
+ }
+
+ /// Drain the scanner into all of its batches.
+ fn collect_all_batches<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py,
PyAny>> {
+ let inner = Arc::clone(&self.inner);
+ future_into_py(py, async move {
+ let mut scanner = inner.lock().await;
+ let batches = scanner
+ .collect_all_batches()
+ .await
+ .map_err(|e| FlussError::from_core_error(&e))?;
+ Python::attach(|py| {
+ batches
+ .into_iter()
+ .map(|sb| Py::new(py, RecordBatch::from_scan_batch(sb)))
+ .collect::<PyResult<Vec<_>>>()
+ })
+ })
+ }
+
+ /// Drain the scanner into a PyArrow Table (empty, with the projected
schema,
+ /// when the scan yields nothing).
+ fn to_arrow<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+ let inner = Arc::clone(&self.inner);
+ let projected_schema = self.projected_schema.clone();
+ future_into_py(py, async move {
+ let mut scanner = inner.lock().await;
+ let batches = scanner
+ .collect_all_batches()
+ .await
+ .map_err(|e| FlussError::from_core_error(&e))?
+ .into_iter()
+ .map(|sb| Arc::new(sb.into_batch()))
+ .collect();
+ Python::attach(|py| LogScanner::batches_to_arrow_table(py,
batches, &projected_schema))
+ })
+ }
+
+ /// Drain the scanner into a Pandas DataFrame.
+ fn to_pandas<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
Review Comment:
Sure, we can 👍
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
