feat(ir): make impure ibis.random() and ibis.uuid() functions return unique node instances

Author: kszucs

ibis/backends/bigquery/compiler.py

@@ -120,8 +120,6 @@ class BigQueryCompiler(SQLGlotCompiler):
         ops.RPad: "rpad",
         ops.Levenshtein: "edit_distance",
         ops.Modulus: "mod",
-        ops.RandomScalar: "rand",
-        ops.RandomUUID: "generate_uuid",
         ops.RegexReplace: "regexp_replace",
         ops.RegexSearch: "regexp_contains",
         ops.Time: "time",
@@ -698,3 +696,6 @@ def visit_CountDistinct(self, op, *, arg, where):
         if where is not None:
             arg = self.if_(where, arg, NULL)
         return self.f.count(sge.Distinct(expressions=[arg]))
+
+    def visit_RandomUUID(self, op, **kwargs):
+        return self.f.generate_uuid()

ibis/backends/clickhouse/compiler.py

@@ -98,8 +98,6 @@ class ClickHouseCompiler(SQLGlotCompiler):
         ops.NotNull: "isNotNull",
         ops.NullIf: "nullIf",
         ops.RStrip: "trimRight",
-        ops.RandomScalar: "randCanonical",
-        ops.RandomUUID: "generateUUIDv4",
         ops.RegexReplace: "replaceRegexpAll",
         ops.RowNumber: "row_number",
         ops.StartsWith: "startsWith",
@@ -637,6 +635,12 @@ def visit_TimestampRange(self, op, *, start, stop, step):
     def visit_RegexSplit(self, op, *, arg, pattern):
         return self.f.splitByRegexp(pattern, self.cast(arg, dt.String(nullable=False)))
 
+    def visit_RandomScalar(self, op, **kwargs):
+        return self.f.randCanonical()
+
+    def visit_RandomUUID(self, op, **kwargs):
+        return self.f.generateUUIDv4()
+
     @staticmethod
     def _generate_groups(groups):
         return groups

ibis/backends/datafusion/compiler.py

@@ -75,7 +75,6 @@ class DataFusionCompiler(SQLGlotCompiler):
         ops.Last: "last_value",
         ops.Median: "median",
         ops.StringLength: "character_length",
-        ops.RandomUUID: "uuid",
         ops.RegexSplit: "regex_split",
         ops.EndsWith: "ends_with",
         ops.ArrayIntersect: "array_intersect",

ibis/backends/duckdb/compiler.py

@@ -48,7 +48,6 @@ class DuckDBCompiler(SQLGlotCompiler):
         ops.MapMerge: "map_concat",
         ops.MapValues: "map_values",
         ops.Mode: "mode",
-        ops.RandomUUID: "uuid",
         ops.TimeFromHMS: "make_time",
         ops.TypeOf: "typeof",
         ops.GeoPoint: "st_point",
@@ -418,3 +417,9 @@ def visit_StructField(self, op, *, arg, field):
                 expression=sg.to_identifier(field, quoted=self.quoted),
             )
         return super().visit_StructField(op, arg=arg, field=field)
+
+    def visit_RandomScalar(self, op, **kwargs):
+        return self.f.random()
+
+    def visit_RandomUUID(self, op, **kwargs):
+        return self.f.uuid()

ibis/backends/flink/compiler.py

@@ -79,8 +79,6 @@ class FlinkCompiler(SQLGlotCompiler):
         ops.MapKeys: "map_keys",
         ops.MapValues: "map_values",
         ops.Power: "power",
-        ops.RandomScalar: "rand",
-        ops.RandomUUID: "uuid",
         ops.RegexSearch: "regexp",
         ops.StrRight: "right",
         ops.StringLength: "char_length",

ibis/backends/impala/compiler.py

@@ -76,7 +76,6 @@ class ImpalaCompiler(SQLGlotCompiler):
         ops.Hash: "fnv_hash",
         ops.LStrip: "ltrim",
         ops.Ln: "ln",
-        ops.RandomUUID: "uuid",
         ops.RStrip: "rtrim",
         ops.Strip: "trim",
         ops.TypeOf: "typeof",

ibis/backends/mssql/compiler.py

