feat: Use data types defined in the Rust Wren Engine to do type mapping (#913)

Author: onlyjackfrost

ibis-server/app/model/metadata/bigquery.py

@@ -4,6 +4,7 @@
     Column,
     Constraint,
     ConstraintType,
+    RustWrenEngineColumnType,
     Table,
     TableProperties,
 )
@@ -17,6 +18,8 @@ def __init__(self, connection_info: BigQueryConnectionInfo):
 
     def get_table_list(self) -> list[Table]:
         dataset_id = self.connection_info.dataset_id.get_secret_value()
+
+        # filter out columns with GEOGRAPHY & RANGE types
         sql = f"""
             SELECT 
                 c.table_catalog,
@@ -46,14 +49,11 @@ def get_table_list(self) -> list[Table]:
                 AND cf.column_name = c.column_name
             LEFT JOIN {dataset_id}.INFORMATION_SCHEMA.TABLE_OPTIONS table_options
                 ON c.table_name = table_options.table_name
+            WHERE cf.data_type != 'GEOGRAPHY'
+                AND cf.data_type NOT LIKE 'RANGE%'
             """
         response = self.connection.sql(sql).to_pandas().to_dict(orient="records")
 
-        def get_data_type(data_type) -> str:
-            if "STRUCT" in data_type:
-                return "RECORD"
-            return data_type
-
         def get_column(row, nestedColumns=None) -> Column:
             return Column(
                 # field_path supports both column & nested column
@@ -139,3 +139,30 @@ def get_constraints(self) -> list[Constraint]:
 
     def get_version(self) -> str:
         return "Follow BigQuery release version"
+
+    def _transform_column_type(self, data_type):
+        # lower case the data_type
+        data_type = data_type.lower()
+
+        # if data_type start with "array" or "struct", by pass it
+        if data_type.startswith(("array", "struct")):
+            return data_type
+
+        # Map BigQuery types to RustWrenEngineColumnType
+        switcher = {
+            # GEOGRAPHY and RANGE columns were filtered out
+            "bytes": RustWrenEngineColumnType.BYTES,
+            "date": RustWrenEngineColumnType.DATE,
+            "datetime": RustWrenEngineColumnType.DATETIME,
+            "interval": RustWrenEngineColumnType.INTERVAL,
+            "json": RustWrenEngineColumnType.JSON,
+            "int64": RustWrenEngineColumnType.INT64,
+            "numeric": RustWrenEngineColumnType.NUMERIC,
+            "bignumeric": RustWrenEngineColumnType.BIGNUMERIC,
+            "float64": RustWrenEngineColumnType.FLOAT64,
+            "string": RustWrenEngineColumnType.STRING,
+            "time": RustWrenEngineColumnType.TIME,
+            "timestamp": RustWrenEngineColumnType.TIMESTAMPTZ,
+        }
+
+        return switcher.get(data_type, RustWrenEngineColumnType.UNKNOWN)

ibis-server/app/model/metadata/canner.py

@@ -1,3 +1,4 @@
+import re
 from urllib.parse import urlparse
 
 from gql import Client, gql
@@ -7,6 +8,7 @@
 from app.model.metadata.dto import (
     Column,
     Constraint,
+    RustWrenEngineColumnType,
     Table,
     TableProperties,
 )
@@ -181,7 +183,7 @@ def _build_columns(cls, columns: list[dict]) -> list[Column]:
         return [
             Column(
                 name=column["originalColumn"]["name"],
-                type=column["originalColumn"]["type"],
+                type=cls._transform_column_type(column["originalColumn"]["type"]),
                 notNull=column["originalColumn"]["properties"].get(
                     "jdbc-nullable", False
                 ),
@@ -190,3 +192,45 @@ def _build_columns(cls, columns: list[dict]) -> list[Column]:
             )
             for column in columns
         ]
+
+    @classmethod
+    def _transform_column_type(self, data_type):
+        # all possible types listed here: https://trino.io/docs/current/language/types.html
+        # trim the (all characters) at the end of the data_type if exists
+        data_type = re.sub(r"\(.*\)", "", data_type).strip()
+
+        switcher = {
+            # String Types (ignore Binary and Spatial Types for now)
+            "char": RustWrenEngineColumnType.CHAR,
+            "varchar": RustWrenEngineColumnType.VARCHAR,
+            "tinytext": RustWrenEngineColumnType.TEXT,
+            "text": RustWrenEngineColumnType.TEXT,
+            "mediumtext": RustWrenEngineColumnType.TEXT,
+            "longtext": RustWrenEngineColumnType.TEXT,
+            "enum": RustWrenEngineColumnType.VARCHAR,
+            "set": RustWrenEngineColumnType.VARCHAR,
+            # Numeric Types(https://dev.mysql.com/doc/refman/8.4/en/numeric-types.html)
+            "bit": RustWrenEngineColumnType.TINYINT,
+            "tinyint": RustWrenEngineColumnType.TINYINT,
+            "smallint": RustWrenEngineColumnType.SMALLINT,
+            "mediumint": RustWrenEngineColumnType.INTEGER,
+            "int": RustWrenEngineColumnType.INTEGER,
+            "integer": RustWrenEngineColumnType.INTEGER,
+            "bigint": RustWrenEngineColumnType.BIGINT,
+            # boolean
+            "bool": RustWrenEngineColumnType.BOOL,
+            "boolean": RustWrenEngineColumnType.BOOL,
+            # Decimal
+            "float": RustWrenEngineColumnType.FLOAT8,
+            "double": RustWrenEngineColumnType.DOUBLE,
+            "decimal": RustWrenEngineColumnType.DECIMAL,
+            "numeric": RustWrenEngineColumnType.NUMERIC,
+            # Date and Time Types(https://dev.mysql.com/doc/refman/8.4/en/date-and-time-types.html)
+            "date": RustWrenEngineColumnType.DATE,
+            "datetime": RustWrenEngineColumnType.TIMESTAMP,
+            "timestamp": RustWrenEngineColumnType.TIMESTAMPTZ,
+            # JSON Type
+            "json": RustWrenEngineColumnType.JSON,
+        }
+
+        return switcher.get(data_type.lower(), RustWrenEngineColumnType.UNKNOWN)

ibis-server/app/model/metadata/clickhouse.py

@@ -3,9 +3,9 @@
 from app.model.metadata.dto import (
     Column,
     Constraint,
+    RustWrenEngineColumnType,
     Table,
     TableProperties,
-    WrenEngineColumnType,
 )
 from app.model.metadata.metadata import Metadata
 
@@ -80,29 +80,29 @@ def _transform_column_type(self, data_type):
         # lower case the data_type
         data_type = data_type.lower()
 
-        # Map ClickHouse types to WrenEngineColumnType
+        # Map ClickHouse types to RustWrenEngineColumnType
         switcher = {
-            "boolean": WrenEngineColumnType.BOOLEAN,
-            "int8": WrenEngineColumnType.TINYINT,
-            "uint8": WrenEngineColumnType.INT2,
-            "int16": WrenEngineColumnType.INT2,
-            "uint16": WrenEngineColumnType.INT2,
-            "int32": WrenEngineColumnType.INT4,
-            "uint32": WrenEngineColumnType.INT4,
-            "int64": WrenEngineColumnType.INT8,
-            "uint64": WrenEngineColumnType.INT8,
-            "float32": WrenEngineColumnType.FLOAT4,
-            "float64": WrenEngineColumnType.FLOAT8,
-            "decimal": WrenEngineColumnType.DECIMAL,
-            "date": WrenEngineColumnType.DATE,
-            "datetime": WrenEngineColumnType.TIMESTAMP,
-            "string": WrenEngineColumnType.VARCHAR,
-            "fixedstring": WrenEngineColumnType.CHAR,
-            "uuid": WrenEngineColumnType.UUID,
-            "enum8": WrenEngineColumnType.STRING,  # Enums can be mapped to strings
-            "enum16": WrenEngineColumnType.STRING,  # Enums can be mapped to strings
-            "ipv4": WrenEngineColumnType.INET,
-            "ipv6": WrenEngineColumnType.INET,
+            "boolean": RustWrenEngineColumnType.BOOL,
+            "int8": RustWrenEngineColumnType.TINYINT,
+            "uint8": RustWrenEngineColumnType.INT2,
+            "int16": RustWrenEngineColumnType.INT2,
+            "uint16": RustWrenEngineColumnType.INT2,
+            "int32": RustWrenEngineColumnType.INT4,
+            "uint32": RustWrenEngineColumnType.INT4,
+            "int64": RustWrenEngineColumnType.INT8,
+            "uint64": RustWrenEngineColumnType.INT8,
+            "float32": RustWrenEngineColumnType.FLOAT4,
+            "float64": RustWrenEngineColumnType.FLOAT8,
+            "decimal": RustWrenEngineColumnType.DECIMAL,
+            "date": RustWrenEngineColumnType.DATE,
+            "datetime": RustWrenEngineColumnType.TIMESTAMP,
+            "string": RustWrenEngineColumnType.VARCHAR,
+            "fixedstring": RustWrenEngineColumnType.CHAR,
+            "uuid": RustWrenEngineColumnType.UUID,
+            "enum8": RustWrenEngineColumnType.STRING,  # Enums can be mapped to strings
+            "enum16": RustWrenEngineColumnType.STRING,  # Enums can be mapped to strings
+            "ipv4": RustWrenEngineColumnType.INET,
+            "ipv6": RustWrenEngineColumnType.INET,
         }
 
-        return switcher.get(data_type, WrenEngineColumnType.UNKNOWN)
+        return switcher.get(data_type, RustWrenEngineColumnType.UNKNOWN)

ibis-server/app/model/metadata/dto.py

@@ -10,59 +10,46 @@ class MetadataDTO(BaseModel):
     connection_info: ConnectionInfo = Field(alias="connectionInfo")
 
 
-class WrenEngineColumnType(Enum):
-    # Boolean Types
-    BOOLEAN = "BOOLEAN"
-
-    # Numeric Types
+class RustWrenEngineColumnType(Enum):
+    BOOL = "BOOL"
     TINYINT = "TINYINT"
     INT2 = "INT2"
-    SMALLINT = "SMALLINT"  # alias for INT2
+    SMALLINT = "SMALLINT"
     INT4 = "INT4"
-    INTEGER = "INTEGER"  # alias for INT4
+    INT = "INT"
+    INTEGER = "INTEGER"
     INT8 = "INT8"
-    BIGINT = "BIGINT"  # alias for INT8
+    BIGINT = "BIGINT"
     NUMERIC = "NUMERIC"
     DECIMAL = "DECIMAL"
-
-    # Floating-Point Types
-    FLOAT4 = "FLOAT4"
-    REAL = "REAL"  # alias for FLOAT4
-    FLOAT8 = "FLOAT8"
-    DOUBLE = "DOUBLE"  # alias for FLOAT8
-
-    # Character Types
     VARCHAR = "VARCHAR"
     CHAR = "CHAR"
-    BPCHAR = "BPCHAR"  # BPCHAR is fixed-length blank padded string
-    TEXT = "TEXT"  # alias for VARCHAR
-    STRING = "STRING"  # alias for VARCHAR
-    NAME = "NAME"  # alias for VARCHAR
-
-    # Date/Time Types
+    BPCHAR = "BPCHAR"
+    TEXT = "TEXT"
+    STRING = "STRING"
+    NAME = "NAME"
+    FLOAT4 = "FLOAT4"
+    REAL = "REAL"
+    FLOAT = "FLOAT"
+    FLOAT8 = "FLOAT8"
+    DOUBLE = "DOUBLE"
     TIMESTAMP = "TIMESTAMP"
-    TIMESTAMPTZ = "TIMESTAMP WITH TIME ZONE"
+    TIMESTAMPTZ = "TIMESTAMPTZ"
     DATE = "DATE"
     INTERVAL = "INTERVAL"
-
-    # JSON Types
     JSON = "JSON"
-
-    # Object identifiers (OIDs) are used internally by PostgreSQL as primary keys for various system tables.
-    # https:#www.postgresql.org/docs/current/datatype-oid.html
     OID = "OID"
-
-    # Binary Data Types
     BYTEA = "BYTEA"
-
-    # UUID Type
     UUID = "UUID"
-
-    # Network Address Types
     INET = "INET"
-
-    # Unknown Type
     UNKNOWN = "UNKNOWN"
+    BIGNUMERIC = "BIGNUMERIC"
+    BYTES = "BYTES"
+    DATETIME = "DATETIME"
+    FLOAT64 = "FLOAT64"
+    INT64 = "INT64"
+    TIME = "TIME"
+    NULL = "NULL"
 
 
 class Column(BaseModel):

ibis-server/app/model/metadata/mssql.py

@@ -4,9 +4,9 @@
     Column,
     Constraint,
     ConstraintType,
+    RustWrenEngineColumnType,
     Table,
     TableProperties,
-    WrenEngineColumnType,
 )
 from app.model.metadata.metadata import Metadata
 
