rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-05-18 05:50:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: make table scan return physical plan (#326)

Browse Source 
* refactor: return PhysicalPlan in Table trait's scan method, to support partitioned execution in Frontend's distribute read

* refactor: pub use necessary DataFusion types

* refactor: replace old "PhysicalPlan" and its adapters

Co-authored-by: luofucong <luofucong@greptime.com>
Co-authored-by: Yingwen <realevenyag@gmail.com>
This commit is contained in:
LFC
2022-10-25 11:34:53 +08:00
committed by GitHub
parent 64dac51e83
commit 2ca667cbdf
39 changed files with 920 additions and 600 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
src/table/src/error.rs
View File

@@ -28,9 +28,6 @@ pub enum InnerError {
#[snafu(display("Missing column when insert, column: {}", name))]
MissingColumn { name: String, backtrace: Backtrace },
#[snafu(display("Not expected to run ExecutionPlan more than once"))]
ExecuteRepeatedly { backtrace: Backtrace },
#[snafu(display("Poll stream failed, source: {}", source))]
PollStream {
source: ArrowError,
@@ -58,7 +55,6 @@ impl ErrorExt for InnerError {
| InnerError::SchemaConversion { .. }
| InnerError::TableProjection { .. } => StatusCode::EngineExecuteQuery,
InnerError::MissingColumn { .. } => StatusCode::InvalidArguments,
InnerError::ExecuteRepeatedly { .. } => StatusCode::Unexpected,
}
}
@@ -97,10 +93,6 @@ mod tests {
Err(DataFusionError::NotImplemented("table test".to_string())).context(DatafusionSnafu)?
}
fn throw_repeatedly() -> Result<()> {
ExecuteRepeatedlySnafu {}.fail()?
}
fn throw_missing_column_inner() -> std::result::Result<(), InnerError> {
MissingColumnSnafu { name: "test" }.fail()
}
@@ -119,10 +111,6 @@ mod tests {
assert!(err.backtrace_opt().is_some());
assert_eq!(StatusCode::EngineExecuteQuery, err.status_code());
let err = throw_repeatedly().err().unwrap();
assert!(err.backtrace_opt().is_some());
assert_eq!(StatusCode::Unexpected, err.status_code());
let err = throw_missing_column().err().unwrap();
assert!(err.backtrace_opt().is_some());
assert_eq!(StatusCode::InvalidArguments, err.status_code());

8
src/table/src/table.rs
View File

@@ -1,11 +1,13 @@
pub mod adapter;
pub mod numbers;
pub mod scan;
use std::any::Any;
use std::sync::Arc;
use async_trait::async_trait;
use common_query::logical_plan::Expr;
use common_recordbatch::SendableRecordBatchStream;
use common_query::physical_plan::PhysicalPlanRef;
use datatypes::schema::SchemaRef;
use crate::error::Result;
@@ -13,7 +15,7 @@ use crate::metadata::{FilterPushDownType, TableInfoRef, TableType};
use crate::requests::{AlterTableRequest, InsertRequest};
/// Table abstraction.
#[async_trait::async_trait]
#[async_trait]
pub trait Table: Send + Sync {
/// Returns the table as [`Any`](std::any::Any) so that it can be
/// downcast to a specific implementation.
@@ -45,7 +47,7 @@ pub trait Table: Send + Sync {
// If set, it contains the amount of rows needed by the `LogicalPlan`,
// The datasource should return *at least* this number of rows if available.
limit: Option<usize>,
) -> Result<SendableRecordBatchStream>;
) -> Result<PhysicalPlanRef>;
/// Tests whether the table provider can make use of a filter expression
/// to optimise data retrieval.

210
src/table/src/table/adapter.rs
View File

@@ -1,14 +1,9 @@
use core::fmt::Formatter;
use core::pin::Pin;
use core::task::{Context, Poll};
use std::any::Any;
use std::fmt::Debug;
use std::mem;
use std::sync::{Arc, Mutex};
use std::sync::Arc;
use common_query::logical_plan::Expr;
use common_recordbatch::error::Result as RecordBatchResult;
use common_recordbatch::{RecordBatch, RecordBatchStream, SendableRecordBatchStream};
use common_query::physical_plan::{DfPhysicalPlanAdapter, PhysicalPlanAdapter, PhysicalPlanRef};
use common_query::DfPhysicalPlan;
use common_telemetry::debug;
use datafusion::arrow::datatypes::SchemaRef as DfSchemaRef;
/// Datafusion table adpaters
@@ -17,92 +12,14 @@ use datafusion::datasource::{
TableType as DfTableType,
};
use datafusion::error::Result as DfResult;
use datafusion::execution::runtime_env::RuntimeEnv;
use datafusion::logical_plan::Expr as DfExpr;
use datafusion::physical_plan::{
expressions::PhysicalSortExpr, ExecutionPlan, Partitioning,
RecordBatchStream as DfRecordBatchStream,
SendableRecordBatchStream as DfSendableRecordBatchStream, Statistics,
};
use datafusion_common::record_batch::RecordBatch as DfRecordBatch;
use datatypes::arrow::error::{ArrowError, Result as ArrowResult};
use datatypes::schema::SchemaRef;
use datatypes::schema::{Schema, SchemaRef as TableSchemaRef};
use futures::Stream;
use datatypes::schema::{SchemaRef as TableSchemaRef, SchemaRef};
use snafu::prelude::*;
use crate::error::{self, Result};
use crate::metadata::TableInfoRef;
use crate::table::{FilterPushDownType, Table, TableRef, TableType};
/// Greptime SendableRecordBatchStream -> datafusion ExecutionPlan.
struct ExecutionPlanAdapter {
stream: Mutex<Option<SendableRecordBatchStream>>,
schema: SchemaRef,
}
impl Debug for ExecutionPlanAdapter {
fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
f.debug_struct("ExecutionPlanAdapter")
.field("schema", &self.schema)
.finish()
}
}
#[async_trait::async_trait]
impl ExecutionPlan for ExecutionPlanAdapter {
fn as_any(&self) -> &dyn Any {
self
}
fn schema(&self) -> DfSchemaRef {
self.schema.arrow_schema().clone()
}
fn output_partitioning(&self) -> Partitioning {
// FIXME(dennis)
Partitioning::UnknownPartitioning(1)
}
fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
// FIXME(dennis)
None
}
fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
// TODO(dennis)
vec![]
}
fn with_new_children(
&self,
_children: Vec<Arc<dyn ExecutionPlan>>,
) -> DfResult<Arc<dyn ExecutionPlan>> {
// TODO(dennis)
todo!();
}
async fn execute(
&self,
_partition: usize,
_runtime: Arc<RuntimeEnv>,
) -> DfResult<DfSendableRecordBatchStream> {
let mut stream = self.stream.lock().unwrap();
if stream.is_some() {
let stream = mem::replace(&mut *stream, None);
Ok(Box::pin(DfRecordBatchStreamAdapter::new(stream.unwrap())))
} else {
error::ExecuteRepeatedlySnafu.fail()?
}
}
fn statistics(&self) -> Statistics {
//TODO(dennis)
Statistics::default()
}
}
/// Greptime Table -> datafusion TableProvider
pub struct DfTableProviderAdapter {
table: TableRef,
@@ -141,14 +58,10 @@ impl TableProvider for DfTableProviderAdapter {
projection: &Option<Vec<usize>>,
filters: &[DfExpr],
limit: Option<usize>,
) -> DfResult<Arc<dyn ExecutionPlan>> {
) -> DfResult<Arc<dyn DfPhysicalPlan>> {
let filters: Vec<Expr> = filters.iter().map(Clone::clone).map(Into::into).collect();
let stream = self.table.scan(projection, &filters, limit).await?;
Ok(Arc::new(ExecutionPlanAdapter {
schema: stream.schema(),
stream: Mutex::new(Some(stream)),
}))
let inner = self.table.scan(projection, &filters, limit).await?;
Ok(Arc::new(DfPhysicalPlanAdapter(inner)))
}
fn supports_filter_pushdown(&self, filter: &DfExpr) -> DfResult<DfTableProviderFilterPushDown> {
@@ -167,15 +80,18 @@ impl TableProvider for DfTableProviderAdapter {
pub struct TableAdapter {
schema: TableSchemaRef,
table_provider: Arc<dyn TableProvider>,
runtime: Arc<RuntimeEnv>,
}
impl TableAdapter {
pub fn new(table_provider: Arc<dyn TableProvider>, runtime: Arc<RuntimeEnv>) -> Result<Self> {
pub fn new(table_provider: Arc<dyn TableProvider>) -> Result<Self> {
Ok(Self {
schema: Arc::new(table_provider.schema().try_into().unwrap()),
schema: Arc::new(
table_provider
.schema()
.try_into()
.context(error::SchemaConversionSnafu)?,
),
table_provider,
runtime,
})
}
}
@@ -207,7 +123,7 @@ impl Table for TableAdapter {
projection: &Option<Vec<usize>>,
filters: &[Expr],
limit: Option<usize>,
) -> Result<SendableRecordBatchStream> {
) -> Result<PhysicalPlanRef> {
let filters: Vec<DfExpr> = filters.iter().map(|e| e.df_expr().clone()).collect();
debug!("TableScan filter size: {}", filters.len());
let execution_plan = self
@@ -215,14 +131,13 @@ impl Table for TableAdapter {
.scan(projection, &filters, limit)
.await
.context(error::DatafusionSnafu)?;
// FIXME(dennis) Partitioning
let df_stream = execution_plan
.execute(0, self.runtime.clone())
.await
.context(error::DatafusionSnafu)?;
Ok(Box::pin(RecordBatchStreamAdapter::try_new(df_stream)?))
let schema: SchemaRef = Arc::new(
execution_plan
.schema()
.try_into()
.context(error::SchemaConversionSnafu)?,
);
Ok(Arc::new(PhysicalPlanAdapter::new(schema, execution_plan)))
}
fn supports_filter_pushdown(&self, filter: &Expr) -> Result<FilterPushDownType> {
@@ -238,83 +153,6 @@ impl Table for TableAdapter {
}
}
/// Greptime SendableRecordBatchStream -> datafusion RecordBatchStream
pub struct DfRecordBatchStreamAdapter {
stream: SendableRecordBatchStream,
}
impl DfRecordBatchStreamAdapter {
pub fn new(stream: SendableRecordBatchStream) -> Self {
Self { stream }
}
}
impl DfRecordBatchStream for DfRecordBatchStreamAdapter {
fn schema(&self) -> DfSchemaRef {
self.stream.schema().arrow_schema().clone()
}
}
impl Stream for DfRecordBatchStreamAdapter {
type Item = ArrowResult<DfRecordBatch>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
match Pin::new(&mut self.stream).poll_next(cx) {
Poll::Pending => Poll::Pending,
Poll::Ready(Some(recordbatch)) => match recordbatch {
Ok(recordbatch) => Poll::Ready(Some(Ok(recordbatch.df_recordbatch))),
Err(e) => Poll::Ready(Some(Err(ArrowError::External("".to_owned(), Box::new(e))))),
},
Poll::Ready(None) => Poll::Ready(None),
}
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
self.stream.size_hint()
}
}
/// Datafusion SendableRecordBatchStream to greptime RecordBatchStream
pub struct RecordBatchStreamAdapter {
schema: SchemaRef,
stream: DfSendableRecordBatchStream,
}
impl RecordBatchStreamAdapter {
pub fn try_new(stream: DfSendableRecordBatchStream) -> Result<Self> {
let schema =
Arc::new(Schema::try_from(stream.schema()).context(error::SchemaConversionSnafu)?);
Ok(Self { schema, stream })
}
}
impl RecordBatchStream for RecordBatchStreamAdapter {
fn schema(&self) -> SchemaRef {
self.schema.clone()
}
}
impl Stream for RecordBatchStreamAdapter {
type Item = RecordBatchResult<RecordBatch>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
match Pin::new(&mut self.stream).poll_next(cx) {
Poll::Pending => Poll::Pending,
Poll::Ready(Some(df_recordbatch)) => Poll::Ready(Some(Ok(RecordBatch {
schema: self.schema(),
df_recordbatch: df_recordbatch.context(error::PollStreamSnafu)?,
}))),
Poll::Ready(None) => Poll::Ready(None),
}
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
self.stream.size_hint()
}
}
#[cfg(test)]
mod tests {
use datafusion::arrow;
@@ -328,14 +166,14 @@ mod tests {
#[should_panic]
fn test_table_adaptor_info() {
let df_table = Arc::new(EmptyTable::new(Arc::new(arrow::datatypes::Schema::empty())));
let table_adapter = TableAdapter::new(df_table, Arc::new(RuntimeEnv::default())).unwrap();
let table_adapter = TableAdapter::new(df_table).unwrap();
let _ = table_adapter.table_info();
}
#[test]
fn test_table_adaptor_type() {
let df_table = Arc::new(EmptyTable::new(Arc::new(arrow::datatypes::Schema::empty())));
let table_adapter = TableAdapter::new(df_table, Arc::new(RuntimeEnv::default())).unwrap();
let table_adapter = TableAdapter::new(df_table).unwrap();
assert_eq!(Base, table_adapter.table_type());
}
}

11
src/table/src/table/numbers.rs
View File

@@ -2,8 +2,9 @@ use std::any::Any;
use std::pin::Pin;
use std::sync::Arc;
use common_query::physical_plan::PhysicalPlanRef;
use common_recordbatch::error::Result as RecordBatchResult;
use common_recordbatch::{RecordBatch, RecordBatchStream, SendableRecordBatchStream};
use common_recordbatch::{RecordBatch, RecordBatchStream};
use datafusion_common::record_batch::RecordBatch as DfRecordBatch;
use datatypes::arrow::array::UInt32Array;
use datatypes::data_type::ConcreteDataType;
@@ -13,6 +14,7 @@ use futures::Stream;
use crate::error::Result;
use crate::metadata::TableInfoRef;
use crate::table::scan::SimpleTableScan;
use crate::table::{Expr, Table};
/// numbers table for test
@@ -53,12 +55,13 @@ impl Table for NumbersTable {
_projection: &Option<Vec<usize>>,
_filters: &[Expr],
limit: Option<usize>,
) -> Result<SendableRecordBatchStream> {
Ok(Box::pin(NumbersStream {
) -> Result<PhysicalPlanRef> {
let stream = Box::pin(NumbersStream {
limit: limit.unwrap_or(100) as u32,
schema: self.schema.clone(),
already_run: false,
}))
});
Ok(Arc::new(SimpleTableScan::new(stream)))
}
}

123
src/table/src/table/scan.rs Normal file
View File

@@ -0,0 +1,123 @@
use std::any::Any;
use std::fmt::{Debug, Formatter};
use std::sync::{Arc, Mutex};
use async_trait::async_trait;
use common_query::error as query_error;
use common_query::error::Result as QueryResult;
use common_query::physical_plan::Partitioning;
use common_query::physical_plan::RuntimeEnv;
use common_query::physical_plan::{PhysicalPlan, PhysicalPlanRef};
use common_recordbatch::SendableRecordBatchStream;
use datatypes::schema::SchemaRef;
use snafu::OptionExt;
pub struct SimpleTableScan {
stream: Mutex<Option<SendableRecordBatchStream>>,
schema: SchemaRef,
}
impl Debug for SimpleTableScan {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SimpleTableScan")
.field("stream", &"<SendableRecordBatchStream>")
.field("schema", &self.schema)
.finish()
}
}
impl SimpleTableScan {
pub fn new(stream: SendableRecordBatchStream) -> Self {
let schema = stream.schema();
Self {
stream: Mutex::new(Some(stream)),
schema,
}
}
}
#[async_trait]
impl PhysicalPlan for SimpleTableScan {
fn as_any(&self) -> &dyn Any {
self
}
fn schema(&self) -> SchemaRef {
self.schema.clone()
}
fn output_partitioning(&self) -> Partitioning {
Partitioning::UnknownPartitioning(1)
}
fn children(&self) -> Vec<PhysicalPlanRef> {
vec![]
}
fn with_new_children(&self, _children: Vec<PhysicalPlanRef>) -> QueryResult<PhysicalPlanRef> {
unimplemented!()
}
async fn execute(
&self,
_partition: usize,
_runtime: Arc<RuntimeEnv>,
) -> QueryResult<SendableRecordBatchStream> {
let mut stream = self.stream.lock().unwrap();
Ok(stream.take().context(query_error::ExecuteRepeatedlySnafu)?)
}
}
#[cfg(test)]
mod test {
use common_recordbatch::util;
use common_recordbatch::{RecordBatch, RecordBatches};
use datatypes::data_type::ConcreteDataType;
use datatypes::schema::{ColumnSchema, Schema};
use datatypes::vectors::Int32Vector;
use super::*;
#[tokio::test]
async fn test_simple_table_scan() {
let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
"a",
ConcreteDataType::int32_datatype(),
false,
)]));
let batch1 = RecordBatch::new(
schema.clone(),
vec![Arc::new(Int32Vector::from_slice(&[1, 2])) as _],
)
.unwrap();
let batch2 = RecordBatch::new(
schema.clone(),
vec![Arc::new(Int32Vector::from_slice(&[3, 4, 5])) as _],
)
.unwrap();
let recordbatches =
RecordBatches::try_new(schema.clone(), vec![batch1.clone(), batch2.clone()]).unwrap();
let stream = recordbatches.as_stream();
let scan = SimpleTableScan::new(stream);
assert_eq!(scan.schema(), schema);
let runtime = Arc::new(RuntimeEnv::default());
let stream = scan.execute(0, runtime.clone()).await.unwrap();
let recordbatches = util::collect(stream).await.unwrap();
assert_eq!(recordbatches[0], batch1);
assert_eq!(recordbatches[1], batch2);
let result = scan.execute(0, runtime).await;
assert!(result.is_err());
match result {
Err(e) => assert!(e
.to_string()
.contains("Not expected to run ExecutionPlan more than once")),
_ => unreachable!(),
}
}
}

