Coverage for python / lsst / dax / apdb / cassandra / apdbCassandra.py: 8%
855 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-22 02:00 -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 execute_concurrent,
76 literal,
77 quote_id,
78 select_concurrent,
79)
80from .config import ApdbCassandraConfig, ApdbCassandraConnectionConfig, ApdbCassandraTimePartitionRange
81from .connectionContext import ConnectionContext, DbVersions
82from .exceptions import CassandraMissingError
83from .partitioner import Partitioner
84from .sessionFactory import SessionContext, SessionFactory
86if TYPE_CHECKING:
87 from ..apdbMetadata import ApdbMetadata
88 from ..apdbUpdateRecord import ApdbUpdateRecord
90_LOG = logging.getLogger(__name__)
92_MON = MonAgent(__name__)
94VERSION = VersionTuple(1, 3, 0)
95"""Version for the code controlling non-replication tables. This needs to be
96updated following compatibility rules when schema produced by this code
97changes.
98"""
101class ApdbCassandra(Apdb):
102 """Implementation of APDB database with Apache Cassandra backend.
104 Parameters
105 ----------
106 config : `ApdbCassandraConfig`
107 Configuration object.
108 """
110 def __init__(self, config: ApdbCassandraConfig):
111 if not CASSANDRA_IMPORTED:
112 raise CassandraMissingError()
114 self._config = config
115 self._keyspace = config.keyspace
116 self._schema = ApdbSchema(config.schema_file, config.ss_schema_file)
118 self._session_factory = SessionFactory(config)
119 self._connection_context: ConnectionContext | None = None
121 @property
122 def _context(self) -> ConnectionContext:
123 """Establish connection if not established and return context."""
124 if self._connection_context is None:
125 current_versions = DbVersions(
126 schema_version=self.schema.schemaVersion(),
127 code_version=self.apdbImplementationVersion(),
128 replica_version=ApdbCassandraReplica.apdbReplicaImplementationVersion(),
129 )
130 _LOG.debug("Current versions: %s", current_versions)
132 session = self._session_factory.session()
133 self._connection_context = ConnectionContext(
134 session, self._config, self.schema.tableSchemas, current_versions
135 )
137 if _LOG.isEnabledFor(logging.DEBUG):
138 _LOG.debug("ApdbCassandra Configuration: %s", self._connection_context.config.model_dump())
140 return self._connection_context
142 def _timer(self, name: str, *, tags: Mapping[str, str | int] | None = None) -> Timer:
143 """Create `Timer` instance given its name."""
144 return Timer(name, _MON, tags=tags)
146 @classmethod
147 def apdbImplementationVersion(cls) -> VersionTuple:
148 """Return version number for current APDB implementation.
150 Returns
151 -------
152 version : `VersionTuple`
153 Version of the code defined in implementation class.
154 """
155 return VERSION
157 def getConfig(self) -> ApdbCassandraConfig:
158 # docstring is inherited from a base class
159 return self._context.config
161 def tableDef(self, table: ApdbTables) -> Table | None:
162 # docstring is inherited from a base class
163 return self.schema.tableSchemas.get(table)
165 @classmethod
166 def init_database(
167 cls,
168 hosts: tuple[str, ...],
169 keyspace: str,
170 *,
171 schema_file: str | None = None,
172 ss_schema_file: str | None = None,
173 read_sources_months: int | None = None,
174 read_forced_sources_months: int | None = None,
175 enable_replica: bool = False,
176 replica_skips_diaobjects: bool = False,
177 port: int | None = None,
178 username: str | None = None,
179 prefix: str | None = None,
180 part_pixelization: str | None = None,
181 part_pix_level: int | None = None,
182 time_partition_tables: bool = True,
183 time_partition_start: str | None = None,
184 time_partition_end: str | None = None,
185 read_consistency: str | None = None,
186 write_consistency: str | None = None,
187 read_timeout: int | None = None,
188 write_timeout: int | None = None,
189 ra_dec_columns: tuple[str, str] | None = None,
190 replication_factor: int | None = None,
191 drop: bool = False,
192 table_options: CreateTableOptions | None = None,
193 ) -> ApdbCassandraConfig:
194 """Initialize new APDB instance and make configuration object for it.
196 Parameters
197 ----------
198 hosts : `tuple` [`str`, ...]
199 List of host names or IP addresses for Cassandra cluster.
200 keyspace : `str`
201 Name of the keyspace for APDB tables.
202 schema_file : `str`, optional
203 Location of (YAML) configuration file with APDB schema. If not
204 specified then default location will be used.
205 ss_schema_file : `str`, optional
206 Location of (YAML) configuration file with SSO schema. If not
207 specified then default location will be used.
208 read_sources_months : `int`, optional
209 Number of months of history to read from DiaSource.
210 read_forced_sources_months : `int`, optional
211 Number of months of history to read from DiaForcedSource.
212 enable_replica : `bool`, optional
213 If True, make additional tables used for replication to PPDB.
214 replica_skips_diaobjects : `bool`, optional
215 If `True` then do not fill regular ``DiaObject`` table when
216 ``enable_replica`` is `True`.
217 port : `int`, optional
218 Port number to use for Cassandra connections.
219 username : `str`, optional
220 User name for Cassandra connections.
221 prefix : `str`, optional
222 Optional prefix for all table names.
223 part_pixelization : `str`, optional
224 Name of the MOC pixelization used for partitioning.
225 part_pix_level : `int`, optional
226 Pixelization level.
227 time_partition_tables : `bool`, optional
228 Create per-partition tables.
229 time_partition_start : `str`, optional
230 Starting time for per-partition tables, in yyyy-mm-ddThh:mm:ss
231 format, in TAI.
232 time_partition_end : `str`, optional
233 Ending time for per-partition tables, in yyyy-mm-ddThh:mm:ss
234 format, in TAI.
235 read_consistency : `str`, optional
236 Name of the consistency level for read operations.
237 write_consistency : `str`, optional
238 Name of the consistency level for write operations.
239 read_timeout : `int`, optional
240 Read timeout in seconds.
241 write_timeout : `int`, optional
242 Write timeout in seconds.
243 ra_dec_columns : `tuple` [`str`, `str`], optional
244 Names of ra/dec columns in DiaObject table.
245 replication_factor : `int`, optional
246 Replication factor used when creating new keyspace, if keyspace
247 already exists its replication factor is not changed.
248 drop : `bool`, optional
249 If `True` then drop existing tables before re-creating the schema.
250 table_options : `CreateTableOptions`, optional
251 Options used when creating Cassandra tables.
253 Returns
254 -------
255 config : `ApdbCassandraConfig`
256 Resulting configuration object for a created APDB instance.
257 """
258 # Some non-standard defaults for connection parameters, these can be
259 # changed later in generated config. Check Cassandra driver
260 # documentation for what these parameters do. These parameters are not
261 # used during database initialization, but they will be saved with
262 # generated config.
263 connection_config = ApdbCassandraConnectionConfig(
264 extra_parameters={
265 "idle_heartbeat_interval": 0,
266 "idle_heartbeat_timeout": 30,
267 "control_connection_timeout": 100,
268 },
269 )
270 config = ApdbCassandraConfig(
271 contact_points=hosts,
272 keyspace=keyspace,
273 enable_replica=enable_replica,
274 replica_skips_diaobjects=replica_skips_diaobjects,
275 connection_config=connection_config,
276 )
277 config.partitioning.time_partition_tables = time_partition_tables
278 if schema_file is not None:
279 config.schema_file = schema_file
280 if ss_schema_file is not None:
281 config.ss_schema_file = ss_schema_file
282 if read_sources_months is not None:
283 config.read_sources_months = read_sources_months
284 if read_forced_sources_months is not None:
285 config.read_forced_sources_months = read_forced_sources_months
286 if port is not None:
287 config.connection_config.port = port
288 if username is not None:
289 config.connection_config.username = username
290 if prefix is not None:
291 config.prefix = prefix
292 if part_pixelization is not None:
293 config.partitioning.part_pixelization = part_pixelization
294 if part_pix_level is not None:
295 config.partitioning.part_pix_level = part_pix_level
296 if time_partition_start is not None:
297 config.partitioning.time_partition_start = time_partition_start
298 if time_partition_end is not None:
299 config.partitioning.time_partition_end = time_partition_end
300 if read_consistency is not None:
301 config.connection_config.read_consistency = read_consistency
302 if write_consistency is not None:
303 config.connection_config.write_consistency = write_consistency
304 if read_timeout is not None:
305 config.connection_config.read_timeout = read_timeout
306 if write_timeout is not None:
307 config.connection_config.write_timeout = write_timeout
308 if ra_dec_columns is not None:
309 config.ra_dec_columns = ra_dec_columns
311 cls._makeSchema(config, drop=drop, replication_factor=replication_factor, table_options=table_options)
313 return config
315 def get_replica(self) -> ApdbCassandraReplica:
316 """Return `ApdbReplica` instance for this database."""
317 # Note that this instance has to stay alive while replica exists, so
318 # we pass reference to self.
319 return ApdbCassandraReplica(self)
321 @classmethod
322 def _makeSchema(
323 cls,
324 config: ApdbConfig,
325 *,
326 drop: bool = False,
327 replication_factor: int | None = None,
328 table_options: CreateTableOptions | None = None,
329 ) -> None:
330 # docstring is inherited from a base class
332 if not isinstance(config, ApdbCassandraConfig):
333 raise TypeError(f"Unexpected type of configuration object: {type(config)}")
335 simple_schema = ApdbSchema(config.schema_file, config.ss_schema_file)
337 with SessionContext(config) as session:
338 schema = ApdbCassandraSchema(
339 session=session,
340 keyspace=config.keyspace,
341 table_schemas=simple_schema.tableSchemas,
342 prefix=config.prefix,
343 time_partition_tables=config.partitioning.time_partition_tables,
344 enable_replica=config.enable_replica,
345 replica_skips_diaobjects=config.replica_skips_diaobjects,
346 )
348 # Ask schema to create all tables.
349 part_range_config: ApdbCassandraTimePartitionRange | None = None
350 if config.partitioning.time_partition_tables:
351 partitioner = Partitioner(config)
352 time_partition_start = astropy.time.Time(
353 config.partitioning.time_partition_start, format="isot", scale="tai"
354 )
355 time_partition_end = astropy.time.Time(
356 config.partitioning.time_partition_end, format="isot", scale="tai"
357 )
358 part_range_config = ApdbCassandraTimePartitionRange(
359 start=partitioner.time_partition(time_partition_start),
360 end=partitioner.time_partition(time_partition_end),
361 )
362 schema.makeSchema(
363 drop=drop,
364 part_range=part_range_config,
365 replication_factor=replication_factor,
366 table_options=table_options,
367 )
368 else:
369 schema.makeSchema(
370 drop=drop, replication_factor=replication_factor, table_options=table_options
371 )
373 meta_table_name = ApdbTables.metadata.table_name(config.prefix)
374 metadata = ApdbMetadataCassandra(
375 session, meta_table_name, config.keyspace, "read_tuples", "write"
376 )
378 # Fill version numbers, overrides if they existed before.
379 metadata.set(
380 ConnectionContext.metadataSchemaVersionKey, str(simple_schema.schemaVersion()), force=True
381 )
382 metadata.set(
383 ConnectionContext.metadataCodeVersionKey, str(cls.apdbImplementationVersion()), force=True
384 )
386 if config.enable_replica:
387 # Only store replica code version if replica is enabled.
388 metadata.set(
389 ConnectionContext.metadataReplicaVersionKey,
390 str(ApdbCassandraReplica.apdbReplicaImplementationVersion()),
391 force=True,
392 )
394 # Store frozen part of a configuration in metadata.
395 freezer = ApdbConfigFreezer[ApdbCassandraConfig](ConnectionContext.frozen_parameters)
396 metadata.set(ConnectionContext.metadataConfigKey, freezer.to_json(config), force=True)
398 # Store time partition range.
399 if part_range_config:
400 part_range_config.save_to_meta(metadata)
402 def getDiaObjects(self, region: sphgeom.Region) -> pandas.DataFrame:
403 # docstring is inherited from a base class
404 context = self._context
405 config = context.config
407 sp_where, num_sp_part = context.partitioner.spatial_where(region, for_prepare=True)
408 _LOG.debug("getDiaObjects: #partitions: %s", len(sp_where))
410 # We need to exclude extra partitioning columns from result.
411 column_names = context.schema.apdbColumnNames(ApdbTables.DiaObjectLast)
412 what = ",".join(quote_id(column) for column in column_names)
414 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
415 query = f'SELECT {what} from "{self._keyspace}"."{table_name}"'
416 statements: list[tuple] = []
417 for where, params in sp_where:
418 full_query = f"{query} WHERE {where}"
419 if params:
420 statement = context.preparer.prepare(full_query)
421 else:
422 # If there are no params then it is likely that query has a
423 # bunch of literals rendered already, no point trying to
424 # prepare it because it's not reusable.
425 statement = cassandra.query.SimpleStatement(full_query)
426 statements.append((statement, params))
427 _LOG.debug("getDiaObjects: #queries: %s", len(statements))
429 with self._timer("select_time", tags={"table": "DiaObject", "method": "getDiaObjects"}) as timer:
430 objects = cast(
431 pandas.DataFrame,
432 select_concurrent(
433 context.session,
434 statements,
435 "read_pandas_multi",
436 config.connection_config.read_concurrency,
437 ),
438 )
439 timer.add_values(row_count=len(objects), num_sp_part=num_sp_part, num_queries=len(statements))
441 _LOG.debug("found %s DiaObjects", objects.shape[0])
442 return objects
444 def getDiaSources(
445 self,
446 region: sphgeom.Region,
447 object_ids: Iterable[int] | None,
448 visit_time: astropy.time.Time,
449 start_time: astropy.time.Time | None = None,
450 ) -> pandas.DataFrame | None:
451 # docstring is inherited from a base class
452 context = self._context
453 config = context.config
455 months = config.read_sources_months
456 if start_time is None and months == 0:
457 return None
459 mjd_end = float(visit_time.tai.mjd)
460 if start_time is None:
461 mjd_start = mjd_end - months * 30
462 else:
463 mjd_start = float(start_time.tai.mjd)
465 return self._getSources(region, object_ids, mjd_start, mjd_end, ApdbTables.DiaSource)
467 def getDiaForcedSources(
468 self,
469 region: sphgeom.Region,
470 object_ids: Iterable[int] | None,
471 visit_time: astropy.time.Time,
472 start_time: astropy.time.Time | None = None,
473 ) -> pandas.DataFrame | None:
474 # docstring is inherited from a base class
475 context = self._context
476 config = context.config
478 months = config.read_forced_sources_months
479 if start_time is None and months == 0:
480 return None
482 mjd_end = float(visit_time.tai.mjd)
483 if start_time is None:
484 mjd_start = mjd_end - months * 30
485 else:
486 mjd_start = float(start_time.tai.mjd)
488 return self._getSources(region, object_ids, mjd_start, mjd_end, ApdbTables.DiaForcedSource)
490 def getDiaObjectsForDedup(self, since: astropy.time.Time | None = None) -> pandas.DataFrame:
491 # docstring is inherited from a base class
492 context = self._context
493 config = context.config
495 if not context.has_dedup_table:
496 raise TypeError("DiaObjectDedup table does not exist in this APDB instance.")
498 if since is None:
499 # Read last deduplication time from metadata.
500 dedup_str = context.metadata.get(context.metadataDedupKey)
501 if dedup_str is not None:
502 dedup_state = json.loads(dedup_str)
503 dedup_time_str = dedup_state["dedup_time_iso_tai"]
504 since = astropy.time.Time(dedup_time_str, format="iso", scale="tai")
506 column_names = context.schema.apdbColumnNames(ExtraTables.DiaObjectDedup)
507 what = ",".join(quote_id(column) for column in column_names)
509 validity_start_column = self._timestamp_column_name("validityStart")
510 timestamp = None if since is None else self._timestamp_column_value(since)
512 table_name = context.schema.tableName(ExtraTables.DiaObjectDedup)
513 query = f'SELECT {what} FROM "{self._keyspace}"."{table_name}" WHERE dedup_part = %s'
514 if since is not None:
515 query += f' AND "{validity_start_column}" >= %s'
516 query += " ALLOW FILTERING"
518 statement = cassandra.query.SimpleStatement(query)
520 num_part = config.partitioning.num_part_dedup
521 statements = []
522 for dedup_part in range(num_part):
523 params = (dedup_part,) if timestamp is None else (dedup_part, timestamp)
524 statements.append((statement, params))
526 with self._timer(
527 "select_time", tags={"table": "DiaObjectDedup", "method": "getDiaObjectsForDedup"}
528 ) as timer:
529 objects = cast(
530 pandas.DataFrame,
531 select_concurrent(
532 context.session,
533 statements,
534 "read_pandas_multi_dedup",
535 config.connection_config.read_concurrency,
536 ),
537 )
538 timer.add_values(row_count=len(objects), num_queries=num_part)
540 _LOG.debug("found %s DiaObjectDedup records", objects.shape[0])
541 return objects
543 def getDiaSourcesForDiaObjects(
544 self, objects: list[DiaObjectId], start_time: astropy.time.Time, max_dist_arcsec: float = 1.0
545 ) -> pandas.DataFrame:
546 # docstring is inherited from a base class
547 context = self._context
548 config = context.config
550 # Which tables to query and temporal constraints.
551 end_time = self._current_time()
552 tables, temporal_where = context.partitioner.temporal_where(
553 ApdbTables.DiaSource,
554 start_time,
555 end_time,
556 for_prepare=False,
557 partitons_range=context.time_partitions_range,
558 query_per_time_part=False,
559 )
560 if not tables:
561 warnings.warn(
562 f"Query time range ({start_time.isot} - {end_time.isot}) does not overlap database "
563 "time partitions."
564 )
566 # Group DiaObjects by partition.
567 partitioned_object_ids = self._group_dia_objects_by_partition(
568 context.partitioner, objects, max_dist_arcsec
569 )
571 # Columns to return.
572 column_names = context.schema.apdbColumnNames(ApdbTables.DiaSource)
573 what = ",".join(quote_id(column) for column in column_names)
575 # Make a bunch of queries.
576 statements = []
577 suffix = "ALLOW FILTERING"
578 for apdb_part, diaObjectIds in partitioned_object_ids.items():
579 spatial_where = [(f"apdb_part = {apdb_part}", ())]
580 for table in tables:
581 prefix = f'SELECT {what} FROM "{self._keyspace}"."{table}"'
582 for id_chunk in chunk_iterable(diaObjectIds, 10_000):
583 chunk_str = ",".join(str(diaObjectId) for diaObjectId in id_chunk)
584 id_where = (f'"diaObjectId" IN ({chunk_str})', ())
585 statements += list(
586 self._combine_where(prefix, spatial_where, temporal_where, id_where, suffix=suffix)
587 )
589 _LOG.debug("getDiaSourcesForDiaObjects #queries: %s", len(statements))
591 with self._timer(
592 "select_time", tags={"table": "DiaSource", "method": "getDiaSourcesForDiaObjects"}
593 ) as timer:
594 catalog = cast(
595 pandas.DataFrame,
596 select_concurrent(
597 context.session,
598 statements,
599 "read_pandas_multi",
600 config.connection_config.read_concurrency,
601 ),
602 )
603 timer.add_values(row_count_from_db=len(catalog), num_queries=len(statements))
605 # precise filtering on midpointMjdTai
606 catalog = cast(pandas.DataFrame, catalog[catalog["midpointMjdTai"] >= start_time.tai.mjd])
608 timer.add_values(row_count=len(catalog))
610 _LOG.debug("found %d DiaSources", len(catalog))
611 return catalog
613 def containsVisitDetector(
614 self,
615 visit: int,
616 detector: int,
617 region: sphgeom.Region,
618 visit_time: astropy.time.Time,
619 ) -> bool:
620 # docstring is inherited from a base class
621 context = self._context
622 config = context.config
624 # If ApdbDetectorVisit table exists just check it.
625 if context.has_visit_detector_table:
626 table_name = context.schema.tableName(ExtraTables.ApdbVisitDetector)
627 query = (
628 f'SELECT count(*) FROM "{self._keyspace}"."{table_name}" WHERE visit = %s AND detector = %s'
629 )
630 with self._timer("contains_visit_detector_time", tags={"table": table_name}):
631 result = context.session.execute(query, (visit, detector))
632 return bool(result.one()[0])
634 # The order of checks corresponds to order in store(), on potential
635 # store failure earlier tables have higher probability containing
636 # stored records. With per-partition tables there will be many tables
637 # in the list, but it is unlikely that we'll use that setup in
638 # production.
639 sp_where, _ = context.partitioner.spatial_where(region, use_ranges=True, for_prepare=True)
640 visit_detector_where = ("visit = ? AND detector = ?", (visit, detector))
642 # Sources are partitioned on their midPointMjdTai. To avoid precision
643 # issues add some fuzziness to visit time.
644 mjd_start = float(visit_time.tai.mjd) - 1.0 / 24
645 mjd_end = float(visit_time.tai.mjd) + 1.0 / 24
647 statements: list[tuple] = []
648 for table_type in ApdbTables.DiaSource, ApdbTables.DiaForcedSource:
649 tables, temporal_where = context.partitioner.temporal_where(
650 table_type, mjd_start, mjd_end, query_per_time_part=True, for_prepare=True
651 )
652 for table in tables:
653 prefix = f'SELECT apdb_part FROM "{self._keyspace}"."{table}"'
654 # Needs ALLOW FILTERING as there is no PK constraint.
655 suffix = "PER PARTITION LIMIT 1 LIMIT 1 ALLOW FILTERING"
656 statements += list(
657 self._combine_where(prefix, sp_where, temporal_where, visit_detector_where, suffix)
658 )
660 with self._timer("contains_visit_detector_time", tags={"table": "DiaSource"}):
661 result = cast(
662 list[tuple[int] | None],
663 select_concurrent(
664 context.session,
665 statements,
666 "read_tuples",
667 config.connection_config.read_concurrency,
668 ),
669 )
670 return bool(result)
672 def store(
673 self,
674 visit_time: astropy.time.Time,
675 objects: pandas.DataFrame,
676 sources: pandas.DataFrame | None = None,
677 forced_sources: pandas.DataFrame | None = None,
678 ) -> None:
679 # docstring is inherited from a base class
680 context = self._context
681 config = context.config
683 if context.has_visit_detector_table:
684 # Store visit/detector in a special table, this has to be done
685 # before all other writes so if there is a failure at any point
686 # later we still have a record for attempted write.
687 visit_detector: set[tuple[int, int]] = set()
688 for df in sources, forced_sources:
689 if df is not None and not df.empty:
690 df = df[["visit", "detector"]]
691 for visit, detector in df.itertuples(index=False):
692 visit_detector.add((visit, detector))
694 if visit_detector:
695 # Typically there is only one entry, do not bother with
696 # concurrency.
697 table_name = context.schema.tableName(ExtraTables.ApdbVisitDetector)
698 query = f'INSERT INTO "{self._keyspace}"."{table_name}" (visit, detector) VALUES (%s, %s)'
699 for item in visit_detector:
700 context.session.execute(query, item, execution_profile="write")
702 objects = self._fix_input_timestamps(objects)
703 if sources is not None:
704 sources = self._fix_input_timestamps(sources)
705 if forced_sources is not None:
706 forced_sources = self._fix_input_timestamps(forced_sources)
708 replica_chunk: ReplicaChunk | None = None
709 if context.schema.replication_enabled:
710 replica_chunk = ReplicaChunk.make_replica_chunk(visit_time, config.replica_chunk_seconds)
711 self._storeReplicaChunk(replica_chunk)
713 # fill region partition column for DiaObjects
714 objects = self._add_apdb_part(objects)
715 self._storeDiaObjects(objects, visit_time, replica_chunk)
717 if sources is not None and len(sources) > 0:
718 # copy apdb_part column from DiaObjects to DiaSources
719 sources = self._add_apdb_part(sources)
720 subchunk = self._storeDiaSources(ApdbTables.DiaSource, sources, replica_chunk)
721 self._storeDiaSourcesPartitions(sources, visit_time, replica_chunk, subchunk)
723 if forced_sources is not None and len(forced_sources) > 0:
724 forced_sources = self._add_apdb_part(forced_sources)
725 self._storeDiaSources(ApdbTables.DiaForcedSource, forced_sources, replica_chunk)
727 def reassignDiaSourcesToDiaObjects(
728 self,
729 idMap: Mapping[DiaSourceId, int],
730 *,
731 increment_nDiaSources: bool = True,
732 decrement_nDiaSources: bool = True,
733 ) -> None:
734 # docstring is inherited from a base class
735 context = self._context
736 config = context.config
738 source_ids = {source_id.diaSourceId for source_id in idMap}
740 # Find all DiaSources.
741 found_sources = self._get_diasource_data(
742 idMap, "apdb_part", "diaObjectId", "ra", "dec", "midpointMjdTai"
743 )
745 if missing_ids := (source_ids - {row.diaSourceId for row in found_sources}):
746 raise LookupError(f"Some source IDs were not found in DiaSource table: {missing_ids}")
748 found_sources_by_id = {row.diaSourceId: row for row in found_sources}
750 # Make sure that all DiaObjects exist, we also want to know
751 # nDiaSources count for current and new records because we want to
752 # send updated values to replica.
753 current_object_ids = {
754 DiaObjectId(diaObjectId=row.diaObjectId, ra=row.ra, dec=row.dec) for row in found_sources
755 }
756 # Assume that DiaSource ra/dec are very close to re-assigned objects.
757 new_object_ids = {
758 DiaObjectId(diaObjectId=diaObjectId, ra=source_id.ra, dec=source_id.dec)
759 for source_id, diaObjectId in idMap.items()
760 }
761 all_object_ids = new_object_ids | current_object_ids
762 found_objects = self._get_diaobject_data(all_object_ids, "apdb_part", "ra", "dec", "nDiaSources")
764 if missing_ids := (
765 {row.diaObjectId for row in all_object_ids} - {row.diaObjectId for row in found_objects}
766 ):
767 raise LookupError(f"Some object IDs were not found in DiaObjectLast table: {missing_ids}")
769 update_records: list[ApdbUpdateRecord] = []
770 update_order = 0
771 current_time = self._current_time()
772 current_time_ns = int(current_time.unix_tai * 1e9)
774 # Update DiaSources.
775 table_name = context.schema.tableName(ApdbTables.DiaSource)
776 statements: list[tuple] = []
777 for source_id, diaObjectId in idMap.items():
778 source_row = found_sources_by_id[source_id.diaSourceId]
779 apdb_part = source_row.apdb_part
780 time_part = context.partitioner.time_partition(source_row.midpointMjdTai)
782 if config.partitioning.time_partition_tables:
783 statement = context.preparer.prepare(
784 f'UPDATE "{self._keyspace}"."{table_name}_{time_part}" SET "diaObjectId" = ? '
785 'WHERE apdb_part = ? AND "diaSourceId" = ?'
786 )
787 statements.append((statement, (diaObjectId, apdb_part, source_id.diaSourceId)))
788 else:
789 statement = context.preparer.prepare(
790 f'UPDATE "{self._keyspace}"."{table_name}" SET "diaObjectId" = ? '
791 'WHERE apdb_part = ? AND apdb_time_part = ? AND "diaSourceId" = ?'
792 )
793 statements.append((statement, (diaObjectId, apdb_part, time_part, source_id.diaSourceId)))
795 if context.schema.replication_enabled:
796 update_records.append(
797 ApdbReassignDiaSourceToDiaObjectRecord(
798 diaSourceId=source_id.diaSourceId,
799 ra=source_id.ra,
800 dec=source_id.dec,
801 midpointMjdTai=source_id.midpointMjdTai,
802 diaObjectId=diaObjectId,
803 update_time_ns=current_time_ns,
804 update_order=update_order,
805 )
806 )
807 update_order += 1
809 with self._timer(
810 "update_time", tags={"table": table_name, "method": "reassignDiaSourcesToDiaObjects"}
811 ) as timer:
812 execute_concurrent(context.session, statements, execution_profile="write")
813 timer.add_values(num_queries=len(statements))
815 # Update nDiaSources in DiaObjectLast. We do not update DiaObject table
816 # here because it may not even exist. PPDB updates DiaObject from
817 # update records.
818 if increment_nDiaSources or decrement_nDiaSources:
819 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
820 statement = context.preparer.prepare(
821 f'UPDATE "{self._keyspace}"."{table_name}" SET "nDiaSources" = ? '
822 'WHERE apdb_part = ? AND "diaObjectId" = ?'
823 )
824 statements = []
826 # Calculate increments/decrements for all affected DiaObjects.
827 increments: Counter = Counter()
828 if increment_nDiaSources:
829 increments.update(idMap.values())
830 if decrement_nDiaSources:
831 increments.subtract(row.diaObjectId for row in found_sources)
833 for row in found_objects:
834 if increments.get(row.diaObjectId):
835 nDiaSources = row.nDiaSources + increments[row.diaObjectId]
836 statements.append((statement, (nDiaSources, row.apdb_part, row.diaObjectId)))
838 # Also send updated values to replica.
839 if context.schema.replication_enabled:
840 update_records.append(
841 ApdbUpdateNDiaSourcesRecord(
842 diaObjectId=row.diaObjectId,
843 ra=row.ra,
844 dec=row.dec,
845 nDiaSources=nDiaSources,
846 update_time_ns=current_time_ns,
847 update_order=update_order,
848 )
849 )
850 update_order += 1
852 if statements:
853 with self._timer(
854 "update_time", tags={"table": table_name, "method": "reassignDiaSourcesToDiaObjects"}
855 ) as timer:
856 execute_concurrent(context.session, statements, execution_profile="write")
857 timer.add_values(num_queries=len(statements))
859 if update_records:
860 replica_chunk = ReplicaChunk.make_replica_chunk(current_time, config.replica_chunk_seconds)
861 self._storeUpdateRecords(update_records, replica_chunk, store_chunk=True)
863 def setValidityEnd(
864 self, objects: list[DiaObjectId], validityEnd: astropy.time.Time, raise_on_missing_id: bool = False
865 ) -> int:
866 # docstring is inherited from a base class
867 if not objects:
868 return 0
870 context = self._context
871 config = context.config
873 pad_arcsec = 1.0
874 partitioned_object_ids = self._group_dia_objects_by_partition(
875 context.partitioner, objects, pad_arcsec
876 )
878 # Check that all objects exist.
879 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
880 statements: list[tuple] = []
881 for apdb_part, diaObjectIds in partitioned_object_ids.items():
882 id_str = ",".join(str(diaObjectId) for diaObjectId in diaObjectIds)
883 query = (
884 f'SELECT apdb_part, "diaObjectId" FROM "{self._keyspace}"."{table_name}" '
885 f'WHERE apdb_part = %s AND "diaObjectId" IN ({id_str})'
886 )
887 statement = cassandra.query.SimpleStatement(query)
888 statements.append((statement, (apdb_part,)))
890 with self._timer("select_time", tags={"table": table_name, "method": "setValidityEnd"}) as timer:
891 records = cast(
892 list[tuple[int, int]],
893 select_concurrent(
894 context.session,
895 statements,
896 "read_tuples",
897 config.connection_config.read_concurrency,
898 ),
899 )
900 timer.add_values(row_count=len(objects))
902 requested_ids = {obj.diaObjectId for obj in objects}
903 found_ids = {rec[1] for rec in records}
904 if extra_ids := (found_ids - requested_ids):
905 raise RuntimeError(f"Consistency error - found duplicate records for object IDs: {extra_ids}")
906 if raise_on_missing_id:
907 if missing_ids := (requested_ids - found_ids):
908 raise LookupError(f"Some object IDs are missing from DiaObjectLast table: {missing_ids}")
910 # Filter existing records.
911 if len(objects) != len(found_ids):
912 objects = [obj for obj in objects if obj.diaObjectId in found_ids]
914 if not objects:
915 return 0
917 # Group by partitions again.
918 grouped_object_ids: dict[int, list[int]] = defaultdict(list)
919 for apdb_part, diaObjectId in records:
920 grouped_object_ids[apdb_part].append(diaObjectId)
922 # Remove all matching rows from DiaObjectLast.
923 statements = []
924 for apdb_part, diaObjectIds in grouped_object_ids.items():
925 id_str = ",".join(str(diaObjectId) for diaObjectId in diaObjectIds)
926 query = (
927 f'DELETE FROM "{self._keyspace}"."{table_name}" '
928 f'WHERE apdb_part = %s AND "diaObjectId" IN ({id_str})'
929 )
930 statement = cassandra.query.SimpleStatement(query)
931 statements.append((statement, (apdb_part,)))
933 # Also remove from DiaObjectLastToPartition.
934 reverse_table_name = context.schema.tableName(ExtraTables.DiaObjectLastToPartition)
935 id_str = ",".join(str(diaObjectId) for _, diaObjectId in records)
936 query = f'DELETE FROM "{self._keyspace}"."{reverse_table_name}" WHERE "diaObjectId" IN ({id_str})'
937 statement = cassandra.query.SimpleStatement(query)
938 statements.append((statement, ()))
940 with self._timer("delete_time", tags={"table": table_name, "method": "setValidityEnd"}) as timer:
941 execute_concurrent(context.session, statements, execution_profile="write")
942 timer.add_values(row_count=len(records))
944 # If repication is enabled then send all updates.
945 if context.schema.replication_enabled:
946 current_time = self._current_time()
947 current_time_ns = int(current_time.unix_tai * 1e9)
948 replica_chunk = ReplicaChunk.make_replica_chunk(current_time, config.replica_chunk_seconds)
950 update_records = [
951 ApdbCloseDiaObjectValidityRecord(
952 diaObjectId=obj.diaObjectId,
953 ra=obj.ra,
954 dec=obj.dec,
955 update_time_ns=current_time_ns,
956 update_order=index,
957 validityEndMjdTai=float(validityEnd.tai.mjd),
958 nDiaSources=None,
959 )
960 for index, obj in enumerate(objects)
961 ]
963 self._storeUpdateRecords(update_records, replica_chunk, store_chunk=True)
965 return len(objects)
967 def resetDedup(self, dedup_time: astropy.time.Time | None = None) -> None:
968 # docstring is inherited from a base class
969 context = self._context
971 if not context.has_dedup_table:
972 raise TypeError("DiaObjectDedup table does not exist in this APDB instance.")
974 if dedup_time is None:
975 dedup_time = self._current_time()
977 validity_start_column = self._timestamp_column_name("validityStart")
979 # Find latest timestamp in deduplication table.
980 table_name = context.schema.tableName(ExtraTables.DiaObjectDedup)
981 query = f'SELECT MAX("{validity_start_column}") FROM "{self._keyspace}"."{table_name}"'
982 result = context.session.execute(query, execution_profile="read_tuples")
983 max_value = result.one()[0]
984 if self._schema.has_mjd_timestamps:
985 max_validity_start = astropy.time.Time(max_value, format="mjd", scale="tai")
986 else:
987 max_validity_start = astropy.time.Time(max_value, format="datetime", scale="tai")
989 # If max time is lower than dedup time we can do TRUNCATE.
990 if dedup_time >= max_validity_start:
991 query = f'TRUNCATE TABLE "{self._keyspace}"."{table_name}"'
992 context.session.execute(query, execution_profile="write")
993 else:
994 dedup_time_value = self._timestamp_column_value(dedup_time)
995 query = f'DELETE FROM "{self._keyspace}"."{table_name}" WHERE "{validity_start_column}" < %s'
996 context.session.execute(query, (dedup_time_value,), execution_profile="write")
998 # Store dedup time.
999 data = {"dedup_time_iso_tai": dedup_time.tai.to_value("iso")}
1000 data_json = json.dumps(data)
1001 context.metadata.set(context.metadataDedupKey, data_json, force=True)
1003 def reassignDiaSources(self, idMap: Mapping[int, int]) -> None:
1004 # docstring is inherited from a base class
1005 context = self._context
1006 config = context.config
1008 now = self._current_time()
1009 reassign_time_column = self._timestamp_column_name("ssObjectReassocTime")
1010 reassignTime = self._timestamp_column_value(now)
1012 # To update a record we need to know its exact primary key (including
1013 # partition key) so we start by querying for diaSourceId to find the
1014 # primary keys.
1016 table_name = context.schema.tableName(ExtraTables.DiaSourceToPartition)
1017 # split it into 1k IDs per query
1018 selects: list[tuple] = []
1019 for ids in chunk_iterable(idMap.keys(), 1_000):
1020 ids_str = ",".join(str(item) for item in ids)
1021 selects.append(
1022 (
1023 (
1024 'SELECT "diaSourceId", "apdb_part", "apdb_time_part", "apdb_replica_chunk" '
1025 f'FROM "{self._keyspace}"."{table_name}" WHERE "diaSourceId" IN ({ids_str})'
1026 ),
1027 {},
1028 )
1029 )
1031 # No need for DataFrame here, read data as tuples.
1032 result = cast(
1033 list[tuple[int, int, int, int | None]],
1034 select_concurrent(
1035 context.session, selects, "read_tuples", config.connection_config.read_concurrency
1036 ),
1037 )
1039 # Make mapping from source ID to its partition.
1040 id2partitions: dict[int, tuple[int, int]] = {}
1041 id2chunk_id: dict[int, int] = {}
1042 for row in result:
1043 id2partitions[row[0]] = row[1:3]
1044 if row[3] is not None:
1045 id2chunk_id[row[0]] = row[3]
1047 # make sure we know partitions for each ID
1048 if set(id2partitions) != set(idMap):
1049 missing = ",".join(str(item) for item in set(idMap) - set(id2partitions))
1050 raise ValueError(f"Following DiaSource IDs do not exist in the database: {missing}")
1052 # Reassign in standard tables
1053 queries: list[tuple[cassandra.query.PreparedStatement, tuple]] = []
1054 for diaSourceId, ssObjectId in idMap.items():
1055 apdb_part, apdb_time_part = id2partitions[diaSourceId]
1056 values: tuple
1057 if config.partitioning.time_partition_tables:
1058 table_name = context.schema.tableName(ApdbTables.DiaSource, apdb_time_part)
1059 query = (
1060 f'UPDATE "{self._keyspace}"."{table_name}"'
1061 f' SET "ssObjectId" = ?, "diaObjectId" = NULL, "{reassign_time_column}" = ?'
1062 ' WHERE "apdb_part" = ? AND "diaSourceId" = ?'
1063 )
1064 values = (ssObjectId, reassignTime, apdb_part, diaSourceId)
1065 else:
1066 table_name = context.schema.tableName(ApdbTables.DiaSource)
1067 query = (
1068 f'UPDATE "{self._keyspace}"."{table_name}"'
1069 f' SET "ssObjectId" = ?, "diaObjectId" = NULL, "{reassign_time_column}" = ?'
1070 ' WHERE "apdb_part" = ? AND "apdb_time_part" = ? AND "diaSourceId" = ?'
1071 )
1072 values = (ssObjectId, reassignTime, apdb_part, apdb_time_part, diaSourceId)
1073 queries.append((context.preparer.prepare(query), values))
1075 # TODO: (DM-50190) Replication for updated records is not implemented.
1076 if id2chunk_id:
1077 warnings.warn("Replication of reassigned DiaSource records is not implemented.", stacklevel=2)
1079 _LOG.debug("%s: will update %d records", table_name, len(idMap))
1080 with self._timer("source_reassign_time") as timer:
1081 execute_concurrent(context.session, queries, execution_profile="write")
1082 timer.add_values(source_count=len(idMap))
1084 def countUnassociatedObjects(self) -> int:
1085 # docstring is inherited from a base class
1087 # It's too inefficient to implement it for Cassandra in current schema.
1088 raise NotImplementedError()
1090 @property
1091 def schema(self) -> ApdbSchema:
1092 # docstring is inherited from a base class
1093 return self._schema
1095 @property
1096 def metadata(self) -> ApdbMetadata:
1097 # docstring is inherited from a base class
1098 context = self._context
1099 return context.metadata
1101 @property
1102 def admin(self) -> ApdbCassandraAdmin:
1103 # docstring is inherited from a base class
1104 return ApdbCassandraAdmin(self)
1106 def _getSources(
1107 self,
1108 region: sphgeom.Region,
1109 object_ids: Iterable[int] | None,
1110 mjd_start: float,
1111 mjd_end: float,
1112 table_name: ApdbTables,
1113 ) -> pandas.DataFrame:
1114 """Return catalog of DiaSource instances given set of DiaObject IDs.
1116 Parameters
1117 ----------
1118 region : `lsst.sphgeom.Region`
1119 Spherical region.
1120 object_ids :
1121 Collection of DiaObject IDs
1122 mjd_start : `float`
1123 Lower bound of time interval.
1124 mjd_end : `float`
1125 Upper bound of time interval.
1126 table_name : `ApdbTables`
1127 Name of the table.
1129 Returns
1130 -------
1131 catalog : `pandas.DataFrame`, or `None`
1132 Catalog containing DiaSource records. Empty catalog is returned if
1133 ``object_ids`` is empty.
1134 """
1135 context = self._context
1136 config = context.config
1138 object_id_set: Set[int] = set()
1139 if object_ids is not None:
1140 object_id_set = set(object_ids)
1141 if len(object_id_set) == 0:
1142 return self._make_empty_catalog(table_name)
1144 sp_where, num_sp_part = context.partitioner.spatial_where(region, for_prepare=True)
1145 tables, temporal_where = context.partitioner.temporal_where(
1146 table_name, mjd_start, mjd_end, for_prepare=True, partitons_range=context.time_partitions_range
1147 )
1148 if not tables:
1149 start = astropy.time.Time(mjd_start, format="mjd", scale="tai")
1150 end = astropy.time.Time(mjd_end, format="mjd", scale="tai")
1151 warnings.warn(
1152 f"Query time range ({start.isot} - {end.isot}) does not overlap database time partitions."
1153 )
1155 # We need to exclude extra partitioning columns from result.
1156 column_names = context.schema.apdbColumnNames(table_name)
1157 what = ",".join(quote_id(column) for column in column_names)
1159 # Build all queries
1160 statements: list[tuple] = []
1161 for table in tables:
1162 prefix = f'SELECT {what} from "{self._keyspace}"."{table}"'
1163 statements += list(self._combine_where(prefix, sp_where, temporal_where))
1164 _LOG.debug("_getSources %s: #queries: %s", table_name, len(statements))
1166 with self._timer("select_time", tags={"table": table_name.name, "method": "_getSources"}) as timer:
1167 catalog = cast(
1168 pandas.DataFrame,
1169 select_concurrent(
1170 context.session,
1171 statements,
1172 "read_pandas_multi",
1173 config.connection_config.read_concurrency,
1174 ),
1175 )
1176 timer.add_values(
1177 row_count_from_db=len(catalog), num_sp_part=num_sp_part, num_queries=len(statements)
1178 )
1180 # filter by given object IDs
1181 if len(object_id_set) > 0:
1182 catalog = cast(pandas.DataFrame, catalog[catalog["diaObjectId"].isin(object_id_set)])
1184 # precise filtering on midpointMjdTai
1185 catalog = cast(pandas.DataFrame, catalog[catalog["midpointMjdTai"] > mjd_start])
1187 timer.add_values(row_count=len(catalog))
1189 _LOG.debug("found %d %ss", catalog.shape[0], table_name.name)
1190 return catalog
1192 def _storeReplicaChunk(self, replica_chunk: ReplicaChunk) -> None:
1193 context = self._context
1194 config = context.config
1196 # Cassandra timestamp uses milliseconds since epoch
1197 timestamp = int(replica_chunk.last_update_time.unix_tai * 1000)
1199 # everything goes into a single partition
1200 partition = 0
1202 table_name = context.schema.tableName(ExtraTables.ApdbReplicaChunks)
1204 columns = ["partition", "apdb_replica_chunk", "last_update_time", "unique_id"]
1205 values = [partition, replica_chunk.id, timestamp, replica_chunk.unique_id]
1206 if context.has_chunk_sub_partitions:
1207 columns.append("has_subchunks")
1208 values.append(True)
1210 column_list = ", ".join(columns)
1211 placeholders = ",".join(["%s"] * len(columns))
1212 query = f'INSERT INTO "{self._keyspace}"."{table_name}" ({column_list}) VALUES ({placeholders})'
1214 context.session.execute(
1215 query,
1216 values,
1217 timeout=config.connection_config.write_timeout,
1218 execution_profile="write",
1219 )
1221 def _queryDiaObjectLastPartitions(self, ids: Iterable[int]) -> Mapping[int, int]:
1222 """Return existing mapping of diaObjectId to its last partition."""
1223 context = self._context
1224 config = context.config
1226 table_name = context.schema.tableName(ExtraTables.DiaObjectLastToPartition)
1227 queries = []
1228 object_count = 0
1229 for id_chunk in chunk_iterable(ids, 10_000):
1230 id_chunk_list = list(id_chunk)
1231 query = (
1232 f'SELECT "diaObjectId", apdb_part FROM "{self._keyspace}"."{table_name}" '
1233 f'WHERE "diaObjectId" in ({",".join(str(oid) for oid in id_chunk_list)})'
1234 )
1235 queries.append((query, ()))
1236 object_count += len(id_chunk_list)
1238 with self._timer("query_object_last_partitions", tags={"table": table_name}) as timer:
1239 data = cast(
1240 ApdbTableData,
1241 select_concurrent(
1242 context.session,
1243 queries,
1244 "read_raw_multi",
1245 config.connection_config.read_concurrency,
1246 ),
1247 )
1248 timer.add_values(object_count=object_count, row_count=len(data.rows()))
1250 if data.column_names() != ["diaObjectId", "apdb_part"]:
1251 raise RuntimeError(f"Unexpected column names in query result: {data.column_names()}")
1253 return {row[0]: row[1] for row in data.rows()}
1255 def _deleteMovingObjects(self, objs: pandas.DataFrame) -> None:
1256 """Objects in DiaObjectsLast can move from one spatial partition to
1257 another. For those objects inserting new version does not replace old
1258 one, so we need to explicitly remove old versions before inserting new
1259 ones.
1260 """
1261 context = self._context
1263 # Extract all object IDs.
1264 new_partitions = dict(zip(objs["diaObjectId"], objs["apdb_part"]))
1265 old_partitions = self._queryDiaObjectLastPartitions(objs["diaObjectId"])
1267 moved_oids: dict[int, tuple[int, int]] = {}
1268 for oid, old_part in old_partitions.items():
1269 new_part = new_partitions.get(oid, old_part)
1270 if new_part != old_part:
1271 moved_oids[oid] = (old_part, new_part)
1272 _LOG.debug("DiaObject IDs that moved to new partition: %s", moved_oids)
1274 if moved_oids:
1275 # Delete old records from DiaObjectLast.
1276 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
1277 query = f'DELETE FROM "{self._keyspace}"."{table_name}" WHERE apdb_part = ? AND "diaObjectId" = ?'
1278 statement = context.preparer.prepare(query)
1279 queries = []
1280 for oid, (old_part, _) in moved_oids.items():
1281 queries.append((statement, (old_part, oid)))
1282 with self._timer("delete_object_last", tags={"table": table_name}) as timer:
1283 execute_concurrent(context.session, queries, execution_profile="write")
1284 timer.add_values(row_count=len(moved_oids))
1286 # Add all new records to the map.
1287 table_name = context.schema.tableName(ExtraTables.DiaObjectLastToPartition)
1288 query = f'INSERT INTO "{self._keyspace}"."{table_name}" ("diaObjectId", apdb_part) VALUES (?,?)'
1289 statement = context.preparer.prepare(query)
1291 queries = []
1292 for oid, new_part in new_partitions.items():
1293 queries.append((statement, (oid, new_part)))
1295 with self._timer("update_object_last_partition", tags={"table": table_name}) as timer:
1296 execute_concurrent(context.session, queries, execution_profile="write")
1297 timer.add_values(row_count=len(queries))
1299 def _storeDiaObjects(
1300 self, objs: pandas.DataFrame, visit_time: astropy.time.Time, replica_chunk: ReplicaChunk | None
1301 ) -> None:
1302 """Store catalog of DiaObjects from current visit.
1304 Parameters
1305 ----------
1306 objs : `pandas.DataFrame`
1307 Catalog with DiaObject records
1308 visit_time : `astropy.time.Time`
1309 Time of the current visit.
1310 replica_chunk : `ReplicaChunk` or `None`
1311 Replica chunk identifier if replication is configured.
1312 """
1313 if len(objs) == 0:
1314 _LOG.debug("No objects to write to database.")
1315 return
1317 context = self._context
1318 config = context.config
1320 if context.has_dia_object_last_to_partition:
1321 self._deleteMovingObjects(objs)
1323 validity_start_column = self._timestamp_column_name("validityStart")
1324 timestamp = self._timestamp_column_value(visit_time)
1326 # DiaObjectLast did not have this column in the past.
1327 extra_columns: dict[str, Any] = {}
1328 if context.schema.check_column(ApdbTables.DiaObjectLast, validity_start_column):
1329 extra_columns[validity_start_column] = timestamp
1331 self._storeObjectsPandas(objs, ApdbTables.DiaObjectLast, extra_columns=extra_columns)
1333 extra_columns[validity_start_column] = timestamp
1334 visit_time_part = context.partitioner.time_partition(visit_time)
1335 time_part: int | None = visit_time_part
1336 if (time_partitions_range := context.time_partitions_range) is not None:
1337 self._check_time_partitions([visit_time_part], time_partitions_range)
1338 if not config.partitioning.time_partition_tables:
1339 extra_columns["apdb_time_part"] = time_part
1340 time_part = None
1342 # Only store DiaObects if not doing replication or explicitly
1343 # configured to always store them.
1344 if replica_chunk is None or not config.replica_skips_diaobjects:
1345 self._storeObjectsPandas(
1346 objs, ApdbTables.DiaObject, extra_columns=extra_columns, time_part=time_part
1347 )
1349 if replica_chunk is not None:
1350 extra_columns = {"apdb_replica_chunk": replica_chunk.id, validity_start_column: timestamp}
1351 table = ExtraTables.DiaObjectChunks
1352 if context.has_chunk_sub_partitions:
1353 table = ExtraTables.DiaObjectChunks2
1354 # Use a random number for a second part of partitioning key so
1355 # that different clients could wrtite to different partitions.
1356 # This makes it not exactly reproducible.
1357 extra_columns["apdb_replica_subchunk"] = random.randrange(config.replica_sub_chunk_count)
1358 self._storeObjectsPandas(objs, table, extra_columns=extra_columns)
1360 # Store copy of the records in dedup table.
1361 if context.has_dedup_table:
1362 table = ExtraTables.DiaObjectDedup
1363 extra_columns = {
1364 "dedup_part": random.randrange(config.partitioning.num_part_dedup),
1365 validity_start_column: timestamp,
1366 }
1367 self._storeObjectsPandas(objs, table, extra_columns=extra_columns)
1369 def _storeDiaSources(
1370 self,
1371 table_name: ApdbTables,
1372 sources: pandas.DataFrame,
1373 replica_chunk: ReplicaChunk | None,
1374 ) -> int | None:
1375 """Store catalog of DIASources or DIAForcedSources from current visit.
1377 Parameters
1378 ----------
1379 table_name : `ApdbTables`
1380 Table where to store the data.
1381 sources : `pandas.DataFrame`
1382 Catalog containing DiaSource records
1383 visit_time : `astropy.time.Time`
1384 Time of the current visit.
1385 replica_chunk : `ReplicaChunk` or `None`
1386 Replica chunk identifier if replication is configured.
1388 Returns
1389 -------
1390 subchunk : `int` or `None`
1391 Subchunk number for resulting replica data, `None` if relication is
1392 not enabled ot subchunking is not enabled.
1393 """
1394 context = self._context
1395 config = context.config
1397 # Time partitioning has to be based on midpointMjdTai, not visit_time
1398 # as visit_time is not really a visit time.
1399 tp_sources = sources.copy(deep=False)
1400 tp_sources["apdb_time_part"] = tp_sources["midpointMjdTai"].apply(context.partitioner.time_partition)
1401 if (time_partitions_range := context.time_partitions_range) is not None:
1402 self._check_time_partitions(tp_sources["apdb_time_part"], time_partitions_range)
1403 extra_columns: dict[str, Any] = {}
1404 if not config.partitioning.time_partition_tables:
1405 self._storeObjectsPandas(tp_sources, table_name)
1406 else:
1407 # Group by time partition
1408 partitions = set(tp_sources["apdb_time_part"])
1409 if len(partitions) == 1:
1410 # Single partition - just save the whole thing.
1411 time_part = partitions.pop()
1412 self._storeObjectsPandas(sources, table_name, time_part=time_part)
1413 else:
1414 # group by time partition.
1415 for time_part, sub_frame in tp_sources.groupby(by="apdb_time_part"):
1416 sub_frame.drop(columns="apdb_time_part", inplace=True)
1417 self._storeObjectsPandas(sub_frame, table_name, time_part=time_part)
1419 subchunk: int | None = None
1420 if replica_chunk is not None:
1421 extra_columns = {"apdb_replica_chunk": replica_chunk.id}
1422 if context.has_chunk_sub_partitions:
1423 subchunk = random.randrange(config.replica_sub_chunk_count)
1424 extra_columns["apdb_replica_subchunk"] = subchunk
1425 if table_name is ApdbTables.DiaSource:
1426 extra_table = ExtraTables.DiaSourceChunks2
1427 else:
1428 extra_table = ExtraTables.DiaForcedSourceChunks2
1429 else:
1430 if table_name is ApdbTables.DiaSource:
1431 extra_table = ExtraTables.DiaSourceChunks
1432 else:
1433 extra_table = ExtraTables.DiaForcedSourceChunks
1434 self._storeObjectsPandas(sources, extra_table, extra_columns=extra_columns)
1436 return subchunk
1438 def _check_time_partitions(
1439 self, partitions: Iterable[int], time_partitions_range: ApdbCassandraTimePartitionRange
1440 ) -> None:
1441 """Check that time partitons for new data actually exist.
1443 Parameters
1444 ----------
1445 partitions : `~collections.abc.Iterable` [`int`]
1446 Time partitions for new data.
1447 time_partitions_range : `ApdbCassandraTimePartitionRange`
1448 Currrent time partition range.
1449 """
1450 partitions = set(partitions)
1451 min_part = min(partitions)
1452 max_part = max(partitions)
1453 if min_part < time_partitions_range.start or max_part > time_partitions_range.end:
1454 raise ValueError(
1455 "Attempt to store data for time partitions that do not yet exist. "
1456 f"Partitons for new records: {min_part}-{max_part}. "
1457 f"Database partitons: {time_partitions_range.start}-{time_partitions_range.end}."
1458 )
1459 # Make a noise when writing to the last partition.
1460 if max_part == time_partitions_range.end:
1461 warnings.warn(
1462 "Writing into the last temporal partition. Partition range needs to be extended soon.",
1463 stacklevel=3,
1464 )
1466 def _storeDiaSourcesPartitions(
1467 self,
1468 sources: pandas.DataFrame,
1469 visit_time: astropy.time.Time,
1470 replica_chunk: ReplicaChunk | None,
1471 subchunk: int | None,
1472 ) -> None:
1473 """Store mapping of diaSourceId to its partitioning values.
1475 Parameters
1476 ----------
1477 sources : `pandas.DataFrame`
1478 Catalog containing DiaSource records
1479 visit_time : `astropy.time.Time`
1480 Time of the current visit.
1481 replica_chunk : `ReplicaChunk` or `None`
1482 Replication chunk, or `None` when replication is disabled.
1483 subchunk : `int` or `None`
1484 Replication sub-chunk, or `None` when replication is disabled or
1485 sub-chunking is not used.
1486 """
1487 context = self._context
1489 id_map = cast(pandas.DataFrame, sources[["diaSourceId", "apdb_part"]])
1490 extra_columns = {
1491 "apdb_time_part": context.partitioner.time_partition(visit_time),
1492 "apdb_replica_chunk": replica_chunk.id if replica_chunk is not None else None,
1493 }
1494 if context.has_chunk_sub_partitions:
1495 extra_columns["apdb_replica_subchunk"] = subchunk
1497 self._storeObjectsPandas(
1498 id_map, ExtraTables.DiaSourceToPartition, extra_columns=extra_columns, time_part=None
1499 )
1501 def _storeObjectsPandas(
1502 self,
1503 records: pandas.DataFrame,
1504 table_name: ApdbTables | ExtraTables,
1505 extra_columns: Mapping | None = None,
1506 time_part: int | None = None,
1507 ) -> None:
1508 """Store generic objects.
1510 Takes Pandas catalog and stores a bunch of records in a table.
1512 Parameters
1513 ----------
1514 records : `pandas.DataFrame`
1515 Catalog containing object records
1516 table_name : `ApdbTables`
1517 Name of the table as defined in APDB schema.
1518 extra_columns : `dict`, optional
1519 Mapping (column_name, column_value) which gives fixed values for
1520 columns in each row, overrides values in ``records`` if matching
1521 columns exist there.
1522 time_part : `int`, optional
1523 If not `None` then insert into a per-partition table.
1525 Notes
1526 -----
1527 If Pandas catalog contains additional columns not defined in table
1528 schema they are ignored. Catalog does not have to contain all columns
1529 defined in a table, but partition and clustering keys must be present
1530 in a catalog or ``extra_columns``.
1531 """
1532 context = self._context
1534 # use extra columns if specified
1535 if extra_columns is None:
1536 extra_columns = {}
1537 extra_fields = list(extra_columns.keys())
1539 # Fields that will come from dataframe.
1540 df_fields = [column for column in records.columns if column not in extra_fields]
1542 column_map = context.schema.getColumnMap(table_name)
1543 # list of columns (as in felis schema)
1544 fields = [column_map[field].name for field in df_fields if field in column_map]
1545 fields += extra_fields
1547 # check that all partitioning and clustering columns are defined
1548 partition_columns = context.schema.partitionColumns(table_name)
1549 required_columns = partition_columns + context.schema.clusteringColumns(table_name)
1550 missing_columns = [column for column in required_columns if column not in fields]
1551 if missing_columns:
1552 raise ValueError(f"Primary key columns are missing from catalog: {missing_columns}")
1554 qfields = [quote_id(field) for field in fields]
1555 qfields_str = ",".join(qfields)
1557 batch_size = self._batch_size(table_name)
1559 with self._timer("insert_build_time", tags={"table": table_name.name}):
1560 # Multi-partition batches are problematic in general, so we want to
1561 # group records in a batch by their partition key.
1562 values_by_key: dict[tuple, list[list]] = defaultdict(list)
1563 for rec in records.itertuples(index=False):
1564 values = []
1565 partitioning_values: dict[str, Any] = {}
1566 for field in df_fields:
1567 if field not in column_map:
1568 continue
1569 value = getattr(rec, field)
1570 if column_map[field].datatype is felis.datamodel.DataType.timestamp:
1571 if isinstance(value, pandas.Timestamp):
1572 value = value.to_pydatetime()
1573 elif value is pandas.NaT:
1574 value = None
1575 else:
1576 # Assume it's seconds since epoch, Cassandra
1577 # datetime is in milliseconds
1578 value = int(value * 1000)
1579 value = literal(value)
1580 values.append(UNSET_VALUE if value is None else value)
1581 if field in partition_columns:
1582 partitioning_values[field] = value
1583 for field in extra_fields:
1584 value = literal(extra_columns[field])
1585 values.append(UNSET_VALUE if value is None else value)
1586 if field in partition_columns:
1587 partitioning_values[field] = value
1589 key = tuple(partitioning_values[field] for field in partition_columns)
1590 values_by_key[key].append(values)
1592 table = context.schema.tableName(table_name, time_part)
1594 holders = ",".join(["?"] * len(qfields))
1595 query = f'INSERT INTO "{self._keyspace}"."{table}" ({qfields_str}) VALUES ({holders})'
1596 statement = context.preparer.prepare(query)
1597 # Cassandra has 64k limit on batch size, normally that should be
1598 # enough but some tests generate too many forced sources.
1599 queries = []
1600 for key_values in values_by_key.values():
1601 for values_chunk in chunk_iterable(key_values, batch_size):
1602 batch = cassandra.query.BatchStatement()
1603 for row_values in values_chunk:
1604 batch.add(statement, row_values)
1605 queries.append((batch, None))
1606 assert batch.routing_key is not None and batch.keyspace is not None
1608 _LOG.debug("%s: will store %d records", context.schema.tableName(table_name), records.shape[0])
1609 with self._timer(
1610 "insert_time", tags={"table": table_name.name, "method": "_storeObjectsPandas"}
1611 ) as timer:
1612 execute_concurrent(context.session, queries, execution_profile="write")
1613 timer.add_values(row_count=len(records), num_batches=len(queries))
1615 def _storeUpdateRecords(
1616 self, records: Iterable[ApdbUpdateRecord], chunk: ReplicaChunk, *, store_chunk: bool = False
1617 ) -> None:
1618 """Store ApdbUpdateRecords in the replica table for those records.
1620 Parameters
1621 ----------
1622 records : `list` [`ApdbUpdateRecord`]
1623 Records to store.
1624 chunk : `ReplicaChunk`
1625 Replica chunk for these records.
1626 store_chunk : `bool`
1627 If True then also store replica chunk.
1629 Raises
1630 ------
1631 TypeError
1632 Raised if replication is not enabled for this instance.
1633 """
1634 context = self._context
1635 config = context.config
1637 if not context.schema.replication_enabled:
1638 raise TypeError("Replication is not enabled for this APDB instance.")
1640 if store_chunk:
1641 self._storeReplicaChunk(chunk)
1643 apdb_replica_chunk = chunk.id
1644 # Do not use unique_if from ReplicaChunk as it could be reused in
1645 # multiple calls to this method.
1646 update_unique_id = uuid.uuid4()
1648 rows = []
1649 for record in records:
1650 rows.append(
1651 [
1652 apdb_replica_chunk,
1653 record.update_time_ns,
1654 record.update_order,
1655 update_unique_id,
1656 record.to_json(),
1657 ]
1658 )
1659 columns = [
1660 "apdb_replica_chunk",
1661 "update_time_ns",
1662 "update_order",
1663 "update_unique_id",
1664 "update_payload",
1665 ]
1666 if context.has_chunk_sub_partitions:
1667 subchunk = random.randrange(config.replica_sub_chunk_count)
1668 for row in rows:
1669 row.append(subchunk)
1670 columns.append("apdb_replica_subchunk")
1672 table_name = context.schema.tableName(ExtraTables.ApdbUpdateRecordChunks)
1673 placeholders = ", ".join(["%s"] * len(columns))
1674 columns_str = ", ".join(columns)
1675 query = f'INSERT INTO "{self._keyspace}"."{table_name}" ({columns_str}) VALUES ({placeholders})'
1676 queries = [(query, row) for row in rows]
1678 with self._timer("store_update_record", tags={"table": table_name}) as timer:
1679 execute_concurrent(context.session, queries, execution_profile="write")
1680 timer.add_values(row_count=len(queries))
1682 def _add_apdb_part(self, df: pandas.DataFrame) -> pandas.DataFrame:
1683 """Calculate spatial partition for each record and add it to a
1684 DataFrame.
1686 Parameters
1687 ----------
1688 df : `pandas.DataFrame`
1689 DataFrame which has to contain ra/dec columns, names of these
1690 columns are defined by configuration ``ra_dec_columns`` field.
1692 Returns
1693 -------
1694 df : `pandas.DataFrame`
1695 DataFrame with ``apdb_part`` column which contains pixel index
1696 for ra/dec coordinates.
1698 Notes
1699 -----
1700 This overrides any existing column in a DataFrame with the same name
1701 (``apdb_part``). Original DataFrame is not changed, copy of a DataFrame
1702 is returned.
1703 """
1704 context = self._context
1705 config = context.config
1707 # Calculate pixelization index for every record.
1708 apdb_part = np.zeros(df.shape[0], dtype=np.int64)
1709 ra_col, dec_col = config.ra_dec_columns
1710 for i, (ra, dec) in enumerate(zip(df[ra_col], df[dec_col])):
1711 uv3d = sphgeom.UnitVector3d(sphgeom.LonLat.fromDegrees(ra, dec))
1712 idx = context.partitioner.pixel(uv3d)
1713 apdb_part[i] = idx
1714 df = df.copy()
1715 df["apdb_part"] = apdb_part
1716 return df
1718 def _make_empty_catalog(self, table_name: ApdbTables) -> pandas.DataFrame:
1719 """Make an empty catalog for a table with a given name.
1721 Parameters
1722 ----------
1723 table_name : `ApdbTables`
1724 Name of the table.
1726 Returns
1727 -------
1728 catalog : `pandas.DataFrame`
1729 An empty catalog.
1730 """
1731 table = self.schema.tableSchemas[table_name]
1733 data = {
1734 columnDef.name: pandas.Series(dtype=self.schema.column_dtype(columnDef.datatype))
1735 for columnDef in table.columns
1736 }
1737 return pandas.DataFrame(data)
1739 def _combine_where(
1740 self,
1741 prefix: str,
1742 where1: list[tuple[str, tuple]],
1743 where2: list[tuple[str, tuple]],
1744 where3: tuple[str, tuple] | None = None,
1745 suffix: str | None = None,
1746 ) -> Iterator[tuple[cassandra.query.Statement, tuple]]:
1747 """Make cartesian product of two parts of WHERE clause into a series
1748 of statements to execute.
1750 Parameters
1751 ----------
1752 prefix : `str`
1753 Initial statement prefix that comes before WHERE clause, e.g.
1754 "SELECT * from Table"
1755 """
1756 context = self._context
1758 # If lists are empty use special sentinels.
1759 if not where1:
1760 where1 = [("", ())]
1761 if not where2:
1762 where2 = [("", ())]
1764 for expr1, params1 in where1:
1765 for expr2, params2 in where2:
1766 full_query = prefix
1767 wheres = []
1768 params = params1 + params2
1769 if expr1:
1770 wheres.append(expr1)
1771 if expr2:
1772 wheres.append(expr2)
1773 if where3:
1774 wheres.append(where3[0])
1775 params += where3[1]
1776 if wheres:
1777 full_query += " WHERE " + " AND ".join(wheres)
1778 if suffix:
1779 full_query += " " + suffix
1780 if params:
1781 statement = context.preparer.prepare(full_query)
1782 else:
1783 # If there are no params then it is likely that query
1784 # has a bunch of literals rendered already, no point
1785 # trying to prepare it.
1786 statement = cassandra.query.SimpleStatement(full_query)
1787 yield (statement, params)
1789 def _fix_input_timestamps(self, df: pandas.DataFrame) -> pandas.DataFrame:
1790 """Update timestamp columns in input DataFrame to be naive datetime
1791 type.
1793 Clients may or may not generate aware timestamps, code in this class
1794 assumes that timestamps are naive, so we convert them to UTC and
1795 drop timezone.
1796 """
1797 # Find all columns with aware timestamps.
1798 columns = [column for column, dtype in df.dtypes.items() if isinstance(dtype, pandas.DatetimeTZDtype)]
1799 for column in columns:
1800 # tz_convert(None) will convert to UTC and drop timezone.
1801 df[column] = df[column].dt.tz_convert(None)
1802 return df
1804 def _batch_size(self, table: ApdbTables | ExtraTables) -> int:
1805 """Calculate batch size based on config parameters."""
1806 context = self._context
1807 config = context.config
1809 # Cassandra limit on number of statements in a batch is 64k.
1810 batch_size = 65_535
1811 if 0 < config.batch_statement_limit < batch_size:
1812 batch_size = config.batch_statement_limit
1813 if config.batch_size_limit > 0:
1814 # The purpose of this limit is to try not to exceed batch size
1815 # threshold which is set on server side. Cassandra wire protocol
1816 # for prepared queries (and batches) only sends column values with
1817 # with an additional 4 bytes per value specifying size. Value is
1818 # not included for NULL or NOT_SET values, but the size is always
1819 # there. There is additional small per-query overhead, which we
1820 # ignore.
1821 row_size = context.schema.table_row_size(table)
1822 row_size += 4 * len(context.schema.getColumnMap(table))
1823 batch_size = min(batch_size, (config.batch_size_limit // row_size) + 1)
1824 return batch_size
1826 def _group_dia_objects_by_partition(
1827 self, partitioner: Partitioner, objects: list[DiaObjectId], pad_arcsec: float
1828 ) -> Mapping[int, list[int]]:
1829 """Group DiaObjects by partition.
1831 Parameters
1832 ----------
1833 partitioner : `Partitioner`
1834 Objects which knows how to partition things.
1835 objects : `list` [`DiaObjectId`]
1836 Collection of objects to partition.
1837 pad_arcsec : `float`
1838 Additional padding around object position.
1840 Returns
1841 -------
1842 grouped_objects
1843 Mapping of spatial patition ID to list ob object IDs that it
1844 contains. Some objects may belong to more than one partition.
1845 """
1846 partitioned_object_ids: dict[int, list[int]] = defaultdict(list)
1847 for obj_id in objects:
1848 partitions = partitioner.pixelization.circle_pixels(obj_id.ra, obj_id.dec, pad_arcsec)
1849 for pixel in partitions:
1850 partitioned_object_ids[pixel].append(obj_id.diaObjectId)
1851 return partitioned_object_ids
1853 def _timestamp_column_name(self, column: str) -> str:
1854 """Return column name before/after schema migration to MJD TAI."""
1855 return self._schema.timestamp_column_name(column)
1857 def _timestamp_column_value(self, time: astropy.time.Time) -> float | int:
1858 """Return column value before/after schema migration to MJD TAI."""
1859 if self._schema.has_mjd_timestamps:
1860 return float(time.tai.mjd)
1861 else:
1862 return int(time.datetime.astimezone(tz=datetime.UTC).timestamp() * 1000)
1864 def _get_diasource_data(self, source_ids: Iterable[DiaSourceId], *columns: str) -> list:
1865 """Select records from DiaSource table by diaSourceId and return all
1866 records as a list of named tuples.
1867 """
1868 context = self._context
1869 config = context.config
1870 partitioner = context.partitioner
1872 columns = ("diaSourceId",) + columns
1873 what = ",".join(quote_id(column) for column in columns)
1875 # Allow some uncertainty for coordinates and time when calculating
1876 # partitions.
1877 statements: list[tuple] = []
1878 pad_arcsec = 1.0
1879 pad_time_day = 10 / (24 * 3600)
1880 for source_id in source_ids:
1881 center = sphgeom.UnitVector3d(sphgeom.LonLat.fromDegrees(source_id.ra, source_id.dec))
1882 region = sphgeom.Circle(center, sphgeom.Angle.fromDegrees(pad_arcsec / 3600.0))
1883 spatial_where, _ = partitioner.spatial_where(region, for_prepare=True)
1885 tables, temporal_where = partitioner.temporal_where(
1886 ApdbTables.DiaSource,
1887 source_id.midpointMjdTai - pad_time_day,
1888 source_id.midpointMjdTai + pad_time_day,
1889 for_prepare=True,
1890 partitons_range=context.time_partitions_range,
1891 query_per_time_part=True,
1892 )
1894 id_where = ('"diaSourceId" = ?', (source_id.diaSourceId,))
1896 for table in tables:
1897 prefix = f'SELECT {what} FROM "{self._keyspace}"."{table}"'
1898 statements += list(self._combine_where(prefix, spatial_where, temporal_where, id_where))
1900 with self._timer(
1901 "select_time", tags={"table": "DiaSource", "method": "_get_diasource_data"}
1902 ) as timer:
1903 result = cast(
1904 list[tuple],
1905 select_concurrent(
1906 context.session,
1907 statements,
1908 "read_named_tuples",
1909 config.connection_config.read_concurrency,
1910 ),
1911 )
1912 timer.add_values(row_count=len(result), num_queries=len(statements))
1914 return result
1916 def _get_diaobject_data(self, object_ids: Iterable[DiaObjectId], *columns: str) -> list:
1917 """Select records from DiaObjectLast table by diaObjectId and return
1918 all records as a list of named tuples.
1919 """
1920 context = self._context
1921 config = context.config
1922 partitioner = context.partitioner
1924 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
1925 columns = ("diaObjectId",) + columns
1926 what = ",".join(quote_id(column) for column in columns)
1928 # Allow some uncertainty for coordinates when calculating partitions.
1929 pad_arcsec = 1.0
1930 ids_by_partition = defaultdict(list)
1931 for object_id in object_ids:
1932 pixels = partitioner.pixelization.circle_pixels(object_id.ra, object_id.dec, pad_arcsec)
1933 for pixel in pixels:
1934 ids_by_partition[pixel].append(object_id.diaObjectId)
1936 statements: list[tuple] = []
1937 for apdb_part, diaObjectIds in ids_by_partition.items():
1938 ids_str = ",".join(str(diaObjectId) for diaObjectId in diaObjectIds)
1939 statement = cassandra.query.SimpleStatement(
1940 f'SELECT {what} FROM "{self._keyspace}"."{table_name}" '
1941 f'WHERE apdb_part = %s AND "diaObjectId" IN ({ids_str})'
1942 )
1943 statements.append((statement, (apdb_part,)))
1945 with self._timer(
1946 "select_time", tags={"table": "DiaObjectLast", "method": "_get_diaobject_data"}
1947 ) as timer:
1948 result = cast(
1949 list[tuple],
1950 select_concurrent(
1951 context.session,
1952 statements,
1953 "read_named_tuples",
1954 config.connection_config.read_concurrency,
1955 ),
1956 )
1957 timer.add_values(row_count=len(result), num_queries=len(statements))
1959 return result