refactor(flink): port to sqlglot

.github/workflows/ibis-backends.yml

@@ -177,22 +177,28 @@ jobs:
               - oracle
             services:
               - oracle
-          # - name: flink
-          #   title: Flink
-          #   serial: true
-          #   extras:
-          #     - flink
-          #   additional_deps:
-          #     - apache-flink
-          #     - pytest-split
-          #   services:
-          #     - flink
-          # - name: risingwave
-          #   title: Risingwave
-          #   services:
-          #     - risingwave
-          #   extras:
-          #     - risingwave
+          - name: flink
+            title: Flink
+            serial: true
+            extras:
+              - flink
+            additional_deps:
+              - apache-flink
+            services:
+              - flink
+        include:
+          - os: ubuntu-latest
+            python-version: "3.10"
+            backend:
+              name: flink
+              title: Flink
+              serial: true
+              extras:
+                - flink
+              additional_deps:
+                - apache-flink
+              services:
+                - flink
         exclude:
           - os: windows-latest
             backend:
@@ -296,32 +302,29 @@ jobs:
                 - oracle
               services:
                 - oracle
-          # - os: windows-latest
-          #   backend:
-          #     name: flink
-          #     title: Flink
-          #     serial: true
-          #     extras:
-          #       - flink
-          #     services:
-          #       - flink
-          # - python-version: "3.11"
-          #   backend:
-          #     name: flink
-          #     title: Flink
-          #     serial: true
-          #     extras:
-          #       - flink
-          #     services:
-          #       - flink
-          # - os: windows-latest
-          #   backend:
-          #     name: risingwave
-          #     title: Risingwave
-          #     services:
-          #       - risingwave
-          #     extras:
-          #       - risingwave
+          - os: windows-latest
+            backend:
+              name: flink
+              title: Flink
+              serial: true
+              extras:
+                - flink
+              additional_deps:
+                - apache-flink
+              services:
+                - flink
+          - os: ubuntu-latest
+            python-version: "3.11"
+            backend:
+              name: flink
+              title: Flink
+              serial: true
+              extras:
+                - flink
+              additional_deps:
+                - apache-flink
+              services:
+                - flink
           - os: windows-latest
             backend:
               name: exasol
@@ -390,29 +393,18 @@ jobs:
           IBIS_TEST_IMPALA_PORT: 21050
           IBIS_EXAMPLES_DATA: ${{ runner.temp }}/examples-${{ matrix.backend.name }}-${{ matrix.os }}-${{ steps.install_python.outputs.python-version }}
 
-      # FIXME(deepyaman): If some backend-specific test, in test_ddl.py,
-      #   executes before common tests, they will fail with:
-      #   org.apache.flink.table.api.ValidationException: Table `default_catalog`.`default_database`.`functional_alltypes` was not found.
-      #   Therefore, we run backend-specific tests second to avoid this.
-      # - name: "run serial tests: ${{ matrix.backend.name }}"
-      #   if: matrix.backend.serial && matrix.backend.name == 'flink'
-      #   run: |
-      #     just ci-check -m ${{ matrix.backend.name }} ibis/backends/tests
-      #     just ci-check -m ${{ matrix.backend.name }} ibis/backends/flink/tests
-      #   env:
-      #     IBIS_EXAMPLES_DATA: ${{ runner.temp }}/examples-${{ matrix.backend.name }}-${{ matrix.os }}-${{ steps.install_python.outputs.python-version }}
-      #     FLINK_REMOTE_CLUSTER_ADDR: localhost
-      #     FLINK_REMOTE_CLUSTER_PORT: "8081"
-      #
       - name: "run serial tests: ${{ matrix.backend.name }}"
-        if: matrix.backend.serial # && matrix.backend.name != 'flink'
+        if: matrix.backend.serial
         run: just ci-check -m ${{ matrix.backend.name }}
         env:
