ftmq 4.1.0__tar.gz → 4.3.1__tar.gz

{ftmq-4.1.0 → ftmq-4.3.1}/PKG-INFO

@@ -1,8 +1,10 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: ftmq
-Version: 4.1.0
+Version: 4.3.1
 Summary: followthemoney query dsl and io helpers
 License: AGPLv3+
+License-File: LICENSE
+License-File: NOTICE
 Author: Simon Wörpel
 Author-email: simon.woerpel@pm.me
 Requires-Python: >=3.11,<4
@@ -19,24 +21,25 @@ Provides-Extra: postgres
 Provides-Extra: redis
 Provides-Extra: sql
 Requires-Dist: alephclient (>=2.6.0,<3.0.0) ; extra == "aleph"
-Requires-Dist: anystore (>=0.3.9,<0.4.0)
+Requires-Dist: anystore (>=0.4.0,<0.5.0)
 Requires-Dist: click (>=8.2.1,<9.0.0)
 Requires-Dist: click-default-group (>=1.2.4,<2.0.0)
-Requires-Dist: deltalake (>=1.1.2,<2.0.0) ; extra == "lake"
-Requires-Dist: duckdb (>=1.3.2,<2.0.0) ; extra == "lake"
+Requires-Dist: deltalake (>=1.2.1,<2.0.0) ; extra == "lake"
+Requires-Dist: duckdb (>=1.4.1,<2.0.0) ; extra == "lake"
 Requires-Dist: fakeredis (>=2.26.2,<3.0.0) ; extra == "redis"
-Requires-Dist: followthemoney (>=4.1.1,<5.0.0)
+Requires-Dist: followthemoney (>=4.3.2,<5.0.0)
 Requires-Dist: furl (>=2.1.4,<3.0.0) ; extra == "aleph"
-Requires-Dist: nomenklatura (>=4.1.0,<5.0.0)
+Requires-Dist: nomenklatura (>=4.1.10,<5.0.0)
 Requires-Dist: orjson (>=3.10.18,<4.0.0)
-Requires-Dist: pandas (>=2.3.1,<3.0.0) ; extra == "lake"
+Requires-Dist: pandas (>=2.3.3,<3.0.0) ; extra == "lake"
 Requires-Dist: plyvel (>=1.5.1,<2.0.0) ; extra == "level"
-Requires-Dist: psycopg2 (>=2.9.10,<3.0.0) ; extra == "postgres"
+Requires-Dist: psycopg[pool] (>=3.2.9,<4.0.0) ; extra == "postgres"
 Requires-Dist: pyarrow (>=21.0.0,<22.0.0) ; extra == "lake"
 Requires-Dist: pycountry (>=24.6.1,<25.0.0)
 Requires-Dist: pydantic (>=2.11.3,<3.0.0)
 Requires-Dist: pyicu (>=2.15.2,<3.0.0)
 Requires-Dist: redis (>=5.2.1,<6.0.0) ; extra == "redis"
+Requires-Dist: rigour (>=1.4.0,<2.0.0)
 Requires-Dist: sqlalchemy (>=2.0.36,<3.0.0) ; extra == "postgres"
 Requires-Dist: sqlalchemy (>=2.0.36,<3.0.0) ; extra == "sql"
 Project-URL: Documentation, https://docs.investigraph.dev/lib/ftmq

{ftmq-4.1.0 → ftmq-4.3.1}/ftmq/__init__.py

@@ -2,7 +2,7 @@ from ftmq.io import smart_read_proxies, smart_write_proxies
 from ftmq.query import Query
 from ftmq.util import make_entity
 
