Skip to content

Commit f103fa8

Browse files
fix: Don't panic in SQL temporal string check; raise suitable ColumnNotFound error (#19473)
1 parent 4ddae71 commit f103fa8

File tree

2 files changed

+35
-11
lines changed

2 files changed

+35
-11
lines changed

crates/polars-sql/src/sql_expr.rs

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -374,34 +374,35 @@ impl SQLExprVisitor<'_> {
374374
},
375375
// identify "CAST(expr AS type) <op> string" and/or "expr::type <op> string" expressions
376376
(Expr::Cast { expr, dtype, .. }, Expr::Literal(LiteralValue::String(s))) => {
377-
if let Expr::Column(name) = &**expr {
378-
(Some(name.clone()), Some(s), Some(dtype))
379-
} else {
380-
(None, Some(s), Some(dtype))
377+
match &**expr {
378+
Expr::Column(name) => (Some(name.clone()), Some(s), Some(dtype)),
379+
_ => (None, Some(s), Some(dtype)),
381380
}
382381
},
383382
_ => (None, None, None),
384383
} {
385384
if expr_dtype.is_none() && self.active_schema.is_none() {
386385
right.clone()
387386
} else {
388-
let left_dtype = expr_dtype
389-
.unwrap_or_else(|| self.active_schema.as_ref().unwrap().get(&name).unwrap());
390-
387+
let left_dtype = expr_dtype.or_else(|| {
388+
self.active_schema
389+
.as_ref()
390+
.and_then(|schema| schema.get(&name))
391+
});
391392
match left_dtype {
392-
DataType::Time if is_iso_time(s) => {
393+
Some(DataType::Time) if is_iso_time(s) => {
393394
right.clone().str().to_time(StrptimeOptions {
394395
strict: true,
395396
..Default::default()
396397
})
397398
},
398-
DataType::Date if is_iso_date(s) => {
399+
Some(DataType::Date) if is_iso_date(s) => {
399400
right.clone().str().to_date(StrptimeOptions {
400401
strict: true,
401402
..Default::default()
402403
})
403404
},
404-
DataType::Datetime(tu, tz) if is_iso_datetime(s) || is_iso_date(s) => {
405+
Some(DataType::Datetime(tu, tz)) if is_iso_datetime(s) || is_iso_date(s) => {
405406
if s.len() == 10 {
406407
// handle upcast from ISO date string (10 chars) to datetime
407408
lit(format!("{}T00:00:00", s))

py-polars/tests/unit/sql/test_miscellaneous.py

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
import pytest
88

99
import polars as pl
10-
from polars.exceptions import SQLInterfaceError, SQLSyntaxError
10+
from polars.exceptions import ColumnNotFoundError, SQLInterfaceError, SQLSyntaxError
1111
from polars.testing import assert_frame_equal
1212

1313
if TYPE_CHECKING:
@@ -362,3 +362,26 @@ def test_global_variable_inference_17398() -> None:
362362
eager=True,
363363
)
364364
assert_frame_equal(res, users)
365+
366+
367+
@pytest.mark.parametrize(
368+
"query",
369+
[
370+
"SELECT invalid_column FROM self",
371+
"SELECT key, invalid_column FROM self",
372+
"SELECT invalid_column * 2 FROM self",
373+
"SELECT * FROM self ORDER BY invalid_column",
374+
"SELECT * FROM self WHERE invalid_column = 200",
375+
"SELECT * FROM self WHERE invalid_column = '200'",
376+
"SELECT key, SUM(n) AS sum_n FROM self GROUP BY invalid_column",
377+
],
378+
)
379+
def test_invalid_cols(query: str) -> None:
380+
df = pl.DataFrame(
381+
{
382+
"key": ["xx", "xx", "yy"],
383+
"n": ["100", "200", "300"],
384+
}
385+
)
386+
with pytest.raises(ColumnNotFoundError, match="invalid_column"):
387+
df.sql(query)

0 commit comments

Comments
 (0)