diff --git a/crates/polars-sql/src/context.rs b/crates/polars-sql/src/context.rs index b46bee0847a9..a799640060bd 100644 --- a/crates/polars-sql/src/context.rs +++ b/crates/polars-sql/src/context.rs @@ -346,10 +346,6 @@ impl SQLContext { } }, op => { - let op = match op { - SetExpr::SetOperation { op, .. } => op, - _ => unreachable!(), - }; polars_bail!(SQLInterface: "'{}' operation is currently unsupported", op) }, } diff --git a/py-polars/tests/unit/sql/test_set_ops.py b/py-polars/tests/unit/sql/test_set_ops.py index 0ad46db5daa9..2472e39474d0 100644 --- a/py-polars/tests/unit/sql/test_set_ops.py +++ b/py-polars/tests/unit/sql/test_set_ops.py @@ -89,6 +89,54 @@ def test_except_intersect_all_unsupported(op: str, op_subtype: str) -> None: pl.sql(f"SELECT * FROM df1 {op} {op_subtype} SELECT * FROM df2") +def test_update_statement_error() -> None: + df_large = pl.DataFrame( + { + "FQDN": ["c.ORG.na", "a.COM.na"], + "NS1": ["ns1.c.org.na", "ns1.d.net.na"], + "NS2": ["ns2.c.org.na", "ns2.d.net.na"], + "NS3": ["ns3.c.org.na", "ns3.d.net.na"], + } + ) + df_small = pl.DataFrame( + { + "FQDN": ["c.org.na"], + "NS1": ["ns1.c.org.na|127.0.0.1"], + "NS2": ["ns2.c.org.na|127.0.0.1"], + "NS3": ["ns3.c.org.na|127.0.0.1"], + } + ) + + # Create a context and register the tables + ctx = pl.SQLContext() + ctx.register("large", df_large) + ctx.register("small", df_small) + + with pytest.raises( + SQLInterfaceError, + match="'UPDATE large SET FQDN = u.FQDN, NS1 = u.NS1, NS2 = u.NS2, NS3 = u.NS3 FROM u WHERE large.FQDN = u.FQDN' operation is currently unsupported", + ): + ctx.execute(""" + WITH u AS ( + SELECT + small.FQDN, + small.NS1, + small.NS2, + small.NS3 + FROM small + INNER JOIN large ON small.FQDN = large.FQDN + ) + UPDATE large + SET + FQDN = u.FQDN, + NS1 = u.NS1, + NS2 = u.NS2, + NS3 = u.NS3 + FROM u + WHERE large.FQDN = u.FQDN + """) + + @pytest.mark.parametrize("op", ["EXCEPT", "INTERSECT", "UNION"]) def test_except_intersect_errors(op: str) -> None: df1 = pl.DataFrame({"x": [1, 9, 1, 1], "y": [2, 3, 4, 4], "z": [5, 5, 5, 5]}) # noqa: F841