Repository: ibis-project/ibis

Issue:
feat: support amending/overriding reflected types in SQLite backend
I ran into a few problems trying to load data from [PyPI-Data](https://github.com/sethmlarson/pypi-data) that I'm not sure how best to solve:

- The database schemas use `STRING` instead of `TEXT` for string types (this has since been fixed, but [doesn't seem to be part of the latest release](https://github.com/sethmlarson/pypi-data/issues/12)). It would be nice if we could just force sqlalchemy to treat these columns as `TEXT` types rather than falling back to `NUMERIC`. The SQLAlchemy docs illustrate one way of [overriding reflected types](https://docs.sqlalchemy.org/en/14/core/reflection.html#overriding-reflected-columns), but like most things in SQLAlchemy there are probably other ways of handling this.
- The timestamp columns are using RFC3339 format (`YYYY-mm-ddTHH:MM:SS`) which isn't handled by SQLAlchemy's standard datetime parsing routine (it doesn't like the `T` delimiter). The way around this in SQLAlchemy is to pass [additional options to the DATETIME type constructor](https://docs.sqlalchemy.org/en/14/dialects/sqlite.html#sqlalchemy.dialects.sqlite.DATETIME), but ibis currently AFAICT doesn't expose this functionality.

I'm honestly not sure what a good API here would be.

Brainstorming a bit, perhaps additional arguments to the `ibis.sqlite.connect` providing:

- A type mapping from SQLITE type (e.g. `TIMESTAMP`) to ibis type (`typemap={...}`?). This would let us map `{"STRING": ibis.expr.datatypes.String()}`
- A datetime format string in strftime syntax to use when parsing `TIMESTAMP` columns to datetime objects (`timestamp_format=...`?). This assumes the database uses the same datetime format for all tables, but hopefully that's the case.

A few things to like about this approach:

- It doesn't expose that SQLAlchemy is used to the user anywhere. Everything is a mapping between SQLite concepts and ibis itself. This would give us the freedom to drop SQLAlchemy if we wanted without forcing an API change.
- It doesn't require listing the full table schemas for all poorly typed columns - the user only needs to pass in overrides for the specific SQLite type names that require it.

For my use case the connect call might look like:

```python
con = ibis.sqlite.connect("pypi.db", typemap={"STRING": ibis.expr.datatypes.String()}, timestamp_format="%Y-%m-%dT%H:%M:%S")
```

There's probably a better way of handling this though.


Discussion:
Honestly

```python
con = ibis.sqlite.connect("pypi.db", typemap={"STRING": ibis.expr.datatypes.String()}, timestamp_format="%Y-%m-%dT%H:%M:%S")
```

is pretty concise. The only thing I might change is the ibis type reference:

```python
con = ibis.sqlite.connect("pypi.db", typemap={"STRING": dt.string}, timestamp_format="%Y-%m-%dT%H:%M:%S")
```
Looking into this a bit more - I wonder if we should always modify the `DATETIME` type used by sqlalchemy's sqlite dialect to support parsing more of the [time formats](https://sqlite.org/lang_datefunc.html#time_values) sqlite supports natively (by changing the `regexp` configuration). This wouldn't change execution, and shouldn't be a breaking change since it would just strictly error less often when loading results with a `TIMESTAMP` column.

The `typemap` option would likely still be useful though.

Resolving patch:
diff --git a/ibis/backends/sqlite/__init__.py b/ibis/backends/sqlite/__init__.py
--- a/ibis/backends/sqlite/__init__.py
+++ b/ibis/backends/sqlite/__init__.py
@@ -26,11 +26,13 @@
 
 if TYPE_CHECKING:
     import ibis.expr.types as ir
+    import ibis.expr.datatypes as dt
 
 from ibis.backends.base import Database
-from ibis.backends.base.sql.alchemy import BaseAlchemyBackend
+from ibis.backends.base.sql.alchemy import BaseAlchemyBackend, to_sqla_type
 from ibis.backends.sqlite import udf
 from ibis.backends.sqlite.compiler import SQLiteCompiler
+from ibis.expr.schema import datatype
 
 
 def to_datetime(value: str | None) -> datetime.datetime | None:
@@ -86,6 +88,7 @@ def do_connect(
         self,
         database: str | Path | None = None,
         path: str | Path | None = None,
+        type_map: dict[str, str | dt.DataType] | None = None,
     ) -> None:
         """Create an Ibis client connected to a SQLite database.
 
@@ -97,6 +100,10 @@ def do_connect(
             File path to the SQLite database file. If `None`, creates an
             in-memory transient database and you can use attach() to add more
             files
+        type_map
+            An optional mapping from a string name of a SQLite "type" to the
+            corresponding ibis DataType that it represents. This can be used
+            to override schema inference for a given SQLite database.
 
         Examples
         --------
@@ -115,6 +122,24 @@ def do_connect(
             f"sqlite:///{database if database is not None else ':memory:'}"
         )
 
+        if type_map:
+            # Patch out ischema_names for the instantiated dialect. This
+            # attribute is required for all SQLAlchemy dialects, but as no
+            # public way of modifying it for a given dialect. Patching seems
+            # easier than subclassing the builtin SQLite dialect, and achieves
+            # the same desired behavior.
+            def _to_ischema_val(t):
+                sa_type = to_sqla_type(datatype(t))
+                if isinstance(sa_type, sa.types.TypeEngine):
+                    # SQLAlchemy expects a callable here, rather than an
+                    # instance. Use a lambda to work around this.
+                    return lambda: sa_type
+                return sa_type
+
+            overrides = {k: _to_ischema_val(v) for k, v in type_map.items()}
+            engine.dialect.ischema_names = engine.dialect.ischema_names.copy()
+            engine.dialect.ischema_names.update(overrides)
+
         sqlite3.register_adapter(pd.Timestamp, lambda value: value.isoformat())
 
         @sa.event.listens_for(engine, "connect")


Test patch:
diff --git a/ibis/backends/sqlite/tests/test_types.py b/ibis/backends/sqlite/tests/test_types.py
--- a/ibis/backends/sqlite/tests/test_types.py
+++ b/ibis/backends/sqlite/tests/test_types.py
@@ -36,6 +36,7 @@ def db(tmp_path_factory):
     con = sqlite3.connect(path)
     con.execute("CREATE TABLE timestamps (ts TIMESTAMP)")
     con.execute("CREATE TABLE timestamps_tz (ts TIMESTAMP)")
+    con.execute("CREATE TABLE weird (str_col STRING, date_col ITSADATE)")
     with con:
         con.executemany(
             "INSERT INTO timestamps VALUES (?)", [(t,) for t in TIMESTAMPS]
@@ -44,6 +45,15 @@ def db(tmp_path_factory):
             "INSERT INTO timestamps_tz VALUES (?)",
             [(t,) for t in TIMESTAMPS_TZ],
         )
+        con.executemany(
+            "INSERT INTO weird VALUES (?, ?)",
+            [
+                ("a", "2022-01-01"),
+                ("b", "2022-01-02"),
+                ("c", "2022-01-03"),
+                ("d", "2022-01-04"),
+            ],
+        )
     con.close()
     return path
 
@@ -59,3 +69,20 @@ def test_timestamps(db, table, data):
     res = t.ts.execute()
     sol = pd.Series([to_datetime(s) for s in data]).dt.tz_localize(None)
     assert res.equals(sol)
+
+
+def test_type_map(db):
+    con = ibis.sqlite.connect(
+        db, type_map={"STRING": dt.string, "ITSADATE": "date"}
+    )
+    t = con.tables.weird
+    expected_schema = ibis.schema({"str_col": "string", "date_col": "date"})
+    assert t.schema() == expected_schema
+    res = t.filter(t.str_col == "a").execute()
+    sol = pd.DataFrame(
+        {
+            "str_col": ["a"],
+            "date_col": pd.Series(["2022-01-01"], dtype="M8[ns]"),
+        }
+    )
+    assert res.equals(sol)