Partitioned scan for RegionEngine

[evenyag]

What type of enhancement is this?

API improvement

What does the enhancement do?

The RegionEngine trait provides a handle_query() method to scan a region and returns a stream of RecordBatch.

greptimedb/src/store-api/src/region_engine.rs

Lines 136 to 140 in d997463

	async fn handle_query(
	&self,
	region_id: RegionId,
	request: ScanRequest,
	) -> Result<SendableRecordBatchStream, BoxedError>;

This method is easy to use but has some limitations:

• The output concurrency is always 1
• The engine can't return more information about the query to callers

To maximize parallelism in #2806, the engine should provide a way to return multiple streams to scan different partitions of a region concurrently.

Implementation challenges

This issue proposes to add a new method to the region engine which supports partitioned scan. The method returns a trait object that can create a stream according to a partition index.

pub struct ScannerProperties {
    // Properties of the scanner
    // e.g. number of partitions, range of partitions
}

pub trait RegionScanner {
    fn properties(&self) -> &ScannerProperties;

    fn scan_partition(&self, partition: usize) -> Result<SendableRecordBatchStream, BoxedError>;
}

pub type RegionScannerRef = Arc<dyn RegionScanner>;

pub trait RegionEngine {
    async fn handle_partitioned_query(
        &self,
        region_id: RegionId,
        request: ScanRequest,
    ) -> Result<RegionScannerRef, BoxedError>;
}

We could then use the scanner to implement a PhysicalPlan and let the query engine process multiple partitions. We might need to refactor the StreamScanAdapter as it assumes there is only one partition.

greptimedb/src/table/src/table/scan.rs

Lines 103 to 119 in a6a702d

	fn execute(
	&self,
	partition: usize,
	context: Arc<TaskContext>,
	) -> QueryResult<SendableRecordBatchStream> {
	let tracing_context = TracingContext::from_json(context.session_id().as_str());
	let span = tracing_context.attach(common_telemetry::tracing::info_span!("stream_adapter"));
	let mut stream = self.stream.lock().unwrap();
	let stream = stream.take().context(query_error::ExecuteRepeatedlySnafu)?;
	let mem_usage_metrics = MemoryUsageMetrics::new(&self.metric, partition);
	Ok(Box::pin(StreamWithMetricWrapper {
	stream,
	metric: mem_usage_metrics,
	span,
	}))
	}

[evenyag]

closed by #3948