diff --git a/documentation/changelog.rst b/documentation/changelog.rst index 1e2f90a7ee..509a6eb88e 100644 --- a/documentation/changelog.rst +++ b/documentation/changelog.rst @@ -49,6 +49,7 @@ New features Infrastructure / Support ---------------------- +* Look up which data sources recorded for which sensors from a small summary table instead of scanning the beliefs table [see `PR #2382 `_] * Speed up sensor data queries and free up disk space by reordering the ``timed_belief`` primary key to lead with ``sensor_id`` and dropping the indexes it makes redundant, in a migration that runs online and so needs no maintenance window (though it can take a while on a large database) [see `PR #2378 `_] * The database migration for this release splits each stored flex-context's ``inflexible-device-sensors`` field into ``inflexible-consumption``/``inflexible-production`` sensor references, classifying each sensor by its ``consumption_is_positive`` attribute (behavior-preserving; sensor attributes themselves are kept). Downgrading merges them back into bare sensor IDs, dropping any source filters added in the meantime [see `PR #2358 `_] * Speed up listing assets: eager-load each asset's sensors instead of lazy-loading them one query per asset during serialization, and skip loading sensors entirely for field-filtered responses that do not include them [see `PR #2363 `_] diff --git a/flexmeasures/data/migrations/versions/f1c8a3d75e29_add_sensor_data_source_table.py b/flexmeasures/data/migrations/versions/f1c8a3d75e29_add_sensor_data_source_table.py new file mode 100644 index 0000000000..078e754169 --- /dev/null +++ b/flexmeasures/data/migrations/versions/f1c8a3d75e29_add_sensor_data_source_table.py @@ -0,0 +1,147 @@ +"""add the sensor_data_source summary table + +Records which data sources have recorded beliefs for which sensors. + +The same information is already implicit in ``timed_belief``, +but getting it from there costs a scan of the largest table in the database, +to produce a relation bounded by sensors times sources, +which in practice is a few thousand rows. +It also has to be read in ``source_id`` order for ``DataSource.sensors``, +which no index serves once the primary key leads with ``sensor_id``. + +The table is kept current by a statement-level trigger on ``timed_belief``, rather than by application code. +A trigger cannot be bypassed: +bulk inserts, ``COPY``, plugins and raw SQL all maintain the summary, +whereas a hook in the save path only covers the callers that happen to use it. +Doing it per statement rather than per row means one small upsert per insert statement, +however many rows that statement carries. + +The table is a superset: +pairs are added when beliefs are inserted and are not removed when beliefs are deleted, +because deciding whether a pair went stale needs exactly the scan this avoids. +``Sensor.search_data_sources`` still reads ``timed_belief`` whenever time filters are given, +so time-bounded questions remain exact. + +The backfill reads every belief row once. +It takes a plain ACCESS SHARE lock, so reads and writes continue, +but on a large database expect it to take a few minutes. + +Revision ID: f1c8a3d75e29 +Revises: b7e5a2c40f18 +Create Date: 2026-08-03 + +""" + +import logging + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = "f1c8a3d75e29" +down_revision = "b7e5a2c40f18" +branch_labels = None +depends_on = None + +logger = logging.getLogger("alembic.runtime.migration") + + +CREATE_FUNCTION = """ +CREATE OR REPLACE FUNCTION record_sensor_data_sources() RETURNS trigger +LANGUAGE plpgsql AS $$ +BEGIN + INSERT INTO sensor_data_source (sensor_id, source_id) + SELECT DISTINCT sensor_id, source_id FROM inserted_beliefs + ON CONFLICT DO NOTHING; + RETURN NULL; +END; +$$ +""" + +CREATE_TRIGGER = """ +CREATE TRIGGER timed_belief_record_sensor_data_sources +AFTER INSERT ON timed_belief +REFERENCING NEW TABLE AS inserted_beliefs +FOR EACH STATEMENT +EXECUTE FUNCTION record_sensor_data_sources() +""" + +BACKFILL = """ +INSERT INTO sensor_data_source (sensor_id, source_id) +SELECT DISTINCT sensor_id, source_id FROM timed_belief +ON CONFLICT DO NOTHING +""" + + +def upgrade(): + connection = op.get_bind() + + # Guarded, because the steps below commit as they go: + # a run that fails during the backfill leaves the table and trigger in place, + # and the retry has to get past this point. + if not sa.inspect(connection).has_table("sensor_data_source"): + op.create_table( + "sensor_data_source", + sa.Column("sensor_id", sa.Integer(), nullable=False), + sa.Column("source_id", sa.Integer(), nullable=False), + sa.ForeignKeyConstraint( + ["sensor_id"], + ["sensor.id"], + name=op.f("sensor_data_source_sensor_id_sensor_fkey"), + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["source_id"], + ["data_source.id"], + name=op.f("sensor_data_source_source_id_data_source_fkey"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint( + "sensor_id", "source_id", name=op.f("sensor_data_source_pkey") + ), + ) + + # Install the trigger *before* the backfill, and commit it, so that it is live for + # other sessions while the backfill runs. + # The other order loses data: + # the backfill reads a snapshot taken when its statement began, + # so beliefs inserted while it runs are not in it, + # and a trigger created afterwards in the same transaction was not visible to those + # inserting sessions either, so nothing would ever record them. + # With this order the two overlap instead of leaving a gap, + # and ON CONFLICT DO NOTHING absorbs the overlap. + with op.get_context().autocommit_block(): + op.execute(sa.text(CREATE_FUNCTION)) + op.execute( + sa.text( + "DROP TRIGGER IF EXISTS timed_belief_record_sensor_data_sources" + " ON timed_belief" + ) + ) + op.execute(sa.text(CREATE_TRIGGER)) + + # Backfill the beliefs that predate the trigger. + # DISTINCT over the whole table is the one expensive step here, + # and it is also the last time we ever have to ask this question that way. + n_rows = connection.execute( + sa.text("SELECT reltuples::bigint FROM pg_class WHERE relname = 'timed_belief'") + ).scalar_one_or_none() + if n_rows is not None and n_rows > 1_000_000: + message = ( + f"Summarising which sources recorded for which sensors" + f" (~{n_rows:,} belief rows to scan once): this may take a few minutes." + ) + # Also print: FlexMeasures' logging setup does not surface the alembic logger. + print(message, flush=True) + logger.info(message) + + op.execute(sa.text(BACKFILL)) + + +def downgrade(): + op.execute( + "DROP TRIGGER IF EXISTS timed_belief_record_sensor_data_sources ON timed_belief" + ) + op.execute("DROP FUNCTION IF EXISTS record_sensor_data_sources()") + op.drop_table("sensor_data_source") diff --git a/flexmeasures/data/models/data_sources.py b/flexmeasures/data/models/data_sources.py index 642b51dc66..7450f10108 100644 --- a/flexmeasures/data/models/data_sources.py +++ b/flexmeasures/data/models/data_sources.py @@ -268,6 +268,42 @@ def _clean_parameters(self, parameters: dict) -> dict: ] +class SensorDataSource(db.Model): + """Records that a data source has recorded beliefs for a sensor. + + This is a summary of ``timed_belief``, not an independent fact: + the same information could be had with ``SELECT DISTINCT sensor_id, source_id FROM timed_belief``. + Keeping it separately matters because that query costs a scan of the beliefs table, + which is typically the largest in the database, + to produce a relation bounded by the number of sensors times the number of sources. + In practice that is a few thousand rows at most. + + .. note:: This is a *superset*. + A pair is added when beliefs are saved, and is not removed when those beliefs are deleted, + because deciding whether a pair has become stale needs exactly the scan this table exists to avoid. + So read a row as "this source has recorded for this sensor at some point", + not "this source has beliefs stored for this sensor right now". + ``Sensor.search_data_sources`` still consults ``timed_belief`` directly whenever time filters are given, + so time-bounded questions stay exact. + """ + + __tablename__ = "sensor_data_source" + + sensor_id = db.Column( + db.Integer, + db.ForeignKey("sensor.id", ondelete="CASCADE"), + primary_key=True, + ) + source_id = db.Column( + db.Integer, + db.ForeignKey("data_source.id", ondelete="CASCADE"), + primary_key=True, + ) + + def __repr__(self) -> str: + return f"" + + class DataSource(db.Model, tb.BeliefSourceDBMixin): """Each data source is a data-providing entity.""" @@ -317,22 +353,25 @@ class DataSource(db.Model, tb.BeliefSourceDBMixin): @property def sensors(self) -> list: - """Return all Sensor objects that have beliefs recorded by this data source. + """Return all Sensor objects that this data source has recorded beliefs for. + + Reads the ``sensor_data_source`` summary rather than ``timed_belief``. + Answering this from the beliefs table would mean scanning it in ``source_id`` order, + which no index serves, to produce a handful of rows. - Uses a two-step subquery (distinct sensor IDs → Sensor rows) so that - it scales to very large timed_belief tables without fetching every belief row. - Mirrors the approach of ``Sensor.data_sources``. + See :class:`SensorDataSource` for the superset semantics: + a sensor stays listed after its beliefs from this source are deleted. """ - from flexmeasures.data.models.time_series import Sensor, TimedBelief + from flexmeasures.data.models.time_series import Sensor - sensor_id_subq = ( - select(TimedBelief.sensor_id) - .where(TimedBelief.source_id == self.id) - .distinct() - .subquery() - ) return db.session.scalars( - select(Sensor).where(Sensor.id.in_(select(sensor_id_subq))) + select(Sensor).where( + Sensor.id.in_( + select(SensorDataSource.sensor_id).where( + SensorDataSource.source_id == self.id + ) + ) + ) ).all() _data_generator: ClassVar[DataGenerator | None] = None diff --git a/flexmeasures/data/models/time_series.py b/flexmeasures/data/models/time_series.py index 845d1bc167..a9e9cb2c0f 100644 --- a/flexmeasures/data/models/time_series.py +++ b/flexmeasures/data/models/time_series.py @@ -9,7 +9,7 @@ import numpy as np import pandas as pd -from sqlalchemy import exists, select +from sqlalchemy import event, exists, select, text as sa_text from sqlalchemy.ext.declarative import declared_attr from sqlalchemy.ext.mutable import MutableDict from sqlalchemy.schema import UniqueConstraint @@ -44,7 +44,7 @@ to_annotation_frame, ) from flexmeasures.data.models.charts import chart_type_to_chart_specs -from flexmeasures.data.models.data_sources import DataSource +from flexmeasures.data.models.data_sources import DataSource, SensorDataSource from flexmeasures.data.models.generic_assets import GenericAsset from flexmeasures.data.models.validation_utils import check_required_attributes from flexmeasures.data.queries.annotations import filter_by_belief_time @@ -717,7 +717,15 @@ def search_data_sources( exclude_source_types: list[str] | None = None, check_exists: bool = False, ) -> list[DataSource] | bool: - """ + """Find the data sources that have recorded beliefs for this sensor. + + Where the answer comes from depends on whether time filters are given. + With them, the beliefs table is consulted, so the answer is exact for that window. + Without them, the ``sensor_data_source`` summary is read instead, + which avoids scanning the beliefs table but is a superset: + a source stays listed after its beliefs for this sensor are deleted. + See :class:`~flexmeasures.data.models.data_sources.SensorDataSource`. + :returns: list of Data Source objects, or, if check_exists, True if any such sources exist, False if none do. """ @@ -774,17 +782,16 @@ def search_data_sources( q = select(DataSource).where(DataSource.id.in_(belief_q.distinct())) else: - # No time filters: retrieve distinct source IDs for this sensor via a - # lightweight index-only scan, then fetch those DataSource rows. This - # avoids a full join across potentially hundreds of millions of belief - # rows just to enumerate a handful of sources. - source_id_subq = ( - select(TimedBelief.source_id) - .where(TimedBelief.sensor_id == self.id) - .distinct() - .subquery() + # No time filters: read the sensor_data_source summary instead of the beliefs table, + # which turns a scan over very many rows into a lookup of a handful. + # See SensorDataSource for the superset semantics this accepts. + q = select(DataSource).where( + DataSource.id.in_( + select(SensorDataSource.source_id).where( + SensorDataSource.sensor_id == self.id + ) + ) ) - q = select(DataSource).where(DataSource.id.in_(select(source_id_subq))) if source_types: q = q.where(DataSource.type.in_(source_types)) @@ -800,9 +807,11 @@ def search_data_sources( def data_sources(self) -> list[DataSource]: """Return all DataSource objects that have recorded beliefs for this sensor. - Uses a two-step subquery (distinct source IDs → DataSource rows) so that - it scales to very large timed_belief tables without fetching every belief row. - Equivalent to ``search_data_sources()`` with no filters. + Equivalent to ``search_data_sources()`` with no filters, + which reads the ``sensor_data_source`` summary rather than the beliefs table. + + See :class:`~flexmeasures.data.models.data_sources.SensorDataSource` for the superset semantics that implies: + a source stays listed after its beliefs for this sensor are deleted. """ return self.search_data_sources() @@ -1216,3 +1225,86 @@ def add( def __repr__(self) -> str: """timely-beliefs representation of timed beliefs.""" return tb.TimedBelief.__repr__(self) + + +# How the sensor_data_source summary is kept up to date. +# +# A database trigger does it, rather than FlexMeasures code. +# The reason is that beliefs reach timed_belief by several routes: +# save_to_db, bulk inserts, plugins, and raw SQL. +# Code added to one of those routes would only ever see the beliefs that took it, +# leaving the summary quietly incomplete for all the others. +# A trigger sits on the table itself, so it sees every insert whatever the route. +# +# The trigger runs once per INSERT *statement* rather than once per row. +# It reads that statement's new rows in one go, through what PostgreSQL calls a +# transition table (named inserted_beliefs below). +# So saving a million beliefs in one statement adds one small insert, not a million. +# +# Migration f1c8a3d75e29 creates the same function and trigger. +# The statements are written out in both places rather than shared, +# because a migration should keep doing what it did when it was written, +# even if this file later changes. +RECORD_SENSOR_DATA_SOURCES_FUNCTION = """ +CREATE OR REPLACE FUNCTION record_sensor_data_sources() RETURNS trigger +LANGUAGE plpgsql AS $$ +BEGIN + INSERT INTO sensor_data_source (sensor_id, source_id) + SELECT DISTINCT sensor_id, source_id FROM inserted_beliefs + ON CONFLICT DO NOTHING; + RETURN NULL; +END; +$$ +""" + +RECORD_SENSOR_DATA_SOURCES_TRIGGER = """ +CREATE TRIGGER timed_belief_record_sensor_data_sources +AFTER INSERT ON timed_belief +REFERENCING NEW TABLE AS inserted_beliefs +FOR EACH STATEMENT +EXECUTE FUNCTION record_sensor_data_sources() +""" + + +def create_sensor_data_source_trigger(target, connection, **kwargs) -> None: + """Install the trigger, for databases whose schema is built by create_all(). + + A FlexMeasures database gets its schema in one of two ways, + and each needs its own route to the trigger: + + - built by Alembic migrations, as in production: + migration f1c8a3d75e29 installs it there. + - built by ``db.create_all()``, as in the test suite and some development setups: + this function installs it there. + + SQLAlchemy fires "after_create" whenever it has created something. + This listens for that on the whole metadata rather than on the timed_belief table, + so that it runs after *all* tables exist. + Listening on timed_belief alone would be a bet on creation order, + because the trigger's function reads sensor_data_source, + and nothing guarantees that table gets created first. + PostgreSQL would in fact tolerate the wrong order today, + since a function written in its procedural language does not look up the tables it names until it first runs, + but depending on that is fragile. + + Does nothing unless the database is PostgreSQL and both tables are present, + so another backend, or a create_all() that made only some tables, is left alone. + """ + if connection.dialect.name != "postgresql": + return + inspector = inspect(connection) + if not inspector.has_table("timed_belief") or not inspector.has_table( + "sensor_data_source" + ): + return + connection.execute(sa_text(RECORD_SENSOR_DATA_SOURCES_FUNCTION)) + connection.execute( + sa_text( + "DROP TRIGGER IF EXISTS timed_belief_record_sensor_data_sources" + " ON timed_belief" + ) + ) + connection.execute(sa_text(RECORD_SENSOR_DATA_SOURCES_TRIGGER)) + + +event.listen(db.metadata, "after_create", create_sensor_data_source_trigger) diff --git a/flexmeasures/data/tests/test_sensor_data_source.py b/flexmeasures/data/tests/test_sensor_data_source.py new file mode 100644 index 0000000000..b9d4823eb9 --- /dev/null +++ b/flexmeasures/data/tests/test_sensor_data_source.py @@ -0,0 +1,204 @@ +"""Tests for the sensor_data_source summary table.""" + +import pandas as pd +from sqlalchemy import func, select +from timely_beliefs import BeliefsDataFrame + +from flexmeasures.data.models.data_sources import DataSource, SensorDataSource +from flexmeasures.data.models.time_series import TimedBelief +from flexmeasures.data.utils import save_to_db +from flexmeasures.tests.utils import get_test_sensor + + +def pairs(db) -> set: + return set( + db.session.execute( + select(SensorDataSource.sensor_id, SensorDataSource.source_id) + ).all() + ) + + +def make_belief(sensor, source, hours_offset: int = 0) -> BeliefsDataFrame: + return BeliefsDataFrame( + [ + TimedBelief( + sensor=sensor, + source=source, + event_start=pd.Timestamp("2021-03-28 16:00:00+00:00") + + pd.Timedelta(hours=hours_offset), + belief_time=pd.Timestamp("2021-03-27 08:00:00+00:00"), + event_value=1.0, + ) + ] + ) + + +def test_saving_beliefs_records_the_pair(setup_beliefs, db): + """Saving beliefs must record which source recorded for which sensor.""" + sensor = get_test_sensor(db) + source = DataSource(name="Association test source", type="demo script") + db.session.add(source) + db.session.commit() + + assert (sensor.id, source.id) not in pairs(db) + + save_to_db(make_belief(sensor, source)) + db.session.commit() + + assert (sensor.id, source.id) in pairs(db) + + +def test_inserts_that_bypass_the_save_path_still_record_the_pair(setup_beliefs, db): + """A raw insert must maintain the summary just as a normal save does. + + This is the case that decided the design. + Beliefs reach timed_belief by more routes than ``save_to_db``: + bulk inserts, plugins and raw SQL among them, and several of this repo's own fixtures. + A hook in the save path would leave the summary silently incomplete for all of them, + so a database trigger maintains it instead. + """ + sensor = get_test_sensor(db) + source = DataSource(name="Raw insert source", type="demo script") + db.session.add(source) + db.session.commit() + + db.session.execute( + TimedBelief.__table__.insert().values( + sensor_id=sensor.id, + source_id=source.id, + event_start=pd.Timestamp("2021-03-28 20:00:00+00:00"), + belief_horizon=pd.Timedelta(hours=1), + cumulative_probability=0.5, + event_value=3.0, + ) + ) + db.session.commit() + + assert (sensor.id, source.id) in pairs(db) + assert sensor in source.sensors + + +def test_recording_the_pair_is_idempotent(setup_beliefs, db): + """Saving more beliefs from the same source must not raise or duplicate.""" + sensor = get_test_sensor(db) + source = DataSource(name="Idempotent source", type="demo script") + db.session.add(source) + db.session.commit() + + save_to_db(make_belief(sensor, source, hours_offset=0)) + db.session.commit() + save_to_db(make_belief(sensor, source, hours_offset=1)) + db.session.commit() + + matching = [p for p in pairs(db) if p == (sensor.id, source.id)] + assert len(matching) == 1 + + +def test_accessors_read_the_summary_and_not_the_beliefs(setup_beliefs, db): + """Both accessors must answer from sensor_data_source, not from timed_belief. + + Shown by writing a pair into the summary for which no belief exists at all. + An accessor reading timed_belief could not return it; + one reading the summary must. + """ + sensor = get_test_sensor(db) + source = DataSource(name="Source with no beliefs at all", type="demo script") + db.session.add(source) + db.session.commit() + + assert sensor not in source.sensors + assert source not in sensor.data_sources + + db.session.add(SensorDataSource(sensor_id=sensor.id, source_id=source.id)) + db.session.commit() + + # No belief ties these two together, so only the summary can be the answer's source. + belief_count = db.session.execute( + select(func.count()) + .select_from(TimedBelief) + .where(TimedBelief.sensor_id == sensor.id, TimedBelief.source_id == source.id) + ).scalar() + assert belief_count == 0 + + assert sensor in source.sensors + assert source in sensor.data_sources + + +def test_data_source_sensors_reflects_a_saved_belief(setup_beliefs, db): + """Saving a belief must make the sensor show up on the source.""" + sensor = get_test_sensor(db) + source = DataSource(name="Source listing sensors", type="demo script") + db.session.add(source) + db.session.commit() + + assert sensor not in source.sensors + + save_to_db(make_belief(sensor, source)) + db.session.commit() + + assert sensor in source.sensors + + +def test_sensor_data_sources_reflects_a_saved_belief(setup_beliefs, db): + """Saving a belief must make the source show up on the sensor, mirroring the above.""" + sensor = get_test_sensor(db) + source = DataSource(name="Source found from sensor", type="demo script") + db.session.add(source) + db.session.commit() + + save_to_db(make_belief(sensor, source)) + db.session.commit() + + assert source in sensor.data_sources + + +def test_summary_is_a_superset_after_deleting_beliefs(setup_beliefs, db): + """A pair survives deletion of the beliefs that created it, by design. + + Deciding whether a pair went stale would need the scan over timed_belief that this table exists to avoid, + so the summary is deliberately a superset. + This test pins that, so the behaviour is a documented choice rather than a surprise. + """ + sensor = get_test_sensor(db) + source = DataSource(name="Source whose beliefs go away", type="demo script") + db.session.add(source) + db.session.commit() + + save_to_db(make_belief(sensor, source)) + db.session.commit() + assert (sensor.id, source.id) in pairs(db) + + db.session.execute( + TimedBelief.__table__.delete().where( + TimedBelief.source_id == source.id, TimedBelief.sensor_id == sensor.id + ) + ) + db.session.commit() + + assert (sensor.id, source.id) in pairs(db), "pair should survive, by design" + assert sensor in source.sensors + + +def test_time_filtered_source_search_still_reads_beliefs(setup_beliefs, db): + """With time filters, the answer must stay exact rather than use the summary. + + The summary knows nothing about when beliefs were recorded, + so a time-bounded question has to go to timed_belief. + If it did not, a source whose beliefs all fall outside the window would be wrongly reported. + """ + sensor = get_test_sensor(db) + source = DataSource(name="Source outside the window", type="demo script") + db.session.add(source) + db.session.commit() + + save_to_db(make_belief(sensor, source)) + db.session.commit() + + # Unfiltered: found via the summary + assert source in sensor.search_data_sources() + + # Filtered to a window containing no beliefs from this source: must not appear + found = sensor.search_data_sources( + event_starts_after=pd.Timestamp("2030-01-01 00:00:00+00:00") + ) + assert source not in found