diff --git a/README.md b/README.md index d4ce3c0ea135..ae9ccd0d5b86 100644 --- a/README.md +++ b/README.md @@ -159,7 +159,7 @@ Ibis supports more than 20 backends: - [SQLite](https://ibis-project.org/backends/sqlite/) - [Snowflake](https://ibis-project.org/backends/snowflake) - [Trino](https://ibis-project.org/backends/trino/) - +- [Db2](https://ibis-project.org/backends/db2/) ## How it works Most Python dataframes are tightly coupled to their execution engine. And many databases only support SQL, with no Python API. Ibis solves this problem by providing a common API for data manipulation in Python, and compiling that API into the backend’s native language. This means you can learn a single API and use it across any supported backend (execution engine). diff --git a/docs/backends/db2.qmd b/docs/backends/db2.qmd new file mode 100644 index 000000000000..6f5ee75e32ae --- /dev/null +++ b/docs/backends/db2.qmd @@ -0,0 +1,173 @@ +# IBM DB2 + +[https://www.ibm.com/products/db2](https://www.ibm.com/products/db2) + +![](https://img.shields.io/badge/memtables-fallback-yellow?style=flat-square) ![](https://img.shields.io/badge/inputs-DB2 tables-blue?style=flat-square) ![](https://img.shields.io/badge/outputs-DB2 tables | CSV | pandas | Parquet | PyArrow-orange?style=flat-square) + +## Install + +Install Ibis and dependencies for the DB2 backend: + +::: {.panel-tabset} + +## `pip` + +Install with the `db2` extra: + +```{.bash} +pip install 'ibis-framework[db2]' +``` + +And connect: + +```{.python} +import ibis + +con = ibis.db2.connect() # <1> +``` + +1. Adjust connection parameters as needed. + +## `conda` + +Install for DB2: + +```{.bash} +conda install -c conda-forge ibis-db2 +``` + +And connect: + +```{.python} +import ibis + +con = ibis.db2.connect() # <1> +``` + +1. Adjust connection parameters as needed. + +## `mamba` + +Install for DB2: + +```{.bash} +mamba install -c conda-forge ibis-db2 +``` + +And connect: + +```{.python} +import ibis + +con = ibis.db2.connect() # <1> +``` + +1. Adjust connection parameters as needed. + +::: + +## Connect + +### `ibis.db2.connect` + +```python +con = ibis.db2.connect( + database="SAMPLE", + hostname="localhost", + port=50000, + username="db2inst1", + password="password", + schema="MYSCHEMA", +) +``` + +::: {.callout-note} +`ibis.db2.connect` is a thin wrapper around [`ibis.backends.db2.Backend.do_connect`](#ibis.backends.db2.Backend.do_connect). +::: + +### Connection Parameters + +```{python} +#| echo: false +#| output: asis +from _utils import render_do_connect + +render_do_connect("db2") +``` + +### `ibis.connect` URL format + +In addition to `ibis.db2.connect`, you can also connect to DB2 by +passing a properly-formatted DB2 connection URL to `ibis.connect`: + +```python +con = ibis.connect(f"db2://{username}:{password}@{hostname}:{port}/{database}") +``` + +## DB2 Client Requirements + +The DB2 backend requires the IBM DB2 client libraries to be installed. You can install them using: + +```bash +pip install ibm_db +``` + +::: {.callout-note} +The `ibm_db` package includes `ibm_db_dbi` (the DB-API 2.0 compliant interface) and requires the DB2 client libraries to be installed on your system. +For installation instructions, see the [IBM DB2 documentation](https://www.ibm.com/docs/en/db2). +::: + +## Features + +The DB2 backend supports: + +- **Temporary tables**: Create temporary tables for intermediate results +- **Schema operations**: List databases (schemas), tables, and get table schemas +- **Data manipulation**: Insert, create, and drop tables +- **Query execution**: Execute SQL queries and return results as pandas DataFrames or PyArrow tables +- **Type mapping**: Automatic conversion between Ibis types and DB2 types + +## Limitations + +- **Python UDFs**: Not currently supported +- **Streaming**: Not supported + +## Example Usage + +```python +import ibis + +# Connect to DB2 +con = ibis.db2.connect( + database="SAMPLE", + hostname="localhost", + port=50000, + username="db2inst1", + password="password" +) + +# List available tables +tables = con.list_tables() +print(tables) + +# Load a table +employee = con.table("EMPLOYEE") + +# Query the table +result = employee.filter(employee.SALARY > 50000).execute() +print(result) + +# Create a new table +con.create_table( + "new_table", + obj=result, + overwrite=True +) +``` + +```{python} +#| echo: false +BACKEND = "DB2" +``` + +{{< include ./_templates/api.qmd >}} \ No newline at end of file diff --git a/ibis/backends/db2/__init__.py b/ibis/backends/db2/__init__.py new file mode 100644 index 000000000000..c32d09c5122e --- /dev/null +++ b/ibis/backends/db2/__init__.py @@ -0,0 +1,720 @@ +"""IBM DB2 backend for Ibis.""" + +from __future__ import annotations + +import contextlib +from typing import TYPE_CHECKING, Any +from urllib.parse import unquote_plus + +import sqlglot as sg + +import ibis +import ibis.backends.sql.compilers as sc +import ibis.common.exceptions as exc +import ibis.expr.operations as ops +import ibis.expr.schema as sch +import ibis.expr.types as ir +from ibis.backends.sql import SQLBackend +from ibis.backends.db2.datatypes import DB2PandasData, parse_db2_type, ibis_type_to_db2_type + +if TYPE_CHECKING: + from collections.abc import Mapping + from urllib.parse import ParseResult + + import pandas as pd + from typing_extensions import Self + + +class Backend(SQLBackend): + """IBM DB2 backend for Ibis.""" + + name = "db2" + supports_temporary_tables = True + supports_python_udfs = False + + @property + def compiler(self): + """Lazy load the compiler to avoid circular imports.""" + from ibis.backends.sql.compilers.db2 import DB2Compiler + if not hasattr(self, '_compiler'): + self._compiler = DB2Compiler() + return self._compiler + + def __init__(self, *args, **kwargs): + """Initialize DB2 backend.""" + super().__init__(*args, **kwargs) + self._connection = None + self._cursor = None + + @property + def version(self) -> str: + """Return the version of the DB2 server.""" + with self._safe_raw_sql("SELECT SERVICE_LEVEL FROM SYSIBMADM.ENV_INST_INFO") as cur: + result = cur.fetchone() + return result[0] if result else "unknown" + + def do_connect( + self, + database: str, + hostname: str = "localhost", + port: int = 50000, + username: str | None = None, + password: str | None = None, + schema: str | None = None, + **kwargs: Any, + ) -> None: + """ + Connect to a DB2 database. + + Parameters + ---------- + database : str + Database name + hostname : str, default "localhost" + Hostname of DB2 server + port : int, default 50000 + Port number + username : str, optional + Username for authentication + password : str, optional + Password for authentication + schema : str, optional + Default schema + **kwargs + Additional connection parameters + """ + import ibm_db + import ibm_db_dbi + + # Build connection string + conn_str_parts = [ + f"DATABASE={database}", + f"HOSTNAME={hostname}", + f"PORT={port}", + "PROTOCOL=TCPIP", + ] + + if username: + conn_str_parts.append(f"UID={username}") + if password: + conn_str_parts.append(f"PWD={password}") + + # Add any additional connection parameters + for key, value in kwargs.items(): + if key.upper() not in ("DATABASE", "HOSTNAME", "PORT", "UID", "PWD"): + conn_str_parts.append(f"{key.upper()}={value}") + + conn_str = ";".join(conn_str_parts) + + try: + # Connect using ibm_db + ibm_db_conn = ibm_db.connect(conn_str, "", "") + # Wrap with DBI-compliant interface + self._connection = ibm_db_dbi.Connection(ibm_db_conn) + self._cursor = self._connection.cursor() + + # Set schema if provided + if schema: + self._cursor.execute(f"SET SCHEMA {schema}") + + except Exception as e: + raise exc.OperationNotDefinedError( + f"Failed to connect to DB2: {e}" + ) from e + + def _from_url(self, url: ParseResult, **kwarg_overrides): + """Create a DB2 backend from a URL. + + Parameters + ---------- + url : ParseResult + Parsed URL object + **kwarg_overrides + Additional keyword arguments to override URL parameters + + Returns + ------- + Self + Connected DB2 backend instance + """ + kwargs = {} + database, *schema = url.path[1:].split("/", 1) + if url.username: + kwargs["username"] = url.username + if url.password: + kwargs["password"] = unquote_plus(url.password) + if url.hostname: + kwargs["hostname"] = url.hostname + if database: + kwargs["database"] = database + if url.port: + kwargs["port"] = url.port + if schema: + kwargs["schema"] = schema[0] + kwargs.update(kwarg_overrides) + return self.connect(**kwargs) + + def disconnect(self) -> None: + """Disconnect from the database.""" + if self._cursor: + self._cursor.close() + self._cursor = None + if self._connection: + self._connection.close() + self._connection = None + + @contextlib.contextmanager + def _safe_raw_sql(self, query: str, **kwargs: Any): + """Execute raw SQL safely with cursor management.""" + cursor = self._connection.cursor() + try: + cursor.execute(query, **kwargs) + yield cursor + finally: + cursor.close() + + def raw_sql(self, query: str | sg.Expression, **kwargs: Any) -> Any: + """ + Execute a raw SQL query. + + Parameters + ---------- + query : str | sg.Expression + SQL query to execute + **kwargs + Additional parameters + + Returns + ------- + Any + Query results (cursor) + """ + if isinstance(query, sg.exp.Expression): + query = query.sql(dialect=self.compiler.dialect) + + # Don't use context manager as it closes the cursor + cursor = self._connection.cursor() + try: + cursor.execute(query, **kwargs) + except Exception: + cursor.close() + raise + else: + return cursor + + def list_tables( + self, like: str | None = None, database: str | None = None + ) -> list[str]: + """ + List tables in the database. + + Parameters + ---------- + like : str, optional + Pattern to filter table names + database : str, optional + Database/schema name + + Returns + ------- + list[str] + List of table names + """ + query = """ + SELECT TABNAME + FROM SYSCAT.TABLES + WHERE TABSCHEMA = CURRENT SCHEMA + AND TYPE = 'T' + """ + + if like: + query += f" AND TABNAME LIKE '{like}'" + + query += " ORDER BY TABNAME" + + with self._safe_raw_sql(query) as cursor: + return [row[0] for row in cursor.fetchall()] + + def list_databases(self, like: str | None = None) -> list[str]: + """ + List schemas in the database. + + Parameters + ---------- + like : str, optional + Pattern to filter schema names + + Returns + ------- + list[str] + List of schema names + """ + query = """ + SELECT SCHEMANAME + FROM SYSCAT.SCHEMATA + WHERE SCHEMANAME NOT LIKE 'SYS%' + """ + + if like: + query += f" AND SCHEMANAME LIKE '{like}'" + + query += " ORDER BY SCHEMANAME" + + with self._safe_raw_sql(query) as cursor: + return [row[0] for row in cursor.fetchall()] + + def get_schema( + self, + table_name: str, + *, + catalog: str | None = None, + database: str | None = None, + ) -> sch.Schema: + """ + Get the schema of a table. + + Parameters + ---------- + table_name : str + Name of the table + catalog : str, optional + Catalog name (unused in DB2) + database : str, optional + Schema name + + Returns + ------- + sch.Schema + Table schema + """ + query = """ + SELECT COLNAME, TYPENAME, LENGTH, SCALE, NULLS + FROM SYSCAT.COLUMNS + WHERE TABNAME = ? + AND TABSCHEMA = COALESCE(?, CURRENT SCHEMA) + ORDER BY COLNO + """ + + schema_name = database or self.current_database + + cursor = self._connection.cursor() + try: + cursor.execute(query, (table_name.upper(), schema_name)) + rows = cursor.fetchall() + finally: + cursor.close() + + if not rows: + raise exc.IbisError(f"Table not found: {table_name}") + + fields = {} + for col_name, type_name, length, scale, nulls in rows: + # Build type string with parameters + if type_name in ("DECIMAL", "NUMERIC"): + type_str = f"{type_name}({length},{scale})" + elif type_name in ("VARCHAR", "CHAR", "VARBINARY"): + type_str = f"{type_name}({length})" + else: + type_str = type_name + + ibis_type = parse_db2_type(type_str) + # Set nullable based on NULLS column + ibis_type = ibis_type(nullable=(nulls == "Y")) + # Keep column names in uppercase as DB2 stores them + fields[col_name] = ibis_type + + return sch.Schema(fields) + + @property + def current_database(self) -> str: + """Return the current schema.""" + with self._safe_raw_sql("VALUES CURRENT SCHEMA") as cursor: + result = cursor.fetchone() + return result[0] if result else None + + def create_table( + self, + name: str, + obj: ir.Table | pd.DataFrame | None = None, + *, + schema: sch.Schema | None = None, + database: str | None = None, + temp: bool = False, + overwrite: bool = False, + ) -> ir.Table: + """ + Create a new table. + + Parameters + ---------- + name : str + Table name + obj : ir.Table | pd.DataFrame, optional + Data to insert + schema : sch.Schema, optional + Table schema + database : str, optional + Schema name + temp : bool, default False + Create temporary table + overwrite : bool, default False + Overwrite if exists + + Returns + ------- + ir.Table + Table expression + """ + import pandas as pd + + if obj is None and schema is None: + raise exc.IbisError("Either obj or schema must be provided") + + if schema is None: + if isinstance(obj, pd.DataFrame): + schema = sch.infer(obj) + else: + schema = obj.schema() + + # Build CREATE TABLE statement + temp_clause = "GLOBAL TEMPORARY " if temp else "" + db_prefix = f"{database}." if database else "" + full_name = f"{db_prefix}{name}" + + if overwrite: + self.drop_table(name, database=database, force=True) + + # Build column definitions + col_defs = [] + for col_name, col_type in schema.items(): + db2_type = ibis_type_to_db2_type(col_type) + nullable = "NULL" if col_type.nullable else "NOT NULL" + col_defs.append(f"{col_name} {db2_type} {nullable}") + + columns_sql = ", ".join(col_defs) + create_sql = f"CREATE {temp_clause}TABLE {full_name} ({columns_sql})" + + self.raw_sql(create_sql) + # Commit the CREATE TABLE statement + self._connection.commit() + + # Insert data if provided + if obj is not None: + if isinstance(obj, pd.DataFrame): + self.insert(name, obj, database=database) + else: + # Insert from table expression + insert_sql = f"INSERT INTO {full_name} {self.compile(obj)}" + self.raw_sql(insert_sql) + + return self.table(name, database=database) + + def drop_table( + self, + name: str, + *, + database: str | None = None, + force: bool = False, + ) -> None: + """ + Drop a table. + + Parameters + ---------- + name : str + Table name + database : str, optional + Schema name + force : bool, default False + Suppress errors if table doesn't exist + """ + db_prefix = f"{database}." if database else "" + full_name = f"{db_prefix}{name}" + + if force: + # Check if table exists first using parameterized query + cursor = self._connection.cursor() + try: + if database: + check_sql = """ + SELECT COUNT(*) + FROM SYSCAT.TABLES + WHERE TABNAME = ? + AND TABSCHEMA = ? + """ + cursor.execute(check_sql, (name.upper(), database.upper())) + else: + check_sql = """ + SELECT COUNT(*) + FROM SYSCAT.TABLES + WHERE TABNAME = ? + AND TABSCHEMA = CURRENT SCHEMA + """ + cursor.execute(check_sql, (name.upper(),)) + + exists = cursor.fetchone()[0] > 0 + finally: + cursor.close() + + if not exists: + return + + drop_sql = f"DROP TABLE {full_name}" + self.raw_sql(drop_sql) + # Commit the DROP TABLE statement + self._connection.commit() + + def insert( + self, + table_name: str, + obj: pd.DataFrame | ir.Table, + *, + database: str | None = None, + overwrite: bool = False, + ) -> None: + """ + Insert data into a table. + + Parameters + ---------- + table_name : str + Target table name + obj : pd.DataFrame | ir.Table + Data to insert + database : str, optional + Schema name + overwrite : bool, default False + Truncate table before insert + """ + import pandas as pd + + db_prefix = f"{database}." if database else "" + full_name = f"{db_prefix}{table_name}" + + if overwrite: + self.raw_sql(f"TRUNCATE TABLE {full_name} IMMEDIATE") + + if isinstance(obj, pd.DataFrame): + # Batch insert from DataFrame + if obj.empty: + return + + columns = ", ".join(obj.columns) + placeholders = ", ".join(["?" for _ in obj.columns]) + insert_sql = f"INSERT INTO {full_name} ({columns}) VALUES ({placeholders})" + + cursor = self._connection.cursor() + try: + # Insert in batches + batch_size = 1000 + for i in range(0, len(obj), batch_size): + batch = obj.iloc[i : i + batch_size] + cursor.executemany(insert_sql, batch.values.tolist()) + self._connection.commit() + finally: + cursor.close() + else: + # Insert from table expression + insert_sql = f"INSERT INTO {full_name} {self.compile(obj)}" + self.raw_sql(insert_sql) + + def to_pyarrow( + self, + expr: ir.Expr, + *, + params: Mapping[ir.Scalar, Any] | None = None, + limit: int | str | None = None, + **kwargs: Any, + ): + """ + Execute expression and return results as PyArrow table. + + Parameters + ---------- + expr : ir.Expr + Ibis expression + params : Mapping[ir.Scalar, Any], optional + Query parameters + limit : int | str, optional + Result limit + **kwargs + Additional arguments + + Returns + ------- + pyarrow.Table + Query results + """ + import pyarrow as pa + + df = self.to_pandas(expr, params=params, limit=limit, **kwargs) + return pa.Table.from_pandas(df) + + def to_pandas( + self, + expr: ir.Expr, + *, + params: Mapping[ir.Scalar, Any] | None = None, + limit: int | str | None = None, + **kwargs: Any, + ) -> pd.DataFrame: + """ + Execute expression and return results as pandas DataFrame. + + Parameters + ---------- + expr : ir.Expr + Ibis expression + params : Mapping[ir.Scalar, Any], optional + Query parameters + limit : int | str, optional + Result limit + **kwargs + Additional arguments + + Returns + ------- + pd.DataFrame + Query results + """ + import pandas as pd + + sql = self.compile(expr, params=params, limit=limit) + + with self._safe_raw_sql(sql) as cursor: + # Fetch column names and types + columns = [desc[0].lower() for desc in cursor.description] + + # Fetch all rows + rows = cursor.fetchall() + + # Convert to DataFrame + df = pd.DataFrame(rows, columns=columns) + + # Pandas handles most type conversions automatically + # Additional type conversions can be added here if needed + + return df + + def execute(self, expr: ir.Expr, **kwargs: Any) -> Any: + """ + Execute an Ibis expression. + + Parameters + ---------- + expr : ir.Expr + Expression to execute + **kwargs + Additional arguments + + Returns + ------- + Any + Execution result + """ + return self.to_pandas(expr, **kwargs) + + def _get_schema_using_query(self, query: str) -> sch.Schema: + """ + Get schema from a SQL query. + + Parameters + ---------- + query : str + SQL query + + Returns + ------- + sch.Schema + Query result schema + """ + with self._safe_raw_sql(query) as cursor: + if not cursor.description: + return sch.Schema({}) + + fields = {} + for col_desc in cursor.description: + col_name = col_desc[0].lower() + # Use a simple string type for now, can be enhanced later + fields[col_name] = parse_db2_type("VARCHAR") + + return sch.Schema(fields) + + def _register_in_memory_table(self, op: ops.InMemoryTable) -> None: + """ + Register an in-memory table. + + Parameters + ---------- + op : ops.InMemoryTable + In-memory table operation + """ + import pandas as pd + + # Create a temporary table from the in-memory data + name = op.name + data = op.data.to_frame() + + if isinstance(data, pd.DataFrame): + schema = sch.infer(data) + self.create_table(name, data, schema=schema, temp=True) + + +def connect( + database: str, + hostname: str = "localhost", + port: int = 50000, + username: str | None = None, + password: str | None = None, + schema: str | None = None, + **kwargs, +): + """ + Connect to a DB2 database. + + Parameters + ---------- + database : str + Database name to connect to + hostname : str, default "localhost" + Hostname of the DB2 server + port : int, default 50000 + Port number of the DB2 server + username : str, optional + Username for authentication + password : str, optional + Password for authentication + schema : str, optional + Default schema to use + **kwargs + Additional connection parameters + + Returns + ------- + Backend + An Ibis DB2 backend instance + + Examples + -------- + >>> import ibis + >>> con = ibis.db2.connect( + ... database="SAMPLE", + ... hostname="localhost", + ... port=50000, + ... username="db2inst1", + ... password="password" + ... ) # doctest: +SKIP + >>> con.list_tables() # doctest: +SKIP + ['EMPLOYEE', 'DEPARTMENT', 'PROJECT'] + """ + backend = Backend() + backend.do_connect( + database=database, + hostname=hostname, + port=port, + username=username, + password=password, + schema=schema, + **kwargs, + ) + return backend diff --git a/ibis/backends/db2/datatypes.py b/ibis/backends/db2/datatypes.py new file mode 100644 index 000000000000..04d3209b5441 --- /dev/null +++ b/ibis/backends/db2/datatypes.py @@ -0,0 +1,265 @@ +"""DB2 data type mappings for Ibis.""" + +from __future__ import annotations + +import ibis.expr.datatypes as dt +from ibis.formats.pandas import PandasData + +# Mapping from DB2 type names to Ibis types +DB2_TO_IBIS_TYPE = { + # Integer types + "SMALLINT": dt.int16, + "INTEGER": dt.int32, + "INT": dt.int32, + "BIGINT": dt.int64, + # Floating point types + "REAL": dt.float32, + "FLOAT": dt.float64, + "DOUBLE": dt.float64, + "DOUBLE PRECISION": dt.float64, + "DECFLOAT": dt.float64, + # Decimal types + "DECIMAL": dt.Decimal, + "NUMERIC": dt.Decimal, + "DEC": dt.Decimal, + # String types + "CHARACTER": dt.string, + "CHAR": dt.string, + "VARCHAR": dt.string, + "CHARACTER VARYING": dt.string, + "LONG VARCHAR": dt.string, + "CLOB": dt.string, + "CHAR () FOR BIT DATA": dt.binary, + "VARCHAR () FOR BIT DATA": dt.binary, + "LONG VARCHAR FOR BIT DATA": dt.binary, + # Binary types + "BINARY": dt.binary, + "VARBINARY": dt.binary, + "BLOB": dt.binary, + # Date/Time types + "DATE": dt.date, + "TIME": dt.time, + "TIMESTAMP": dt.timestamp, + # Boolean type + "BOOLEAN": dt.boolean, + # XML type + "XML": dt.string, + # Graphic types (for DBCS) + "GRAPHIC": dt.string, + "VARGRAPHIC": dt.string, + "LONG VARGRAPHIC": dt.string, + "DBCLOB": dt.string, +} + + +def parse_db2_type(type_string: str) -> dt.DataType: + """ + Parse a DB2 type string into an Ibis data type. + + Parameters + ---------- + type_string : str + DB2 type string (e.g., "VARCHAR(100)", "DECIMAL(10,2)") + + Returns + ------- + dt.DataType + Corresponding Ibis data type + + Examples + -------- + >>> parse_db2_type("VARCHAR(100)") + String(nullable=True) + >>> parse_db2_type("DECIMAL(10,2)") + Decimal(precision=10, scale=2, nullable=True) + """ + type_string = type_string.upper().strip() + + # Handle parameterized types + if "(" in type_string: + base_type = type_string.split("(")[0].strip() + params = type_string.split("(")[1].rstrip(")").split(",") + params = [p.strip() for p in params] + + if base_type in ("DECIMAL", "NUMERIC", "DEC"): + precision = int(params[0]) if params else 10 + scale = int(params[1]) if len(params) > 1 else 0 + return dt.Decimal(precision=precision, scale=scale, nullable=True) + elif base_type in ("VARCHAR", "CHARACTER VARYING", "CHAR", "CHARACTER"): + # Return string type regardless of length + return dt.string + elif base_type in ("VARBINARY", "BINARY"): + return dt.binary + elif base_type == "TIMESTAMP": + # DB2 timestamps can have precision + return dt.timestamp + else: + # For other parameterized types, use base type + return DB2_TO_IBIS_TYPE.get(base_type, dt.string) + + # Handle non-parameterized types + return DB2_TO_IBIS_TYPE.get(type_string, dt.string) + + +def ibis_type_to_db2_type(ibis_type: dt.DataType) -> str: + """ + Convert an Ibis data type to a DB2 type string. + + Parameters + ---------- + ibis_type : dt.DataType + Ibis data type + + Returns + ------- + str + DB2 type string + + Examples + -------- + >>> ibis_type_to_db2_type(dt.int32) + 'INTEGER' + >>> ibis_type_to_db2_type(dt.Decimal(10, 2)) + 'DECIMAL(10, 2)' + """ + if isinstance(ibis_type, dt.Int8): + return "SMALLINT" + elif isinstance(ibis_type, dt.Int16): + return "SMALLINT" + elif isinstance(ibis_type, dt.Int32): + return "INTEGER" + elif isinstance(ibis_type, dt.Int64): + return "BIGINT" + elif isinstance(ibis_type, dt.UInt8): + return "SMALLINT" + elif isinstance(ibis_type, dt.UInt16): + return "INTEGER" + elif isinstance(ibis_type, dt.UInt32): + return "BIGINT" + elif isinstance(ibis_type, dt.UInt64): + return "BIGINT" + elif isinstance(ibis_type, dt.Float32): + return "REAL" + elif isinstance(ibis_type, dt.Float64): + return "DOUBLE" + elif isinstance(ibis_type, dt.Decimal): + precision = ibis_type.precision or 10 + scale = ibis_type.scale or 0 + return f"DECIMAL({precision}, {scale})" + elif isinstance(ibis_type, dt.String): + return "VARCHAR(32672)" # Max VARCHAR length in DB2 + elif isinstance(ibis_type, dt.Binary): + return "VARBINARY(32672)" + elif isinstance(ibis_type, dt.Date): + return "DATE" + elif isinstance(ibis_type, dt.Time): + return "TIME" + elif isinstance(ibis_type, dt.Timestamp): + return "TIMESTAMP" + elif isinstance(ibis_type, dt.Boolean): + return "BOOLEAN" + elif isinstance(ibis_type, dt.JSON): + return "CLOB" # Store JSON as CLOB + elif isinstance(ibis_type, dt.UUID): + return "CHAR(36)" + elif isinstance(ibis_type, dt.Array): + # DB2 doesn't have native array type, use CLOB for JSON representation + return "CLOB" + elif isinstance(ibis_type, dt.Map): + # DB2 doesn't have native map type, use CLOB for JSON representation + return "CLOB" + elif isinstance(ibis_type, dt.Struct): + # DB2 doesn't have native struct type, use CLOB for JSON representation + return "CLOB" + else: + # Default to VARCHAR for unknown types + return "VARCHAR(32672)" + + +class DB2PandasData(PandasData): + """DB2-specific pandas data handler.""" + + @classmethod + def convert_Boolean(cls, s, dtype, pandas_type): + """Convert boolean columns.""" + # DB2 boolean values might come as 0/1 or True/False + if pandas_type is object: + return s.map({0: False, 1: True, "0": False, "1": True}).astype(bool) + return s.astype(bool) + + @classmethod + def convert_Timestamp(cls, s, dtype, pandas_type): + """Convert timestamp columns.""" + import pandas as pd + + # Handle DB2 timestamp format + if pandas_type is object: + return pd.to_datetime(s, errors="coerce") + return pd.to_datetime(s) + + @classmethod + def convert_Date(cls, s, dtype, pandas_type): + """Convert date columns.""" + import pandas as pd + + if pandas_type is object: + return pd.to_datetime(s, errors="coerce").dt.date + return pd.to_datetime(s).dt.date + + @classmethod + def convert_Time(cls, s, dtype, pandas_type): + """Convert time columns.""" + import pandas as pd + + if pandas_type is object: + return pd.to_datetime(s, errors="coerce").dt.time + return pd.to_datetime(s).dt.time + + +# Type code mappings for ibm_db_dbi +DB2_TYPE_CODES = { + 384: dt.date, # DATE + 388: dt.time, # TIME + 392: dt.timestamp, # TIMESTAMP + 448: dt.string, # VARCHAR + 452: dt.string, # CHAR + 456: dt.string, # LONG VARCHAR + 460: dt.string, # CLOB + 464: dt.binary, # BLOB + 468: dt.binary, # BINARY + 472: dt.binary, # VARBINARY + 480: dt.float64, # FLOAT + 484: dt.float64, # DOUBLE + 492: dt.Decimal, # DECIMAL + 496: dt.int32, # INTEGER + 500: dt.int16, # SMALLINT + 504: dt.int64, # BIGINT + 908: dt.boolean, # BOOLEAN +} + + +def type_code_to_ibis_type(type_code: int, precision: int = 0, scale: int = 0) -> dt.DataType: + """ + Convert DB2 type code to Ibis data type. + + Parameters + ---------- + type_code : int + DB2 type code from cursor description + precision : int, default 0 + Precision for decimal types + scale : int, default 0 + Scale for decimal types + + Returns + ------- + dt.DataType + Corresponding Ibis data type + """ + base_type = DB2_TYPE_CODES.get(type_code, dt.string) + + if type_code == 492 and precision > 0: # DECIMAL + return dt.Decimal(precision=precision, scale=scale, nullable=True) + + return base_type + diff --git a/ibis/backends/db2/tests/__init__.py b/ibis/backends/db2/tests/__init__.py new file mode 100644 index 000000000000..632649c82399 --- /dev/null +++ b/ibis/backends/db2/tests/__init__.py @@ -0,0 +1 @@ +"""Tests for ibis-db2 backend.""" diff --git a/ibis/backends/db2/tests/conftest.py b/ibis/backends/db2/tests/conftest.py new file mode 100644 index 000000000000..98da7ba68380 --- /dev/null +++ b/ibis/backends/db2/tests/conftest.py @@ -0,0 +1,175 @@ +"""Pytest configuration and fixtures for DB2 backend tests.""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING + +import pytest + +if TYPE_CHECKING: + pass + + +@pytest.fixture(scope="session") +def db2_config(): + """ + Get DB2 connection configuration from environment variables. + + Returns + ------- + dict + Connection configuration + """ + return { + "database": os.getenv("DB2_DATABASE", "SAMPLE"), + "hostname": os.getenv("DB2_HOSTNAME", "localhost"), + "port": int(os.getenv("DB2_PORT", "50000")), + "username": os.getenv("DB2_USERNAME", "db2inst1"), + "password": os.getenv("DB2_PASSWORD", "password"), + "schema": os.getenv("DB2_SCHEMA"), + } + + +@pytest.fixture(scope="session") +def con(db2_config): + """ + Create a DB2 backend connection for testing. + + Parameters + ---------- + db2_config : dict + Connection configuration + + Returns + ------- + Backend + Connected backend instance + """ + from ibis.backends import db2 + + try: + backend = db2.connect(**db2_config) + yield backend + backend.disconnect() + except Exception as e: + pytest.skip(f"Could not connect to DB2: {e}") + + +@pytest.fixture +def test_table_name(): + """Generate a unique test table name.""" + import uuid + + return f"test_table_{uuid.uuid4().hex[:8]}" + + +@pytest.fixture +def sample_dataframe(): + """Create a sample pandas DataFrame for testing.""" + import pandas as pd + + return pd.DataFrame( + { + "id": [1, 2, 3, 4, 5], + "name": ["Alice", "Bob", "Charlie", "David", "Eve"], + "age": [25, 30, 35, 40, 45], + "salary": [50000.0, 60000.0, 70000.0, 80000.0, 90000.0], + "is_active": [True, True, False, True, False], + } + ) + + +@pytest.fixture +def temp_table(con, test_table_name, sample_dataframe): + """ + Create a temporary test table. + + Parameters + ---------- + con : Backend + Backend connection + test_table_name : str + Table name + sample_dataframe : pd.DataFrame + Sample data + + Yields + ------ + str + Table name + """ + # Create table + con.create_table(test_table_name, sample_dataframe) + + yield test_table_name + + # Cleanup + try: + con.drop_table(test_table_name, force=True) + except Exception: + pass + + +@pytest.fixture +def alltypes_table(con): + """ + Create a table with all supported data types. + + Parameters + ---------- + con : Backend + Backend connection + + Yields + ------ + str + Table name + """ + from datetime import date, datetime + + import pandas as pd + + table_name = "test_alltypes" + + df = pd.DataFrame( + { + "int_col": [1, 2, 3], + "bigint_col": [1000000, 2000000, 3000000], + "float_col": [1.1, 2.2, 3.3], + "double_col": [1.111, 2.222, 3.333], + "string_col": ["a", "b", "c"], + "bool_col": [True, False, True], + "date_col": [date(2024, 1, 1), date(2024, 1, 2), date(2024, 1, 3)], + "timestamp_col": [ + datetime(2024, 1, 1, 12, 0, 0), + datetime(2024, 1, 2, 12, 0, 0), + datetime(2024, 1, 3, 12, 0, 0), + ], + } + ) + + try: + con.create_table(table_name, df, overwrite=True) + yield table_name + finally: + try: + con.drop_table(table_name, force=True) + except Exception: + pass + + +def pytest_configure(config): + """Configure pytest with custom markers.""" + config.addinivalue_line( + "markers", "integration: mark test as integration test (requires DB2)" + ) + config.addinivalue_line("markers", "slow: mark test as slow running") + + +def pytest_collection_modifyitems(config, items): + """Automatically mark tests based on their location.""" + for item in items: + # Mark all tests in test_integration.py as integration tests + if "test_integration" in str(item.fspath): + item.add_marker(pytest.mark.integration) diff --git a/ibis/backends/db2/tests/test_backend.py b/ibis/backends/db2/tests/test_backend.py new file mode 100644 index 000000000000..ad4ac26254be --- /dev/null +++ b/ibis/backends/db2/tests/test_backend.py @@ -0,0 +1,298 @@ +"""Tests for DB2 backend functionality.""" + +from __future__ import annotations + +import pytest + +from ibis.backends import db2 + + +class TestConnection: + """Tests for connection functionality.""" + + def test_connect_with_all_params(self, db2_config): + """Test connection with all parameters.""" + con = db2.connect(**db2_config) + assert con is not None + assert isinstance(con, db2.Backend) + con.disconnect() + + def test_connect_minimal_params(self): + """Test connection with minimal parameters.""" + # This will fail without a real DB2 instance + import ibis.common.exceptions as exc + + with pytest.raises(exc.OperationNotDefinedError): + db2.connect(database="SAMPLE") + + def test_disconnect(self, db2_config): + """Test disconnection.""" + con = db2.connect(**db2_config) + con.disconnect() + # After disconnect, connection should be None + assert con._connection is None + + +@pytest.mark.integration +class TestBackendOperations: + """Integration tests for backend operations.""" + + def test_list_tables(self, con): + """Test listing tables.""" + tables = con.list_tables() + assert isinstance(tables, list) + + def test_list_databases(self, con): + """Test listing databases/schemas.""" + databases = con.list_databases() + assert isinstance(databases, list) + assert len(databases) > 0 + + def test_current_database(self, con): + """Test getting current database.""" + current_db = con.current_database + assert isinstance(current_db, str) + assert len(current_db) > 0 + + def test_version(self, con): + """Test getting DB2 version.""" + version = con.version + assert isinstance(version, str) + assert len(version) > 0 + + +@pytest.mark.integration +class TestTableOperations: + """Integration tests for table operations.""" + + def test_create_table_from_dataframe(self, con, test_table_name, sample_dataframe): + """Test creating a table from a DataFrame.""" + table = con.create_table(test_table_name, sample_dataframe) + assert table is not None + + # Verify table exists + tables = con.list_tables() + assert test_table_name.upper() in [t.upper() for t in tables] + + # Cleanup + con.drop_table(test_table_name, force=True) + + def test_create_table_with_schema(self, con, test_table_name): + """Test creating a table with explicit schema.""" + import ibis.expr.datatypes as dt + import ibis.expr.schema as sch + + schema = sch.Schema( + { + "id": dt.int32, + "name": dt.string, + "value": dt.float64, + } + ) + + table = con.create_table(test_table_name, schema=schema) + assert table is not None + + # Verify schema + result_schema = con.get_schema(test_table_name) + assert "id" in result_schema + assert "name" in result_schema + assert "value" in result_schema + + # Cleanup + con.drop_table(test_table_name, force=True) + + def test_create_temp_table(self, con, test_table_name, sample_dataframe): + """Test creating a temporary table.""" + table = con.create_table(test_table_name, sample_dataframe, temp=True) + assert table is not None + + # Cleanup + con.drop_table(test_table_name, force=True) + + def test_create_table_overwrite(self, con, test_table_name, sample_dataframe): + """Test creating a table with overwrite.""" + # Create table first time + con.create_table(test_table_name, sample_dataframe) + + # Create again with overwrite + table = con.create_table(test_table_name, sample_dataframe, overwrite=True) + assert table is not None + + # Cleanup + con.drop_table(test_table_name, force=True) + + def test_drop_table(self, con, test_table_name, sample_dataframe): + """Test dropping a table.""" + con.create_table(test_table_name, sample_dataframe) + con.drop_table(test_table_name) + + # Verify table doesn't exist + tables = con.list_tables() + assert test_table_name.upper() not in [t.upper() for t in tables] + + def test_drop_table_force(self, con, test_table_name): + """Test dropping a non-existent table with force.""" + # Should not raise error + con.drop_table(test_table_name, force=True) + + def test_get_schema(self, con, temp_table): + """Test getting table schema.""" + schema = con.get_schema(temp_table) + assert "id" in schema + assert "name" in schema + assert "age" in schema + assert "salary" in schema + assert "is_active" in schema + + def test_table_expression(self, con, temp_table): + """Test creating a table expression.""" + table = con.table(temp_table) + assert table is not None + assert hasattr(table, "schema") + + +@pytest.mark.integration +class TestDataOperations: + """Integration tests for data operations.""" + + def test_insert_dataframe(self, con, test_table_name, sample_dataframe): + """Test inserting data from DataFrame.""" + import ibis.expr.schema as sch + + # Create empty table + schema = sch.infer(sample_dataframe) + con.create_table(test_table_name, schema=schema) + + # Insert data + con.insert(test_table_name, sample_dataframe) + + # Verify data + table = con.table(test_table_name) + result = table.execute() + assert len(result) == len(sample_dataframe) + + # Cleanup + con.drop_table(test_table_name, force=True) + + def test_insert_overwrite(self, con, temp_table, sample_dataframe): + """Test inserting with overwrite.""" + # Insert new data with overwrite + con.insert(temp_table, sample_dataframe, overwrite=True) + + # Verify data + table = con.table(temp_table) + result = table.execute() + assert len(result) == len(sample_dataframe) + + def test_execute_query(self, con, temp_table): + """Test executing a query.""" + table = con.table(temp_table) + result = table.filter(table.age > 30).execute() + assert len(result) > 0 + assert all(result["age"] > 30) + + def test_to_pandas(self, con, temp_table): + """Test converting to pandas DataFrame.""" + table = con.table(temp_table) + df = con.to_pandas(table) + assert df is not None + assert len(df) > 0 + + def test_to_pyarrow(self, con, temp_table): + """Test converting to PyArrow table.""" + table = con.table(temp_table) + arrow_table = con.to_pyarrow(table) + assert arrow_table is not None + assert len(arrow_table) > 0 + + +@pytest.mark.integration +class TestQueryOperations: + """Integration tests for query operations.""" + + def test_select(self, con, temp_table): + """Test SELECT operation.""" + table = con.table(temp_table) + result = table.select("name", "age").execute() + assert "name" in result.columns + assert "age" in result.columns + assert "salary" not in result.columns + + def test_filter(self, con, temp_table): + """Test WHERE/filter operation.""" + table = con.table(temp_table) + result = table.filter(table.age >= 35).execute() + assert all(result["age"] >= 35) + + def test_order_by(self, con, temp_table): + """Test ORDER BY operation.""" + table = con.table(temp_table) + result = table.order_by(table.age.desc()).execute() + ages = result["age"].tolist() + assert ages == sorted(ages, reverse=True) + + def test_limit(self, con, temp_table): + """Test LIMIT operation.""" + table = con.table(temp_table) + result = table.limit(3).execute() + assert len(result) == 3 + + def test_aggregate(self, con, temp_table): + """Test aggregation operations.""" + table = con.table(temp_table) + result = table.aggregate( + avg_age=table.age.mean(), + max_salary=table.salary.max(), + count=table.count(), + ).execute() + assert "avg_age" in result.columns + assert "max_salary" in result.columns + assert "count" in result.columns + + def test_group_by(self, con, temp_table): + """Test GROUP BY operation.""" + table = con.table(temp_table) + result = table.group_by("is_active").aggregate( + count=table.count(), + avg_age=table.age.mean(), + ).execute() + assert len(result) <= 2 # True and False + + +@pytest.mark.integration +class TestComplexQueries: + """Integration tests for complex queries.""" + + def test_join(self, con, test_table_name, sample_dataframe): + """Test JOIN operation.""" + # Create two tables + table1_name = f"{test_table_name}_1" + table2_name = f"{test_table_name}_2" + + con.create_table(table1_name, sample_dataframe) + con.create_table(table2_name, sample_dataframe) + + try: + table1 = con.table(table1_name) + table2 = con.table(table2_name) + + result = table1.join(table2, table1.id == table2.id).execute() + assert len(result) > 0 + finally: + con.drop_table(table1_name, force=True) + con.drop_table(table2_name, force=True) + + def test_union(self, con, temp_table): + """Test UNION operation.""" + table = con.table(temp_table) + result = table.union(table).execute() + # Union should have twice the rows + assert len(result) == len(table.execute()) * 2 + + def test_subquery(self, con, temp_table): + """Test subquery.""" + table = con.table(temp_table) + avg_age = table.age.mean().name("avg_age") + result = table.filter(table.age > avg_age).execute() + assert len(result) > 0 diff --git a/ibis/backends/db2/tests/test_datatypes.py b/ibis/backends/db2/tests/test_datatypes.py new file mode 100644 index 000000000000..9ab4738a7908 --- /dev/null +++ b/ibis/backends/db2/tests/test_datatypes.py @@ -0,0 +1,219 @@ +"""Tests for DB2 data type mappings.""" + +from __future__ import annotations + +import ibis.expr.datatypes as dt + +from ibis.backends.db2.datatypes import ( + ibis_type_to_db2_type, + parse_db2_type, + type_code_to_ibis_type, +) + + +class TestParseDB2Type: + """Tests for parse_db2_type function.""" + + def test_integer_types(self): + """Test parsing of integer types.""" + assert parse_db2_type("SMALLINT") == dt.int16 + assert parse_db2_type("INTEGER") == dt.int32 + assert parse_db2_type("INT") == dt.int32 + assert parse_db2_type("BIGINT") == dt.int64 + + def test_float_types(self): + """Test parsing of floating point types.""" + assert parse_db2_type("REAL") == dt.float32 + assert parse_db2_type("FLOAT") == dt.float64 + assert parse_db2_type("DOUBLE") == dt.float64 + assert parse_db2_type("DOUBLE PRECISION") == dt.float64 + + def test_decimal_types(self): + """Test parsing of decimal types.""" + result = parse_db2_type("DECIMAL(10,2)") + assert isinstance(result, dt.Decimal) + assert result.precision == 10 + assert result.scale == 2 + + result = parse_db2_type("NUMERIC(15,3)") + assert isinstance(result, dt.Decimal) + assert result.precision == 15 + assert result.scale == 3 + + def test_string_types(self): + """Test parsing of string types.""" + assert parse_db2_type("VARCHAR(100)") == dt.string + assert parse_db2_type("CHAR(10)") == dt.string + assert parse_db2_type("CLOB") == dt.string + assert parse_db2_type("CHARACTER VARYING(50)") == dt.string + + def test_binary_types(self): + """Test parsing of binary types.""" + assert parse_db2_type("BINARY(10)") == dt.binary + assert parse_db2_type("VARBINARY(100)") == dt.binary + assert parse_db2_type("BLOB") == dt.binary + + def test_datetime_types(self): + """Test parsing of date/time types.""" + assert parse_db2_type("DATE") == dt.date + assert parse_db2_type("TIME") == dt.time + assert parse_db2_type("TIMESTAMP") == dt.timestamp + assert parse_db2_type("TIMESTAMP(6)") == dt.timestamp + + def test_boolean_type(self): + """Test parsing of boolean type.""" + assert parse_db2_type("BOOLEAN") == dt.boolean + + def test_case_insensitive(self): + """Test that parsing is case-insensitive.""" + assert parse_db2_type("varchar(100)") == dt.string + assert parse_db2_type("INTEGER") == dt.int32 + assert parse_db2_type("integer") == dt.int32 + + def test_unknown_type(self): + """Test that unknown types default to string.""" + assert parse_db2_type("UNKNOWN_TYPE") == dt.string + + +class TestIbisTypeToDB2Type: + """Tests for ibis_type_to_db2_type function.""" + + def test_integer_types(self): + """Test conversion of integer types.""" + assert ibis_type_to_db2_type(dt.int8) == "SMALLINT" + assert ibis_type_to_db2_type(dt.int16) == "SMALLINT" + assert ibis_type_to_db2_type(dt.int32) == "INTEGER" + assert ibis_type_to_db2_type(dt.int64) == "BIGINT" + + def test_unsigned_integer_types(self): + """Test conversion of unsigned integer types.""" + assert ibis_type_to_db2_type(dt.uint8) == "SMALLINT" + assert ibis_type_to_db2_type(dt.uint16) == "INTEGER" + assert ibis_type_to_db2_type(dt.uint32) == "BIGINT" + assert ibis_type_to_db2_type(dt.uint64) == "BIGINT" + + def test_float_types(self): + """Test conversion of floating point types.""" + assert ibis_type_to_db2_type(dt.float32) == "REAL" + assert ibis_type_to_db2_type(dt.float64) == "DOUBLE" + + def test_decimal_types(self): + """Test conversion of decimal types.""" + assert ibis_type_to_db2_type(dt.Decimal(10, 2)) == "DECIMAL(10, 2)" + assert ibis_type_to_db2_type(dt.Decimal(15, 3)) == "DECIMAL(15, 3)" + assert ibis_type_to_db2_type(dt.Decimal()) == "DECIMAL(10, 0)" + + def test_string_type(self): + """Test conversion of string type.""" + assert ibis_type_to_db2_type(dt.string) == "VARCHAR(32672)" + + def test_binary_type(self): + """Test conversion of binary type.""" + assert ibis_type_to_db2_type(dt.binary) == "VARBINARY(32672)" + + def test_datetime_types(self): + """Test conversion of date/time types.""" + assert ibis_type_to_db2_type(dt.date) == "DATE" + assert ibis_type_to_db2_type(dt.time) == "TIME" + assert ibis_type_to_db2_type(dt.timestamp) == "TIMESTAMP" + + def test_boolean_type(self): + """Test conversion of boolean type.""" + assert ibis_type_to_db2_type(dt.boolean) == "BOOLEAN" + + def test_json_type(self): + """Test conversion of JSON type.""" + assert ibis_type_to_db2_type(dt.json) == "CLOB" + + def test_uuid_type(self): + """Test conversion of UUID type.""" + assert ibis_type_to_db2_type(dt.uuid) == "CHAR(36)" + + def test_complex_types(self): + """Test conversion of complex types.""" + # Array, Map, and Struct should be stored as CLOB (JSON) + assert ibis_type_to_db2_type(dt.Array(dt.int32)) == "CLOB" + assert ibis_type_to_db2_type(dt.Map(dt.string, dt.int32)) == "CLOB" + assert ibis_type_to_db2_type(dt.Struct({"a": dt.int32, "b": dt.string})) == "CLOB" + + +class TestTypeCodeToIbisType: + """Tests for type_code_to_ibis_type function.""" + + def test_date_type(self): + """Test conversion of DATE type code.""" + assert type_code_to_ibis_type(384) == dt.date + + def test_time_type(self): + """Test conversion of TIME type code.""" + assert type_code_to_ibis_type(388) == dt.time + + def test_timestamp_type(self): + """Test conversion of TIMESTAMP type code.""" + assert type_code_to_ibis_type(392) == dt.timestamp + + def test_string_types(self): + """Test conversion of string type codes.""" + assert type_code_to_ibis_type(448) == dt.string # VARCHAR + assert type_code_to_ibis_type(452) == dt.string # CHAR + assert type_code_to_ibis_type(460) == dt.string # CLOB + + def test_binary_types(self): + """Test conversion of binary type codes.""" + assert type_code_to_ibis_type(464) == dt.binary # BLOB + assert type_code_to_ibis_type(468) == dt.binary # BINARY + assert type_code_to_ibis_type(472) == dt.binary # VARBINARY + + def test_numeric_types(self): + """Test conversion of numeric type codes.""" + assert type_code_to_ibis_type(480) == dt.float64 # FLOAT + assert type_code_to_ibis_type(484) == dt.float64 # DOUBLE + assert type_code_to_ibis_type(496) == dt.int32 # INTEGER + assert type_code_to_ibis_type(500) == dt.int16 # SMALLINT + assert type_code_to_ibis_type(504) == dt.int64 # BIGINT + + def test_decimal_type_with_precision(self): + """Test conversion of DECIMAL type code with precision.""" + result = type_code_to_ibis_type(492, precision=10, scale=2) + assert isinstance(result, dt.Decimal) + assert result.precision == 10 + assert result.scale == 2 + + def test_boolean_type(self): + """Test conversion of BOOLEAN type code.""" + assert type_code_to_ibis_type(908) == dt.boolean + + def test_unknown_type_code(self): + """Test that unknown type codes default to string.""" + assert type_code_to_ibis_type(9999) == dt.string + + +class TestRoundTripConversion: + """Tests for round-trip type conversions.""" + + def test_integer_round_trip(self): + """Test round-trip conversion of integer types.""" + for ibis_type in [dt.int16, dt.int32, dt.int64]: + db2_type = ibis_type_to_db2_type(ibis_type) + result = parse_db2_type(db2_type) + assert result == ibis_type + + def test_float_round_trip(self): + """Test round-trip conversion of float types.""" + for ibis_type in [dt.float32, dt.float64]: + db2_type = ibis_type_to_db2_type(ibis_type) + result = parse_db2_type(db2_type) + assert result == ibis_type + + def test_datetime_round_trip(self): + """Test round-trip conversion of datetime types.""" + for ibis_type in [dt.date, dt.time, dt.timestamp]: + db2_type = ibis_type_to_db2_type(ibis_type) + result = parse_db2_type(db2_type) + assert result == ibis_type + + def test_boolean_round_trip(self): + """Test round-trip conversion of boolean type.""" + db2_type = ibis_type_to_db2_type(dt.boolean) + result = parse_db2_type(db2_type) + assert result == dt.boolean diff --git a/ibis/backends/sql/compilers/__init__.py b/ibis/backends/sql/compilers/__init__.py index 4a85c54594cc..48c6251b295f 100644 --- a/ibis/backends/sql/compilers/__init__.py +++ b/ibis/backends/sql/compilers/__init__.py @@ -6,6 +6,7 @@ "ClickHouseCompiler", "DataFusionCompiler", "DatabricksCompiler", + "DB2Compiler", "DruidCompiler", "DuckDBCompiler", "ExasolCompiler", @@ -28,6 +29,7 @@ from ibis.backends.sql.compilers.clickhouse import ClickHouseCompiler from ibis.backends.sql.compilers.databricks import DatabricksCompiler from ibis.backends.sql.compilers.datafusion import DataFusionCompiler +from ibis.backends.sql.compilers.db2 import DB2Compiler from ibis.backends.sql.compilers.druid import DruidCompiler from ibis.backends.sql.compilers.duckdb import DuckDBCompiler from ibis.backends.sql.compilers.exasol import ExasolCompiler diff --git a/ibis/backends/sql/compilers/db2.py b/ibis/backends/sql/compilers/db2.py new file mode 100644 index 000000000000..b8769e53c16a --- /dev/null +++ b/ibis/backends/sql/compilers/db2.py @@ -0,0 +1,525 @@ +"""DB2 SQL compiler for Ibis expressions.""" + +from __future__ import annotations + +import sqlglot as sg +import sqlglot.expressions as sge +from ibis.backends.sql.compilers.base import SQLGlotCompiler +from ibis.backends.sql.rewrites import ( + exclude_unsupported_window_frame_from_ops, + lower_sample, +) +from sqlglot import exp +from sqlglot.dialects.dialect import Dialect +from sqlglot.generator import Generator + + +from db2_sqlglot import Db2 as DB2BaseDialect + +# Import type mapper from local datatypes module +from ibis.backends.db2.datatypes import ibis_type_to_db2_type + + +class DB2Generator(Generator): + """Custom SQL generator for DB2.""" + + TRANSFORMS = { + **Generator.TRANSFORMS, + exp.DateAdd: lambda self, e: self.func( + "DATE_ADD", e.this, self.sql(e, "expression"), self.sql(e, "unit") + ), + exp.DateDiff: lambda self, e: self.func( + "DAYS_BETWEEN", e.this, self.sql(e, "expression") + ), + exp.StrPosition: lambda self, e: self.func( + "LOCATE", self.sql(e, "substr"), self.sql(e, "this") + ), + exp.GroupConcat: lambda self, e: self.func( + "LISTAGG", + e.this, + self.sql(e, "separator") if e.args.get("separator") else self.sql(exp.Literal.string(",")), + ), + } + + TYPE_MAPPING = { + **Generator.TYPE_MAPPING, + exp.DataType.Type.BOOLEAN: "BOOLEAN", + exp.DataType.Type.TINYINT: "SMALLINT", + exp.DataType.Type.SMALLINT: "SMALLINT", + exp.DataType.Type.INT: "INTEGER", + exp.DataType.Type.BIGINT: "BIGINT", + exp.DataType.Type.FLOAT: "REAL", + exp.DataType.Type.DOUBLE: "DOUBLE", + exp.DataType.Type.DECIMAL: "DECIMAL", + exp.DataType.Type.VARCHAR: "VARCHAR", + exp.DataType.Type.CHAR: "CHAR", + exp.DataType.Type.TEXT: "CLOB", + exp.DataType.Type.BINARY: "VARBINARY", + exp.DataType.Type.VARBINARY: "VARBINARY", + exp.DataType.Type.DATE: "DATE", + exp.DataType.Type.DATETIME: "TIMESTAMP", + exp.DataType.Type.TIMESTAMP: "TIMESTAMP", + exp.DataType.Type.TIMESTAMPTZ: "TIMESTAMP", + } + + def limit_sql(self, expression: exp.Limit, top: bool = False) -> str: + """Generate LIMIT clause for DB2.""" + if top: + return f"TOP {self.sql(expression, 'expression')}" + return f"FETCH FIRST {self.sql(expression, 'expression')} ROWS ONLY" + + def offset_sql(self, expression: exp.Offset) -> str: + """Generate OFFSET clause for DB2.""" + return f"OFFSET {self.sql(expression, 'expression')} ROWS" + + def fetch_sql(self, expression: exp.Fetch) -> str: + """Generate FETCH clause for DB2.""" + return f"FETCH FIRST {self.sql(expression, 'count')} ROWS ONLY" + + def select_sql(self, expression: exp.Select) -> str: + """Generate SELECT statement for DB2 with proper clause ordering.""" + import re + + # Check if we have both OFFSET and LIMIT + has_offset = expression.args.get("offset") is not None + has_limit = expression.args.get("limit") is not None + + if has_offset and has_limit: + # Temporarily remove both limit and offset to generate SQL without them + limit_expr = expression.args.pop("limit") + offset_expr = expression.args.pop("offset") + + # Generate SQL with everything except limit and offset + sql = super().select_sql(expression) + + # Restore limit and offset + expression.args["limit"] = limit_expr + expression.args["offset"] = offset_expr + + # Now manually add OFFSET and FETCH FIRST in the correct order + offset_val = self.sql(offset_expr, "expression") + limit_val = self.sql(limit_expr, "expression") + + sql = f"{sql} OFFSET {offset_val} ROWS FETCH FIRST {limit_val} ROWS ONLY" + else: + # Normal generation for cases without both OFFSET and LIMIT + sql = super().select_sql(expression) + + # Fix spacing issues in generated SQL + # DB2 requires proper spacing between clauses + sql = sql.replace("LASTFETCH", "LAST FETCH") + sql = sql.replace("DESCFETCH", "DESC FETCH") + sql = sql.replace("ASCFETCH", "ASC FETCH") + sql = sql.replace("ONLYOFFSET", "ONLY OFFSET") + sql = re.sub(r'(\d+)\s*ROWSFETCH', r'\1 ROWS FETCH', sql) + + # Ensure space before FETCH FIRST and OFFSET + sql = re.sub(r'(\S)FETCH\s+FIRST', r'\1 FETCH FIRST', sql) + sql = re.sub(r'(\S)OFFSET\s+', r'\1 OFFSET ', sql) + + # Remove ROWS BETWEEN clauses from window functions that don't support them + sql = re.sub( + r'(ROW_NUMBER\(\)|RANK\(\)|DENSE_RANK\(\)|PERCENT_RANK\(\)|CUME_DIST\(\)|NTILE\([^)]+\)|LAG\([^)]+\)|LEAD\([^)]+\)|FIRST_VALUE\([^)]+\)|LAST_VALUE\([^)]+\))\s+OVER\s+\(([^)]*?)\s+ROWS\s+BETWEEN[^)]+\)', + r'\1 OVER (\2)', + sql + ) + + return sql + + def tablesample_sql(self, expression: exp.TableSample) -> str: + """DB2 doesn't support TABLESAMPLE, return empty string.""" + return "" + + def cast_sql(self, expression: exp.Cast) -> str: + """Generate CAST expression for DB2.""" + return f"CAST({self.sql(expression, 'this')} AS {self.sql(expression, 'to')})" + + def trycast_sql(self, expression: exp.TryCast) -> str: + """DB2 doesn't have TRY_CAST, use regular CAST.""" + return self.cast_sql(expression) + + def boolean_sql(self, expression: exp.Boolean) -> str: + """Generate boolean literal for DB2.""" + return "TRUE" if expression.this else "FALSE" + + def concat_sql(self, expression: exp.Concat) -> str: + """Generate CONCAT expression for DB2.""" + # DB2 uses || operator or CONCAT function + return self.func("CONCAT", *expression.expressions) + + def substring_sql(self, expression: exp.Substring) -> str: + """Generate SUBSTRING expression for DB2.""" + # DB2 uses SUBSTR function + args = [expression.this] + if expression.args.get("start"): + args.append(expression.args["start"]) + if expression.args.get("length"): + args.append(expression.args["length"]) + return self.func("SUBSTR", *args) + + def dateadd_sql(self, expression: exp.DateAdd) -> str: + """Generate date addition for DB2.""" + unit = self.sql(expression, "unit") + value = self.sql(expression, "expression") + date = self.sql(expression, "this") + + # DB2 uses specific functions for date arithmetic + unit_map = { + "DAY": f"{date} + {value} DAYS", + "MONTH": f"{date} + {value} MONTHS", + "YEAR": f"{date} + {value} YEARS", + "HOUR": f"{date} + {value} HOURS", + "MINUTE": f"{date} + {value} MINUTES", + "SECOND": f"{date} + {value} SECONDS", + } + return unit_map.get(unit.upper(), f"{date} + {value} {unit}") + + def extract_sql(self, expression: exp.Extract) -> str: + """Generate EXTRACT expression for DB2.""" + # DB2 supports EXTRACT function + return f"EXTRACT({self.sql(expression, 'this')} FROM {self.sql(expression, 'expression')})" + + def regexp_like_sql(self, expression: exp.RegexpLike) -> str: + """Generate REGEXP_LIKE for DB2.""" + return self.func("REGEXP_LIKE", expression.this, expression.expression) + + def if_sql(self, expression: exp.If) -> str: + """Generate IF/CASE expression for DB2.""" + # DB2 uses CASE WHEN for conditional logic + return f"CASE WHEN {self.sql(expression, 'this')} THEN {self.sql(expression, 'true')} ELSE {self.sql(expression, 'false')} END" + + def div_sql(self, expression: exp.Div) -> str: + """Generate division for DB2.""" + # DB2 integer division needs special handling + return f"({self.sql(expression, 'this')} / {self.sql(expression, 'expression')})" + + def mod_sql(self, expression: exp.Mod) -> str: + """Generate modulo for DB2.""" + return self.func("MOD", expression.this, expression.expression) + + def log_sql(self, expression: exp.Log) -> str: + """Generate logarithm for DB2.""" + if expression.args.get("base"): + # DB2 LOG function with base + return self.func("LOG", expression.args["base"], expression.this) + # Natural log + return self.func("LN", expression.this) + + def sqrt_sql(self, expression: exp.Sqrt) -> str: + """Generate square root for DB2.""" + return self.func("SQRT", expression.this) + + def power_sql(self, expression: exp.Pow) -> str: + """Generate power function for DB2.""" + return self.func("POWER", expression.this, expression.expression) + + +class DB2Dialect(DB2BaseDialect): + """DB2 SQL dialect for SQLGlot, extending the base dialect from sqlglot-db2.""" + + class Generator(DB2Generator): + """Extended DB2 generator with Ibis-specific customizations.""" + + # Inherit from both the base DB2 generator and our custom generator + TYPE_MAPPING = { + **DB2BaseDialect.Generator.TYPE_MAPPING, + **DB2Generator.TYPE_MAPPING, + } + + TRANSFORMS = { + **DB2BaseDialect.Generator.TRANSFORMS, + **DB2Generator.TRANSFORMS, + } + + class Parser(DB2BaseDialect.Parser): + """Extended DB2 parser with Ibis-specific customizations.""" + + FUNCTIONS = { + **DB2BaseDialect.Parser.FUNCTIONS, + "LOCATE": exp.StrPosition.from_arg_list, + "LISTAGG": exp.GroupConcat.from_arg_list, + "DAYS_BETWEEN": exp.DateDiff.from_arg_list, + "REGEXP_LIKE": lambda args: exp.RegexpLike( + this=args[0], expression=args[1] + ), + } + + +class DB2Compiler(SQLGlotCompiler): + """SQL compiler for DB2 backend.""" + + __slots__ = () + + dialect = DB2Dialect + type_mapper = ibis_type_to_db2_type + + rewrites = ( + exclude_unsupported_window_frame_from_ops | lower_sample(), + *SQLGlotCompiler.rewrites, + ) + + @staticmethod + def _generate_groups(groups): + """Generate GROUP BY clause.""" + return groups + + def visit_Cast(self, op, *, arg, to, **kwargs): + """Visit a Cast operation.""" + # DB2 uses CAST syntax + return self.cast(arg, to) + + def visit_TryCast(self, op, *, arg, to, **kwargs): + """Visit a TryCast operation.""" + # DB2 doesn't have TRY_CAST, use CAST with error handling + return self.cast(arg, to) + + def visit_Sample(self, op, *, table, fraction, method, seed, **kwargs): + """Visit a Sample operation.""" + # DB2 doesn't have native TABLESAMPLE, use WHERE RAND() < fraction + if method == "row": + # Use RAND() function for sampling + condition = sge.LT( + this=sge.Anonymous(this="RAND", expressions=[]), + expression=sge.convert(fraction) + ) + return sge.Where(this=table, expression=condition) + return table + + def visit_StringContains(self, op, *, haystack, needle, **kwargs): + """Visit a StringContains operation.""" + # DB2 uses LOCATE function + return sge.GT( + this=sge.Anonymous( + this="LOCATE", + expressions=[needle, haystack] + ), + expression=sge.convert(0) + ) + + def visit_StringFind(self, op, *, haystack, needle, start, end, **kwargs): + """Visit a StringFind operation.""" + # DB2 uses LOCATE function + if start is not None: + return sge.Anonymous( + this="LOCATE", + expressions=[needle, haystack, start] + ) + return sge.Anonymous( + this="LOCATE", + expressions=[needle, haystack] + ) + + def visit_RegexSearch(self, op, *, arg, pattern, **kwargs): + """Visit a RegexSearch operation.""" + # DB2 uses REGEXP_LIKE function + return sge.Anonymous( + this="REGEXP_LIKE", + expressions=[arg, pattern] + ) + + def visit_RegexExtract(self, op, *, arg, pattern, index, **kwargs): + """Visit a RegexExtract operation.""" + # DB2 uses REGEXP_SUBSTR function + return sge.Anonymous( + this="REGEXP_SUBSTR", + expressions=[arg, pattern] + ) + + def visit_RegexReplace(self, op, *, arg, pattern, replacement, **kwargs): + """Visit a RegexReplace operation.""" + # DB2 uses REGEXP_REPLACE function + return sge.Anonymous( + this="REGEXP_REPLACE", + expressions=[arg, pattern, replacement] + ) + + def visit_StringSplit(self, op, *, arg, delimiter): + """Visit a StringSplit operation.""" + # DB2 doesn't have native string split, return as-is + # This would need to be handled at a higher level + return arg + + def visit_ArrayCollect(self, op, *, arg): + """Visit an ArrayCollect operation.""" + # DB2 uses LISTAGG for array aggregation + return sge.Anonymous( + this="LISTAGG", + expressions=[arg, sge.convert(",")] + ) + + def visit_Median(self, op, *, arg, where): + """Visit a Median operation.""" + # DB2 uses MEDIAN function + func = sge.Anonymous(this="MEDIAN", expressions=[arg]) + if where is not None: + return sge.Filter(this=func, expression=sge.Where(this=where)) + return func + + def visit_Mode(self, op, *, arg, where): + """Visit a Mode operation.""" + # DB2 doesn't have native MODE, use workaround with COUNT and GROUP BY + # This is a simplified version + return sge.Anonymous(this="MODE", expressions=[arg]) + + def visit_ArgMin(self, op, *, arg, key, where): + """Visit an ArgMin operation.""" + # DB2 uses MIN_BY or equivalent window function + return sge.Anonymous(this="MIN_BY", expressions=[arg, key]) + + def visit_ArgMax(self, op, *, arg, key, where): + """Visit an ArgMax operation.""" + # DB2 uses MAX_BY or equivalent window function + return sge.Anonymous(this="MAX_BY", expressions=[arg, key]) + + def visit_CountDistinct(self, op, *, arg, where): + """Visit a CountDistinct operation.""" + func = sge.Count(this=arg, distinct=True) + if where is not None: + return sge.Filter(this=func, expression=sge.Where(this=where)) + return func + + def visit_ApproxCountDistinct(self, op, *, arg, where): + """Visit an ApproxCountDistinct operation.""" + # DB2 doesn't have APPROX_COUNT_DISTINCT, use COUNT(DISTINCT) + return self.visit_CountDistinct(op, arg=arg, where=where) + + def visit_First(self, op, *, arg, where): + """Visit a First operation.""" + # DB2 uses FIRST_VALUE window function + return sge.Anonymous(this="FIRST_VALUE", expressions=[arg]) + + def visit_Last(self, op, *, arg, where): + """Visit a Last operation.""" + # DB2 uses LAST_VALUE window function + return sge.Anonymous(this="LAST_VALUE", expressions=[arg]) + + def visit_Lag(self, op, *, arg, offset, default): + """Visit a Lag operation.""" + # DB2 LAG function + expressions = [arg] + if offset is not None: + expressions.append(offset) + if default is not None: + expressions.append(default) + return sge.Anonymous(this="LAG", expressions=expressions) + + def visit_Lead(self, op, *, arg, offset, default): + """Visit a Lead operation.""" + # DB2 LEAD function + expressions = [arg] + if offset is not None: + expressions.append(offset) + if default is not None: + expressions.append(default) + return sge.Anonymous(this="LEAD", expressions=expressions) + + def visit_RowNumber(self, op): + """Visit a RowNumber operation.""" + # DB2 ROW_NUMBER function + return sge.Anonymous(this="ROW_NUMBER", expressions=[]) + + def visit_Rank(self, op): + """Visit a Rank operation.""" + # DB2 RANK function + return sge.Anonymous(this="RANK", expressions=[]) + + def visit_DenseRank(self, op): + """Visit a DenseRank operation.""" + # DB2 DENSE_RANK function + return sge.Anonymous(this="DENSE_RANK", expressions=[]) + + def visit_PercentRank(self, op): + """Visit a PercentRank operation.""" + # DB2 PERCENT_RANK function + return sge.Anonymous(this="PERCENT_RANK", expressions=[]) + + def visit_CumeDist(self, op): + """Visit a CumeDist operation.""" + # DB2 CUME_DIST function + return sge.Anonymous(this="CUME_DIST", expressions=[]) + + def visit_NTile(self, op, *, buckets): + """Visit an NTile operation.""" + # DB2 NTILE function + return sge.Anonymous(this="NTILE", expressions=[buckets]) + + def visit_DateTruncate(self, op, *, arg, unit): + """Visit a DateTruncate operation.""" + # DB2 uses TRUNC function for dates + unit_map = { + "Y": "YEAR", + "Q": "QUARTER", + "M": "MONTH", + "W": "WEEK", + "D": "DAY", + } + db2_unit = unit_map.get(unit, unit) + return sge.Anonymous( + this="TRUNC", + expressions=[arg, sge.convert(db2_unit)] + ) + + def visit_TimestampTruncate(self, op, *, arg, unit): + """Visit a TimestampTruncate operation.""" + # Similar to DateTruncate but for timestamps + return self.visit_DateTruncate(op, arg=arg, unit=unit) + + def visit_ExtractYear(self, op, *, arg): + """Visit an ExtractYear operation.""" + return sge.Anonymous(this="YEAR", expressions=[arg]) + + def visit_ExtractMonth(self, op, *, arg): + """Visit an ExtractMonth operation.""" + return sge.Anonymous(this="MONTH", expressions=[arg]) + + def visit_ExtractDay(self, op, *, arg): + """Visit an ExtractDay operation.""" + return sge.Anonymous(this="DAY", expressions=[arg]) + + def visit_ExtractHour(self, op, *, arg): + """Visit an ExtractHour operation.""" + return sge.Anonymous(this="HOUR", expressions=[arg]) + + def visit_ExtractMinute(self, op, *, arg): + """Visit an ExtractMinute operation.""" + return sge.Anonymous(this="MINUTE", expressions=[arg]) + + def visit_ExtractSecond(self, op, *, arg): + """Visit an ExtractSecond operation.""" + return sge.Anonymous(this="SECOND", expressions=[arg]) + + def visit_DayOfWeekIndex(self, op, *, arg): + """Visit a DayOfWeekIndex operation.""" + # DB2 uses DAYOFWEEK function (1=Sunday, 7=Saturday) + # Adjust to 0-based index (0=Monday) + return sge.Sub( + this=sge.Anonymous(this="DAYOFWEEK", expressions=[arg]), + expression=sge.convert(1) + ) + + def visit_DayOfWeekName(self, op, *, arg): + """Visit a DayOfWeekName operation.""" + # DB2 uses DAYNAME function + return sge.Anonymous(this="DAYNAME", expressions=[arg]) + + def visit_RandomScalar(self, op): + """Visit a RandomScalar operation.""" + # DB2 uses RAND() function + return sge.Anonymous(this="RAND", expressions=[]) + + def visit_RandomUUID(self, op): + """Visit a RandomUUID operation.""" + # DB2 doesn't have native UUID generation + # Use combination of functions to generate UUID-like string + return sge.Anonymous(this="GENERATE_UNIQUE", expressions=[]) + + def visit_Hash(self, op, *, arg): + """Visit a Hash operation.""" + # DB2 uses HASH function + return sge.Anonymous(this="HASH", expressions=[arg]) + + def visit_HashBytes(self, op, *, arg, how): + """Visit a HashBytes operation.""" + # DB2 uses HASH function with algorithm + return sge.Anonymous(this="HASH", expressions=[arg, sge.convert(how)]) + diff --git a/pyproject.toml b/pyproject.toml index b1890f7057cd..279e6d5f08bc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,7 +52,7 @@ dependencies = [ "atpublic>=2.3", "parsy>=2", "python-dateutil>=2.8.2", - "sqlglot>=23.4,!=26.32.0", + "sqlglot>=28.5.0", "toolz>=0.11", "typing-extensions>=4.3.0", "tzdata>=2022.7", # fallback time zone data on Windows @@ -96,6 +96,16 @@ clickhouse = [ "pandas>=1.5.3,<4", "rich>=12.4.4", ] +db2 = [ + "ibm_db>=3.0.0", + "db2-sqlglot-dialect>=1.0.0", + "pyarrow>=10.0.1", + "pyarrow-hotfix>=0.4", + "numpy>=1.23.2,<3", + "pandas>=1.5.3,<3", + "rich>=12.4.4", + "packaging>=21.3" +] databricks = [ "databricks-sql-connector>=4", "pyarrow>=10.0.1", @@ -332,6 +342,7 @@ bigquery = "ibis.backends.bigquery" clickhouse = "ibis.backends.clickhouse" databricks = "ibis.backends.databricks" datafusion = "ibis.backends.datafusion" +db2 = "ibis.backends.db2" druid = "ibis.backends.druid" duckdb = "ibis.backends.duckdb" exasol = "ibis.backends.exasol" diff --git a/requirements-local.txt b/requirements-local.txt new file mode 100644 index 000000000000..b8131cd1db2d --- /dev/null +++ b/requirements-local.txt @@ -0,0 +1,5 @@ +# Local dependencies for ibis-db2 +# Install with: pip install -r requirements-local.txt + +# SQLGlot DB2 dialect (local package) +-e /Users/amitkumar/Documents/Project/sqlglot-db2 \ No newline at end of file