Speed up supplemental attribute queries

NREL-Sienna · Jan 10, 2025 · 2f6eeec · 2f6eeec
1 parent 0561aa1
commit 2f6eeec
Show file tree

Hide file tree

Showing 4 changed files with 152 additions and 130 deletions.
diff --git 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,83 +189,102 @@ 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.
 """
 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,
     attribute::SupplementalAttribute,
 )
     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.
 """
 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
@@ -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
@@ -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},