22from __future__
import annotations
24__all__ = [
"ApdbCassandraConfig",
"ApdbCassandra"]
30from typing
import TYPE_CHECKING, Any, cast
39 import cassandra.query
40 from cassandra.auth
import AuthProvider, PlainTextAuthProvider
41 from cassandra.cluster
import EXEC_PROFILE_DEFAULT, Cluster, ExecutionProfile, Session
42 from cassandra.policies
import AddressTranslator, RoundRobinPolicy, WhiteListRoundRobinPolicy
44 CASSANDRA_IMPORTED =
True
46 CASSANDRA_IMPORTED =
False
50from lsst
import sphgeom
52from lsst.utils.db_auth
import DbAuth, DbAuthNotFoundError
53from lsst.utils.iteration
import chunk_iterable
55from ..apdb
import Apdb, ApdbConfig
56from ..apdbConfigFreezer
import ApdbConfigFreezer
57from ..apdbReplica
import ReplicaChunk
58from ..apdbSchema
import ApdbTables
59from ..pixelization
import Pixelization
60from ..schema_model
import Table
61from ..timer
import Timer
62from ..versionTuple
import IncompatibleVersionError, VersionTuple
63from .apdbCassandraReplica
import ApdbCassandraReplica
64from .apdbCassandraSchema
import ApdbCassandraSchema, ExtraTables
65from .apdbMetadataCassandra
import ApdbMetadataCassandra
66from .cassandra_utils
import (
67 PreparedStatementCache,
69 pandas_dataframe_factory,
76 from ..apdbMetadata
import ApdbMetadata
78_LOG = logging.getLogger(__name__)
81"""Version for the code controlling non-replication tables. This needs to be
82updated following compatibility rules when schema produced by this code
87DB_AUTH_ENVVAR =
"LSST_DB_AUTH"
88"""Default name of the environmental variable that will be used to locate DB
89credentials configuration file. """
91DB_AUTH_PATH =
"~/.lsst/db-auth.yaml"
92"""Default path at which it is expected that DB credentials are found."""
97 super().
__init__(
"cassandra-driver module cannot be imported")
101 """Configuration class for Cassandra-based APDB implementation."""
103 contact_points = ListField[str](
104 doc=
"The list of contact points to try connecting for cluster discovery.", default=[
"127.0.0.1"]
106 private_ips = ListField[str](doc=
"List of internal IP addresses for contact_points.", default=[])
107 port = Field[int](doc=
"Port number to connect to.", default=9042)
108 keyspace = Field[str](doc=
"Default keyspace for operations.", default=
"apdb")
109 username = Field[str](
110 doc=f
"Cassandra user name, if empty then {DB_AUTH_PATH} has to provide it with password.",
113 read_consistency = Field[str](
114 doc=
"Name for consistency level of read operations, default: QUORUM, can be ONE.", default=
"QUORUM"
116 write_consistency = Field[str](
117 doc=
"Name for consistency level of write operations, default: QUORUM, can be ONE.", default=
"QUORUM"
119 read_timeout = Field[float](doc=
"Timeout in seconds for read operations.", default=120.0)
120 write_timeout = Field[float](doc=
"Timeout in seconds for write operations.", default=10.0)
121 remove_timeout = Field[float](doc=
"Timeout in seconds for remove operations.", default=600.0)
122 read_concurrency = Field[int](doc=
"Concurrency level for read operations.", default=500)
123 protocol_version = Field[int](
124 doc=
"Cassandra protocol version to use, default is V4",
125 default=cassandra.ProtocolVersion.V4
if CASSANDRA_IMPORTED
else 0,
127 dia_object_columns = ListField[str](
128 doc=
"List of columns to read from DiaObject[Last], by default read all columns", default=[]
130 prefix = Field[str](doc=
"Prefix to add to table names", default=
"")
131 part_pixelization = ChoiceField[str](
132 allowed=dict(htm=
"HTM pixelization", q3c=
"Q3C pixelization", mq3c=
"MQ3C pixelization"),
133 doc=
"Pixelization used for partitioning index.",
136 part_pix_level = Field[int](doc=
"Pixelization level used for partitioning index.", default=10)
137 part_pix_max_ranges = Field[int](doc=
"Max number of ranges in pixelization envelope", default=64)
138 ra_dec_columns = ListField[str](default=[
"ra",
"dec"], doc=
"Names of ra/dec columns in DiaObject table")
139 timer = Field[bool](doc=
"If True then print/log timing information", default=
False)
140 time_partition_tables = Field[bool](
141 doc=
"Use per-partition tables for sources instead of partitioning by time", default=
False
143 time_partition_days = Field[int](
145 "Time partitioning granularity in days, this value must not be changed after database is "
150 time_partition_start = Field[str](
152 "Starting time for per-partition tables, in yyyy-mm-ddThh:mm:ss format, in TAI. "
153 "This is used only when time_partition_tables is True."
155 default=
"2018-12-01T00:00:00",
157 time_partition_end = Field[str](
159 "Ending time for per-partition tables, in yyyy-mm-ddThh:mm:ss format, in TAI. "
160 "This is used only when time_partition_tables is True."
162 default=
"2030-01-01T00:00:00",
164 query_per_time_part = Field[bool](
167 "If True then build separate query for each time partition, otherwise build one single query. "
168 "This is only used when time_partition_tables is False in schema config."
171 query_per_spatial_part = Field[bool](
173 doc=
"If True then build one query per spatial partition, otherwise build single query.",
175 use_insert_id_skips_diaobjects = Field[bool](
178 "If True then do not store DiaObjects when use_insert_id is True "
179 "(DiaObjectsChunks has the same data)."
184@dataclasses.dataclass
186 """Part of the configuration that is saved in metadata table and read back.
188 The attributes are a subset of attributes in `ApdbCassandraConfig` class.
192 config : `ApdbSqlConfig`
193 Configuration used to copy initial values of attributes.
197 part_pixelization: str
199 ra_dec_columns: list[str]
200 time_partition_tables: bool
201 time_partition_days: int
202 use_insert_id_skips_diaobjects: bool
214 """Convert this instance to JSON representation."""
215 return json.dumps(dataclasses.asdict(self))
218 """Update attribute values from a JSON string.
223 String containing JSON representation of configuration.
225 data = json.loads(json_str)
226 if not isinstance(data, dict):
227 raise TypeError(f
"JSON string must be convertible to object: {json_str!r}")
228 allowed_names = {field.name
for field
in dataclasses.fields(self)}
229 for key, value
in data.items():
230 if key
not in allowed_names:
231 raise ValueError(f
"JSON object contains unknown key: {key}")
232 setattr(self, key, value)
235if CASSANDRA_IMPORTED:
238 """Translate internal IP address to external.
240 Only used for docker-based setup, not viable long-term solution.
243 def __init__(self, public_ips: list[str], private_ips: list[str]):
244 self.
_map = dict((k, v)
for k, v
in zip(private_ips, public_ips))
247 return self.
_map.get(private_ip, private_ip)
251 """Implementation of APDB database on to of Apache Cassandra.
253 The implementation is configured via standard ``pex_config`` mechanism
254 using `ApdbCassandraConfig` configuration class. For an example of
255 different configurations check config/ folder.
259 config : `ApdbCassandraConfig`
260 Configuration object.
263 metadataSchemaVersionKey =
"version:schema"
264 """Name of the metadata key to store schema version number."""
266 metadataCodeVersionKey =
"version:ApdbCassandra"
267 """Name of the metadata key to store code version number."""
269 metadataReplicaVersionKey =
"version:ApdbCassandraReplica"
270 """Name of the metadata key to store replica code version number."""
272 metadataConfigKey =
"config:apdb-cassandra.json"
273 """Name of the metadata key to store code version number."""
275 _frozen_parameters = (
280 "time_partition_tables",
281 "time_partition_days",
282 "use_insert_id_skips_diaobjects",
284 """Names of the config parameters to be frozen in metadata table."""
286 partition_zero_epoch = astropy.time.Time(0, format=
"unix_tai")
287 """Start time for partition 0, this should never be changed."""
290 if not CASSANDRA_IMPORTED:
297 meta_table_name = ApdbTables.metadata.table_name(config.prefix)
299 self.
_session, meta_table_name, config.keyspace,
"read_tuples",
"write"
304 if config_json
is not None:
307 self.
config = freezer.update(config, config_json)
313 self.
config.part_pixelization,
314 self.
config.part_pix_level,
315 config.part_pix_max_ranges,
321 schema_file=self.
config.schema_file,
322 schema_name=self.
config.schema_name,
323 prefix=self.
config.prefix,
324 time_partition_tables=self.
config.time_partition_tables,
325 enable_replica=self.
config.use_insert_id,
335 _LOG.debug(
"ApdbCassandra Configuration:")
337 _LOG.debug(
" %s: %s", key, value)
340 if hasattr(self,
"_cluster"):
344 def _make_session(cls, config: ApdbCassandraConfig) -> tuple[Cluster, Session]:
345 """Make Cassandra session."""
346 addressTranslator: AddressTranslator |
None =
None
347 if config.private_ips:
348 addressTranslator =
_AddressTranslator(list(config.contact_points), list(config.private_ips))
352 contact_points=config.contact_points,
354 address_translator=addressTranslator,
355 protocol_version=config.protocol_version,
358 session = cluster.connect()
360 session.default_fetch_size =
None
362 return cluster, session
366 """Make Cassandra authentication provider instance."""
368 dbauth = DbAuth(DB_AUTH_PATH, DB_AUTH_ENVVAR)
369 except DbAuthNotFoundError:
373 empty_username =
True
375 for hostname
in config.contact_points:
377 username, password = dbauth.getAuth(
378 "cassandra", config.username, hostname, config.port, config.keyspace
383 empty_username =
True
385 return PlainTextAuthProvider(username=username, password=password)
386 except DbAuthNotFoundError:
391 f
"Credentials file ({DB_AUTH_PATH} or ${DB_AUTH_ENVVAR}) provided password but not "
392 f
"user name, anonymous Cassandra logon will be attempted."
398 """Check schema version compatibility."""
400 def _get_version(key: str, default: VersionTuple) -> VersionTuple:
401 """Retrieve version number from given metadata key."""
402 if metadata.table_exists():
403 version_str = metadata.get(key)
404 if version_str
is None:
406 raise RuntimeError(f
"Version key {key!r} does not exist in metadata table.")
407 return VersionTuple.fromString(version_str)
418 if not self.
_schema.schemaVersion().checkCompatibility(db_schema_version,
True):
420 f
"Configured schema version {self._schema.schemaVersion()} "
421 f
"is not compatible with database version {db_schema_version}"
425 f
"Current code version {self.apdbImplementationVersion()} "
426 f
"is not compatible with database version {db_code_version}"
430 if self.
_schema.has_replica_chunks:
432 code_replica_version = ApdbCassandraReplica.apdbReplicaImplementationVersion()
433 if not code_replica_version.checkCompatibility(db_replica_version,
True):
435 f
"Current replication code version {code_replica_version} "
436 f
"is not compatible with database version {db_replica_version}"
446 return self.
_schema.schemaVersion()
448 def tableDef(self, table: ApdbTables) -> Table |
None:
450 return self.
_schema.tableSchemas.get(table)
458 schema_file: str |
None =
None,
459 schema_name: str |
None =
None,
460 read_sources_months: int |
None =
None,
461 read_forced_sources_months: int |
None =
None,
462 use_insert_id: bool =
False,
463 replica_skips_diaobjects: bool =
False,
464 port: int |
None =
None,
465 username: str |
None =
None,
466 prefix: str |
None =
None,
467 part_pixelization: str |
None =
None,
468 part_pix_level: int |
None =
None,
469 time_partition_tables: bool =
True,
470 time_partition_start: str |
None =
None,
471 time_partition_end: str |
None =
None,
472 read_consistency: str |
None =
None,
473 write_consistency: str |
None =
None,
474 read_timeout: int |
None =
None,
475 write_timeout: int |
None =
None,
476 ra_dec_columns: list[str] |
None =
None,
477 replication_factor: int |
None =
None,
479 ) -> ApdbCassandraConfig:
480 """Initialize new APDB instance and make configuration object for it.
484 hosts : `list` [`str`]
485 List of host names or IP addresses for Cassandra cluster.
487 Name of the keyspace for APDB tables.
488 schema_file : `str`, optional
489 Location of (YAML) configuration file with APDB schema. If not
490 specified then default location will be used.
491 schema_name : `str`, optional
492 Name of the schema in YAML configuration file. If not specified
493 then default name will be used.
494 read_sources_months : `int`, optional
495 Number of months of history to read from DiaSource.
496 read_forced_sources_months : `int`, optional
497 Number of months of history to read from DiaForcedSource.
498 use_insert_id : `bool`, optional
499 If True, make additional tables used for replication to PPDB.
500 replica_skips_diaobjects : `bool`, optional
501 If `True` then do not fill regular ``DiaObject`` table when
502 ``use_insert_id`` is `True`.
503 port : `int`, optional
504 Port number to use for Cassandra connections.
505 username : `str`, optional
506 User name for Cassandra connections.
507 prefix : `str`, optional
508 Optional prefix for all table names.
509 part_pixelization : `str`, optional
510 Name of the MOC pixelization used for partitioning.
511 part_pix_level : `int`, optional
513 time_partition_tables : `bool`, optional
514 Create per-partition tables.
515 time_partition_start : `str`, optional
516 Starting time for per-partition tables, in yyyy-mm-ddThh:mm:ss
518 time_partition_end : `str`, optional
519 Ending time for per-partition tables, in yyyy-mm-ddThh:mm:ss
521 read_consistency : `str`, optional
522 Name of the consistency level for read operations.
523 write_consistency : `str`, optional
524 Name of the consistency level for write operations.
525 read_timeout : `int`, optional
526 Read timeout in seconds.
527 write_timeout : `int`, optional
528 Write timeout in seconds.
529 ra_dec_columns : `list` [`str`], optional
530 Names of ra/dec columns in DiaObject table.
531 replication_factor : `int`, optional
532 Replication factor used when creating new keyspace, if keyspace
533 already exists its replication factor is not changed.
534 drop : `bool`, optional
535 If `True` then drop existing tables before re-creating the schema.
539 config : `ApdbCassandraConfig`
540 Resulting configuration object for a created APDB instance.
543 contact_points=hosts,
545 use_insert_id=use_insert_id,
546 use_insert_id_skips_diaobjects=replica_skips_diaobjects,
547 time_partition_tables=time_partition_tables,
549 if schema_file
is not None:
550 config.schema_file = schema_file
551 if schema_name
is not None:
552 config.schema_name = schema_name
553 if read_sources_months
is not None:
554 config.read_sources_months = read_sources_months
555 if read_forced_sources_months
is not None:
556 config.read_forced_sources_months = read_forced_sources_months
559 if username
is not None:
560 config.username = username
561 if prefix
is not None:
562 config.prefix = prefix
563 if part_pixelization
is not None:
564 config.part_pixelization = part_pixelization
565 if part_pix_level
is not None:
566 config.part_pix_level = part_pix_level
567 if time_partition_start
is not None:
568 config.time_partition_start = time_partition_start
569 if time_partition_end
is not None:
570 config.time_partition_end = time_partition_end
571 if read_consistency
is not None:
572 config.read_consistency = read_consistency
573 if write_consistency
is not None:
574 config.write_consistency = write_consistency
575 if read_timeout
is not None:
576 config.read_timeout = read_timeout
577 if write_timeout
is not None:
578 config.write_timeout = write_timeout
579 if ra_dec_columns
is not None:
580 config.ra_dec_columns = ra_dec_columns
582 cls.
_makeSchema(config, drop=drop, replication_factor=replication_factor)
587 """Return `ApdbReplica` instance for this database."""
594 cls, config: ApdbConfig, *, drop: bool =
False, replication_factor: int |
None =
None
598 if not isinstance(config, ApdbCassandraConfig):
599 raise TypeError(f
"Unexpected type of configuration object: {type(config)}")
605 keyspace=config.keyspace,
606 schema_file=config.schema_file,
607 schema_name=config.schema_name,
608 prefix=config.prefix,
609 time_partition_tables=config.time_partition_tables,
610 enable_replica=config.use_insert_id,
614 if config.time_partition_tables:
615 time_partition_start = astropy.time.Time(config.time_partition_start, format=
"isot", scale=
"tai")
616 time_partition_end = astropy.time.Time(config.time_partition_end, format=
"isot", scale=
"tai")
618 part_days = config.time_partition_days
623 schema.makeSchema(drop=drop, part_range=part_range, replication_factor=replication_factor)
625 schema.makeSchema(drop=drop, replication_factor=replication_factor)
627 meta_table_name = ApdbTables.metadata.table_name(config.prefix)
631 if metadata.table_exists():
635 if config.use_insert_id:
639 str(ApdbCassandraReplica.apdbReplicaImplementationVersion()),
653 _LOG.debug(
"getDiaObjects: #partitions: %s", len(sp_where))
656 column_names = self.
_schema.apdbColumnNames(ApdbTables.DiaObjectLast)
657 what =
",".join(quote_id(column)
for column
in column_names)
659 table_name = self.
_schema.tableName(ApdbTables.DiaObjectLast)
660 query = f
'SELECT {what} from "{self._keyspace}"."{table_name}"'
661 statements: list[tuple] = []
662 for where, params
in sp_where:
663 full_query = f
"{query} WHERE {where}"
665 statement = self.
_preparer.prepare(full_query)
670 statement = cassandra.query.SimpleStatement(full_query)
671 statements.append((statement, params))
672 _LOG.debug(
"getDiaObjects: #queries: %s", len(statements))
678 self.
_session, statements,
"read_pandas_multi", self.
config.read_concurrency
682 _LOG.debug(
"found %s DiaObjects", objects.shape[0])
686 self, region: sphgeom.Region, object_ids: Iterable[int] |
None, visit_time: astropy.time.Time
687 ) -> pandas.DataFrame |
None:
689 months = self.
config.read_sources_months
692 mjd_end = visit_time.mjd
693 mjd_start = mjd_end - months * 30
695 return self.
_getSources(region, object_ids, mjd_start, mjd_end, ApdbTables.DiaSource)
698 self, region: sphgeom.Region, object_ids: Iterable[int] |
None, visit_time: astropy.time.Time
699 ) -> pandas.DataFrame |
None:
701 months = self.
config.read_forced_sources_months
704 mjd_end = visit_time.mjd
705 mjd_start = mjd_end - months * 30
707 return self.
_getSources(region, object_ids, mjd_start, mjd_end, ApdbTables.DiaForcedSource)
711 raise NotImplementedError()
715 tableName = self.
_schema.tableName(ApdbTables.SSObject)
716 query = f
'SELECT * from "{self._keyspace}"."{tableName}"'
720 result = self.
_session.execute(query, execution_profile=
"read_pandas")
721 objects = result._current_rows
723 _LOG.debug(
"found %s DiaObjects", objects.shape[0])
728 visit_time: astropy.time.Time,
729 objects: pandas.DataFrame,
730 sources: pandas.DataFrame |
None =
None,
731 forced_sources: pandas.DataFrame |
None =
None,
735 replica_chunk: ReplicaChunk |
None =
None
736 if self.
_schema.has_replica_chunks:
737 replica_chunk = ReplicaChunk.make_replica_chunk(visit_time, self.
config.replica_chunk_seconds)
744 if sources
is not None:
747 self.
_storeDiaSources(ApdbTables.DiaSource, sources, visit_time, replica_chunk)
750 if forced_sources
is not None:
752 self.
_storeDiaSources(ApdbTables.DiaForcedSource, forced_sources, visit_time, replica_chunk)
765 table_name = self.
_schema.tableName(ExtraTables.DiaSourceToPartition)
767 selects: list[tuple] = []
768 for ids
in chunk_iterable(idMap.keys(), 1_000):
769 ids_str =
",".join(str(item)
for item
in ids)
773 'SELECT "diaSourceId", "apdb_part", "apdb_time_part", "apdb_replica_chunk" '
774 f
'FROM "{self._keyspace}"."{table_name}" WHERE "diaSourceId" IN ({ids_str})'
782 list[tuple[int, int, int, int |
None]],
783 select_concurrent(self.
_session, selects,
"read_tuples", self.
config.read_concurrency),
787 id2partitions: dict[int, tuple[int, int]] = {}
788 id2chunk_id: dict[int, int] = {}
790 id2partitions[row[0]] = row[1:3]
791 if row[3]
is not None:
792 id2chunk_id[row[0]] = row[3]
795 if set(id2partitions) !=
set(idMap):
796 missing =
",".join(str(item)
for item
in set(idMap) -
set(id2partitions))
797 raise ValueError(f
"Following DiaSource IDs do not exist in the database: {missing}")
800 queries = cassandra.query.BatchStatement()
801 table_name = self.
_schema.tableName(ApdbTables.DiaSource)
802 for diaSourceId, ssObjectId
in idMap.items():
803 apdb_part, apdb_time_part = id2partitions[diaSourceId]
805 if self.
config.time_partition_tables:
807 f
'UPDATE "{self._keyspace}"."{table_name}_{apdb_time_part}"'
808 ' SET "ssObjectId" = ?, "diaObjectId" = NULL'
809 ' WHERE "apdb_part" = ? AND "diaSourceId" = ?'
811 values = (ssObjectId, apdb_part, diaSourceId)
814 f
'UPDATE "{self._keyspace}"."{table_name}"'
815 ' SET "ssObjectId" = ?, "diaObjectId" = NULL'
816 ' WHERE "apdb_part" = ? AND "apdb_time_part" = ? AND "diaSourceId" = ?'
818 values = (ssObjectId, apdb_part, apdb_time_part, diaSourceId)
819 queries.add(self.
_preparer.prepare(query), values)
827 if replica_chunks := self.
get_replica().getReplicaChunks():
828 known_ids =
set(replica_chunk.id
for replica_chunk
in replica_chunks)
829 id2chunk_id = {key: value
for key, value
in id2chunk_id.items()
if value
in known_ids}
831 table_name = self.
_schema.tableName(ExtraTables.DiaSourceChunks)
832 for diaSourceId, ssObjectId
in idMap.items():
833 if replica_chunk := id2chunk_id.get(diaSourceId):
835 f
'UPDATE "{self._keyspace}"."{table_name}" '
836 ' SET "ssObjectId" = ?, "diaObjectId" = NULL '
837 'WHERE "apdb_replica_chunk" = ? AND "diaSourceId" = ?'
839 values = (ssObjectId, replica_chunk, diaSourceId)
840 queries.add(self.
_preparer.prepare(query), values)
842 _LOG.debug(
"%s: will update %d records", table_name, len(idMap))
843 with Timer(table_name +
" update", self.
config.timer):
844 self.
_session.execute(queries, execution_profile=
"write")
854 raise NotImplementedError()
860 raise RuntimeError(
"Database schema was not initialized.")
864 def _makeProfiles(cls, config: ApdbCassandraConfig) -> Mapping[Any, ExecutionProfile]:
865 """Make all execution profiles used in the code."""
866 if config.private_ips:
867 loadBalancePolicy = WhiteListRoundRobinPolicy(hosts=config.contact_points)
869 loadBalancePolicy = RoundRobinPolicy()
871 read_tuples_profile = ExecutionProfile(
872 consistency_level=getattr(cassandra.ConsistencyLevel, config.read_consistency),
873 request_timeout=config.read_timeout,
874 row_factory=cassandra.query.tuple_factory,
875 load_balancing_policy=loadBalancePolicy,
877 read_pandas_profile = ExecutionProfile(
878 consistency_level=getattr(cassandra.ConsistencyLevel, config.read_consistency),
879 request_timeout=config.read_timeout,
880 row_factory=pandas_dataframe_factory,
881 load_balancing_policy=loadBalancePolicy,
883 read_raw_profile = ExecutionProfile(
884 consistency_level=getattr(cassandra.ConsistencyLevel, config.read_consistency),
885 request_timeout=config.read_timeout,
886 row_factory=raw_data_factory,
887 load_balancing_policy=loadBalancePolicy,
890 read_pandas_multi_profile = ExecutionProfile(
891 consistency_level=getattr(cassandra.ConsistencyLevel, config.read_consistency),
892 request_timeout=config.read_timeout,
893 row_factory=pandas_dataframe_factory,
894 load_balancing_policy=loadBalancePolicy,
898 read_raw_multi_profile = ExecutionProfile(
899 consistency_level=getattr(cassandra.ConsistencyLevel, config.read_consistency),
900 request_timeout=config.read_timeout,
901 row_factory=raw_data_factory,
902 load_balancing_policy=loadBalancePolicy,
904 write_profile = ExecutionProfile(
905 consistency_level=getattr(cassandra.ConsistencyLevel, config.write_consistency),
906 request_timeout=config.write_timeout,
907 load_balancing_policy=loadBalancePolicy,
910 default_profile = ExecutionProfile(
911 load_balancing_policy=loadBalancePolicy,
914 "read_tuples": read_tuples_profile,
915 "read_pandas": read_pandas_profile,
916 "read_raw": read_raw_profile,
917 "read_pandas_multi": read_pandas_multi_profile,
918 "read_raw_multi": read_raw_multi_profile,
919 "write": write_profile,
920 EXEC_PROFILE_DEFAULT: default_profile,
925 region: sphgeom.Region,
926 object_ids: Iterable[int] |
None,
929 table_name: ApdbTables,
930 ) -> pandas.DataFrame:
931 """Return catalog of DiaSource instances given set of DiaObject IDs.
935 region : `lsst.sphgeom.Region`
938 Collection of DiaObject IDs
940 Lower bound of time interval.
942 Upper bound of time interval.
943 table_name : `ApdbTables`
948 catalog : `pandas.DataFrame`, or `None`
949 Catalog containing DiaSource records. Empty catalog is returned if
950 ``object_ids`` is empty.
952 object_id_set: Set[int] =
set()
953 if object_ids
is not None:
954 object_id_set =
set(object_ids)
955 if len(object_id_set) == 0:
959 tables, temporal_where = self.
_temporal_where(table_name, mjd_start, mjd_end)
962 column_names = self.
_schema.apdbColumnNames(table_name)
963 what =
",".join(quote_id(column)
for column
in column_names)
966 statements: list[tuple] = []
968 prefix = f
'SELECT {what} from "{self._keyspace}"."{table}"'
969 statements += list(self.
_combine_where(prefix, sp_where, temporal_where))
970 _LOG.debug(
"_getSources %s: #queries: %s", table_name, len(statements))
972 with Timer(table_name.name +
" select", self.
config.timer):
976 self.
_session, statements,
"read_pandas_multi", self.
config.read_concurrency
981 if len(object_id_set) > 0:
982 catalog = cast(pandas.DataFrame, catalog[catalog[
"diaObjectId"].isin(object_id_set)])
985 catalog = cast(pandas.DataFrame, catalog[catalog[
"midpointMjdTai"] > mjd_start])
987 _LOG.debug(
"found %d %ss", catalog.shape[0], table_name.name)
992 timestamp = int(replica_chunk.last_update_time.unix_tai * 1000)
997 table_name = self.
_schema.tableName(ExtraTables.ApdbReplicaChunks)
999 f
'INSERT INTO "{self._keyspace}"."{table_name}" '
1000 "(partition, apdb_replica_chunk, last_update_time, unique_id) "
1001 "VALUES (?, ?, ?, ?)"
1006 (partition, replica_chunk.id, timestamp, replica_chunk.unique_id),
1007 timeout=self.
config.write_timeout,
1008 execution_profile=
"write",
1012 self, objs: pandas.DataFrame, visit_time: astropy.time.Time, replica_chunk: ReplicaChunk |
None
1014 """Store catalog of DiaObjects from current visit.
1018 objs : `pandas.DataFrame`
1019 Catalog with DiaObject records
1020 visit_time : `astropy.time.Time`
1021 Time of the current visit.
1022 replica_chunk : `ReplicaChunk` or `None`
1023 Replica chunk identifier if replication is configured.
1026 _LOG.debug(
"No objects to write to database.")
1029 visit_time_dt = visit_time.datetime
1030 extra_columns = dict(lastNonForcedSource=visit_time_dt)
1033 extra_columns[
"validityStart"] = visit_time_dt
1035 if not self.
config.time_partition_tables:
1036 extra_columns[
"apdb_time_part"] = time_part
1041 if replica_chunk
is None or not self.
config.use_insert_id_skips_diaobjects:
1043 objs, ApdbTables.DiaObject, extra_columns=extra_columns, time_part=time_part
1046 if replica_chunk
is not None:
1047 extra_columns = dict(apdb_replica_chunk=replica_chunk.id, validityStart=visit_time_dt)
1052 table_name: ApdbTables,
1053 sources: pandas.DataFrame,
1054 visit_time: astropy.time.Time,
1055 replica_chunk: ReplicaChunk |
None,
1057 """Store catalog of DIASources or DIAForcedSources from current visit.
1061 table_name : `ApdbTables`
1062 Table where to store the data.
1063 sources : `pandas.DataFrame`
1064 Catalog containing DiaSource records
1065 visit_time : `astropy.time.Time`
1066 Time of the current visit.
1067 replica_chunk : `ReplicaChunk` or `None`
1068 Replica chunk identifier if replication is configured.
1071 extra_columns: dict[str, Any] = {}
1072 if not self.
config.time_partition_tables:
1073 extra_columns[
"apdb_time_part"] = time_part
1076 self.
_storeObjectsPandas(sources, table_name, extra_columns=extra_columns, time_part=time_part)
1078 if replica_chunk
is not None:
1079 extra_columns = dict(apdb_replica_chunk=replica_chunk.id)
1080 if table_name
is ApdbTables.DiaSource:
1081 extra_table = ExtraTables.DiaSourceChunks
1083 extra_table = ExtraTables.DiaForcedSourceChunks
1087 self, sources: pandas.DataFrame, visit_time: astropy.time.Time, replica_chunk: ReplicaChunk |
None
1089 """Store mapping of diaSourceId to its partitioning values.
1093 sources : `pandas.DataFrame`
1094 Catalog containing DiaSource records
1095 visit_time : `astropy.time.Time`
1096 Time of the current visit.
1098 id_map = cast(pandas.DataFrame, sources[[
"diaSourceId",
"apdb_part"]])
1101 "apdb_replica_chunk": replica_chunk.id
if replica_chunk
is not None else None,
1105 id_map, ExtraTables.DiaSourceToPartition, extra_columns=extra_columns, time_part=
None
1110 records: pandas.DataFrame,
1111 table_name: ApdbTables | ExtraTables,
1112 extra_columns: Mapping |
None =
None,
1113 time_part: int |
None =
None,
1115 """Store generic objects.
1117 Takes Pandas catalog and stores a bunch of records in a table.
1121 records : `pandas.DataFrame`
1122 Catalog containing object records
1123 table_name : `ApdbTables`
1124 Name of the table as defined in APDB schema.
1125 extra_columns : `dict`, optional
1126 Mapping (column_name, column_value) which gives fixed values for
1127 columns in each row, overrides values in ``records`` if matching
1128 columns exist there.
1129 time_part : `int`, optional
1130 If not `None` then insert into a per-partition table.
1134 If Pandas catalog contains additional columns not defined in table
1135 schema they are ignored. Catalog does not have to contain all columns
1136 defined in a table, but partition and clustering keys must be present
1137 in a catalog or ``extra_columns``.
1140 if extra_columns
is None:
1142 extra_fields = list(extra_columns.keys())
1145 df_fields = [column
for column
in records.columns
if column
not in extra_fields]
1147 column_map = self.
_schema.getColumnMap(table_name)
1149 fields = [column_map[field].name
for field
in df_fields
if field
in column_map]
1150 fields += extra_fields
1153 required_columns = self.
_schema.partitionColumns(table_name) + self.
_schema.clusteringColumns(
1156 missing_columns = [column
for column
in required_columns
if column
not in fields]
1158 raise ValueError(f
"Primary key columns are missing from catalog: {missing_columns}")
1160 qfields = [quote_id(field)
for field
in fields]
1161 qfields_str =
",".join(qfields)
1163 with Timer(table_name.name +
" query build", self.
config.timer):
1164 table = self.
_schema.tableName(table_name)
1165 if time_part
is not None:
1166 table = f
"{table}_{time_part}"
1168 holders =
",".join([
"?"] * len(qfields))
1169 query = f
'INSERT INTO "{self._keyspace}"."{table}" ({qfields_str}) VALUES ({holders})'
1170 statement = self.
_preparer.prepare(query)
1171 queries = cassandra.query.BatchStatement()
1172 for rec
in records.itertuples(index=
False):
1174 for field
in df_fields:
1175 if field
not in column_map:
1177 value = getattr(rec, field)
1178 if column_map[field].datatype
is felis.datamodel.DataType.timestamp:
1179 if isinstance(value, pandas.Timestamp):
1180 value = literal(value.to_pydatetime())
1184 value = int(value * 1000)
1185 values.append(literal(value))
1186 for field
in extra_fields:
1187 value = extra_columns[field]
1188 values.append(literal(value))
1189 queries.add(statement, values)
1191 _LOG.debug(
"%s: will store %d records", self.
_schema.tableName(table_name), records.shape[0])
1192 with Timer(table_name.name +
" insert", self.
config.timer):
1193 self.
_session.execute(queries, timeout=self.
config.write_timeout, execution_profile=
"write")
1196 """Calculate spatial partition for each record and add it to a
1201 This overrides any existing column in a DataFrame with the same name
1202 (apdb_part). Original DataFrame is not changed, copy of a DataFrame is
1206 apdb_part = np.zeros(df.shape[0], dtype=np.int64)
1207 ra_col, dec_col = self.
config.ra_dec_columns
1208 for i, (ra, dec)
in enumerate(zip(df[ra_col], df[dec_col])):
1213 df[
"apdb_part"] = apdb_part
1216 def _add_src_part(self, sources: pandas.DataFrame, objs: pandas.DataFrame) -> pandas.DataFrame:
1217 """Add apdb_part column to DiaSource catalog.
1221 This method copies apdb_part value from a matching DiaObject record.
1222 DiaObject catalog needs to have a apdb_part column filled by
1223 ``_add_obj_part`` method and DiaSource records need to be
1224 associated to DiaObjects via ``diaObjectId`` column.
1226 This overrides any existing column in a DataFrame with the same name
1227 (apdb_part). Original DataFrame is not changed, copy of a DataFrame is
1230 pixel_id_map: dict[int, int] = {
1231 diaObjectId: apdb_part
for diaObjectId, apdb_part
in zip(objs[
"diaObjectId"], objs[
"apdb_part"])
1233 apdb_part = np.zeros(sources.shape[0], dtype=np.int64)
1234 ra_col, dec_col = self.
config.ra_dec_columns
1235 for i, (diaObjId, ra, dec)
in enumerate(
1236 zip(sources[
"diaObjectId"], sources[ra_col], sources[dec_col])
1246 apdb_part[i] = pixel_id_map[diaObjId]
1247 sources = sources.copy()
1248 sources[
"apdb_part"] = apdb_part
1251 def _add_fsrc_part(self, sources: pandas.DataFrame, objs: pandas.DataFrame) -> pandas.DataFrame:
1252 """Add apdb_part column to DiaForcedSource catalog.
1256 This method copies apdb_part value from a matching DiaObject record.
1257 DiaObject catalog needs to have a apdb_part column filled by
1258 ``_add_obj_part`` method and DiaSource records need to be
1259 associated to DiaObjects via ``diaObjectId`` column.
1261 This overrides any existing column in a DataFrame with the same name
1262 (apdb_part). Original DataFrame is not changed, copy of a DataFrame is
1265 pixel_id_map: dict[int, int] = {
1266 diaObjectId: apdb_part
for diaObjectId, apdb_part
in zip(objs[
"diaObjectId"], objs[
"apdb_part"])
1268 apdb_part = np.zeros(sources.shape[0], dtype=np.int64)
1269 for i, diaObjId
in enumerate(sources[
"diaObjectId"]):
1270 apdb_part[i] = pixel_id_map[diaObjId]
1271 sources = sources.copy()
1272 sources[
"apdb_part"] = apdb_part
1277 """Calculate time partition number for a given time.
1281 time : `float` or `astropy.time.Time`
1282 Time for which to calculate partition number. Can be float to mean
1283 MJD or `astropy.time.Time`
1285 Epoch time for partition 0.
1287 Number of days per partition.
1292 Partition number for a given time.
1294 if isinstance(time, astropy.time.Time):
1295 mjd = float(time.mjd)
1298 days_since_epoch = mjd - epoch_mjd
1299 partition = int(days_since_epoch) // part_days
1303 """Calculate time partition number for a given time.
1307 time : `float` or `astropy.time.Time`
1308 Time for which to calculate partition number. Can be float to mean
1309 MJD or `astropy.time.Time`
1314 Partition number for a given time.
1316 if isinstance(time, astropy.time.Time):
1317 mjd = float(time.mjd)
1321 partition = int(days_since_epoch) // self.
config.time_partition_days
1325 """Make an empty catalog for a table with a given name.
1329 table_name : `ApdbTables`
1334 catalog : `pandas.DataFrame`
1337 table = self.
_schema.tableSchemas[table_name]
1340 columnDef.name: pandas.Series(dtype=self.
_schema.column_dtype(columnDef.datatype))
1341 for columnDef
in table.columns
1343 return pandas.DataFrame(data)
1348 where1: list[tuple[str, tuple]],
1349 where2: list[tuple[str, tuple]],
1350 suffix: str |
None =
None,
1351 ) -> Iterator[tuple[cassandra.query.Statement, tuple]]:
1352 """Make cartesian product of two parts of WHERE clause into a series
1353 of statements to execute.
1358 Initial statement prefix that comes before WHERE clause, e.g.
1359 "SELECT * from Table"
1367 for expr1, params1
in where1:
1368 for expr2, params2
in where2:
1372 wheres.append(expr1)
1374 wheres.append(expr2)
1376 full_query +=
" WHERE " +
" AND ".join(wheres)
1378 full_query +=
" " + suffix
1379 params = params1 + params2
1381 statement = self.
_preparer.prepare(full_query)
1386 statement = cassandra.query.SimpleStatement(full_query)
1387 yield (statement, params)
1390 self, region: sphgeom.Region |
None, use_ranges: bool =
False
1391 ) -> list[tuple[str, tuple]]:
1392 """Generate expressions for spatial part of WHERE clause.
1396 region : `sphgeom.Region`
1397 Spatial region for query results.
1399 If True then use pixel ranges ("apdb_part >= p1 AND apdb_part <=
1400 p2") instead of exact list of pixels. Should be set to True for
1401 large regions covering very many pixels.
1405 expressions : `list` [ `tuple` ]
1406 Empty list is returned if ``region`` is `None`, otherwise a list
1407 of one or more (expression, parameters) tuples
1413 expressions: list[tuple[str, tuple]] = []
1414 for lower, upper
in pixel_ranges:
1417 expressions.append((
'"apdb_part" = ?', (lower,)))
1419 expressions.append((
'"apdb_part" >= ? AND "apdb_part" <= ?', (lower, upper)))
1423 if self.
config.query_per_spatial_part:
1424 return [(
'"apdb_part" = ?', (pixel,))
for pixel
in pixels]
1426 pixels_str =
",".join([str(pix)
for pix
in pixels])
1427 return [(f
'"apdb_part" IN ({pixels_str})', ())]
1432 start_time: float | astropy.time.Time,
1433 end_time: float | astropy.time.Time,
1434 query_per_time_part: bool |
None =
None,
1435 ) -> tuple[list[str], list[tuple[str, tuple]]]:
1436 """Generate table names and expressions for temporal part of WHERE
1441 table : `ApdbTables`
1442 Table to select from.
1443 start_time : `astropy.time.Time` or `float`
1444 Starting Datetime of MJD value of the time range.
1445 end_time : `astropy.time.Time` or `float`
1446 Starting Datetime of MJD value of the time range.
1447 query_per_time_part : `bool`, optional
1448 If None then use ``query_per_time_part`` from configuration.
1452 tables : `list` [ `str` ]
1453 List of the table names to query.
1454 expressions : `list` [ `tuple` ]
1455 A list of zero or more (expression, parameters) tuples.
1458 temporal_where: list[tuple[str, tuple]] = []
1459 table_name = self.
_schema.tableName(table)
1462 time_parts = list(range(time_part_start, time_part_end + 1))
1463 if self.
config.time_partition_tables:
1464 tables = [f
"{table_name}_{part}" for part
in time_parts]
1466 tables = [table_name]
1467 if query_per_time_part
is None:
1468 query_per_time_part = self.
config.query_per_time_part
1469 if query_per_time_part:
1470 temporal_where = [(
'"apdb_time_part" = ?', (time_part,))
for time_part
in time_parts]
1472 time_part_list =
",".join([str(part)
for part
in time_parts])
1473 temporal_where = [(f
'"apdb_time_part" IN ({time_part_list})', ())]
1475 return tables, temporal_where
std::vector< SchemaItem< Flag > > * items
VersionTuple apdbImplementationVersion(cls)
str translate(self, str private_ip)
__init__(self, list[str] public_ips, list[str] private_ips)
bool use_insert_id_skips_diaobjects
use_insert_id_skips_diaobjects
None update(self, str json_str)
bool time_partition_tables
__init__(self, ApdbCassandraConfig config)
None _storeDiaSources(self, ApdbTables table_name, pandas.DataFrame sources, astropy.time.Time visit_time, ReplicaChunk|None replica_chunk)
None _storeReplicaChunk(self, ReplicaChunk replica_chunk, astropy.time.Time visit_time)
str metadataReplicaVersionKey
tuple[Cluster, Session] _make_session(cls, ApdbCassandraConfig config)
pandas.DataFrame _add_obj_part(self, pandas.DataFrame df)
_partition_zero_epoch_mjd
bool containsVisitDetector(self, int visit, int detector)
str metadataSchemaVersionKey
None reassignDiaSources(self, Mapping[int, int] idMap)
pandas.DataFrame getDiaObjects(self, sphgeom.Region region)
None storeSSObjects(self, pandas.DataFrame objects)
None _storeDiaObjects(self, pandas.DataFrame objs, astropy.time.Time visit_time, ReplicaChunk|None replica_chunk)
ApdbCassandraConfig init_database(cls, list[str] hosts, str keyspace, *str|None schema_file=None, str|None schema_name=None, int|None read_sources_months=None, int|None read_forced_sources_months=None, bool use_insert_id=False, bool replica_skips_diaobjects=False, int|None port=None, str|None username=None, str|None prefix=None, str|None part_pixelization=None, int|None part_pix_level=None, bool time_partition_tables=True, str|None time_partition_start=None, str|None time_partition_end=None, str|None read_consistency=None, str|None write_consistency=None, int|None read_timeout=None, int|None write_timeout=None, list[str]|None ra_dec_columns=None, int|None replication_factor=None, bool drop=False)
pandas.DataFrame _make_empty_catalog(self, ApdbTables table_name)
AuthProvider|None _make_auth_provider(cls, ApdbCassandraConfig config)
None _versionCheck(self, ApdbMetadataCassandra metadata)
__init__(self, ApdbCassandraConfig config)
Mapping[Any, ExecutionProfile] _makeProfiles(cls, ApdbCassandraConfig config)
ApdbMetadata metadata(self)
tuple[list[str], list[tuple[str, tuple]]] _temporal_where(self, ApdbTables table, float|astropy.time.Time start_time, float|astropy.time.Time end_time, bool|None query_per_time_part=None)
None _storeDiaSourcesPartitions(self, pandas.DataFrame sources, astropy.time.Time visit_time, ReplicaChunk|None replica_chunk)
VersionTuple apdbImplementationVersion(cls)
None store(self, astropy.time.Time visit_time, pandas.DataFrame objects, pandas.DataFrame|None sources=None, pandas.DataFrame|None forced_sources=None)
None _storeObjectsPandas(self, pandas.DataFrame records, ApdbTables|ExtraTables table_name, Mapping|None extra_columns=None, int|None time_part=None)
pandas.DataFrame|None getDiaForcedSources(self, sphgeom.Region region, Iterable[int]|None object_ids, astropy.time.Time visit_time)
str metadataCodeVersionKey
metadataReplicaVersionKey
None _makeSchema(cls, ApdbConfig config, *bool drop=False, int|None replication_factor=None)
ApdbCassandraReplica get_replica(self)
list[tuple[str, tuple]] _spatial_where(self, sphgeom.Region|None region, bool use_ranges=False)
pandas.DataFrame getSSObjects(self)
VersionTuple apdbSchemaVersion(self)
int _time_partition_cls(cls, float|astropy.time.Time time, float epoch_mjd, int part_days)
Iterator[tuple[cassandra.query.Statement, tuple]] _combine_where(self, str prefix, list[tuple[str, tuple]] where1, list[tuple[str, tuple]] where2, str|None suffix=None)
pandas.DataFrame _add_fsrc_part(self, pandas.DataFrame sources, pandas.DataFrame objs)
pandas.DataFrame _getSources(self, sphgeom.Region region, Iterable[int]|None object_ids, float mjd_start, float mjd_end, ApdbTables table_name)
pandas.DataFrame|None getDiaSources(self, sphgeom.Region region, Iterable[int]|None object_ids, astropy.time.Time visit_time)
int _time_partition(self, float|astropy.time.Time time)
pandas.DataFrame _add_src_part(self, pandas.DataFrame sources, pandas.DataFrame objs)
int countUnassociatedObjects(self)
Table|None tableDef(self, ApdbTables table)
UnitVector3d is a unit vector in ℝ³ with components stored in double precision.
daf::base::PropertySet * set