9
src/table/src/test_util/empty_table.rs
View File

@@ -1,11 +1,13 @@
use std::sync::Arc;
use async_trait::async_trait;
use common_recordbatch::{EmptyRecordBatchStream, SendableRecordBatchStream};
use common_query::physical_plan::PhysicalPlanRef;
use common_recordbatch::EmptyRecordBatchStream;
use crate::metadata::TableInfoBuilder;
use crate::metadata::TableInfoRef;
use crate::requests::InsertRequest;
use crate::table::scan::SimpleTableScan;
use crate::Result;
use crate::{
metadata::{TableMetaBuilder, TableType},
@@ -64,7 +66,8 @@ impl Table for EmptyTable {
_projection: &Option<Vec<usize>>,
_filters: &[common_query::prelude::Expr],
_limit: Option<usize>,
) -> Result<SendableRecordBatchStream> {
Ok(Box::pin(EmptyRecordBatchStream::new(self.schema())))
) -> Result<PhysicalPlanRef> {
let scan = SimpleTableScan::new(Box::pin(EmptyRecordBatchStream::new(self.schema())));
Ok(Arc::new(scan))
}
}

19
src/table/src/test_util/memtable.rs
View File

@@ -3,9 +3,10 @@ use std::pin::Pin;
use std::sync::Arc;
use async_trait::async_trait;
use common_query::physical_plan::PhysicalPlanRef;
use common_query::prelude::Expr;
use common_recordbatch::error::Result as RecordBatchResult;
use common_recordbatch::{RecordBatch, RecordBatchStream, SendableRecordBatchStream};
use common_recordbatch::{RecordBatch, RecordBatchStream};
use datatypes::prelude::*;
use datatypes::schema::{ColumnSchema, Schema, SchemaRef};
use datatypes::vectors::UInt32Vector;
@@ -15,6 +16,7 @@ use snafu::prelude::*;
use crate::error::{Result, SchemaConversionSnafu, TableProjectionSnafu};
use crate::metadata::TableInfoRef;
use crate::table::scan::SimpleTableScan;
use crate::Table;
#[derive(Debug, Clone)]
@@ -71,7 +73,7 @@ impl Table for MemTable {
projection: &Option<Vec<usize>>,
_filters: &[Expr],
limit: Option<usize>,
) -> Result<SendableRecordBatchStream> {
) -> Result<PhysicalPlanRef> {
let df_recordbatch = if let Some(indices) = projection {
self.recordbatch
.df_recordbatch
@@ -95,10 +97,10 @@ impl Table for MemTable {
),
df_recordbatch,
};
Ok(Box::pin(MemtableStream {
Ok(Arc::new(SimpleTableScan::new(Box::pin(MemtableStream {
schema: recordbatch.schema.clone(),
recordbatch: Some(recordbatch),
}))
}))))
}
}
@@ -126,6 +128,7 @@ impl Stream for MemtableStream {
#[cfg(test)]
mod test {
use common_query::physical_plan::RuntimeEnv;
use common_recordbatch::util;
use datatypes::prelude::*;
use datatypes::schema::ColumnSchema;
@@ -138,6 +141,10 @@ mod test {
let table = build_testing_table();
let scan_stream = table.scan(&Some(vec![1]), &[], None).await.unwrap();
let scan_stream = scan_stream
.execute(0, Arc::new(RuntimeEnv::default()))
.await
.unwrap();
let recordbatch = util::collect(scan_stream).await.unwrap();
assert_eq!(1, recordbatch.len());
let columns = recordbatch[0].df_recordbatch.columns();
@@ -157,6 +164,10 @@ mod test {
let table = build_testing_table();
let scan_stream = table.scan(&None, &[], Some(2)).await.unwrap();
let scan_stream = scan_stream
.execute(0, Arc::new(RuntimeEnv::default()))
.await
.unwrap();
let recordbatch = util::collect(scan_stream).await.unwrap();
assert_eq!(1, recordbatch.len());
let columns = recordbatch[0].df_recordbatch.columns();