From 2f6eeec9dc97cebab72e1c7f0f3197536b0b007e Mon Sep 17 00:00:00 2001 From: Daniel Thom Date: Tue, 7 Jan 2025 11:51:17 -0700 Subject: [PATCH] Speed up supplemental attribute queries --- src/supplemental_attribute_associations.jl | 168 +++++++++++---------- src/supplemental_attribute_manager.jl | 15 +- src/system_data.jl | 10 ++ src/time_series_metadata_store.jl | 89 +++++------ 4 files changed, 152 insertions(+), 130 deletions(-) diff --git a/src/supplemental_attribute_associations.jl b/src/supplemental_attribute_associations.jl index 8a5bd15f9..a5abd405e 100644 --- a/src/supplemental_attribute_associations.jl +++ b/src/supplemental_attribute_associations.jl @@ -31,11 +31,13 @@ end """ Construct a new SupplementalAttributeAssociations with an in-memory database. """ -function SupplementalAttributeAssociations() +function SupplementalAttributeAssociations(; create_indexes = true) associations = SupplementalAttributeAssociations(SQLite.DB(), Dict{String, SQLite.Stmt}()) _create_attribute_associations_table!(associations) - _create_indexes!(associations) + if create_indexes + _create_indexes!(associations) + end @debug "Initialized new supplemental attributes association table" _group = LOG_GROUP_SUPPLEMENTAL_ATTRIBUTES return associations @@ -64,20 +66,14 @@ function _create_indexes!(associations::SupplementalAttributeAssociations) SQLite.createindex!( associations.db, SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME, - "by_attribute_and_component", + "by_component", [ - "attribute_uuid", "component_uuid", + "attribute_uuid", + "attribute_type", ]; unique = false, ) - SQLite.createindex!( - associations.db, - SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME, - "by_component", - "component_uuid"; - unique = false, - ) return end @@ -110,8 +106,8 @@ function add_association!( string(nameof(typeof(component))), ) params = chop(repeat("?,", length(row))) - SQLite.DBInterface.execute( - associations.db, + _execute_cached( + associations, "INSERT INTO $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME VALUES($params)", row, ) @@ -193,6 +189,13 @@ function get_num_components_with_attributes(associations::SupplementalAttributeA return _execute_count(associations, query) end +const _HAS_ASSOCIATION_BY_ATTRIBUTE = """ + SELECT attribute_uuid + FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME + WHERE attribute_uuid = ? + LIMIT 1 +""" + """ Return true if there is at least one association matching the inputs. """ @@ -200,9 +203,21 @@ function has_association( associations::SupplementalAttributeAssociations, attribute::SupplementalAttribute, ) - return _has_association(associations, attribute, "attribute_uuid") + # Note: Unlike the other has_association methods, this is not covered by an index. + params = (string(get_uuid(attribute)),) + return !isempty( + Tables.rowtable( + _execute_cached(associations, _HAS_ASSOCIATION_BY_ATTRIBUTE, params), + ), + ) end +const _HAS_ASSOCIATION_BY_COMPONENT_ATTRIBUTE = """ + SELECT attribute_uuid + FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME + WHERE attribute_uuid = ? AND component_uuid = ? + LIMIT 1 +""" function has_association( associations::SupplementalAttributeAssociations, component::InfrastructureSystemsComponent, @@ -210,53 +225,55 @@ function has_association( ) a_uuid = get_uuid(attribute) c_uuid = get_uuid(component) - query = """ - SELECT attribute_uuid - FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME - WHERE attribute_uuid = ? AND component_uuid = ? - LIMIT 1 - """ params = (string(a_uuid), string(c_uuid)) - return !isempty(_execute_cached(associations, query, params)) + return !isempty( + _execute_cached(associations, _HAS_ASSOCIATION_BY_COMPONENT_ATTRIBUTE, params), + ) end +const _HAS_ASSOCIATION_BY_COMPONENT = """ + SELECT attribute_uuid + FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME + WHERE component_uuid = ? + LIMIT 1 +""" function has_association( associations::SupplementalAttributeAssociations, component::InfrastructureSystemsComponent, ) params = (string(get_uuid(component)),) - query = """ - SELECT attribute_uuid - FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME - WHERE component_uuid = ? - LIMIT 1 - """ - return !isempty(Tables.rowtable(_execute_cached(associations, query, params))) + return !isempty( + Tables.rowtable( + _execute_cached(associations, _HAS_ASSOCIATION_BY_COMPONENT, params), + ), + ) end +const _HAS_ASSOCIATION_BY_COMP_ATTR_TYPE = """ + SELECT attribute_uuid + FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME + WHERE component_uuid = ? AND attribute_type = ? + LIMIT 1 +""" function has_association( associations::SupplementalAttributeAssociations, component::InfrastructureSystemsComponent, attribute_type::Type{<:SupplementalAttribute}, ) - c_str = "component_uuid = ?" - params = [string(get_uuid(component))] - where_clause = if isnothing(attribute_type) - c_str - else - a_str = _get_attribute_type_string!(params, attribute_type) - "$c_str AND $a_str" - end - - query = """ - SELECT attribute_uuid - FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME - WHERE $where_clause - LIMIT 1 - """ - return !isempty(Tables.rowtable(_execute_cached(associations, query, params))) + params = (string(get_uuid(component)), string(nameof(attribute_type))) + return !isempty( + Tables.rowtable( + _execute_cached(associations, _HAS_ASSOCIATION_BY_COMP_ATTR_TYPE, params), + ), + ) end +const _LIST_ASSOCIATED_COMP_UUIDS = """ + SELECT component_uuid + FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME + WHERE attribute_uuid = ? +""" + """ Return the component UUIDs associated with the attribute. """ @@ -264,12 +281,10 @@ function list_associated_component_uuids( associations::SupplementalAttributeAssociations, attribute::SupplementalAttribute, ) - query = """ - SELECT component_uuid - FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME - WHERE attribute_uuid = '$(get_uuid(attribute))' - """ - table = Tables.columntable(_execute_cached(associations, query)) + params = (string(get_uuid(attribute)),) + table = Tables.columntable( + _execute_cached(associations, _LIST_ASSOCIATED_COMP_UUIDS, params), + ) return Base.UUID.(table.component_uuid) end @@ -306,10 +321,9 @@ function remove_association!( component::InfrastructureSystemsComponent, attribute::SupplementalAttribute, ) - a_uuid = get_uuid(attribute) - c_uuid = get_uuid(component) - where_clause = "WHERE attribute_uuid = '$a_uuid' AND component_uuid = '$c_uuid'" - num_deleted = _remove_associations!(associations, where_clause) + where_clause = "WHERE attribute_uuid = ? AND component_uuid = ?" + params = (string(get_uuid(attribute)), string(get_uuid(component))) + num_deleted = _remove_associations!(associations, where_clause, params) if num_deleted != 1 error("Bug: unexpected number of deletions: $num_deleted. Should have been 1.") end @@ -322,13 +336,20 @@ function remove_associations!( associations::SupplementalAttributeAssociations, type::Type{<:SupplementalAttribute}, ) - where_clause = "WHERE attribute_type = '$(nameof(type))'" - num_deleted = _remove_associations!(associations, where_clause) + where_clause = "WHERE attribute_type = ?" + params = (string(nameof(type)),) + num_deleted = _remove_associations!(associations, where_clause, params) @debug "Deleted $num_deleted supplemental attribute associations" _group = LOG_GROUP_SUPPLEMENTAL_ATTRIBUTES return end +const _REPLACE_COMP_UUID_SA = """ + UPDATE $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME + SET component_uuid = ? + WHERE component_uuid = ? +""" + """ Replace the component UUID in the table. """ @@ -337,12 +358,8 @@ function replace_component_uuid!( old_uuid::Base.UUID, new_uuid::Base.UUID, ) - query = """ - UPDATE $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME - SET component_uuid = '$new_uuid' - WHERE component_uuid = '$old_uuid' - """ - _execute(associations, query) + params = (string(new_uuid), string(old_uuid)) + _execute_cached(associations, _REPLACE_COMP_UUID_SA, params) return end @@ -368,8 +385,9 @@ end """ Add records to the database. Expects output from [`to_records`](@ref). """ -function load_records!(associations::SupplementalAttributeAssociations, records) - isempty(records) && return +function from_records(::Type{SupplementalAttributeAssociations}, records) + associations = SupplementalAttributeAssociations(; create_indexes = false) + isempty(records) && return associations columns = ("attribute_uuid", "attribute_type", "component_uuid", "component_type") num_rows = length(records) @@ -386,31 +404,19 @@ function load_records!(associations::SupplementalAttributeAssociations, records) "INSERT INTO $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME VALUES($params)", NamedTuple(Symbol(k) => v for (k, v) in data), ) - return -end - -function _has_association( - associations::SupplementalAttributeAssociations, - val::Union{InfrastructureSystemsComponent, SupplementalAttribute}, - column::String, -) - uuid = get_uuid(val) - query = """ - SELECT attribute_uuid - FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME - WHERE $column = '$uuid' - LIMIT 1 - """ - return !isempty(Tables.rowtable(_execute_cached(associations, query))) + _create_indexes!(associations) + return associations end function _remove_associations!( associations::SupplementalAttributeAssociations, where_clause::AbstractString, + params, ) - _execute( + _execute_cached( associations, "DELETE FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME $where_clause", + params, ) table = Tables.rowtable(_execute(associations, "SELECT CHANGES() AS changes")) @assert_op length(table) == 1 diff --git a/src/supplemental_attribute_manager.jl b/src/supplemental_attribute_manager.jl index 6acb4f56b..58fb64ce8 100644 --- a/src/supplemental_attribute_manager.jl +++ b/src/supplemental_attribute_manager.jl @@ -6,12 +6,11 @@ struct SupplementalAttributeManager <: InfrastructureSystemsContainer associations::SupplementalAttributeAssociations end -function SupplementalAttributeManager(data::SupplementalAttributesByType) - return SupplementalAttributeManager(data, SupplementalAttributeAssociations()) -end - function SupplementalAttributeManager() - return SupplementalAttributeManager(SupplementalAttributesByType()) + return SupplementalAttributeManager( + SupplementalAttributesByType(), + SupplementalAttributeAssociations(), + ) end get_member_string(::SupplementalAttributeManager) = "supplemental attributes" @@ -244,7 +243,9 @@ function deserialize( @debug "Deserialized $(summary(attr))" _group = LOG_GROUP_SERIALIZATION end - mgr = SupplementalAttributeManager(SupplementalAttributesByType(attributes)) - load_records!(mgr.associations, data["associations"]) + mgr = SupplementalAttributeManager( + SupplementalAttributesByType(attributes), + from_records(SupplementalAttributeAssociations, data["associations"]), + ) return mgr end diff --git a/src/system_data.jl b/src/system_data.jl index d97e3b2bf..5896b95f8 100644 --- a/src/system_data.jl +++ b/src/system_data.jl @@ -1152,6 +1152,16 @@ function add_supplemental_attribute!(data::SystemData, component, attribute; kwa return end +""" +Begin a transaction to add supplemental attributes. Use this function when adding +many supplemental attributes in order to improve performance. +""" +function begin_supplemental_attributes_transaction(func::Function, data::SystemData) + SQLite.transaction(data.supplemental_attribute_manager.associations.db) do + func() + end +end + function get_supplemental_attributes( filter_func::Function, ::Type{T}, diff --git a/src/time_series_metadata_store.jl b/src/time_series_metadata_store.jl index c9920205e..e62a3661b 100644 --- a/src/time_series_metadata_store.jl +++ b/src/time_series_metadata_store.jl @@ -295,40 +295,40 @@ function add_metadata!( metadata::TimeSeriesMetadata, ) TimerOutputs.@timeit_debug SYSTEM_TIMERS "add time series metadata single" begin - SQLite.transaction(store.db) do - owner_category = _get_owner_category(owner) - time_series_type = time_series_metadata_to_data(metadata) - ts_category = _get_time_series_category(time_series_type) - features = make_features_string(metadata.features) - vals = _create_row( - metadata, - owner, - owner_category, - _convert_ts_type_to_string(time_series_type), - ts_category, - features, - ) - params = chop(repeat("?,", length(vals))) - _execute_cached( - store, - "INSERT INTO $ASSOCIATIONS_TABLE_NAME VALUES($params)", - vals, - ) - metadata_uuid = get_uuid(metadata) - metadata_row = - (missing, string(metadata_uuid), JSON3.write(serialize(metadata))) - metadata_params = ("?,?,jsonb(?)") + owner_category = _get_owner_category(owner) + time_series_type = time_series_metadata_to_data(metadata) + ts_category = _get_time_series_category(time_series_type) + features = make_features_string(metadata.features) + vals = _create_row( + metadata, + owner, + owner_category, + _convert_ts_type_to_string(time_series_type), + ts_category, + features, + ) + params = chop(repeat("?,", length(vals))) + _execute_cached( + store, + "INSERT INTO $ASSOCIATIONS_TABLE_NAME VALUES($params)", + vals, + ) + metadata_uuid = get_uuid(metadata) + metadata_row = + (missing, string(metadata_uuid), JSON3.write(serialize(metadata))) + metadata_params = ("?,?,jsonb(?)") + TimerOutputs.@timeit_debug SYSTEM_TIMERS "add ts_metadata row" begin _execute_cached( store, "INSERT OR IGNORE INTO $METADATA_TABLE_NAME VALUES($metadata_params)", metadata_row, ) - if !haskey(store.metadata_uuids, metadata_uuid) - store.metadata_uuids[metadata_uuid] = metadata - end - @debug "Added metadata = $metadata to $(summary(owner))" _group = - LOG_GROUP_TIME_SERIES end + if !haskey(store.metadata_uuids, metadata_uuid) + store.metadata_uuids[metadata_uuid] = metadata + end + @debug "Added metadata = $metadata to $(summary(owner))" _group = + LOG_GROUP_TIME_SERIES end return end @@ -806,15 +806,17 @@ function has_metadata( return _has_metadata(stmt, params2) end +const _HAS_METADATA_BY_TS_UUID = "SELECT id FROM $ASSOCIATIONS_TABLE_NAME WHERE time_series_uuid = ?" + """ Return True if there is time series matching the UUID. """ function has_metadata(store::TimeSeriesMetadataStore, time_series_uuid::Base.UUID) - query = "SELECT id FROM $ASSOCIATIONS_TABLE_NAME WHERE time_series_uuid = ?" params = (string(time_series_uuid),) - return _has_metadata(store, query, params) + return _has_metadata(store, _HAS_METADATA_BY_TS_UUID, params) end +const _BASE_HAS_METADATA = "SELECT id FROM $ASSOCIATIONS_TABLE_NAME WHERE owner_uuid = ?" function _make_has_metadata_statement!( store::TimeSeriesMetadataStore, key::HasMetadataQueryKey; @@ -824,7 +826,6 @@ function _make_has_metadata_statement!( return stmt end - query = "SELECT id FROM $ASSOCIATIONS_TABLE_NAME WHERE owner_uuid = ?" where_clauses = String[] if key.has_name push!(where_clauses, "name = ?") @@ -847,10 +848,10 @@ function _make_has_metadata_statement!( end if isempty(where_clauses) - final = "$query LIMIT 1" + final = "$_BASE_HAS_METADATA LIMIT 1" else where_clause = join(where_clauses, " AND ") - final = "$query AND $where_clause LIMIT 1" + final = "$_BASE_HAS_METADATA AND $where_clause LIMIT 1" end stmt = SQLite.Stmt(store.db, final) @@ -1162,18 +1163,18 @@ function _handle_removed_metadata(store::TimeSeriesMetadataStore, metadata_uuid: end end +const _REPLACE_COMP_UUID_TS = """ + UPDATE $ASSOCIATIONS_TABLE_NAME + SET owner_uuid = ? + WHERE owner_uuid = ? +""" function replace_component_uuid!( store::TimeSeriesMetadataStore, old_uuid::Base.UUID, new_uuid::Base.UUID, ) - query = """ - UPDATE $ASSOCIATIONS_TABLE_NAME - SET owner_uuid = ? - WHERE owner_uuid = ? - """ - params = [string(new_uuid), string(old_uuid)] - _execute(store, query, params) + params = (string(new_uuid), string(old_uuid)) + _execute(store, _REPLACE_COMP_UUID_TS, params) return end @@ -1306,10 +1307,14 @@ function _try_get_time_series_metadata_by_full_params( end end +const _GET_METADATA_BY_UUID = """ + SELECT json(metadata) AS metadata FROM $METADATA_TABLE_NAME WHERE metadata_uuid = ? +""" function _get_metadata_by_uuid!(store::TimeSeriesMetadataStore, metadata_uuid::Base.UUID) return get!(store.metadata_uuids, metadata_uuid) do - query = "SELECT json(metadata) AS metadata FROM $METADATA_TABLE_NAME WHERE metadata_uuid = ?" - rows = Tables.rowtable(_execute_cached(store, query, (string(metadata_uuid),))) + rows = Tables.rowtable( + _execute_cached(store, _GET_METADATA_BY_UUID, (string(metadata_uuid),)), + ) if length(rows) != 1 error( "Bug: _get_metadata returned $(length(rows)) entries instead of 1 for $metadata_uuid",