@@ -173,40 +173,40 @@ def _format_constraint_name(
         return f"{table_name}_{column_name}_{referenced_table_name}_{referenced_column_name}"
 
     def _transform_column_type(self, data_type):
-        # Define the mapping of MSSQL data types to WrenEngineColumnType
+        # Define the mapping of MSSQL data types to RustWrenEngineColumnType
         # ref: https://learn.microsoft.com/en-us/sql/t-sql/data-types/data-types-transact-sql?view=sql-server-ver15#exact-numerics
         switcher = {
             # String Types
-            "char": WrenEngineColumnType.CHAR,
-            "varchar": WrenEngineColumnType.VARCHAR,
-            "text": WrenEngineColumnType.TEXT,
-            "nchar": WrenEngineColumnType.CHAR,
-            "nvarchar": WrenEngineColumnType.VARCHAR,
-            "ntext": WrenEngineColumnType.TEXT,
+            "char": RustWrenEngineColumnType.CHAR,
+            "varchar": RustWrenEngineColumnType.VARCHAR,
+            "text": RustWrenEngineColumnType.TEXT,
+            "nchar": RustWrenEngineColumnType.CHAR,
+            "nvarchar": RustWrenEngineColumnType.VARCHAR,
+            "ntext": RustWrenEngineColumnType.TEXT,
             # Numeric Types
-            "bit": WrenEngineColumnType.TINYINT,
-            "tinyint": WrenEngineColumnType.TINYINT,
-            "smallint": WrenEngineColumnType.SMALLINT,
-            "int": WrenEngineColumnType.INTEGER,
-            "bigint": WrenEngineColumnType.BIGINT,
+            "bit": RustWrenEngineColumnType.TINYINT,
+            "tinyint": RustWrenEngineColumnType.TINYINT,
+            "smallint": RustWrenEngineColumnType.SMALLINT,
+            "int": RustWrenEngineColumnType.INTEGER,
+            "bigint": RustWrenEngineColumnType.BIGINT,
             # Boolean
-            "boolean": WrenEngineColumnType.BOOLEAN,
+            "boolean": RustWrenEngineColumnType.BOOL,
             # Decimal
-            "float": WrenEngineColumnType.FLOAT8,
-            "real": WrenEngineColumnType.FLOAT8,
-            "decimal": WrenEngineColumnType.DECIMAL,
-            "numeric": WrenEngineColumnType.NUMERIC,
-            "money": WrenEngineColumnType.DECIMAL,
-            "smallmoney": WrenEngineColumnType.DECIMAL,
+            "float": RustWrenEngineColumnType.FLOAT8,
+            "real": RustWrenEngineColumnType.FLOAT8,
+            "decimal": RustWrenEngineColumnType.DECIMAL,
+            "numeric": RustWrenEngineColumnType.NUMERIC,
+            "money": RustWrenEngineColumnType.DECIMAL,
+            "smallmoney": RustWrenEngineColumnType.DECIMAL,
             # Date and Time Types
-            "date": WrenEngineColumnType.DATE,
-            "datetime": WrenEngineColumnType.TIMESTAMP,
-            "datetime2": WrenEngineColumnType.TIMESTAMPTZ,
-            "smalldatetime": WrenEngineColumnType.TIMESTAMP,
-            "time": WrenEngineColumnType.INTERVAL,
-            "datetimeoffset": WrenEngineColumnType.TIMESTAMPTZ,
+            "date": RustWrenEngineColumnType.DATE,
+            "datetime": RustWrenEngineColumnType.TIMESTAMP,
+            "datetime2": RustWrenEngineColumnType.TIMESTAMP,
+            "smalldatetime": RustWrenEngineColumnType.TIMESTAMP,
+            "time": RustWrenEngineColumnType.INTERVAL,
+            "datetimeoffset": RustWrenEngineColumnType.TIMESTAMPTZ,
             # JSON Type (Note: MSSQL supports JSON natively as a string type)
-            "json": WrenEngineColumnType.JSON,
+            "json": RustWrenEngineColumnType.JSON,
         }
 
-        return switcher.get(data_type.lower(), WrenEngineColumnType.UNKNOWN)
+        return switcher.get(data_type.lower(), RustWrenEngineColumnType.UNKNOWN)