Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Update for API changes
Browse files Browse the repository at this point in the history
alamb committed Jan 19, 2025
1 parent 65f92e4 commit 4b06bc8
Showing 11 changed files with 214 additions and 104 deletions.
1 change: 1 addition & 0 deletions datafusion/expr/src/logical_plan/statement.rs
Original file line number Diff line number Diff line change
@@ -153,6 +153,7 @@ pub enum TransactionIsolationLevel {
ReadCommitted,
RepeatableRead,
Serializable,
Snapshot,
}

/// Indicator that the following statements should be committed or rolled back atomically
219 changes: 138 additions & 81 deletions datafusion/sql/src/expr/mod.rs
Original file line number Diff line number Diff line change
@@ -21,14 +21,14 @@ use datafusion_expr::planner::{
PlannerResult, RawBinaryExpr, RawDictionaryExpr, RawFieldAccessExpr,
};
use sqlparser::ast::{
BinaryOperator, CastFormat, CastKind, DataType as SQLDataType, DictionaryField,
Expr as SQLExpr, ExprWithAlias as SQLExprWithAlias, MapEntry, StructField, Subscript,
TrimWhereField, Value,
AccessExpr, BinaryOperator, CastFormat, CastKind, DataType as SQLDataType,
DictionaryField, Expr as SQLExpr, ExprWithAlias as SQLExprWithAlias, MapEntry,
StructField, Subscript, TrimWhereField, Value,
};

use datafusion_common::{
internal_datafusion_err, internal_err, not_impl_err, plan_err, DFSchema, Result,
ScalarValue,
internal_datafusion_err, internal_err, not_impl_err, plan_err, Column, DFSchema,
Result, ScalarValue,
};
use datafusion_expr::expr::ScalarFunction;
use datafusion_expr::expr::{InList, WildcardOptions};
@@ -236,14 +236,14 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
self.sql_identifier_to_expr(id, schema, planner_context)
}

SQLExpr::MapAccess { .. } => {
not_impl_err!("Map Access")
}

// <expr>["foo"], <expr>[4] or <expr>[4:5]
SQLExpr::Subscript { expr, subscript } => {
self.sql_subscript_to_expr(*expr, subscript, schema, planner_context)
}
SQLExpr::CompoundFieldAccess { root, access_chain } => self
.sql_compound_field_access_to_expr(
*root,
access_chain,
schema,
planner_context,
),

SQLExpr::CompoundIdentifier(ids) => {
self.sql_compound_identifier_to_expr(ids, schema, planner_context)
@@ -984,84 +984,141 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
Ok(Expr::Cast(Cast::new(Box::new(expr), dt)))
}