+          FLINK_REMOTE_CLUSTER_ADDR: localhost
+          FLINK_REMOTE_CLUSTER_PORT: "8081"
           IBIS_EXAMPLES_DATA: ${{ runner.temp }}/examples-${{ matrix.backend.name }}-${{ matrix.os }}-${{ steps.install_python.outputs.python-version }}
 
       - name: check that no untracked files were produced
         shell: bash
-        run: git checkout poetry.lock pyproject.toml && ! git status --porcelain | tee /dev/stderr | grep .
+        run: |
+          ! git status --porcelain | tee /dev/stderr | grep .
 
       - name: upload code coverage
         if: success()

ibis/backends/base/__init__.py

@@ -1277,9 +1277,14 @@
 
 
 @functools.cache
-def _get_backend_names() -> frozenset[str]:
+def _get_backend_names(*, exclude: tuple[str] = ()) -> frozenset[str]:
     """Return the set of known backend names.
 
+    Parameters
+    ----------
+    exclude
+        Exclude these backend names from the result
+
     Notes
     -----
     This function returns a frozenset to prevent cache pollution.
@@ -1293,7 +1298,7 @@
         entrypoints = importlib.metadata.entry_points()["ibis.backends"]
     else:
         entrypoints = importlib.metadata.entry_points(group="ibis.backends")
-    return frozenset(ep.name for ep in entrypoints)
+    return frozenset(ep.name for ep in entrypoints).difference(exclude)
 
 
 def connect(resource: Path | str, **kwargs: Any) -> BaseBackend:
@@ -1431,11 +1436,11 @@
 
         for name, value in query_params.items():
             if len(value) > 1:
                 kwargs[name] = value
             elif len(value) == 1:
                 kwargs[name] = value[0]
             else:
                 raise exc.IbisError(f"Invalid URL parameter: {name}")
 
         self._convert_kwargs(kwargs)
         return self.connect(database=database, **kwargs)

ibis/backends/base/sqlglot/compiler.py

@@ -71,7 +71,7 @@
         )
 
     def __getitem__(self, key: str) -> Callable[..., sge.Anonymous]:
         return getattr(self, key)
 
 
 class FuncGen:
@@ -124,17 +124,17 @@
         return sg.column(name, table=self.table)
 
     def __getitem__(self, key: str) -> sge.Column:
         return sg.column(key, table=self.table)
 
 
 def paren(expr):
     """Wrap a sqlglot expression in parentheses."""
-    return sge.Paren(this=expr)
+    return sge.Paren(this=sge.convert(expr))
 
 
 def parenthesize(op, arg):
     if isinstance(op, (ops.Binary, ops.Unary)):
         return paren(arg)
     # function calls don't need parens
     return arg
 
@@ -297,7 +297,7 @@
 
     @singledispatchmethod
     def visit_node(self, op: ops.Node, **_):
         raise com.OperationNotDefinedError(
             f"No translation rule for {type(op).__name__}"
         )
 
@@ -340,7 +340,7 @@
         if value is None:
             if dtype.nullable:
                 return NULL if dtype.is_null() else self.cast(NULL, dtype)
             raise com.UnsupportedOperationError(
                 f"Unsupported NULL for non-nullable type: {dtype!r}"
             )
         else:
@@ -636,7 +636,7 @@
     @visit_node.register(ops.StringFind)
     def visit_StringFind(self, op, *, arg, substr, start, end):
         if end is not None:
             raise com.UnsupportedOperationError(
                 "String find doesn't support `end` argument"
             )
 
@@ -965,7 +965,7 @@
     def visit_Select(self, op, *, parent, selections, predicates, sort_keys):
         # if we've constructed a useless projection return the parent relation
         if not selections and not predicates and not sort_keys:
             return parent
 
         result = parent
 
@@ -1022,7 +1022,7 @@
 
         for link in rest:
             if isinstance(link, sge.Alias):
                 link = link.this
             result = result.join(link)
         return result
 
@@ -1085,17 +1085,17 @@
             self._add_parens(raw_predicate, predicate)
             for raw_predicate, predicate in zip(op.predicates, predicates)
         )
         try:
             return parent.where(*predicates)
         except AttributeError:
             return sg.select(STAR).from_(parent).where(*predicates)
 
     @visit_node.register(ops.Sort)
     def visit_Sort(self, op, *, parent, keys):
         try:
             return parent.order_by(*keys)
         except AttributeError:
             return sg.select(STAR).from_(parent).order_by(*keys)
 
     @visit_node.register(ops.Union)
     def visit_Union(self, op, *, left, right, distinct):
@@ -1114,10 +1114,10 @@
     @visit_node.register(ops.Intersection)
     def visit_Intersection(self, op, *, left, right, distinct):
         if isinstance(left, sge.Table):
             left = sg.select(STAR).from_(left)
 
         if isinstance(right, sge.Table):
             right = sg.select(STAR).from_(right)
 
         return sg.intersect(
             left.args.get("this", left),
@@ -1128,10 +1128,10 @@
     @visit_node.register(ops.Difference)
     def visit_Difference(self, op, *, left, right, distinct):
         if isinstance(left, sge.Table):
             left = sg.select(STAR).from_(left)
 
         if isinstance(right, sge.Table):
             right = sg.select(STAR).from_(right)
 
         return sg.except_(
             left.args.get("this", left),
@@ -1156,7 +1156,7 @@
         else:
             assert n is None, n
             if self.no_limit_value is not None:
                 result = result.limit(self.no_limit_value)
 
         assert offset is not None, "offset is None"
 

ibis/backends/base/sqlglot/datatypes.py

@@ -254,7 +254,7 @@
                 unit = precision_or_span.this.this
             return dt.Interval(unit=unit, nullable=nullable)
         elif isinstance(precision_or_span, sge.Var):
             return dt.Interval(unit=precision_or_span.this, nullable=nullable)
         elif precision_or_span is None:
             raise com.IbisTypeError("Interval precision is None")
         else:
@@ -334,7 +334,7 @@
 
         if scale is not None:
             if precision is None:
                 raise com.IbisTypeError(
                     "Decimal scale cannot be specified without precision"
                 )
             expressions.append(sge.DataTypeParam(this=sge.Literal.number(scale)))
@@ -390,7 +390,7 @@
     @classmethod
     def _from_ibis_Map(cls, dtype: dt.Map) -> sge.DataType:
         if not dtype.key_type.is_string():
             raise com.IbisTypeError("Postgres only supports string keys in maps")
         if not dtype.value_type.is_string():
             raise com.IbisTypeError("Postgres only supports string values in maps")
         return sge.DataType(this=typecode.HSTORE)
@@ -511,9 +511,9 @@
 
     @classmethod
     def _from_ibis_Interval(cls, dtype: dt.Interval) -> sge.DataType:
         assert dtype.unit is not None, "interval unit cannot be None"
         if (short := dtype.unit.short) in ("Y", "Q", "M"):
             return sge.DataType(
                 this=typecode.INTERVAL,
                 expressions=[
                     sge.IntervalSpan(
@@ -521,8 +521,8 @@
                     )
                 ],
             )
         elif short in ("D", "h", "m", "s", "ms", "us", "ns"):
             return sge.DataType(
                 this=typecode.INTERVAL,
                 expressions=[
                     sge.IntervalSpan(
@@ -537,7 +537,7 @@
 
     @classmethod
     def _from_sqlglot_UBIGINT(cls):
         return dt.Decimal(precision=19, scale=0, nullable=cls.default_nullable)
 
     @classmethod
     def _from_ibis_UInt64(cls, dtype):
@@ -551,7 +551,7 @@
 
     @classmethod
     def _from_sqlglot_UINT(cls):
         return dt.Int64(nullable=cls.default_nullable)
 
     @classmethod
     def _from_ibis_UInt32(cls, dtype):
@@ -559,7 +559,7 @@
 
     @classmethod
     def _from_sqlglot_USMALLINT(cls):
         return dt.Int32(nullable=cls.default_nullable)
 
     @classmethod
     def _from_ibis_UInt16(cls, dtype):
@@ -567,7 +567,7 @@
 
     @classmethod
     def _from_sqlglot_UTINYINT(cls):
         return dt.Int16(nullable=cls.default_nullable)
 
     @classmethod
     def _from_ibis_UInt8(cls, dtype):
@@ -599,7 +599,7 @@
         if scale is None or int(scale.this.this) == 0:
             return dt.Int64(nullable=cls.default_nullable)
         else:
             return super()._from_sqlglot_DECIMAL(precision, scale)
 
     @classmethod
     def _from_ibis_String(cls, dtype: dt.String) -> sge.DataType:
@@ -633,19 +633,19 @@
 
     @classmethod
     def _from_ibis_JSON(cls, dtype: dt.JSON) -> sge.DataType:
         return sge.DataType(this=sge.DataType.Type.VARIANT)
 
     @classmethod
     def _from_ibis_Array(cls, dtype: dt.Array) -> sge.DataType:
         return sge.DataType(this=sge.DataType.Type.ARRAY, nested=True)
 
     @classmethod
     def _from_ibis_Map(cls, dtype: dt.Map) -> sge.DataType:
         return sge.DataType(this=sge.DataType.Type.OBJECT, nested=True)
 
     @classmethod
     def _from_ibis_Struct(cls, dtype: dt.Struct) -> sge.DataType:
         return sge.DataType(this=sge.DataType.Type.OBJECT, nested=True)
 
 
 class SQLiteType(SqlglotType):
@@ -684,7 +684,7 @@
 
     @classmethod
     def _from_ibis_Map(cls, dtype: dt.Map) -> NoReturn:
         raise com.UnsupportedBackendType("Map types aren't supported in Impala")
 
     @classmethod
     def _from_ibis_Struct(cls, dtype: dt.Struct) -> sge.DataType:
@@ -706,7 +706,7 @@
 
     @classmethod
     def _from_sqlglot_NUMERIC(cls) -> dt.Decimal:
         return dt.Decimal(
             cls.default_decimal_precision,
             cls.default_decimal_scale,
             nullable=cls.default_nullable,
@@ -714,25 +714,25 @@
 
     @classmethod
     def _from_sqlglot_BIGNUMERIC(cls) -> dt.Decimal:
         return dt.Decimal(76, 38, nullable=cls.default_nullable)
 
     @classmethod
     def _from_sqlglot_DATETIME(cls) -> dt.Timestamp:
         return dt.Timestamp(timezone=None, nullable=cls.default_nullable)
 
     @classmethod
     def _from_sqlglot_TIMESTAMP(cls) -> dt.Timestamp:
         return dt.Timestamp(timezone="UTC", nullable=cls.default_nullable)
 
     @classmethod
     def _from_sqlglot_GEOGRAPHY(cls) -> dt.GeoSpatial:
         return dt.GeoSpatial(
             geotype="geography", srid=4326, nullable=cls.default_nullable
         )
 
     @classmethod
     def _from_sqlglot_TINYINT(cls) -> dt.Int64:
         return dt.Int64(nullable=cls.default_nullable)
 
     _from_sqlglot_UINT = (
         _from_sqlglot_USMALLINT
@@ -742,43 +742,43 @@
 
     @classmethod
     def _from_sqlglot_UBIGINT(cls) -> NoReturn:
         raise com.UnsupportedBackendType(
             "Unsigned BIGINT isn't representable in BigQuery INT64"
         )
 
     @classmethod
     def _from_sqlglot_FLOAT(cls) -> dt.Float64:
         return dt.Float64(nullable=cls.default_nullable)
 
     @classmethod
     def _from_sqlglot_MAP(cls) -> NoReturn:
         raise com.UnsupportedBackendType("Maps are not supported in BigQuery")
 
     @classmethod
     def _from_ibis_Map(cls, dtype: dt.Map) -> NoReturn:
         raise com.UnsupportedBackendType("Maps are not supported in BigQuery")
 
     @classmethod
     def _from_ibis_Timestamp(cls, dtype: dt.Timestamp) -> sge.DataType:
         if dtype.timezone is None:
             return sge.DataType(this=sge.DataType.Type.DATETIME)
         elif dtype.timezone == "UTC":
             return sge.DataType(this=sge.DataType.Type.TIMESTAMPTZ)
         else:
             raise com.UnsupportedBackendType(
                 "BigQuery does not support timestamps with timezones other than 'UTC'"
             )
 
     @classmethod
     def _from_ibis_Decimal(cls, dtype: dt.Decimal) -> sge.DataType:
         precision = dtype.precision
         scale = dtype.scale
         if (precision, scale) == (76, 38):
             return sge.DataType(this=sge.DataType.Type.BIGDECIMAL)
         elif (precision, scale) in ((38, 9), (None, None)):
             return sge.DataType(this=sge.DataType.Type.DECIMAL)
         else:
             raise com.UnsupportedBackendType(
                 "BigQuery only supports decimal types with precision of 38 and "
                 f"scale of 9 (NUMERIC) or precision of 76 and scale of 38 (BIGNUMERIC). "
                 f"Current precision: {dtype.precision}. Current scale: {dtype.scale}"
@@ -786,22 +786,22 @@
 
     @classmethod
     def _from_ibis_UInt64(cls, dtype: dt.UInt64) -> NoReturn:
         raise com.UnsupportedBackendType(
             f"Conversion from {dtype} to BigQuery integer type (Int64) is lossy"
         )
 
     @classmethod
     def _from_ibis_UInt32(cls, dtype: dt.UInt32) -> sge.DataType:
         return sge.DataType(this=sge.DataType.Type.BIGINT)
 
     _from_ibis_UInt8 = _from_ibis_UInt16 = _from_ibis_UInt32
 
     @classmethod
     def _from_ibis_GeoSpatial(cls, dtype: dt.GeoSpatial) -> sge.DataType:
         if (dtype.geotype, dtype.srid) == ("geography", 4326):
             return sge.DataType(this=sge.DataType.Type.GEOGRAPHY)
         else:
             raise com.UnsupportedBackendType(
                 "BigQuery geography uses points on WGS84 reference ellipsoid."
                 f"Current geotype: {dtype.geotype}, Current srid: {dtype.srid}"
             )
@@ -810,7 +810,7 @@
 class BigQueryUDFType(BigQueryType):
     @classmethod
     def _from_ibis_Int64(cls, dtype: dt.Int64) -> NoReturn:
         raise com.UnsupportedBackendType(
             "int64 is not a supported input or output type in BigQuery UDFs; use float64 instead"
         )
 
@@ -837,18 +837,18 @@
         scale: sge.DataTypeParam | None = None,
     ) -> dt.Decimal:
         if precision is None:
             precision = cls.default_decimal_precision
         else:
             precision = int(precision.this.this)
 
         if scale is None:
             scale = cls.default_decimal_scale
         else:
             scale = int(scale.this.this)
 
         if not scale:
             if 0 < precision <= 3:
                 return dt.Int8(nullable=cls.default_nullable)
             elif 3 < precision <= 9:
                 return dt.Int16(nullable=cls.default_nullable)
             elif 9 < precision <= 18:
@@ -856,10 +856,10 @@
             elif 18 < precision <= 36:
                 return dt.Int64(nullable=cls.default_nullable)
             else:
                 raise com.UnsupportedBackendType(
                     "Decimal precision is too large; Exasol supports precision up to 36."
                 )
         return dt.Decimal(precision, scale, nullable=cls.default_nullable)
 
     @classmethod
     def _from_ibis_Array(cls, dtype: dt.Array) -> NoReturn:
@@ -880,15 +880,15 @@
 
     @classmethod
     def _from_sqlglot_ARRAY(cls, value_type: sge.DataType) -> NoReturn:
         raise com.UnsupportedBackendType("Arrays not supported in Exasol")
 
     @classmethod
     def _from_sqlglot_MAP(cls, key: sge.DataType, value: sge.DataType) -> NoReturn:
         raise com.UnsupportedBackendType("Maps not supported in Exasol")
 
     @classmethod
     def _from_sqlglot_STRUCT(cls, *cols: sge.ColumnDef) -> NoReturn:
         raise com.UnsupportedBackendType("Structs not supported in Exasol")
 
 
 class MSSQLType(SqlglotType):
@@ -933,15 +933,15 @@
 
     @classmethod
     def _from_sqlglot_ARRAY(cls) -> sge.DataType:
         raise com.UnsupportedBackendType("SQL Server does not support arrays")
 
     @classmethod
     def _from_sqlglot_MAP(cls) -> sge.DataType:
         raise com.UnsupportedBackendType("SQL Server does not support map")
 
     @classmethod
     def _from_sqlglot_STRUCT(cls) -> sge.DataType:
         raise com.UnsupportedBackendType("SQL Server does not support structs")
 
 
 class ClickHouseType(SqlglotType):
@@ -1029,3 +1029,13 @@
         key_type = cls.from_ibis(dtype.key_type.copy(nullable=False))
         value_type = cls.from_ibis(dtype.value_type)
         return sge.DataType(this=typecode.MAP, expressions=[key_type, value_type])
+
+
+class FlinkType(SqlglotType):
+    dialect = "flink"
+    default_decimal_precision = 38
+    default_decimal_scale = 18
+
+    @classmethod
+    def _from_ibis_Binary(cls, dtype: dt.Binary) -> sge.DataType:
+        return sge.DataType(this=sge.DataType.Type.VARBINARY)

ibis/backends/base/sqlglot/dialects.py

@@ -70,6 +70,10 @@
 
 class Flink(Hive):
     class Generator(Hive.Generator):
+        TYPE_MAPPING = Hive.Generator.TYPE_MAPPING.copy() | {
+            sge.DataType.Type.TIME: "TIME",
+        }
+
         TRANSFORMS = Hive.Generator.TRANSFORMS.copy() | {
             sge.Stddev: rename_func("stddev_samp"),
             sge.StddevPop: rename_func("stddev_pop"),
@@ -82,8 +86,21 @@
             ),
             sge.ArrayConcat: rename_func("array_concat"),
             sge.Length: rename_func("char_length"),
+            sge.TryCast: lambda self,
+            e: f"TRY_CAST({e.this.sql(self.dialect)} AS {e.to.sql(self.dialect)})",
+            sge.DayOfYear: rename_func("dayofyear"),
+            sge.DayOfWeek: rename_func("dayofweek"),
+            sge.DayOfMonth: rename_func("dayofmonth"),
         }
 
+    class Tokenizer(Hive.Tokenizer):
+        # In Flink, embedded single quotes are escaped like most other SQL
+        # dialects: doubling up the single quote
+        #
+        # We override it here because we inherit from Hive's dialect and Hive
+        # uses a backslash to escape single quotes
+        STRING_ESCAPES = ["'"]
+
 
 class Impala(Hive):
     class Generator(Hive.Generator):
@@ -137,7 +154,7 @@
     kind = expression.args["kind"]
     if (obj := kind.upper()) in ("TABLE", "VIEW") and temporary:
         if expression.expression:
             return f"CREATE GLOBAL TEMPORARY {obj} {self.sql(expression, 'this')} AS {self.sql(expression, 'expression')}"
         else:
             # TODO: why does autocommit not work here?  need to specify the ON COMMIT part...
             return f"CREATE GLOBAL TEMPORARY {obj} {self.sql(expression, 'this')} ON COMMIT PRESERVE ROWS"