Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat!: Dedicated SQLInterface and SQLSyntax errors #16635

Merged
merged 1 commit into from
Jun 4, 2024
Merged

feat!: Dedicated SQLInterface and SQLSyntax errors #16635

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions crates/polars-error/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,10 @@ pub enum PolarsError {
SchemaMismatch(ErrString),
#[error("lengths don't match: {0}")]
ShapeMismatch(ErrString),
#[error("{0}")]
SQLInterface(ErrString),
#[error("{0}")]
SQLSyntax(ErrString),
#[error("string caches don't match: {0}")]
StringCacheMismatch(ErrString),
#[error("field not found: {0}")]
Expand Down Expand Up @@ -198,6 +202,8 @@ impl PolarsError {
ShapeMismatch(msg) => ShapeMismatch(func(msg).into()),
StringCacheMismatch(msg) => StringCacheMismatch(func(msg).into()),
StructFieldNotFound(msg) => StructFieldNotFound(func(msg).into()),
SQLInterface(msg) => SQLInterface(func(msg).into()),
SQLSyntax(msg) => SQLSyntax(func(msg).into()),
_ => unreachable!(),
}
}
Expand Down
78 changes: 38 additions & 40 deletions crates/polars-sql/src/context.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,7 @@ impl SQLContext {
.parse_statements()
.map_err(to_compute_err)?;

polars_ensure!(ast.len() == 1, ComputeError: "One (and only one) statement at a time please");
polars_ensure!(ast.len() == 1, SQLInterface: "one (and only one) statement at a time please");
let res = self.execute_statement(ast.first().unwrap())?;

// Ensure the result uses the proper arenas.
Expand Down Expand Up @@ -168,7 +168,7 @@ impl SQLContext {
stmt @ Statement::Explain { .. } => self.execute_explain(stmt)?,
stmt @ Statement::Truncate { .. } => self.execute_truncate_table(stmt)?,
_ => polars_bail!(
ComputeError: "SQL statement type {:?} is not supported", ast,
SQLInterface: "statement type {:?} is not supported", ast,
),
})
}
Expand Down Expand Up @@ -221,9 +221,9 @@ impl SQLContext {
right,
} => self.process_union(left, right, set_quantifier, query),
SetExpr::SetOperation { op, .. } => {
polars_bail!(InvalidOperation: "'{}' operation not yet supported", op)
polars_bail!(SQLInterface: "'{}' operation not yet supported", op)
},
op => polars_bail!(InvalidOperation: "'{}' operation not yet supported", op),
op => polars_bail!(SQLInterface: "'{}' operation not yet supported", op),
}
}

Expand Down Expand Up @@ -259,7 +259,7 @@ impl SQLContext {
concatenated.map(|lf| lf.unique(None, UniqueKeepStrategy::Any))
},
#[allow(unreachable_patterns)]
_ => polars_bail!(InvalidOperation: "'UNION {}' is not yet supported", quantifier),
_ => polars_bail!(SQLInterface: "'UNION {}' is not yet supported", quantifier),
}
}

