[rust] Add into_stream() to bounded readers

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

@@ -39,6 +39,7 @@ use arrow::array::RecordBatch;
 use arrow::compute::concat_batches;
 use byteorder::{ByteOrder, LittleEndian};
 use bytes::Bytes;
+use futures::Stream;
 use std::collections::HashMap;
 use std::ops::Range;
 use std::sync::Arc;
@@ -103,6 +104,15 @@ impl LimitBatchScanner {
         Ok(batches)
     }
 
+    /// Consume this scanner into a [`Stream`] that yields its single
+    /// [`ScanBatch`], then ends. Shares the streaming contract of
+    /// [`RecordBatchLogReader::into_stream`](crate::client::RecordBatchLogReader::into_stream).
+    pub fn into_stream(self) -> impl Stream<Item = Result<ScanBatch>> + Send {
+        futures::stream::try_unfold(self, |mut scanner| async move {
+            Ok(scanner.next_batch().await?.map(|batch| (batch, scanner)))
+        })
+    }
+
     /// The bucket scanned by this `LimitBatchScanner`.
     pub fn bucket(&self) -> &TableBucket {
         &self.bucket

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

@@ -38,6 +38,7 @@ use crate::record::ScanBatch;
 use crate::rpc::message::OffsetSpec;
 use arrow::record_batch::RecordBatch;
 use arrow_schema::SchemaRef;
+use futures::Stream;
 use log::warn;
 use std::collections::{HashMap, VecDeque};
 use std::time::Duration;
@@ -208,6 +209,18 @@ impl RecordBatchLogReader {
         }
     }
 
+    /// Consume this reader into a [`Stream`] of [`ScanBatch`]es, one per
+    /// [`next_batch`](Self::next_batch) call, ending when all buckets reach
+    /// their stopping offsets or on the first error.
+    ///
+    /// Dropping the stream early runs the reader's [`Drop`] cleanup. The stream
+    /// is `Send` but `!Unpin`; pin it before polling.
+    pub fn into_stream(self) -> impl Stream<Item = Result<ScanBatch>> + Send {
+        futures::stream::try_unfold(self, |mut reader| async move {
+            Ok(reader.next_batch().await?.map(|batch| (batch, reader)))
+        })
+    }
+
     /// Convert this async reader into a synchronous [`arrow::record_batch::RecordBatchReader`].
     ///
     /// The returned adapter calls [`tokio::runtime::Handle::block_on`] on each

fluss-rust/crates/fluss/src/client/write/accumulator.rs

