Add support for LISTOF function, array return, and EXPLODE operator in PyDough

pydough/database_connectors/builtin_databases.py

@@ -90,7 +90,7 @@ def load_sqlite_connection(**kwargs) -> DatabaseConnection:
     if "database" not in kwargs:
         raise PyDoughSessionException("SQLite connection requires a database path.")
     connection: sqlite3.Connection = sqlite3.connect(**kwargs)
-    return DatabaseConnection(connection)
+    return DatabaseConnection(connection, DatabaseDialect.SQLITE)
 
 
 def load_snowflake_connection(**kwargs) -> DatabaseConnection:
@@ -122,7 +122,7 @@ def load_snowflake_connection(**kwargs) -> DatabaseConnection:
     connection: snowflake.connector.connection.SnowflakeConnection
     if connection := kwargs.pop("connection", None):
         # If a connection object is provided, return it wrapped in DatabaseConnection
-        return DatabaseConnection(connection)
+        return DatabaseConnection(connection, DatabaseDialect.SNOWFLAKE)
     # Snowflake connection requires specific parameters:
     # user, password, account.
     # Raise an error if any of these are missing.
@@ -137,7 +137,7 @@ def load_snowflake_connection(**kwargs) -> DatabaseConnection:
         )
     # Create a Snowflake connection using the provided keyword arguments
     connection = snowflake.connector.connect(**kwargs)
-    return DatabaseConnection(connection)
+    return DatabaseConnection(connection, DatabaseDialect.SNOWFLAKE)
 
 
 def load_trino_connection(**kwargs) -> DatabaseConnection:
@@ -168,7 +168,7 @@ def load_trino_connection(**kwargs) -> DatabaseConnection:
 
     connection = kwargs.pop("connection", None)
     if connection:
-        return DatabaseConnection(connection)
+        return DatabaseConnection(connection, dialect=DatabaseDialect.TRINO)
     required_keys = ["user", "host", "port"]
     if not all(key in kwargs for key in required_keys):
         raise ValueError(
@@ -177,7 +177,7 @@ def load_trino_connection(**kwargs) -> DatabaseConnection:
         )
     # Create a Trino connection using the provided keyword arguments
     connection = trino.dbapi.connect(**kwargs)
-    return DatabaseConnection(connection)
+    return DatabaseConnection(connection, dialect=DatabaseDialect.TRINO)
 
 
 def load_mysql_connection(**kwargs) -> DatabaseConnection:
@@ -222,7 +222,7 @@ def load_mysql_connection(**kwargs) -> DatabaseConnection:
     if connection := kwargs.pop("connection", None):
         # If a connection object is provided, return it wrapped in
         # DatabaseConnection
-        return DatabaseConnection(connection)
+        return DatabaseConnection(connection, DatabaseDialect.MYSQL)
 
     # MySQL connection requires specific parameters:
     # user, password, database.
@@ -257,7 +257,7 @@ def load_mysql_connection(**kwargs) -> DatabaseConnection:
     while attempt <= attempts:
         try:
             connection = mysql.connector.connect(**kwargs)
-            return DatabaseConnection(connection)
+            return DatabaseConnection(connection, DatabaseDialect.MYSQL)
 
         except (OSError, mysql.connector.Error) as err:
             if attempt >= attempts:
@@ -313,7 +313,7 @@ def load_postgres_connection(**kwargs) -> DatabaseConnection:
     if connection := kwargs.pop("connection", None):
         # If a connection object is provided, return it wrapped in
         # DatabaseConnection
-        return DatabaseConnection(connection)
+        return DatabaseConnection(connection, DatabaseDialect.POSTGRES)
 
     # Postgres connection requires specific parameters:
     # user, password, dbname.
@@ -347,7 +347,7 @@ def load_postgres_connection(**kwargs) -> DatabaseConnection:
     while attempt <= attempts:
         try:
             connection = psycopg2.connect(**kwargs)
-            return DatabaseConnection(connection)
+            return DatabaseConnection(connection, DatabaseDialect.POSTGRES)
 
         except (OSError, psycopg2.Error) as err:
             if attempt >= attempts:
@@ -401,7 +401,7 @@ def load_oracle_connection(**kwargs) -> DatabaseConnection:
         # If a connection object is provided, return it wrapped in
         # DatabaseConnection
         assert isinstance(connection, oracledb.Connection)
-        return DatabaseConnection(connection)
+        return DatabaseConnection(connection, DatabaseDialect.ORACLE)
 
     # Oracle connection requires specific parameters:
     # user, password, host and service_name.
@@ -439,7 +439,7 @@ def load_oracle_connection(**kwargs) -> DatabaseConnection:
     while attempt <= attempts:
         try:
             connection = oracledb.connect(**kwargs)
-            return DatabaseConnection(connection)
+            return DatabaseConnection(connection, DatabaseDialect.ORACLE)
 
         except (OSError, oracledb.Error) as err:
             if attempt >= attempts:

pydough/database_connectors/database_connector.py

@@ -11,6 +11,7 @@
     "DatabaseDialect",
 ]
 
+import json
 from dataclasses import dataclass
 from enum import Enum
 from typing import TYPE_CHECKING, Union, cast
@@ -36,10 +37,12 @@ class DatabaseConnection:
     # with Python.
     _connection: DBConnection
     _cursor: DBCursor | None
+    _dialect: "DatabaseDialect"
 
-    def __init__(self, connection: DBConnection) -> None:
+    def __init__(self, connection: DBConnection, dialect: "DatabaseDialect") -> None:
         self._connection = connection
         self._cursor = None