@@ -129,8 +129,6 @@ class MSSQLCompiler(SQLGlotCompiler):
         ops.Ln: "log",
         ops.Log10: "log10",
         ops.Power: "power",
-        ops.RandomScalar: "rand",
-        ops.RandomUUID: "newid",
         ops.Repeat: "replicate",
         ops.Reverse: "reverse",
         ops.StringAscii: "ascii",
@@ -172,6 +170,9 @@ def _minimize_spec(start, end, spec):
             return None
         return spec
 
+    def visit_RandomUUID(self, op, **kwargs):
+        return self.f.newid()
+
     def visit_StringLength(self, op, *, arg):
         """The MSSQL LEN function doesn't count trailing spaces.
 

ibis/backends/postgres/compiler.py

@@ -100,7 +100,6 @@ class PostgresCompiler(SQLGlotCompiler):
         ops.MapContains: "exist",
         ops.MapKeys: "akeys",
         ops.MapValues: "avals",
-        ops.RandomUUID: "gen_random_uuid",
         ops.RegexSearch: "regexp_like",
         ops.TimeFromHMS: "make_time",
     }
@@ -111,6 +110,9 @@ def _aggregate(self, funcname: str, *args, where):
             return sge.Filter(this=expr, expression=sge.Where(this=where))
         return expr
 
+    def visit_RandomUUID(self, op, **kwargs):
+        return self.f.gen_random_uuid()
+
     def visit_Mode(self, op, *, arg, where):
         expr = self.f.mode()
         expr = sge.WithinGroup(

ibis/backends/snowflake/compiler.py

@@ -78,7 +78,6 @@ class SnowflakeCompiler(SQLGlotCompiler):
         ops.Hash: "hash",
         ops.Median: "median",
         ops.Mode: "mode",
-        ops.RandomUUID: "uuid_string",
         ops.StringToTimestamp: "to_timestamp_tz",
         ops.TimeFromHMS: "time_from_parts",
         ops.TimestampFromYMDHMS: "timestamp_from_parts",
@@ -241,11 +240,14 @@ def visit_MapLength(self, op, *, arg):
     def visit_Log(self, op, *, arg, base):
         return self.f.log(base, arg, dialect=self.dialect)
 
-    def visit_RandomScalar(self, op):
+    def visit_RandomScalar(self, op, **kwargs):
         return self.f.uniform(
             self.f.to_double(0.0), self.f.to_double(1.0), self.f.random()
         )
 
+    def visit_RandomUUID(self, op, **kwargs):
+        return self.f.uuid_string()
+
     def visit_ApproxMedian(self, op, *, arg, where):
         return self.agg.approx_percentile(arg, 0.5, where=where)
 

ibis/backends/sql/compiler.py

@@ -262,7 +262,6 @@ class SQLGlotCompiler(abc.ABC):
         ops.Power: "pow",
         ops.RPad: "rpad",
         ops.Radians: "radians",
-        ops.RandomScalar: "random",
         ops.RegexSearch: "regexp_like",
         ops.RegexSplit: "regexp_split",
         ops.Repeat: "repeat",
@@ -687,6 +686,14 @@ def visit_Round(self, op, *, arg, digits):
             return sge.Round(this=arg, decimals=digits)
         return sge.Round(this=arg)
 
+    ### Random Noise
+
+    def visit_RandomScalar(self, op, **kwargs):
+        return self.f.rand()
+
+    def visit_RandomUUID(self, op, **kwargs):
+        return self.f.uuid()
+
     ### Dtype Dysmorphia
 
     def visit_TryCast(self, op, *, arg, to):

ibis/backends/sql/rewrites.py

@@ -152,7 +152,13 @@ def merge_select_select(_, **kwargs):
     from the inner Select are inlined into the outer Select.
     """
     # don't merge if either the outer or the inner select has window functions
-    blocking = (ops.WindowFunction, ops.ExistsSubquery, ops.InSubquery, ops.Unnest)
+    blocking = (
+        ops.WindowFunction,
+        ops.ExistsSubquery,
+        ops.InSubquery,
+        ops.Unnest,
+        ops.Impure,
+    )
     if _.find_below(blocking, filter=ops.Value):
         return _
     if _.parent.find_below(blocking, filter=ops.Value):

ibis/backends/sqlite/compiler.py

