fix(duckdb): avoid literals casts that might defeat optimization

ibis-project · Jul 23, 2024 · e4ff1bd · e4ff1bd
1 parent feb12f4
commit e4ff1bd
Show file tree

Hide file tree

Showing 72 changed files with 156 additions and 177 deletions.
diff --git a/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_geospatial_dwithin/out.sql b/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_geospatial_dwithin/out.sql
@@ -1,3 +1,3 @@
 SELECT
-  ST_DWITHIN("t0"."geom", "t0"."geom", CAST(3.0 AS DOUBLE)) AS "tmp"
+  ST_DWITHIN("t0"."geom", "t0"."geom", 3.0) AS "tmp"
 FROM "t" AS "t0"
diff --git a/ibis/backends/sql/compilers/duckdb.py b/ibis/backends/sql/compilers/duckdb.py
@@ -201,10 +201,17 @@ def visit_ArrayZip(self, op, *, arg):
         any_arg_null = sg.or_(*(arr.is_(NULL) for arr in arg))
         return self.if_(any_arg_null, NULL, zipped_arrays)
 
+    def visit_Array(self, op, *, exprs):
+        return self.cast(self.f.array(*exprs), op.dtype)
+
     def visit_Map(self, op, *, keys, values):
         # workaround for https://github.com/ibis-project/ibis/issues/8632
         return self.if_(
-            sg.or_(keys.is_(NULL), values.is_(NULL)), NULL, self.f.map(keys, values)
+            sg.or_(keys.is_(NULL), values.is_(NULL)),
+            NULL,
+            self.f.map(
+                self.cast(keys, op.keys.dtype), self.cast(values, op.values.dtype)
+            ),
         )
 
     def visit_MapGet(self, op, *, arg, key, default):
@@ -378,6 +385,8 @@ def visit_NonNullLiteral(self, op, *, value, dtype):
                 return self.cast(
                     str(value), to=dt.float32 if dtype.is_decimal() else dtype
                 )