fn sql_subscript_to_expr(
fn sql_compound_field_access_to_expr(
&self,
expr: SQLExpr,
subscript: Box<Subscript>,
root: SQLExpr,
access_chain: Vec<AccessExpr>,
schema: &DFSchema,
planner_context: &mut PlannerContext,
) -> Result<Expr> {
let expr = self.sql_expr_to_logical_expr(expr, schema, planner_context)?;

let field_access = match *subscript {
Subscript::Index { index } => {
// index can be a name, in which case it is a named field access
match index {
SQLExpr::Value(
Value::SingleQuotedString(s) | Value::DoubleQuotedString(s),
) => GetFieldAccess::NamedStructField {
name: ScalarValue::from(s),
},
SQLExpr::JsonAccess { .. } => {
return not_impl_err!("JsonAccess");
let mut root = self.sql_expr_to_logical_expr(root, schema, planner_context)?;
let fields = access_chain
.into_iter()
.map(|field| match field {
AccessExpr::Subscript(subscript) => {
match subscript {
Subscript::Index { index } => {
// index can be a name, in which case it is a named field access
match index {
SQLExpr::Value(
Value::SingleQuotedString(s)
| Value::DoubleQuotedString(s),
) => Ok(Some(GetFieldAccess::NamedStructField {
name: ScalarValue::from(s),
})),
SQLExpr::JsonAccess { .. } => {
not_impl_err!("JsonAccess")
}
// otherwise treat like a list index
_ => Ok(Some(GetFieldAccess::ListIndex {
key: Box::new(self.sql_expr_to_logical_expr(
index,
schema,
planner_context,
)?),
})),
}
}
Subscript::Slice {
lower_bound,
upper_bound,
stride,
} => {
// Means access like [:2]
let lower_bound = if let Some(lower_bound) = lower_bound {
self.sql_expr_to_logical_expr(
lower_bound,
schema,
planner_context,
)
} else {
not_impl_err!("Slice subscript requires a lower bound")
}?;

// means access like [2:]
let upper_bound = if let Some(upper_bound) = upper_bound {
self.sql_expr_to_logical_expr(
upper_bound,
schema,
planner_context,
)
} else {
not_impl_err!("Slice subscript requires an upper bound")
}?;

// stride, default to 1
let stride = if let Some(stride) = stride {
self.sql_expr_to_logical_expr(
stride,
schema,
planner_context,
)?
} else {
lit(1i64)
};

Ok(Some(GetFieldAccess::ListRange {
start: Box::new(lower_bound),
stop: Box::new(upper_bound),
stride: Box::new(stride),
}))
}
}
// otherwise treat like a list index
_ => GetFieldAccess::ListIndex {
key: Box::new(self.sql_expr_to_logical_expr(
index,
schema,
planner_context,
)?),
},
}
}
Subscript::Slice {
lower_bound,
upper_bound,
stride,
} => {
// Means access like [:2]
let lower_bound = if let Some(lower_bound) = lower_bound {
self.sql_expr_to_logical_expr(lower_bound, schema, planner_context)
} else {
not_impl_err!("Slice subscript requires a lower bound")
}?;

// means access like [2:]
let upper_bound = if let Some(upper_bound) = upper_bound {
self.sql_expr_to_logical_expr(upper_bound, schema, planner_context)
} else {
not_impl_err!("Slice subscript requires an upper bound")
}?;

// stride, default to 1
let stride = if let Some(stride) = stride {
self.sql_expr_to_logical_expr(stride, schema, planner_context)?
} else {
lit(1i64)
};

GetFieldAccess::ListRange {
start: Box::new(lower_bound),
stop: Box::new(upper_bound),
stride: Box::new(stride),
AccessExpr::Dot(expr) => {
let expr =
self.sql_expr_to_logical_expr(expr, schema, planner_context)?;
match expr {
Expr::Column(Column { name, relation }) => {
if let Some(relation) = &relation {
// If the first part of the dot access is a column reference, we should
// check if the column is from the same table as the root expression.
// If it is, we should replace the root expression with the column reference.
// Otherwise, we should treat the dot access as a named field access.
if relation.table() == root.schema_name().to_string() {
root = Expr::Column(Column {
name,
relation: Some(relation.clone()),
});
Ok(None)
} else {
plan_err!(
"table name mismatch: {} != {}",
relation.table(),
root.schema_name()
)
}
} else {
Ok(Some(GetFieldAccess::NamedStructField {
name: ScalarValue::from(name),
}))
}
}
_ => not_impl_err!(
"Dot access not supported for non-column expr: {expr:?}"
),
}
}
}
};
})
.collect::<Result<Vec<_>>>()?;

let mut field_access_expr = RawFieldAccessExpr { expr, field_access };
for planner in self.context_provider.get_expr_planners() {
match planner.plan_field_access(field_access_expr, schema)? {
PlannerResult::Planned(expr) => return Ok(expr),
PlannerResult::Original(expr) => {
field_access_expr = expr;
fields
.into_iter()
.flatten()
.try_fold(root, |expr, field_access| {
let mut field_access_expr = RawFieldAccessExpr { expr, field_access };
for planner in self.context_provider.get_expr_planners() {
match planner.plan_field_access(field_access_expr, schema)? {
PlannerResult::Planned(expr) => return Ok(expr),
PlannerResult::Original(expr) => {
field_access_expr = expr;
}
}
}
}
}

not_impl_err!(
"GetFieldAccess not supported by ExprPlanner: {field_access_expr:?}"
)
not_impl_err!(
"GetFieldAccess not supported by ExprPlanner: {field_access_expr:?}"
)
})
}
}

6 changes: 3 additions & 3 deletions datafusion/sql/src/parser.rs
Original file line number Diff line number Diff line change
@@ -563,7 +563,7 @@ impl<'a> DFParser<'a> {

loop {
if let Token::Word(_) = self.parser.peek_token().token {
let identifier = self.parser.parse_identifier(false)?;
let identifier = self.parser.parse_identifier()?;
partitions.push(identifier.to_string());
} else {
return self.expected("partition name", self.parser.peek_token());
@@ -666,7 +666,7 @@ impl<'a> DFParser<'a> {
}

fn parse_column_def(&mut self) -> Result<ColumnDef, ParserError> {
let name = self.parser.parse_identifier(false)?;
let name = self.parser.parse_identifier()?;
let data_type = self.parser.parse_data_type()?;
let collation = if self.parser.parse_keyword(Keyword::COLLATE) {
Some(self.parser.parse_object_name(false)?)
@@ -676,7 +676,7 @@ impl<'a> DFParser<'a> {
let mut options = vec![];
loop {
if self.parser.parse_keyword(Keyword::CONSTRAINT) {
let name = Some(self.parser.parse_identifier(false)?);
let name = Some(self.parser.parse_identifier()?);
if let Some(option) = self.parser.parse_optional_column_option()? {
options.push(ColumnOptionDef { name, option });
} else {
9 changes: 7 additions & 2 deletions datafusion/sql/src/planner.rs
Original file line number Diff line number Diff line change
@@ -430,7 +430,10 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
SQLDataType::UnsignedBigInt(_) | SQLDataType::UnsignedInt8(_) => Ok(DataType::UInt64),
SQLDataType::Float(_) => Ok(DataType::Float32),
SQLDataType::Real | SQLDataType::Float4 => Ok(DataType::Float32),
SQLDataType::Double | SQLDataType::DoublePrecision | SQLDataType::Float8 => Ok(DataType::Float64),
SQLDataType::Double(ExactNumberInfo::None) | SQLDataType::DoublePrecision | SQLDataType::Float8 => Ok(DataType::Float64),
SQLDataType::Double(ExactNumberInfo::Precision(_)|ExactNumberInfo::PrecisionAndScale(_, _)) => {
not_impl_err!("Unsupported SQL type (precision/scale not supported) {sql_type}")
}
SQLDataType::Char(_)
| SQLDataType::Text
| SQLDataType::String(_) => Ok(DataType::Utf8),
@@ -566,7 +569,9 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
| SQLDataType::MediumText
| SQLDataType::LongText
| SQLDataType::Bit(_)
|SQLDataType::BitVarying(_)
| SQLDataType::BitVarying(_)
// BIG Query UDFs
| SQLDataType::AnyType
=> not_impl_err!(
"Unsupported SQL type {sql_type:?}"
),
8 changes: 7 additions & 1 deletion datafusion/sql/src/relation/join.rs
Original file line number Diff line number Diff line change
@@ -123,7 +123,13 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
.join_on(right, join_type, Some(expr))?
.build()
}
JoinConstraint::Using(idents) => {
JoinConstraint::Using(mut object_names) => {
if object_names.len() != 1 {
return not_impl_err!(
"Only one column name is supported in USING clause"
);
};
let idents = object_names.pop().unwrap().0;
let keys: Vec<Column> = idents
.into_iter()
.map(|x| Column::from_name(self.ident_normalizer.normalize(x)))
3 changes: 3 additions & 0 deletions datafusion/sql/src/set_expr.rs
Original file line number Diff line number Diff line change
@@ -88,6 +88,9 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
(SetOperator::Except, false) => {
LogicalPlanBuilder::except(left_plan, right_plan, false)
}
(SetOperator::Minus, _) => {
not_impl_err!("MINUS Set Operator not implemented")
}
}
}
}
Loading

0 comments on commit 4b06bc8

Please sign in to comment.