Skip to content

Commit

Permalink
feat(rust,python,cli): add SQL engine support for string cast to `jso…
Browse files Browse the repository at this point in the history
  • Loading branch information
alexander-beedie authored Jan 11, 2024
1 parent a8bdc76 commit c105516
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 1 deletion.
6 changes: 5 additions & 1 deletion crates/polars-sql/src/sql_expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -439,9 +439,13 @@ impl SQLExprVisitor<'_> {
if format.is_some() {
return Err(polars_err!(ComputeError: "unsupported use of FORMAT in CAST expression"));
}
let polars_type = map_sql_polars_datatype(data_type)?;
let expr = self.visit_expr(expr)?;

#[cfg(feature = "json")]
if data_type == &SQLDataType::JSON {
return Ok(expr.str().json_decode(None, None));
}
let polars_type = map_sql_polars_datatype(data_type)?;
Ok(expr.cast(polars_type))
}

Expand Down
29 changes: 29 additions & 0 deletions py-polars/tests/unit/sql/test_cast.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

import polars as pl
from polars.exceptions import ComputeError
from polars.testing import assert_frame_equal


def test_cast() -> None:
Expand Down Expand Up @@ -65,3 +66,31 @@ def test_cast() -> None:
pl.SQLContext(df=df, eager_execution=True).execute(
"SELECT CAST(a AS STRING FORMAT 'HEX') FROM df"
)


def test_cast_json() -> None:
df = pl.DataFrame({"txt": ['{"a":[1,2,3],"b":["x","y","z"],"c":5.0}']})

with pl.SQLContext(df=df, eager_execution=True) as ctx:
for json_cast in ("txt::json", "CAST(txt AS JSON)"):
res = ctx.execute(f"SELECT {json_cast} AS j FROM df")

assert res.schema == {
"j": pl.Struct(
{
"a": pl.List(pl.Int64),
"b": pl.List(pl.String),
"c": pl.Float64,
},
)
}
assert_frame_equal(
res.unnest("j"),
pl.DataFrame(
{
"a": [[1, 2, 3]],
"b": [["x", "y", "z"]],
"c": [5.0],
}
),
)

0 comments on commit c105516

Please sign in to comment.