+            if dtype.is_floating() or dtype.is_integer():
+                return sge.convert(value)
             return self.cast(value, dtype)
         elif dtype.is_time():
             return self.f.make_time(
@@ -401,16 +410,16 @@ def visit_NonNullLiteral(self, op, *, value, dtype):
 
             return self.f[funcname](*args)
         elif dtype.is_struct():
-            return sge.Struct.from_arg_list(
-                [
-                    sge.PropertyEQ(
-                        this=sg.to_identifier(k, quoted=self.quoted),
-                        expression=self.visit_Literal(
+            return self.cast(
+                sge.Struct.from_arg_list(
+                    [
+                        self.visit_Literal(
                             ops.Literal(v, field_dtype), value=v, dtype=field_dtype
-                        ),
-                    )
-                    for field_dtype, (k, v) in zip(dtype.types, value.items())
-                ]
+                        )
+                        for field_dtype, v in zip(dtype.types, value.values())
+                    ]
+                ),
+                op.dtype,
             )
         else:
             return None

diff --git a/ibis/backends/tests/snapshots/test_generic/test_many_subqueries/duckdb/out.sql b/ibis/backends/tests/snapshots/test_generic/test_many_subqueries/duckdb/out.sql
@@ -1,12 +1,12 @@
 WITH "t1" AS (
   SELECT
     "t0"."street",
-    ROW_NUMBER() OVER (ORDER BY "t0"."street" ASC ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) - CAST(1 AS TINYINT) AS "key"
+    ROW_NUMBER() OVER (ORDER BY "t0"."street" ASC ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) - 1 AS "key"
   FROM "data" AS "t0"
 ), "t7" AS (
   SELECT
     "t6"."street",
-    ROW_NUMBER() OVER (ORDER BY "t6"."street" ASC ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) - CAST(1 AS TINYINT) AS "key"
+    ROW_NUMBER() OVER (ORDER BY "t6"."street" ASC ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) - 1 AS "key"
   FROM (
     SELECT
       "t3"."street",

diff --git a/ibis/backends/tests/snapshots/test_sql/test_isin_bug/duckdb/out.sql b/ibis/backends/tests/snapshots/test_sql/test_isin_bug/duckdb/out.sql
@@ -4,6 +4,6 @@ SELECT
       *
     FROM "t" AS "t0"
     WHERE
-      "t0"."x" > CAST(2 AS TINYINT)
+      "t0"."x" > 2
   ) AS "InSubquery(x)"
 FROM "t" AS "t0"
diff --git a/ibis/backends/tests/snapshots/test_sql/test_union_aliasing/duckdb/out.sql b/ibis/backends/tests/snapshots/test_sql/test_union_aliasing/duckdb/out.sql
@@ -26,26 +26,26 @@ WITH "t5" AS (
           SELECT
             "t0"."field_of_study",
             UNNEST(
-              [
-                {'years': '1970-71', 'degrees': "t0"."1970-71"},
-                {'years': '1975-76', 'degrees': "t0"."1975-76"},
-                {'years': '1980-81', 'degrees': "t0"."1980-81"},
-                {'years': '1985-86', 'degrees': "t0"."1985-86"},
-                {'years': '1990-91', 'degrees': "t0"."1990-91"},
-                {'years': '1995-96', 'degrees': "t0"."1995-96"},
-                {'years': '2000-01', 'degrees': "t0"."2000-01"},
-                {'years': '2005-06', 'degrees': "t0"."2005-06"},
-                {'years': '2010-11', 'degrees': "t0"."2010-11"},
-                {'years': '2011-12', 'degrees': "t0"."2011-12"},
-                {'years': '2012-13', 'degrees': "t0"."2012-13"},
-                {'years': '2013-14', 'degrees': "t0"."2013-14"},
-                {'years': '2014-15', 'degrees': "t0"."2014-15"},
-                {'years': '2015-16', 'degrees': "t0"."2015-16"},
-                {'years': '2016-17', 'degrees': "t0"."2016-17"},
-                {'years': '2017-18', 'degrees': "t0"."2017-18"},
-                {'years': '2018-19', 'degrees': "t0"."2018-19"},
-                {'years': '2019-20', 'degrees': "t0"."2019-20"}
-              ]
+              CAST([
+                ROW('1970-71', "t0"."1970-71"),
+                ROW('1975-76', "t0"."1975-76"),
+                ROW('1980-81', "t0"."1980-81"),
+                ROW('1985-86', "t0"."1985-86"),
+                ROW('1990-91', "t0"."1990-91"),
+                ROW('1995-96', "t0"."1995-96"),
+                ROW('2000-01', "t0"."2000-01"),
+                ROW('2005-06', "t0"."2005-06"),
+                ROW('2010-11', "t0"."2010-11"),
+                ROW('2011-12', "t0"."2011-12"),
+                ROW('2012-13', "t0"."2012-13"),
+                ROW('2013-14', "t0"."2013-14"),
+                ROW('2014-15', "t0"."2014-15"),
+                ROW('2015-16', "t0"."2015-16"),
+                ROW('2016-17', "t0"."2016-17"),
+                ROW('2017-18', "t0"."2017-18"),
+                ROW('2018-19', "t0"."2018-19"),
+                ROW('2019-20', "t0"."2019-20")
+              ] AS STRUCT("years" TEXT, "degrees" BIGINT)[])
             ) AS "__pivoted__"
           FROM "humanities" AS "t0"
         ) AS "t1"
@@ -73,7 +73,7 @@ FROM (
     *
   FROM "t5" AS "t6"
   WHERE
-    "t6"."diff" < CAST(0 AS TINYINT)
+    "t6"."diff" < 0
   ORDER BY
     "t6"."diff" ASC
   LIMIT 10

diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_agg_and_non_agg_filter/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_agg_and_non_agg_filter/out.sql
@@ -2,7 +2,7 @@ SELECT
   *
 FROM "my_table" AS "t0"
 WHERE
-  "t0"."a" < CAST(100 AS TINYINT)
+  "t0"."a" < 100
   AND "t0"."a" = (
     SELECT
       MAX("t1"."a") AS "Max(a)"
@@ -11,7 +11,7 @@ WHERE
         *
       FROM "my_table" AS "t0"
       WHERE
-        "t0"."a" < CAST(100 AS TINYINT)
+        "t0"."a" < 100
     ) AS "t1"
   )
   AND "t0"."b" = 'a'
diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_agg_filter/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_agg_filter/out.sql
@@ -1,18 +1,18 @@
 SELECT
   "t0"."a",
-  "t0"."b" * CAST(2 AS TINYINT) AS "b2"
+  "t0"."b" * 2 AS "b2"
 FROM "my_table" AS "t0"
 WHERE
-  "t0"."a" < CAST(100 AS TINYINT)
+  "t0"."a" < 100
   AND "t0"."a" = (
     SELECT
       MAX("t1"."a") AS "Max(a)"
     FROM (
       SELECT
         "t0"."a",
-        "t0"."b" * CAST(2 AS TINYINT) AS "b2"
+        "t0"."b" * 2 AS "b2"
       FROM "my_table" AS "t0"
       WHERE
-        "t0"."a" < CAST(100 AS TINYINT)
+        "t0"."a" < 100
     ) AS "t1"
   )
diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_agg_filter_with_alias/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_agg_filter_with_alias/out.sql
@@ -1,18 +1,18 @@
 SELECT
   "t0"."a",
-  "t0"."b" * CAST(2 AS TINYINT) AS "b2"
+  "t0"."b" * 2 AS "b2"
 FROM "my_table" AS "t0"
 WHERE
-  "t0"."a" < CAST(100 AS TINYINT)
+  "t0"."a" < 100
   AND "t0"."a" = (
     SELECT
       MAX("t1"."a") AS "Max(a)"
     FROM (
       SELECT
         "t0"."a",
-        "t0"."b" * CAST(2 AS TINYINT) AS "b2"
+        "t0"."b" * 2 AS "b2"
       FROM "my_table" AS "t0"
       WHERE
-        "t0"."a" < CAST(100 AS TINYINT)
+        "t0"."a" < 100
     ) AS "t1"
   )
diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_column_expr_default_name/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_column_expr_default_name/out.sql
@@ -1,3 +1,3 @@
 SELECT
-  "t0"."int_col" + CAST(4 AS TINYINT) AS "Add(int_col, 4)"
+  "t0"."int_col" + 4 AS "Add(int_col, 4)"
 FROM "int_col_table" AS "t0"
diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_column_expr_retains_name/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_column_expr_retains_name/out.sql
@@ -1,3 +1,3 @@
 SELECT
-  "t0"."int_col" + CAST(4 AS TINYINT) AS "foo"
+  "t0"."int_col" + 4 AS "foo"
 FROM "int_col_table" AS "t0"
diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_count_distinct/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_count_distinct/out.sql
@@ -6,7 +6,7 @@ FROM (
     *
   FROM "functional_alltypes" AS "t0"
   WHERE
-    "t0"."bigint_col" > CAST(0 AS TINYINT)
+    "t0"."bigint_col" > 0
 ) AS "t1"
 GROUP BY
   1
diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_difference_project_column/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_difference_project_column/out.sql
@@ -9,7 +9,7 @@ FROM (
       CAST("t0"."float_col" AS DOUBLE) AS "value"
     FROM "functional_alltypes" AS "t0"
     WHERE
-      "t0"."int_col" > CAST(0 AS TINYINT)
+      "t0"."int_col" > 0
   ) AS "t1"
   EXCEPT
   SELECT
@@ -20,6 +20,6 @@ FROM (
       "t0"."double_col" AS "value"
     FROM "functional_alltypes" AS "t0"
     WHERE
-      "t0"."int_col" <= CAST(0 AS TINYINT)
+      "t0"."int_col" <= 0
   ) AS "t2"
 ) AS "t3"
diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_having_from_filter/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_having_from_filter/out.sql
@@ -17,4 +17,4 @@ FROM (
     1
 ) AS "t2"
 WHERE
-  "t2"."Max(a)" = CAST(2 AS TINYINT)
+  "t2"."Max(a)" = 2
diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_having_size/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_having_size/out.sql
@@ -11,4 +11,4 @@ FROM (
     1
 ) AS "t1"
 WHERE
-  "t1"."Max(double_col)" = CAST(1 AS TINYINT)
+  "t1"."Max(double_col)" = 1
diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_intersect_project_column/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_intersect_project_column/out.sql
@@ -9,7 +9,7 @@ FROM (
       CAST("t0"."float_col" AS DOUBLE) AS "value"
     FROM "functional_alltypes" AS "t0"
     WHERE
-      "t0"."int_col" > CAST(0 AS TINYINT)
+      "t0"."int_col" > 0
   ) AS "t1"
   INTERSECT
   SELECT
@@ -20,6 +20,6 @@ FROM (
       "t0"."double_col" AS "value"
     FROM "functional_alltypes" AS "t0"
     WHERE
-      "t0"."int_col" <= CAST(0 AS TINYINT)
+      "t0"."int_col" <= 0
   ) AS "t2"
 ) AS "t3"
diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_pushdown_with_or/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_pushdown_with_or/out.sql
@@ -2,15 +2,14 @@ SELECT
   *
 FROM "functional_alltypes" AS "t0"
 WHERE
-  "t0"."double_col" > CAST(3.14 AS DOUBLE)
+  "t0"."double_col" > 3.14
   AND CONTAINS("t0"."string_col", 'foo')
   AND (
     (
       (
-        "t0"."int_col" - CAST(1 AS TINYINT)
-      ) = CAST(0 AS TINYINT)
-    )
-    OR (
-      "t0"."float_col" <= CAST(1.34 AS DOUBLE)
+        "t0"."int_col" - 1
+      ) = 0
+    ) OR (
+      "t0"."float_col" <= 1.34
     )
   )
diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_simple_agg_filter/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_simple_agg_filter/out.sql
@@ -2,7 +2,7 @@ SELECT
   *
 FROM "my_table" AS "t0"
 WHERE
-  "t0"."a" < CAST(100 AS TINYINT)
+  "t0"."a" < 100
   AND "t0"."a" = (
     SELECT
       MAX("t1"."a") AS "Max(a)"
@@ -11,6 +11,6 @@ WHERE
         *
       FROM "my_table" AS "t0"
       WHERE
-        "t0"."a" < CAST(100 AS TINYINT)
+        "t0"."a" < 100
     ) AS "t1"
   )
diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_table_difference/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_table_difference/out.sql
@@ -6,7 +6,7 @@ FROM (
     CAST("t0"."float_col" AS DOUBLE) AS "value"
   FROM "functional_alltypes" AS "t0"
   WHERE
-    "t0"."int_col" > CAST(0 AS TINYINT)
+    "t0"."int_col" > 0
 ) AS "t1"
 EXCEPT
 SELECT
@@ -17,5 +17,5 @@ FROM (
     "t0"."double_col" AS "value"
   FROM "functional_alltypes" AS "t0"
   WHERE
-    "t0"."int_col" <= CAST(0 AS TINYINT)
+    "t0"."int_col" <= 0
 ) AS "t2"
diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_table_drop_with_filter/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_table_drop_with_filter/out.sql
@@ -16,4 +16,4 @@ FROM (
     ON "t4"."b" = "t2"."b"
 ) AS "t5"
 WHERE
-  "t5"."a" < CAST(1.0 AS DOUBLE)
+  "t5"."a" < 1.0
diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_table_intersect/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_table_intersect/out.sql
@@ -6,7 +6,7 @@ FROM (
     CAST("t0"."float_col" AS DOUBLE) AS "value"
   FROM "functional_alltypes" AS "t0"
   WHERE
-    "t0"."int_col" > CAST(0 AS TINYINT)
+    "t0"."int_col" > 0
 ) AS "t1"
 INTERSECT
 SELECT
@@ -17,5 +17,5 @@ FROM (
     "t0"."double_col" AS "value"
   FROM "functional_alltypes" AS "t0"
   WHERE
-    "t0"."int_col" <= CAST(0 AS TINYINT)
+    "t0"."int_col" <= 0
 ) AS "t2"
diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_union/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_union/out.sql
@@ -6,7 +6,7 @@ FROM (
     CAST("t0"."float_col" AS DOUBLE) AS "value"
   FROM "functional_alltypes" AS "t0"
   WHERE
-    "t0"."int_col" > CAST(0 AS TINYINT)
+    "t0"."int_col" > 0
 ) AS "t1"
 UNION
 SELECT
@@ -17,5 +17,5 @@ FROM (
     "t0"."double_col" AS "value"
   FROM "functional_alltypes" AS "t0"
   WHERE
-    "t0"."int_col" <= CAST(0 AS TINYINT)
+    "t0"."int_col" <= 0
 ) AS "t2"
diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_union_project_column/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_union_project_column/out.sql
@@ -9,7 +9,7 @@ FROM (
       CAST("t0"."float_col" AS DOUBLE) AS "value"
     FROM "functional_alltypes" AS "t0"
     WHERE
-      "t0"."int_col" > CAST(0 AS TINYINT)
+      "t0"."int_col" > 0
   ) AS "t1"
   UNION ALL
   SELECT
@@ -20,6 +20,6 @@ FROM (
       "t0"."double_col" AS "value"
     FROM "functional_alltypes" AS "t0"
     WHERE
-      "t0"."int_col" <= CAST(0 AS TINYINT)
+      "t0"."int_col" <= 0
   ) AS "t2"
 ) AS "t3"