Add support for Trino dialect and database context in PyDough

.github/workflows/pr_testing.yml

@@ -40,6 +40,11 @@ on:
         type: boolean
         required: false
         default: false
+      run-trino:
+        description: "Run Trino Tests"
+        type: boolean
+        required: false
+        default: false
       run-mysql:
         description: "Run MySQL Tests"
         type: boolean
@@ -207,6 +212,25 @@ jobs:
                       && needs.get-py-ver-matrix.outputs.matrix
                       || '["3.10", "3.11", "3.12", "3.13"]' }}
 
+  run-trino-tests:
+    name: Trino Tests
+    needs: [run-defog-daily-update, get-msg, get-py-ver-matrix]
+    if: |
+      (github.event_name == 'pull_request' && contains(needs.get-msg.outputs.commitMsg, '[run all]')) ||
+      (github.event_name == 'pull_request' && contains(needs.get-msg.outputs.commitMsg, '[run dialects]')) ||
+      (github.event_name == 'pull_request' && contains(needs.get-msg.outputs.commitMsg, '[run trino]')) ||
+      (github.event_name == 'workflow_dispatch' && (inputs.run-all || inputs.run-dialects || inputs.run-trino))
+    uses: ./.github/workflows/trino_testing.yml
+    secrets:
+      MYSQL_USERNAME: ${{ secrets.MYSQL_USERNAME }}
+      MYSQL_PASSWORD: ${{ secrets.MYSQL_PASSWORD }}
+      POSTGRES_USER: ${{ secrets.POSTGRES_USER }}
+      POSTGRES_PASSWORD: ${{ secrets.POSTGRES_PASSWORD }}
+    with:
+      python-versions: ${{ github.event_name == 'workflow_dispatch'
+                      && needs.get-py-ver-matrix.outputs.matrix
+                      || '["3.10", "3.11", "3.12", "3.13"]' }}
+
   run-mysql-tests:
     name: MySQL Tests
     needs: [get-msg, get-py-ver-matrix]

.github/workflows/trino_testing.yml