@@ -103,7 +103,6 @@ class SQLiteCompiler(SQLGlotCompiler):
         ops.Mode: "_ibis_mode",
         ops.Time: "time",
         ops.Date: "date",
-        ops.RandomUUID: "uuid",
     }
 
     def _aggregate(self, funcname: str, *args, where):
@@ -213,7 +212,7 @@ def visit_Clip(self, op, *, arg, lower, upper):
 
         return arg
 
-    def visit_RandomScalar(self, op):
+    def visit_RandomScalar(self, op, **kwargs):
         return 0.5 + self.f.random() / sge.Literal.number(float(-1 << 64))
 
     def visit_Cot(self, op, *, arg):

ibis/backends/tests/snapshots/test_sql/test_selects_with_impure_operations_not_merged/clickhouse-value0/out.sql

@@ -0,0 +1,12 @@
+SELECT
+  "t1"."x",
+  "t1"."y",
+  "t1"."z",
+  CASE WHEN "t1"."y" = "t1"."z" THEN 'big' ELSE 'small' END AS "size"
+FROM (
+  SELECT
+    "t0"."x",
+    randCanonical() AS "y",
+    randCanonical() AS "z"
+  FROM "t" AS "t0"
+) AS "t1"

ibis/backends/tests/snapshots/test_sql/test_selects_with_impure_operations_not_merged/clickhouse-value1/out.sql

@@ -0,0 +1,12 @@
+SELECT
+  "t1"."x",
+  "t1"."y",
+  "t1"."z",
+  CASE WHEN "t1"."y" = "t1"."z" THEN 'big' ELSE 'small' END AS "size"
+FROM (
+  SELECT
+    "t0"."x",
+    generateUUIDv4() AS "y",
+    generateUUIDv4() AS "z"
+  FROM "t" AS "t0"
+) AS "t1"

ibis/backends/tests/snapshots/test_sql/test_selects_with_impure_operations_not_merged/datafusion-value0/out.sql

@@ -0,0 +1,12 @@
+SELECT
+  "t1"."x",
+  "t1"."y",
+  "t1"."z",
+  CASE WHEN "t1"."y" = "t1"."z" THEN 'big' ELSE 'small' END AS "size"
+FROM (
+  SELECT
+    "t0"."x",
+    RANDOM() AS "y",
+    RANDOM() AS "z"
+  FROM "t" AS "t0"
+) AS "t1"

ibis/backends/tests/snapshots/test_sql/test_selects_with_impure_operations_not_merged/datafusion-value1/out.sql

@@ -0,0 +1,12 @@
+SELECT
+  "t1"."x",
+  "t1"."y",
+  "t1"."z",
+  CASE WHEN "t1"."y" = "t1"."z" THEN 'big' ELSE 'small' END AS "size"
+FROM (
+  SELECT
+    "t0"."x",
+    UUID() AS "y",
+    UUID() AS "z"
+  FROM "t" AS "t0"
+) AS "t1"

ibis/backends/tests/snapshots/test_sql/test_selects_with_impure_operations_not_merged/duckdb-value0/out.sql

@@ -0,0 +1,12 @@
+SELECT
+  "t1"."x",
+  "t1"."y",
+  "t1"."z",
+  CASE WHEN "t1"."y" = "t1"."z" THEN 'big' ELSE 'small' END AS "size"
+FROM (
+  SELECT
+    "t0"."x",
+    RANDOM() AS "y",
+    RANDOM() AS "z"
+  FROM "t" AS "t0"
+) AS "t1"

ibis/backends/tests/snapshots/test_sql/test_selects_with_impure_operations_not_merged/duckdb-value1/out.sql

@@ -0,0 +1,12 @@
+SELECT
+  "t1"."x",
+  "t1"."y",
+  "t1"."z",
+  CASE WHEN "t1"."y" = "t1"."z" THEN 'big' ELSE 'small' END AS "size"
+FROM (
+  SELECT
+    "t0"."x",
+    UUID() AS "y",
+    UUID() AS "z"
+  FROM "t" AS "t0"
+) AS "t1"

ibis/backends/tests/snapshots/test_sql/test_selects_with_impure_operations_not_merged/flink-value0/out.sql