+        self._dialect = dialect
 
     def execute_query_df(self, sql: str) -> pd.DataFrame:
         """Create a cursor object using the connection and execute the query,
@@ -63,9 +66,40 @@ def execute_query_df(self, sql: str) -> pd.DataFrame:
             # check at run-time if the cursor has the method.
             if TYPE_CHECKING:
                 _ = cast(SnowflakeCursor, self.cursor).fetch_pandas_all
+
+            # Identify which columns must be coerced from strings into native
+            # types using json.loads based on the cursor description.
+            semi_structured_cols: set[int] = set()
+            match self._dialect:
+                case DatabaseDialect.SNOWFLAKE:
+                    # Snowflake returns array/object/variant types as JSON
+                    # strings (types 5/9/10), so we need to parse those back
+                    # into Python types.
+                    semi_structured_cols.update(
+                        {
+                            idx
+                            for idx, desc in enumerate(self.cursor.description)
+                            if desc[1] in (5, 9, 10)
+                        }
+                    )
+                case DatabaseDialect.MYSQL:
+                    # Snowflake returns JSON data (type 245) as strings, so we
+                    # need to parse those back into Python types.
+                    semi_structured_cols.update(
+                        {
+                            idx
+                            for idx, desc in enumerate(self.cursor.description)
+                            if desc[1] == 245
+                        }
+                    )
+                case _:
+                    pass
             # At run-time check and run the fetch.
             if hasattr(self.cursor, "fetch_pandas_all"):
-                return self.cursor.fetch_pandas_all()
+                pd_table: pd.DataFrame = self.cursor.fetch_pandas_all()
+                for idx in semi_structured_cols:
+                    pd_table.iloc[:, idx] = pd_table.iloc[:, idx].apply(json.loads)
+                return pd_table
             else:
                 # Assume sqlite3
                 column_names: list[str] = [
@@ -74,7 +108,10 @@ def execute_query_df(self, sql: str) -> pd.DataFrame:
                 # TODO: (gh #174) Cache the cursor?
                 # TODO: (gh #175) enable typed DataFrames.
                 data = self.cursor.fetchall()
-                return pd.DataFrame(data, columns=column_names)
+                pd_table = pd.DataFrame(data, columns=column_names)
+                for idx in semi_structured_cols:
+                    pd_table.iloc[:, idx] = pd_table.iloc[:, idx].apply(json.loads)
+                return pd_table
         except Exception as e:
             print(f"ERROR WHILE EXECUTING QUERY:\n{sql}")
             raise pydough.active_session.error_builder.sql_runtime_failure(

pydough/database_connectors/empty_connection.py

@@ -11,7 +11,7 @@
 
 from pydough.errors import PyDoughSessionException
 
-from .database_connector import DatabaseConnection
+from .database_connector import DatabaseConnection, DatabaseDialect
 
 
 class EmptyConnection(Connection):
@@ -36,4 +36,6 @@ def cursor(self, *args, **kwargs):
         raise PyDoughSessionException("No SQL Database is specified.")
 
 
-empty_connection: DatabaseConnection = DatabaseConnection(EmptyConnection())
+empty_connection: DatabaseConnection = DatabaseConnection(
+    EmptyConnection(), DatabaseDialect.ANSI
+)

pydough/errors/pydough_error_builder.py

@@ -255,6 +255,21 @@ def sql_call_conversion_error(
             f"Failed to convert expression {call.to_string(True)} to SQL: {error}"
         )
 
+    def sql_call_dialect_unsupported(
+        self, operator: "PyDoughOperator", dialect: str
+    ) -> PyDoughException:
+        """
+        Creates an exception for when a SQL dialect does not allow converting
+        a certain feature to SQL.
+
+        Args:
+            `operator`: The function operator that is not supported.
+            `dialect`: The SQL dialect in which the feature is not supported.
+        """
+        return PyDoughSQLException(
+            f"Cannot convert function {operator} to SQL using dialect {dialect}"
+        )
+
     def undefined_function_call(
         self, node: "UnqualifiedNode", *args, **kwargs
     ) -> PyDoughException:

pydough/pydough_operators/__init__.py

@@ -51,6 +51,7 @@
     "LEQ",
     "LET",
     "LIKE",
+    "LISTOF",
     "LOWER",
     "LPAD",
     "MAX",
@@ -156,6 +157,7 @@
     LEQ,
     LET,
     LIKE,
+    LISTOF,
     LOWER,
     LPAD,
     MAX,

pydough/pydough_operators/expression_operators/__init__.py

@@ -48,6 +48,7 @@
     "LEQ",
     "LET",
     "LIKE",
+    "LISTOF",
     "LOWER",
     "LPAD",
     "MAX",
@@ -149,6 +150,7 @@
     LEQ,
     LET,
     LIKE,
+    LISTOF,
     LOWER,
     LPAD,
     MAX,

pydough/pydough_operators/expression_operators/registered_expression_operators.py

@@ -42,6 +42,7 @@
     "LEQ",
     "LET",
     "LIKE",
+    "LISTOF",
     "LOWER",
     "LPAD",
     "MAX",
@@ -167,6 +168,9 @@
 QUANTILE = ExpressionFunctionOperator(
     "QUANTILE", True, RequireNumArgs(2), ConstantType(NumericType())
 )
+LISTOF = ExpressionFunctionOperator(
+    "LISTOF", True, RequireNumArgs(1), SelectArgumentType(0)
+)
 POWER = ExpressionFunctionOperator(
     "POWER", False, RequireNumArgs(2), ConstantType(NumericType())
 )

pydough/sqlglot/transform_bindings/base_transform_bindings.py

@@ -23,11 +23,13 @@
     LiteralExpression,
 )
 from pydough.types import (
+    ArrayType,
     BooleanType,
     DatetimeType,
     NumericType,
     PyDoughType,
     StringType,
+    UnknownType,
 )
 from pydough.user_collections.dataframe_collection import DataframeGeneratedCollection
 from pydough.user_collections.range_collection import RangeGeneratedCollection
@@ -265,6 +267,8 @@ def convert_call_to_sqlglot(
                 return sqlglot_expressions.Count(
                     this=sqlglot_expressions.Distinct(expressions=[args[0]])
                 )
+            case pydop.LISTOF:
+                return self.convert_listof(args, types)
             case pydop.STARTSWITH:
                 return self.convert_startswith(args, types)
             case pydop.ENDSWITH:
@@ -405,6 +409,17 @@ def convert_sum(
         """
         return sqlglot_expressions.Sum.from_arg_list(args)
 