@@ -0,0 +1,129 @@
+name: Run Trino Tests
+
+on:
+  workflow_call:
+    inputs:
+      python-versions:
+        description: "JSON string of Python versions"
+        type: string
+        required: true
+    secrets:
+      POSTGRES_USER:
+        required: true
+      POSTGRES_PASSWORD:
+        required: true
+
+jobs:
+  trino-test:
+    name: Trino Tests (Python ${{ matrix.python-version }})
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ${{ fromJSON(inputs.python-versions) }}
+
+    # Define services here to run Docker containers alongside your job
+    services:
+      postgres:
+        image: bodoai1/pydough-postgres-tpch:latest
+        env:
+          # Set environment variables for Postgres container
+          POSTGRES_USER: ${{ secrets.POSTGRES_USER }}
+          POSTGRES_PASSWORD: ${{ secrets.POSTGRES_PASSWORD }}
+          POSTGRES_DB: "pydough_test"
+        ports:
+          - 5432:5432
+      mysql:
+        image: bodoai1/pydough-mysql-tpch:latest
+        env:
+          # Set environment variables for MySQL container
+          MYSQL_ROOT_PASSWORD: ${{ secrets.MYSQL_PASSWORD }}
+          MYSQL_DATABASE: tpch
+        ports:
+          - 3306:3306
+      mongo:
+        image: bodoai1/pydough-mongo-tpch:latest
+        env:
+          # Set environment variables for MongoDB container
+          MONGO_HOST: 127.0.0.1
+          MONGO_USER: "root_user"
+          MONGO_PASSWORD: "mongo_pwd"
+        ports:
+          - 27017:27017
+      cassandra:
+        image: bodoai1/pydough-cassandra-tpch:latest
+        command: ["--memory=4g", "--memory-swap=6g"]
+        env:
+          # Set environment variables for Cassandra container
+          MAX_HEAP_SIZE: "2G"
+          HEAP_NEWSIZE: "512M"
+          CASSANDRA_CLUSTER_NAME: "test"
+          CASSANDRA_DC: "datacenter1"
+          CASSANDRA_RACK: "rack1"
+          CASSANDRA_HOST: 127.0.0.1
+          FILTER_TABLES: "true"
+          TPCH_TABLES: "partsupp,nation"
+          DEFOG_TABLES: "sbDailyPrice,customers,sales,patients,treatments,adverse_events,coupons,wallet_merchant_balance_daily,user_setting_snapshot,conference,domain_conference,domain_publication,organization,writes,restaurant"
+        ports:
+          - 9042:9042
+      trino:
+        image: bodoai1/pydough-trino:latest
+        env:
+            POSTGRES_HOST: postgres
+            POSTGRES_DB: "pydough_test"
+            POSTGRES_USER: ${{ secrets.POSTGRES_USER }}
+            POSTGRES_PASSWORD: ${{ secrets.POSTGRES_PASSWORD }}
+            MYSQL_HOST: mysql
+            MYSQL_DB: test
+            MYSQL_USERNAME: ${{ secrets.MYSQL_USERNAME }}
+            MYSQL_PASSWORD: ${{ secrets.MYSQL_PASSWORD }}
+        ports:
+            - "8080:8080"
+        depends_on:
+            - postgres
+            - mysql
+            - mongo
+            - cassandra
+    env:
+      POSTGRES_USER: ${{ secrets.POSTGRES_USER }}
+      POSTGRES_PASSWORD: ${{ secrets.POSTGRES_PASSWORD }}
+      POSTGRES_DB: pydough_test
+      POSTGRES_PORT: 5432
+      POSTGRES_HOST: 127.0.0.1
+      MYSQL_USERNAME: ${{ secrets.MYSQL_USERNAME }}
+      MYSQL_PASSWORD: ${{ secrets.MYSQL_PASSWORD }}
+      MYSQL_PORT: 3306
+      MYSQL_HOST: 127.0.0.1
+      MONGO_HOST: 127.0.0.1
+      CASSANDRA_HOST: 127.0.0.1
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Setup Python ${{ matrix.python-version }}
+        id: setup-python
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v3
+        with:
+          version: "0.4.23"
+
+      - name: Create virtual environment
+        run: uv venv
+
+      - name: Install dependencies
+        run: uv pip install -e ".[trino]" 
+
+      - name: Wait for Trino to be ready
+        run: |
+          for i in $(seq 1 36); do
+            pg_isready -h ${{ env.POSTGRES_HOST }} -U ${{ env.POSTGRES_USER }} -d ${{ env.POSTGRES_DB }} && break || sleep 5;
+          done
+          pg_isready -h ${{ env.POSTGRES_HOST }} -U ${{ env.POSTGRES_USER }} -d ${{ env.POSTGRES_DB }} && echo "PostgreSQL is running" || exit 1;
+      - name: Confirm Trino connector is installed
+        run: uv run python -c "import trino; print('Trino connector installed')"
+
+      - name: Run Trino Tests
+        run: uv run pytest -m trino tests/ -rs

pydough/conversion/hybrid_translator.py

@@ -1749,6 +1749,10 @@ def convert_qdag_to_hybrid(self, node: PyDoughCollectionQDAG) -> HybridTree:
         """
         # 1. Run the initial conversion from QDAG to Hybrid
         hybrid: HybridTree = self.make_hybrid_tree(node, None)
+        print()
+        print(node.to_tree_string())
+        print()
+        print(hybrid)
         # 2. Eject any aggregate inputs from the hybrid tree.
         self.eject_aggregate_inputs(hybrid)
         # 3. Syncretize any children of the hybrid tree that share a common

pydough/database_connectors/builtin_databases.py

@@ -22,6 +22,7 @@
     "load_postgres_connection",
     "load_snowflake_connection",
     "load_sqlite_connection",
+    "load_trino_connection",
 ]
 
 
@@ -55,6 +56,9 @@ def load_database_context(database_name: str, **kwargs) -> DatabaseContext:
         case "snowflake":
             connection = load_snowflake_connection(**kwargs)
             dialect = DatabaseDialect.SNOWFLAKE
+        case "trino":
+            connection = load_trino_connection(**kwargs)
+            dialect = DatabaseDialect.TRINO
         case "mysql":
             connection = load_mysql_connection(**kwargs)
             dialect = DatabaseDialect.MYSQL
@@ -136,6 +140,46 @@ def load_snowflake_connection(**kwargs) -> DatabaseConnection:
     return DatabaseConnection(connection)
 
 
+def load_trino_connection(**kwargs) -> DatabaseConnection:
+    """
+    Loads a Trino database connection. This is done by providing a
+    wrapper around the DB 2.0 connect API.
+
+    Args:
+        **kwargs: The keyword arguments which are expected to include either a
+        connection object for Trino or the required connection parameters:
+        - user: Trino username (str)
+        - host: Trino server host (str, default: "127.0.01"/"localhost")
+        - port: Trino server port (int, default: 8080)
+
+    Raises:
+        ImportError: If the Trino connector is not installed.
+        ValueError: If required connection parameters are missing.
+
+    Returns:
+        A DatabaseContext object for Trino.
+    """
+    try:
+        import trino
+    except ImportError:
+        raise ImportError(
+            "Trino connector is not installed. Please install it with `pip install trino`."
+        )
+
+    connection = kwargs.pop("connection", None)
+    if connection:
+        return DatabaseConnection(connection)
+    required_keys = ["user", "host", "port"]
+    if not all(key in kwargs for key in required_keys):
+        raise ValueError(
+            "Trino connection requires the following arguments: "
+            + ", ".join(required_keys)
+        )
+    # Create a Trino connection using the provided keyword arguments
+    connection = trino.dbapi.connect(**kwargs)
+    return DatabaseConnection(connection)
+
+
 def load_mysql_connection(**kwargs) -> DatabaseConnection:
     """
     Loads a MySQL database connection. This is done by providing a wrapper

