diff --git a/crates/fluss/src/client/table/batch_scanner.rs b/crates/fluss/src/client/table/batch_scanner.rs index 090a7428..13ee78c7 100644 --- a/crates/fluss/src/client/table/batch_scanner.rs +++ b/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, + schema: SchemaRef, } /// Request inputs captured at creation, consumed by the first `next_batch`. @@ -71,8 +74,13 @@ impl LimitBatchScanner { projected_fields: Option>, bucket: TableBucket, limit: i32, - ) -> Self { - Self { + ) -> Result { + 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; + + fn next(&mut self) -> Option { + 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 diff --git a/crates/fluss/src/client/table/mod.rs b/crates/fluss/src/client/table/mod.rs index 41cbadab..e6d4674e 100644 --- a/crates/fluss/src/client/table/mod.rs +++ b/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)] diff --git a/crates/fluss/src/client/table/reader.rs b/crates/fluss/src/client/table/reader.rs index 518c68a2..5bb04c5d 100644 --- a/crates/fluss/src/client/table/reader.rs +++ b/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 diff --git a/crates/fluss/src/client/table/scanner.rs b/crates/fluss/src/client/table/scanner.rs index aadd1cb8..d50d69df 100644 --- a/crates/fluss/src/client/table/scanner.rs +++ b/crates/fluss/src/client/table/scanner.rs @@ -172,7 +172,7 @@ impl<'a> TableScan<'a> { self.conn.get_admin()?, latest, )); - Ok(LimitBatchScanner::new( + LimitBatchScanner::new( self.conn.get_connections(), self.metadata.clone(), self.table_info, @@ -180,7 +180,7 @@ impl<'a> TableScan<'a> { self.projected_fields, table_bucket, limit, - )) + ) } /// Projects the scan to only include specified columns by their indices. diff --git a/crates/fluss/src/client/table/stream.rs b/crates/fluss/src/client/table/stream.rs new file mode 100644 index 00000000..a0feb825 --- /dev/null +++ b/crates/fluss/src/client/table/stream.rs @@ -0,0 +1,323 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Stream adapters and bridges for bounded readers ([`RecordBatchLogReader`] and [`LimitBatchScanner`]). +//! +//! Provides: +//! - [`ScanBatchStream`]: An asynchronous [`Stream`] yielding [`ScanBatch`] items with schema metadata. +//! - [`RecordBatchStream`]: An asynchronous [`Stream`] yielding bare Arrow [`RecordBatch`] items, +//! suitable for DataFusion (`RecordBatchStreamAdapter`) and PyArrow interop. +//! - [`SyncRecordBatchStreamReader`]: A synchronous [`arrow::record_batch::RecordBatchReader`] +//! adapter driving an async stream using a Tokio runtime handle. + +use crate::error::Result; +use crate::record::ScanBatch; +use arrow::record_batch::RecordBatch; +use arrow_schema::SchemaRef; +use futures::stream::BoxStream; +use futures::{Stream, StreamExt}; +use std::fmt; +use std::pin::Pin; +use std::task::{Context, Poll}; + +/// An asynchronous [`Stream`] yielding [`ScanBatch`] items from a bounded scan or reader. +/// +/// Created via [`RecordBatchLogReader::into_stream`] or [`LimitBatchScanner::into_stream`]. +/// +/// Provides access to the Arrow schema via [`schema`](Self::schema) and can be converted +/// to a bare-`RecordBatch` stream via [`into_record_batch_stream`](Self::into_record_batch_stream) +/// or to a synchronous [`arrow::record_batch::RecordBatchReader`] via +/// [`to_record_batch_reader`](Self::to_record_batch_reader). +pub struct ScanBatchStream { + inner: BoxStream<'static, Result>, + schema: SchemaRef, +} + +impl ScanBatchStream { + /// Creates a new `ScanBatchStream` from a boxed stream and its Arrow schema. + pub fn new(inner: BoxStream<'static, Result>, schema: SchemaRef) -> Self { + Self { inner, schema } + } + + /// Returns the Arrow schema of batches produced by this stream. + pub fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + /// Converts this [`ScanBatchStream`] into a [`RecordBatchStream`], stripping + /// Fluss metadata and yielding bare Arrow [`RecordBatch`]es suitable for + /// DataFusion or PyArrow. + pub fn into_record_batch_stream(self) -> RecordBatchStream { + RecordBatchStream::new(self) + } + + /// Converts this stream into a synchronous [`arrow::record_batch::RecordBatchReader`]. + /// + /// The returned adapter blocks on the provided Tokio runtime handle. + pub fn to_record_batch_reader( + self, + handle: tokio::runtime::Handle, + ) -> SyncRecordBatchStreamReader { + self.into_record_batch_stream().to_record_batch_reader(handle) + } +} + +impl Stream for ScanBatchStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.inner.as_mut().poll_next(cx) + } +} + +impl fmt::Debug for ScanBatchStream { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ScanBatchStream") + .field("schema", &self.schema) + .finish_non_exhaustive() + } +} + +/// An asynchronous [`Stream`] yielding bare Arrow [`RecordBatch`] items. +/// +/// Created via [`ScanBatchStream::into_record_batch_stream`], +/// [`RecordBatchLogReader::into_record_batch_stream`], or +/// [`LimitBatchScanner::into_record_batch_stream`]. +/// +/// Strips Fluss-specific metadata (bucket ID, log offsets) from each [`ScanBatch`], +/// yielding standard Arrow `Result` items. +/// This type exposes [`schema`](Self::schema), allowing it to plug directly into +/// DataFusion (e.g. `RecordBatchStreamAdapter`) and PyArrow. +pub struct RecordBatchStream { + inner: BoxStream<'static, std::result::Result>, + schema: SchemaRef, +} + +impl RecordBatchStream { + /// Creates a new `RecordBatchStream` by wrapping a [`ScanBatchStream`]. + pub fn new(stream: ScanBatchStream) -> Self { + let schema = stream.schema(); + let inner = Box::pin(stream.map(|res| { + res.map(|scan_batch| scan_batch.into_batch()) + .map_err(|e| arrow::error::ArrowError::ExternalError(Box::new(e))) + })); + Self { inner, schema } + } + + /// Creates a new `RecordBatchStream` from an existing boxed stream and its Arrow schema. + pub fn from_stream( + inner: BoxStream<'static, std::result::Result>, + schema: SchemaRef, + ) -> Self { + Self { inner, schema } + } + + /// Returns the Arrow schema of batches produced by this stream. + pub fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + /// Converts this stream into a synchronous [`arrow::record_batch::RecordBatchReader`]. + /// + /// The returned adapter blocks on the provided Tokio runtime handle. + pub fn to_record_batch_reader( + self, + handle: tokio::runtime::Handle, + ) -> SyncRecordBatchStreamReader { + SyncRecordBatchStreamReader { + stream: self, + handle, + } + } +} + +impl Stream for RecordBatchStream { + type Item = std::result::Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.inner.as_mut().poll_next(cx) + } +} + +impl From for RecordBatchStream { + fn from(stream: ScanBatchStream) -> Self { + stream.into_record_batch_stream() + } +} + +impl fmt::Debug for RecordBatchStream { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("RecordBatchStream") + .field("schema", &self.schema) + .finish_non_exhaustive() + } +} + +/// Synchronous adapter that wraps a [`RecordBatchStream`] to implement +/// [`arrow::record_batch::RecordBatchReader`]. +/// +/// Created via [`RecordBatchStream::to_record_batch_reader`] or +/// [`ScanBatchStream::to_record_batch_reader`]. +/// Blocks the current thread on each `next()` call using the provided +/// Tokio runtime handle. +pub struct SyncRecordBatchStreamReader { + stream: RecordBatchStream, + handle: tokio::runtime::Handle, +} + +impl SyncRecordBatchStreamReader { + /// Creates a new synchronous reader wrapping the given stream. + pub fn new(stream: RecordBatchStream, handle: tokio::runtime::Handle) -> Self { + Self { stream, handle } + } + + /// Returns the Arrow schema of batches produced by this reader. + pub fn schema(&self) -> SchemaRef { + self.stream.schema() + } +} + +impl Iterator for SyncRecordBatchStreamReader { + type Item = std::result::Result; + + fn next(&mut self) -> Option { + self.handle.block_on(self.stream.next()) + } +} + +impl arrow::record_batch::RecordBatchReader for SyncRecordBatchStreamReader { + fn schema(&self) -> SchemaRef { + self.stream.schema() + } +} + +impl fmt::Debug for SyncRecordBatchStreamReader { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("SyncRecordBatchStreamReader") + .field("schema", &self.stream.schema()) + .finish_non_exhaustive() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::error::Error; + use crate::metadata::TableBucket; + use arrow::array::{ArrayRef, Int32Array, StringArray}; + use arrow_schema::{DataType as ArrowDataType, Field, Schema as ArrowSchema}; + use std::sync::Arc; + + fn test_schema() -> SchemaRef { + Arc::new(ArrowSchema::new(vec![ + Field::new("id", ArrowDataType::Int32, false), + Field::new("name", ArrowDataType::Utf8, false), + ])) + } + + fn test_scan_batch(schema: SchemaRef, ids: Vec, names: Vec<&str>) -> ScanBatch { + let columns: Vec = vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(StringArray::from(names)), + ]; + let record_batch = RecordBatch::try_new(schema, columns).unwrap(); + ScanBatch::new(TableBucket::new(1, 0), record_batch, 0) + } + + #[tokio::test] + async fn test_scan_batch_stream_consumption() { + let schema = test_schema(); + let batch1 = test_scan_batch(schema.clone(), vec![1, 2], vec!["a", "b"]); + let batch2 = test_scan_batch(schema.clone(), vec![3], vec!["c"]); + + let items = vec![Ok(batch1), Ok(batch2)]; + let boxed = Box::pin(futures::stream::iter(items)); + let mut stream = ScanBatchStream::new(boxed, schema.clone()); + + assert_eq!(stream.schema(), schema); + + let first = stream.next().await.unwrap().unwrap(); + assert_eq!(first.num_records(), 2); + + let second = stream.next().await.unwrap().unwrap(); + assert_eq!(second.num_records(), 1); + + assert!(stream.next().await.is_none()); + } + + #[tokio::test] + async fn test_into_record_batch_stream() { + let schema = test_schema(); + let batch1 = test_scan_batch(schema.clone(), vec![1, 2], vec!["a", "b"]); + + let items = vec![Ok(batch1)]; + let boxed = Box::pin(futures::stream::iter(items)); + let scan_stream = ScanBatchStream::new(boxed, schema.clone()); + let mut rb_stream = scan_stream.into_record_batch_stream(); + + assert_eq!(rb_stream.schema(), schema); + + let first = rb_stream.next().await.unwrap().unwrap(); + assert_eq!(first.num_rows(), 2); + assert_eq!(first.num_columns(), 2); + + assert!(rb_stream.next().await.is_none()); + } + + #[tokio::test] + async fn test_record_batch_stream_error_propagation() { + let schema = test_schema(); + let items: Vec> = vec![ + Err(Error::UnexpectedError { + message: "scan failure".to_string(), + source: None, + }), + ]; + let boxed = Box::pin(futures::stream::iter(items)); + let scan_stream = ScanBatchStream::new(boxed, schema.clone()); + let mut rb_stream = scan_stream.into_record_batch_stream(); + + let err = rb_stream.next().await.unwrap().unwrap_err(); + match err { + arrow::error::ArrowError::ExternalError(e) => { + assert!(e.to_string().contains("scan failure")); + } + _ => panic!("Expected ExternalError, got {:?}", err), + } + + assert!(rb_stream.next().await.is_none()); + } + + #[test] + fn test_sync_record_batch_stream_reader() { + let rt = tokio::runtime::Runtime::new().unwrap(); + let handle = rt.handle().clone(); + + let schema = test_schema(); + let batch1 = test_scan_batch(schema.clone(), vec![10, 20], vec!["x", "y"]); + let boxed = Box::pin(futures::stream::iter(vec![Ok(batch1)])); + let scan_stream = ScanBatchStream::new(boxed, schema.clone()); + let mut reader = scan_stream.to_record_batch_reader(handle); + + assert_eq!(arrow::record_batch::RecordBatchReader::schema(&reader), schema); + + let batch = reader.next().unwrap().unwrap(); + assert_eq!(batch.num_rows(), 2); + + assert!(reader.next().is_none()); + } +} diff --git a/crates/fluss/tests/integration/batch_scanner.rs b/crates/fluss/tests/integration/batch_scanner.rs index 443d0518..1ad6e7d7 100644 --- a/crates/fluss/tests/integration/batch_scanner.rs +++ b/crates/fluss/tests/integration/batch_scanner.rs @@ -420,4 +420,143 @@ mod batch_scanner_test { "create_record_batch_log_scanner must reject a configured limit" ); } + + /// Verify `into_stream()` on `LimitBatchScanner`. + #[tokio::test] + async fn batch_scanner_into_stream() { + use futures::StreamExt; + + 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"); + + let batch = record_batch!( + ("c1", Int32, [10, 20, 30]), + ("c2", Utf8, ["x", "y", "z"]) + ) + .unwrap(); + writer.append_arrow_batch(batch).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(2) + .expect("limit") + .create_bucket_batch_scanner(bucket.clone()) + .expect("create batch scanner"); + + assert_eq!(scanner.schema().fields().len(), 2); + let mut stream = scanner.into_stream(); + assert_eq!(stream.schema().fields().len(), 2); + + let first = stream + .next() + .await + .expect("poll") + .expect("first batch should be Ok"); + + assert_eq!(first.bucket(), &bucket); + assert!(first.num_records() > 0 && first.num_records() <= 2); + + assert!(stream.next().await.is_none(), "stream must end after bounded scan"); + } + + /// Verify `into_record_batch_stream()` and `to_record_batch_reader()` on `LimitBatchScanner`. + #[tokio::test] + async fn batch_scanner_bare_record_batch_bridges() { + use arrow::record_batch::RecordBatchReader; + use futures::StreamExt; + + 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_bridges"); + 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"); + + let batch = record_batch!( + ("c1", Int32, [1, 2, 3]), + ("c2", Utf8, ["a", "b", "c"]) + ) + .unwrap(); + writer.append_arrow_batch(batch).expect("append batch"); + writer.flush().await.expect("flush"); + + let table_info = table.get_table_info(); + let bucket = TableBucket::new(table_info.table_id, 0); + + // Test into_record_batch_stream() + let scanner = table + .new_scan() + .limit(3) + .expect("limit") + .create_bucket_batch_scanner(bucket.clone()) + .expect("create batch scanner"); + + let mut rb_stream = scanner.into_record_batch_stream(); + assert_eq!(rb_stream.schema().fields().len(), 2); + let rb = rb_stream + .next() + .await + .expect("stream poll") + .expect("stream batch Ok"); + assert!(rb.num_rows() > 0 && rb.num_rows() <= 3); + assert!(rb_stream.next().await.is_none()); + + // Test to_record_batch_reader() + let scanner2 = table + .new_scan() + .limit(3) + .expect("limit") + .create_bucket_batch_scanner(bucket.clone()) + .expect("create batch scanner"); + + let mut sync_reader = scanner2.to_record_batch_reader(tokio::runtime::Handle::current()); + assert_eq!(sync_reader.schema().fields().len(), 2); + let sync_rb = sync_reader.next().expect("sync reader batch Some").expect("sync reader batch Ok"); + assert!(sync_rb.num_rows() > 0 && sync_rb.num_rows() <= 3); + assert!(sync_reader.next().is_none()); + } } diff --git a/crates/fluss/tests/integration/record_batch_log_reader.rs b/crates/fluss/tests/integration/record_batch_log_reader.rs index 6c8d5392..0b5b305f 100644 --- a/crates/fluss/tests/integration/record_batch_log_reader.rs +++ b/crates/fluss/tests/integration/record_batch_log_reader.rs @@ -521,4 +521,124 @@ mod reader_test { .await .expect("Failed to drop table"); } + + #[tokio::test] + async fn record_batch_log_reader_stream_and_bridges() { + use arrow::record_batch::RecordBatchReader; + use futures::StreamExt; + + 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_stream_bridges"); + 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]), + ("name", Utf8, ["a", "b", "c"]) + ) + .unwrap(), + ) + .expect("Failed to append batch"); + writer.flush().await.expect("Failed to flush"); + + // 1. Test into_stream() + 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"); + + let reader = RecordBatchLogReader::new_until_latest(scanner, &admin) + .await + .expect("Failed to create reader"); + + assert_eq!(reader.schema().fields().len(), 2); + let mut stream = reader.into_stream(); + assert_eq!(stream.schema().fields().len(), 2); + + let mut batches = Vec::new(); + while let Some(batch_res) = stream.next().await { + batches.push(batch_res.expect("stream batch Ok")); + } + assert_eq!(extract_ids_from_batches(&batches), vec![1, 2, 3]); + + // 2. Test into_record_batch_stream() + let scanner2 = table + .new_scan() + .create_record_batch_log_scanner() + .expect("Failed to create scanner"); + scanner2 + .subscribe(0, EARLIEST_OFFSET) + .await + .expect("Failed to subscribe"); + + let reader2 = RecordBatchLogReader::new_until_latest(scanner2, &admin) + .await + .expect("Failed to create reader"); + + let mut rb_stream = reader2.into_record_batch_stream(); + assert_eq!(rb_stream.schema().fields().len(), 2); + + let mut rb_count = 0; + while let Some(rb_res) = rb_stream.next().await { + let rb = rb_res.expect("rb stream batch Ok"); + rb_count += rb.num_rows(); + } + assert_eq!(rb_count, 3); + + // 3. Test to_record_batch_reader() + let scanner3 = table + .new_scan() + .create_record_batch_log_scanner() + .expect("Failed to create scanner"); + scanner3 + .subscribe(0, EARLIEST_OFFSET) + .await + .expect("Failed to subscribe"); + + let reader3 = RecordBatchLogReader::new_until_latest(scanner3, &admin) + .await + .expect("Failed to create reader"); + + let mut sync_reader = reader3.to_record_batch_reader(tokio::runtime::Handle::current()); + assert_eq!(sync_reader.schema().fields().len(), 2); + + let mut sync_count = 0; + for rb_res in &mut sync_reader { + let rb = rb_res.expect("sync reader batch Ok"); + sync_count += rb.num_rows(); + } + assert_eq!(sync_count, 3); + + admin + .drop_table(&table_path, false) + .await + .expect("Failed to drop table"); + } }