refactor(sql-compiler): pull up argmin and argmax into the base compiler

ibis-project · Nov 2, 2024 · ef6634c · ef6634c
1 parent 9b7a377
commit ef6634c
Show file tree

Hide file tree

Showing 9 changed files with 16 additions and 2 deletions.
diff --git a/ibis/backends/sql/compilers/base.py b/ibis/backends/sql/compilers/base.py
@@ -323,6 +323,8 @@ class SQLGlotCompiler(abc.ABC):
         ops.ArrayLength: "array_size",
         ops.ArraySort: "array_sort",
         ops.ArrayStringJoin: "array_to_string",
+        ops.ArgMax: "max_by",
+        ops.ArgMin: "min_by",
         ops.Asin: "asin",
         ops.Atan2: "atan2",
         ops.Atan: "atan",

diff --git a/ibis/backends/sql/compilers/druid.py b/ibis/backends/sql/compilers/druid.py
@@ -33,6 +33,8 @@ class DruidCompiler(SQLGlotCompiler):
         ops.ArraySort,
         ops.ArrayUnion,
         ops.ArrayZip,
+        ops.ArgMax,
+        ops.ArgMin,
         ops.CountDistinctStar,
         ops.Covariance,
         ops.Date,

diff --git a/ibis/backends/sql/compilers/exasol.py b/ibis/backends/sql/compilers/exasol.py
@@ -41,6 +41,8 @@ class ExasolCompiler(SQLGlotCompiler):
         ops.ArrayStringJoin,
         ops.ArrayUnion,
         ops.ArrayZip,
+        ops.ArgMax,
+        ops.ArgMin,
         ops.BitwiseNot,
         ops.CumeDist,
         ops.DateAdd,

diff --git a/ibis/backends/sql/compilers/flink.py b/ibis/backends/sql/compilers/flink.py
@@ -71,6 +71,8 @@ class FlinkCompiler(SQLGlotCompiler):
         ops.ApproxMedian,
         ops.ArrayFlatten,
         ops.ArrayStringJoin,
+        ops.ArgMax,
+        ops.ArgMin,
         ops.Correlation,
         ops.CountDistinctStar,
         ops.Covariance,

diff --git a/ibis/backends/sql/compilers/impala.py b/ibis/backends/sql/compilers/impala.py
@@ -32,6 +32,8 @@ class ImpalaCompiler(SQLGlotCompiler):
     UNSUPPORTED_OPS = (
         ops.ArrayPosition,
         ops.Array,
+        ops.ArgMax,
+        ops.ArgMin,
         ops.Covariance,
         ops.ExtractDayOfYear,
         ops.Levenshtein,

diff --git a/ibis/backends/sql/compilers/mssql.py b/ibis/backends/sql/compilers/mssql.py
@@ -90,6 +90,8 @@ class MSSQLCompiler(SQLGlotCompiler):
         ops.ArrayMap,
         ops.ArraySort,
         ops.ArrayUnion,
+        ops.ArgMax,
+        ops.ArgMin,
         ops.BitAnd,
         ops.BitOr,
         ops.BitXor,

diff --git a/ibis/backends/sql/compilers/mysql.py b/ibis/backends/sql/compilers/mysql.py
@@ -68,6 +68,8 @@ def POS_INF(self):
         ops.Array,
         ops.ArrayFlatten,
         ops.ArrayMap,
+        ops.ArgMax,
+        ops.ArgMin,
         ops.Covariance,
         ops.Levenshtein,
         ops.Median,

diff --git a/ibis/backends/sql/compilers/oracle.py b/ibis/backends/sql/compilers/oracle.py
@@ -55,6 +55,8 @@ class OracleCompiler(SQLGlotCompiler):
         ops.ArrayFlatten,
         ops.ArrayMap,
         ops.ArrayStringJoin,
+        ops.ArgMax,
+        ops.ArgMin,
         ops.MultiQuantile,
         ops.RegexSplit,
         ops.StringSplit,

diff --git a/ibis/backends/sql/compilers/snowflake.py b/ibis/backends/sql/compilers/snowflake.py
@@ -111,8 +111,6 @@ class SnowflakeCompiler(SQLGlotCompiler):
     SIMPLE_OPS = {
         ops.All: "min",
         ops.Any: "max",
-        ops.ArgMax: "max_by",
-        ops.ArgMin: "min_by",
         ops.ArrayDistinct: "array_distinct",
         ops.ArrayFlatten: "array_flatten",
         ops.ArrayIndex: "get",