Coverage for python / lsst / dax / apdb / cassandra / apdbCassandra.py: 8%
859 statements
« prev ^ index » next coverage.py v7.13.2, created at 2026-02-04 01:08 -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
626 config = context.config
628 # If ApdbDetectorVisit table exists just check it.
629 if context.has_visit_detector_table:
630 table_name = context.schema.tableName(ExtraTables.ApdbVisitDetector)
631 query = (
632 f'SELECT count(*) FROM "{self._keyspace}"."{table_name}" WHERE visit = %s AND detector = %s'
633 )
634 with self._timer("contains_visit_detector_time", tags={"table": table_name}):
635 result = context.session.execute(query, (visit, detector))
636 return bool(result.one()[0])
638 # The order of checks corresponds to order in store(), on potential
639 # store failure earlier tables have higher probability containing
640 # stored records. With per-partition tables there will be many tables
641 # in the list, but it is unlikely that we'll use that setup in
642 # production.
643 sp_where, _ = context.partitioner.spatial_where(region, use_ranges=True, for_prepare=True)
644 visit_detector_where = ("visit = ? AND detector = ?", (visit, detector))
646 # Sources are partitioned on their midPointMjdTai. To avoid precision
647 # issues add some fuzziness to visit time.
648 mjd_start = float(visit_time.tai.mjd) - 1.0 / 24
649 mjd_end = float(visit_time.tai.mjd) + 1.0 / 24
651 statements: list[tuple] = []
652 for table_type in ApdbTables.DiaSource, ApdbTables.DiaForcedSource:
653 tables, temporal_where = context.partitioner.temporal_where(
654 table_type, mjd_start, mjd_end, query_per_time_part=True, for_prepare=True
655 )
656 for table in tables:
657 prefix = f'SELECT apdb_part FROM "{self._keyspace}"."{table}"'
658 # Needs ALLOW FILTERING as there is no PK constraint.
659 suffix = "PER PARTITION LIMIT 1 LIMIT 1 ALLOW FILTERING"
660 statements += list(
661 self._combine_where(prefix, sp_where, temporal_where, visit_detector_where, suffix)
662 )
664 with self._timer("contains_visit_detector_time", tags={"table": "DiaSource"}):
665 result = cast(
666 list[tuple[int] | None],
667 select_concurrent(
668 context.session,
669 statements,
670 "read_tuples",
671 config.connection_config.read_concurrency,
672 ),
673 )
674 return bool(result)
676 def store(
677 self,
678 visit_time: astropy.time.Time,
679 objects: pandas.DataFrame,
680 sources: pandas.DataFrame | None = None,
681 forced_sources: pandas.DataFrame | None = None,
682 ) -> None:
683 # docstring is inherited from a base class
684 context = self._context
685 config = context.config
687 if context.has_visit_detector_table:
688 # Store visit/detector in a special table, this has to be done
689 # before all other writes so if there is a failure at any point
690 # later we still have a record for attempted write.
691 visit_detector: set[tuple[int, int]] = set()
692 for df in sources, forced_sources:
693 if df is not None and not df.empty:
694 df = df[["visit", "detector"]]
695 for visit, detector in df.itertuples(index=False):
696 visit_detector.add((visit, detector))
698 if visit_detector:
699 # Typically there is only one entry, do not bother with
700 # concurrency.
701 table_name = context.schema.tableName(ExtraTables.ApdbVisitDetector)
702 query = f'INSERT INTO "{self._keyspace}"."{table_name}" (visit, detector) VALUES (%s, %s)'
703 for item in visit_detector:
704 context.session.execute(query, item, execution_profile="write")
706 objects = self._fix_input_timestamps(objects)
707 if sources is not None:
708 sources = self._fix_input_timestamps(sources)
709 if forced_sources is not None:
710 forced_sources = self._fix_input_timestamps(forced_sources)
712 replica_chunk: ReplicaChunk | None = None
713 if context.schema.replication_enabled:
714 replica_chunk = ReplicaChunk.make_replica_chunk(visit_time, config.replica_chunk_seconds)
715 self._storeReplicaChunk(replica_chunk)
717 # fill region partition column for DiaObjects
718 objects = self._add_apdb_part(objects)
719 self._storeDiaObjects(objects, visit_time, replica_chunk)
721 if sources is not None and len(sources) > 0:
722 # copy apdb_part column from DiaObjects to DiaSources
723 sources = self._add_apdb_part(sources)
724 subchunk = self._storeDiaSources(ApdbTables.DiaSource, sources, replica_chunk)
725 self._storeDiaSourcesPartitions(sources, visit_time, replica_chunk, subchunk)
727 if forced_sources is not None and len(forced_sources) > 0:
728 forced_sources = self._add_apdb_part(forced_sources)
729 self._storeDiaSources(ApdbTables.DiaForcedSource, forced_sources, replica_chunk)
731 def reassignDiaSourcesToDiaObjects(
732 self,
733 idMap: Mapping[DiaSourceId, int],
734 *,
735 increment_nDiaSources: bool = True,
736 decrement_nDiaSources: bool = True,
737 ) -> None:
738 # docstring is inherited from a base class
739 context = self._context
740 config = context.config
742 source_ids = {source_id.diaSourceId for source_id in idMap}
744 # Find all DiaSources.
745 found_sources = self._get_diasource_data(
746 idMap, "apdb_part", "diaObjectId", "ra", "dec", "midpointMjdTai"
747 )
749 if missing_ids := (source_ids - {row.diaSourceId for row in found_sources}):
750 raise LookupError(f"Some source IDs were not found in DiaSource table: {missing_ids}")
752 found_sources_by_id = {row.diaSourceId: row for row in found_sources}
754 # Make sure that all DiaObjects exist, we also want to know
755 # nDiaSources count for current and new records because we want to
756 # send updated values to replica.
757 current_object_ids = {
758 DiaObjectId(diaObjectId=row.diaObjectId, ra=row.ra, dec=row.dec) for row in found_sources
759 }
760 # Assume that DiaSource ra/dec are very close to re-assigned objects.
761 new_object_ids = {
762 DiaObjectId(diaObjectId=diaObjectId, ra=source_id.ra, dec=source_id.dec)
763 for source_id, diaObjectId in idMap.items()
764 }
765 all_object_ids = new_object_ids | current_object_ids
766 found_objects = self._get_diaobject_data(all_object_ids, "apdb_part", "ra", "dec", "nDiaSources")
768 if missing_ids := (
769 {row.diaObjectId for row in all_object_ids} - {row.diaObjectId for row in found_objects}
770 ):
771 raise LookupError(f"Some object IDs were not found in DiaObjectLast table: {missing_ids}")
773 update_records: list[ApdbUpdateRecord] = []
774 update_order = 0
775 current_time = self._current_time()
776 current_time_ns = int(current_time.unix_tai * 1e9)
778 # Update DiaSources.
779 table_name = context.schema.tableName(ApdbTables.DiaSource)
780 statements: list[tuple] = []
781 for source_id, diaObjectId in idMap.items():
782 source_row = found_sources_by_id[source_id.diaSourceId]
783 apdb_part = source_row.apdb_part
784 time_part = context.partitioner.time_partition(source_row.midpointMjdTai)
786 if config.partitioning.time_partition_tables:
787 statement = context.preparer.prepare(
788 f'UPDATE "{self._keyspace}"."{table_name}_{time_part}" SET "diaObjectId" = ? '
789 'WHERE apdb_part = ? AND "diaSourceId" = ?'
790 )
791 statements.append((statement, (diaObjectId, apdb_part, source_id.diaSourceId)))
792 else:
793 statement = context.preparer.prepare(
794 f'UPDATE "{self._keyspace}"."{table_name}" SET "diaObjectId" = ? '
795 'WHERE apdb_part = ? AND apdb_time_part = ? AND "diaSourceId" = ?'
796 )
797 statements.append((statement, (diaObjectId, apdb_part, time_part, source_id.diaSourceId)))
799 if context.schema.replication_enabled:
800 update_records.append(
801 ApdbReassignDiaSourceToDiaObjectRecord(
802 diaSourceId=source_id.diaSourceId,
803 ra=source_id.ra,
804 dec=source_id.dec,
805 midpointMjdTai=source_id.midpointMjdTai,
806 diaObjectId=diaObjectId,
807 update_time_ns=current_time_ns,
808 update_order=update_order,
809 )
810 )
811 update_order += 1
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 # Update nDiaSources in DiaObjectLast. We do not update DiaObject table
820 # here because it may not even exist. PPDB updates DiaObject from
821 # update records.
822 if increment_nDiaSources or decrement_nDiaSources:
823 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
824 statement = context.preparer.prepare(
825 f'UPDATE "{self._keyspace}"."{table_name}" SET "nDiaSources" = ? '
826 'WHERE apdb_part = ? AND "diaObjectId" = ?'
827 )
828 statements = []
830 # Calculate increments/decrements for all affected DiaObjects.
831 increments: Counter = Counter()
832 if increment_nDiaSources:
833 increments.update(idMap.values())
834 if decrement_nDiaSources:
835 increments.subtract(row.diaObjectId for row in found_sources)
837 for row in found_objects:
838 if increments.get(row.diaObjectId):
839 nDiaSources = row.nDiaSources + increments[row.diaObjectId]
840 statements.append((statement, (nDiaSources, row.apdb_part, row.diaObjectId)))
842 # Also send updated values to replica.
843 if context.schema.replication_enabled:
844 update_records.append(
845 ApdbUpdateNDiaSourcesRecord(
846 diaObjectId=row.diaObjectId,
847 ra=row.ra,
848 dec=row.dec,
849 nDiaSources=nDiaSources,
850 update_time_ns=current_time_ns,
851 update_order=update_order,
852 )
853 )
854 update_order += 1
856 if statements:
857 with self._timer(
858 "update_time", tags={"table": table_name, "method": "reassignDiaSourcesToDiaObjects"}
859 ) as timer:
860 execute_concurrent(context.session, statements, execution_profile="write")
861 timer.add_values(num_queries=len(statements))
863 if update_records:
864 replica_chunk = ReplicaChunk.make_replica_chunk(current_time, config.replica_chunk_seconds)
865 self._storeUpdateRecords(update_records, replica_chunk, store_chunk=True)
867 def setValidityEnd(
868 self, objects: list[DiaObjectId], validityEnd: astropy.time.Time, raise_on_missing_id: bool = False
869 ) -> int:
870 # docstring is inherited from a base class
871 if not objects:
872 return 0
874 context = self._context
875 config = context.config
877 pad_arcsec = 1.0
878 partitioned_object_ids = self._group_dia_objects_by_partition(
879 context.partitioner, objects, pad_arcsec
880 )
882 # Check that all objects exist.
883 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
884 statements: list[tuple] = []
885 for apdb_part, diaObjectIds in partitioned_object_ids.items():
886 id_str = ",".join(str(diaObjectId) for diaObjectId in diaObjectIds)
887 query = (
888 f'SELECT apdb_part, "diaObjectId" FROM "{self._keyspace}"."{table_name}" '
889 f'WHERE apdb_part = %s AND "diaObjectId" IN ({id_str})'
890 )
891 statement = cassandra.query.SimpleStatement(query)
892 statements.append((statement, (apdb_part,)))
894 with self._timer("select_time", tags={"table": table_name, "method": "setValidityEnd"}) as timer:
895 records = cast(
896 list[tuple[int, int]],
897 select_concurrent(
898 context.session,
899 statements,
900 "read_tuples",
901 config.connection_config.read_concurrency,
902 ),
903 )
904 timer.add_values(row_count=len(objects))
906 requested_ids = {obj.diaObjectId for obj in objects}
907 found_ids = {rec[1] for rec in records}
908 if extra_ids := (found_ids - requested_ids):
909 raise RuntimeError(f"Consistency error - found duplicate records for object IDs: {extra_ids}")
910 if raise_on_missing_id:
911 if missing_ids := (requested_ids - found_ids):
912 raise LookupError(f"Some object IDs are missing from DiaObjectLast table: {missing_ids}")
914 # Filter existing records.
915 if len(objects) != len(found_ids):
916 objects = [obj for obj in objects if obj.diaObjectId in found_ids]
918 if not objects:
919 return 0
921 # Group by partitions again.
922 grouped_object_ids: dict[int, list[int]] = defaultdict(list)
923 for apdb_part, diaObjectId in records:
924 grouped_object_ids[apdb_part].append(diaObjectId)
926 # Remove all matching rows from DiaObjectLast.
927 statements = []
928 for apdb_part, diaObjectIds in grouped_object_ids.items():
929 id_str = ",".join(str(diaObjectId) for diaObjectId in diaObjectIds)
930 query = (
931 f'DELETE FROM "{self._keyspace}"."{table_name}" '
932 f'WHERE apdb_part = %s AND "diaObjectId" IN ({id_str})'
933 )
934 statement = cassandra.query.SimpleStatement(query)
935 statements.append((statement, (apdb_part,)))
937 # Also remove from DiaObjectLastToPartition.
938 reverse_table_name = context.schema.tableName(ExtraTables.DiaObjectLastToPartition)
939 id_str = ",".join(str(diaObjectId) for _, diaObjectId in records)
940 query = f'DELETE FROM "{self._keyspace}"."{reverse_table_name}" WHERE "diaObjectId" IN ({id_str})'
941 statement = cassandra.query.SimpleStatement(query)
942 statements.append((statement, ()))
944 with self._timer("delete_time", tags={"table": table_name, "method": "setValidityEnd"}) as timer:
945 execute_concurrent(context.session, statements, execution_profile="write")
946 timer.add_values(row_count=len(records))
948 # If repication is enabled then send all updates.
949 if context.schema.replication_enabled:
950 current_time = self._current_time()
951 current_time_ns = int(current_time.unix_tai * 1e9)
952 replica_chunk = ReplicaChunk.make_replica_chunk(current_time, config.replica_chunk_seconds)
954 update_records = [
955 ApdbCloseDiaObjectValidityRecord(
956 diaObjectId=obj.diaObjectId,
957 ra=obj.ra,
958 dec=obj.dec,
959 update_time_ns=current_time_ns,
960 update_order=index,
961 validityEndMjdTai=float(validityEnd.tai.mjd),
962 nDiaSources=None,
963 )
964 for index, obj in enumerate(objects)
965 ]
967 self._storeUpdateRecords(update_records, replica_chunk, store_chunk=True)
969 return len(objects)
971 def resetDedup(self, dedup_time: astropy.time.Time | None = None) -> None:
972 # docstring is inherited from a base class
973 context = self._context
975 if not context.has_dedup_table:
976 raise TypeError("DiaObjectDedup table does not exist in this APDB instance.")
978 if dedup_time is None:
979 dedup_time = self._current_time()
981 validity_start_column = self._timestamp_column_name("validityStart")
983 # Find latest timestamp in deduplication table.
984 table_name = context.schema.tableName(ExtraTables.DiaObjectDedup)
985 query = f'SELECT MAX("{validity_start_column}") FROM "{self._keyspace}"."{table_name}"'
986 result = context.session.execute(query, execution_profile="read_tuples")
987 max_value = result.one()[0]
988 if self._schema.has_mjd_timestamps:
989 max_validity_start = astropy.time.Time(max_value, format="mjd", scale="tai")
990 else:
991 max_validity_start = astropy.time.Time(max_value, format="datetime", scale="tai")
993 # If max time is lower than dedup time we can do TRUNCATE.
994 if dedup_time >= max_validity_start:
995 query = f'TRUNCATE TABLE "{self._keyspace}"."{table_name}"'
996 context.session.execute(query, execution_profile="write")
997 else:
998 dedup_time_value = self._timestamp_column_value(dedup_time)
999 query = f'DELETE FROM "{self._keyspace}"."{table_name}" WHERE "{validity_start_column}" < %s'
1000 context.session.execute(query, (dedup_time_value,), execution_profile="write")
1002 # Store dedup time.
1003 data = {"dedup_time_iso_tai": dedup_time.tai.to_value("iso")}
1004 data_json = json.dumps(data)
1005 context.metadata.set(context.metadataDedupKey, data_json, force=True)
1007 def reassignDiaSources(self, idMap: Mapping[int, int]) -> None:
1008 # docstring is inherited from a base class
1009 context = self._context
1010 config = context.config
1012 now = self._current_time()
1013 reassign_time_column = self._timestamp_column_name("ssObjectReassocTime")
1014 reassignTime = self._timestamp_column_value(now)
1016 # To update a record we need to know its exact primary key (including
1017 # partition key) so we start by querying for diaSourceId to find the
1018 # primary keys.
1020 table_name = context.schema.tableName(ExtraTables.DiaSourceToPartition)
1021 # split it into 1k IDs per query
1022 selects: list[tuple] = []
1023 for ids in chunk_iterable(idMap.keys(), 1_000):
1024 ids_str = ",".join(str(item) for item in ids)
1025 selects.append(
1026 (
1027 (
1028 'SELECT "diaSourceId", "apdb_part", "apdb_time_part", "apdb_replica_chunk" '
1029 f'FROM "{self._keyspace}"."{table_name}" WHERE "diaSourceId" IN ({ids_str})'
1030 ),
1031 {},
1032 )
1033 )
1035 # No need for DataFrame here, read data as tuples.
1036 result = cast(
1037 list[tuple[int, int, int, int | None]],
1038 select_concurrent(
1039 context.session, selects, "read_tuples", config.connection_config.read_concurrency
1040 ),
1041 )
1043 # Make mapping from source ID to its partition.
1044 id2partitions: dict[int, tuple[int, int]] = {}
1045 id2chunk_id: dict[int, int] = {}
1046 for row in result:
1047 id2partitions[row[0]] = row[1:3]
1048 if row[3] is not None:
1049 id2chunk_id[row[0]] = row[3]
1051 # make sure we know partitions for each ID
1052 if set(id2partitions) != set(idMap):
1053 missing = ",".join(str(item) for item in set(idMap) - set(id2partitions))
1054 raise ValueError(f"Following DiaSource IDs do not exist in the database: {missing}")
1056 # Reassign in standard tables
1057 queries: list[tuple[cassandra.query.PreparedStatement, tuple]] = []
1058 for diaSourceId, ssObjectId in idMap.items():
1059 apdb_part, apdb_time_part = id2partitions[diaSourceId]
1060 values: tuple
1061 if config.partitioning.time_partition_tables:
1062 table_name = context.schema.tableName(ApdbTables.DiaSource, apdb_time_part)
1063 query = (
1064 f'UPDATE "{self._keyspace}"."{table_name}"'
1065 f' SET "ssObjectId" = ?, "diaObjectId" = NULL, "{reassign_time_column}" = ?'
1066 ' WHERE "apdb_part" = ? AND "diaSourceId" = ?'
1067 )
1068 values = (ssObjectId, reassignTime, apdb_part, diaSourceId)
1069 else:
1070 table_name = context.schema.tableName(ApdbTables.DiaSource)
1071 query = (
1072 f'UPDATE "{self._keyspace}"."{table_name}"'
1073 f' SET "ssObjectId" = ?, "diaObjectId" = NULL, "{reassign_time_column}" = ?'
1074 ' WHERE "apdb_part" = ? AND "apdb_time_part" = ? AND "diaSourceId" = ?'
1075 )
1076 values = (ssObjectId, reassignTime, apdb_part, apdb_time_part, diaSourceId)
1077 queries.append((context.preparer.prepare(query), values))
1079 # TODO: (DM-50190) Replication for updated records is not implemented.
1080 if id2chunk_id:
1081 warnings.warn("Replication of reassigned DiaSource records is not implemented.", stacklevel=2)
1083 _LOG.debug("%s: will update %d records", table_name, len(idMap))
1084 with self._timer("source_reassign_time") as timer:
1085 execute_concurrent(context.session, queries, execution_profile="write")
1086 timer.add_values(source_count=len(idMap))
1088 def countUnassociatedObjects(self) -> int:
1089 # docstring is inherited from a base class
1091 # It's too inefficient to implement it for Cassandra in current schema.
1092 raise NotImplementedError()
1094 @property
1095 def schema(self) -> ApdbSchema:
1096 # docstring is inherited from a base class
1097 return self._schema
1099 @property
1100 def metadata(self) -> ApdbMetadata:
1101 # docstring is inherited from a base class
1102 context = self._context
1103 return context.metadata
1105 @property
1106 def admin(self) -> ApdbCassandraAdmin:
1107 # docstring is inherited from a base class
1108 return ApdbCassandraAdmin(self)
1110 def _getSources(
1111 self,
1112 region: sphgeom.Region,
1113 object_ids: Iterable[int] | None,
1114 mjd_start: float,
1115 mjd_end: float,
1116 table_name: ApdbTables,
1117 ) -> pandas.DataFrame:
1118 """Return catalog of DiaSource instances given set of DiaObject IDs.
1120 Parameters
1121 ----------
1122 region : `lsst.sphgeom.Region`
1123 Spherical region.
1124 object_ids :
1125 Collection of DiaObject IDs
1126 mjd_start : `float`
1127 Lower bound of time interval.
1128 mjd_end : `float`
1129 Upper bound of time interval.
1130 table_name : `ApdbTables`
1131 Name of the table.
1133 Returns
1134 -------
1135 catalog : `pandas.DataFrame`, or `None`
1136 Catalog containing DiaSource records. Empty catalog is returned if
1137 ``object_ids`` is empty.
1138 """
1139 context = self._context
1140 config = context.config
1142 object_id_set: Set[int] = set()
1143 if object_ids is not None:
1144 object_id_set = set(object_ids)
1145 if len(object_id_set) == 0:
1146 return self._make_empty_catalog(table_name)
1148 sp_where, num_sp_part = context.partitioner.spatial_where(region, for_prepare=True)
1149 tables, temporal_where = context.partitioner.temporal_where(
1150 table_name, mjd_start, mjd_end, for_prepare=True, partitons_range=context.time_partitions_range
1151 )
1152 if not tables:
1153 start = astropy.time.Time(mjd_start, format="mjd", scale="tai")
1154 end = astropy.time.Time(mjd_end, format="mjd", scale="tai")
1155 warnings.warn(
1156 f"Query time range ({start.isot} - {end.isot}) does not overlap database time partitions."
1157 )
1159 # We need to exclude extra partitioning columns from result.
1160 column_names = context.schema.apdbColumnNames(table_name)
1161 what = ",".join(quote_id(column) for column in column_names)
1163 # Build all queries
1164 statements: list[tuple] = []
1165 for table in tables:
1166 prefix = f'SELECT {what} from "{self._keyspace}"."{table}"'
1167 statements += list(self._combine_where(prefix, sp_where, temporal_where))
1168 _LOG.debug("_getSources %s: #queries: %s", table_name, len(statements))
1170 with self._timer("select_time", tags={"table": table_name.name, "method": "_getSources"}) as timer:
1171 table_data_raw = cast(
1172 ApdbCassandraTableData,
1173 select_concurrent(
1174 context.session,
1175 statements,
1176 "read_raw_multi",
1177 config.connection_config.read_concurrency,
1178 ),
1179 )
1180 catalog = table_data_raw.to_pandas(context.schema._table_schema(table_name))
1181 timer.add_values(
1182 row_count_from_db=len(catalog), num_sp_part=num_sp_part, num_queries=len(statements)
1183 )
1185 # filter by given object IDs
1186 if len(object_id_set) > 0:
1187 catalog = cast(pandas.DataFrame, catalog[catalog["diaObjectId"].isin(object_id_set)])
1189 # precise filtering on midpointMjdTai
1190 catalog = cast(pandas.DataFrame, catalog[catalog["midpointMjdTai"] > mjd_start])
1192 timer.add_values(row_count=len(catalog))
1194 _LOG.debug("found %d %ss", catalog.shape[0], table_name.name)
1195 return catalog
1197 def _storeReplicaChunk(self, replica_chunk: ReplicaChunk) -> None:
1198 context = self._context
1199 config = context.config
1201 # Cassandra timestamp uses milliseconds since epoch
1202 timestamp = int(replica_chunk.last_update_time.unix_tai * 1000)
1204 # everything goes into a single partition
1205 partition = 0
1207 table_name = context.schema.tableName(ExtraTables.ApdbReplicaChunks)
1209 columns = ["partition", "apdb_replica_chunk", "last_update_time", "unique_id"]
1210 values = [partition, replica_chunk.id, timestamp, replica_chunk.unique_id]
1211 if context.has_chunk_sub_partitions:
1212 columns.append("has_subchunks")
1213 values.append(True)
1215 column_list = ", ".join(columns)
1216 placeholders = ",".join(["%s"] * len(columns))
1217 query = f'INSERT INTO "{self._keyspace}"."{table_name}" ({column_list}) VALUES ({placeholders})'
1219 context.session.execute(
1220 query,
1221 values,
1222 timeout=config.connection_config.write_timeout,
1223 execution_profile="write",
1224 )
1226 def _queryDiaObjectLastPartitions(self, ids: Iterable[int]) -> Mapping[int, int]:
1227 """Return existing mapping of diaObjectId to its last partition."""
1228 context = self._context
1229 config = context.config
1231 table_name = context.schema.tableName(ExtraTables.DiaObjectLastToPartition)
1232 queries = []
1233 object_count = 0
1234 for id_chunk in chunk_iterable(ids, 10_000):
1235 id_chunk_list = list(id_chunk)
1236 query = (
1237 f'SELECT "diaObjectId", apdb_part FROM "{self._keyspace}"."{table_name}" '
1238 f'WHERE "diaObjectId" in ({",".join(str(oid) for oid in id_chunk_list)})'
1239 )
1240 queries.append((query, ()))
1241 object_count += len(id_chunk_list)
1243 with self._timer("query_object_last_partitions", tags={"table": table_name}) as timer:
1244 data = cast(
1245 ApdbTableData,
1246 select_concurrent(
1247 context.session,
1248 queries,
1249 "read_raw_multi",
1250 config.connection_config.read_concurrency,
1251 ),
1252 )
1253 timer.add_values(object_count=object_count, row_count=len(data.rows()))
1255 if data.column_names() != ["diaObjectId", "apdb_part"]:
1256 raise RuntimeError(f"Unexpected column names in query result: {data.column_names()}")
1258 return {row[0]: row[1] for row in data.rows()}
1260 def _deleteMovingObjects(self, objs: pandas.DataFrame) -> None:
1261 """Objects in DiaObjectsLast can move from one spatial partition to
1262 another. For those objects inserting new version does not replace old
1263 one, so we need to explicitly remove old versions before inserting new
1264 ones.
1265 """
1266 context = self._context
1268 # Extract all object IDs.
1269 new_partitions = dict(zip(objs["diaObjectId"], objs["apdb_part"]))
1270 old_partitions = self._queryDiaObjectLastPartitions(objs["diaObjectId"])
1272 moved_oids: dict[int, tuple[int, int]] = {}
1273 for oid, old_part in old_partitions.items():
1274 new_part = new_partitions.get(oid, old_part)
1275 if new_part != old_part:
1276 moved_oids[oid] = (old_part, new_part)
1277 _LOG.debug("DiaObject IDs that moved to new partition: %s", moved_oids)
1279 if moved_oids:
1280 # Delete old records from DiaObjectLast.
1281 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
1282 query = f'DELETE FROM "{self._keyspace}"."{table_name}" WHERE apdb_part = ? AND "diaObjectId" = ?'
1283 statement = context.preparer.prepare(query)
1284 queries = []
1285 for oid, (old_part, _) in moved_oids.items():
1286 queries.append((statement, (old_part, oid)))
1287 with self._timer("delete_object_last", tags={"table": table_name}) as timer:
1288 execute_concurrent(context.session, queries, execution_profile="write")
1289 timer.add_values(row_count=len(moved_oids))
1291 # Add all new records to the map.
1292 table_name = context.schema.tableName(ExtraTables.DiaObjectLastToPartition)
1293 query = f'INSERT INTO "{self._keyspace}"."{table_name}" ("diaObjectId", apdb_part) VALUES (?,?)'
1294 statement = context.preparer.prepare(query)
1296 queries = []
1297 for oid, new_part in new_partitions.items():
1298 queries.append((statement, (oid, new_part)))
1300 with self._timer("update_object_last_partition", tags={"table": table_name}) as timer:
1301 execute_concurrent(context.session, queries, execution_profile="write")
1302 timer.add_values(row_count=len(queries))
1304 def _storeDiaObjects(
1305 self, objs: pandas.DataFrame, visit_time: astropy.time.Time, replica_chunk: ReplicaChunk | None
1306 ) -> None:
1307 """Store catalog of DiaObjects from current visit.
1309 Parameters
1310 ----------
1311 objs : `pandas.DataFrame`
1312 Catalog with DiaObject records
1313 visit_time : `astropy.time.Time`
1314 Time of the current visit.
1315 replica_chunk : `ReplicaChunk` or `None`
1316 Replica chunk identifier if replication is configured.
1317 """
1318 if len(objs) == 0:
1319 _LOG.debug("No objects to write to database.")
1320 return
1322 context = self._context
1323 config = context.config
1325 if context.has_dia_object_last_to_partition:
1326 self._deleteMovingObjects(objs)
1328 validity_start_column = self._timestamp_column_name("validityStart")
1329 timestamp = self._timestamp_column_value(visit_time)
1331 # DiaObjectLast did not have this column in the past.
1332 extra_columns: dict[str, Any] = {}
1333 if context.schema.check_column(ApdbTables.DiaObjectLast, validity_start_column):
1334 extra_columns[validity_start_column] = timestamp
1336 self._storeObjectsPandas(objs, ApdbTables.DiaObjectLast, extra_columns=extra_columns)
1338 extra_columns[validity_start_column] = timestamp
1339 visit_time_part = context.partitioner.time_partition(visit_time)
1340 time_part: int | None = visit_time_part
1341 if (time_partitions_range := context.time_partitions_range) is not None:
1342 self._check_time_partitions([visit_time_part], time_partitions_range)
1343 if not config.partitioning.time_partition_tables:
1344 extra_columns["apdb_time_part"] = time_part
1345 time_part = None
1347 # Only store DiaObects if not doing replication or explicitly
1348 # configured to always store them.
1349 if replica_chunk is None or not config.replica_skips_diaobjects:
1350 self._storeObjectsPandas(
1351 objs, ApdbTables.DiaObject, extra_columns=extra_columns, time_part=time_part
1352 )
1354 if replica_chunk is not None:
1355 extra_columns = {"apdb_replica_chunk": replica_chunk.id, validity_start_column: timestamp}
1356 table = ExtraTables.DiaObjectChunks
1357 if context.has_chunk_sub_partitions:
1358 table = ExtraTables.DiaObjectChunks2
1359 # Use a random number for a second part of partitioning key so
1360 # that different clients could wrtite to different partitions.
1361 # This makes it not exactly reproducible.
1362 extra_columns["apdb_replica_subchunk"] = random.randrange(config.replica_sub_chunk_count)
1363 self._storeObjectsPandas(objs, table, extra_columns=extra_columns)
1365 # Store copy of the records in dedup table.
1366 if context.has_dedup_table:
1367 table = ExtraTables.DiaObjectDedup
1368 extra_columns = {
1369 "dedup_part": random.randrange(config.partitioning.num_part_dedup),
1370 validity_start_column: timestamp,
1371 }
1372 self._storeObjectsPandas(objs, table, extra_columns=extra_columns)
1374 def _storeDiaSources(
1375 self,
1376 table_name: ApdbTables,
1377 sources: pandas.DataFrame,
1378 replica_chunk: ReplicaChunk | None,
1379 ) -> int | None:
1380 """Store catalog of DIASources or DIAForcedSources from current visit.
1382 Parameters
1383 ----------
1384 table_name : `ApdbTables`
1385 Table where to store the data.
1386 sources : `pandas.DataFrame`
1387 Catalog containing DiaSource records
1388 visit_time : `astropy.time.Time`
1389 Time of the current visit.
1390 replica_chunk : `ReplicaChunk` or `None`
1391 Replica chunk identifier if replication is configured.
1393 Returns
1394 -------
1395 subchunk : `int` or `None`
1396 Subchunk number for resulting replica data, `None` if relication is
1397 not enabled ot subchunking is not enabled.
1398 """
1399 context = self._context
1400 config = context.config
1402 # Time partitioning has to be based on midpointMjdTai, not visit_time
1403 # as visit_time is not really a visit time.
1404 tp_sources = sources.copy(deep=False)
1405 tp_sources["apdb_time_part"] = tp_sources["midpointMjdTai"].apply(context.partitioner.time_partition)
1406 if (time_partitions_range := context.time_partitions_range) is not None:
1407 self._check_time_partitions(tp_sources["apdb_time_part"], time_partitions_range)
1408 extra_columns: dict[str, Any] = {}
1409 if not config.partitioning.time_partition_tables:
1410 self._storeObjectsPandas(tp_sources, table_name)
1411 else:
1412 # Group by time partition
1413 partitions = set(tp_sources["apdb_time_part"])
1414 if len(partitions) == 1:
1415 # Single partition - just save the whole thing.
1416 time_part = partitions.pop()
1417 self._storeObjectsPandas(sources, table_name, time_part=time_part)
1418 else:
1419 # group by time partition.
1420 for time_part, sub_frame in tp_sources.groupby(by="apdb_time_part"):
1421 sub_frame.drop(columns="apdb_time_part", inplace=True)
1422 self._storeObjectsPandas(sub_frame, table_name, time_part=time_part)
1424 subchunk: int | None = None
1425 if replica_chunk is not None:
1426 extra_columns = {"apdb_replica_chunk": replica_chunk.id}
1427 if context.has_chunk_sub_partitions:
1428 subchunk = random.randrange(config.replica_sub_chunk_count)
1429 extra_columns["apdb_replica_subchunk"] = subchunk
1430 if table_name is ApdbTables.DiaSource:
1431 extra_table = ExtraTables.DiaSourceChunks2
1432 else:
1433 extra_table = ExtraTables.DiaForcedSourceChunks2
1434 else:
1435 if table_name is ApdbTables.DiaSource:
1436 extra_table = ExtraTables.DiaSourceChunks
1437 else:
1438 extra_table = ExtraTables.DiaForcedSourceChunks
1439 self._storeObjectsPandas(sources, extra_table, extra_columns=extra_columns)
1441 return subchunk
1443 def _check_time_partitions(
1444 self, partitions: Iterable[int], time_partitions_range: ApdbCassandraTimePartitionRange
1445 ) -> None:
1446 """Check that time partitons for new data actually exist.
1448 Parameters
1449 ----------
1450 partitions : `~collections.abc.Iterable` [`int`]
1451 Time partitions for new data.
1452 time_partitions_range : `ApdbCassandraTimePartitionRange`
1453 Currrent time partition range.
1454 """
1455 partitions = set(partitions)
1456 min_part = min(partitions)
1457 max_part = max(partitions)
1458 if min_part < time_partitions_range.start or max_part > time_partitions_range.end:
1459 raise ValueError(
1460 "Attempt to store data for time partitions that do not yet exist. "
1461 f"Partitons for new records: {min_part}-{max_part}. "
1462 f"Database partitons: {time_partitions_range.start}-{time_partitions_range.end}."
1463 )
1464 # Make a noise when writing to the last partition.
1465 if max_part == time_partitions_range.end:
1466 warnings.warn(
1467 "Writing into the last temporal partition. Partition range needs to be extended soon.",
1468 stacklevel=3,
1469 )
1471 def _storeDiaSourcesPartitions(
1472 self,
1473 sources: pandas.DataFrame,
1474 visit_time: astropy.time.Time,
1475 replica_chunk: ReplicaChunk | None,
1476 subchunk: int | None,
1477 ) -> None:
1478 """Store mapping of diaSourceId to its partitioning values.
1480 Parameters
1481 ----------
1482 sources : `pandas.DataFrame`
1483 Catalog containing DiaSource records
1484 visit_time : `astropy.time.Time`
1485 Time of the current visit.
1486 replica_chunk : `ReplicaChunk` or `None`
1487 Replication chunk, or `None` when replication is disabled.
1488 subchunk : `int` or `None`
1489 Replication sub-chunk, or `None` when replication is disabled or
1490 sub-chunking is not used.
1491 """
1492 context = self._context
1494 id_map = cast(pandas.DataFrame, sources[["diaSourceId", "apdb_part"]])
1495 extra_columns = {
1496 "apdb_time_part": context.partitioner.time_partition(visit_time),
1497 "apdb_replica_chunk": replica_chunk.id if replica_chunk is not None else None,
1498 }
1499 if context.has_chunk_sub_partitions:
1500 extra_columns["apdb_replica_subchunk"] = subchunk
1502 self._storeObjectsPandas(
1503 id_map, ExtraTables.DiaSourceToPartition, extra_columns=extra_columns, time_part=None
1504 )
1506 def _storeObjectsPandas(
1507 self,
1508 records: pandas.DataFrame,
1509 table_name: ApdbTables | ExtraTables,
1510 extra_columns: Mapping | None = None,
1511 time_part: int | None = None,
1512 ) -> None:
1513 """Store generic objects.
1515 Takes Pandas catalog and stores a bunch of records in a table.
1517 Parameters
1518 ----------
1519 records : `pandas.DataFrame`
1520 Catalog containing object records
1521 table_name : `ApdbTables`
1522 Name of the table as defined in APDB schema.
1523 extra_columns : `dict`, optional
1524 Mapping (column_name, column_value) which gives fixed values for
1525 columns in each row, overrides values in ``records`` if matching
1526 columns exist there.
1527 time_part : `int`, optional
1528 If not `None` then insert into a per-partition table.
1530 Notes
1531 -----
1532 If Pandas catalog contains additional columns not defined in table
1533 schema they are ignored. Catalog does not have to contain all columns
1534 defined in a table, but partition and clustering keys must be present
1535 in a catalog or ``extra_columns``.
1536 """
1537 context = self._context
1539 # use extra columns if specified
1540 if extra_columns is None:
1541 extra_columns = {}
1542 extra_fields = list(extra_columns.keys())
1544 # Fields that will come from dataframe.
1545 df_fields = [column for column in records.columns if column not in extra_fields]
1547 column_map = context.schema.getColumnMap(table_name)
1548 # list of columns (as in felis schema)
1549 fields = [column_map[field].name for field in df_fields if field in column_map]
1550 fields += extra_fields
1552 # check that all partitioning and clustering columns are defined
1553 partition_columns = context.schema.partitionColumns(table_name)
1554 required_columns = partition_columns + context.schema.clusteringColumns(table_name)
1555 missing_columns = [column for column in required_columns if column not in fields]
1556 if missing_columns:
1557 raise ValueError(f"Primary key columns are missing from catalog: {missing_columns}")
1559 qfields = [quote_id(field) for field in fields]
1560 qfields_str = ",".join(qfields)
1562 batch_size = self._batch_size(table_name)
1564 with self._timer("insert_build_time", tags={"table": table_name.name}):
1565 # Multi-partition batches are problematic in general, so we want to
1566 # group records in a batch by their partition key.
1567 values_by_key: dict[tuple, list[list]] = defaultdict(list)
1568 for rec in records.itertuples(index=False):
1569 values = []
1570 partitioning_values: dict[str, Any] = {}
1571 for field in df_fields:
1572 if field not in column_map:
1573 continue
1574 value = getattr(rec, field)
1575 if column_map[field].datatype is felis.datamodel.DataType.timestamp:
1576 if isinstance(value, pandas.Timestamp):
1577 value = value.to_pydatetime()
1578 elif value is pandas.NaT:
1579 value = None
1580 else:
1581 # Assume it's seconds since epoch, Cassandra
1582 # datetime is in milliseconds
1583 value = int(value * 1000)
1584 value = literal(value)
1585 values.append(UNSET_VALUE if value is None else value)
1586 if field in partition_columns:
1587 partitioning_values[field] = value
1588 for field in extra_fields:
1589 value = literal(extra_columns[field])
1590 values.append(UNSET_VALUE if value is None else value)
1591 if field in partition_columns:
1592 partitioning_values[field] = value
1594 key = tuple(partitioning_values[field] for field in partition_columns)
1595 values_by_key[key].append(values)
1597 table = context.schema.tableName(table_name, time_part)
1599 holders = ",".join(["?"] * len(qfields))
1600 query = f'INSERT INTO "{self._keyspace}"."{table}" ({qfields_str}) VALUES ({holders})'
1601 statement = context.preparer.prepare(query)
1602 # Cassandra has 64k limit on batch size, normally that should be
1603 # enough but some tests generate too many forced sources.
1604 queries = []
1605 for key_values in values_by_key.values():
1606 for values_chunk in chunk_iterable(key_values, batch_size):
1607 batch = cassandra.query.BatchStatement()
1608 for row_values in values_chunk:
1609 batch.add(statement, row_values)
1610 queries.append((batch, None))
1611 assert batch.routing_key is not None and batch.keyspace is not None
1613 _LOG.debug("%s: will store %d records", context.schema.tableName(table_name), records.shape[0])
1614 with self._timer(
1615 "insert_time", tags={"table": table_name.name, "method": "_storeObjectsPandas"}
1616 ) as timer:
1617 execute_concurrent(context.session, queries, execution_profile="write")
1618 timer.add_values(row_count=len(records), num_batches=len(queries))
1620 def _storeUpdateRecords(
1621 self, records: Iterable[ApdbUpdateRecord], chunk: ReplicaChunk, *, store_chunk: bool = False
1622 ) -> None:
1623 """Store ApdbUpdateRecords in the replica table for those records.
1625 Parameters
1626 ----------
1627 records : `list` [`ApdbUpdateRecord`]
1628 Records to store.
1629 chunk : `ReplicaChunk`
1630 Replica chunk for these records.
1631 store_chunk : `bool`
1632 If True then also store replica chunk.
1634 Raises
1635 ------
1636 TypeError
1637 Raised if replication is not enabled for this instance.
1638 """
1639 context = self._context
1640 config = context.config
1642 if not context.schema.replication_enabled:
1643 raise TypeError("Replication is not enabled for this APDB instance.")
1645 if store_chunk:
1646 self._storeReplicaChunk(chunk)
1648 apdb_replica_chunk = chunk.id
1649 # Do not use unique_if from ReplicaChunk as it could be reused in
1650 # multiple calls to this method.
1651 update_unique_id = uuid.uuid4()
1653 rows = []
1654 for record in records:
1655 rows.append(
1656 [
1657 apdb_replica_chunk,
1658 record.update_time_ns,
1659 record.update_order,
1660 update_unique_id,
1661 record.to_json(),
1662 ]
1663 )
1664 columns = [
1665 "apdb_replica_chunk",
1666 "update_time_ns",
1667 "update_order",
1668 "update_unique_id",
1669 "update_payload",
1670 ]
1671 if context.has_chunk_sub_partitions:
1672 subchunk = random.randrange(config.replica_sub_chunk_count)
1673 for row in rows:
1674 row.append(subchunk)
1675 columns.append("apdb_replica_subchunk")
1677 table_name = context.schema.tableName(ExtraTables.ApdbUpdateRecordChunks)
1678 placeholders = ", ".join(["%s"] * len(columns))
1679 columns_str = ", ".join(columns)
1680 query = f'INSERT INTO "{self._keyspace}"."{table_name}" ({columns_str}) VALUES ({placeholders})'
1681 queries = [(query, row) for row in rows]
1683 with self._timer("store_update_record", tags={"table": table_name}) as timer:
1684 execute_concurrent(context.session, queries, execution_profile="write")
1685 timer.add_values(row_count=len(queries))
1687 def _add_apdb_part(self, df: pandas.DataFrame) -> pandas.DataFrame:
1688 """Calculate spatial partition for each record and add it to a
1689 DataFrame.
1691 Parameters
1692 ----------
1693 df : `pandas.DataFrame`
1694 DataFrame which has to contain ra/dec columns, names of these
1695 columns are defined by configuration ``ra_dec_columns`` field.
1697 Returns
1698 -------
1699 df : `pandas.DataFrame`
1700 DataFrame with ``apdb_part`` column which contains pixel index
1701 for ra/dec coordinates.
1703 Notes
1704 -----
1705 This overrides any existing column in a DataFrame with the same name
1706 (``apdb_part``). Original DataFrame is not changed, copy of a DataFrame
1707 is returned.
1708 """
1709 context = self._context
1710 config = context.config
1712 # Calculate pixelization index for every record.
1713 apdb_part = np.zeros(df.shape[0], dtype=np.int64)
1714 ra_col, dec_col = config.ra_dec_columns
1715 for i, (ra, dec) in enumerate(zip(df[ra_col], df[dec_col])):
1716 uv3d = sphgeom.UnitVector3d(sphgeom.LonLat.fromDegrees(ra, dec))
1717 idx = context.partitioner.pixel(uv3d)
1718 apdb_part[i] = idx
1719 df = df.copy()
1720 df["apdb_part"] = apdb_part
1721 return df
1723 def _make_empty_catalog(self, table_name: ApdbTables) -> pandas.DataFrame:
1724 """Make an empty catalog for a table with a given name.
1726 Parameters
1727 ----------
1728 table_name : `ApdbTables`
1729 Name of the table.
1731 Returns
1732 -------
1733 catalog : `pandas.DataFrame`
1734 An empty catalog.
1735 """
1736 table = self.schema.tableSchemas[table_name]
1738 data = {columnDef.name: pandas.Series(dtype=columnDef.pandas_type) for columnDef in table.columns}
1739 return pandas.DataFrame(data)
1741 def _combine_where(
1742 self,
1743 prefix: str,
1744 where1: list[tuple[str, tuple]],
1745 where2: list[tuple[str, tuple]],
1746 where3: tuple[str, tuple] | None = None,
1747 suffix: str | None = None,
1748 ) -> Iterator[tuple[cassandra.query.Statement, tuple]]:
1749 """Make cartesian product of two parts of WHERE clause into a series
1750 of statements to execute.
1752 Parameters
1753 ----------
1754 prefix : `str`
1755 Initial statement prefix that comes before WHERE clause, e.g.
1756 "SELECT * from Table"
1757 """
1758 context = self._context
1760 # If lists are empty use special sentinels.
1761 if not where1:
1762 where1 = [("", ())]
1763 if not where2:
1764 where2 = [("", ())]
1766 for expr1, params1 in where1:
1767 for expr2, params2 in where2:
1768 full_query = prefix
1769 wheres = []
1770 params = params1 + params2
1771 if expr1:
1772 wheres.append(expr1)
1773 if expr2:
1774 wheres.append(expr2)
1775 if where3:
1776 wheres.append(where3[0])
1777 params += where3[1]
1778 if wheres:
1779 full_query += " WHERE " + " AND ".join(wheres)
1780 if suffix:
1781 full_query += " " + suffix
1782 if params:
1783 statement = context.preparer.prepare(full_query)
1784 else:
1785 # If there are no params then it is likely that query
1786 # has a bunch of literals rendered already, no point
1787 # trying to prepare it.
1788 statement = cassandra.query.SimpleStatement(full_query)
1789 yield (statement, params)
1791 def _fix_input_timestamps(self, df: pandas.DataFrame) -> pandas.DataFrame:
1792 """Update timestamp columns in input DataFrame to be naive datetime
1793 type.
1795 Clients may or may not generate aware timestamps, code in this class
1796 assumes that timestamps are naive, so we convert them to UTC and
1797 drop timezone.
1798 """
1799 # Find all columns with aware timestamps.
1800 columns = [column for column, dtype in df.dtypes.items() if isinstance(dtype, pandas.DatetimeTZDtype)]
1801 for column in columns:
1802 # tz_convert(None) will convert to UTC and drop timezone.
1803 df[column] = df[column].dt.tz_convert(None)
1804 return df
1806 def _batch_size(self, table: ApdbTables | ExtraTables) -> int:
1807 """Calculate batch size based on config parameters."""
1808 context = self._context
1809 config = context.config
1811 # Cassandra limit on number of statements in a batch is 64k.
1812 batch_size = 65_535
1813 if 0 < config.batch_statement_limit < batch_size:
1814 batch_size = config.batch_statement_limit
1815 if config.batch_size_limit > 0:
1816 # The purpose of this limit is to try not to exceed batch size
1817 # threshold which is set on server side. Cassandra wire protocol
1818 # for prepared queries (and batches) only sends column values with
1819 # with an additional 4 bytes per value specifying size. Value is
1820 # not included for NULL or NOT_SET values, but the size is always
1821 # there. There is additional small per-query overhead, which we
1822 # ignore.
1823 row_size = context.schema.table_row_size(table)
1824 row_size += 4 * len(context.schema.getColumnMap(table))
1825 batch_size = min(batch_size, (config.batch_size_limit // row_size) + 1)
1826 return batch_size
1828 def _group_dia_objects_by_partition(
1829 self, partitioner: Partitioner, objects: list[DiaObjectId], pad_arcsec: float
1830 ) -> Mapping[int, list[int]]:
1831 """Group DiaObjects by partition.
1833 Parameters
1834 ----------
1835 partitioner : `Partitioner`
1836 Objects which knows how to partition things.
1837 objects : `list` [`DiaObjectId`]
1838 Collection of objects to partition.
1839 pad_arcsec : `float`
1840 Additional padding around object position.
1842 Returns
1843 -------
1844 grouped_objects
1845 Mapping of spatial patition ID to list ob object IDs that it
1846 contains. Some objects may belong to more than one partition.
1847 """
1848 partitioned_object_ids: dict[int, list[int]] = defaultdict(list)
1849 for obj_id in objects:
1850 partitions = partitioner.pixelization.circle_pixels(obj_id.ra, obj_id.dec, pad_arcsec)
1851 for pixel in partitions:
1852 partitioned_object_ids[pixel].append(obj_id.diaObjectId)
1853 return partitioned_object_ids
1855 def _timestamp_column_name(self, column: str) -> str:
1856 """Return column name before/after schema migration to MJD TAI."""
1857 return self._schema.timestamp_column_name(column)
1859 def _timestamp_column_value(self, time: astropy.time.Time) -> float | int:
1860 """Return column value before/after schema migration to MJD TAI."""
1861 if self._schema.has_mjd_timestamps:
1862 return float(time.tai.mjd)
1863 else:
1864 return int(time.datetime.astimezone(tz=datetime.UTC).timestamp() * 1000)
1866 def _get_diasource_data(self, source_ids: Iterable[DiaSourceId], *columns: str) -> list:
1867 """Select records from DiaSource table by diaSourceId and return all
1868 records as a list of named tuples.
1869 """
1870 context = self._context
1871 config = context.config
1872 partitioner = context.partitioner
1874 columns = ("diaSourceId",) + columns
1875 what = ",".join(quote_id(column) for column in columns)
1877 # Allow some uncertainty for coordinates and time when calculating
1878 # partitions.
1879 statements: list[tuple] = []
1880 pad_arcsec = 1.0
1881 pad_time_day = 10 / (24 * 3600)
1882 for source_id in source_ids:
1883 center = sphgeom.UnitVector3d(sphgeom.LonLat.fromDegrees(source_id.ra, source_id.dec))
1884 region = sphgeom.Circle(center, sphgeom.Angle.fromDegrees(pad_arcsec / 3600.0))
1885 spatial_where, _ = partitioner.spatial_where(region, for_prepare=True)
1887 tables, temporal_where = partitioner.temporal_where(
1888 ApdbTables.DiaSource,
1889 source_id.midpointMjdTai - pad_time_day,
1890 source_id.midpointMjdTai + pad_time_day,
1891 for_prepare=True,
1892 partitons_range=context.time_partitions_range,
1893 query_per_time_part=True,
1894 )
1896 id_where = ('"diaSourceId" = ?', (source_id.diaSourceId,))
1898 for table in tables:
1899 prefix = f'SELECT {what} FROM "{self._keyspace}"."{table}"'
1900 statements += list(self._combine_where(prefix, spatial_where, temporal_where, id_where))
1902 with self._timer(
1903 "select_time", tags={"table": "DiaSource", "method": "_get_diasource_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
1918 def _get_diaobject_data(self, object_ids: Iterable[DiaObjectId], *columns: str) -> list:
1919 """Select records from DiaObjectLast table by diaObjectId and return
1920 all records as a list of named tuples.
1921 """
1922 context = self._context
1923 config = context.config
1924 partitioner = context.partitioner
1926 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
1927 columns = ("diaObjectId",) + columns
1928 what = ",".join(quote_id(column) for column in columns)
1930 # Allow some uncertainty for coordinates when calculating partitions.
1931 pad_arcsec = 1.0
1932 ids_by_partition = defaultdict(list)
1933 for object_id in object_ids:
1934 pixels = partitioner.pixelization.circle_pixels(object_id.ra, object_id.dec, pad_arcsec)
1935 for pixel in pixels:
1936 ids_by_partition[pixel].append(object_id.diaObjectId)
1938 statements: list[tuple] = []
1939 for apdb_part, diaObjectIds in ids_by_partition.items():
1940 ids_str = ",".join(str(diaObjectId) for diaObjectId in diaObjectIds)
1941 statement = cassandra.query.SimpleStatement(
1942 f'SELECT {what} FROM "{self._keyspace}"."{table_name}" '
1943 f'WHERE apdb_part = %s AND "diaObjectId" IN ({ids_str})'
1944 )
1945 statements.append((statement, (apdb_part,)))
1947 with self._timer(
1948 "select_time", tags={"table": "DiaObjectLast", "method": "_get_diaobject_data"}
1949 ) as timer:
1950 result = cast(
1951 list[tuple],
1952 select_concurrent(
1953 context.session,
1954 statements,
1955 "read_named_tuples",
1956 config.connection_config.read_concurrency,
1957 ),
1958 )
1959 timer.add_values(row_count=len(result), num_queries=len(statements))
1961 return result