diff --git a/src/mito2/src/extension.rs b/src/mito2/src/extension.rs index 58a16f7c4d3..944344b1cfc 100644 --- a/src/mito2/src/extension.rs +++ b/src/mito2/src/extension.rs @@ -8,10 +8,10 @@ use common_time::range::TimestampRange; use store_api::storage::{ScanRequest, SequenceNumber}; use crate::error::Result; -use crate::read::BoxedBatchStream; use crate::read::range::RowGroupIndex; use crate::read::scan_region::StreamContext; use crate::read::scan_util::PartitionMetrics; +use crate::read::{BoxedBatchStream, BoxedRecordBatchStream}; use crate::region::MitoRegionRef; pub type InclusiveTimeRange = (Timestamp, Timestamp); @@ -32,6 +32,9 @@ pub trait ExtensionRange: Debug + Send + Sync { /// Create the reader for reading this range. fn reader(&self, context: &StreamContext) -> BoxedExtensionRangeReader; + + /// Create the flat reader for reading this range in flat format. + fn flat_reader(&self, context: &StreamContext) -> BoxedExtensionFlatRangeReader; } pub type BoxedExtensionRange = Box; @@ -50,6 +53,20 @@ pub trait ExtensionRangeReader: Send { pub type BoxedExtensionRangeReader = Box; +/// The reader to read an extension range in flat format (producing [`RecordBatch`]). +#[async_trait] +pub trait ExtensionFlatRangeReader: Send { + /// Read the extension range by creating a stream that produces [`RecordBatch`]. + async fn read( + self: Box, + context: Arc, + metrics: PartitionMetrics, + index: RowGroupIndex, + ) -> Result; +} + +pub type BoxedExtensionFlatRangeReader = Box; + /// The provider to feed the extension ranges into the mito scanner. #[async_trait] pub trait ExtensionRangeProvider: Send + Sync { diff --git a/src/mito2/src/read/scan_util.rs b/src/mito2/src/read/scan_util.rs index a1bd0552424..0ce957a6160 100644 --- a/src/mito2/src/read/scan_util.rs +++ b/src/mito2/src/read/scan_util.rs @@ -1558,19 +1558,45 @@ pub(crate) async fn maybe_scan_other_ranges( } } +/// Build the stream of scanning the extension range in flat format denoted by the [`RowGroupIndex`]. +#[cfg(feature = "enterprise")] +pub(crate) async fn scan_flat_extension_range( + context: Arc, + index: RowGroupIndex, + partition_metrics: PartitionMetrics, +) -> Result { + use snafu::ResultExt; + + let range = context.input.extension_range(index.index); + let reader = range.flat_reader(context.as_ref()); + let stream = reader + .read(context, partition_metrics, index) + .await + .context(crate::error::ScanExternalRangeSnafu)?; + Ok(stream) +} + pub(crate) async fn maybe_scan_flat_other_ranges( context: &Arc, index: RowGroupIndex, metrics: &PartitionMetrics, ) -> Result { - let _ = context; - let _ = index; - let _ = metrics; + #[cfg(feature = "enterprise")] + { + scan_flat_extension_range(context.clone(), index, metrics.clone()).await + } - crate::error::UnexpectedSnafu { - reason: "no other ranges scannable in flat format", + #[cfg(not(feature = "enterprise"))] + { + let _ = context; + let _ = index; + let _ = metrics; + + crate::error::UnexpectedSnafu { + reason: "no other ranges scannable in flat format", + } + .fail() } - .fail() } /// A stream wrapper that splits record batches from an inner stream.