-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
11 changed files
with
397 additions
and
348 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,11 +1,125 @@ | ||
| use std::sync::Arc; | ||
| use std::{collections::HashSet, fmt::Display, sync::Arc}; | ||
|
|
||
| use super::table_relation::TableRelation; | ||
| use arrow::datatypes::SchemaRef; | ||
| use crate::{ | ||
| error::{Error, Result}, | ||
| internal_err, | ||
| logical::expr::Column, | ||
| }; | ||
| use arrow::datatypes::{Schema, SchemaRef}; | ||
|
|
||
| pub type TableSchemaRef = Arc<TableSchema>; | ||
|
|
||
| #[derive(Debug, Clone, PartialEq, Eq, Hash)] | ||
| pub struct TableSchema { | ||
| pub schema: SchemaRef, | ||
| pub qualified_name: Vec<Option<TableRelation>>, | ||
| pub field_qualifiers: Vec<Option<TableRelation>>, | ||
| } | ||
|
|
||
| impl TableSchema { | ||
| pub fn try_from_qualified_schema(relation: impl Into<TableRelation>, schema: SchemaRef) -> Result<Self> { | ||
| Ok(Self { | ||
| field_qualifiers: vec![Some(relation.into()); schema.fields().len()], | ||
| schema, | ||
| }) | ||
| } | ||
|
|
||
| pub fn empty() -> Self { | ||
| Self { | ||
| schema: Arc::new(Schema::empty()), | ||
| field_qualifiers: vec![], | ||
| } | ||
| } | ||
|
|
||
| pub fn arrow_schema(&self) -> SchemaRef { | ||
| self.schema.clone() | ||
| } | ||
|
|
||
| pub fn has_field(&self, qualifier: Option<TableRelation>, name: &str) -> bool { | ||
| match (self.schema.index_of(name).ok(), qualifier) { | ||
| (Some(i), Some(q)) => self.field_qualifiers[i] == Some(q), | ||
| (Some(_), None) => true, | ||
| _ => false, | ||
| } | ||
| } | ||
|
|
||
| pub fn columns(&self) -> Vec<Column> { | ||
| self.schema | ||
| .fields() | ||
| .iter() | ||
| .zip(self.field_qualifiers.iter()) | ||
| .map(|(f, q)| Column::new(f.name(), q.clone())) | ||
| .collect() | ||
| } | ||
| } | ||
|
|
||
| impl TableSchema { | ||
| pub fn merge(schemas: Vec<TableSchemaRef>) -> Result<Self> { | ||
| let fields = schemas | ||
| .into_iter() | ||
| .map(|s| { | ||
| let s = Arc::unwrap_or_clone(s); | ||
| let fields = s.schema.fields().iter().map(|f| f.as_ref().clone()); | ||
| let field_qualifiers = s.field_qualifiers.into_iter(); | ||
| fields.zip(field_qualifiers).collect::<Vec<_>>() | ||
| }) | ||
| .flatten() | ||
| .collect::<Vec<_>>(); | ||
|
|
||
| // check if the number of fields and qualifiers are the same | ||
| let mut new_fields = HashSet::new(); | ||
| for (f, q) in &fields { | ||
| if !new_fields.insert((f, q)) { | ||
| return internal_err!( | ||
| "Try merge schema failed, column [{}] is ambiguous, please use qualified name to disambiguate", | ||
| f.name() | ||
| ); | ||
| } | ||
| } | ||
|
|
||
| let (fields, field_qualifiers): (Vec<_>, Vec<_>) = fields.into_iter().unzip(); | ||
|
|
||
| Ok(TableSchema { | ||
| schema: Arc::new(Schema::new(fields)), | ||
| field_qualifiers, | ||
| }) | ||
| } | ||
| } | ||
|
|
||
| impl Display for TableSchema { | ||
| fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { | ||
| write!( | ||
| f, | ||
| "{}", | ||
| self.schema | ||
| .fields() | ||
| .iter() | ||
| .zip(self.field_qualifiers.iter()) | ||
| .map(|(f, q)| qualified_name(q, &f.name())) | ||
| .collect::<Vec<_>>() | ||
| .join(", ") | ||
| ) | ||
| } | ||
| } | ||
|
|
||
| impl From<SchemaRef> for TableSchema { | ||
| fn from(value: SchemaRef) -> Self { | ||
| TableSchema { | ||
| schema: value, | ||
| field_qualifiers: vec![], | ||
| } | ||
| } | ||
| } | ||
|
|
||
| impl From<Schema> for TableSchema { | ||
| fn from(value: Schema) -> Self { | ||
| TableSchema::from(Arc::new(value)) | ||
| } | ||
| } | ||
|
|
||
| pub fn qualified_name(qualifier: &Option<TableRelation>, name: &str) -> String { | ||
| match qualifier { | ||
| Some(q) => format!("{}.{}", q, name), | ||
| None => name.to_string(), | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,113 +1,69 @@ | ||
| use std::{fmt::Display, sync::Arc, hash::{Hash, Hasher}}; | ||
| use std::{ | ||
| fmt::Display, | ||
| hash::{Hash, Hasher}, | ||
| sync::Arc, | ||
| }; | ||
|
|
||
| use arrow::datatypes::{Schema, SchemaRef}; | ||
| use arrow::datatypes::SchemaRef; | ||
|
|
||
| use crate::arrow_err; | ||
| use crate::common::table_relation::TableRelation; | ||
| use crate::error::{Error, Result}; | ||
| use crate::common::table_schema::TableSchemaRef; | ||
| use crate::common::{table_relation::TableRelation, table_schema::TableSchema}; | ||
| use crate::error::Result; | ||
| use crate::logical::expr::LogicalExpr; | ||
| use crate::provider::table::TableProvider; | ||
|
|
||
| use super::LogicalPlan; | ||
|
|
||
| #[derive(Debug, Clone)] | ||
| pub struct TableScan { | ||
| pub relation: TableRelation, | ||
| pub table_name: TableRelation, | ||
| pub source: Arc<dyn TableProvider>, | ||
| pub projections: Option<Vec<String>>, | ||
| pub projected_schema: SchemaRef, | ||
| pub filter: Option<LogicalExpr>, | ||
| pub schema: TableSchemaRef, | ||
| } | ||
|
|
||
| impl TableScan { | ||
| pub fn try_new( | ||
| relation: impl Into<TableRelation>, | ||
| source: Arc<dyn TableProvider>, | ||
| projections: Option<Vec<String>>, | ||
| filter: Option<LogicalExpr>, | ||
| ) -> Result<Self> { | ||
| let relation = relation.into(); | ||
|
|
||
| let projected_schema = projections | ||
| .as_ref() | ||
| .map(|pj| { | ||
| pj.iter() | ||
| .map(|name| { | ||
| source | ||
| .schema() | ||
| .field_with_name(name) | ||
| .map_err(|err| arrow_err!(err)) | ||
| .cloned() | ||
| }) | ||
| .collect::<Result<Vec<_>>>() | ||
| }) | ||
| .unwrap_or(Ok(source | ||
| .schema() | ||
| .fields() | ||
| .iter() | ||
| .map(|f| f.as_ref().clone()) | ||
| .collect())) | ||
| .map(|fields| Arc::new(Schema::new(fields)))?; | ||
|
|
||
| let table_name = relation.into(); | ||
| Ok(Self { | ||
| relation, | ||
| source, | ||
| projections, | ||
| projected_schema, | ||
| filter, | ||
| schema: TableSchema::try_from_qualified_schema(table_name.clone(), source.schema()).map(Arc::new)?, | ||
| source, | ||
| table_name, | ||
| }) | ||
| } | ||
|
|
||
| pub fn schema(&self) -> SchemaRef { | ||
| self.projected_schema.clone() | ||
| self.schema.arrow_schema() | ||
| } | ||
|
|
||
| pub fn children(&self) -> Option<Vec<&LogicalPlan>> { | ||
| None | ||
| } | ||
|
|
||
| pub fn set_metadata(&mut self, k: &str, v: &str) { | ||
| let schema = self.projected_schema.as_ref().clone(); | ||
| let mut metadata = schema.metadata; | ||
| metadata.insert(k.to_owned(), v.to_owned()); | ||
|
|
||
| self.projected_schema = Arc::new(Schema::new_with_metadata(schema.fields, metadata)); | ||
| } | ||
| } | ||
|
|
||
| impl Display for TableScan { | ||
| fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { | ||
| if let Some(projections) = &self.projections { | ||
| write!( | ||
| f, | ||
| "TableScan: {} Projection: [{}]", | ||
| self.relation.to_quanlify_name(), | ||
| projections.join(", ") | ||
| ) | ||
| } else { | ||
| write!(f, "TableScan: {}", self.relation.to_quanlify_name()) | ||
| } | ||
| write!(f, "TableScan: {}", self.table_name.to_quanlify_name(),) | ||
| } | ||
| } | ||
|
|
||
| impl PartialEq for TableScan { | ||
| fn eq(&self, other: &Self) -> bool { | ||
| self.relation == other.relation | ||
| && Arc::ptr_eq(&self.source, &other.source) | ||
| && self.projections == other.projections | ||
| && Arc::ptr_eq(&self.projected_schema, &other.projected_schema) | ||
| && self.filter == other.filter | ||
| self.table_name == other.table_name && Arc::ptr_eq(&self.source, &other.source) | ||
| } | ||
| } | ||
|
|
||
| impl Eq for TableScan {} | ||
|
|
||
| impl Hash for TableScan { | ||
| fn hash<H: Hasher>(&self, state: &mut H) { | ||
| self.relation.hash(state); | ||
| Arc::as_ptr(&self.source).hash(state); | ||
| self.projections.hash(state); | ||
| Arc::as_ptr(&self.projected_schema).hash(state); | ||
| self.table_name.hash(state); | ||
| self.schema.hash(state); | ||
| self.filter.hash(state); | ||
| } | ||
| } |
Oops, something went wrong.