@@ -0,0 +1,12 @@
+SELECT
+  `t1`.`x`,
+  `t1`.`y`,
+  `t1`.`z`,
+  IF(`t1`.`y` = `t1`.`z`, 'big', 'small') AS `size`
+FROM (
+  SELECT
+    `t0`.`x`,
+    RAND() AS `y`,
+    RAND() AS `z`
+  FROM `t` AS `t0`
+) AS `t1`

ibis/backends/tests/snapshots/test_sql/test_selects_with_impure_operations_not_merged/flink-value1/out.sql

@@ -0,0 +1,12 @@
+SELECT
+  `t1`.`x`,
+  `t1`.`y`,
+  `t1`.`z`,
+  IF(`t1`.`y` = `t1`.`z`, 'big', 'small') AS `size`
+FROM (
+  SELECT
+    `t0`.`x`,
+    UUID() AS `y`,
+    UUID() AS `z`
+  FROM `t` AS `t0`
+) AS `t1`

ibis/backends/tests/snapshots/test_sql/test_selects_with_impure_operations_not_merged/mssql-value0/out.sql

@@ -0,0 +1,12 @@
+SELECT
+  [t1].[x],
+  [t1].[y],
+  [t1].[z],
+  IIF([t1].[y] = [t1].[z], 'big', 'small') AS [size]
+FROM (
+  SELECT
+    [t0].[x],
+    RAND() AS [y],
+    RAND() AS [z]
+  FROM [t] AS [t0]
+) AS [t1]

ibis/backends/tests/snapshots/test_sql/test_selects_with_impure_operations_not_merged/mssql-value1/out.sql

@@ -0,0 +1,12 @@
+SELECT
+  [t1].[x],
+  [t1].[y],
+  [t1].[z],
+  IIF([t1].[y] = [t1].[z], 'big', 'small') AS [size]
+FROM (
+  SELECT
+    [t0].[x],
+    NEWID() AS [y],
+    NEWID() AS [z]
+  FROM [t] AS [t0]
+) AS [t1]

ibis/backends/tests/snapshots/test_sql/test_selects_with_impure_operations_not_merged/mysql-value0/out.sql

@@ -0,0 +1,12 @@
+SELECT
+  `t1`.`x`,
+  `t1`.`y`,
+  `t1`.`z`,
+  CASE WHEN `t1`.`y` = `t1`.`z` THEN 'big' ELSE 'small' END AS `size`
+FROM (
+  SELECT
+    `t0`.`x`,
+    RAND() AS `y`,
+    RAND() AS `z`
+  FROM `t` AS `t0`
+) AS `t1`

ibis/backends/tests/snapshots/test_sql/test_selects_with_impure_operations_not_merged/mysql-value1/out.sql

@@ -0,0 +1,12 @@
+SELECT
+  `t1`.`x`,
+  `t1`.`y`,
+  `t1`.`z`,
+  CASE WHEN `t1`.`y` = `t1`.`z` THEN 'big' ELSE 'small' END AS `size`
+FROM (
+  SELECT
+    `t0`.`x`,
+    UUID() AS `y`,
+    UUID() AS `z`
+  FROM `t` AS `t0`
+) AS `t1`

ibis/backends/tests/snapshots/test_sql/test_selects_with_impure_operations_not_merged/postgres-value0/out.sql

@@ -0,0 +1,12 @@
+SELECT
+  "t1"."x",
+  "t1"."y",
+  "t1"."z",
+  CASE WHEN "t1"."y" = "t1"."z" THEN 'big' ELSE 'small' END AS "size"
+FROM (
+  SELECT
+    "t0"."x",
+    RANDOM() AS "y",
+    RANDOM() AS "z"
+  FROM "t" AS "t0"
+) AS "t1"

ibis/backends/tests/snapshots/test_sql/test_selects_with_impure_operations_not_merged/postgres-value1/out.sql

@@ -0,0 +1,12 @@
+SELECT
+  "t1"."x",
+  "t1"."y",
+  "t1"."z",
+  CASE WHEN "t1"."y" = "t1"."z" THEN 'big' ELSE 'small' END AS "size"
+FROM (
+  SELECT
+    "t0"."x",
+    GEN_RANDOM_UUID() AS "y",
+    GEN_RANDOM_UUID() AS "z"
+  FROM "t" AS "t0"
+) AS "t1"