diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py index 7a21563fde..e578a676a1 100644 --- a/sqlglot/expressions.py +++ b/sqlglot/expressions.py @@ -3077,12 +3077,36 @@ def lock(self, update: bool = True, copy: bool = True) -> Select: Returns: The modified expression. """ - inst = _maybe_copy(self, copy) inst.set("locks", [Lock(update=update)]) return inst + def hint(self, *hints: ExpOrStr, dialect: DialectType = None, copy: bool = True) -> Select: + """ + Set hints for this expression. + + Examples: + >>> Select().select("x").from_("tbl").hint("BROADCAST(y)").sql(dialect="spark") + 'SELECT /*+ BROADCAST(y) */ x FROM tbl' + + Args: + hints: The SQL code strings to parse as the hints. + If an `Expression` instance is passed, it will be used as-is. + dialect: The dialect used to parse the hints. + copy: If `False`, modify this expression instance in-place. + + Returns: + The modified expression. + """ + inst = _maybe_copy(self, copy) + inst.set( + "hint", + Hint(expressions=[maybe_parse(h, copy=copy, dialect=dialect) for h in hints if h]), + ) + + return inst + @property def named_selects(self) -> t.List[str]: return [e.output_name for e in self.expressions if e.alias_or_name] diff --git a/tests/test_build.py b/tests/test_build.py index 1e286899fa..f3546405ac 100644 --- a/tests/test_build.py +++ b/tests/test_build.py @@ -127,6 +127,16 @@ def test_build(self): "SELECT x FROM tbl WHERE x > 0 FOR SHARE", "postgres", ), + ( + lambda: select("x").from_("tbl").hint("repartition(100)"), + "SELECT /*+ REPARTITION(100) */ x FROM tbl", + "spark", + ), + ( + lambda: select("x").from_("tbl").hint("coalesce(3)", "broadcast(x)"), + "SELECT /*+ COALESCE(3), BROADCAST(x) */ x FROM tbl", + "spark", + ), ( lambda: select("x", "y").from_("tbl").group_by("x"), "SELECT x, y FROM tbl GROUP BY x",