Coverage for python/lsst/dax/apdb/cassandra/apdbCassandra.py: 8%
926 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-11-08 01:40 -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, 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 ApdbWithdrawDiaForcedSourceRecord,
65 ApdbWithdrawDiaSourceRecord,
66)
67from ..monitor import MonAgent
68from ..recordIds import DiaForcedSourceId, DiaObjectId, DiaSourceId
69from ..schema_model import Table
70from ..timer import Timer
71from ..versionTuple import VersionTuple
72from .apdbCassandraAdmin import ApdbCassandraAdmin
73from .apdbCassandraReplica import ApdbCassandraReplica
74from .apdbCassandraSchema import ApdbCassandraSchema, CreateTableOptions, ExtraTables
75from .apdbMetadataCassandra import ApdbMetadataCassandra
76from .cassandra_utils import (
77 ApdbCassandraTableData,
78 execute_concurrent,
79 literal,
80 select_concurrent,
81)
82from .config import ApdbCassandraConfig, ApdbCassandraConnectionConfig, ApdbCassandraTimePartitionRange
83from .connectionContext import ConnectionContext, DbVersions
84from .exceptions import CassandraMissingError
85from .partitioner import Partitioner
86from .queries import Column as C # noqa: N817
87from .queries import ColumnExpr, Delete, Insert, QExpr, Select, Update
88from .sessionFactory import SessionContext, SessionFactory
90if TYPE_CHECKING:
91 from ..apdbMetadata import ApdbMetadata
92 from ..apdbUpdateRecord import ApdbUpdateRecord
94_LOG = logging.getLogger(__name__)
96_MON = MonAgent(__name__)
98VERSION = VersionTuple(1, 3, 0)
99"""Version for the code controlling non-replication tables. This needs to be
100updated following compatibility rules when schema produced by this code
101changes.
102"""
105class ApdbCassandra(Apdb):
106 """Implementation of APDB database with Apache Cassandra backend.
108 Parameters
109 ----------
110 config : `ApdbCassandraConfig`
111 Configuration object.
112 """
114 def __init__(self, config: ApdbCassandraConfig):
115 if not CASSANDRA_IMPORTED:
116 raise CassandraMissingError()
118 self._config = config
119 self._keyspace = config.keyspace
120 self._schema = ApdbSchema(config.schema_file, config.ss_schema_file)
122 self._session_factory = SessionFactory(config)
123 self._connection_context: ConnectionContext | None = None
125 def close(self) -> None:
126 # docstring is inherited from a base class
127 self._session_factory.close()
129 @property
130 def _context(self) -> ConnectionContext:
131 """Establish connection if not established and return context."""
132 if self._connection_context is None:
133 current_versions = DbVersions(
134 schema_version=self.schema.schemaVersion(),
135 code_version=self.apdbImplementationVersion(),
136 replica_version=ApdbCassandraReplica.apdbReplicaImplementationVersion(),
137 )
138 _LOG.debug("Current versions: %s", current_versions)
140 session = self._session_factory.session()
141 self._connection_context = ConnectionContext(
142 session, self._config, self.schema.tableSchemas, current_versions
143 )
145 if _LOG.isEnabledFor(logging.DEBUG):
146 _LOG.debug("ApdbCassandra Configuration: %s", self._connection_context.config.model_dump())
148 return self._connection_context
150 def _timer(self, name: str, *, tags: Mapping[str, str | int] | None = None) -> Timer:
151 """Create `Timer` instance given its name."""
152 return Timer(name, _MON, tags=tags)
154 @classmethod
155 def apdbImplementationVersion(cls) -> VersionTuple:
156 """Return version number for current APDB implementation.
158 Returns
159 -------
160 version : `VersionTuple`
161 Version of the code defined in implementation class.
162 """
163 return VERSION
165 def getConfig(self) -> ApdbCassandraConfig:
166 # docstring is inherited from a base class
167 return self._context.config
169 def tableDef(self, table: ApdbTables) -> Table | None:
170 # docstring is inherited from a base class
171 return self.schema.tableSchemas.get(table)
173 @classmethod
174 def init_database(
175 cls,
176 hosts: tuple[str, ...],
177 keyspace: str,
178 *,
179 schema_file: str | None = None,
180 ss_schema_file: str | None = None,
181 read_sources_months: int | None = None,
182 read_forced_sources_months: int | None = None,
183 enable_replica: bool = False,
184 replica_skips_diaobjects: bool = False,
185 port: int | None = None,
186 username: str | None = None,
187 dbauth_alias: str | None = None,
188 prefix: str | None = None,
189 part_pixelization: str | None = None,
190 part_pix_level: int | None = None,
191 time_partition_tables: bool = True,
192 time_partition_start: str | None = None,
193 time_partition_end: str | None = None,
194 read_consistency: str | None = None,
195 write_consistency: str | None = None,
196 read_timeout: int | None = None,
197 write_timeout: int | None = None,
198 ra_dec_columns: tuple[str, str] | None = None,
199 replication_factor: int | None = None,
200 drop: bool = False,
201 table_options: CreateTableOptions | None = None,
202 ) -> ApdbCassandraConfig:
203 """Initialize new APDB instance and make configuration object for it.
205 Parameters
206 ----------
207 hosts : `tuple` [`str`, ...]
208 List of host names or IP addresses for Cassandra cluster.
209 keyspace : `str`
210 Name of the keyspace for APDB tables.
211 schema_file : `str`, optional
212 Location of (YAML) configuration file with APDB schema. If not
213 specified then default location will be used.
214 ss_schema_file : `str`, optional
215 Location of (YAML) configuration file with SSO schema. If not
216 specified then default location will be used.
217 read_sources_months : `int`, optional
218 Number of months of history to read from DiaSource.
219 read_forced_sources_months : `int`, optional
220 Number of months of history to read from DiaForcedSource.
221 enable_replica : `bool`, optional
222 If True, make additional tables used for replication to PPDB.
223 replica_skips_diaobjects : `bool`, optional
224 If `True` then do not fill regular ``DiaObject`` table when
225 ``enable_replica`` is `True`.
226 port : `int`, optional
227 Port number to use for Cassandra connections.
228 username : `str`, optional
229 User name for Cassandra connections.
230 dbauth_alias : `str`, optional
231 If specified then this string will be used to as a host name when
232 checking credentials in db-auth.yaml in addition to regular host
233 names in contact_points. For example if
234 dbauth_alias='pp_apdb_prod_cluster' then the entry
235 'cassandra://pp_apdb_prod_cluster/' will match. Port number should
236 not be used in that entry. Alias has higher priority than host
237 names.
238 prefix : `str`, optional
239 Optional prefix for all table names.
240 part_pixelization : `str`, optional
241 Name of the MOC pixelization used for partitioning.
242 part_pix_level : `int`, optional
243 Pixelization level.
244 time_partition_tables : `bool`, optional
245 Create per-partition tables.
246 time_partition_start : `str`, optional
247 Starting time for per-partition tables, in yyyy-mm-ddThh:mm:ss
248 format, in TAI.
249 time_partition_end : `str`, optional
250 Ending time for per-partition tables, in yyyy-mm-ddThh:mm:ss
251 format, in TAI.
252 read_consistency : `str`, optional
253 Name of the consistency level for read operations.
254 write_consistency : `str`, optional
255 Name of the consistency level for write operations.
256 read_timeout : `int`, optional
257 Read timeout in seconds.
258 write_timeout : `int`, optional
259 Write timeout in seconds.
260 ra_dec_columns : `tuple` [`str`, `str`], optional
261 Names of ra/dec columns in DiaObject table.
262 replication_factor : `int`, optional
263 Replication factor used when creating new keyspace, if keyspace
264 already exists its replication factor is not changed.
265 drop : `bool`, optional
266 If `True` then drop existing tables before re-creating the schema.
267 table_options : `CreateTableOptions`, optional
268 Options used when creating Cassandra tables.
270 Returns
271 -------
272 config : `ApdbCassandraConfig`
273 Resulting configuration object for a created APDB instance.
274 """
275 # Some non-standard defaults for connection parameters, these can be
276 # changed later in generated config. Check Cassandra driver
277 # documentation for what these parameters do. These parameters are not
278 # used during database initialization, but they will be saved with
279 # generated config.
280 connection_config = ApdbCassandraConnectionConfig(
281 extra_parameters={
282 "idle_heartbeat_interval": 0,
283 "idle_heartbeat_timeout": 30,
284 "control_connection_timeout": 100,
285 },
286 )
287 config = ApdbCassandraConfig(
288 contact_points=hosts,
289 keyspace=keyspace,
290 enable_replica=enable_replica,
291 replica_skips_diaobjects=replica_skips_diaobjects,
292 connection_config=connection_config,
293 )
294 config.partitioning.time_partition_tables = time_partition_tables
295 if schema_file is not None:
296 config.schema_file = schema_file
297 if ss_schema_file is not None:
298 config.ss_schema_file = ss_schema_file
299 if read_sources_months is not None:
300 config.read_sources_months = read_sources_months
301 if read_forced_sources_months is not None:
302 config.read_forced_sources_months = read_forced_sources_months
303 if port is not None:
304 config.connection_config.port = port
305 if username is not None:
306 config.connection_config.username = username
307 if dbauth_alias is not None:
308 config.connection_config.dbauth_alias = dbauth_alias
309 if prefix is not None:
310 config.prefix = prefix
311 if part_pixelization is not None:
312 config.partitioning.part_pixelization = part_pixelization
313 if part_pix_level is not None:
314 config.partitioning.part_pix_level = part_pix_level
315 if time_partition_start is not None:
316 config.partitioning.time_partition_start = time_partition_start
317 if time_partition_end is not None:
318 config.partitioning.time_partition_end = time_partition_end
319 if read_consistency is not None:
320 config.connection_config.read_consistency = read_consistency
321 if write_consistency is not None:
322 config.connection_config.write_consistency = write_consistency
323 if read_timeout is not None:
324 config.connection_config.read_timeout = read_timeout
325 if write_timeout is not None:
326 config.connection_config.write_timeout = write_timeout
327 if ra_dec_columns is not None:
328 config.ra_dec_columns = ra_dec_columns
330 cls._makeSchema(config, drop=drop, replication_factor=replication_factor, table_options=table_options)
332 return config
334 def get_replica(self) -> ApdbCassandraReplica:
335 """Return `ApdbReplica` instance for this database."""
336 # Note that this instance has to stay alive while replica exists, so
337 # we pass reference to self.
338 return ApdbCassandraReplica(self)
340 @classmethod
341 def _makeSchema(
342 cls,
343 config: ApdbConfig,
344 *,
345 drop: bool = False,
346 replication_factor: int | None = None,
347 table_options: CreateTableOptions | None = None,
348 ) -> None:
349 # docstring is inherited from a base class
351 if not isinstance(config, ApdbCassandraConfig):
352 raise TypeError(f"Unexpected type of configuration object: {type(config)}")
354 simple_schema = ApdbSchema(config.schema_file, config.ss_schema_file)
356 with SessionContext(config) as session:
357 schema = ApdbCassandraSchema(
358 session=session,
359 keyspace=config.keyspace,
360 table_schemas=simple_schema.tableSchemas,
361 prefix=config.prefix,
362 time_partition_tables=config.partitioning.time_partition_tables,
363 enable_replica=config.enable_replica,
364 replica_skips_diaobjects=config.replica_skips_diaobjects,
365 )
367 # Ask schema to create all tables.
368 part_range_config: ApdbCassandraTimePartitionRange | None = None
369 if config.partitioning.time_partition_tables:
370 partitioner = Partitioner(config)
371 time_partition_start = astropy.time.Time(
372 config.partitioning.time_partition_start, format="isot", scale="tai"
373 )
374 time_partition_end = astropy.time.Time(
375 config.partitioning.time_partition_end, format="isot", scale="tai"
376 )
377 part_range_config = ApdbCassandraTimePartitionRange(
378 start=partitioner.time_partition(time_partition_start),
379 end=partitioner.time_partition(time_partition_end),
380 )
381 schema.makeSchema(
382 drop=drop,
383 part_range=part_range_config,
384 replication_factor=replication_factor,
385 table_options=table_options,
386 )
387 else:
388 schema.makeSchema(
389 drop=drop, replication_factor=replication_factor, table_options=table_options
390 )
392 meta_table_name = ApdbTables.metadata.table_name(config.prefix)
393 metadata = ApdbMetadataCassandra(
394 session, meta_table_name, config.keyspace, "read_tuples", "write"
395 )
397 # Fill version numbers, overrides if they existed before.
398 metadata.set(
399 ConnectionContext.metadataSchemaVersionKey, str(simple_schema.schemaVersion()), force=True
400 )
401 metadata.set(
402 ConnectionContext.metadataCodeVersionKey, str(cls.apdbImplementationVersion()), force=True
403 )
405 if config.enable_replica:
406 # Only store replica code version if replica is enabled.
407 metadata.set(
408 ConnectionContext.metadataReplicaVersionKey,
409 str(ApdbCassandraReplica.apdbReplicaImplementationVersion()),
410 force=True,
411 )
413 # Store frozen part of a configuration in metadata.
414 freezer = ApdbConfigFreezer[ApdbCassandraConfig](ConnectionContext.frozen_parameters)
415 metadata.set(ConnectionContext.metadataConfigKey, freezer.to_json(config), force=True)
417 # Store time partition range.
418 if part_range_config:
419 part_range_config.save_to_meta(metadata)
421 def getDiaObjects(self, region: sphgeom.Region) -> pandas.DataFrame:
422 # docstring is inherited from a base class
423 context = self._context
424 config = context.config
426 sp_where, num_sp_part = context.partitioner.spatial_where(region)
427 _LOG.debug("getDiaObjects: #partitions: %s", len(sp_where))
429 # We need to exclude extra partitioning columns from result.
430 column_names = context.schema.apdbColumnNames(ApdbTables.DiaObjectLast)
431 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
432 query = Select(self._keyspace, table_name, column_names)
433 statements: list[tuple] = []
434 for where_clause in sp_where:
435 full_query = query.where(where_clause)
436 statements.append(context.stmt_factory.with_params(full_query, prepare=True))
437 _LOG.debug("getDiaObjects: #queries: %s", len(statements))
439 with self._timer("select_time", tags={"table": "DiaObject", "method": "getDiaObjects"}) as timer:
440 raw_objects = cast(
441 ApdbCassandraTableData,
442 select_concurrent(
443 context.session,
444 statements,
445 "read_raw_multi",
446 config.connection_config.read_concurrency,
447 ),
448 )
449 objects = raw_objects.to_pandas(context.schema._table_schema(ApdbTables.DiaObjectLast))
450 timer.add_values(row_count=len(objects), num_sp_part=num_sp_part, num_queries=len(statements))
452 self._fix_null_columns(ApdbTables.DiaObjectLast, objects)
454 _LOG.debug("found %s DiaObjects", objects.shape[0])
455 return objects
457 def getDiaSources(
458 self,
459 region: sphgeom.Region,
460 object_ids: Iterable[int] | None,
461 visit_time: astropy.time.Time,
462 start_time: astropy.time.Time | None = None,
463 ) -> pandas.DataFrame | None:
464 # docstring is inherited from a base class
465 context = self._context
466 config = context.config
468 months = config.read_sources_months
469 if start_time is None and months == 0:
470 return None
472 mjd_end = float(visit_time.tai.mjd)
473 if start_time is None:
474 mjd_start = mjd_end - months * 30
475 else:
476 mjd_start = float(start_time.tai.mjd)
478 sources = self._getSources(region, object_ids, mjd_start, mjd_end, ApdbTables.DiaSource)
480 self._fix_null_columns(ApdbTables.DiaSource, sources)
482 return sources
484 def getDiaForcedSources(
485 self,
486 region: sphgeom.Region,
487 object_ids: Iterable[int] | None,
488 visit_time: astropy.time.Time,
489 start_time: astropy.time.Time | None = None,
490 ) -> pandas.DataFrame | None:
491 # docstring is inherited from a base class
492 context = self._context
493 config = context.config
495 months = config.read_forced_sources_months
496 if start_time is None and months == 0:
497 return None
499 mjd_end = float(visit_time.tai.mjd)
500 if start_time is None:
501 mjd_start = mjd_end - months * 30
502 else:
503 mjd_start = float(start_time.tai.mjd)
505 sources = self._getSources(region, object_ids, mjd_start, mjd_end, ApdbTables.DiaForcedSource)
507 self._fix_null_columns(ApdbTables.DiaForcedSource, sources)
509 return sources
511 def getDiaObjectsForDedup(self, since: astropy.time.Time | None = None) -> pandas.DataFrame:
512 # docstring is inherited from a base class
513 context = self._context
514 config = context.config
516 if not context.has_dedup_table:
517 raise TypeError("DiaObjectDedup table does not exist in this APDB instance.")
519 if since is None:
520 # Read last deduplication time from metadata.
521 dedup_str = context.metadata.get(context.metadataDedupKey)
522 if dedup_str is not None:
523 dedup_state = json.loads(dedup_str)
524 dedup_time_str = dedup_state["dedup_time_iso_tai"]
525 since = astropy.time.Time(dedup_time_str, format="iso", scale="tai")
527 column_names = context.schema.apdbColumnNames(ExtraTables.DiaObjectDedup)
529 validity_start_column = self._timestamp_column_name("validityStart")
530 timestamp = None if since is None else self._timestamp_column_value(since)
532 table_name = context.schema.tableName(ExtraTables.DiaObjectDedup)
533 query = Select(self._keyspace, table_name, column_names, extra_clause="ALLOW FILTERING")
534 query = query.where(C("dedup_part") == 0)
535 if since is not None:
536 query = query.where(C(validity_start_column) >= 0)
538 statement = context.stmt_factory(query, prepare=False)
540 num_part = config.partitioning.num_part_dedup
541 statements = []
542 for dedup_part in range(num_part):
543 params = (dedup_part,) if timestamp is None else (dedup_part, timestamp)
544 statements.append((statement, params))
546 with self._timer(
547 "select_time", tags={"table": "DiaObjectDedup", "method": "getDiaObjectsForDedup"}
548 ) as timer:
549 objects_raw = cast(
550 ApdbCassandraTableData,
551 select_concurrent(
552 context.session,
553 statements,
554 "read_raw_multi_dedup",
555 config.connection_config.read_concurrency,
556 ),
557 )
558 objects = objects_raw.to_pandas(context.schema._table_schema(ExtraTables.DiaObjectDedup))
559 timer.add_values(row_count=len(objects), num_queries=num_part)
561 # DiaObjectDedup columns are copies of DiaObject columns.
562 self._fix_null_columns(ApdbTables.DiaObject, objects)
564 _LOG.debug("found %s DiaObjectDedup records", objects.shape[0])
565 return objects
567 def getDiaSourcesForDiaObjects(
568 self, objects: list[DiaObjectId], start_time: astropy.time.Time, max_dist_arcsec: float = 1.0
569 ) -> pandas.DataFrame:
570 # docstring is inherited from a base class
571 context = self._context
572 config = context.config
574 # Which tables to query and temporal constraints.
575 end_time = self._current_time()
576 tables, temporal_where = context.partitioner.temporal_where(
577 ApdbTables.DiaSource,
578 start_time,
579 end_time,
580 partitons_range=context.time_partitions_range,
581 query_per_time_part=False,
582 )
583 if not tables:
584 warnings.warn(
585 f"Query time range ({start_time.isot} - {end_time.isot}) does not overlap database "
586 "time partitions."
587 )
589 # Group DiaObjects by partition.
590 partitioned_object_ids = self._group_dia_objects_by_partition(
591 context.partitioner, objects, max_dist_arcsec
592 )
594 # Columns to return.
595 column_names = context.schema.apdbColumnNames(ApdbTables.DiaSource)
597 # Make a bunch of queries.
598 statements = []
599 for apdb_part, diaObjectIds in partitioned_object_ids.items():
600 spatial_where = [C("apdb_part") == apdb_part]
601 for table in tables:
602 query = Select(self._keyspace, table, column_names, extra_clause="ALLOW FILTERING")
603 for id_chunk in chunk_iterable(diaObjectIds, 10_000):
604 id_where = C("diaObjectId").in_(id_chunk)
605 for clause in QExpr.combine(spatial_where, temporal_where, extra=id_where):
606 statements.append(
607 context.stmt_factory.with_params(query.where(clause), prepare=False)
608 )
610 _LOG.debug("getDiaSourcesForDiaObjects #queries: %s", len(statements))
612 with self._timer(
613 "select_time", tags={"table": "DiaSource", "method": "getDiaSourcesForDiaObjects"}
614 ) as timer:
615 table_data_raw = cast(
616 ApdbCassandraTableData,
617 select_concurrent(
618 context.session,
619 statements,
620 "read_raw_multi",
621 config.connection_config.read_concurrency,
622 ),
623 )
624 catalog = table_data_raw.to_pandas(context.schema._table_schema(ApdbTables.DiaSource))
625 timer.add_values(row_count_from_db=len(catalog), num_queries=len(statements))
627 # precise filtering on midpointMjdTai
628 catalog = cast(pandas.DataFrame, catalog[catalog["midpointMjdTai"] >= start_time.tai.mjd])
630 timer.add_values(row_count=len(catalog))
632 self._fix_null_columns(ApdbTables.DiaSource, catalog)
634 _LOG.debug("found %d DiaSources", len(catalog))
635 return catalog
637 def containsVisitDetector(
638 self,
639 visit: int,
640 detector: int,
641 region: sphgeom.Region | None = None,
642 visit_time: astropy.time.Time | None = None,
643 ) -> bool:
644 # docstring is inherited from a base class
645 context = self._context
647 table_name = context.schema.tableName(ExtraTables.ApdbVisitDetector)
648 query = Select(self._keyspace, table_name, [ColumnExpr("count(*)")])
649 query = query.where((C("visit") == visit) & (C("detector") == detector))
650 stmt, params = context.stmt_factory.with_params(query, prepare=False)
652 with self._timer("contains_visit_detector_time", tags={"table": table_name}):
653 result = context.session.execute(stmt, params)
654 return bool(result.one()[0])
656 def store(
657 self,
658 visit_time: astropy.time.Time,
659 objects: pandas.DataFrame,
660 sources: pandas.DataFrame | None = None,
661 forced_sources: pandas.DataFrame | None = None,
662 ) -> None:
663 # docstring is inherited from a base class
664 context = self._context
665 config = context.config
667 # Store visit/detector in a special table, this has to be done
668 # before all other writes so if there is a failure at any point
669 # later we still have a record for attempted write.
670 visit_detector: set[tuple[int, int]] = set()
671 for df in sources, forced_sources:
672 if df is not None and not df.empty:
673 df = df[["visit", "detector"]]
674 for visit, detector in df.itertuples(index=False):
675 visit_detector.add((visit, detector))
677 if visit_detector:
678 # Typically there is only one entry, do not bother with
679 # concurrency.
680 table_name = context.schema.tableName(ExtraTables.ApdbVisitDetector)
681 query = Insert(self._keyspace, table_name, ("visit", "detector"))
682 stmt = context.stmt_factory(query)
683 for item in visit_detector:
684 context.session.execute(stmt, item, execution_profile="write")
686 objects = self._fix_input_timestamps(objects)
687 if sources is not None:
688 sources = self._fix_input_timestamps(sources)
689 if forced_sources is not None:
690 forced_sources = self._fix_input_timestamps(forced_sources)
692 replica_chunk: ReplicaChunk | None = None
693 if context.schema.replication_enabled:
694 replica_chunk = ReplicaChunk.make_replica_chunk(visit_time, config.replica_chunk_seconds)
695 self._storeReplicaChunk(replica_chunk)
697 # fill region partition column for DiaObjects
698 objects = self._add_apdb_part(objects)
699 self._storeDiaObjects(objects, visit_time, replica_chunk)
701 if sources is not None and len(sources) > 0:
702 # copy apdb_part column from DiaObjects to DiaSources
703 sources = self._add_apdb_part(sources)
704 subchunk = self._storeDiaSources(ApdbTables.DiaSource, sources, replica_chunk)
705 self._storeDiaSourcesPartitions(sources, visit_time, replica_chunk, subchunk)
707 if forced_sources is not None and len(forced_sources) > 0:
708 forced_sources = self._add_apdb_part(forced_sources)
709 self._storeDiaSources(ApdbTables.DiaForcedSource, forced_sources, replica_chunk)
711 def reassignDiaSourcesToDiaObjects(
712 self,
713 idMap: Mapping[DiaSourceId, int],
714 *,
715 increment_nDiaSources: bool = True,
716 decrement_nDiaSources: bool = True,
717 ) -> None:
718 # docstring is inherited from a base class
719 context = self._context
720 config = context.config
722 source_ids = {source_id.diaSourceId for source_id in idMap}
724 # Find all DiaSources.
725 found_sources = self._get_diasource_data(
726 idMap, "apdb_part", "diaObjectId", "ra", "dec", "midpointMjdTai"
727 )
729 if missing_ids := (source_ids - {row.diaSourceId for row in found_sources}):
730 raise LookupError(f"Some source IDs were not found in DiaSource table: {missing_ids}")
732 found_sources_by_id = {row.diaSourceId: row for row in found_sources}
734 # Make sure that all DiaObjects exist, we also want to know
735 # nDiaSources count for current and new records because we want to
736 # send updated values to replica.
737 current_object_ids = {
738 DiaObjectId(diaObjectId=row.diaObjectId, ra=row.ra, dec=row.dec) for row in found_sources
739 }
740 # Assume that DiaSource ra/dec are very close to re-assigned objects.
741 new_object_ids = {
742 DiaObjectId(diaObjectId=diaObjectId, ra=source_id.ra, dec=source_id.dec)
743 for source_id, diaObjectId in idMap.items()
744 }
745 all_object_ids = new_object_ids | current_object_ids
746 found_objects = self._get_diaobject_data(all_object_ids, "apdb_part", "ra", "dec", "nDiaSources")
748 if missing_ids := (
749 {row.diaObjectId for row in all_object_ids} - {row.diaObjectId for row in found_objects}
750 ):
751 raise LookupError(f"Some object IDs were not found in DiaObjectLast table: {missing_ids}")
753 update_records: list[ApdbUpdateRecord] = []
754 update_order = 0
755 current_time = self._current_time()
756 current_time_ns = int(current_time.unix_tai * 1e9)
758 # Update DiaSources.
759 statements: list[tuple] = []
760 for source_id, diaObjectId in idMap.items():
761 source_row = found_sources_by_id[source_id.diaSourceId]
762 apdb_part = source_row.apdb_part
763 time_part = context.partitioner.time_partition(source_row.midpointMjdTai)
765 if config.partitioning.time_partition_tables:
766 table_name = context.schema.tableName(ApdbTables.DiaSource, time_part)
767 update = (
768 Update(self._keyspace, table_name)
769 .values(C("diaObjectId").update(diaObjectId))
770 .where(C("apdb_part") == apdb_part)
771 .where(C("diaSourceId") == source_id.diaSourceId)
772 )
773 else:
774 table_name = context.schema.tableName(ApdbTables.DiaSource)
775 update = (
776 Update(self._keyspace, table_name)
777 .values(C("diaObjectId").update(diaObjectId))
778 .where(C("apdb_part") == apdb_part)
779 .where(C("apdb_time_part") == time_part)
780 .where(C("diaSourceId") == source_id.diaSourceId)
781 )
782 statements.append(context.stmt_factory.with_params(update, prepare=True))
784 if context.schema.replication_enabled:
785 update_records.append(
786 ApdbReassignDiaSourceToDiaObjectRecord(
787 diaSourceId=source_id.diaSourceId,
788 ra=source_id.ra,
789 dec=source_id.dec,
790 midpointMjdTai=source_id.midpointMjdTai,
791 diaObjectId=diaObjectId,
792 update_time_ns=current_time_ns,
793 update_order=update_order,
794 )
795 )
796 update_order += 1
798 with self._timer(
799 "update_time", tags={"table": "DiaSource", "method": "reassignDiaSourcesToDiaObjects"}
800 ) as timer:
801 execute_concurrent(context.session, statements, execution_profile="write")
802 timer.add_values(num_queries=len(statements))
804 # Update nDiaSources in DiaObjectLast. We do not update DiaObject table
805 # here because it may not even exist. PPDB updates DiaObject from
806 # update records.
807 if increment_nDiaSources or decrement_nDiaSources:
808 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
809 update = (
810 Update(self._keyspace, table_name)
811 .values(C("nDiaSources").update(-1))
812 .where(C("apdb_part") == -1)
813 .where(C("diaObjectId") == -1)
814 )
815 statement = context.stmt_factory(update, prepare=True)
816 statements = []
818 # Calculate increments/decrements for all affected DiaObjects.
819 increments: Counter = Counter()
820 if increment_nDiaSources:
821 increments.update(idMap.values())
822 if decrement_nDiaSources:
823 increments.subtract(row.diaObjectId for row in found_sources)
825 for row in found_objects:
826 if increments.get(row.diaObjectId):
827 nDiaSources = row.nDiaSources + increments[row.diaObjectId]
828 statements.append((statement, (nDiaSources, row.apdb_part, row.diaObjectId)))
830 # Also send updated values to replica.
831 if context.schema.replication_enabled:
832 update_records.append(
833 ApdbUpdateNDiaSourcesRecord(
834 diaObjectId=row.diaObjectId,
835 ra=row.ra,
836 dec=row.dec,
837 nDiaSources=nDiaSources,
838 update_time_ns=current_time_ns,
839 update_order=update_order,
840 )
841 )
842 update_order += 1
844 if statements:
845 with self._timer(
846 "update_time", tags={"table": table_name, "method": "reassignDiaSourcesToDiaObjects"}
847 ) as timer:
848 execute_concurrent(context.session, statements, execution_profile="write")
849 timer.add_values(num_queries=len(statements))
851 if update_records:
852 replica_chunk = ReplicaChunk.make_replica_chunk(current_time, config.replica_chunk_seconds)
853 self._storeUpdateRecords(update_records, replica_chunk, store_chunk=True)
855 def setValidityEnd(
856 self, objects: list[DiaObjectId], validityEnd: astropy.time.Time, raise_on_missing_id: bool = False
857 ) -> int:
858 # docstring is inherited from a base class
859 if not objects:
860 return 0
862 context = self._context
863 config = context.config
865 pad_arcsec = 1.0
866 partitioned_object_ids = self._group_dia_objects_by_partition(
867 context.partitioner, objects, pad_arcsec
868 )
870 # Check that all objects exist.
871 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
872 statements: list[tuple] = []
873 for apdb_part, diaObjectIds in partitioned_object_ids.items():
874 query = Select(self._keyspace, table_name, ["apdb_part", "diaObjectId"])
875 query = query.where(C("apdb_part") == apdb_part)
876 query = query.where(C("diaObjectId").in_(diaObjectIds))
877 statements.append(context.stmt_factory.with_params(query, prepare=False))
879 with self._timer("select_time", tags={"table": table_name, "method": "setValidityEnd"}) as timer:
880 records = cast(
881 list[tuple[int, int]],
882 select_concurrent(
883 context.session,
884 statements,
885 "read_tuples",
886 config.connection_config.read_concurrency,
887 ),
888 )
889 timer.add_values(row_count=len(objects))
891 requested_ids = {obj.diaObjectId for obj in objects}
892 found_ids = {rec[1] for rec in records}
893 if extra_ids := (found_ids - requested_ids):
894 raise RuntimeError(f"Consistency error - found duplicate records for object IDs: {extra_ids}")
895 if raise_on_missing_id:
896 if missing_ids := (requested_ids - found_ids):
897 raise LookupError(f"Some object IDs are missing from DiaObjectLast table: {missing_ids}")
899 # Filter existing records.
900 if len(objects) != len(found_ids):
901 objects = [obj for obj in objects if obj.diaObjectId in found_ids]
903 if not objects:
904 return 0
906 # Group by partitions again.
907 grouped_object_ids: dict[int, list[int]] = defaultdict(list)
908 for apdb_part, diaObjectId in records:
909 grouped_object_ids[apdb_part].append(diaObjectId)
911 # Remove all matching rows from DiaObjectLast.
912 statements = []
913 for apdb_part, diaObjectIds in grouped_object_ids.items():
914 delete = (
915 Delete(self._keyspace, table_name)
916 .where(C("apdb_part") == apdb_part)
917 .where(C("diaObjectId").in_(diaObjectIds))
918 )
919 statements.append(context.stmt_factory.with_params(delete))
921 # Also remove from DiaObjectLastToPartition.
922 reverse_table_name = context.schema.tableName(ExtraTables.DiaObjectLastToPartition)
923 delete = Delete(self._keyspace, reverse_table_name).where(
924 C("diaObjectId").in_([rec[1] for rec in records])
925 )
926 statements.append(context.stmt_factory.with_params(delete))
928 with self._timer("delete_time", tags={"table": table_name, "method": "setValidityEnd"}) as timer:
929 execute_concurrent(context.session, statements, execution_profile="write")
930 timer.add_values(row_count=len(records))
932 # If repication is enabled then send all updates.
933 if context.schema.replication_enabled:
934 current_time = self._current_time()
935 current_time_ns = int(current_time.unix_tai * 1e9)
936 replica_chunk = ReplicaChunk.make_replica_chunk(current_time, config.replica_chunk_seconds)
938 update_records = [
939 ApdbCloseDiaObjectValidityRecord(
940 diaObjectId=obj.diaObjectId,
941 ra=obj.ra,
942 dec=obj.dec,
943 update_time_ns=current_time_ns,
944 update_order=index,
945 validityEndMjdTai=float(validityEnd.tai.mjd),
946 nDiaSources=None,
947 )
948 for index, obj in enumerate(objects)
949 ]
951 self._storeUpdateRecords(update_records, replica_chunk, store_chunk=True)
953 return len(objects)
955 def resetDedup(self, dedup_time: astropy.time.Time | None = None) -> None:
956 # docstring is inherited from a base class
957 context = self._context
959 if not context.has_dedup_table:
960 raise TypeError("DiaObjectDedup table does not exist in this APDB instance.")
962 if dedup_time is None:
963 dedup_time = self._current_time()
965 validity_start_column = self._timestamp_column_name("validityStart")
967 # Find latest timestamp in deduplication table.
968 table_name = context.schema.tableName(ExtraTables.DiaObjectDedup)
969 query = Select(self._keyspace, table_name, [ColumnExpr(f'MAX("{validity_start_column}")')])
970 stmt = context.stmt_factory(query, prepare=False)
972 result = context.session.execute(stmt, execution_profile="read_tuples")
973 max_value = result.one()[0]
974 if self._schema.has_mjd_timestamps:
975 max_validity_start = astropy.time.Time(max_value, format="mjd", scale="tai")
976 else:
977 max_validity_start = astropy.time.Time(max_value, format="datetime", scale="tai")
979 # If max time is lower than dedup time we can do TRUNCATE.
980 if dedup_time >= max_validity_start:
981 query_str = f'TRUNCATE TABLE "{self._keyspace}"."{table_name}"'
982 context.session.execute(query_str, execution_profile="write")
983 else:
984 dedup_time_value = self._timestamp_column_value(dedup_time)
985 delete = Delete(self._keyspace, table_name).where(C(validity_start_column) < dedup_time_value)
986 stmt, params = context.stmt_factory.with_params(delete)
987 context.session.execute(stmt, params, execution_profile="write")
989 # Store dedup time.
990 data = {"dedup_time_iso_tai": dedup_time.tai.to_value("iso")}
991 data_json = json.dumps(data)
992 context.metadata.set(context.metadataDedupKey, data_json, force=True)
994 def reassignDiaSources(self, idMap: Mapping[int, int]) -> None:
995 # docstring is inherited from a base class
996 context = self._context
997 config = context.config
999 now = self._current_time()
1000 reassign_time_column = self._timestamp_column_name("ssObjectReassocTime")
1001 reassignTime = self._timestamp_column_value(now)
1003 # To update a record we need to know its exact primary key (including
1004 # partition key) so we start by querying for diaSourceId to find the
1005 # primary keys.
1007 table_name = context.schema.tableName(ExtraTables.DiaSourceToPartition)
1008 # split it into 1k IDs per query
1009 selects: list[tuple] = []
1010 columns = ["diaSourceId", "apdb_part", "apdb_time_part", "apdb_replica_chunk"]
1011 query = Select(self._keyspace, table_name, columns)
1012 for ids in chunk_iterable(idMap.keys(), 1_000):
1013 full_query = query.where(C("diaSourceId").in_(ids))
1014 selects.append(context.stmt_factory.with_params(full_query, prepare=False))
1016 # No need for DataFrame here, read data as tuples.
1017 result = cast(
1018 list[tuple[int, int, int, int | None]],
1019 select_concurrent(
1020 context.session, selects, "read_tuples", config.connection_config.read_concurrency
1021 ),
1022 )
1024 # Make mapping from source ID to its partition.
1025 id2partitions: dict[int, tuple[int, int]] = {}
1026 id2chunk_id: dict[int, int] = {}
1027 for row in result:
1028 id2partitions[row[0]] = row[1:3]
1029 if row[3] is not None:
1030 id2chunk_id[row[0]] = row[3]
1032 # make sure we know partitions for each ID
1033 if set(id2partitions) != set(idMap):
1034 missing = ",".join(str(item) for item in set(idMap) - set(id2partitions))
1035 raise ValueError(f"Following DiaSource IDs do not exist in the database: {missing}")
1037 # Reassign in standard tables
1038 queries: list[tuple[cassandra.query.PreparedStatement, tuple]] = []
1039 for diaSourceId, ssObjectId in idMap.items():
1040 apdb_part, apdb_time_part = id2partitions[diaSourceId]
1041 if config.partitioning.time_partition_tables:
1042 table_name = context.schema.tableName(ApdbTables.DiaSource, apdb_time_part)
1043 update = (
1044 Update(self._keyspace, table_name)
1045 .values(
1046 C("ssObjectId").update(ssObjectId),
1047 C("diaObjectId").update(None),
1048 C(reassign_time_column).update(reassignTime),
1049 )
1050 .where(C("apdb_part") == apdb_part)
1051 .where(C("diaSourceId") == diaSourceId)
1052 )
1053 else:
1054 table_name = context.schema.tableName(ApdbTables.DiaSource)
1055 update = (
1056 Update(self._keyspace, table_name)
1057 .values(
1058 C("ssObjectId").update(ssObjectId),
1059 C("diaObjectId").update(None),
1060 C(reassign_time_column).update(reassignTime),
1061 )
1062 .where(C("apdb_part") == apdb_part)
1063 .where(C("apdb_time_part") == apdb_time_part)
1064 .where(C("diaSourceId") == diaSourceId)
1065 )
1066 queries.append(context.stmt_factory.with_params(update, prepare=True))
1068 # TODO: (DM-50190) Replication for updated records is not implemented.
1069 if id2chunk_id:
1070 warnings.warn("Replication of reassigned DiaSource records is not implemented.", stacklevel=2)
1072 _LOG.debug("%s: will update %d records", table_name, len(idMap))
1073 with self._timer("source_reassign_time") as timer:
1074 execute_concurrent(context.session, queries, execution_profile="write")
1075 timer.add_values(source_count=len(idMap))
1077 def withdrawDiaSources(
1078 self,
1079 diaSourceIds: Iterable[DiaSourceId],
1080 *,
1081 timeWithdrawn: astropy.time.Time | None = None,
1082 ) -> None:
1083 # docstring is inherited from a base class
1084 context = self._context
1085 config = context.config
1087 if timeWithdrawn is None:
1088 timeWithdrawn = self._current_time()
1089 time_value = self._timestamp_column_value(timeWithdrawn)
1090 column_name = self._timestamp_column_name("time_withdrawn")
1092 diaSourceIds = list(diaSourceIds)
1093 source_ids = {source_id.diaSourceId for source_id in diaSourceIds}
1095 # Find all DiaSources.
1096 found_sources = self._get_diasource_data(
1097 diaSourceIds, "apdb_part", "diaObjectId", "ra", "dec", "midpointMjdTai", column_name
1098 )
1100 if missing_ids := (source_ids - {row.diaSourceId for row in found_sources}):
1101 raise LookupError(f"Some source IDs were not found in DiaSource table: {missing_ids}")
1103 found_sources_by_id = {row.diaSourceId: row for row in found_sources}
1105 update_records: list[ApdbUpdateRecord] = []
1106 update_order = 0
1107 current_time = self._current_time()
1108 current_time_ns = int(current_time.unix_tai * 1e9)
1110 # Update DiaSources.
1111 statements: list[tuple] = []
1112 for source_id in diaSourceIds:
1113 source_row = found_sources_by_id[source_id.diaSourceId]
1114 # Ignore sources already withdrawn.
1115 if getattr(source_row, column_name) is not None:
1116 continue
1118 apdb_part = source_row.apdb_part
1119 time_part = context.partitioner.time_partition(source_row.midpointMjdTai)
1121 if config.partitioning.time_partition_tables:
1122 table_name = context.schema.tableName(ApdbTables.DiaSource, time_part)
1123 update = (
1124 Update(self._keyspace, table_name)
1125 .values(C(column_name).update(time_value))
1126 .where(C("apdb_part") == apdb_part)
1127 .where(C("diaSourceId") == source_id.diaSourceId)
1128 )
1129 else:
1130 table_name = context.schema.tableName(ApdbTables.DiaSource)
1131 update = (
1132 Update(self._keyspace, table_name)
1133 .values(C(column_name).update(time_value))
1134 .where(C("apdb_part") == apdb_part)
1135 .where(C("apdb_time_part") == time_part)
1136 .where(C("diaSourceId") == source_id.diaSourceId)
1137 )
1138 statements.append(context.stmt_factory.with_params(update, prepare=True))
1140 if context.schema.replication_enabled:
1141 update_records.append(
1142 ApdbWithdrawDiaSourceRecord(
1143 diaSourceId=source_id.diaSourceId,
1144 ra=source_id.ra,
1145 dec=source_id.dec,
1146 midpointMjdTai=source_id.midpointMjdTai,
1147 update_time_ns=current_time_ns,
1148 update_order=update_order,
1149 timeWithdrawnMjdTai=float(timeWithdrawn.tai.mjd),
1150 )
1151 )
1152 update_order += 1
1154 with self._timer("update_time", tags={"table": "DiaSource", "method": "withdrawDiaSources"}) as timer:
1155 execute_concurrent(context.session, statements, execution_profile="write")
1156 timer.add_values(num_queries=len(statements))
1158 if update_records:
1159 replica_chunk = ReplicaChunk.make_replica_chunk(current_time, config.replica_chunk_seconds)
1160 self._storeUpdateRecords(update_records, replica_chunk, store_chunk=True)
1162 def withdrawDiaForcedSources(
1163 self,
1164 diaForcedSourceIds: Iterable[DiaForcedSourceId],
1165 *,
1166 timeWithdrawn: astropy.time.Time | None = None,
1167 ) -> None:
1168 # docstring is inherited from a base class
1169 context = self._context
1170 config = context.config
1172 def _fsrc_id(fsource: Any) -> tuple[int, int, int]:
1173 return (fsource.diaObjectId, fsource.visit, fsource.detector)
1175 if timeWithdrawn is None:
1176 timeWithdrawn = self._current_time()
1177 time_value = self._timestamp_column_value(timeWithdrawn)
1178 column_name = self._timestamp_column_name("time_withdrawn")
1180 diaForcedSourceIds = list(diaForcedSourceIds)
1181 fsource_keys = {_fsrc_id(source) for source in diaForcedSourceIds}
1183 found_fsources = self._get_diaforcedsource_data(
1184 diaForcedSourceIds, "apdb_part", "ra", "dec", "midpointMjdTai", column_name
1185 )
1187 found_keys = {_fsrc_id(row) for row in found_fsources}
1188 if missing_ids := (fsource_keys - found_keys):
1189 raise LookupError(f"Some source IDs were not found in DiaForcedSource table: {missing_ids}")
1191 statements: list[tuple] = []
1192 update_records = []
1193 update_order = 0
1194 current_time = self._current_time()
1195 current_time_ns = int(current_time.unix_tai * 1e9)
1197 for source_row in found_fsources:
1198 # Ignore sources already withdrawn.
1199 if getattr(source_row, column_name) is not None:
1200 continue
1202 apdb_part = source_row.apdb_part
1203 time_part = context.partitioner.time_partition(source_row.midpointMjdTai)
1205 if config.partitioning.time_partition_tables:
1206 table_name = context.schema.tableName(ApdbTables.DiaForcedSource, time_part)
1207 update = (
1208 Update(self._keyspace, table_name)
1209 .values(C(column_name).update(time_value))
1210 .where(C("apdb_part") == apdb_part)
1211 .where(C("diaObjectId") == source_row.diaObjectId)
1212 .where(C("visit") == source_row.visit)
1213 .where(C("detector") == source_row.detector)
1214 )
1215 else:
1216 table_name = context.schema.tableName(ApdbTables.DiaForcedSource)
1217 update = (
1218 Update(self._keyspace, table_name)
1219 .values(C(column_name).update(time_value))
1220 .where(C("apdb_part") == apdb_part)
1221 .where(C("apdb_time_part") == time_part)
1222 .where(C("diaObjectId") == source_row.diaObjectId)
1223 .where(C("visit") == source_row.visit)
1224 .where(C("detector") == source_row.detector)
1225 )
1226 statements.append(context.stmt_factory.with_params(update, prepare=True))
1228 if context.schema.replication_enabled:
1229 update_records.append(
1230 ApdbWithdrawDiaForcedSourceRecord(
1231 diaObjectId=source_row.diaObjectId,
1232 visit=source_row.visit,
1233 detector=source_row.detector,
1234 ra=source_row.ra,
1235 dec=source_row.dec,
1236 midpointMjdTai=source_row.midpointMjdTai,
1237 update_time_ns=current_time_ns,
1238 update_order=update_order,
1239 timeWithdrawnMjdTai=float(timeWithdrawn.tai.mjd),
1240 )
1241 )
1242 update_order += 1
1244 if statements:
1245 with self._timer(
1246 "update_time", tags={"table": "DiaForcedSource", "method": "withdrawDiaForcedSources"}
1247 ) as timer:
1248 execute_concurrent(context.session, statements, execution_profile="write")
1249 timer.add_values(num_queries=len(statements))
1251 if update_records:
1252 replica_chunk = ReplicaChunk.make_replica_chunk(current_time, config.replica_chunk_seconds)
1253 self._storeUpdateRecords(update_records, replica_chunk, store_chunk=True)
1255 def countUnassociatedObjects(self) -> int:
1256 # docstring is inherited from a base class
1258 # It's too inefficient to implement it for Cassandra in current schema.
1259 raise NotImplementedError()
1261 @property
1262 def schema(self) -> ApdbSchema:
1263 # docstring is inherited from a base class
1264 return self._schema
1266 @property
1267 def metadata(self) -> ApdbMetadata:
1268 # docstring is inherited from a base class
1269 context = self._context
1270 return context.metadata
1272 @property
1273 def admin(self) -> ApdbCassandraAdmin:
1274 # docstring is inherited from a base class
1275 return ApdbCassandraAdmin(self)
1277 def _getSources(
1278 self,
1279 region: sphgeom.Region,
1280 object_ids: Iterable[int] | None,
1281 mjd_start: float,
1282 mjd_end: float,
1283 table_name: ApdbTables,
1284 ) -> pandas.DataFrame:
1285 """Return catalog of DiaSource instances given set of DiaObject IDs.
1287 Parameters
1288 ----------
1289 region : `lsst.sphgeom.Region`
1290 Spherical region.
1291 object_ids :
1292 Collection of DiaObject IDs
1293 mjd_start : `float`
1294 Lower bound of time interval.
1295 mjd_end : `float`
1296 Upper bound of time interval.
1297 table_name : `ApdbTables`
1298 Name of the table.
1300 Returns
1301 -------
1302 catalog : `pandas.DataFrame`, or `None`
1303 Catalog containing DiaSource records. Empty catalog is returned if
1304 ``object_ids`` is empty.
1305 """
1306 context = self._context
1307 config = context.config
1309 object_id_set: Set[int] = set()
1310 if object_ids is not None:
1311 object_id_set = set(object_ids)
1312 if len(object_id_set) == 0:
1313 return self._make_empty_catalog(table_name)
1315 sp_where, num_sp_part = context.partitioner.spatial_where(region)
1316 tables, temporal_where = context.partitioner.temporal_where(
1317 table_name, mjd_start, mjd_end, partitons_range=context.time_partitions_range
1318 )
1319 if not tables:
1320 start = astropy.time.Time(mjd_start, format="mjd", scale="tai")
1321 end = astropy.time.Time(mjd_end, format="mjd", scale="tai")
1322 warnings.warn(
1323 f"Query time range ({start.isot} - {end.isot}) does not overlap database time partitions."
1324 )
1326 # We need to exclude extra partitioning columns from result.
1327 column_names = context.schema.apdbColumnNames(table_name)
1329 # Build all queries
1330 statements: list[tuple] = []
1331 for table in tables:
1332 query = Select(self._keyspace, table, column_names)
1333 for clause in QExpr.combine(sp_where, temporal_where):
1334 statements.append(context.stmt_factory.with_params(query.where(clause), prepare=True))
1335 _LOG.debug("_getSources %s: #queries: %s", table_name, len(statements))
1337 with self._timer("select_time", tags={"table": table_name.name, "method": "_getSources"}) as timer:
1338 table_data_raw = cast(
1339 ApdbCassandraTableData,
1340 select_concurrent(
1341 context.session,
1342 statements,
1343 "read_raw_multi",
1344 config.connection_config.read_concurrency,
1345 ),
1346 )
1347 catalog = table_data_raw.to_pandas(context.schema._table_schema(table_name))
1348 timer.add_values(
1349 row_count_from_db=len(catalog), num_sp_part=num_sp_part, num_queries=len(statements)
1350 )
1352 # filter by given object IDs
1353 if len(object_id_set) > 0:
1354 catalog = cast(pandas.DataFrame, catalog[catalog["diaObjectId"].isin(object_id_set)])
1356 # precise filtering on midpointMjdTai
1357 catalog = cast(pandas.DataFrame, catalog[catalog["midpointMjdTai"] > mjd_start])
1359 timer.add_values(row_count=len(catalog))
1361 _LOG.debug("found %d %ss", catalog.shape[0], table_name.name)
1362 return catalog
1364 def _storeReplicaChunk(self, replica_chunk: ReplicaChunk) -> None:
1365 context = self._context
1366 config = context.config
1368 # Cassandra timestamp uses milliseconds since epoch
1369 timestamp = int(replica_chunk.last_update_time.unix_tai * 1000)
1371 # everything goes into a single partition
1372 partition = 0
1374 table_name = context.schema.tableName(ExtraTables.ApdbReplicaChunks)
1376 columns = ["partition", "apdb_replica_chunk", "last_update_time", "unique_id"]
1377 values = [partition, replica_chunk.id, timestamp, replica_chunk.unique_id]
1378 if context.has_chunk_sub_partitions:
1379 columns.append("has_subchunks")
1380 values.append(True)
1382 query = Insert(self._keyspace, table_name, columns)
1383 stmt = context.stmt_factory(query)
1385 context.session.execute(
1386 stmt,
1387 values,
1388 timeout=config.connection_config.write_timeout,
1389 execution_profile="write",
1390 )
1392 def _queryDiaObjectLastPartitions(self, ids: Iterable[int]) -> Mapping[int, int]:
1393 """Return existing mapping of diaObjectId to its last partition."""
1394 context = self._context
1395 config = context.config
1397 table_name = context.schema.tableName(ExtraTables.DiaObjectLastToPartition)
1398 queries = []
1399 object_count = 0
1400 for id_chunk in chunk_iterable(ids, 10_000):
1401 id_chunk_list = tuple(id_chunk)
1402 query = Select(self._keyspace, table_name, ("diaObjectId", "apdb_part"))
1403 query = query.where(C("diaObjectId").in_(id_chunk_list))
1404 queries.append(context.stmt_factory.with_params(query, prepare=False))
1405 object_count += len(id_chunk_list)
1407 with self._timer("query_object_last_partitions", tags={"table": table_name}) as timer:
1408 data = cast(
1409 ApdbTableData,
1410 select_concurrent(
1411 context.session,
1412 queries,
1413 "read_raw_multi",
1414 config.connection_config.read_concurrency,
1415 ),
1416 )
1417 timer.add_values(object_count=object_count, row_count=len(data.rows()))
1419 if data.column_names() != ["diaObjectId", "apdb_part"]:
1420 raise RuntimeError(f"Unexpected column names in query result: {data.column_names()}")
1422 return {row[0]: row[1] for row in data.rows()}
1424 def _deleteMovingObjects(self, objs: pandas.DataFrame) -> None:
1425 """Objects in DiaObjectsLast can move from one spatial partition to
1426 another. For those objects inserting new version does not replace old
1427 one, so we need to explicitly remove old versions before inserting new
1428 ones.
1429 """
1430 context = self._context
1432 # Extract all object IDs.
1433 new_partitions = dict(zip(objs["diaObjectId"], objs["apdb_part"]))
1434 old_partitions = self._queryDiaObjectLastPartitions(objs["diaObjectId"])
1436 moved_oids: dict[int, tuple[int, int]] = {}
1437 for oid, old_part in old_partitions.items():
1438 new_part = new_partitions.get(oid, old_part)
1439 if new_part != old_part:
1440 moved_oids[oid] = (old_part, new_part)
1441 _LOG.debug("DiaObject IDs that moved to new partition: %s", moved_oids)
1443 if moved_oids:
1444 # Delete old records from DiaObjectLast.
1445 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
1446 query = Delete(self._keyspace, table_name)
1447 query = query.where('apdb_part = {} AND "diaObjectId" = {}', (-1, -1))
1448 statement = context.stmt_factory(query, prepare=True)
1449 queries = []
1450 for oid, (old_part, _) in moved_oids.items():
1451 queries.append((statement, (old_part, oid)))
1452 with self._timer("delete_object_last", tags={"table": table_name}) as timer:
1453 execute_concurrent(context.session, queries, execution_profile="write")
1454 timer.add_values(row_count=len(moved_oids))
1456 # Add all new records to the map.
1457 table_name = context.schema.tableName(ExtraTables.DiaObjectLastToPartition)
1458 insert = Insert(self._keyspace, table_name, ("diaObjectId", "apdb_part"))
1459 statement = context.stmt_factory(insert, prepare=True)
1461 queries = []
1462 for oid, new_part in new_partitions.items():
1463 queries.append((statement, (oid, new_part)))
1465 with self._timer("update_object_last_partition", tags={"table": table_name}) as timer:
1466 execute_concurrent(context.session, queries, execution_profile="write")
1467 timer.add_values(row_count=len(queries))
1469 def _storeDiaObjects(
1470 self, objs: pandas.DataFrame, visit_time: astropy.time.Time, replica_chunk: ReplicaChunk | None
1471 ) -> None:
1472 """Store catalog of DiaObjects from current visit.
1474 Parameters
1475 ----------
1476 objs : `pandas.DataFrame`
1477 Catalog with DiaObject records
1478 visit_time : `astropy.time.Time`
1479 Time of the current visit.
1480 replica_chunk : `ReplicaChunk` or `None`
1481 Replica chunk identifier if replication is configured.
1482 """
1483 if len(objs) == 0:
1484 _LOG.debug("No objects to write to database.")
1485 return
1487 context = self._context
1488 config = context.config
1490 self._deleteMovingObjects(objs)
1492 validity_start_column = self._timestamp_column_name("validityStart")
1493 timestamp = self._timestamp_column_value(visit_time)
1495 # DiaObjectLast did not have this column in the past.
1496 extra_columns: dict[str, Any] = {}
1497 if context.schema.check_column(ApdbTables.DiaObjectLast, validity_start_column):
1498 extra_columns[validity_start_column] = timestamp
1500 self._storeObjectsPandas(objs, ApdbTables.DiaObjectLast, extra_columns=extra_columns)
1502 extra_columns[validity_start_column] = timestamp
1503 visit_time_part = context.partitioner.time_partition(visit_time)
1504 time_part: int | None = visit_time_part
1505 if (time_partitions_range := context.time_partitions_range) is not None:
1506 self._check_time_partitions([visit_time_part], time_partitions_range)
1507 if not config.partitioning.time_partition_tables:
1508 extra_columns["apdb_time_part"] = time_part
1509 time_part = None
1511 # Only store DiaObects if not doing replication or explicitly
1512 # configured to always store them.
1513 if replica_chunk is None or not config.replica_skips_diaobjects:
1514 self._storeObjectsPandas(
1515 objs, ApdbTables.DiaObject, extra_columns=extra_columns, time_part=time_part
1516 )
1518 if replica_chunk is not None:
1519 extra_columns = {"apdb_replica_chunk": replica_chunk.id, validity_start_column: timestamp}
1520 table = ExtraTables.DiaObjectChunks
1521 if context.has_chunk_sub_partitions:
1522 table = ExtraTables.DiaObjectChunks2
1523 # Use a random number for a second part of partitioning key so
1524 # that different clients could wrtite to different partitions.
1525 # This makes it not exactly reproducible.
1526 extra_columns["apdb_replica_subchunk"] = random.randrange(config.replica_sub_chunk_count)
1527 self._storeObjectsPandas(objs, table, extra_columns=extra_columns)
1529 # Store copy of the records in dedup table.
1530 if context.has_dedup_table:
1531 table = ExtraTables.DiaObjectDedup
1532 extra_columns = {
1533 "dedup_part": random.randrange(config.partitioning.num_part_dedup),
1534 validity_start_column: timestamp,
1535 }
1536 self._storeObjectsPandas(objs, table, extra_columns=extra_columns)
1538 def _storeDiaSources(
1539 self,
1540 table_name: ApdbTables,
1541 sources: pandas.DataFrame,
1542 replica_chunk: ReplicaChunk | None,
1543 ) -> int | None:
1544 """Store catalog of DIASources or DIAForcedSources from current visit.
1546 Parameters
1547 ----------
1548 table_name : `ApdbTables`
1549 Table where to store the data.
1550 sources : `pandas.DataFrame`
1551 Catalog containing DiaSource records
1552 visit_time : `astropy.time.Time`
1553 Time of the current visit.
1554 replica_chunk : `ReplicaChunk` or `None`
1555 Replica chunk identifier if replication is configured.
1557 Returns
1558 -------
1559 subchunk : `int` or `None`
1560 Subchunk number for resulting replica data, `None` if relication is
1561 not enabled ot subchunking is not enabled.
1562 """
1563 context = self._context
1564 config = context.config
1566 # Time partitioning has to be based on midpointMjdTai, not visit_time
1567 # as visit_time is not really a visit time.
1568 tp_sources = sources.copy(deep=False)
1569 tp_sources["apdb_time_part"] = tp_sources["midpointMjdTai"].apply(context.partitioner.time_partition)
1570 if (time_partitions_range := context.time_partitions_range) is not None:
1571 self._check_time_partitions(tp_sources["apdb_time_part"], time_partitions_range)
1572 extra_columns: dict[str, Any] = {}
1573 if not config.partitioning.time_partition_tables:
1574 self._storeObjectsPandas(tp_sources, table_name)
1575 else:
1576 # Group by time partition
1577 partitions = set(tp_sources["apdb_time_part"])
1578 if len(partitions) == 1:
1579 # Single partition - just save the whole thing.
1580 time_part = partitions.pop()
1581 self._storeObjectsPandas(sources, table_name, time_part=time_part)
1582 else:
1583 # group by time partition.
1584 for time_part, sub_frame in tp_sources.groupby(by="apdb_time_part"):
1585 sub_frame.drop(columns="apdb_time_part", inplace=True)
1586 self._storeObjectsPandas(sub_frame, table_name, time_part=time_part)
1588 subchunk: int | None = None
1589 if replica_chunk is not None:
1590 extra_columns = {"apdb_replica_chunk": replica_chunk.id}
1591 if context.has_chunk_sub_partitions:
1592 subchunk = random.randrange(config.replica_sub_chunk_count)
1593 extra_columns["apdb_replica_subchunk"] = subchunk
1594 if table_name is ApdbTables.DiaSource:
1595 extra_table = ExtraTables.DiaSourceChunks2
1596 else:
1597 extra_table = ExtraTables.DiaForcedSourceChunks2
1598 else:
1599 if table_name is ApdbTables.DiaSource:
1600 extra_table = ExtraTables.DiaSourceChunks
1601 else:
1602 extra_table = ExtraTables.DiaForcedSourceChunks
1603 self._storeObjectsPandas(sources, extra_table, extra_columns=extra_columns)
1605 return subchunk
1607 def _check_time_partitions(
1608 self, partitions: Iterable[int], time_partitions_range: ApdbCassandraTimePartitionRange
1609 ) -> None:
1610 """Check that time partitons for new data actually exist.
1612 Parameters
1613 ----------
1614 partitions : `~collections.abc.Iterable` [`int`]
1615 Time partitions for new data.
1616 time_partitions_range : `ApdbCassandraTimePartitionRange`
1617 Currrent time partition range.
1618 """
1619 partitions = set(partitions)
1620 min_part = min(partitions)
1621 max_part = max(partitions)
1622 if min_part < time_partitions_range.start or max_part > time_partitions_range.end:
1623 raise ValueError(
1624 "Attempt to store data for time partitions that do not yet exist. "
1625 f"Partitons for new records: {min_part}-{max_part}. "
1626 f"Database partitons: {time_partitions_range.start}-{time_partitions_range.end}."
1627 )
1628 # Make a noise when writing to the last partition.
1629 if max_part == time_partitions_range.end:
1630 warnings.warn(
1631 "Writing into the last temporal partition. Partition range needs to be extended soon.",
1632 stacklevel=3,
1633 )
1635 def _storeDiaSourcesPartitions(
1636 self,
1637 sources: pandas.DataFrame,
1638 visit_time: astropy.time.Time,
1639 replica_chunk: ReplicaChunk | None,
1640 subchunk: int | None,
1641 ) -> None:
1642 """Store mapping of diaSourceId to its partitioning values.
1644 Parameters
1645 ----------
1646 sources : `pandas.DataFrame`
1647 Catalog containing DiaSource records
1648 visit_time : `astropy.time.Time`
1649 Time of the current visit.
1650 replica_chunk : `ReplicaChunk` or `None`
1651 Replication chunk, or `None` when replication is disabled.
1652 subchunk : `int` or `None`
1653 Replication sub-chunk, or `None` when replication is disabled or
1654 sub-chunking is not used.
1655 """
1656 context = self._context
1658 id_map = cast(pandas.DataFrame, sources[["diaSourceId", "apdb_part"]])
1659 extra_columns = {
1660 "apdb_time_part": context.partitioner.time_partition(visit_time),
1661 "apdb_replica_chunk": replica_chunk.id if replica_chunk is not None else None,
1662 }
1663 if context.has_chunk_sub_partitions:
1664 extra_columns["apdb_replica_subchunk"] = subchunk
1666 self._storeObjectsPandas(
1667 id_map, ExtraTables.DiaSourceToPartition, extra_columns=extra_columns, time_part=None
1668 )
1670 def _storeObjectsPandas(
1671 self,
1672 records: pandas.DataFrame,
1673 table_name: ApdbTables | ExtraTables,
1674 extra_columns: Mapping | None = None,
1675 time_part: int | None = None,
1676 ) -> None:
1677 """Store generic objects.
1679 Takes Pandas catalog and stores a bunch of records in a table.
1681 Parameters
1682 ----------
1683 records : `pandas.DataFrame`
1684 Catalog containing object records
1685 table_name : `ApdbTables`
1686 Name of the table as defined in APDB schema.
1687 extra_columns : `dict`, optional
1688 Mapping (column_name, column_value) which gives fixed values for
1689 columns in each row, overrides values in ``records`` if matching
1690 columns exist there.
1691 time_part : `int`, optional
1692 If not `None` then insert into a per-partition table.
1694 Notes
1695 -----
1696 If Pandas catalog contains additional columns not defined in table
1697 schema they are ignored. Catalog does not have to contain all columns
1698 defined in a table, but partition and clustering keys must be present
1699 in a catalog or ``extra_columns``.
1700 """
1701 context = self._context
1703 # use extra columns if specified
1704 if extra_columns is None:
1705 extra_columns = {}
1706 extra_fields = list(extra_columns.keys())
1708 # Fields that will come from dataframe.
1709 df_fields = [column for column in records.columns if column not in extra_fields]
1711 column_map = context.schema.getColumnMap(table_name)
1712 # list of columns (as in felis schema)
1713 fields = [column_map[field].name for field in df_fields if field in column_map]
1714 fields += extra_fields
1716 # check that all partitioning and clustering columns are defined
1717 partition_columns = context.schema.partitionColumns(table_name)
1718 required_columns = partition_columns + context.schema.clusteringColumns(table_name)
1719 missing_columns = [column for column in required_columns if column not in fields]
1720 if missing_columns:
1721 raise ValueError(f"Primary key columns are missing from catalog: {missing_columns}")
1723 batch_size = self._batch_size(table_name)
1725 with self._timer("insert_build_time", tags={"table": table_name.name}):
1726 # Multi-partition batches are problematic in general, so we want to
1727 # group records in a batch by their partition key.
1728 values_by_key: dict[tuple, list[list]] = defaultdict(list)
1729 for rec in records.itertuples(index=False):
1730 values = []
1731 partitioning_values: dict[str, Any] = {}
1732 for field in df_fields:
1733 if field not in column_map:
1734 continue
1735 value = getattr(rec, field)
1736 if column_map[field].datatype is felis.datamodel.DataType.timestamp:
1737 if isinstance(value, pandas.Timestamp):
1738 value = value.to_pydatetime()
1739 elif value is pandas.NaT:
1740 value = None
1741 else:
1742 # Assume it's seconds since epoch, Cassandra
1743 # datetime is in milliseconds
1744 value = int(value * 1000)
1745 value = literal(value)
1746 values.append(UNSET_VALUE if value is None else value)
1747 if field in partition_columns:
1748 partitioning_values[field] = value
1749 for field in extra_fields:
1750 value = literal(extra_columns[field])
1751 values.append(UNSET_VALUE if value is None else value)
1752 if field in partition_columns:
1753 partitioning_values[field] = value
1755 key = tuple(partitioning_values[field] for field in partition_columns)
1756 values_by_key[key].append(values)
1758 table = context.schema.tableName(table_name, time_part)
1760 query = Insert(self._keyspace, table, fields)
1761 statement = context.stmt_factory(query, prepare=True)
1762 # Cassandra has 64k limit on batch size, normally that should be
1763 # enough but some tests generate too many forced sources.
1764 queries = []
1765 for key_values in values_by_key.values():
1766 for values_chunk in chunk_iterable(key_values, batch_size):
1767 batch = cassandra.query.BatchStatement()
1768 for row_values in values_chunk:
1769 batch.add(statement, row_values)
1770 queries.append((batch, None))
1771 assert batch.routing_key is not None and batch.keyspace is not None
1773 _LOG.debug("%s: will store %d records", context.schema.tableName(table_name), records.shape[0])
1774 with self._timer(
1775 "insert_time", tags={"table": table_name.name, "method": "_storeObjectsPandas"}
1776 ) as timer:
1777 execute_concurrent(context.session, queries, execution_profile="write")
1778 timer.add_values(row_count=len(records), num_batches=len(queries))
1780 def _storeUpdateRecords(
1781 self, records: Iterable[ApdbUpdateRecord], chunk: ReplicaChunk, *, store_chunk: bool = False
1782 ) -> None:
1783 """Store ApdbUpdateRecords in the replica table for those records.
1785 Parameters
1786 ----------
1787 records : `list` [`ApdbUpdateRecord`]
1788 Records to store.
1789 chunk : `ReplicaChunk`
1790 Replica chunk for these records.
1791 store_chunk : `bool`
1792 If True then also store replica chunk.
1794 Raises
1795 ------
1796 TypeError
1797 Raised if replication is not enabled for this instance.
1798 """
1799 context = self._context
1800 config = context.config
1802 if not context.schema.replication_enabled:
1803 raise TypeError("Replication is not enabled for this APDB instance.")
1805 if store_chunk:
1806 self._storeReplicaChunk(chunk)
1808 apdb_replica_chunk = chunk.id
1809 # Do not use unique_if from ReplicaChunk as it could be reused in
1810 # multiple calls to this method.
1811 update_unique_id = uuid.uuid4()
1813 rows = []
1814 for record in records:
1815 rows.append(
1816 [
1817 apdb_replica_chunk,
1818 record.update_time_ns,
1819 record.update_order,
1820 update_unique_id,
1821 record.to_json(),
1822 ]
1823 )
1824 columns = [
1825 "apdb_replica_chunk",
1826 "update_time_ns",
1827 "update_order",
1828 "update_unique_id",
1829 "update_payload",
1830 ]
1831 if context.has_chunk_sub_partitions:
1832 subchunk = random.randrange(config.replica_sub_chunk_count)
1833 for row in rows:
1834 row.append(subchunk)
1835 columns.append("apdb_replica_subchunk")
1837 table_name = context.schema.tableName(ExtraTables.ApdbUpdateRecordChunks)
1838 query = Insert(self._keyspace, table_name, columns)
1839 stmt = context.stmt_factory(query)
1840 queries = [(stmt, row) for row in rows]
1842 with self._timer("store_update_record", tags={"table": table_name}) as timer:
1843 execute_concurrent(context.session, queries, execution_profile="write")
1844 timer.add_values(row_count=len(queries))
1846 def _add_apdb_part(self, df: pandas.DataFrame) -> pandas.DataFrame:
1847 """Calculate spatial partition for each record and add it to a
1848 DataFrame.
1850 Parameters
1851 ----------
1852 df : `pandas.DataFrame`
1853 DataFrame which has to contain ra/dec columns, names of these
1854 columns are defined by configuration ``ra_dec_columns`` field.
1856 Returns
1857 -------
1858 df : `pandas.DataFrame`
1859 DataFrame with ``apdb_part`` column which contains pixel index
1860 for ra/dec coordinates.
1862 Notes
1863 -----
1864 This overrides any existing column in a DataFrame with the same name
1865 (``apdb_part``). Original DataFrame is not changed, copy of a DataFrame
1866 is returned.
1867 """
1868 context = self._context
1869 config = context.config
1871 # Calculate pixelization index for every record.
1872 apdb_part = np.zeros(df.shape[0], dtype=np.int64)
1873 ra_col, dec_col = config.ra_dec_columns
1874 for i, (ra, dec) in enumerate(zip(df[ra_col], df[dec_col])):
1875 idx = context.partitioner.pixel(ra, dec)
1876 apdb_part[i] = idx
1877 df = df.copy()
1878 df["apdb_part"] = apdb_part
1879 return df
1881 def _make_empty_catalog(self, table_name: ApdbTables) -> pandas.DataFrame:
1882 """Make an empty catalog for a table with a given name.
1884 Parameters
1885 ----------
1886 table_name : `ApdbTables`
1887 Name of the table.
1889 Returns
1890 -------
1891 catalog : `pandas.DataFrame`
1892 An empty catalog.
1893 """
1894 table = self.schema.tableSchemas[table_name]
1896 data = {columnDef.name: pandas.Series(dtype=columnDef.pandas_type) for columnDef in table.columns}
1897 return pandas.DataFrame(data)
1899 def _fix_input_timestamps(self, df: pandas.DataFrame) -> pandas.DataFrame:
1900 """Update timestamp columns in input DataFrame to be naive datetime
1901 type.
1903 Clients may or may not generate aware timestamps, code in this class
1904 assumes that timestamps are naive, so we convert them to UTC and
1905 drop timezone.
1906 """
1907 # Find all columns with aware timestamps.
1908 columns = [column for column, dtype in df.dtypes.items() if isinstance(dtype, pandas.DatetimeTZDtype)]
1909 for column in columns:
1910 # tz_convert(None) will convert to UTC and drop timezone.
1911 df[column] = df[column].dt.tz_convert(None)
1912 return df
1914 def _batch_size(self, table: ApdbTables | ExtraTables) -> int:
1915 """Calculate batch size based on config parameters."""
1916 context = self._context
1917 config = context.config
1919 # Cassandra limit on number of statements in a batch is 64k.
1920 batch_size = 65_535
1921 if 0 < config.batch_statement_limit < batch_size:
1922 batch_size = config.batch_statement_limit
1923 if config.batch_size_limit > 0:
1924 # The purpose of this limit is to try not to exceed batch size
1925 # threshold which is set on server side. Cassandra wire protocol
1926 # for prepared queries (and batches) only sends column values with
1927 # with an additional 4 bytes per value specifying size. Value is
1928 # not included for NULL or NOT_SET values, but the size is always
1929 # there. There is additional small per-query overhead, which we
1930 # ignore.
1931 row_size = context.schema.table_row_size(table)
1932 row_size += 4 * len(context.schema.getColumnMap(table))
1933 batch_size = min(batch_size, (config.batch_size_limit // row_size) + 1)
1934 return batch_size
1936 def _group_dia_objects_by_partition(
1937 self, partitioner: Partitioner, objects: list[DiaObjectId], pad_arcsec: float
1938 ) -> Mapping[int, list[int]]:
1939 """Group DiaObjects by partition.
1941 Parameters
1942 ----------
1943 partitioner : `Partitioner`
1944 Objects which knows how to partition things.
1945 objects : `list` [`DiaObjectId`]
1946 Collection of objects to partition.
1947 pad_arcsec : `float`
1948 Additional padding around object position.
1950 Returns
1951 -------
1952 grouped_objects
1953 Mapping of spatial patition ID to list ob object IDs that it
1954 contains. Some objects may belong to more than one partition.
1955 """
1956 partitioned_object_ids: dict[int, list[int]] = defaultdict(list)
1957 for obj_id in objects:
1958 partitions = partitioner.pixelization.circle_pixels(obj_id.ra, obj_id.dec, pad_arcsec)
1959 for pixel in partitions:
1960 partitioned_object_ids[pixel].append(obj_id.diaObjectId)
1961 return partitioned_object_ids
1963 def _timestamp_column_name(self, column: str) -> str:
1964 """Return column name before/after schema migration to MJD TAI."""
1965 return self._schema.timestamp_column_name(column)
1967 def _timestamp_column_value(self, time: astropy.time.Time) -> float | int:
1968 """Return column value before/after schema migration to MJD TAI."""
1969 if self._schema.has_mjd_timestamps:
1970 return float(time.tai.mjd)
1971 else:
1972 return int(time.datetime.astimezone(tz=datetime.UTC).timestamp() * 1000)
1974 def _get_diasource_data(self, source_ids: Iterable[DiaSourceId], *columns: str) -> list:
1975 """Select records from DiaSource table by diaSourceId and return all
1976 records as a list of named tuples.
1977 """
1978 context = self._context
1979 config = context.config
1980 partitioner = context.partitioner
1982 columns = ("diaSourceId",) + columns
1984 # Allow some uncertainty for coordinates and time when calculating
1985 # partitions.
1986 statements: list[tuple] = []
1987 pad_arcsec = 1.0
1988 pad_time_day = 10 / (24 * 3600)
1989 for source_id in source_ids:
1990 center = sphgeom.UnitVector3d(sphgeom.LonLat.fromDegrees(source_id.ra, source_id.dec))
1991 region = sphgeom.Circle(center, sphgeom.Angle.fromDegrees(pad_arcsec / 3600.0))
1992 spatial_where, _ = partitioner.spatial_where(region)
1994 tables, temporal_where = partitioner.temporal_where(
1995 ApdbTables.DiaSource,
1996 source_id.midpointMjdTai - pad_time_day,
1997 source_id.midpointMjdTai + pad_time_day,
1998 partitons_range=context.time_partitions_range,
1999 query_per_time_part=True,
2000 )
2002 id_where = QExpr('"diaSourceId" = {}', (source_id.diaSourceId,))
2004 for table in tables:
2005 query = Select(self._keyspace, table, columns)
2006 for clause in QExpr.combine(spatial_where, temporal_where, extra=id_where):
2007 statements.append(context.stmt_factory.with_params(query.where(clause), prepare=True))
2009 with self._timer(
2010 "select_time", tags={"table": "DiaSource", "method": "_get_diasource_data"}
2011 ) as timer:
2012 result = cast(
2013 list[tuple],
2014 select_concurrent(
2015 context.session,
2016 statements,
2017 "read_named_tuples",
2018 config.connection_config.read_concurrency,
2019 ),
2020 )
2021 timer.add_values(row_count=len(result), num_queries=len(statements))
2023 return result
2025 def _get_diaforcedsource_data(self, source_ids: Iterable[DiaForcedSourceId], *columns: str) -> list:
2026 """Select records from DiaForcedSource table by (diaObjectId, visit,
2027 detector) and return all records as a list of named tuples.
2028 """
2029 context = self._context
2030 config = context.config
2031 partitioner = context.partitioner
2033 columns = ("diaObjectId", "visit", "detector") + columns
2035 # Allow some uncertainty for coordinates and time when calculating
2036 # partitions.
2037 statements: list[tuple] = []
2038 pad_arcsec = 1.0
2039 pad_time_day = 10 / (24 * 3600)
2040 for source_id in source_ids:
2041 center = sphgeom.UnitVector3d(sphgeom.LonLat.fromDegrees(source_id.ra, source_id.dec))
2042 region = sphgeom.Circle(center, sphgeom.Angle.fromDegrees(pad_arcsec / 3600.0))
2043 spatial_where, _ = partitioner.spatial_where(region)
2045 tables, temporal_where = partitioner.temporal_where(
2046 ApdbTables.DiaForcedSource,
2047 source_id.midpointMjdTai - pad_time_day,
2048 source_id.midpointMjdTai + pad_time_day,
2049 partitons_range=context.time_partitions_range,
2050 query_per_time_part=True,
2051 )
2053 id_where = (
2054 (C("diaObjectId") == source_id.diaObjectId)
2055 & (C("visit") == source_id.visit)
2056 & (C("detector") == source_id.detector)
2057 )
2059 for table in tables:
2060 query = Select(self._keyspace, table, columns)
2061 for clause in QExpr.combine(spatial_where, temporal_where, extra=id_where):
2062 statements.append(context.stmt_factory.with_params(query.where(clause), prepare=True))
2064 with self._timer(
2065 "select_time", tags={"table": "DiaForcedSource", "method": "_get_diaforcedsource_data"}
2066 ) as timer:
2067 result = cast(
2068 list[tuple],
2069 select_concurrent(
2070 context.session,
2071 statements,
2072 "read_named_tuples",
2073 config.connection_config.read_concurrency,
2074 ),
2075 )
2076 timer.add_values(row_count=len(result), num_queries=len(statements))
2078 return result
2080 def _get_diaobject_data(self, object_ids: Iterable[DiaObjectId], *columns: str) -> list:
2081 """Select records from DiaObjectLast table by diaObjectId and return
2082 all records as a list of named tuples.
2083 """
2084 context = self._context
2085 config = context.config
2086 partitioner = context.partitioner
2088 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
2089 columns = ("diaObjectId",) + columns
2091 # Allow some uncertainty for coordinates when calculating partitions.
2092 pad_arcsec = 1.0
2093 ids_by_partition = defaultdict(list)
2094 for object_id in object_ids:
2095 pixels = partitioner.pixelization.circle_pixels(object_id.ra, object_id.dec, pad_arcsec)
2096 for pixel in pixels:
2097 ids_by_partition[pixel].append(object_id.diaObjectId)
2099 statements: list[tuple] = []
2100 for apdb_part, diaObjectIds in ids_by_partition.items():
2101 query = Select(self._keyspace, table_name, columns)
2102 query = query.where(C("apdb_part") == apdb_part)
2103 query = query.where(C("diaObjectId").in_(diaObjectIds))
2104 statements.append(context.stmt_factory.with_params(query, prepare=False))
2106 with self._timer(
2107 "select_time", tags={"table": "DiaObjectLast", "method": "_get_diaobject_data"}
2108 ) as timer:
2109 result = cast(
2110 list[tuple],
2111 select_concurrent(
2112 context.session,
2113 statements,
2114 "read_named_tuples",
2115 config.connection_config.read_concurrency,
2116 ),
2117 )
2118 timer.add_values(row_count=len(result), num_queries=len(statements))
2120 return result
2122 def _fix_null_columns(self, table: ApdbTables, df: pandas.DataFrame) -> None:
2123 """For non-null columns replace NULL values with column defaults.
2125 When new non-null columns are added to the schema we do not populate
2126 values for existing records. Instead when reading old data back we
2127 fill NULLs with the defaults. The dataflame is updated in-place.
2128 """
2129 # Find all non-null columns which also have default values defined.
2130 table_schema = self._schema.tableSchemas[table]
2131 defaults = {
2132 column.name: column.value
2133 for column in table_schema.columns
2134 if not column.nullable and column.value is not None
2135 }
2136 if defaults:
2137 df.fillna(defaults, inplace=True)