Coverage for python / lsst / dax / apdb / cassandra / apdbCassandra.py: 8%
841 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-05 02:05 -0800
1# This file is part of dax_apdb.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <http://www.gnu.org/licenses/>.
22from __future__ import annotations
24__all__ = ["ApdbCassandra"]
26import datetime
27import json
28import logging
29import random
30import uuid
31import warnings
32from collections import Counter, defaultdict
33from collections.abc import Iterable, Iterator, Mapping, Set
34from typing import TYPE_CHECKING, Any, cast
36import numpy as np
37import pandas
39# If cassandra-driver is not there the module can still be imported
40# but ApdbCassandra cannot be instantiated.
41try:
42 import cassandra
43 import cassandra.query
44 from cassandra.query import UNSET_VALUE
46 CASSANDRA_IMPORTED = True
47except ImportError:
48 CASSANDRA_IMPORTED = False
50import astropy.time
51import felis.datamodel
53from lsst import sphgeom
54from lsst.utils.iteration import chunk_iterable
56from ..apdb import Apdb, ApdbConfig
57from ..apdbConfigFreezer import ApdbConfigFreezer
58from ..apdbReplica import ApdbTableData, ReplicaChunk
59from ..apdbSchema import ApdbSchema, ApdbTables
60from ..apdbUpdateRecord import (
61 ApdbCloseDiaObjectValidityRecord,
62 ApdbReassignDiaSourceToDiaObjectRecord,
63 ApdbUpdateNDiaSourcesRecord,
64)
65from ..monitor import MonAgent
66from ..recordIds import DiaObjectId, DiaSourceId
67from ..schema_model import Table
68from ..timer import Timer
69from ..versionTuple import VersionTuple
70from .apdbCassandraAdmin import ApdbCassandraAdmin
71from .apdbCassandraReplica import ApdbCassandraReplica
72from .apdbCassandraSchema import ApdbCassandraSchema, CreateTableOptions, ExtraTables
73from .apdbMetadataCassandra import ApdbMetadataCassandra
74from .cassandra_utils import (
75 ApdbCassandraTableData,
76 execute_concurrent,
77 literal,
78 quote_id,
79 select_concurrent,
80)
81from .config import ApdbCassandraConfig, ApdbCassandraConnectionConfig, ApdbCassandraTimePartitionRange
82from .connectionContext import ConnectionContext, DbVersions
83from .exceptions import CassandraMissingError
84from .partitioner import Partitioner
85from .sessionFactory import SessionContext, SessionFactory
87if TYPE_CHECKING:
88 from ..apdbMetadata import ApdbMetadata
89 from ..apdbUpdateRecord import ApdbUpdateRecord
91_LOG = logging.getLogger(__name__)
93_MON = MonAgent(__name__)
95VERSION = VersionTuple(1, 3, 0)
96"""Version for the code controlling non-replication tables. This needs to be
97updated following compatibility rules when schema produced by this code
98changes.
99"""
102class ApdbCassandra(Apdb):
103 """Implementation of APDB database with Apache Cassandra backend.
105 Parameters
106 ----------
107 config : `ApdbCassandraConfig`
108 Configuration object.
109 """
111 def __init__(self, config: ApdbCassandraConfig):
112 if not CASSANDRA_IMPORTED:
113 raise CassandraMissingError()
115 self._config = config
116 self._keyspace = config.keyspace
117 self._schema = ApdbSchema(config.schema_file, config.ss_schema_file)
119 self._session_factory = SessionFactory(config)
120 self._connection_context: ConnectionContext | None = None
122 @property
123 def _context(self) -> ConnectionContext:
124 """Establish connection if not established and return context."""
125 if self._connection_context is None:
126 current_versions = DbVersions(
127 schema_version=self.schema.schemaVersion(),
128 code_version=self.apdbImplementationVersion(),
129 replica_version=ApdbCassandraReplica.apdbReplicaImplementationVersion(),
130 )
131 _LOG.debug("Current versions: %s", current_versions)
133 session = self._session_factory.session()
134 self._connection_context = ConnectionContext(
135 session, self._config, self.schema.tableSchemas, current_versions
136 )
138 if _LOG.isEnabledFor(logging.DEBUG):
139 _LOG.debug("ApdbCassandra Configuration: %s", self._connection_context.config.model_dump())
141 return self._connection_context
143 def _timer(self, name: str, *, tags: Mapping[str, str | int] | None = None) -> Timer:
144 """Create `Timer` instance given its name."""
145 return Timer(name, _MON, tags=tags)
147 @classmethod
148 def apdbImplementationVersion(cls) -> VersionTuple:
149 """Return version number for current APDB implementation.
151 Returns
152 -------
153 version : `VersionTuple`
154 Version of the code defined in implementation class.
155 """
156 return VERSION
158 def getConfig(self) -> ApdbCassandraConfig:
159 # docstring is inherited from a base class
160 return self._context.config
162 def tableDef(self, table: ApdbTables) -> Table | None:
163 # docstring is inherited from a base class
164 return self.schema.tableSchemas.get(table)
166 @classmethod
167 def init_database(
168 cls,
169 hosts: tuple[str, ...],
170 keyspace: str,
171 *,
172 schema_file: str | None = None,
173 ss_schema_file: str | None = None,
174 read_sources_months: int | None = None,
175 read_forced_sources_months: int | None = None,
176 enable_replica: bool = False,
177 replica_skips_diaobjects: bool = False,
178 port: int | None = None,
179 username: str | None = None,
180 prefix: str | None = None,
181 part_pixelization: str | None = None,
182 part_pix_level: int | None = None,
183 time_partition_tables: bool = True,
184 time_partition_start: str | None = None,
185 time_partition_end: str | None = None,
186 read_consistency: str | None = None,
187 write_consistency: str | None = None,
188 read_timeout: int | None = None,
189 write_timeout: int | None = None,
190 ra_dec_columns: tuple[str, str] | None = None,
191 replication_factor: int | None = None,
192 drop: bool = False,
193 table_options: CreateTableOptions | None = None,
194 ) -> ApdbCassandraConfig:
195 """Initialize new APDB instance and make configuration object for it.
197 Parameters
198 ----------
199 hosts : `tuple` [`str`, ...]
200 List of host names or IP addresses for Cassandra cluster.
201 keyspace : `str`
202 Name of the keyspace for APDB tables.
203 schema_file : `str`, optional
204 Location of (YAML) configuration file with APDB schema. If not
205 specified then default location will be used.
206 ss_schema_file : `str`, optional
207 Location of (YAML) configuration file with SSO schema. If not
208 specified then default location will be used.
209 read_sources_months : `int`, optional
210 Number of months of history to read from DiaSource.
211 read_forced_sources_months : `int`, optional
212 Number of months of history to read from DiaForcedSource.
213 enable_replica : `bool`, optional
214 If True, make additional tables used for replication to PPDB.
215 replica_skips_diaobjects : `bool`, optional
216 If `True` then do not fill regular ``DiaObject`` table when
217 ``enable_replica`` is `True`.
218 port : `int`, optional
219 Port number to use for Cassandra connections.
220 username : `str`, optional
221 User name for Cassandra connections.
222 prefix : `str`, optional
223 Optional prefix for all table names.
224 part_pixelization : `str`, optional
225 Name of the MOC pixelization used for partitioning.
226 part_pix_level : `int`, optional
227 Pixelization level.
228 time_partition_tables : `bool`, optional
229 Create per-partition tables.
230 time_partition_start : `str`, optional
231 Starting time for per-partition tables, in yyyy-mm-ddThh:mm:ss
232 format, in TAI.
233 time_partition_end : `str`, optional
234 Ending time for per-partition tables, in yyyy-mm-ddThh:mm:ss
235 format, in TAI.
236 read_consistency : `str`, optional
237 Name of the consistency level for read operations.
238 write_consistency : `str`, optional
239 Name of the consistency level for write operations.
240 read_timeout : `int`, optional
241 Read timeout in seconds.
242 write_timeout : `int`, optional
243 Write timeout in seconds.
244 ra_dec_columns : `tuple` [`str`, `str`], optional
245 Names of ra/dec columns in DiaObject table.
246 replication_factor : `int`, optional
247 Replication factor used when creating new keyspace, if keyspace
248 already exists its replication factor is not changed.
249 drop : `bool`, optional
250 If `True` then drop existing tables before re-creating the schema.
251 table_options : `CreateTableOptions`, optional
252 Options used when creating Cassandra tables.
254 Returns
255 -------
256 config : `ApdbCassandraConfig`
257 Resulting configuration object for a created APDB instance.
258 """
259 # Some non-standard defaults for connection parameters, these can be
260 # changed later in generated config. Check Cassandra driver
261 # documentation for what these parameters do. These parameters are not
262 # used during database initialization, but they will be saved with
263 # generated config.
264 connection_config = ApdbCassandraConnectionConfig(
265 extra_parameters={
266 "idle_heartbeat_interval": 0,
267 "idle_heartbeat_timeout": 30,
268 "control_connection_timeout": 100,
269 },
270 )
271 config = ApdbCassandraConfig(
272 contact_points=hosts,
273 keyspace=keyspace,
274 enable_replica=enable_replica,
275 replica_skips_diaobjects=replica_skips_diaobjects,
276 connection_config=connection_config,
277 )
278 config.partitioning.time_partition_tables = time_partition_tables
279 if schema_file is not None:
280 config.schema_file = schema_file
281 if ss_schema_file is not None:
282 config.ss_schema_file = ss_schema_file
283 if read_sources_months is not None:
284 config.read_sources_months = read_sources_months
285 if read_forced_sources_months is not None:
286 config.read_forced_sources_months = read_forced_sources_months
287 if port is not None:
288 config.connection_config.port = port
289 if username is not None:
290 config.connection_config.username = username
291 if prefix is not None:
292 config.prefix = prefix
293 if part_pixelization is not None:
294 config.partitioning.part_pixelization = part_pixelization
295 if part_pix_level is not None:
296 config.partitioning.part_pix_level = part_pix_level
297 if time_partition_start is not None:
298 config.partitioning.time_partition_start = time_partition_start
299 if time_partition_end is not None:
300 config.partitioning.time_partition_end = time_partition_end
301 if read_consistency is not None:
302 config.connection_config.read_consistency = read_consistency
303 if write_consistency is not None:
304 config.connection_config.write_consistency = write_consistency
305 if read_timeout is not None:
306 config.connection_config.read_timeout = read_timeout
307 if write_timeout is not None:
308 config.connection_config.write_timeout = write_timeout
309 if ra_dec_columns is not None:
310 config.ra_dec_columns = ra_dec_columns
312 cls._makeSchema(config, drop=drop, replication_factor=replication_factor, table_options=table_options)
314 return config
316 def get_replica(self) -> ApdbCassandraReplica:
317 """Return `ApdbReplica` instance for this database."""
318 # Note that this instance has to stay alive while replica exists, so
319 # we pass reference to self.
320 return ApdbCassandraReplica(self)
322 @classmethod
323 def _makeSchema(
324 cls,
325 config: ApdbConfig,
326 *,
327 drop: bool = False,
328 replication_factor: int | None = None,
329 table_options: CreateTableOptions | None = None,
330 ) -> None:
331 # docstring is inherited from a base class
333 if not isinstance(config, ApdbCassandraConfig):
334 raise TypeError(f"Unexpected type of configuration object: {type(config)}")
336 simple_schema = ApdbSchema(config.schema_file, config.ss_schema_file)
338 with SessionContext(config) as session:
339 schema = ApdbCassandraSchema(
340 session=session,
341 keyspace=config.keyspace,
342 table_schemas=simple_schema.tableSchemas,
343 prefix=config.prefix,
344 time_partition_tables=config.partitioning.time_partition_tables,
345 enable_replica=config.enable_replica,
346 replica_skips_diaobjects=config.replica_skips_diaobjects,
347 )
349 # Ask schema to create all tables.
350 part_range_config: ApdbCassandraTimePartitionRange | None = None
351 if config.partitioning.time_partition_tables:
352 partitioner = Partitioner(config)
353 time_partition_start = astropy.time.Time(
354 config.partitioning.time_partition_start, format="isot", scale="tai"
355 )
356 time_partition_end = astropy.time.Time(
357 config.partitioning.time_partition_end, format="isot", scale="tai"
358 )
359 part_range_config = ApdbCassandraTimePartitionRange(
360 start=partitioner.time_partition(time_partition_start),
361 end=partitioner.time_partition(time_partition_end),
362 )
363 schema.makeSchema(
364 drop=drop,
365 part_range=part_range_config,
366 replication_factor=replication_factor,
367 table_options=table_options,
368 )
369 else:
370 schema.makeSchema(
371 drop=drop, replication_factor=replication_factor, table_options=table_options
372 )
374 meta_table_name = ApdbTables.metadata.table_name(config.prefix)
375 metadata = ApdbMetadataCassandra(
376 session, meta_table_name, config.keyspace, "read_tuples", "write"
377 )
379 # Fill version numbers, overrides if they existed before.
380 metadata.set(
381 ConnectionContext.metadataSchemaVersionKey, str(simple_schema.schemaVersion()), force=True
382 )
383 metadata.set(
384 ConnectionContext.metadataCodeVersionKey, str(cls.apdbImplementationVersion()), force=True
385 )
387 if config.enable_replica:
388 # Only store replica code version if replica is enabled.
389 metadata.set(
390 ConnectionContext.metadataReplicaVersionKey,
391 str(ApdbCassandraReplica.apdbReplicaImplementationVersion()),
392 force=True,
393 )
395 # Store frozen part of a configuration in metadata.
396 freezer = ApdbConfigFreezer[ApdbCassandraConfig](ConnectionContext.frozen_parameters)
397 metadata.set(ConnectionContext.metadataConfigKey, freezer.to_json(config), force=True)
399 # Store time partition range.
400 if part_range_config:
401 part_range_config.save_to_meta(metadata)
403 def getDiaObjects(self, region: sphgeom.Region) -> pandas.DataFrame:
404 # docstring is inherited from a base class
405 context = self._context
406 config = context.config
408 sp_where, num_sp_part = context.partitioner.spatial_where(region, for_prepare=True)
409 _LOG.debug("getDiaObjects: #partitions: %s", len(sp_where))
411 # We need to exclude extra partitioning columns from result.
412 column_names = context.schema.apdbColumnNames(ApdbTables.DiaObjectLast)
413 what = ",".join(quote_id(column) for column in column_names)
415 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
416 query = f'SELECT {what} from "{self._keyspace}"."{table_name}"'
417 statements: list[tuple] = []
418 for where, params in sp_where:
419 full_query = f"{query} WHERE {where}"
420 if params:
421 statement = context.preparer.prepare(full_query)
422 else:
423 # If there are no params then it is likely that query has a
424 # bunch of literals rendered already, no point trying to
425 # prepare it because it's not reusable.
426 statement = cassandra.query.SimpleStatement(full_query)
427 statements.append((statement, params))
428 _LOG.debug("getDiaObjects: #queries: %s", len(statements))
430 with self._timer("select_time", tags={"table": "DiaObject", "method": "getDiaObjects"}) as timer:
431 raw_objects = cast(
432 ApdbCassandraTableData,
433 select_concurrent(
434 context.session,
435 statements,
436 "read_raw_multi",
437 config.connection_config.read_concurrency,
438 ),
439 )
440 objects = raw_objects.to_pandas(context.schema._table_schema(ApdbTables.DiaObjectLast))
441 timer.add_values(row_count=len(objects), num_sp_part=num_sp_part, num_queries=len(statements))
443 _LOG.debug("found %s DiaObjects", objects.shape[0])
444 return objects
446 def getDiaSources(
447 self,
448 region: sphgeom.Region,
449 object_ids: Iterable[int] | None,
450 visit_time: astropy.time.Time,
451 start_time: astropy.time.Time | None = None,
452 ) -> pandas.DataFrame | None:
453 # docstring is inherited from a base class
454 context = self._context
455 config = context.config
457 months = config.read_sources_months
458 if start_time is None and months == 0:
459 return None
461 mjd_end = float(visit_time.tai.mjd)
462 if start_time is None:
463 mjd_start = mjd_end - months * 30
464 else:
465 mjd_start = float(start_time.tai.mjd)
467 return self._getSources(region, object_ids, mjd_start, mjd_end, ApdbTables.DiaSource)
469 def getDiaForcedSources(
470 self,
471 region: sphgeom.Region,
472 object_ids: Iterable[int] | None,
473 visit_time: astropy.time.Time,
474 start_time: astropy.time.Time | None = None,
475 ) -> pandas.DataFrame | None:
476 # docstring is inherited from a base class
477 context = self._context
478 config = context.config
480 months = config.read_forced_sources_months
481 if start_time is None and months == 0:
482 return None
484 mjd_end = float(visit_time.tai.mjd)
485 if start_time is None:
486 mjd_start = mjd_end - months * 30
487 else:
488 mjd_start = float(start_time.tai.mjd)
490 return self._getSources(region, object_ids, mjd_start, mjd_end, ApdbTables.DiaForcedSource)
492 def getDiaObjectsForDedup(self, since: astropy.time.Time | None = None) -> pandas.DataFrame:
493 # docstring is inherited from a base class
494 context = self._context
495 config = context.config
497 if not context.has_dedup_table:
498 raise TypeError("DiaObjectDedup table does not exist in this APDB instance.")
500 if since is None:
501 # Read last deduplication time from metadata.
502 dedup_str = context.metadata.get(context.metadataDedupKey)
503 if dedup_str is not None:
504 dedup_state = json.loads(dedup_str)
505 dedup_time_str = dedup_state["dedup_time_iso_tai"]
506 since = astropy.time.Time(dedup_time_str, format="iso", scale="tai")
508 column_names = context.schema.apdbColumnNames(ExtraTables.DiaObjectDedup)
509 what = ",".join(quote_id(column) for column in column_names)
511 validity_start_column = self._timestamp_column_name("validityStart")
512 timestamp = None if since is None else self._timestamp_column_value(since)
514 table_name = context.schema.tableName(ExtraTables.DiaObjectDedup)
515 query = f'SELECT {what} FROM "{self._keyspace}"."{table_name}" WHERE dedup_part = %s'
516 if since is not None:
517 query += f' AND "{validity_start_column}" >= %s'
518 query += " ALLOW FILTERING"
520 statement = cassandra.query.SimpleStatement(query)
522 num_part = config.partitioning.num_part_dedup
523 statements = []
524 for dedup_part in range(num_part):
525 params = (dedup_part,) if timestamp is None else (dedup_part, timestamp)
526 statements.append((statement, params))
528 with self._timer(
529 "select_time", tags={"table": "DiaObjectDedup", "method": "getDiaObjectsForDedup"}
530 ) as timer:
531 objects_raw = cast(
532 ApdbCassandraTableData,
533 select_concurrent(
534 context.session,
535 statements,
536 "read_raw_multi_dedup",
537 config.connection_config.read_concurrency,
538 ),
539 )
540 objects = objects_raw.to_pandas(context.schema._table_schema(ExtraTables.DiaObjectDedup))
541 timer.add_values(row_count=len(objects), num_queries=num_part)
543 _LOG.debug("found %s DiaObjectDedup records", objects.shape[0])
544 return objects
546 def getDiaSourcesForDiaObjects(
547 self, objects: list[DiaObjectId], start_time: astropy.time.Time, max_dist_arcsec: float = 1.0
548 ) -> pandas.DataFrame:
549 # docstring is inherited from a base class
550 context = self._context
551 config = context.config
553 # Which tables to query and temporal constraints.
554 end_time = self._current_time()
555 tables, temporal_where = context.partitioner.temporal_where(
556 ApdbTables.DiaSource,
557 start_time,
558 end_time,
559 for_prepare=False,
560 partitons_range=context.time_partitions_range,
561 query_per_time_part=False,
562 )
563 if not tables:
564 warnings.warn(
565 f"Query time range ({start_time.isot} - {end_time.isot}) does not overlap database "
566 "time partitions."
567 )
569 # Group DiaObjects by partition.
570 partitioned_object_ids = self._group_dia_objects_by_partition(
571 context.partitioner, objects, max_dist_arcsec
572 )
574 # Columns to return.
575 column_names = context.schema.apdbColumnNames(ApdbTables.DiaSource)
576 what = ",".join(quote_id(column) for column in column_names)
578 # Make a bunch of queries.
579 statements = []
580 suffix = "ALLOW FILTERING"
581 for apdb_part, diaObjectIds in partitioned_object_ids.items():
582 spatial_where = [(f"apdb_part = {apdb_part}", ())]
583 for table in tables:
584 prefix = f'SELECT {what} FROM "{self._keyspace}"."{table}"'
585 for id_chunk in chunk_iterable(diaObjectIds, 10_000):
586 chunk_str = ",".join(str(diaObjectId) for diaObjectId in id_chunk)
587 id_where = (f'"diaObjectId" IN ({chunk_str})', ())
588 statements += list(
589 self._combine_where(prefix, spatial_where, temporal_where, id_where, suffix=suffix)
590 )
592 _LOG.debug("getDiaSourcesForDiaObjects #queries: %s", len(statements))
594 with self._timer(
595 "select_time", tags={"table": "DiaSource", "method": "getDiaSourcesForDiaObjects"}
596 ) as timer:
597 table_data_raw = cast(
598 ApdbCassandraTableData,
599 select_concurrent(
600 context.session,
601 statements,
602 "read_raw_multi",
603 config.connection_config.read_concurrency,
604 ),
605 )
606 catalog = table_data_raw.to_pandas(context.schema._table_schema(ApdbTables.DiaSource))
607 timer.add_values(row_count_from_db=len(catalog), num_queries=len(statements))
609 # precise filtering on midpointMjdTai
610 catalog = cast(pandas.DataFrame, catalog[catalog["midpointMjdTai"] >= start_time.tai.mjd])
612 timer.add_values(row_count=len(catalog))
614 _LOG.debug("found %d DiaSources", len(catalog))
615 return catalog
617 def containsVisitDetector(
618 self,
619 visit: int,
620 detector: int,
621 region: sphgeom.Region,
622 visit_time: astropy.time.Time,
623 ) -> bool:
624 # docstring is inherited from a base class
625 context = self._context
627 table_name = context.schema.tableName(ExtraTables.ApdbVisitDetector)
628 query = f'SELECT count(*) FROM "{self._keyspace}"."{table_name}" WHERE visit = %s AND detector = %s'
629 with self._timer("contains_visit_detector_time", tags={"table": table_name}):
630 result = context.session.execute(query, (visit, detector))
631 return bool(result.one()[0])
633 def store(
634 self,
635 visit_time: astropy.time.Time,
636 objects: pandas.DataFrame,
637 sources: pandas.DataFrame | None = None,
638 forced_sources: pandas.DataFrame | None = None,
639 ) -> None:
640 # docstring is inherited from a base class
641 context = self._context
642 config = context.config
644 # Store visit/detector in a special table, this has to be done
645 # before all other writes so if there is a failure at any point
646 # later we still have a record for attempted write.
647 visit_detector: set[tuple[int, int]] = set()
648 for df in sources, forced_sources:
649 if df is not None and not df.empty:
650 df = df[["visit", "detector"]]
651 for visit, detector in df.itertuples(index=False):
652 visit_detector.add((visit, detector))
654 if visit_detector:
655 # Typically there is only one entry, do not bother with
656 # concurrency.
657 table_name = context.schema.tableName(ExtraTables.ApdbVisitDetector)
658 query = f'INSERT INTO "{self._keyspace}"."{table_name}" (visit, detector) VALUES (%s, %s)'
659 for item in visit_detector:
660 context.session.execute(query, item, execution_profile="write")
662 objects = self._fix_input_timestamps(objects)
663 if sources is not None:
664 sources = self._fix_input_timestamps(sources)
665 if forced_sources is not None:
666 forced_sources = self._fix_input_timestamps(forced_sources)
668 replica_chunk: ReplicaChunk | None = None
669 if context.schema.replication_enabled:
670 replica_chunk = ReplicaChunk.make_replica_chunk(visit_time, config.replica_chunk_seconds)
671 self._storeReplicaChunk(replica_chunk)
673 # fill region partition column for DiaObjects
674 objects = self._add_apdb_part(objects)
675 self._storeDiaObjects(objects, visit_time, replica_chunk)
677 if sources is not None and len(sources) > 0:
678 # copy apdb_part column from DiaObjects to DiaSources
679 sources = self._add_apdb_part(sources)
680 subchunk = self._storeDiaSources(ApdbTables.DiaSource, sources, replica_chunk)
681 self._storeDiaSourcesPartitions(sources, visit_time, replica_chunk, subchunk)
683 if forced_sources is not None and len(forced_sources) > 0:
684 forced_sources = self._add_apdb_part(forced_sources)
685 self._storeDiaSources(ApdbTables.DiaForcedSource, forced_sources, replica_chunk)
687 def reassignDiaSourcesToDiaObjects(
688 self,
689 idMap: Mapping[DiaSourceId, int],
690 *,
691 increment_nDiaSources: bool = True,
692 decrement_nDiaSources: bool = True,
693 ) -> None:
694 # docstring is inherited from a base class
695 context = self._context
696 config = context.config
698 source_ids = {source_id.diaSourceId for source_id in idMap}
700 # Find all DiaSources.
701 found_sources = self._get_diasource_data(
702 idMap, "apdb_part", "diaObjectId", "ra", "dec", "midpointMjdTai"
703 )
705 if missing_ids := (source_ids - {row.diaSourceId for row in found_sources}):
706 raise LookupError(f"Some source IDs were not found in DiaSource table: {missing_ids}")
708 found_sources_by_id = {row.diaSourceId: row for row in found_sources}
710 # Make sure that all DiaObjects exist, we also want to know
711 # nDiaSources count for current and new records because we want to
712 # send updated values to replica.
713 current_object_ids = {
714 DiaObjectId(diaObjectId=row.diaObjectId, ra=row.ra, dec=row.dec) for row in found_sources
715 }
716 # Assume that DiaSource ra/dec are very close to re-assigned objects.
717 new_object_ids = {
718 DiaObjectId(diaObjectId=diaObjectId, ra=source_id.ra, dec=source_id.dec)
719 for source_id, diaObjectId in idMap.items()
720 }
721 all_object_ids = new_object_ids | current_object_ids
722 found_objects = self._get_diaobject_data(all_object_ids, "apdb_part", "ra", "dec", "nDiaSources")
724 if missing_ids := (
725 {row.diaObjectId for row in all_object_ids} - {row.diaObjectId for row in found_objects}
726 ):
727 raise LookupError(f"Some object IDs were not found in DiaObjectLast table: {missing_ids}")
729 update_records: list[ApdbUpdateRecord] = []
730 update_order = 0
731 current_time = self._current_time()
732 current_time_ns = int(current_time.unix_tai * 1e9)
734 # Update DiaSources.
735 table_name = context.schema.tableName(ApdbTables.DiaSource)
736 statements: list[tuple] = []
737 for source_id, diaObjectId in idMap.items():
738 source_row = found_sources_by_id[source_id.diaSourceId]
739 apdb_part = source_row.apdb_part
740 time_part = context.partitioner.time_partition(source_row.midpointMjdTai)
742 if config.partitioning.time_partition_tables:
743 statement = context.preparer.prepare(
744 f'UPDATE "{self._keyspace}"."{table_name}_{time_part}" SET "diaObjectId" = ? '
745 'WHERE apdb_part = ? AND "diaSourceId" = ?'
746 )
747 statements.append((statement, (diaObjectId, apdb_part, source_id.diaSourceId)))
748 else:
749 statement = context.preparer.prepare(
750 f'UPDATE "{self._keyspace}"."{table_name}" SET "diaObjectId" = ? '
751 'WHERE apdb_part = ? AND apdb_time_part = ? AND "diaSourceId" = ?'
752 )
753 statements.append((statement, (diaObjectId, apdb_part, time_part, source_id.diaSourceId)))
755 if context.schema.replication_enabled:
756 update_records.append(
757 ApdbReassignDiaSourceToDiaObjectRecord(
758 diaSourceId=source_id.diaSourceId,
759 ra=source_id.ra,
760 dec=source_id.dec,
761 midpointMjdTai=source_id.midpointMjdTai,
762 diaObjectId=diaObjectId,
763 update_time_ns=current_time_ns,
764 update_order=update_order,
765 )
766 )
767 update_order += 1
769 with self._timer(
770 "update_time", tags={"table": table_name, "method": "reassignDiaSourcesToDiaObjects"}
771 ) as timer:
772 execute_concurrent(context.session, statements, execution_profile="write")
773 timer.add_values(num_queries=len(statements))
775 # Update nDiaSources in DiaObjectLast. We do not update DiaObject table
776 # here because it may not even exist. PPDB updates DiaObject from
777 # update records.
778 if increment_nDiaSources or decrement_nDiaSources:
779 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
780 statement = context.preparer.prepare(
781 f'UPDATE "{self._keyspace}"."{table_name}" SET "nDiaSources" = ? '
782 'WHERE apdb_part = ? AND "diaObjectId" = ?'
783 )
784 statements = []
786 # Calculate increments/decrements for all affected DiaObjects.
787 increments: Counter = Counter()
788 if increment_nDiaSources:
789 increments.update(idMap.values())
790 if decrement_nDiaSources:
791 increments.subtract(row.diaObjectId for row in found_sources)
793 for row in found_objects:
794 if increments.get(row.diaObjectId):
795 nDiaSources = row.nDiaSources + increments[row.diaObjectId]
796 statements.append((statement, (nDiaSources, row.apdb_part, row.diaObjectId)))
798 # Also send updated values to replica.
799 if context.schema.replication_enabled:
800 update_records.append(
801 ApdbUpdateNDiaSourcesRecord(
802 diaObjectId=row.diaObjectId,
803 ra=row.ra,
804 dec=row.dec,
805 nDiaSources=nDiaSources,
806 update_time_ns=current_time_ns,
807 update_order=update_order,
808 )
809 )
810 update_order += 1
812 if statements:
813 with self._timer(
814 "update_time", tags={"table": table_name, "method": "reassignDiaSourcesToDiaObjects"}
815 ) as timer:
816 execute_concurrent(context.session, statements, execution_profile="write")
817 timer.add_values(num_queries=len(statements))
819 if update_records:
820 replica_chunk = ReplicaChunk.make_replica_chunk(current_time, config.replica_chunk_seconds)
821 self._storeUpdateRecords(update_records, replica_chunk, store_chunk=True)
823 def setValidityEnd(
824 self, objects: list[DiaObjectId], validityEnd: astropy.time.Time, raise_on_missing_id: bool = False
825 ) -> int:
826 # docstring is inherited from a base class
827 if not objects:
828 return 0
830 context = self._context
831 config = context.config
833 pad_arcsec = 1.0
834 partitioned_object_ids = self._group_dia_objects_by_partition(
835 context.partitioner, objects, pad_arcsec
836 )
838 # Check that all objects exist.
839 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
840 statements: list[tuple] = []
841 for apdb_part, diaObjectIds in partitioned_object_ids.items():
842 id_str = ",".join(str(diaObjectId) for diaObjectId in diaObjectIds)
843 query = (
844 f'SELECT apdb_part, "diaObjectId" FROM "{self._keyspace}"."{table_name}" '
845 f'WHERE apdb_part = %s AND "diaObjectId" IN ({id_str})'
846 )
847 statement = cassandra.query.SimpleStatement(query)
848 statements.append((statement, (apdb_part,)))
850 with self._timer("select_time", tags={"table": table_name, "method": "setValidityEnd"}) as timer:
851 records = cast(
852 list[tuple[int, int]],
853 select_concurrent(
854 context.session,
855 statements,
856 "read_tuples",
857 config.connection_config.read_concurrency,
858 ),
859 )
860 timer.add_values(row_count=len(objects))
862 requested_ids = {obj.diaObjectId for obj in objects}
863 found_ids = {rec[1] for rec in records}
864 if extra_ids := (found_ids - requested_ids):
865 raise RuntimeError(f"Consistency error - found duplicate records for object IDs: {extra_ids}")
866 if raise_on_missing_id:
867 if missing_ids := (requested_ids - found_ids):
868 raise LookupError(f"Some object IDs are missing from DiaObjectLast table: {missing_ids}")
870 # Filter existing records.
871 if len(objects) != len(found_ids):
872 objects = [obj for obj in objects if obj.diaObjectId in found_ids]
874 if not objects:
875 return 0
877 # Group by partitions again.
878 grouped_object_ids: dict[int, list[int]] = defaultdict(list)
879 for apdb_part, diaObjectId in records:
880 grouped_object_ids[apdb_part].append(diaObjectId)
882 # Remove all matching rows from DiaObjectLast.
883 statements = []
884 for apdb_part, diaObjectIds in grouped_object_ids.items():
885 id_str = ",".join(str(diaObjectId) for diaObjectId in diaObjectIds)
886 query = (
887 f'DELETE FROM "{self._keyspace}"."{table_name}" '
888 f'WHERE apdb_part = %s AND "diaObjectId" IN ({id_str})'
889 )
890 statement = cassandra.query.SimpleStatement(query)
891 statements.append((statement, (apdb_part,)))
893 # Also remove from DiaObjectLastToPartition.
894 reverse_table_name = context.schema.tableName(ExtraTables.DiaObjectLastToPartition)
895 id_str = ",".join(str(diaObjectId) for _, diaObjectId in records)
896 query = f'DELETE FROM "{self._keyspace}"."{reverse_table_name}" WHERE "diaObjectId" IN ({id_str})'
897 statement = cassandra.query.SimpleStatement(query)
898 statements.append((statement, ()))
900 with self._timer("delete_time", tags={"table": table_name, "method": "setValidityEnd"}) as timer:
901 execute_concurrent(context.session, statements, execution_profile="write")
902 timer.add_values(row_count=len(records))
904 # If repication is enabled then send all updates.
905 if context.schema.replication_enabled:
906 current_time = self._current_time()
907 current_time_ns = int(current_time.unix_tai * 1e9)
908 replica_chunk = ReplicaChunk.make_replica_chunk(current_time, config.replica_chunk_seconds)
910 update_records = [
911 ApdbCloseDiaObjectValidityRecord(
912 diaObjectId=obj.diaObjectId,
913 ra=obj.ra,
914 dec=obj.dec,
915 update_time_ns=current_time_ns,
916 update_order=index,
917 validityEndMjdTai=float(validityEnd.tai.mjd),
918 nDiaSources=None,
919 )
920 for index, obj in enumerate(objects)
921 ]
923 self._storeUpdateRecords(update_records, replica_chunk, store_chunk=True)
925 return len(objects)
927 def resetDedup(self, dedup_time: astropy.time.Time | None = None) -> None:
928 # docstring is inherited from a base class
929 context = self._context
931 if not context.has_dedup_table:
932 raise TypeError("DiaObjectDedup table does not exist in this APDB instance.")
934 if dedup_time is None:
935 dedup_time = self._current_time()
937 validity_start_column = self._timestamp_column_name("validityStart")
939 # Find latest timestamp in deduplication table.
940 table_name = context.schema.tableName(ExtraTables.DiaObjectDedup)
941 query = f'SELECT MAX("{validity_start_column}") FROM "{self._keyspace}"."{table_name}"'
942 result = context.session.execute(query, execution_profile="read_tuples")
943 max_value = result.one()[0]
944 if self._schema.has_mjd_timestamps:
945 max_validity_start = astropy.time.Time(max_value, format="mjd", scale="tai")
946 else:
947 max_validity_start = astropy.time.Time(max_value, format="datetime", scale="tai")
949 # If max time is lower than dedup time we can do TRUNCATE.
950 if dedup_time >= max_validity_start:
951 query = f'TRUNCATE TABLE "{self._keyspace}"."{table_name}"'
952 context.session.execute(query, execution_profile="write")
953 else:
954 dedup_time_value = self._timestamp_column_value(dedup_time)
955 query = f'DELETE FROM "{self._keyspace}"."{table_name}" WHERE "{validity_start_column}" < %s'
956 context.session.execute(query, (dedup_time_value,), execution_profile="write")
958 # Store dedup time.
959 data = {"dedup_time_iso_tai": dedup_time.tai.to_value("iso")}
960 data_json = json.dumps(data)
961 context.metadata.set(context.metadataDedupKey, data_json, force=True)
963 def reassignDiaSources(self, idMap: Mapping[int, int]) -> None:
964 # docstring is inherited from a base class
965 context = self._context
966 config = context.config
968 now = self._current_time()
969 reassign_time_column = self._timestamp_column_name("ssObjectReassocTime")
970 reassignTime = self._timestamp_column_value(now)
972 # To update a record we need to know its exact primary key (including
973 # partition key) so we start by querying for diaSourceId to find the
974 # primary keys.
976 table_name = context.schema.tableName(ExtraTables.DiaSourceToPartition)
977 # split it into 1k IDs per query
978 selects: list[tuple] = []
979 for ids in chunk_iterable(idMap.keys(), 1_000):
980 ids_str = ",".join(str(item) for item in ids)
981 selects.append(
982 (
983 (
984 'SELECT "diaSourceId", "apdb_part", "apdb_time_part", "apdb_replica_chunk" '
985 f'FROM "{self._keyspace}"."{table_name}" WHERE "diaSourceId" IN ({ids_str})'
986 ),
987 {},
988 )
989 )
991 # No need for DataFrame here, read data as tuples.
992 result = cast(
993 list[tuple[int, int, int, int | None]],
994 select_concurrent(
995 context.session, selects, "read_tuples", config.connection_config.read_concurrency
996 ),
997 )
999 # Make mapping from source ID to its partition.
1000 id2partitions: dict[int, tuple[int, int]] = {}
1001 id2chunk_id: dict[int, int] = {}
1002 for row in result:
1003 id2partitions[row[0]] = row[1:3]
1004 if row[3] is not None:
1005 id2chunk_id[row[0]] = row[3]
1007 # make sure we know partitions for each ID
1008 if set(id2partitions) != set(idMap):
1009 missing = ",".join(str(item) for item in set(idMap) - set(id2partitions))
1010 raise ValueError(f"Following DiaSource IDs do not exist in the database: {missing}")
1012 # Reassign in standard tables
1013 queries: list[tuple[cassandra.query.PreparedStatement, tuple]] = []
1014 for diaSourceId, ssObjectId in idMap.items():
1015 apdb_part, apdb_time_part = id2partitions[diaSourceId]
1016 values: tuple
1017 if config.partitioning.time_partition_tables:
1018 table_name = context.schema.tableName(ApdbTables.DiaSource, apdb_time_part)
1019 query = (
1020 f'UPDATE "{self._keyspace}"."{table_name}"'
1021 f' SET "ssObjectId" = ?, "diaObjectId" = NULL, "{reassign_time_column}" = ?'
1022 ' WHERE "apdb_part" = ? AND "diaSourceId" = ?'
1023 )
1024 values = (ssObjectId, reassignTime, apdb_part, diaSourceId)
1025 else:
1026 table_name = context.schema.tableName(ApdbTables.DiaSource)
1027 query = (
1028 f'UPDATE "{self._keyspace}"."{table_name}"'
1029 f' SET "ssObjectId" = ?, "diaObjectId" = NULL, "{reassign_time_column}" = ?'
1030 ' WHERE "apdb_part" = ? AND "apdb_time_part" = ? AND "diaSourceId" = ?'
1031 )
1032 values = (ssObjectId, reassignTime, apdb_part, apdb_time_part, diaSourceId)
1033 queries.append((context.preparer.prepare(query), values))
1035 # TODO: (DM-50190) Replication for updated records is not implemented.
1036 if id2chunk_id:
1037 warnings.warn("Replication of reassigned DiaSource records is not implemented.", stacklevel=2)
1039 _LOG.debug("%s: will update %d records", table_name, len(idMap))
1040 with self._timer("source_reassign_time") as timer:
1041 execute_concurrent(context.session, queries, execution_profile="write")
1042 timer.add_values(source_count=len(idMap))
1044 def countUnassociatedObjects(self) -> int:
1045 # docstring is inherited from a base class
1047 # It's too inefficient to implement it for Cassandra in current schema.
1048 raise NotImplementedError()
1050 @property
1051 def schema(self) -> ApdbSchema:
1052 # docstring is inherited from a base class
1053 return self._schema
1055 @property
1056 def metadata(self) -> ApdbMetadata:
1057 # docstring is inherited from a base class
1058 context = self._context
1059 return context.metadata
1061 @property
1062 def admin(self) -> ApdbCassandraAdmin:
1063 # docstring is inherited from a base class
1064 return ApdbCassandraAdmin(self)
1066 def _getSources(
1067 self,
1068 region: sphgeom.Region,
1069 object_ids: Iterable[int] | None,
1070 mjd_start: float,
1071 mjd_end: float,
1072 table_name: ApdbTables,
1073 ) -> pandas.DataFrame:
1074 """Return catalog of DiaSource instances given set of DiaObject IDs.
1076 Parameters
1077 ----------
1078 region : `lsst.sphgeom.Region`
1079 Spherical region.
1080 object_ids :
1081 Collection of DiaObject IDs
1082 mjd_start : `float`
1083 Lower bound of time interval.
1084 mjd_end : `float`
1085 Upper bound of time interval.
1086 table_name : `ApdbTables`
1087 Name of the table.
1089 Returns
1090 -------
1091 catalog : `pandas.DataFrame`, or `None`
1092 Catalog containing DiaSource records. Empty catalog is returned if
1093 ``object_ids`` is empty.
1094 """
1095 context = self._context
1096 config = context.config
1098 object_id_set: Set[int] = set()
1099 if object_ids is not None:
1100 object_id_set = set(object_ids)
1101 if len(object_id_set) == 0:
1102 return self._make_empty_catalog(table_name)
1104 sp_where, num_sp_part = context.partitioner.spatial_where(region, for_prepare=True)
1105 tables, temporal_where = context.partitioner.temporal_where(
1106 table_name, mjd_start, mjd_end, for_prepare=True, partitons_range=context.time_partitions_range
1107 )
1108 if not tables:
1109 start = astropy.time.Time(mjd_start, format="mjd", scale="tai")
1110 end = astropy.time.Time(mjd_end, format="mjd", scale="tai")
1111 warnings.warn(
1112 f"Query time range ({start.isot} - {end.isot}) does not overlap database time partitions."
1113 )
1115 # We need to exclude extra partitioning columns from result.
1116 column_names = context.schema.apdbColumnNames(table_name)
1117 what = ",".join(quote_id(column) for column in column_names)
1119 # Build all queries
1120 statements: list[tuple] = []
1121 for table in tables:
1122 prefix = f'SELECT {what} from "{self._keyspace}"."{table}"'
1123 statements += list(self._combine_where(prefix, sp_where, temporal_where))
1124 _LOG.debug("_getSources %s: #queries: %s", table_name, len(statements))
1126 with self._timer("select_time", tags={"table": table_name.name, "method": "_getSources"}) as timer:
1127 table_data_raw = cast(
1128 ApdbCassandraTableData,
1129 select_concurrent(
1130 context.session,
1131 statements,
1132 "read_raw_multi",
1133 config.connection_config.read_concurrency,
1134 ),
1135 )
1136 catalog = table_data_raw.to_pandas(context.schema._table_schema(table_name))
1137 timer.add_values(
1138 row_count_from_db=len(catalog), num_sp_part=num_sp_part, num_queries=len(statements)
1139 )
1141 # filter by given object IDs
1142 if len(object_id_set) > 0:
1143 catalog = cast(pandas.DataFrame, catalog[catalog["diaObjectId"].isin(object_id_set)])
1145 # precise filtering on midpointMjdTai
1146 catalog = cast(pandas.DataFrame, catalog[catalog["midpointMjdTai"] > mjd_start])
1148 timer.add_values(row_count=len(catalog))
1150 _LOG.debug("found %d %ss", catalog.shape[0], table_name.name)
1151 return catalog
1153 def _storeReplicaChunk(self, replica_chunk: ReplicaChunk) -> None:
1154 context = self._context
1155 config = context.config
1157 # Cassandra timestamp uses milliseconds since epoch
1158 timestamp = int(replica_chunk.last_update_time.unix_tai * 1000)
1160 # everything goes into a single partition
1161 partition = 0
1163 table_name = context.schema.tableName(ExtraTables.ApdbReplicaChunks)
1165 columns = ["partition", "apdb_replica_chunk", "last_update_time", "unique_id"]
1166 values = [partition, replica_chunk.id, timestamp, replica_chunk.unique_id]
1167 if context.has_chunk_sub_partitions:
1168 columns.append("has_subchunks")
1169 values.append(True)
1171 column_list = ", ".join(columns)
1172 placeholders = ",".join(["%s"] * len(columns))
1173 query = f'INSERT INTO "{self._keyspace}"."{table_name}" ({column_list}) VALUES ({placeholders})'
1175 context.session.execute(
1176 query,
1177 values,
1178 timeout=config.connection_config.write_timeout,
1179 execution_profile="write",
1180 )
1182 def _queryDiaObjectLastPartitions(self, ids: Iterable[int]) -> Mapping[int, int]:
1183 """Return existing mapping of diaObjectId to its last partition."""
1184 context = self._context
1185 config = context.config
1187 table_name = context.schema.tableName(ExtraTables.DiaObjectLastToPartition)
1188 queries = []
1189 object_count = 0
1190 for id_chunk in chunk_iterable(ids, 10_000):
1191 id_chunk_list = list(id_chunk)
1192 query = (
1193 f'SELECT "diaObjectId", apdb_part FROM "{self._keyspace}"."{table_name}" '
1194 f'WHERE "diaObjectId" in ({",".join(str(oid) for oid in id_chunk_list)})'
1195 )
1196 queries.append((query, ()))
1197 object_count += len(id_chunk_list)
1199 with self._timer("query_object_last_partitions", tags={"table": table_name}) as timer:
1200 data = cast(
1201 ApdbTableData,
1202 select_concurrent(
1203 context.session,
1204 queries,
1205 "read_raw_multi",
1206 config.connection_config.read_concurrency,
1207 ),
1208 )
1209 timer.add_values(object_count=object_count, row_count=len(data.rows()))
1211 if data.column_names() != ["diaObjectId", "apdb_part"]:
1212 raise RuntimeError(f"Unexpected column names in query result: {data.column_names()}")
1214 return {row[0]: row[1] for row in data.rows()}
1216 def _deleteMovingObjects(self, objs: pandas.DataFrame) -> None:
1217 """Objects in DiaObjectsLast can move from one spatial partition to
1218 another. For those objects inserting new version does not replace old
1219 one, so we need to explicitly remove old versions before inserting new
1220 ones.
1221 """
1222 context = self._context
1224 # Extract all object IDs.
1225 new_partitions = dict(zip(objs["diaObjectId"], objs["apdb_part"]))
1226 old_partitions = self._queryDiaObjectLastPartitions(objs["diaObjectId"])
1228 moved_oids: dict[int, tuple[int, int]] = {}
1229 for oid, old_part in old_partitions.items():
1230 new_part = new_partitions.get(oid, old_part)
1231 if new_part != old_part:
1232 moved_oids[oid] = (old_part, new_part)
1233 _LOG.debug("DiaObject IDs that moved to new partition: %s", moved_oids)
1235 if moved_oids:
1236 # Delete old records from DiaObjectLast.
1237 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
1238 query = f'DELETE FROM "{self._keyspace}"."{table_name}" WHERE apdb_part = ? AND "diaObjectId" = ?'
1239 statement = context.preparer.prepare(query)
1240 queries = []
1241 for oid, (old_part, _) in moved_oids.items():
1242 queries.append((statement, (old_part, oid)))
1243 with self._timer("delete_object_last", tags={"table": table_name}) as timer:
1244 execute_concurrent(context.session, queries, execution_profile="write")
1245 timer.add_values(row_count=len(moved_oids))
1247 # Add all new records to the map.
1248 table_name = context.schema.tableName(ExtraTables.DiaObjectLastToPartition)
1249 query = f'INSERT INTO "{self._keyspace}"."{table_name}" ("diaObjectId", apdb_part) VALUES (?,?)'
1250 statement = context.preparer.prepare(query)
1252 queries = []
1253 for oid, new_part in new_partitions.items():
1254 queries.append((statement, (oid, new_part)))
1256 with self._timer("update_object_last_partition", tags={"table": table_name}) as timer:
1257 execute_concurrent(context.session, queries, execution_profile="write")
1258 timer.add_values(row_count=len(queries))
1260 def _storeDiaObjects(
1261 self, objs: pandas.DataFrame, visit_time: astropy.time.Time, replica_chunk: ReplicaChunk | None
1262 ) -> None:
1263 """Store catalog of DiaObjects from current visit.
1265 Parameters
1266 ----------
1267 objs : `pandas.DataFrame`
1268 Catalog with DiaObject records
1269 visit_time : `astropy.time.Time`
1270 Time of the current visit.
1271 replica_chunk : `ReplicaChunk` or `None`
1272 Replica chunk identifier if replication is configured.
1273 """
1274 if len(objs) == 0:
1275 _LOG.debug("No objects to write to database.")
1276 return
1278 context = self._context
1279 config = context.config
1281 self._deleteMovingObjects(objs)
1283 validity_start_column = self._timestamp_column_name("validityStart")
1284 timestamp = self._timestamp_column_value(visit_time)
1286 # DiaObjectLast did not have this column in the past.
1287 extra_columns: dict[str, Any] = {}
1288 if context.schema.check_column(ApdbTables.DiaObjectLast, validity_start_column):
1289 extra_columns[validity_start_column] = timestamp
1291 self._storeObjectsPandas(objs, ApdbTables.DiaObjectLast, extra_columns=extra_columns)
1293 extra_columns[validity_start_column] = timestamp
1294 visit_time_part = context.partitioner.time_partition(visit_time)
1295 time_part: int | None = visit_time_part
1296 if (time_partitions_range := context.time_partitions_range) is not None:
1297 self._check_time_partitions([visit_time_part], time_partitions_range)
1298 if not config.partitioning.time_partition_tables:
1299 extra_columns["apdb_time_part"] = time_part
1300 time_part = None
1302 # Only store DiaObects if not doing replication or explicitly
1303 # configured to always store them.
1304 if replica_chunk is None or not config.replica_skips_diaobjects:
1305 self._storeObjectsPandas(
1306 objs, ApdbTables.DiaObject, extra_columns=extra_columns, time_part=time_part
1307 )
1309 if replica_chunk is not None:
1310 extra_columns = {"apdb_replica_chunk": replica_chunk.id, validity_start_column: timestamp}
1311 table = ExtraTables.DiaObjectChunks
1312 if context.has_chunk_sub_partitions:
1313 table = ExtraTables.DiaObjectChunks2
1314 # Use a random number for a second part of partitioning key so
1315 # that different clients could wrtite to different partitions.
1316 # This makes it not exactly reproducible.
1317 extra_columns["apdb_replica_subchunk"] = random.randrange(config.replica_sub_chunk_count)
1318 self._storeObjectsPandas(objs, table, extra_columns=extra_columns)
1320 # Store copy of the records in dedup table.
1321 if context.has_dedup_table:
1322 table = ExtraTables.DiaObjectDedup
1323 extra_columns = {
1324 "dedup_part": random.randrange(config.partitioning.num_part_dedup),
1325 validity_start_column: timestamp,
1326 }
1327 self._storeObjectsPandas(objs, table, extra_columns=extra_columns)
1329 def _storeDiaSources(
1330 self,
1331 table_name: ApdbTables,
1332 sources: pandas.DataFrame,
1333 replica_chunk: ReplicaChunk | None,
1334 ) -> int | None:
1335 """Store catalog of DIASources or DIAForcedSources from current visit.
1337 Parameters
1338 ----------
1339 table_name : `ApdbTables`
1340 Table where to store the data.
1341 sources : `pandas.DataFrame`
1342 Catalog containing DiaSource records
1343 visit_time : `astropy.time.Time`
1344 Time of the current visit.
1345 replica_chunk : `ReplicaChunk` or `None`
1346 Replica chunk identifier if replication is configured.
1348 Returns
1349 -------
1350 subchunk : `int` or `None`
1351 Subchunk number for resulting replica data, `None` if relication is
1352 not enabled ot subchunking is not enabled.
1353 """
1354 context = self._context
1355 config = context.config
1357 # Time partitioning has to be based on midpointMjdTai, not visit_time
1358 # as visit_time is not really a visit time.
1359 tp_sources = sources.copy(deep=False)
1360 tp_sources["apdb_time_part"] = tp_sources["midpointMjdTai"].apply(context.partitioner.time_partition)
1361 if (time_partitions_range := context.time_partitions_range) is not None:
1362 self._check_time_partitions(tp_sources["apdb_time_part"], time_partitions_range)
1363 extra_columns: dict[str, Any] = {}
1364 if not config.partitioning.time_partition_tables:
1365 self._storeObjectsPandas(tp_sources, table_name)
1366 else:
1367 # Group by time partition
1368 partitions = set(tp_sources["apdb_time_part"])
1369 if len(partitions) == 1:
1370 # Single partition - just save the whole thing.
1371 time_part = partitions.pop()
1372 self._storeObjectsPandas(sources, table_name, time_part=time_part)
1373 else:
1374 # group by time partition.
1375 for time_part, sub_frame in tp_sources.groupby(by="apdb_time_part"):
1376 sub_frame.drop(columns="apdb_time_part", inplace=True)
1377 self._storeObjectsPandas(sub_frame, table_name, time_part=time_part)
1379 subchunk: int | None = None
1380 if replica_chunk is not None:
1381 extra_columns = {"apdb_replica_chunk": replica_chunk.id}
1382 if context.has_chunk_sub_partitions:
1383 subchunk = random.randrange(config.replica_sub_chunk_count)
1384 extra_columns["apdb_replica_subchunk"] = subchunk
1385 if table_name is ApdbTables.DiaSource:
1386 extra_table = ExtraTables.DiaSourceChunks2
1387 else:
1388 extra_table = ExtraTables.DiaForcedSourceChunks2
1389 else:
1390 if table_name is ApdbTables.DiaSource:
1391 extra_table = ExtraTables.DiaSourceChunks
1392 else:
1393 extra_table = ExtraTables.DiaForcedSourceChunks
1394 self._storeObjectsPandas(sources, extra_table, extra_columns=extra_columns)
1396 return subchunk
1398 def _check_time_partitions(
1399 self, partitions: Iterable[int], time_partitions_range: ApdbCassandraTimePartitionRange
1400 ) -> None:
1401 """Check that time partitons for new data actually exist.
1403 Parameters
1404 ----------
1405 partitions : `~collections.abc.Iterable` [`int`]
1406 Time partitions for new data.
1407 time_partitions_range : `ApdbCassandraTimePartitionRange`
1408 Currrent time partition range.
1409 """
1410 partitions = set(partitions)
1411 min_part = min(partitions)
1412 max_part = max(partitions)
1413 if min_part < time_partitions_range.start or max_part > time_partitions_range.end:
1414 raise ValueError(
1415 "Attempt to store data for time partitions that do not yet exist. "
1416 f"Partitons for new records: {min_part}-{max_part}. "
1417 f"Database partitons: {time_partitions_range.start}-{time_partitions_range.end}."
1418 )
1419 # Make a noise when writing to the last partition.
1420 if max_part == time_partitions_range.end:
1421 warnings.warn(
1422 "Writing into the last temporal partition. Partition range needs to be extended soon.",
1423 stacklevel=3,
1424 )
1426 def _storeDiaSourcesPartitions(
1427 self,
1428 sources: pandas.DataFrame,
1429 visit_time: astropy.time.Time,
1430 replica_chunk: ReplicaChunk | None,
1431 subchunk: int | None,
1432 ) -> None:
1433 """Store mapping of diaSourceId to its partitioning values.
1435 Parameters
1436 ----------
1437 sources : `pandas.DataFrame`
1438 Catalog containing DiaSource records
1439 visit_time : `astropy.time.Time`
1440 Time of the current visit.
1441 replica_chunk : `ReplicaChunk` or `None`
1442 Replication chunk, or `None` when replication is disabled.
1443 subchunk : `int` or `None`
1444 Replication sub-chunk, or `None` when replication is disabled or
1445 sub-chunking is not used.
1446 """
1447 context = self._context
1449 id_map = cast(pandas.DataFrame, sources[["diaSourceId", "apdb_part"]])
1450 extra_columns = {
1451 "apdb_time_part": context.partitioner.time_partition(visit_time),
1452 "apdb_replica_chunk": replica_chunk.id if replica_chunk is not None else None,
1453 }
1454 if context.has_chunk_sub_partitions:
1455 extra_columns["apdb_replica_subchunk"] = subchunk
1457 self._storeObjectsPandas(
1458 id_map, ExtraTables.DiaSourceToPartition, extra_columns=extra_columns, time_part=None
1459 )
1461 def _storeObjectsPandas(
1462 self,
1463 records: pandas.DataFrame,
1464 table_name: ApdbTables | ExtraTables,
1465 extra_columns: Mapping | None = None,
1466 time_part: int | None = None,
1467 ) -> None:
1468 """Store generic objects.
1470 Takes Pandas catalog and stores a bunch of records in a table.
1472 Parameters
1473 ----------
1474 records : `pandas.DataFrame`
1475 Catalog containing object records
1476 table_name : `ApdbTables`
1477 Name of the table as defined in APDB schema.
1478 extra_columns : `dict`, optional
1479 Mapping (column_name, column_value) which gives fixed values for
1480 columns in each row, overrides values in ``records`` if matching
1481 columns exist there.
1482 time_part : `int`, optional
1483 If not `None` then insert into a per-partition table.
1485 Notes
1486 -----
1487 If Pandas catalog contains additional columns not defined in table
1488 schema they are ignored. Catalog does not have to contain all columns
1489 defined in a table, but partition and clustering keys must be present
1490 in a catalog or ``extra_columns``.
1491 """
1492 context = self._context
1494 # use extra columns if specified
1495 if extra_columns is None:
1496 extra_columns = {}
1497 extra_fields = list(extra_columns.keys())
1499 # Fields that will come from dataframe.
1500 df_fields = [column for column in records.columns if column not in extra_fields]
1502 column_map = context.schema.getColumnMap(table_name)
1503 # list of columns (as in felis schema)
1504 fields = [column_map[field].name for field in df_fields if field in column_map]
1505 fields += extra_fields
1507 # check that all partitioning and clustering columns are defined
1508 partition_columns = context.schema.partitionColumns(table_name)
1509 required_columns = partition_columns + context.schema.clusteringColumns(table_name)
1510 missing_columns = [column for column in required_columns if column not in fields]
1511 if missing_columns:
1512 raise ValueError(f"Primary key columns are missing from catalog: {missing_columns}")
1514 qfields = [quote_id(field) for field in fields]
1515 qfields_str = ",".join(qfields)
1517 batch_size = self._batch_size(table_name)
1519 with self._timer("insert_build_time", tags={"table": table_name.name}):
1520 # Multi-partition batches are problematic in general, so we want to
1521 # group records in a batch by their partition key.
1522 values_by_key: dict[tuple, list[list]] = defaultdict(list)
1523 for rec in records.itertuples(index=False):
1524 values = []
1525 partitioning_values: dict[str, Any] = {}
1526 for field in df_fields:
1527 if field not in column_map:
1528 continue
1529 value = getattr(rec, field)
1530 if column_map[field].datatype is felis.datamodel.DataType.timestamp:
1531 if isinstance(value, pandas.Timestamp):
1532 value = value.to_pydatetime()
1533 elif value is pandas.NaT:
1534 value = None
1535 else:
1536 # Assume it's seconds since epoch, Cassandra
1537 # datetime is in milliseconds
1538 value = int(value * 1000)
1539 value = literal(value)
1540 values.append(UNSET_VALUE if value is None else value)
1541 if field in partition_columns:
1542 partitioning_values[field] = value
1543 for field in extra_fields:
1544 value = literal(extra_columns[field])
1545 values.append(UNSET_VALUE if value is None else value)
1546 if field in partition_columns:
1547 partitioning_values[field] = value
1549 key = tuple(partitioning_values[field] for field in partition_columns)
1550 values_by_key[key].append(values)
1552 table = context.schema.tableName(table_name, time_part)
1554 holders = ",".join(["?"] * len(qfields))
1555 query = f'INSERT INTO "{self._keyspace}"."{table}" ({qfields_str}) VALUES ({holders})'
1556 statement = context.preparer.prepare(query)
1557 # Cassandra has 64k limit on batch size, normally that should be
1558 # enough but some tests generate too many forced sources.
1559 queries = []
1560 for key_values in values_by_key.values():
1561 for values_chunk in chunk_iterable(key_values, batch_size):
1562 batch = cassandra.query.BatchStatement()
1563 for row_values in values_chunk:
1564 batch.add(statement, row_values)
1565 queries.append((batch, None))
1566 assert batch.routing_key is not None and batch.keyspace is not None
1568 _LOG.debug("%s: will store %d records", context.schema.tableName(table_name), records.shape[0])
1569 with self._timer(
1570 "insert_time", tags={"table": table_name.name, "method": "_storeObjectsPandas"}
1571 ) as timer:
1572 execute_concurrent(context.session, queries, execution_profile="write")
1573 timer.add_values(row_count=len(records), num_batches=len(queries))
1575 def _storeUpdateRecords(
1576 self, records: Iterable[ApdbUpdateRecord], chunk: ReplicaChunk, *, store_chunk: bool = False
1577 ) -> None:
1578 """Store ApdbUpdateRecords in the replica table for those records.
1580 Parameters
1581 ----------
1582 records : `list` [`ApdbUpdateRecord`]
1583 Records to store.
1584 chunk : `ReplicaChunk`
1585 Replica chunk for these records.
1586 store_chunk : `bool`
1587 If True then also store replica chunk.
1589 Raises
1590 ------
1591 TypeError
1592 Raised if replication is not enabled for this instance.
1593 """
1594 context = self._context
1595 config = context.config
1597 if not context.schema.replication_enabled:
1598 raise TypeError("Replication is not enabled for this APDB instance.")
1600 if store_chunk:
1601 self._storeReplicaChunk(chunk)
1603 apdb_replica_chunk = chunk.id
1604 # Do not use unique_if from ReplicaChunk as it could be reused in
1605 # multiple calls to this method.
1606 update_unique_id = uuid.uuid4()
1608 rows = []
1609 for record in records:
1610 rows.append(
1611 [
1612 apdb_replica_chunk,
1613 record.update_time_ns,
1614 record.update_order,
1615 update_unique_id,
1616 record.to_json(),
1617 ]
1618 )
1619 columns = [
1620 "apdb_replica_chunk",
1621 "update_time_ns",
1622 "update_order",
1623 "update_unique_id",
1624 "update_payload",
1625 ]
1626 if context.has_chunk_sub_partitions:
1627 subchunk = random.randrange(config.replica_sub_chunk_count)
1628 for row in rows:
1629 row.append(subchunk)
1630 columns.append("apdb_replica_subchunk")
1632 table_name = context.schema.tableName(ExtraTables.ApdbUpdateRecordChunks)
1633 placeholders = ", ".join(["%s"] * len(columns))
1634 columns_str = ", ".join(columns)
1635 query = f'INSERT INTO "{self._keyspace}"."{table_name}" ({columns_str}) VALUES ({placeholders})'
1636 queries = [(query, row) for row in rows]
1638 with self._timer("store_update_record", tags={"table": table_name}) as timer:
1639 execute_concurrent(context.session, queries, execution_profile="write")
1640 timer.add_values(row_count=len(queries))
1642 def _add_apdb_part(self, df: pandas.DataFrame) -> pandas.DataFrame:
1643 """Calculate spatial partition for each record and add it to a
1644 DataFrame.
1646 Parameters
1647 ----------
1648 df : `pandas.DataFrame`
1649 DataFrame which has to contain ra/dec columns, names of these
1650 columns are defined by configuration ``ra_dec_columns`` field.
1652 Returns
1653 -------
1654 df : `pandas.DataFrame`
1655 DataFrame with ``apdb_part`` column which contains pixel index
1656 for ra/dec coordinates.
1658 Notes
1659 -----
1660 This overrides any existing column in a DataFrame with the same name
1661 (``apdb_part``). Original DataFrame is not changed, copy of a DataFrame
1662 is returned.
1663 """
1664 context = self._context
1665 config = context.config
1667 # Calculate pixelization index for every record.
1668 apdb_part = np.zeros(df.shape[0], dtype=np.int64)
1669 ra_col, dec_col = config.ra_dec_columns
1670 for i, (ra, dec) in enumerate(zip(df[ra_col], df[dec_col])):
1671 uv3d = sphgeom.UnitVector3d(sphgeom.LonLat.fromDegrees(ra, dec))
1672 idx = context.partitioner.pixel(uv3d)
1673 apdb_part[i] = idx
1674 df = df.copy()
1675 df["apdb_part"] = apdb_part
1676 return df
1678 def _make_empty_catalog(self, table_name: ApdbTables) -> pandas.DataFrame:
1679 """Make an empty catalog for a table with a given name.
1681 Parameters
1682 ----------
1683 table_name : `ApdbTables`
1684 Name of the table.
1686 Returns
1687 -------
1688 catalog : `pandas.DataFrame`
1689 An empty catalog.
1690 """
1691 table = self.schema.tableSchemas[table_name]
1693 data = {columnDef.name: pandas.Series(dtype=columnDef.pandas_type) for columnDef in table.columns}
1694 return pandas.DataFrame(data)
1696 def _combine_where(
1697 self,
1698 prefix: str,
1699 where1: list[tuple[str, tuple]],
1700 where2: list[tuple[str, tuple]],
1701 where3: tuple[str, tuple] | None = None,
1702 suffix: str | None = None,
1703 ) -> Iterator[tuple[cassandra.query.Statement, tuple]]:
1704 """Make cartesian product of two parts of WHERE clause into a series
1705 of statements to execute.
1707 Parameters
1708 ----------
1709 prefix : `str`
1710 Initial statement prefix that comes before WHERE clause, e.g.
1711 "SELECT * from Table"
1712 """
1713 context = self._context
1715 # If lists are empty use special sentinels.
1716 if not where1:
1717 where1 = [("", ())]
1718 if not where2:
1719 where2 = [("", ())]
1721 for expr1, params1 in where1:
1722 for expr2, params2 in where2:
1723 full_query = prefix
1724 wheres = []
1725 params = params1 + params2
1726 if expr1:
1727 wheres.append(expr1)
1728 if expr2:
1729 wheres.append(expr2)
1730 if where3:
1731 wheres.append(where3[0])
1732 params += where3[1]
1733 if wheres:
1734 full_query += " WHERE " + " AND ".join(wheres)
1735 if suffix:
1736 full_query += " " + suffix
1737 if params:
1738 statement = context.preparer.prepare(full_query)
1739 else:
1740 # If there are no params then it is likely that query
1741 # has a bunch of literals rendered already, no point
1742 # trying to prepare it.
1743 statement = cassandra.query.SimpleStatement(full_query)
1744 yield (statement, params)
1746 def _fix_input_timestamps(self, df: pandas.DataFrame) -> pandas.DataFrame:
1747 """Update timestamp columns in input DataFrame to be naive datetime
1748 type.
1750 Clients may or may not generate aware timestamps, code in this class
1751 assumes that timestamps are naive, so we convert them to UTC and
1752 drop timezone.
1753 """
1754 # Find all columns with aware timestamps.
1755 columns = [column for column, dtype in df.dtypes.items() if isinstance(dtype, pandas.DatetimeTZDtype)]
1756 for column in columns:
1757 # tz_convert(None) will convert to UTC and drop timezone.
1758 df[column] = df[column].dt.tz_convert(None)
1759 return df
1761 def _batch_size(self, table: ApdbTables | ExtraTables) -> int:
1762 """Calculate batch size based on config parameters."""
1763 context = self._context
1764 config = context.config
1766 # Cassandra limit on number of statements in a batch is 64k.
1767 batch_size = 65_535
1768 if 0 < config.batch_statement_limit < batch_size:
1769 batch_size = config.batch_statement_limit
1770 if config.batch_size_limit > 0:
1771 # The purpose of this limit is to try not to exceed batch size
1772 # threshold which is set on server side. Cassandra wire protocol
1773 # for prepared queries (and batches) only sends column values with
1774 # with an additional 4 bytes per value specifying size. Value is
1775 # not included for NULL or NOT_SET values, but the size is always
1776 # there. There is additional small per-query overhead, which we
1777 # ignore.
1778 row_size = context.schema.table_row_size(table)
1779 row_size += 4 * len(context.schema.getColumnMap(table))
1780 batch_size = min(batch_size, (config.batch_size_limit // row_size) + 1)
1781 return batch_size
1783 def _group_dia_objects_by_partition(
1784 self, partitioner: Partitioner, objects: list[DiaObjectId], pad_arcsec: float
1785 ) -> Mapping[int, list[int]]:
1786 """Group DiaObjects by partition.
1788 Parameters
1789 ----------
1790 partitioner : `Partitioner`
1791 Objects which knows how to partition things.
1792 objects : `list` [`DiaObjectId`]
1793 Collection of objects to partition.
1794 pad_arcsec : `float`
1795 Additional padding around object position.
1797 Returns
1798 -------
1799 grouped_objects
1800 Mapping of spatial patition ID to list ob object IDs that it
1801 contains. Some objects may belong to more than one partition.
1802 """
1803 partitioned_object_ids: dict[int, list[int]] = defaultdict(list)
1804 for obj_id in objects:
1805 partitions = partitioner.pixelization.circle_pixels(obj_id.ra, obj_id.dec, pad_arcsec)
1806 for pixel in partitions:
1807 partitioned_object_ids[pixel].append(obj_id.diaObjectId)
1808 return partitioned_object_ids
1810 def _timestamp_column_name(self, column: str) -> str:
1811 """Return column name before/after schema migration to MJD TAI."""
1812 return self._schema.timestamp_column_name(column)
1814 def _timestamp_column_value(self, time: astropy.time.Time) -> float | int:
1815 """Return column value before/after schema migration to MJD TAI."""
1816 if self._schema.has_mjd_timestamps:
1817 return float(time.tai.mjd)
1818 else:
1819 return int(time.datetime.astimezone(tz=datetime.UTC).timestamp() * 1000)
1821 def _get_diasource_data(self, source_ids: Iterable[DiaSourceId], *columns: str) -> list:
1822 """Select records from DiaSource table by diaSourceId and return all
1823 records as a list of named tuples.
1824 """
1825 context = self._context
1826 config = context.config
1827 partitioner = context.partitioner
1829 columns = ("diaSourceId",) + columns
1830 what = ",".join(quote_id(column) for column in columns)
1832 # Allow some uncertainty for coordinates and time when calculating
1833 # partitions.
1834 statements: list[tuple] = []
1835 pad_arcsec = 1.0
1836 pad_time_day = 10 / (24 * 3600)
1837 for source_id in source_ids:
1838 center = sphgeom.UnitVector3d(sphgeom.LonLat.fromDegrees(source_id.ra, source_id.dec))
1839 region = sphgeom.Circle(center, sphgeom.Angle.fromDegrees(pad_arcsec / 3600.0))
1840 spatial_where, _ = partitioner.spatial_where(region, for_prepare=True)
1842 tables, temporal_where = partitioner.temporal_where(
1843 ApdbTables.DiaSource,
1844 source_id.midpointMjdTai - pad_time_day,
1845 source_id.midpointMjdTai + pad_time_day,
1846 for_prepare=True,
1847 partitons_range=context.time_partitions_range,
1848 query_per_time_part=True,
1849 )
1851 id_where = ('"diaSourceId" = ?', (source_id.diaSourceId,))
1853 for table in tables:
1854 prefix = f'SELECT {what} FROM "{self._keyspace}"."{table}"'
1855 statements += list(self._combine_where(prefix, spatial_where, temporal_where, id_where))
1857 with self._timer(
1858 "select_time", tags={"table": "DiaSource", "method": "_get_diasource_data"}
1859 ) as timer:
1860 result = cast(
1861 list[tuple],
1862 select_concurrent(
1863 context.session,
1864 statements,
1865 "read_named_tuples",
1866 config.connection_config.read_concurrency,
1867 ),
1868 )
1869 timer.add_values(row_count=len(result), num_queries=len(statements))
1871 return result
1873 def _get_diaobject_data(self, object_ids: Iterable[DiaObjectId], *columns: str) -> list:
1874 """Select records from DiaObjectLast table by diaObjectId and return
1875 all records as a list of named tuples.
1876 """
1877 context = self._context
1878 config = context.config
1879 partitioner = context.partitioner
1881 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
1882 columns = ("diaObjectId",) + columns
1883 what = ",".join(quote_id(column) for column in columns)
1885 # Allow some uncertainty for coordinates when calculating partitions.
1886 pad_arcsec = 1.0
1887 ids_by_partition = defaultdict(list)
1888 for object_id in object_ids:
1889 pixels = partitioner.pixelization.circle_pixels(object_id.ra, object_id.dec, pad_arcsec)
1890 for pixel in pixels:
1891 ids_by_partition[pixel].append(object_id.diaObjectId)
1893 statements: list[tuple] = []
1894 for apdb_part, diaObjectIds in ids_by_partition.items():
1895 ids_str = ",".join(str(diaObjectId) for diaObjectId in diaObjectIds)
1896 statement = cassandra.query.SimpleStatement(
1897 f'SELECT {what} FROM "{self._keyspace}"."{table_name}" '
1898 f'WHERE apdb_part = %s AND "diaObjectId" IN ({ids_str})'
1899 )
1900 statements.append((statement, (apdb_part,)))
1902 with self._timer(
1903 "select_time", tags={"table": "DiaObjectLast", "method": "_get_diaobject_data"}
1904 ) as timer:
1905 result = cast(
1906 list[tuple],
1907 select_concurrent(
1908 context.session,
1909 statements,
1910 "read_named_tuples",
1911 config.connection_config.read_concurrency,
1912 ),
1913 )
1914 timer.add_values(row_count=len(result), num_queries=len(statements))
1916 return result