feat(snowflake): Add BITSHIFTLEFT/BITSHIFTRIGHT transpilation to DuckDB with INT128 casts and precedence fixes

Author: fivetran-kwoodbeck

sqlglot/dialects/duckdb.py

@@ -691,6 +691,32 @@ def _boolxor_agg_sql(self: DuckDB.Generator, expression: exp.BoolxorAgg) -> str:
     )
 
 
+def _prepare_bitshift_for_duckdb(expression: exp.Expression) -> exp.Expression:
+    """
+    Transform bitwise shift expressions for DuckDB by injecting INT128 casts.
+
+    DuckDB's bitwise shift operators don't work with BLOB/BINARY types, so we cast
+    them to INT128 for integer arithmetic.
+
+    Note: Assumes type annotation has been applied with the source dialect.
+    """
+    # Unwrap BINARY/VARBINARY casts that DuckDB can't handle
+    if isinstance(expression.this, exp.Cast) and _is_binary(expression.this):
+        expression.this.replace(expression.this.this)
+
+    # Check if the input is a BLOB/BINARY type (using is_type) or requires INT128
+    # If so, cast to INT128 for integer arithmetic
+    if _is_binary(expression.this) or expression.args.get("requires_int128"):
+        expression.this.replace(exp.cast(expression.this, exp.DataType.Type.INT128))
+
+    # Wrap in parentheses if parent is a bitwise operator to "fix" DuckDB precedence issue
+    # DuckDB parses: a << b | c << d  as  (a << b | c) << d
+    if isinstance(expression.parent, (exp.BitwiseAnd, exp.BitwiseOr, exp.BitwiseXor)):
+        return exp.paren(expression, copy=False)
+
+    return expression
+
+
 def _scale_rounding_sql(
     self: DuckDB.Generator,
     expression: exp.Expression,
@@ -1230,8 +1256,10 @@ class Generator(generator.Generator):
             ),
             exp.BitwiseAnd: lambda self, e: self._bitwise_op(e, "&"),
             exp.BitwiseAndAgg: _bitwise_agg_sql,
+            exp.BitwiseLeftShift: transforms.preprocess([_prepare_bitshift_for_duckdb]),
             exp.BitwiseOr: lambda self, e: self._bitwise_op(e, "|"),
             exp.BitwiseOrAgg: _bitwise_agg_sql,
+            exp.BitwiseRightShift: transforms.preprocess([_prepare_bitshift_for_duckdb]),
             exp.BitwiseXorAgg: _bitwise_agg_sql,
             exp.CommentColumnConstraint: no_comment_column_constraint_sql,
             exp.CosineDistance: rename_func("LIST_COSINE_DISTANCE"),

sqlglot/dialects/snowflake.py

@@ -163,7 +163,16 @@ def _builder(args: t.List) -> B | exp.Anonymous:
                 )
             return exp.Anonymous(this=name, expressions=args)
 
-        return binary_from_function(expr_type)(args)
+        result = binary_from_function(expr_type)(args)
+
+        # Snowflake specifies INT128 for bitwise shifts
+        if expr_type in (exp.BitwiseLeftShift, exp.BitwiseRightShift):
+            result.set("requires_int128", True)
+            # Mark HexStrings as integers for proper rendering in target dialects
+            for hexstr in result.find_all(exp.HexString):
+                hexstr.set("is_integer", True)
+
+        return result
 
     return _builder
 

sqlglot/expressions.py

@@ -5143,15 +5143,15 @@ class BitwiseAnd(Binary):
 
 
 class BitwiseLeftShift(Binary):
-    pass
+    arg_types = {"this": True, "expression": True, "requires_int128": False}
 
 
 class BitwiseOr(Binary):
     arg_types = {"this": True, "expression": True, "padside": False}
 
 
 class BitwiseRightShift(Binary):
-    pass
+    arg_types = {"this": True, "expression": True, "requires_int128": False}
 
 
 class BitwiseXor(Binary):

tests/dialects/test_snowflake.py

@@ -1691,6 +1691,46 @@ def test_snowflake(self):
         self.validate_identity("SELECT BIT_SHIFTLEFT(a, 1)", "SELECT BITSHIFTLEFT(a, 1)")
         self.validate_identity("SELECT BITSHIFTRIGHT(a, 1)")
         self.validate_identity("SELECT BIT_SHIFTRIGHT(a, 1)", "SELECT BITSHIFTRIGHT(a, 1)")
+
+        # duckdb has an order of operations precedence issue with bitshift and bitwise operators
+        self.validate_all(
+            "SELECT BITOR(BITSHIFTLEFT(5, 16), BITSHIFTLEFT(3, 8))",
+            write={"duckdb": "SELECT (CAST(5 AS INT128) << 16) | (CAST(3 AS INT128) << 8)"},
+        )
+        self.validate_all(
+            "SELECT BITAND(BITSHIFTLEFT(255, 4), BITSHIFTLEFT(15, 2))",
+            write={"duckdb": "SELECT (CAST(255 AS INT128) << 4) & (CAST(15 AS INT128) << 2)"},
+        )
+        self.validate_all(
+            "SELECT BITSHIFTLEFT(255, 4)",
+            write={"duckdb": "SELECT CAST(255 AS INT128) << 4"},
+        )
+        self.validate_all(
+            "SELECT BITSHIFTLEFT(255, 4)",
+            write={"duckdb": "SELECT CAST(255 AS INT128) << 4"},
+        )
+        self.validate_all(
+            "SELECT BITSHIFTLEFT(CAST(255 AS BINARY), 4)",
+            write={"duckdb": "SELECT CAST(255 AS INT128) << 4"},
+        )
+        self.validate_all(
+            "SELECT BITSHIFTLEFT(X'FF', 4)",
+            write={"duckdb": "SELECT CAST(255 AS INT128) << 4"},
+        )
+
+        self.validate_all(
+            "SELECT BITSHIFTRIGHT(1024, 4)",
+            write={"duckdb": "SELECT CAST(1024 AS INT128) >> 4"},
+        )
+        self.validate_all(
+            "SELECT BITSHIFTRIGHT(CAST(255 AS BINARY), 4)",
+            write={"duckdb": "SELECT CAST(255 AS INT128) >> 4"},
+        )
+        self.validate_all(
+            "SELECT BITSHIFTRIGHT(X'FF', 4)",
+            write={"duckdb": "SELECT CAST(255 AS INT128) >> 4"},
+        )
+
         self.validate_all(
             "OCTET_LENGTH('A')",
             read={

tests/fixtures/optimizer/annotate_types.sql

@@ -186,3 +186,35 @@ FLOAT64;
 # dialect: bigquery
 CASE WHEN x < y THEN CAST(3.5 AS BIGNUMERIC) WHEN x > y THEN 3/10 ELSE 2 END;
 FLOAT64;
+
+# dialect: snowflake
+BITSHIFTLEFT(255, 4);
+INT;
+
+# dialect: snowflake
+BITSHIFTRIGHT(1024, 2);
+INT;
+
+# dialect: snowflake
+BITSHIFTLEFT(CAST(255 AS BINARY), 4);
+BINARY;
+
+# dialect: snowflake
+BITSHIFTRIGHT(CAST(255 AS BINARY), 4);
+BINARY;
+
+# dialect: snowflake
+BITSHIFTLEFT(X'FF', 4);
+BINARY;
+
+# dialect: snowflake
+BITSHIFTRIGHT(X'FF', 4);
+BINARY;
+
+# dialect: snowflake
+BITOR(BITSHIFTLEFT(5, 16), BITSHIFTLEFT(3, 8));
+INT;
+
+# dialect: snowflake
+BITAND(BITSHIFTLEFT(255, 4), BITSHIFTLEFT(15, 2));
+INT;