[rust] Stream adapter for bounded readers (DataFusion interop) (#572)

crates/fluss/src/client/table/batch_scanner.rs

@@ -34,7 +34,9 @@ use crate::record::{
 };
 use crate::row::FixedSchemaDecoder;
 use crate::rpc::RpcClient;
+use crate::client::table::stream::{RecordBatchStream, ScanBatchStream};
 use crate::rpc::message::LimitScanRequest;
+use arrow_schema::SchemaRef;
 use arrow::array::RecordBatch;
 use arrow::compute::concat_batches;
 use byteorder::{ByteOrder, LittleEndian};
@@ -50,6 +52,7 @@ pub struct LimitBatchScanner {
     bucket: TableBucket,
     /// Taken on the first `next_batch` to run the scan; `None` afterward.
     pending: Option<PendingScan>,
+    schema: SchemaRef,
 }
 
 /// Request inputs captured at creation, consumed by the first `next_batch`.
@@ -71,8 +74,13 @@ impl LimitBatchScanner {
         projected_fields: Option<Vec<usize>>,
         bucket: TableBucket,
         limit: i32,
-    ) -> Self {
-        Self {
+    ) -> Result<Self> {
+        let full_schema = to_arrow_schema(table_info.get_row_type())?;
+        let schema = match &projected_fields {
+            None => full_schema,
+            Some(fields) => ArrowReadContext::project_schema(full_schema, fields)?,
+        };
+        Ok(Self {
             bucket,
             pending: Some(PendingScan {
                 rpc_client,
@@ -82,7 +90,13 @@ impl LimitBatchScanner {
                 projected_fields,
                 limit,
             }),
-        }
+            schema,
+        })
+    }
+
+    /// Returns the Arrow schema for batches produced by this scanner.
+    pub fn schema(&self) -> SchemaRef {
+        self.schema.clone()
     }
 
     /// Runs the scan on the first call and returns its batch, then `None`. Not
@@ -107,6 +121,82 @@ impl LimitBatchScanner {
     pub fn bucket(&self) -> &TableBucket {
         &self.bucket
     }
+
+    /// Converts this scanner into a [`ScanBatchStream`] yielding [`ScanBatch`] items.
+    ///
+    /// The returned stream is fused: once `next_batch` returns `Ok(None)` or `Err(_)`,
+    /// subsequent polls return `None`.
+    pub fn into_stream(self) -> ScanBatchStream {
+        let schema = self.schema();
+        let stream = futures::stream::unfold((self, false), |(mut scanner, terminated)| async move {
+            if terminated {
+                return None;
+            }
+            match scanner.next_batch().await {
+                Ok(Some(batch)) => Some((Ok(batch), (scanner, false))),
+                Ok(None) => None,
+                Err(e) => Some((Err(e), (scanner, true))),
+            }
+        });
+        ScanBatchStream::new(Box::pin(stream), schema)
+    }
+
+    /// Converts this scanner into a [`RecordBatchStream`] yielding bare Arrow [`RecordBatch`]es.
+    ///
+    /// This is an async bridge for DataFusion (`RecordBatchStreamAdapter`) and PyArrow interop.
+    pub fn into_record_batch_stream(self) -> RecordBatchStream {
+        self.into_stream().into_record_batch_stream()
+    }
+
+    /// Convert this async scanner into a synchronous [`arrow::record_batch::RecordBatchReader`].
+    ///
+    /// The returned adapter calls [`tokio::runtime::Handle::block_on`] on each
+    /// iterator step. **Do not** call this from inside a Tokio worker thread
+    /// while the same runtime is driving async work (nested `block_on` can
+    /// panic or deadlock). Prefer [`next_batch`](LimitBatchScanner::next_batch)
+    /// in async Rust code. This is intended for sync/FFI boundaries (C++, some
+    /// Python call paths).
+    pub fn to_record_batch_reader(
+        self,
+        handle: tokio::runtime::Handle,
+    ) -> SyncRecordBatchLimitScanner {
+        SyncRecordBatchLimitScanner {
+            scanner: self,
+            handle,
+        }
+    }
+}
+
+/// Synchronous adapter that implements [`arrow::record_batch::RecordBatchReader`]
+/// for [`LimitBatchScanner`].
+///
+/// Created via [`LimitBatchScanner::to_record_batch_reader`].
+/// Blocks the current thread on each `next()` call using the provided
+/// Tokio runtime handle.
+///
+/// The iterator yields plain [`RecordBatch`]es (bucket/offset metadata from
+/// [`ScanBatch`] is stripped to satisfy the Arrow trait contract).
+pub struct SyncRecordBatchLimitScanner {
+    scanner: LimitBatchScanner,
+    handle: tokio::runtime::Handle,
+}
+
+impl Iterator for SyncRecordBatchLimitScanner {
+    type Item = std::result::Result<RecordBatch, arrow::error::ArrowError>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self.handle.block_on(self.scanner.next_batch()) {
+            Ok(Some(scan_batch)) => Some(Ok(scan_batch.into_batch())),
+            Ok(None) => None,
+            Err(e) => Some(Err(arrow::error::ArrowError::ExternalError(Box::new(e)))),
+        }
+    }
+}
+
+impl arrow::record_batch::RecordBatchReader for SyncRecordBatchLimitScanner {
+    fn schema(&self) -> SchemaRef {
+        self.scanner.schema()
+    }
 }
 
 /// Resolves the leader, sends the `LimitScanRequest`, and decodes the response

crates/fluss/src/client/table/mod.rs

@@ -34,16 +34,18 @@ mod read_context_resolver;
 mod reader;
 mod remote_log;
 mod scanner;
+mod stream;
 mod upsert;
 
 pub use append::{AppendWriter, TableAppend};
-pub use batch_scanner::LimitBatchScanner;
+pub use batch_scanner::{LimitBatchScanner, SyncRecordBatchLimitScanner};
 pub use lookup::{LookupResult, Lookuper, PrefixKeyLookuper, TableLookup, TablePrefixLookup};
 pub use reader::{RecordBatchLogReader, SyncRecordBatchLogReader};
 pub use remote_log::{
     DEFAULT_REMOTE_FILE_DOWNLOAD_THREAD_NUM, DEFAULT_SCANNER_REMOTE_LOG_PREFETCH_NUM,
 };
 pub use scanner::{LogScanner, RecordBatchLogScanner, TableScan};
+pub use stream::{RecordBatchStream, ScanBatchStream, SyncRecordBatchStreamReader};
 pub use upsert::{TableUpsert, UpsertWriter};
 
 #[allow(dead_code)]

crates/fluss/src/client/table/reader.rs

@@ -32,6 +32,7 @@
 
 use crate::client::admin::FlussAdmin;
 use crate::client::table::RecordBatchLogScanner;
+use crate::client::table::stream::{RecordBatchStream, ScanBatchStream};
 use crate::error::{Error, Result};
 use crate::metadata::TableBucket;
 use crate::record::ScanBatch;
@@ -225,6 +226,33 @@ impl RecordBatchLogReader {
             handle,
         }
     }
