Fix DuckDB outputs and test that the match SQLite

Author: zaneselvans

src/ferc_xbrl_extractor/cli.py

@@ -110,6 +110,7 @@ def write_to_sqlite(sqlite_engine: Engine, table_name: str, table_data: pd.DataF
 
 def write_to_duckdb(duckdb_path: str, table_name: str, table_data: pd.DataFrame):
     """Write one table to a duckdb database."""
+    table_data = table_data.reset_index()
     with duckdb.connect(duckdb_path) as duckdb_conn:
         duckdb_conn.execute(
             f"CREATE OR REPLACE TABLE {table_name} AS SELECT * FROM table_data"  # noqa: S608

src/ferc_xbrl_extractor/datapackage.py

@@ -120,7 +120,7 @@ def __hash__(self):
     "boolean": "boolean",
     "date": "string",
     "duration": "string",
-    "year": "datetime64[ns]",
+    "year": "Int64",
 }
 """
 Pandas data type by schema field type (Data Package `field.type`).
@@ -391,9 +391,7 @@ def __init__(self, schema: Schema, period_type: str):
         """Create FactTable and prepare for constructing dataframe."""
         self.schema = schema
         # Map column names to function to convert parsed values
-        self.columns = {
-            field.name: CONVERT_DTYPES[field.type_] for field in schema.fields
-        }
+        self.columns = {field.name: field.type_ for field in schema.fields}
         self.axes = [name for name in schema.primary_key if name.endswith("axis")]
         self.data_columns = [
             field.name
@@ -424,7 +422,7 @@ def construct_dataframe(self, instance: Instance) -> pd.DataFrame:
                 {
                     "c_id": fact.c_id,
                     "name": fact.name,
-                    "value": self.columns[fact.name](fact.value),
+                    "value": CONVERT_DTYPES[self.columns[fact.name]](fact.value),
                 }
                 for fact in raw_facts
             )
@@ -443,7 +441,14 @@ def construct_dataframe(self, instance: Instance) -> pd.DataFrame:
             )
         )
 
-        return contexts.join(facts).set_index(self.schema.primary_key).dropna(how="all")
+        return (
+            contexts.join(facts)
+            .astype(
+                {name: FIELD_TO_PANDAS[dtype] for name, dtype in self.columns.items()}
+            )
+            .set_index(self.schema.primary_key)
+            .dropna(how="all")
+        )
 
 
 class Datapackage(BaseModel):

tests/integration/console_scripts_test.py

@@ -3,9 +3,8 @@
 from importlib.metadata import entry_points
 
 import duckdb
+import pandas as pd
 import pytest
-from sqlalchemy import create_engine, text
-from sqlalchemy.engine import Connection
 
 # Obtain a list of all deployed entry point scripts to test:
 ENTRY_POINTS = [
@@ -26,22 +25,9 @@ def test_extractor_scripts(script_runner, ep):
     assert ret.success  # nosec: B101
 
 
-def _get_sqlite_tables(sqlite_conn) -> set[str]:
-    """Return set of all tables in SQLITE db."""
-    tables = sqlite_conn.execute(
-        text(
-            "SELECT name FROM sqlite_master "
-            "WHERE type='table' AND name NOT LIKE 'sqlite_%';"
-        )
-    ).fetchall()
-    return {table_name for (table_name,) in tables}
-
-
-def _get_duckdb_tables(duckdb_conn) -> set[str]:
+def _get_tables(conn) -> set[str]:
     """Return set of all tables in duckdb."""
-    tables = duckdb_conn.execute(
-        "SELECT table_name FROM information_schema.tables"
-    ).fetchall()
+    tables = conn.execute("SELECT table_name FROM information_schema.tables").fetchall()
     return {table_name for (table_name,) in tables}
 
 
@@ -50,9 +36,6 @@ def _find_empty_tables(db_conn, tables: set[str]) -> list[str]:
     empty_tables = []
     for table_name in tables:
         query = f"SELECT COUNT(*) FROM '{table_name}';"  # noqa: S608
-        if isinstance(db_conn, Connection):
-            query = text(query)
-
         if db_conn.execute(query).fetchone()[0] == 0:
             empty_tables.append(table_name)
     return empty_tables
@@ -95,15 +78,13 @@ def test_extract_example_filings(script_runner, tmp_path, test_dir):
     assert ret.success
 
     # Sanity check the sqlite/duckdb outputs
-    sqlite_uri = f"sqlite:///{sqlite_path.absolute()}"
-    sqlite_engine = create_engine(sqlite_uri)
     with (
-        sqlite_engine.begin() as sqlite_conn,
+        duckdb.connect(sqlite_path) as sqlite_conn,
         duckdb.connect(duckdb_path) as duckdb_conn,
     ):
         # Check for tables that only exist in either sqlite/duckdb but not both
-        sqlite_tables = _get_sqlite_tables(sqlite_conn)
-        duckdb_tables = _get_duckdb_tables(duckdb_conn)
+        sqlite_tables = _get_tables(sqlite_conn)
+        duckdb_tables = _get_tables(duckdb_conn)
 
         extra_sqlite_tables = sqlite_tables - duckdb_tables
         extra_duckdb_tables = duckdb_tables - sqlite_tables
@@ -118,6 +99,17 @@ def test_extract_example_filings(script_runner, tmp_path, test_dir):
         assert empty_sqlite_tables == []
         assert empty_duckdb_tables == []
 
+        # Check that tables are identical
+        # SQLite/duckdb have nuanced dtype differences, so ignore types
+        for table in sorted(sqlite_tables):
+            pd.testing.assert_frame_equal(
+                sqlite_conn.table(table).df(),
+                duckdb_conn.table(table).df(),
+                check_like=True,
+                check_dtype=False,
+                check_exact=True,
+            )
+
 
 @pytest.mark.script_launch_mode("inprocess")
 def test_extract_example_filings_bad_form(script_runner, tmp_path, test_dir):

tests/integration/datapackage_test.py

@@ -96,6 +96,7 @@ def _create_schema(instant=True, axes=None):
                     'cid_1,EID1,filing,2023-01-01T00:00:01,2021-01-01,2021-12-31,"value 1","value 2",\n'
                     'cid_4,EID1,filing,2023-01-01T00:00:01,2020-01-01,2020-12-31,"value 3","value 4",\n'
                 ),
+                dtype="string",
                 parse_dates=["publication_time"],
             ),
         ),
@@ -109,6 +110,7 @@ def _create_schema(instant=True, axes=None):
                     'cid_4,EID1,filing,2023-01-01T00:00:01,2020-01-01,2020-12-31,total,"value 3","value 4",\n'
                     'cid_5,EID1,filing,2023-01-01T00:00:01,2020-01-01,2020-12-31,"Dim 1 Value","value 9","value 10",\n'
                 ),
+                dtype="string",
                 parse_dates=["publication_time"],
             ),
         ),
@@ -121,6 +123,7 @@ def _create_schema(instant=True, axes=None):
                     'cid_2,EID1,filing,2023-01-01T00:00:01,2021-12-31,total,total,"value 5","value 6",\n'
                     'cid_3,EID1,filing,2023-01-01T00:00:01,2021-12-31,"Dim 1 Value","ferc:Dimension2Value","value 7","value 8",\n'
                 ),
+                dtype="string",
                 parse_dates=["publication_time"],
             ),
         ),