@@ -880,7 +880,7 @@ impl RecordAccumulator {
         self.memory_limiter.close();
         // Complete batches still in deques (not yet drained).
         for mut entry in self.write_batches.iter_mut() {
-            for (_bucket_id, deque) in entry.value_mut().batches.iter_mut() {
+            for deque in entry.value_mut().batches.values_mut() {
                 let mut dq = deque.lock();
                 while let Some(batch) = dq.pop_front() {
                     batch.complete(Err(error.clone()));
@@ -931,7 +931,7 @@ impl RecordAccumulator {
 
     pub fn has_undrained(&self) -> bool {
         for entry in self.write_batches.iter() {
-            for (_, batch_deque) in entry.value().batches.iter() {
+            for batch_deque in entry.value().batches.values() {
                 if !batch_deque.lock().is_empty() {
                     return true;
                 }

fluss-rust/crates/fluss/tests/integration/batch_scanner.rs

@@ -22,6 +22,7 @@ mod batch_scanner_test {
     use arrow::array::{Int32Array, Int64Array, StringArray, record_batch};
     use fluss::metadata::{DataTypes, LogFormat, Schema, TableBucket, TableDescriptor, TablePath};
     use fluss::row::GenericRow;
+    use futures::TryStreamExt;
     use std::collections::HashMap;
 
     /// End-to-end check that the scanner yields the appended rows once and then
@@ -93,6 +94,70 @@ mod batch_scanner_test {
         );
     }
 
+    /// `into_stream` yields the scanner's single batch then ends, mirroring the
+    /// `next_batch` -> `Some` -> `None` sequence.
+    #[tokio::test]
+    async fn batch_scanner_into_stream_yields_single_batch() {
+        let cluster = get_shared_cluster();
+        let connection = cluster.get_fluss_connection().await;
+        let admin = connection.get_admin().expect("admin");
+
+        let table_path = TablePath::new("fluss", "test_batch_scanner_into_stream");
+        let descriptor = TableDescriptor::builder()
+            .schema(
+                Schema::builder()
+                    .column("c1", DataTypes::int())
+                    .column("c2", DataTypes::string())
+                    .build()
+                    .expect("schema"),
+            )
+            .distributed_by(Some(1), vec!["c1".to_string()])
+            .build()
+            .expect("descriptor");
+        create_table(&admin, &table_path, &descriptor).await;
+
+        let table = connection.get_table(&table_path).await.expect("table");
+        let writer = table
+            .new_append()
+            .expect("append")
+            .create_writer()
+            .expect("writer");
+        writer
+            .append_arrow_batch(
+                record_batch!(
+                    ("c1", Int32, [1, 2, 3, 4, 5]),
+                    ("c2", Utf8, ["a", "b", "c", "d", "e"])
+                )
+                .unwrap(),
+            )
+            .expect("append batch");
+        writer.flush().await.expect("flush");
+
+        let table_info = table.get_table_info();
+        let bucket = TableBucket::new(table_info.table_id, 0);
+
+        let scanner = table
+            .new_scan()
+            .limit(3)
+            .expect("limit")
+            .create_bucket_batch_scanner(bucket.clone())
+            .expect("create batch scanner");
+
+        let batches = scanner
+            .into_stream()
+            .try_collect::<Vec<_>>()
+            .await
+            .expect("drain batch scanner stream");
+
+        assert_eq!(batches.len(), 1, "limit scanner yields exactly one batch");
+        assert_eq!(batches[0].bucket(), &bucket);
+        assert!(
+            batches[0].num_records() > 0 && batches[0].num_records() <= 3,
+            "expected 1..=3 records, got {}",
+            batches[0].num_records()
+        );
+    }
+
     /// End-to-end projection skipping the middle `c2` string column.
     #[tokio::test]
     async fn batch_scanner_projects_non_contiguous_columns() {

fluss-rust/crates/fluss/tests/integration/record_batch_log_reader.rs

@@ -28,6 +28,7 @@ mod reader_test {
     use fluss::config::{Config, NoKeyAssigner};
     use fluss::metadata::{DataTypes, Schema, TableBucket, TableDescriptor, TablePath};
     use fluss::rpc::message::OffsetSpec;
+    use futures::TryStreamExt;
     use std::collections::HashMap;
     use std::time::Duration;
 
@@ -521,4 +522,75 @@ mod reader_test {
             .await
             .expect("Failed to drop table");
     }
+
+    #[tokio::test]
+    async fn into_stream_yields_same_batches_as_next_batch() {
+        let cluster = get_shared_cluster();
+        let connection = cluster.get_fluss_connection().await;
+        let admin = connection.get_admin().expect("Failed to get admin");
+
+        let table_path = TablePath::new("fluss", "test_reader_into_stream");
+        let table_descriptor = TableDescriptor::builder()
+            .schema(
+                Schema::builder()
+                    .column("id", DataTypes::int())
+                    .column("name", DataTypes::string())
+                    .build()
+                    .expect("Failed to build schema"),
+            )
+            .build()
+            .expect("Failed to build table");
+        create_table(&admin, &table_path, &table_descriptor).await;
+
+        let table = connection
+            .get_table(&table_path)
+            .await
+            .expect("Failed to get table");
+        let writer = table
+            .new_append()
+            .expect("Failed to create append")
+            .create_writer()
+            .expect("Failed to create writer");
+        writer
+            .append_arrow_batch(
+                record_batch!(
+                    ("id", Int32, [1, 2, 3, 4, 5]),
+                    ("name", Utf8, ["a", "b", "c", "d", "e"])
+                )
+                .unwrap(),
+            )
+            .expect("Failed to append batch");
+        writer.flush().await.expect("Failed to flush");
+
+        let scanner = table
+            .new_scan()
+            .create_record_batch_log_scanner()
+            .expect("Failed to create record batch scanner");
+        scanner
+            .subscribe(0, EARLIEST_OFFSET)
+            .await
+            .expect("Failed to subscribe bucket");
+
+        let reader = RecordBatchLogReader::new_until_latest(scanner, &admin)
+            .await
+            .expect("Failed to create latest-offset reader");
+        let batches = tokio::time::timeout(
+            Duration::from_secs(10),
+            reader.into_stream().try_collect::<Vec<_>>(),
+        )
+        .await
+        .expect("Timed out draining reader stream")
+        .expect("Failed to drain reader stream");
+
+        assert_eq!(
+            extract_ids_from_batches(&batches),
+            vec![1, 2, 3, 4, 5],
+            "into_stream should yield the same records next_batch would return"
+        );
+
+        admin
+            .drop_table(&table_path, false)
+            .await
+            .expect("Failed to drop table");
+    }
 }