Added support for Variant datatype in SQLAlchemy

Author: msrathore-db

src/databricks/sqlalchemy/__init__.py

@@ -5,6 +5,7 @@
     TIMESTAMP_NTZ,
     DatabricksArray,
     DatabricksMap,
+    DatabricksVariant,
 )
 
-__all__ = ["TINYINT", "TIMESTAMP", "TIMESTAMP_NTZ", "DatabricksArray", "DatabricksMap"]
+__all__ = ["TINYINT", "TIMESTAMP", "TIMESTAMP_NTZ", "DatabricksArray", "DatabricksMap", "DatabricksVariant"]

src/databricks/sqlalchemy/_parse.py

@@ -318,6 +318,7 @@ def get_comment_from_dte_output(dte_output: List[Dict[str, str]]) -> Optional[st
     "map": sqlalchemy.types.String,
     "struct": sqlalchemy.types.String,
     "uniontype": sqlalchemy.types.String,
+    "variant": type_overrides.DatabricksVariant,
     "decimal": sqlalchemy.types.Numeric,
     "timestamp": type_overrides.TIMESTAMP,
     "timestamp_ntz": type_overrides.TIMESTAMP_NTZ,

src/databricks/sqlalchemy/_types.py

@@ -9,6 +9,7 @@
 
 from databricks.sql.utils import ParamEscaper
 
+from sqlalchemy.sql import expression
 
 def process_literal_param_hack(value: Any):
     """This method is supposed to accept a Python type and return a string representation of that type.
@@ -397,3 +398,47 @@ def compile_databricks_map(type_, compiler, **kw):
     key_type = compiler.process(type_.key_type, **kw)
     value_type = compiler.process(type_.value_type, **kw)
     return f"MAP<{key_type},{value_type}>"
+
+class DatabricksVariant(UserDefinedType):
+    """
+    A custom variant type for storing semi-structured data including STRUCT, ARRAY, MAP, and scalar types.
+    Note: VARIANT MAP types can only have STRING keys.
+    
+    Examples:
+        DatabricksVariant()  -> VARIANT
+        
+    Usage:
+        Column('data', DatabricksVariant())
+    """
+    cache_ok = True
+
+    def __init__(self):
+        self.pe = ParamEscaper()
+
+    def bind_processor(self, dialect):
+        """Process values before sending to database.
+        """
+
+        def process(value):
+            return value
+
+        return process
+
+    def bind_expression(self, bindvalue):
+        """Wrap with PARSE_JSON() in SQL"""
+        return expression.func.PARSE_JSON(bindvalue)
+
+    def literal_processor(self, dialect):
+        """Process literal values for SQL generation.      
+        For VARIANT columns, use PARSE_JSON() to properly insert data.
+        """
+        def process(value):
+            if value is None:
+                return "NULL"
+            return self.pe.escape_string(value)
+            
+        return f"PARSE_JSON('{process}')"
+
+@compiles(DatabricksVariant, "databricks")
+def compile_variant(type_, compiler, **kw):
+    return "VARIANT"

tests/test_local/e2e/test_complex_types.py

@@ -11,13 +11,14 @@
     DateTime,
 )
 from collections.abc import Sequence
-from databricks.sqlalchemy import TIMESTAMP, TINYINT, DatabricksArray, DatabricksMap
+from databricks.sqlalchemy import TIMESTAMP, TINYINT, DatabricksArray, DatabricksMap, DatabricksVariant
 from sqlalchemy.orm import DeclarativeBase, Session
 from sqlalchemy import select
 from datetime import date, datetime, time, timedelta, timezone
 import pandas as pd
 import numpy as np
 import decimal
+import json
 
 
 class TestComplexTypes(TestSetup):
@@ -46,7 +47,7 @@ def _parse_to_common_type(self, value):
         ):
             return tuple(value)
         elif isinstance(value, dict):
-            return tuple(value.items())
+            return tuple(sorted(value.items()))
         elif isinstance(value, np.generic):
             return value.item()
         elif isinstance(value, decimal.Decimal):
@@ -152,6 +153,35 @@ class MapTable(Base):
 
         return MapTable, sample_data
 
+    def sample_variant_table(self) -> tuple[DeclarativeBase, dict]:
+        class Base(DeclarativeBase):
+            pass
+
+        class VariantTable(Base):
+            __tablename__ = "sqlalchemy_variant_table"
+
+            int_col = Column(Integer, primary_key=True)
+            variant_simple_col = Column(DatabricksVariant())
+            variant_nested_col = Column(DatabricksVariant())
+            variant_array_col = Column(DatabricksVariant())
+            variant_mixed_col = Column(DatabricksVariant())
+
+        sample_data = {
+            "int_col": 1,
+            "variant_simple_col": {"key": "value", "number": 42},
+            "variant_nested_col": {"user": {"name": "John", "age": 30}, "active": True},
+            "variant_array_col": [1, 2, 3, "hello", {"nested": "data"}],
+            "variant_mixed_col": {
+                "string": "test",
+                "number": 123,
+                "boolean": True,
+                "array": [1, 2, 3],
+                "object": {"nested": "value"}
+            }
+        }
+
+        return VariantTable, sample_data
+
     def test_insert_array_table_sqlalchemy(self):
         table, sample_data = self.sample_array_table()
 
@@ -209,3 +239,57 @@ def test_map_table_creation_pandas(self):
             stmt = select(table)
             df_result = pd.read_sql(stmt, engine)
             assert self._recursive_compare(df_result.iloc[0].to_dict(), sample_data)
+
+    def test_insert_variant_table_sqlalchemy(self):
+        table, sample_data = self.sample_variant_table()
+
+        with self.table_context(table) as engine:
+            # Pre-serialize variant data for SQLAlchemy
+            variant_data = sample_data.copy()
+            for key in ['variant_simple_col', 'variant_nested_col', 'variant_array_col', 'variant_mixed_col']:
+                variant_data[key] = None if sample_data[key] is None else json.dumps(sample_data[key])
+            
+            sa_obj = table(**variant_data)
+            session = Session(engine)
+            session.add(sa_obj)
+            session.commit()
+
+            stmt = select(table).where(table.int_col == 1)
+
+            result = session.scalar(stmt)
+
+            compare = {key: getattr(result, key) for key in sample_data.keys()}
+            # Parse JSON values back to original format for comparison
+            for key in ['variant_simple_col', 'variant_nested_col', 'variant_array_col', 'variant_mixed_col']:
+                if compare[key] is not None:
+                    compare[key] = json.loads(compare[key])
+            assert self._recursive_compare(compare, sample_data)
+
+    def test_variant_table_creation_pandas(self):
+        table, sample_data = self.sample_variant_table()
+
+        with self.table_context(table) as engine:
+            # Pre-serialize variant data for pandas
+            variant_data = sample_data.copy()
+            for key in ['variant_simple_col', 'variant_nested_col', 'variant_array_col', 'variant_mixed_col']:
+                variant_data[key] = None if sample_data[key] is None else json.dumps(sample_data[key])
+            
+            # Insert the data into the table
+            df = pd.DataFrame([variant_data])
+            dtype_mapping = {
+                "variant_simple_col": DatabricksVariant,
+                "variant_nested_col": DatabricksVariant,
+                "variant_array_col": DatabricksVariant,
+                "variant_mixed_col": DatabricksVariant
+            }
+            df.to_sql(table.__tablename__, engine, if_exists="append", index=False, dtype=dtype_mapping)
+            
+            # Read the data from the table
+            stmt = select(table)
+            df_result = pd.read_sql(stmt, engine)
+            result_dict = df_result.iloc[0].to_dict()
+            # Parse JSON values back to original format for comparison
+            for key in ['variant_simple_col', 'variant_nested_col', 'variant_array_col', 'variant_mixed_col']:
+                if result_dict[key] is not None:
+                    result_dict[key] = json.loads(result_dict[key])
+            assert self._recursive_compare(result_dict, sample_data)

tests/test_local/test_ddl.py

@@ -7,7 +7,7 @@
     SetColumnComment,
     SetTableComment,
 )
-from databricks.sqlalchemy import DatabricksArray, DatabricksMap
+from databricks.sqlalchemy import DatabricksArray, DatabricksMap, DatabricksVariant
 
 
 class DDLTestBase:
@@ -103,7 +103,8 @@ def metadata(self) -> MetaData:
         metadata = MetaData()
         col1 = Column("array_array_string", DatabricksArray(DatabricksArray(String)))
         col2 = Column("map_string_string", DatabricksMap(String, String))
-        table = Table("complex_type", metadata, col1, col2)
+        col3 = Column("variant_col", DatabricksVariant())
+        table = Table("complex_type", metadata, col1, col2, col3)
         return metadata
 
     def test_create_table_with_complex_type(self, metadata):
@@ -112,3 +113,4 @@ def test_create_table_with_complex_type(self, metadata):
 
         assert "array_array_string ARRAY<ARRAY<STRING>>" in output
         assert "map_string_string MAP<STRING,STRING>" in output
+        assert "variant_col VARIANT" in output

tests/test_local/test_types.py

@@ -4,7 +4,7 @@
 import sqlalchemy
 
 from databricks.sqlalchemy.base import DatabricksDialect
-from databricks.sqlalchemy._types import TINYINT, TIMESTAMP, TIMESTAMP_NTZ
+from databricks.sqlalchemy._types import TINYINT, TIMESTAMP, TIMESTAMP_NTZ, DatabricksVariant
 
 
 class DatabricksDataType(enum.Enum):
@@ -28,6 +28,7 @@ class DatabricksDataType(enum.Enum):
     ARRAY = enum.auto()
     MAP = enum.auto()
     STRUCT = enum.auto()
+    VARIANT = enum.auto()
 
 
 # Defines the way that SQLAlchemy CamelCase types are compiled into Databricks SQL types.
@@ -131,6 +132,7 @@ def test_numeric_renders_as_decimal_with_precision_and_scale(self):
     TINYINT: DatabricksDataType.TINYINT,
     TIMESTAMP: DatabricksDataType.TIMESTAMP,
     TIMESTAMP_NTZ: DatabricksDataType.TIMESTAMP_NTZ,
+    DatabricksVariant: DatabricksDataType.VARIANT,
 }