chg: use Duckdb Python API for GBIF backbone

Author: frafra

src/datasync/gbif_backbone.py

@@ -1,20 +1,17 @@
-from importlib.resources import files
-
-import duckdb
 import fsspec
 import typer
+from duckdb import (
+    DuckDBPyConnection,
+    StarExpression,
+    connect,
+)
 
-from .libs.helpers import pa_read_tsv
+from .libs.helpers import DuckDBAtomicTransaction, create_fts_index, function_expression
 from .settings import (
     env,
     log,
 )
 
-resources = files(__package__).joinpath("gbif_backbone")
-taxon_sql = resources.joinpath("taxon.sql").read_text()
-vernacularname_sql = resources.joinpath("vernacularname.sql").read_text()
-
-
 log.debug("Importing GBIF Backbone settings")
 
 
@@ -29,23 +26,51 @@
 app = typer.Typer(help="export GBIF Backbone data to DuckDB database")
 
 
-def import_taxon(conn, archive):
+def import_taxon(conn: DuckDBPyConnection, archive):
     log.debug("Importing Taxon.tsv")
-    pa_taxon = pa_read_tsv(archive, "Taxon.tsv")  # noqa: F841
-    conn.execute(taxon_sql)
+    conn.sql("DROP TABLE IF EXISTS taxon")
+    conn.from_csv_auto(archive.open("Taxon.tsv")).to_table("taxon")
+    conn.sql("CREATE INDEX taxon_taxonid ON taxon (taxonid)")
+    conn.sql(
+        create_fts_index(
+            "taxon",
+            "taxonID",
+            "canonicalName",
+            overwrite=True,
+        )
+    )
 
 
-def import_vernacular_name(conn, archive):
+def import_vernacular_name(conn: DuckDBPyConnection, archive):
     log.debug("Importing VernacularName.tsv")
-    pa_vernacular_names = pa_read_tsv(archive, "VernacularName.tsv")  # noqa: F841
-    conn.execute(vernacularname_sql)
+    conn.sql("DROP TABLE IF EXISTS vernacular_name")
+    generate_vernacular_id = function_expression(
+        "concat_ws",
+        "|",
+        "taxonID",
+        "language",
+        "vernacularName",
+    )
+    (
+        conn.from_csv_auto(archive.open("VernacularName.tsv"))
+        .select(StarExpression(), generate_vernacular_id.alias("vernacularID"))
+        .to_table("vernacular_name")
+    )
+    conn.sql("CREATE INDEX vernacular_name_taxonid ON vernacular_name (taxonid)")
+    conn.sql("CREATE INDEX vernacular_name_language ON vernacular_name (language)")
+    conn.sql(
+        create_fts_index(
+            "vernacular_name", "vernacularID", "vernacularName", overwrite=True
+        )
+    )
 
 
 @app.command()
 def import_all():
     """Import GBIF Backbone data into a DuckDB database."""
     archive = fsspec.filesystem("zip", fo=GBIF_BACKBONE_URL.geturl(), mode="r")
-    with duckdb.connect(GBIF_BACKBONE_DUCKDB_NAME) as conn:
-        import_taxon(conn, archive)
-        import_vernacular_name(conn, archive)
+    with connect(GBIF_BACKBONE_DUCKDB_NAME) as conn:
+        with DuckDBAtomicTransaction(conn):
+            import_taxon(conn, archive)
+            import_vernacular_name(conn, archive)
     log.info("GBIF Backbone data imported successfully")

src/datasync/gbif_backbone/taxon.sql

@@ -1,7 +1,7 @@
 import re
 
+from duckdb import ConstantExpression, DuckDBPyConnection, FunctionExpression
 from lxml import etree
-from pyarrow import csv
 
 PARSER = etree.XMLParser(resolve_entities=False)
 
@@ -26,13 +26,61 @@ def get_anytext(bag: str) -> str:
         )
 
 
-TSV_PARSE_OPTIONS = csv.ParseOptions(
-    delimiter="\t", quote_char=False, double_quote=False
-)
+class DuckDBAtomicTransaction:
+    def __init__(self, conn: DuckDBPyConnection):
+        self.conn = conn
 
+    def __enter__(self):
+        self.conn.begin()
+        return self.conn
 
-def pa_read_tsv(archive, filename):
-    return csv.read_csv(
-        archive.open(filename),
-        parse_options=TSV_PARSE_OPTIONS,
+    def __exit__(self, exc_type, exc_value, traceback):
+        if exc_type is None:
+            self.conn.commit()
+        else:
+            self.conn.rollback()
+
+
+def parameter_expression(value):
+    return (
+        str(FunctionExpression("", ConstantExpression("").alias(value)))
+        .removeprefix("(")
+        .removesuffix(" := '')")
+    )
+
+
+def identity_expression(value):
+    """SQLIdentifier is not exposed by DuckDB https://github.com/duckdb/duckdb/discussions/20425"""
+    return (
+        str(FunctionExpression("", ConstantExpression(value).alias("")))
+        .removeprefix("(")
+        .removesuffix(")")
+    )
+
+
+def named_argument_expression(name, value):
+    return parameter_expression(name) + " = " + identity_expression(value)
+
+
+def function_expression(name, *args):
+    expressions = [ConstantExpression(arg) for arg in args]
+    function = FunctionExpression(name, *expressions)
+    return function
+
+
+def function_expression_with_named_arguments(name, *args, **kwargs):
+    function = function_expression(name, *args)
+    named_arguments = []
+    for key, value in kwargs.items():
+        named_argument = named_argument_expression(key, value)
+        named_arguments.append(named_argument)
+    return str(function)[:-1] + ", " + ", ".join(named_arguments) + ")"
+
+
+def create_fts_index(*args, **kwargs):
+    expression = function_expression_with_named_arguments(
+        "create_fts_index",
+        *args,
+        **kwargs,
     )
+    return f"PRAGMA {expression}"