pydough/database_connectors/database_connector.py

@@ -193,6 +193,7 @@ class DatabaseDialect(Enum):
     ANSI = "ansi"
     SQLITE = "sqlite"
     SNOWFLAKE = "snowflake"
+    TRINO = "trino"
     MYSQL = "mysql"
     POSTGRES = "postgres"
     ORACLE = "oracle"

pydough/sqlglot/execute_relational.py

@@ -15,9 +15,11 @@
 from sqlglot.dialects import Postgres as PostgresDialect
 from sqlglot.dialects import Snowflake as SnowflakeDialect
 from sqlglot.dialects import SQLite as SQLiteDialect
+from sqlglot.dialects import Trino as TrinoDialect
 from sqlglot.dialects.dialect import rename_func
 from sqlglot.dialects.mysql import MySQL
 from sqlglot.dialects.oracle import Oracle
+from sqlglot.dialects.trino import Trino
 from sqlglot.errors import SqlglotError
 from sqlglot.expressions import (
     Alias,
@@ -552,6 +554,8 @@ def convert_dialect_to_sqlglot(dialect: DatabaseDialect) -> SQLGlotDialect:
             # The BodoSQL dialect is essentially a subset of the Snowflake SQL
             # dialect without many of the extraneous features.
             return SnowflakeDialect()
+        case DatabaseDialect.TRINO:
+            return TrinoDialect()
         case DatabaseDialect.MYSQL:
             return MySQLDialect()
         case DatabaseDialect.POSTGRES:
@@ -602,6 +606,15 @@ def change_sqlglot_dialect_configuration(dialect: DatabaseDialect) -> None:
                 )
             )
 
+        case DatabaseDialect.TRINO:
+            # Replace the DAYOFWEEK override for the Trino generator with the
+            # default version, since PyDough handles the conversion logic
+            # correctly whereas the SQLGlot version gets confused with multiple
+            # rounds of parsing and unparsing.
+            Trino.Generator.TRANSFORMS[sqlglot_expressions.DayOfWeek] = (
+                lambda self, e: f"DAY_OF_WEEK({self.sql(e.this)})"
+            )
+
         case _:
             pass
 
@@ -629,6 +642,15 @@ def reset_sqlglot_dialect_configuration(dialect: DatabaseDialect) -> None:
                     ]
                 )
             )
+        case DatabaseDialect.TRINO:
+            # Replace the DAYOFWEEK override for the Trino generator with the
+            # default version, since PyDough handles the conversion logic
+            # correctly whereas the SQLGlot version gets confused with multiple
+            # rounds of parsing and unparsing.
+            Trino.Generator.TRANSFORMS[sqlglot_expressions.DayOfWeek] = (
+                lambda self, e: f"(({self.func('DAY_OF_WEEK', e.this)} % 7) + 1)",
+            )
+
         case _:
             pass