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

Author: fivetran-kwoodbeck

sqlglot/dialects/duckdb.py

@@ -642,34 +642,47 @@ def _prepare_bitshift_for_duckdb(expression: exp.Expression) -> exp.Expression:
     """
     Transform bitwise shift expressions for DuckDB by injecting INT128 casts.
 
-    Modifies the AST in-place to:
-    1. Unwrap BINARY/BLOB casts and replace with INT128 cast
-    2. Mark HexString literals as integers (sets is_integer flag for proper rendering)
-    3. Apply INT128 cast if requires_int128=True (from Snowflake)
+    DuckDB's bitwise shift operators don't work with BLOB/BINARY types, so we cast
+    them to INT128 for integer arithmetic.
 
-    After transformation, falls back to the base binary generator.
+    Note: Assumes type annotation has been applied with the source dialect.
     """
-    needs_int128 = False
-
-    # Check if left operand is a BINARY/BLOB cast - unwrap it
+    # Unwrap BINARY/VARBINARY casts that DuckDB can't handle
     if isinstance(expression.this, exp.Cast) and _is_binary(expression.this):
-        # Replace Cast(value, BINARY) with just the inner value
         expression.this.replace(expression.this.this)
-        needs_int128 = True
-
-    # Mark HexString as integer so DuckDB renders it as decimal instead of BLOB
-    if isinstance(expression.this, exp.HexString):
-        expression.this.set("is_integer", True)
-        needs_int128 = True
 
-    # Apply INT128 cast if needed
-    if needs_int128 or expression.args.get("requires_int128"):
+    # 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 _floor_sql(self: DuckDB.Generator, expression: exp.Floor) -> str:
+def _scale_rounding_sql(
+    self: DuckDB.Generator,
+    expression: exp.Expression,
+    rounding_func: type[exp.Expression],
+) -> str | None:
+    """
+    Handle scale parameter transformation for rounding functions.
+
+    DuckDB doesn't support the scale parameter for certain functions (e.g., FLOOR, CEIL),
+    so we transform: FUNC(x, n) to ROUND(FUNC(x * 10^n) / 10^n, n)
+
+    Args:
+        self: The DuckDB generator instance
+        expression: The expression to transform (must have 'this', 'decimals', and 'to' args)
+        rounding_func: The rounding function class to use in the transformation
+
+    Returns:
+        The transformed SQL string if decimals parameter exists, None otherwise
+    """
     decimals = expression.args.get("decimals")
 
     if decimals is None or expression.args.get("to") is not None:
@@ -1190,8 +1203,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,13 +163,16 @@ def _builder(args: t.List) -> B | exp.Anonymous:
                 )
             return exp.Anonymous(this=name, expressions=args)
 
+        result = binary_from_function(expr_type)(args)
+
         # Snowflake specifies INT128 for bitwise shifts
         if expr_type in (exp.BitwiseLeftShift, exp.BitwiseRightShift):
-            result = binary_from_function(expr_type)(args)
             result.set("requires_int128", True)
-            return result
+            # Mark HexStrings as integers for proper rendering in target dialects
+            for hexstr in result.find_all(exp.HexString):
+                hexstr.set("is_integer", True)
 
-        return binary_from_function(expr_type)(args)
+        return result
 
     return _builder
 

tests/dialects/test_snowflake.py

@@ -1692,25 +1692,45 @@ def test_snowflake(self):
         self.validate_identity("SELECT BITSHIFTRIGHT(a, 1)")
         self.validate_identity("SELECT BIT_SHIFTRIGHT(a, 1)", "SELECT BITSHIFTRIGHT(a, 1)")
 
-        # Test bitshift transpilation to DuckDB with INT128 casts
+        # 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 BITSHIFTRIGHT(1024, 2)",
-            write={"duckdb": "SELECT CAST(1024 AS INT128) >> 2"},
+            "SELECT BITSHIFTLEFT(255, 4)",
+            write={"duckdb": "SELECT CAST(255 AS INT128) << 4"},
         )
-        # Test BINARY cast unwrapping
         self.validate_all(
-            "SELECT BITSHIFTLEFT(X'FF'::BINARY, 4)",
+            "SELECT BITSHIFTLEFT(CAST(255 AS BINARY), 4)",
             write={"duckdb": "SELECT CAST(255 AS INT128) << 4"},
         )
-        # Test HexString with is_integer flag
         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

@@ -162,3 +162,27 @@ 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;