-__version__ = "4.1.0"
+__version__ = "4.3.1"
 __all__ = [
     "smart_read_proxies",
     "smart_write_proxies",

{ftmq-4.1.0 → ftmq-4.3.1}/ftmq/cli.py

@@ -1,3 +1,5 @@
+from datetime import datetime
+
 import click
 from anystore.io import smart_write, smart_write_json, smart_write_model
 from click_default_group import DefaultGroup
@@ -11,6 +13,9 @@ from ftmq.model.dataset import Catalog, Dataset
 from ftmq.model.stats import Collector
 from ftmq.query import Query
 from ftmq.store import get_store
+from ftmq.store.fragments import get_fragments
+from ftmq.store.fragments import get_store as get_fragments_store
+from ftmq.store.fragments.settings import Settings as FragmentsSettings
 from ftmq.util import apply_dataset, parse_unknown_filters
 
 log = get_logger(__name__)
@@ -311,6 +316,82 @@ def store_iterate(
     smart_write_proxies(output_uri, store.iterate())
 
 
+@cli.group()
+def fragments():
+    pass
+
+
+fragments_settings = FragmentsSettings()
+
+
+@fragments.command("list-datasets")
+@click.option(
+    "-i",
+    "--input-uri",
+    default=fragments_settings.database_uri,
+    show_default=True,
+    help="input file or uri",
+)
+@click.option(
+    "-o", "--output-uri", default="-", show_default=True, help="output file or uri"
+)
+def fragments_list_datasets(
+    input_uri: str = fragments_settings.database_uri,
+    output_uri: str = "-",
+):
+    """
+    List datasets within a fragments store
+    """
+    store = get_fragments_store(input_uri)
+    datasets = [ds.name for ds in store.all()]
+    smart_write(output_uri, "\n".join(datasets).encode() + b"\n")
+
+
+@fragments.command("iterate")
+@click.option(
+    "-i",
+    "--input-uri",
+    default=fragments_settings.database_uri,
+    show_default=True,
+    help="fragments store input uri",
+)
+@click.option(
+    "-o", "--output-uri", default="-", show_default=True, help="output file or uri"
+)
+@click.option("-d", "--dataset", required=True, help="Dataset name to iterate")
+@click.option("-s", "--schema", default=None, help="Filter by schema")
+@click.option(
+    "--since",
+    default=None,
+    help="Filter by timestamp (since), ISO format: YYYY-MM-DDTHH:MM:SS",
+)
+@click.option(
+    "--until",
+    default=None,
+    help="Filter by timestamp (until), ISO format: YYYY-MM-DDTHH:MM:SS",
+)
+def fragments_iterate(
+    input_uri: str = fragments_settings.database_uri,
+    output_uri: str = "-",
+    dataset: str = None,
+    schema: str | None = None,
+    since: str | None = None,
+    until: str | None = None,
+):
+    """
+    Iterate all entities from a fragments dataset
+    """
+    fragments = get_fragments(dataset, database_uri=input_uri)
+
+    # Parse timestamp strings to datetime objects
+    since_dt = datetime.fromisoformat(since) if since else None
+    until_dt = datetime.fromisoformat(until) if until else None
+
+    smart_write_proxies(
+        output_uri, fragments.iterate(schema=schema, since=since_dt, until=until_dt)
+    )
+
+
 @cli.command("aggregate")
 @click.option(
     "-i", "--input-uri", default="-", show_default=True, help="input file or uri"

{ftmq-4.1.0 → ftmq-4.3.1}/ftmq/filters.py

@@ -125,6 +125,21 @@ class DatasetFilter(BaseFilter):
         return False
 
 
+class OriginFilter(BaseFilter):
+    key = "origin"
+
+    def apply(self, entity: Entity) -> bool:
+        if not hasattr(entity, "context"):
+            return False
+        origins = ensure_list(entity.context.get("origin"))
+        if self.comparator == Lookup.EQUALS:
+            return self.value in origins
+        for value in origins:
+            if self.lookup.apply(value):
+                return True
+        return False
+
+
 class SchemaFilter(BaseFilter):
     key = "schema"
 
@@ -232,4 +247,5 @@ FILTERS = {
     "reverse": ReverseFilter,
     "entity_id": EntityIdFilter,
     "canonical_id": CanonicalIdFilter,
+    "origin": OriginFilter,
 }

{ftmq-4.1.0 → ftmq-4.3.1}/ftmq/io.py

@@ -127,8 +127,6 @@ def smart_write_proxies(
             for proxy in proxies:
                 ix += 1
                 bulk.add_entity(proxy)
-                if ix % 1_000 == 0:
-                    log.info("Writing proxy %d ..." % ix)
         return ix
 
     with smart_open(uri, mode=mode) as fh:

{ftmq-4.1.0 → ftmq-4.3.1}/ftmq/model/dataset.py

@@ -19,7 +19,7 @@ class Dataset(BaseModel, _DatasetModel):
     maintainer: DataPublisher | None = None
     stats: DatasetStats = DatasetStats()
     git_repo: AnyUrl | None = None
-    content_type: ContentType | None = "structured"
+    content_type: ContentType | None = None
     uri: str | None = None
 
     def iterate(self) -> Entities:

{ftmq-4.1.0 → ftmq-4.3.1}/ftmq/query.py

@@ -4,6 +4,7 @@ from typing import Any, Self, TypeVar
 
 from banal import ensure_list, is_listish, is_mapping
 from followthemoney import registry
+from sqlalchemy import Table
 
 from ftmq.aggregations import Aggregation, Aggregator
 from ftmq.enums import Aggregations, Properties
@@ -12,6 +13,7 @@ from ftmq.filters import (
     DatasetFilter,
     F,
     IdFilter,
+    OriginFilter,
     PropertyFilter,
     ReverseFilter,
     SchemaFilter,
@@ -57,12 +59,14 @@ class Query:
         aggregator: Aggregator | None = None,
         sort: Sort | None = None,
         slice: Slice | None = None,
+        table: Table | None = None,
     ):
         self.filters = set(ensure_list(filters))
         self.aggregations = set(ensure_list(aggregations))
         self.aggregator = aggregator
         self.sort = sort
         self.slice = slice
+        self.table = table
 
     def __getitem__(self, value: Any) -> Self:
         """
@@ -213,6 +217,23 @@ class Query:
             names.update(ensure_list([s.name for s in f.schemata]))
         return names
 
+    @property
+    def origins(self) -> set[OriginFilter]:
+        """
+        The current filtered origins
+        """
+        return {f for f in self.filters if isinstance(f, OriginFilter)}
+
+    @property
+    def origin_names(self) -> set[str]:
+        """
+        The names of the current filtered origins
+        """
+        names = set()
+        for f in self.origins:
+            names.update(ensure_list(f.value))
+        return names
+
     @property
     def countries(self) -> set[str]:
         """

{ftmq-4.1.0 → ftmq-4.3.1}/ftmq/sql.py

@@ -32,7 +32,7 @@ from ftmq.enums import (
 from ftmq.filters import F
 
 if TYPE_CHECKING:
-    from ftmq.query import Q
+    from ftmq.query import Query
 
 
 Field: TypeAlias = Properties | PropertyTypes | Fields
@@ -50,10 +50,13 @@ class Sql:
         Comparators.lte: "__le__",
     }
 
-    def __init__(self, q: "Q") -> None:
+    def __init__(self, q: "Query") -> None:
         self.q = q
         self.metadata = MetaData()
-        self.table = make_statement_table(self.metadata)
+        if q.table is None:
+            self.table = make_statement_table(self.metadata)
+        else:
+            self.table = q.table
         self.META_COLUMNS = {
             "id": self.table.c.canonical_id,
             "dataset": self.table.c.dataset,
@@ -92,6 +95,13 @@ class Sql:
                     for f in sorted(self.q.schemata)
                 )
             )
+        if self.q.origins:
+            clauses.append(
+                or_(
+                    self.get_expression(self.table.c.origin, f)
+                    for f in sorted(self.q.origins)
+                )
+            )
         if self.q.reversed:
             rclause = or_(
                 and_(

{ftmq-4.1.0 → ftmq-4.3.1}/ftmq/store/__init__.py

@@ -1,7 +1,7 @@
-from functools import cache
 from pathlib import Path
 from urllib.parse import urlparse
 
+from anystore.functools import weakref_cache as cache
 from anystore.types import Uri
 from followthemoney.dataset.dataset import Dataset
 from nomenklatura import Resolver, settings

{ftmq-4.1.0 → ftmq-4.3.1}/ftmq/store/base.py

@@ -1,7 +1,7 @@
-from functools import cache
 from typing import Generator, Iterable
 from urllib.parse import urlparse
 
+from anystore.functools import weakref_cache as cache
 from followthemoney import DefaultDataset
 from followthemoney.dataset.dataset import Dataset
 from nomenklatura import store as nk

{ftmq-4.1.0 → ftmq-4.3.1}/ftmq/store/fragments/__init__.py

@@ -1,4 +1,4 @@
-from functools import cache
+from anystore.functools import weakref_cache as cache
 
 from ftmq.store.fragments.dataset import Fragments
 from ftmq.store.fragments.settings import Settings

{ftmq-4.1.0 → ftmq-4.3.1}/ftmq/store/fragments/dataset.py

@@ -1,4 +1,5 @@
 import logging
+from contextlib import contextmanager
 from datetime import datetime
 from typing import Generator, Iterable, TypeAlias
 
@@ -9,6 +10,7 @@ from normality import slugify
 from sqlalchemy import (
     JSON,
     Column,
+    Connection,
     DateTime,
     String,
     Table,
@@ -22,22 +24,48 @@ from sqlalchemy.exc import OperationalError
 
 from ftmq.store.fragments.loader import BulkLoader
 from ftmq.store.fragments.utils import NULL_ORIGIN
-from ftmq.types import OriginStatements
+from ftmq.types import Statements
 from ftmq.util import make_dataset
 
 log = logging.getLogger(__name__)
 UNDEFINED = (OperationalError,)
 try:
-    from psycopg2.errors import UndefinedTable
+    from psycopg.errors import UndefinedTable
 
     UNDEFINED = (UndefinedTable, *UNDEFINED)
 except ImportError:
-    pass
+    try:
+        from psycopg2.errors import UndefinedTable
+
+        UNDEFINED = (UndefinedTable, *UNDEFINED)
+    except ImportError:
+        pass
 
 
 EntityFragments: TypeAlias = Generator[EntityProxy, None, None]
 
 
+@contextmanager
+def disable_timeout(conn: Connection, store):
+    # for long running iterations (e.g. re-index in OpenAleph), for postgres we
+    # don't want to get cancelled if a idle_in_transaction_timeout is configured
+    # on the server
+    if store.is_postgres:
+        raw_conn = conn.connection.driver_connection
+        with raw_conn.cursor() as cursor:
+            cursor.execute("SET idle_in_transaction_session_timeout = 0")
+    try:
+        yield conn
+    finally:
+        if store.is_postgres:
+            try:
+                raw_conn = conn.connection.driver_connection
+                with raw_conn.cursor() as cursor:
+                    cursor.execute("SET idle_in_transaction_session_timeout = DEFAULT")
+            except Exception:
+                pass  # Connection might be closed
+
+
 class Fragments(object):
     def __init__(self, store, name, origin=NULL_ORIGIN):
         self.store = store
@@ -104,7 +132,9 @@ class Fragments(object):
     def bulk(self, size=1000):
         return BulkLoader(self, size)
 
-    def fragments(self, entity_ids=None, fragment=None):
+    def fragments(
+        self, entity_ids=None, fragment=None, schema=None, since=None, until=None
+    ):
         stmt = self.table.select()
         entity_ids = ensure_list(entity_ids)
         if len(entity_ids) == 1:
@@ -113,25 +143,42 @@ class Fragments(object):
             stmt = stmt.where(self.table.c.id.in_(entity_ids))
         if fragment is not None:
             stmt = stmt.where(self.table.c.fragment == fragment)
+        if schema is not None:
+            if self.store.is_postgres:
+                stmt = stmt.where(self.table.c.entity["schema"].astext == schema)
+            else:
+                # SQLite JSON support - use json_extract function
+                stmt = stmt.where(
+                    func.json_extract(self.table.c.entity, "$.schema") == schema
+                )
+        if since is not None:
+            stmt = stmt.where(self.table.c.timestamp >= since)
+        if until is not None:
+            stmt = stmt.where(self.table.c.timestamp <= until)
         stmt = stmt.order_by(self.table.c.id)
         # stmt = stmt.order_by(self.table.c.origin)
         # stmt = stmt.order_by(self.table.c.fragment)
         conn = self.store.engine.connect()
         try:
-            conn = conn.execution_options(stream_results=True)
-            for ent in conn.execute(stmt):
-                data = {"id": ent.id, "datasets": [self.name], **ent.entity}
-                if ent.origin != NULL_ORIGIN:
-                    data["origin"] = ent.origin
-                yield data
+            with disable_timeout(conn, self.store) as conn:
+                conn = conn.execution_options(stream_results=True)
+                for ent in conn.execute(stmt):
+                    data = {"id": ent.id, "datasets": [self.name], **ent.entity}
+                    if ent.origin != NULL_ORIGIN:
+                        data["origin"] = ent.origin
+                    yield data
         except Exception:
             self.reset()
             raise
         finally:
             conn.close()
 
-    def partials(self, entity_id=None, skip_errors=False) -> EntityFragments:
-        for fragment in self.fragments(entity_ids=entity_id):
+    def partials(
+        self, entity_id=None, skip_errors=False, schema=None, since=None, until=None
+    ) -> EntityFragments:
+        for fragment in self.fragments(
+            entity_ids=entity_id, schema=schema, since=since, until=until
+        ):
             try:
                 yield EntityProxy.from_dict(fragment, cleaned=True)
             except Exception:
@@ -140,18 +187,32 @@ class Fragments(object):
                     continue
                 raise
 
-    def iterate(self, entity_id=None, skip_errors=False) -> EntityFragments:
+    def iterate(
+        self, entity_id=None, skip_errors=False, schema=None, since=None, until=None
+    ) -> EntityFragments:
+        if entity_id is None:
+            log.info("Using batched iteration for complete dataset.")
+            yield from self.iterate_batched(
+                skip_errors=skip_errors, schema=schema, since=since, until=until
+            )
+            return
         entity = None
         invalid = None
         fragments = 1
-        for partial in self.partials(entity_id=entity_id, skip_errors=skip_errors):
+        for partial in self.partials(
+            entity_id=entity_id,
+            skip_errors=skip_errors,
+            schema=schema,
+            since=since,
+            until=until,
+        ):
             if partial.id == invalid:
                 continue
             if entity is not None:
                 if entity.id == partial.id:
                     fragments += 1
                     if fragments % 10000 == 0:
-                        log.debug(
+                        log.warning(
                             "[%s:%s] aggregated %d fragments...",
                             entity.schema.name,
                             entity.id,
@@ -176,13 +237,78 @@ class Fragments(object):
         if entity is not None:
             yield entity
 
-    def origin_statements(
+    def iterate_batched(
+        self, skip_errors=False, batch_size=10_000, schema=None, since=None, until=None
+    ) -> EntityFragments:
+        """
+        For large datasets an overall sort is not feasible, so we iterate in
+        sorted batched IDs.
+        """
+        for entity_ids in self.get_sorted_id_batches(
+            batch_size, schema=schema, since=since, until=until
+        ):
+            yield from self.iterate(
+                entity_id=entity_ids,
+                skip_errors=skip_errors,
+                schema=schema,
+                since=since,
+                until=until,
+            )
+
+    def get_sorted_id_batches(
+        self, batch_size=10_000, schema=None, since=None, until=None
+    ) -> Generator[list[str], None, None]:
+        """
+        Get sorted ID batches to speed up iteration and useful to parallelize
+        processing of iterator Entities
+        """
+        last_id = None
+        with self.store.engine.connect() as conn:
+            while True:
+                stmt = select(self.table.c.id).distinct()
+                if last_id is not None:
+                    stmt = stmt.where(self.table.c.id > last_id)
+                if schema is not None:
+                    if self.store.is_postgres:
+                        stmt = stmt.where(
+                            self.table.c.entity["schema"].astext == schema
+                        )
+                    else:
+                        # SQLite JSON support - use json_extract function
+                        stmt = stmt.where(
+                            func.json_extract(self.table.c.entity, "$.schema") == schema
+                        )
+                if since is not None:
+                    stmt = stmt.where(self.table.c.timestamp >= since)
+                if until is not None:
+                    stmt = stmt.where(self.table.c.timestamp <= until)
+                stmt = stmt.order_by(self.table.c.id).limit(batch_size)
+                try:
+                    res = conn.execute(stmt)
+                    entity_ids = [r.id for r in res.fetchall()]
+                    if not entity_ids:
+                        return
+                    yield entity_ids
+                    last_id = entity_ids[-1]
+                except Exception:
+                    self.reset()
+                    raise
+
+    def get_sorted_ids(
+        self, batch_size=10_000, schema=None, since=None, until=None
+    ) -> Generator[str, None, None]:
+        """Get sorted IDs, optionally filtered by schema"""
+        for batch in self.get_sorted_id_batches(batch_size, schema, since, until):
+            yield from batch
+
+    def statements(
         self,
         entity_ids: Iterable[str] | None = None,
         origin: str | None = None,
         since: datetime | None = None,
-    ) -> OriginStatements:
-        """Iterate unsorted statements with its origins: (Statement, origin)"""
+        until: datetime | None = None,
+    ) -> Statements:
+        """Iterate unsorted statements with its fragment origins"""
         stmt = self.table.select()
         entity_ids = ensure_list(entity_ids)
         if len(entity_ids) == 1:
@@ -192,7 +318,9 @@ class Fragments(object):
         if origin is not None:
             stmt = stmt.where(self.table.c.origin == origin)
         if since is not None:
-            stmt = stmt.where(self.table.c.timestamp > since)
+            stmt = stmt.where(self.table.c.timestamp >= since)
+        if until is not None:
+            stmt = stmt.where(self.table.c.timestamp <= until)
         conn = self.store.engine.connect()
         default_dataset = make_dataset(self.name)
         try:
@@ -204,8 +332,10 @@ class Fragments(object):
                 )
                 for statement in entity.statements:
                     statement.last_seen = fragment.timestamp.isoformat()
-                    origin = fragment.origin if fragment.origin != NULL_ORIGIN else None
-                    yield statement, origin
+                    statement.origin = (
+                        fragment.origin if fragment.origin != NULL_ORIGIN else None
+                    )
+                    yield statement
         except Exception:
             self.reset()
             raise

{ftmq-4.1.0 → ftmq-4.3.1}/ftmq/store/fragments/loader.py

@@ -26,11 +26,16 @@ EXCEPTIONS = (
     TimeoutError,
 )
 try:
-    from psycopg2 import DatabaseError, OperationalError
+    from psycopg import DatabaseError, OperationalError
 
     EXCEPTIONS = (DatabaseError, OperationalError, *EXCEPTIONS)
 except ImportError:
-    pass
+    try:
+        from psycopg2 import DatabaseError, OperationalError
+
+        EXCEPTIONS = (DatabaseError, OperationalError, *EXCEPTIONS)
+    except ImportError:
+        pass
 
 log = logging.getLogger(__name__)
 
@@ -50,9 +55,12 @@ class BulkLoader(object):
         else:
             entity = dict(entity)
         id_ = entity.pop("id")
-        self.buffer[(id_, origin, fragment)] = entity
-        if len(self.buffer) >= self.size:
-            self.flush()
+        if id_:
+            self.buffer[(id_, origin, fragment)] = entity
+            if len(self.buffer) >= self.size:
+                self.flush()
+        else:
+            log.warning("Entity has no ID!")
 
     def _store_values(self, conn, values):
         table = self.dataset.table

ftmq-4.3.1/ftmq/store/fragments/store.py

@@ -0,0 +1,71 @@
+from sqlalchemy import MetaData, create_engine
+from sqlalchemy import inspect as sqlalchemy_inspect
+
+from ftmq.store.fragments.dataset import Fragments
+from ftmq.store.fragments.utils import NULL_ORIGIN
+
+
+class Store(object):
+    """A database containing multiple tables that represent
+    FtM-store datasets."""
+
+    PREFIX = "ftm"
+
+    def _adjust_psycopg3_uri(self, database_uri: str) -> str:
+        """Adjust PostgreSQL URI to use psycopg3 dialect if psycopg is available."""
+        if database_uri.startswith(("postgresql://", "postgres://")):
+            try:
+                import psycopg  # noqa: F401
+
+                # Use psycopg3 dialect for better performance and compatibility
+                if database_uri.startswith("postgresql://"):
+                    return database_uri.replace(
+                        "postgresql://", "postgresql+psycopg://", 1
+                    )
+                elif database_uri.startswith("postgres://"):
+                    return database_uri.replace(
+                        "postgres://", "postgresql+psycopg://", 1
+                    )
+            except ImportError:
+                # Fall back to psycopg2 if psycopg3 is not available
+                pass
+        return database_uri
+
+    def __init__(
+        self,
+        database_uri: str,
+        **config,
+    ):
+        self.database_uri = self._adjust_psycopg3_uri(database_uri)
+
+        # Configure connection pooling for psycopg3
+        config.setdefault("pool_size", 1)
+        if self.database_uri.startswith("postgresql+psycopg://"):
+            config.setdefault("max_overflow", 5)
+            config.setdefault("pool_timeout", 60)
+            config.setdefault("pool_recycle", 3600)
+            config.setdefault("pool_pre_ping", True)
+
+        self.engine = create_engine(self.database_uri, future=True, **config)
+        self.is_postgres = self.engine.dialect.name == "postgresql"
+        self.meta = MetaData()
+
+    def get(self, name, origin=NULL_ORIGIN):
+        return Fragments(self, name, origin=origin)
+
+    def all(self, origin=NULL_ORIGIN):
+        prefix = f"{self.PREFIX}_"
+        inspect = sqlalchemy_inspect(self.engine)
+        for table in inspect.get_table_names():
+            if table.startswith(prefix):
+                name = table[len(prefix) :]
+                yield Fragments(self, name, origin=origin)
+
+    def close(self):
+        self.engine.dispose()
+
+    def __len__(self):
+        return len(list(self.all()))
+
+    def __repr__(self):
+        return "<Store(%r)>" % self.engine

{ftmq-4.1.0 → ftmq-4.3.1}/ftmq/store/lake.py

@@ -18,13 +18,14 @@ Layout:
     ```
 """
 
-from functools import cache
 from pathlib import Path
 from typing import Any, Generator, Iterable
+from urllib.parse import urlparse
 
 import duckdb
 import numpy as np
 import pandas as pd
+from anystore.functools import weakref_cache as cache
 from anystore.lock import Lock
 from anystore.logging import get_logger
 from anystore.store.fs import Store as FSStore
@@ -52,26 +53,35 @@ from sqlalchemy.sql import Select
 from ftmq.query import Query
 from ftmq.store.base import Store
 from ftmq.store.sql import SQLQueryView, SQLStore
-from ftmq.types import OriginStatement, StatementEntities
-from ftmq.util import ensure_entity, get_scope_dataset
+from ftmq.types import StatementEntities
+from ftmq.util import apply_dataset, ensure_entity, get_scope_dataset
 
 log = get_logger(__name__)
 
-Z_ORDER = ["canonical_id", "schema", "prop", "value"]
+Z_ORDER = ["canonical_id", "entity_id", "schema", "prop"]
+TARGET_SIZE = 50 * 10_485_760  # 500 MB
 PARTITION_BY = ["dataset", "bucket", "origin"]
 DEFAULT_ORIGIN = "default"
 BUCKET_DOCUMENT = "document"
 BUCKET_INTERVAL = "interval"
 BUCKET_THING = "thing"
-BLOOM = ColumnProperties(bloom_filter_properties=BloomFilterProperties(True))
+STATISTICS_BLOOM = ColumnProperties(
+    bloom_filter_properties=BloomFilterProperties(True),
+    statistics_enabled="CHUNK",
+    dictionary_enabled=True,
+)
+STATISTICS = ColumnProperties(statistics_enabled="CHUNK", dictionary_enabled=True)
 WRITER = WriterProperties(
+    data_page_size_limit=64 * 1024,
+    dictionary_page_size_limit=512 * 1024,
+    max_row_group_size=500_000,
     compression="SNAPPY",
     column_properties={
-        "canonical_id": BLOOM,
-        "entity_id": BLOOM,
-        "schema": BLOOM,
-        "prop": BLOOM,
-        "value": BLOOM,
+        "canonical_id": STATISTICS,
+        "entity_id": STATISTICS,
+        "schema": STATISTICS,
+        "prop": STATISTICS_BLOOM,
+        "value": STATISTICS_BLOOM,
     },
 )
 
@@ -111,6 +121,13 @@ class StorageSettings(BaseSettings):
             return not self.endpoint.startswith("https")
         return False
 
+    @property
+    def duckdb_endpoint(self) -> str | None:
+        if not self.endpoint:
+            return
+        scheme = urlparse(self.endpoint).scheme
+        return self.endpoint[len(scheme) + len("://") :]
+
 
 storage_settings = StorageSettings()
 
@@ -154,16 +171,15 @@ def get_schema_bucket(schema_name: str) -> str:
     return BUCKET_THING
 
 
-def pack_statement(stmt: Statement, origin: str) -> SDict:
+def pack_statement(stmt: Statement) -> SDict:
     data = stmt.to_db_row()
-    data["origin"] = origin
     data["bucket"] = get_schema_bucket(data["schema"])
     return data
 
 
-def pack_statements(statements: Iterable[OriginStatement]) -> pd.DataFrame:
-    df = pd.DataFrame(pack_statement(*s) for s in statements)
-    df = df.drop_duplicates().sort_values(Z_ORDER)
+def pack_statements(statements: Iterable[Statement]) -> pd.DataFrame:
+    df = pd.DataFrame(map(pack_statement, statements))
+    df = df.drop_duplicates()  # .sort_values(Z_ORDER)
     df = df.fillna(np.nan)
     return df
 
@@ -213,9 +229,10 @@ def ensure_schema_buckets(q: Query) -> Select:
 class LakeQueryView(SQLQueryView):
     def query(self, query: Query | None = None) -> StatementEntities:
         if query:
+            query.table = self.store.table
             query = self.ensure_scoped_query(query)
-            # sql = ensure_schema_buckets(query)
-            yield from self.store._iterate(query.sql.statements)
+            sql = ensure_schema_buckets(query)
+            yield from self.store._iterate(sql)
         else:
             yield from super().query(query)
 
@@ -225,6 +242,7 @@ class LakeStore(SQLStore):
         self._backend: FSStore = FSStore(uri=kwargs.pop("uri"))
         self._partition_by = kwargs.pop("partition_by", PARTITION_BY)
         self._lock: Lock = kwargs.pop("lock", Lock(self._backend))
+        self._enforce_dataset = kwargs.pop("enforce_dataset", False)
         assert isinstance(
             self._backend, FSStore
         ), f"Invalid store backend: `{self._backend.__class__}"
@@ -235,12 +253,13 @@ class LakeStore(SQLStore):
         self.uri = self._backend.uri
         setup_duckdb_storage()
 
-    def get_deltatable(self) -> DeltaTable:
+    @property
+    def deltatable(self) -> DeltaTable:
         return DeltaTable(self.uri, storage_options=storage_options())
 
     def _execute(self, q: Select, stream: bool = True) -> Generator[Any, None, None]:
         try:
-            yield from stream_duckdb(q, self.get_deltatable())
+            yield from stream_duckdb(q, self.deltatable)
         except TableNotFoundError:
             pass
 
@@ -265,7 +284,7 @@ class LakeStore(SQLStore):
 
     def get_origins(self) -> set[str]:
         q = select(self.table.c.origin).distinct()
-        return set([r.origin for r in stream_duckdb(q, self.get_deltatable())])
+        return set([r.origin for r in stream_duckdb(q, self.deltatable)])
 
 
 class LakeWriter(nk.Writer):
@@ -274,21 +293,27 @@ class LakeWriter(nk.Writer):
 
     def __init__(self, store: Store, origin: str | None = DEFAULT_ORIGIN):
         super().__init__(store)
-        self.batch: set[OriginStatement] = set()
+        self.batch: set[Statement] = set()
         self.origin = origin or DEFAULT_ORIGIN
 
-    def add_statement(self, stmt: Statement, origin: str | None) -> None:
+    def add_statement(self, stmt: Statement) -> None:
         if stmt.entity_id is None:
             return
-        origin = origin or self.origin
+        stmt.origin = stmt.origin or self.origin
         canonical_id = self.store.linker.get_canonical(stmt.entity_id)
         stmt.canonical_id = canonical_id
-        self.batch.add((stmt, origin))
+        self.batch.add(stmt)
 
     def add_entity(self, entity: EntityProxy, origin: str | None = None) -> None:
         e = ensure_entity(entity, StatementEntity, self.store.dataset)
+        if self.store._enforce_dataset:
+            e = apply_dataset(e, self.store.dataset, replace=True)
         for stmt in e.statements:
-            self.add_statement(stmt, origin)
+            if origin:
+                stmt.origin = origin
+            self.add_statement(stmt)
+        # we check here instead of in `add_statement` as this will keep entities
+        # together in the same parquet files`
         if len(self.batch) >= self.BATCH_STATEMENTS:
             self.flush()
 
@@ -304,8 +329,10 @@ class LakeWriter(nk.Writer):
                     pack_statements(self.batch),
                     partition_by=self.store._partition_by,
                     mode="append",
-                    writer_properties=WRITER,
                     schema_mode="merge",
+                    writer_properties=WRITER,
+                    target_file_size=TARGET_SIZE,
+                    storage_options=storage_options(),
                 )
 
         self.batch = set()
@@ -317,8 +344,7 @@ class LakeWriter(nk.Writer):
         for row in self.store._execute(q):
             statements.append(Statement.from_db_row(row))
 
-        table = self.store.get_deltatable()
-        table.delete(f"canonical_id = '{entity_id}'")
+        self.store.deltatable.delete(f"canonical_id = '{entity_id}'")
         return statements
 
     def optimize(
@@ -327,10 +353,11 @@ class LakeWriter(nk.Writer):
         """
         Optimize the storage: Z-Ordering and compacting
         """
-        table = self.store.get_deltatable()
-        table.optimize.z_order(Z_ORDER, writer_properties=WRITER)
+        self.store.deltatable.optimize.z_order(
+            Z_ORDER, writer_properties=WRITER, target_size=TARGET_SIZE
+        )
         if vacuum:
-            table.vacuum(
+            self.store.deltatable.vacuum(
                 retention_hours=vacuum_keep_hours,
                 enforce_retention_duration=False,
                 dry_run=False,

{ftmq-4.1.0 → ftmq-4.3.1}/ftmq/types.py

@@ -20,7 +20,3 @@ ValueEntities: TypeAlias = Generator[ValueEntity, None, None]
 # statements
 Statements: TypeAlias = Generator[Statement, None, None]
 """A generator for Statement instances"""
-OriginStatement: TypeAlias = tuple[Statement, str | None]
-"""A statement with its origin"""
-OriginStatements: TypeAlias = Generator[OriginStatement, None, None]
-"""A generator for OriginStatement instances"""

{ftmq-4.1.0 → ftmq-4.3.1}/ftmq/util.py

@@ -1,17 +1,22 @@
-from functools import cache, lru_cache
+from functools import lru_cache
 from typing import Any, Generator, Type
 
 import pycountry
-from anystore.types import SDict
+from anystore.functools import weakref_cache as cache
+from anystore.types import SDict, StrGenerator
 from banal import ensure_list, is_listish
-from followthemoney import E
+from followthemoney import E, model
+from followthemoney.compare import _normalize_names
 from followthemoney.dataset import Dataset
 from followthemoney.entity import ValueEntity
 from followthemoney.proxy import EntityProxy
 from followthemoney.schema import Schema
 from followthemoney.types import registry
 from followthemoney.util import make_entity_id, sanitize_text
-from normality import collapse_spaces, slugify
+from normality import collapse_spaces, latinize_text, slugify
+from rigour.names import Name, Symbol, tag_org_name, tag_person_name
+from rigour.names.tokenize import normalize_name
+from rigour.text.scripts import can_latinize
 
 from ftmq.enums import Comparators
 from ftmq.types import Entity
@@ -91,6 +96,8 @@ def make_entity(
     etype = entity_type or ValueEntity
     if data.get("id") is None:
         raise ValueError("Entity has no ID.")
+    if etype == EntityProxy:
+        return EntityProxy.from_dict(data)
     if etype == ValueEntity:
         if not data.get("datasets"):
             dataset = make_dataset(default_dataset).name
@@ -373,6 +380,24 @@ def make_fingerprint(value: Any) -> str | None:
     return " ".join(sorted(set(slugify(value).split("-"))))
 
 
+def entity_fingerprints(entity: EntityProxy) -> set[str]:
+    """Get the set of entity name fingerprints, latinized if the alphabet allows
+    it and with org / person tags removed depending on entity schema"""
+    return make_fingerprints(*entity.names, schemata={entity.schema})
+
+
+def make_fingerprints(*names: str, schemata: set[Schema] | None = None) -> set[str]:
+    """Get the set of name fingerprints, latinized if the alphabet allows
+    it and with org / person tags removed depending on given schemata"""
+    # FIXME private import
+    schemata = schemata or {model["LegalEntity"]}
+    fps: set[str] = set()
+    for schema in schemata:
+        fps.update(set(_normalize_names(schema, names)))
+    # add latinized if appropriate
+    return {latinize_text(fp) if can_latinize(fp) else fp for fp in fps}
+
+
 def make_string_id(*values: Any) -> str | None:
     """
     Compute a hash id based on values
@@ -458,3 +483,64 @@ def must_str(value: Any) -> str:
     if not value:
         raise ValueError(f"Value invalid: `{value}`")
     return value
+
+
+SELECT_SYMBOLS = "__symbols__"
+SELECT_ANNOTATED = "__annotated__"
+
+
+def get_name_symbols(schema: Schema, *names: str) -> set[Symbol]:
+    """Get the rigour names symbols for the given schema and list of names"""
+    symbols: set[Symbol] = set()
+    if schema.is_a("Person"):
+        taggers = [tag_person_name]
+    elif schema.is_a("Organization"):
+        taggers = [tag_org_name]
+    elif schema.is_a("LegalEntity"):
+        taggers = [tag_org_name, tag_person_name]
+    else:
+        return symbols
+    for name in names:
+        n = Name(name)
+        for tagger in taggers:
+            for symbol in tagger(n, normalize_name).symbols:
+                symbols.add(symbol)
+    return symbols
+
+
+def get_symbols(entity: EntityProxy) -> set[Symbol]:
+    """Get the rigour names symbols for the given entity"""
+    if not entity.schema.is_a("LegalEntity"):
+        return set()
+    names = entity.get_type_values(registry.name, matchable=True)
+    return get_name_symbols(entity.schema, *names)
+
+
+def inline_symbols(entity: EntityProxy) -> None:
+    """Get the rigour names symbols for the given entity and write them to `indexText`"""
+    # clean up old symbols from indexText:
+    for text in entity.pop("indexText"):
+        if not text.startswith(SELECT_SYMBOLS):
+            entity.add("indexText", text)
+    symbols = get_symbols(entity)
+    entity.add("indexText", f"{SELECT_SYMBOLS} {','.join(map(str, symbols))}")
+
+
+def select_data(e: EntityProxy, prefix: str) -> StrGenerator:
+    """Select arbitrary stored data in `indexText` identified by given prefix"""
+    for text in e.get("indexText", quiet=True):
+        if text.startswith(prefix):
+            yield text.replace(prefix, "").strip()
+
+
+def select_symbols(e: EntityProxy) -> set[str]:
+    """Select stored symbols in `indexText`"""
+    symbols: set[str] = set()
+    for data in select_data(e, SELECT_SYMBOLS):
+        symbols.update(data.split(","))
+    return symbols
+
+
+def select_annotations(e: EntityProxy) -> set[str]:
+    """Select stored annotations in `indexText`"""
+    return {s for s in select_data(e, SELECT_ANNOTATED)}

{ftmq-4.1.0 → ftmq-4.3.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ftmq"
-version = "4.1.0"
+version = "4.3.1"
 description = "followthemoney query dsl and io helpers"
 authors = [{ name = "Simon Wörpel", email = "simon.woerpel@pm.me" }]
 license = "AGPLv3+"
@@ -15,9 +15,10 @@ classifiers = [
 ]
 requires-python = ">=3.11,<4"
 dependencies = [
-    "anystore (>=0.3.9,<0.4.0)",
-    "followthemoney (>=4.1.1,<5.0.0)",
-    "nomenklatura (>=4.1.0,<5.0.0)",
+    "anystore (>=0.4.0,<0.5.0)",
+    "followthemoney (>=4.3.2,<5.0.0)",
+    "nomenklatura (>=4.1.10,<5.0.0)",
+    "rigour (>=1.4.0,<2.0.0)",
     "click (>=8.2.1,<9.0.0)",
     "click-default-group (>=1.2.4,<2.0.0)",
     "orjson (>=3.10.18,<4.0.0)",
@@ -29,12 +30,12 @@ dependencies = [
 [project.optional-dependencies]
 level = ["plyvel (>=1.5.1,<2.0.0)"]
 sql = ["sqlalchemy (>=2.0.36,<3.0.0)"]
-postgres = ["sqlalchemy (>=2.0.36,<3.0.0)", "psycopg2 (>=2.9.10,<3.0.0)"]
+postgres = ["sqlalchemy (>=2.0.36,<3.0.0)", "psycopg[pool] (>=3.2.9,<4.0.0)"]
 redis = ["redis (>=5.2.1,<6.0.0)", "fakeredis (>=2.26.2,<3.0.0)"]
 lake = [
-    "duckdb (>=1.3.2,<2.0.0)",
-    "pandas (>=2.3.1,<3.0.0)",
-    "deltalake (>=1.1.2,<2.0.0)",
+    "duckdb (>=1.4.1,<2.0.0)",
+    "pandas (>=2.3.3,<3.0.0)",
+    "deltalake (>=1.2.1,<2.0.0)",
     "pyarrow (>=21.0.0,<22.0.0)",
 ]
 aleph = ["furl (>=2.1.4,<3.0.0)", "alephclient (>=2.6.0,<3.0.0)"]
@@ -50,19 +51,19 @@ Issues = "https://github.com/dataresearchcenter/ftmq/issues"
 
 [tool.poetry.group.dev.dependencies]
 pytest = ">=7.4.3,<9.0.0"
-pytest-cov = ">=4.1,<7.0"
+pytest-cov = ">=4.1,<8.0"
 pytest-env = "^1.1.1"
 black = ">=23.11,<26.0"
-isort = "^6.0.1"
-mypy = "^1.17.0"
+isort = "^7.0.0"
+mypy = "^1.17.1"
 pre-commit = "^4.0.1"
 flake8 = ">=6.1,<8.0"
 ipdb = "^0.13.13"
 bump2version = "^1.0.1"
 mkdocs = "^1.6.1"
-mkdocstrings-python = "^1.16.10"
-mkdocs-autorefs = "^1.4.1"
-mkdocs-material = "^9.6.15"
+mkdocs-autorefs = "^1.4.3"
+mkdocstrings-python = "^1.18.2"
+mkdocs-material = "^9.6.18"
 mkdocs-click = "^0.9.0"
 
 [build-system]

ftmq-4.1.0/ftmq/store/fragments/store.py

@@ -1,43 +0,0 @@
-from sqlalchemy import MetaData, create_engine
-from sqlalchemy import inspect as sqlalchemy_inspect
-
-from ftmq.store.fragments.dataset import Fragments
-from ftmq.store.fragments.utils import NULL_ORIGIN
-
-
-class Store(object):
-    """A database containing multiple tables that represent
-    FtM-store datasets."""
-
-    PREFIX = "ftm"
-
-    def __init__(
-        self,
-        database_uri: str,
-        **config,
-    ):
-        self.database_uri = database_uri
-        # config.setdefault('pool_size', 1)
-        self.engine = create_engine(database_uri, future=True, **config)
-        self.is_postgres = self.engine.dialect.name == "postgresql"
-        self.meta = MetaData()
-
-    def get(self, name, origin=NULL_ORIGIN):
-        return Fragments(self, name, origin=origin)
-
-    def all(self, origin=NULL_ORIGIN):
-        prefix = f"{self.PREFIX}_"
-        inspect = sqlalchemy_inspect(self.engine)
-        for table in inspect.get_table_names():
-            if table.startswith(prefix):
-                name = table[len(prefix) :]
-                yield Fragments(self, name, origin=origin)
-
-    def close(self):
-        self.engine.dispose()
-
-    def __len__(self):
-        return len(list(self.all()))
-
-    def __repr__(self):
-        return "<Store(%r)>" % self.engine