feat: add extension range api for flat format (#7730)

feat: add flat extension api

Signed-off-by: evenyag <realevenyag@gmail.com>

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<dyn ExtensionRange>;
@@ -50,6 +53,20 @@ pub trait ExtensionRangeReader: Send {
 
 pub type BoxedExtensionRangeReader = Box<dyn ExtensionRangeReader>;
 
+/// 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<Self>,
+        context: Arc<StreamContext>,
+        metrics: PartitionMetrics,
+        index: RowGroupIndex,
+    ) -> Result<BoxedRecordBatchStream, BoxedError>;
+}
+
+pub type BoxedExtensionFlatRangeReader = Box<dyn ExtensionFlatRangeReader>;
+
 /// The provider to feed the extension ranges into the mito scanner.
 #[async_trait]
 pub trait ExtensionRangeProvider: Send + Sync {

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<StreamContext>,
+    index: RowGroupIndex,
+    partition_metrics: PartitionMetrics,
+) -> Result<BoxedRecordBatchStream> {
+    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<StreamContext>,
     index: RowGroupIndex,
     metrics: &PartitionMetrics,
 ) -> Result<BoxedRecordBatchStream> {
-    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.