+    def convert_listof(
+        self, args: SQLGlotExpression, types: list[PyDoughType]
+    ) -> SQLGlotExpression:
+        """
+        Converts a LISTOF function call to its SQLGlot equivalent. Some dialects
+        do not support this functionality.
+        """
+        raise self._visitor._session.error_builder.sql_call_dialect_unsupported(
+            pydop.LISTOF, self._visitor._session.database.dialect.name
+        )
+
     def convert_find(
         self,
         args: list[SQLGlotExpression],
@@ -2451,6 +2466,16 @@ def generate_dataframe_item_expression(
         (None, NaN, BooleanType, StringType) meaning that this representation
         works through all current dialects.
         """
+        if isinstance(item, list):
+            inner_type: PyDoughType
+            if isinstance(item_type, ArrayType):
+                inner_type = item_type.elem_type
+            else:
+                inner_type = UnknownType()
+            inner_items: list[SQLGlotExpression] = [
+                self.generate_dataframe_item_expression(i, inner_type) for i in item
+            ]
+            return self.generate_dataframe_array_expression(inner_items, inner_type)
 
         if item is None or pd.isna(item):
             return sqlglot_expressions.Null()
@@ -2465,6 +2490,22 @@ def generate_dataframe_item_expression(
             case _:  # Specific dialect expression
                 return self.generate_dataframe_item_dialect_expression(item, item_type)
 
+    def generate_dataframe_array_expression(
+        self, items: list[SQLGlotExpression], inner_type: PyDoughType
+    ) -> SQLGlotExpression:
+        """
+        Generate the sqlglot expression for an array of items with given pydough type.
+
+        Args:
+            `items` : The list of SQLGlotExpressions representing the items in the array.
+            `inner_type` : The mapped PydDough type for the items in the array.
+        Returns:
+            A SQLGlotExpression representing the array of items.
+        """
+        raise NotImplementedError(
+            f"Array types are not currently supported in dialect {self._visitor._expr_visitor._dialect.name}."
+        )
+
     def generate_dataframe_item_dialect_expression(
         self, item: Any, item_type: PyDoughType
     ) -> SQLGlotExpression:

pydough/sqlglot/transform_bindings/mysql_transform_bindings.py

@@ -88,6 +88,16 @@ def convert_call_to_sqlglot(
 
         return super().convert_call_to_sqlglot(operator, args, types)
 
+    def convert_listof(
+        self, args: SQLGlotExpression, types: list[PyDoughType]
+    ) -> SQLGlotExpression:
+        return sqlglot_expressions.Anonymous(this="JSON_ARRAYAGG", expressions=args)
+
+    def generate_dataframe_array_expression(
+        self, items: list[SQLGlotExpression], inner_type: PyDoughType
+    ) -> SQLGlotExpression:
+        return sqlglot_expressions.JSONArray(expressions=items)
+
     def convert_slice(
         self, args: list[SQLGlotExpression], types: list[PyDoughType]
     ) -> SQLGlotExpression:

pydough/sqlglot/transform_bindings/oracle_transform_bindings.py

@@ -97,6 +97,28 @@ def convert_call_to_sqlglot(
 
         return super().convert_call_to_sqlglot(operator, args, types)
 
+    def convert_listof(
+        self, args: SQLGlotExpression, types: list[PyDoughType]
+    ) -> SQLGlotExpression:
+        return sqlglot_expressions.Anonymous(this="JSON_ARRAYAGG", expressions=args)
+
+    def generate_dataframe_array_expression(
+        self, items: list[SQLGlotExpression], inner_type: PyDoughType
+    ) -> SQLGlotExpression:
+        func: str
+        match inner_type:
+            case StringType():
+                func = "SYS.ODCIVARCHAR2LIST"
+            case NumericType():
+                func = "SYS.ODCINUMBERLIST"
+            case DatetimeType():
+                func = "SYS.ODCIDATELIST"
+            case _:
+                raise ValueError(
+                    f"Cannot support constant array of type {inner_type} in Oracle."
+                )
+        return sqlglot_expressions.Anonymous(this=func, expressions=items)
+
     def convert_default_to(
         self, args: list[SQLGlotExpression], types: list[PyDoughType]
     ) -> SQLGlotExpression: