feat: add bigquery mod op (#3157)

* feat: add bigquery mod op * feedback * add parser func * fix tests
tobymao · Mar 18, 2024 · 706fac3 · 706fac3
1 parent be8c5ea
commit 706fac3
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 12 deletions.
diff --git a/sqlglot/dialects/bigquery.py b/sqlglot/dialects/bigquery.py
@@ -616,6 +616,7 @@ class Generator(generator.Generator):
  exp.IntDiv: rename_func("DIV"),
  exp.JSONFormat: rename_func("TO_JSON_STRING"),
  exp.Max: max_or_greatest,
+ exp.Mod: rename_func("MOD"),
  exp.MD5: lambda self, e: self.func("TO_HEX", self.func("MD5", e.this)),
  exp.MD5Digest: rename_func("MD5"),
  exp.Min: min_or_least,

diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -121,6 +121,7 @@ class Parser(metaclass=_Parser):
  "LOG": build_logarithm,
  "LOG2": lambda args: exp.Log(this=exp.Literal.number(2), expression=seq_get(args, 0)),
  "LOG10": lambda args: exp.Log(this=exp.Literal.number(10), expression=seq_get(args, 0)),
+ "MOD": lambda args: exp.Mod(this=seq_get(args, 0), expression=seq_get(args, 1)),
  "TIME_TO_TIME_STR": lambda args: exp.Cast(
  this=seq_get(args, 0),
  to=exp.DataType(this=exp.DataType.Type.TEXT),

diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py
@@ -1068,28 +1068,36 @@ def test_bigquery(self):
  )
  self.validate_all(
  """SELECT
- `u`.`harness_user_email` AS `harness_user_email`,
- `d`.`harness_user_id` AS `harness_user_id`,
- `harness_account_id` AS `harness_account_id`
-FROM `analytics_staging`.`stg_mongodb__users` AS `u`, UNNEST(`u`.`harness_cluster_details`) AS `d`, UNNEST(`d`.`harness_account_ids`) AS `harness_account_id`
+ `u`.`user_email` AS `user_email`,
+ `d`.`user_id` AS `user_id`,
+ `account_id` AS `account_id`
+FROM `analytics_staging`.`stg_mongodb__users` AS `u`, UNNEST(`u`.`cluster_details`) AS `d`, UNNEST(`d`.`account_ids`) AS `account_id`
 WHERE
- NOT `harness_account_id` IS NULL""",
+ NOT `account_id` IS NULL""",
  read={
  "": """
  SELECT
- "u"."harness_user_email" AS "harness_user_email",
- "_q_0"."d"."harness_user_id" AS "harness_user_id",
- "_q_1"."harness_account_id" AS "harness_account_id"
+ "u"."user_email" AS "user_email",
+ "_q_0"."d"."user_id" AS "user_id",
+ "_q_1"."account_id" AS "account_id"
  FROM
  "analytics_staging"."stg_mongodb__users" AS "u",
- UNNEST("u"."harness_cluster_details") AS "_q_0"("d"),
- UNNEST("_q_0"."d"."harness_account_ids") AS "_q_1"("harness_account_id")
+ UNNEST("u"."cluster_details") AS "_q_0"("d"),
+ UNNEST("_q_0"."d"."account_ids") AS "_q_1"("account_id")
  WHERE
- NOT "_q_1"."harness_account_id" IS NULL
+ NOT "_q_1"."account_id" IS NULL
  """
  },
  pretty=True,
  )
+ self.validate_all(
+ "SELECT MOD(x, 10)",
+ read={"postgres": "SELECT x % 10"},
+ write={
+ "bigquery": "SELECT MOD(x, 10)",
+ "postgres": "SELECT x % 10",
+ },
+ )
 
  def test_errors(self):
  with self.assertRaises(TokenError):

diff --git a/tests/dialects/test_duckdb.py b/tests/dialects/test_duckdb.py
@@ -754,7 +754,7 @@ def test_time(self):
  write={"duckdb": "SELECT (90 * INTERVAL '1' DAY)"},
  )
  self.validate_all(
- "SELECT ((DATE_TRUNC('DAY', CAST(CAST(DATE_TRUNC('DAY', CURRENT_TIMESTAMP) AS DATE) AS TIMESTAMP) + INTERVAL (0 - MOD((DAYOFWEEK(CAST(CAST(DATE_TRUNC('DAY', CURRENT_TIMESTAMP) AS DATE) AS TIMESTAMP)) % 7) - 1 + 7, 7)) DAY) + (7 * INTERVAL (-5) DAY))) AS t1",
+ "SELECT ((DATE_TRUNC('DAY', CAST(CAST(DATE_TRUNC('DAY', CURRENT_TIMESTAMP) AS DATE) AS TIMESTAMP) + INTERVAL (0 - (DAYOFWEEK(CAST(CAST(DATE_TRUNC('DAY', CURRENT_TIMESTAMP) AS DATE) AS TIMESTAMP)) % 7) - 1 + 7 % 7) DAY) + (7 * INTERVAL (-5) DAY))) AS t1",
  read={
  "presto": "SELECT ((DATE_ADD('week', -5, DATE_TRUNC('DAY', DATE_ADD('day', (0 - MOD((DAY_OF_WEEK(CAST(CAST(DATE_TRUNC('DAY', NOW()) AS DATE) AS TIMESTAMP)) % 7) - 1 + 7, 7)), CAST(CAST(DATE_TRUNC('DAY', NOW()) AS DATE) AS TIMESTAMP)))))) AS t1",
  },