+
+    /// Converts this reader into a [ScanBatchStream] yielding [ScanBatch] items.
+    ///
+    /// The returned stream is fused: once 
+ext_batch returns Ok(None) or Err(_),
+    /// subsequent polls return None.
+    pub fn into_stream(self) -> ScanBatchStream {
+        let schema = self.schema();
+        let stream = futures::stream::unfold((self, false), |(mut reader, terminated)| async move {
+            if terminated {
+                return None;
+            }
+            match reader.next_batch().await {
+                Ok(Some(batch)) => Some((Ok(batch), (reader, false))),
+                Ok(None) => None,
+                Err(e) => Some((Err(e), (reader, true))),
+            }
+        });
+        ScanBatchStream::new(Box::pin(stream), schema)
+    }
+
+    /// Converts this reader into a [RecordBatchStream] yielding bare Arrow [RecordBatch]es.
+    ///
+    /// This is an async bridge for DataFusion (RecordBatchStreamAdapter) and PyArrow interop.
+    pub fn into_record_batch_stream(self) -> RecordBatchStream {
+        self.into_stream().into_record_batch_stream()
+    }
 }
 
 /// Best-effort cleanup when the reader is dropped before all buckets reach

crates/fluss/src/client/table/scanner.rs

@@ -172,15 +172,15 @@ impl<'a> TableScan<'a> {
             self.conn.get_admin()?,
             latest,
         ));
-        Ok(LimitBatchScanner::new(
+        LimitBatchScanner::new(
             self.conn.get_connections(),
             self.metadata.clone(),
             self.table_info,
             schema_getter,
             self.projected_fields,
             table_bucket,
             limit,
-        ))
+        )
     }
 
     /// Projects the scan to only include specified columns by their indices.