Expand Down Expand Up @@ -318,10 +318,12 @@ impl SQLContext {
.lazy();
Ok(lf.clone())
} else {
polars_bail!(ComputeError: "table '{}' does not exist", tbl);
polars_bail!(SQLInterface: "table '{}' does not exist", tbl);
}
},
_ => polars_bail!(ComputeError: "TRUNCATE does not support use of 'partitions'"),
_ => {
polars_bail!(SQLInterface: "TRUNCATE does not support use of 'partitions'")
},
}
} else {
unreachable!()
Expand All @@ -335,7 +337,7 @@ impl SQLContext {
fn register_ctes(&mut self, query: &Query) -> PolarsResult<()> {
if let Some(with) = &query.with {
if with.recursive {
polars_bail!(ComputeError: "recursive CTEs are not supported")
polars_bail!(SQLInterface: "recursive CTEs are not supported")
}
for cte in &with.cte_tables {
let cte_name = cte.alias.name.value.clone();
Expand Down Expand Up @@ -386,7 +388,7 @@ impl SQLContext {
JoinOperator::CrossJoin => lf.cross_join(rf, Some(format!(":{}", r_name))),
join_type => {
polars_bail!(
InvalidOperation:
SQLInterface:
"join type '{:?}' not yet supported by polars-sql", join_type
);
},
Expand Down Expand Up @@ -424,7 +426,7 @@ impl SQLContext {
let sql_tbl: &TableWithJoins = select_stmt
.from
.first()
.ok_or_else(|| polars_err!(ComputeError: "no table name provided in query"))?;
.ok_or_else(|| polars_err!(SQLSyntax: "no table name provided in query"))?;

let mut lf = self.execute_from_statement(sql_tbl)?;
let mut contains_wildcard = false;
Expand Down Expand Up @@ -482,8 +484,8 @@ impl SQLContext {
} if matches!(**expr, SQLExpr::Value(SQLValue::Number(_, _))) => {
if let SQLExpr::Value(SQLValue::Number(ref idx, _)) = **expr {
Err(polars_err!(
ComputeError:
"group_by error: expected a positive integer or valid expression; got -{}",
SQLSyntax:
"GROUP BY error: expected a positive integer or valid expression; got -{}",
idx
))
} else {
Expand All @@ -496,8 +498,8 @@ impl SQLContext {
Ok(projections[idx - 1].clone())
},
SQLExpr::Value(v) => Err(polars_err!(
ComputeError:
"group_by error: expected a positive integer or valid expression; got {}", v,
SQLSyntax:
"GROUP BY error: expected a positive integer or valid expression; got {}", v,
)),
_ => parse_sql_expr(e, self, schema.as_deref()),
})
Expand Down Expand Up @@ -621,10 +623,7 @@ impl SQLContext {
if let Expr::Column(name) = expr {
Ok(name.to_string())
} else {
Err(polars_err!(
ComputeError:
"DISTINCT ON only supports column names"
))
Err(polars_err!(SQLSyntax:"DISTINCT ON only supports column names"))
}
})
.collect::<PolarsResult<Vec<_>>>()?;
Expand Down Expand Up @@ -700,7 +699,7 @@ impl SQLContext {
let tbl_name = name.0.first().unwrap().value.as_str();
// CREATE TABLE IF NOT EXISTS
if *if_not_exists && self.table_map.contains_key(tbl_name) {
polars_bail!(ComputeError: "relation {} already exists", tbl_name);
polars_bail!(SQLInterface: "relation {} already exists", tbl_name);
// CREATE OR REPLACE TABLE
}
if let Some(query) = query {
Expand All @@ -713,7 +712,7 @@ impl SQLContext {
.lazy();
Ok(out)
} else {
polars_bail!(ComputeError: "only CREATE TABLE AS SELECT is supported");
polars_bail!(SQLInterface: "only `CREATE TABLE AS SELECT` is currently supported");
}
} else {
unreachable!()
Expand All @@ -740,22 +739,21 @@ impl SQLContext {
None => Ok((tbl_name.to_string(), lf)),
}
} else {
polars_bail!(ComputeError: "relation '{}' was not found", tbl_name);
polars_bail!(SQLInterface: "relation '{}' was not found", tbl_name);
}
},
TableFactor::Derived {
lateral,
subquery,
alias,
} => {
polars_ensure!(!(*lateral), ComputeError: "LATERAL not supported");

polars_ensure!(!(*lateral), SQLInterface: "LATERAL not supported");
if let Some(alias) = alias {
let lf = self.execute_query_no_ctes(subquery)?;
self.table_map.insert(alias.name.value.clone(), lf.clone());
Ok((alias.name.value.clone(), lf))
} else {
polars_bail!(ComputeError: "derived tables must have aliases");
polars_bail!(SQLSyntax: "derived tables must have aliases");
}
},
TableFactor::UNNEST {
Expand Down Expand Up @@ -784,13 +782,13 @@ impl SQLContext {
.collect::<Result<_, _>>()?;

polars_ensure!(!column_names.is_empty(),
ComputeError:
SQLSyntax:
"UNNEST table alias must also declare column names, eg: {} (a,b,c)", alias.name.to_string()
);
if column_names.len() != column_values.len() {
let plural = if column_values.len() > 1 { "s" } else { "" };
polars_bail!(
ComputeError:
SQLSyntax:
"UNNEST table alias requires {} column name{}, found {}", column_values.len(), plural, column_names.len()
);
}
Expand All @@ -810,17 +808,17 @@ impl SQLContext {
if *with_offset {
// TODO: make a PR to `sqlparser-rs` to support 'ORDINALITY'
// (note that 'OFFSET' is BigQuery-specific syntax, not PostgreSQL)
polars_bail!(ComputeError: "UNNEST tables do not (yet) support WITH OFFSET/ORDINALITY");
polars_bail!(SQLInterface: "UNNEST tables do not (yet) support WITH OFFSET/ORDINALITY");
}
self.table_map.insert(table_name.clone(), lf.clone());
Ok((table_name.clone(), lf))
} else {
polars_bail!(ComputeError: "UNNEST table must have an alias");
polars_bail!(SQLSyntax: "UNNEST table must have an alias");
}
},

// Support bare table, optional with alias for now
_ => polars_bail!(ComputeError: "not yet implemented: {}", relation),
_ => polars_bail!(SQLInterface: "not yet implemented: {}", relation),
}
}

Expand Down Expand Up @@ -858,7 +856,7 @@ impl SQLContext {
descending.push(!ob.asc.unwrap_or(true));
polars_ensure!(
ob.nulls_first.is_none(),
ComputeError: "nulls first/last is not yet supported",
SQLInterface: "NULLS FIRST/LAST is not (yet) supported",
);
}

Expand All @@ -879,7 +877,7 @@ impl SQLContext {
) -> PolarsResult<LazyFrame> {
polars_ensure!(
!contains_wildcard,
ComputeError: "group_by error: can't process wildcard in group_by"
SQLSyntax: "GROUP BY error: can't process wildcard in group_by"
);
let schema_before = lf.schema_with_arenas(&mut self.lp_arena, &mut self.expr_arena)?;
let group_by_keys_schema =
Expand Down Expand Up @@ -919,7 +917,7 @@ impl SQLContext {
} else if let Expr::Column(_) = e {
// Non-aggregated columns must be part of the GROUP BY clause
if !group_by_keys_schema.contains(&field.name) {
polars_bail!(ComputeError: "'{}' should participate in the GROUP BY clause or an aggregate function", &field.name);
polars_bail!(SQLSyntax: "'{}' should participate in the GROUP BY clause or an aggregate function", &field.name);
}
}
}
Expand Down Expand Up @@ -962,10 +960,10 @@ impl SQLContext {
) => Ok(lf.slice(
offset
.parse()
.map_err(|e| polars_err!(ComputeError: "OFFSET conversion error: {}", e))?,
.map_err(|e| polars_err!(SQLInterface: "OFFSET conversion error: {}", e))?,
limit
.parse()
.map_err(|e| polars_err!(ComputeError: "LIMIT conversion error: {}", e))?,
.map_err(|e| polars_err!(SQLInterface: "LIMIT conversion error: {}", e))?,
)),
(
Some(Offset {
Expand All @@ -976,17 +974,17 @@ impl SQLContext {
) => Ok(lf.slice(
offset
.parse()
.map_err(|e| polars_err!(ComputeError: "OFFSET conversion error: {}", e))?,
.map_err(|e| polars_err!(SQLInterface: "OFFSET conversion error: {}", e))?,
IdxSize::MAX,
)),
(None, Some(SQLExpr::Value(SQLValue::Number(limit, _)))) => Ok(lf.limit(
limit
.parse()
.map_err(|e| polars_err!(ComputeError: "LIMIT conversion error: {}", e))?,
.map_err(|e| polars_err!(SQLInterface: "LIMIT conversion error: {}", e))?,
)),
(None, None) => Ok(lf),
_ => polars_bail!(
ComputeError: "non-numeric arguments for LIMIT/OFFSET are not supported",
SQLSyntax: "non-numeric arguments for LIMIT/OFFSET are not supported",
),
}
}
Expand All @@ -1002,15 +1000,15 @@ impl SQLContext {
[tbl_name] => {
let lf = self.table_map.get_mut(&tbl_name.value).ok_or_else(|| {
polars_err!(
ComputeError: "no table named '{}' found",
SQLInterface: "no table named '{}' found",
tbl_name
)
})?;
let schema = lf.schema_with_arenas(&mut self.lp_arena, &mut self.expr_arena)?;
cols(schema.iter_names())
},
e => polars_bail!(
ComputeError: "invalid wildcard expression: {:?}",
SQLSyntax: "invalid wildcard expression: {:?}",
e
),
};
Expand All @@ -1024,7 +1022,7 @@ impl SQLContext {
contains_wildcard_exclude: &mut bool,
) -> PolarsResult<Expr> {
if options.opt_except.is_some() {
polars_bail!(InvalidOperation: "EXCEPT not supported; use EXCLUDE instead")
polars_bail!(SQLSyntax: "EXCEPT not supported; use EXCLUDE instead")
}
Ok(match &options.opt_exclude {
Some(ExcludeSelectItem::Single(ident)) => {
Expand Down
Loading