perf(query-compiler): avoid compound selector materialization

Signed-off-by: Alexey Orlenko's AI Agent <robot@aqrln.net>

Author: tensordreams

query-compiler/core/src/query_document/mod.rs

@@ -30,7 +30,7 @@ pub(crate) use parser::*;
 
 use crate::{
     query_ast::{QueryOption, QueryOptions},
-    query_graph_builder::resolve_compound_field,
+    query_graph_builder::is_compound_field,
 };
 use itertools::Itertools;
 use query_structure::Model;
@@ -91,7 +91,7 @@ impl BatchDocument {
 
         where_obj.iter().any(|(key, val)| match val {
             // If it's a compound, then it's still considered as scalar
-            ArgumentValue::Object(_) if resolve_compound_field(key, &model).is_some() => false,
+            ArgumentValue::Object(_) if is_compound_field(key, &model) => false,
             // Otherwise, we just look for a scalar field inside the model. If it's not one, then we break.
             val => match model.fields().find_from_scalar(key) {
                 Ok(sf) => match val {
@@ -319,7 +319,7 @@ fn extract_filter(where_obj: ArgumentValueObject, model: &Model) -> Vec<Selectio
         .into_iter()
         .flat_map(|(key, val)| match val {
             // This means our query has a compound field in the form of: {co1_col2: { col1_col2: { col1: <val>, col2: <val> } }}
-            ArgumentValue::Object(obj) if resolve_compound_field(&key, model).is_some() => obj.into_iter().collect(),
+            ArgumentValue::Object(obj) if is_compound_field(&key, model) => obj.into_iter().collect(),
             // This means our query has a scalar filter in the form of {col1: { equals: <val> }}
             ArgumentValue::Object(obj) => {
                 // This is safe because it's been validated before in the `.can_compact` method.

query-compiler/core/src/query_graph_builder/extractors/filters/mod.rs

@@ -23,7 +23,7 @@ pub fn extract_unique_filter(value_map: ParsedInputMap<'_>, model: &Model) -> Qu
             .fields()
             .find_from_scalar(field_name)
             .is_ok_and(|field| field.unique())
-            || utils::resolve_compound_field(field_name, model).is_some()
+            || utils::is_compound_field(field_name, model)
     }) {
         return internal_extract_unique_filter(value_map, model);
     }
@@ -36,7 +36,7 @@ pub fn extract_unique_filter(value_map: ParsedInputMap<'_>, model: &Model) -> Qu
             .into_iter()
             .partition(|(field_name, _)| match model.fields().find_from_scalar(field_name) {
                 Ok(field) => field.unique(),
-                Err(_) => utils::resolve_compound_field(field_name, model).is_some(),
+                Err(_) => utils::is_compound_field(field_name, model),
             });
     let mut unique_map = ParsedInputMap::from(unique_map);
     let mut rest_map = ParsedInputMap::from(rest_map);

query-compiler/core/src/query_graph_builder/extractors/mod.rs

@@ -6,6 +6,6 @@ mod utils;
 pub(crate) use filters::*;
 pub(crate) use query_arguments::*;
 pub(crate) use rel_aggregations::*;
-pub(crate) use utils::resolve_compound_field;
+pub(crate) use utils::is_compound_field;
 
 use crate::query_document::*;

query-compiler/core/src/query_graph_builder/extractors/utils.rs

@@ -1,26 +1,74 @@
 use query_structure::{Model, ScalarFieldRef};
 
+/// Checks whether a field name resolves to a compound field without materializing the fields.
+pub fn is_compound_field(name: &str, model: &Model) -> bool {
+    is_compound_id(name, model) || is_index_field(name, model)
+}
+
 /// Attempts to resolve a field name to a compound field.
 pub fn resolve_compound_field(name: &str, model: &Model) -> Option<Vec<ScalarFieldRef>> {
     resolve_compound_id(name, model).or_else(|| resolve_index_fields(name, model))
 }
 
 /// Attempts to match a given name to the (schema) name of a compound id field on the model.
 pub fn resolve_compound_id(name: &str, model: &Model) -> Option<Vec<ScalarFieldRef>> {
-    model.fields().compound_id().and_then(|pk| {
-        (name == schema::compound_id_field_name(model.walker().primary_key().unwrap())).then(|| pk.collect())
-    })
+    model
+        .fields()
+        .compound_id()
+        .and_then(|pk| is_compound_id(name, model).then(|| pk.collect()))
 }
 
 /// Attempts to match a given name to the (schema) name of a compound indexes on the model and returns the first match.
 pub fn resolve_index_fields(name: &str, model: &Model) -> Option<Vec<ScalarFieldRef>> {
     model
         .unique_indexes()
-        .find(|index| schema::compound_index_field_name(*index) == name)
+        .find(|index| index_field_name_matches(name, *index))
         .map(|index| {
             index
                 .fields()
                 .map(|f| ScalarFieldRef::from((model.dm.clone(), f)))
                 .collect()
         })
 }
+
+fn is_compound_id(name: &str, model: &Model) -> bool {
+    model.fields().compound_id().is_some()
+        && model.walker().primary_key().is_some_and(|pk| match pk.name() {
+            Some(pk_name) => name == pk_name,
+            None => compound_field_name_matches(name, pk.fields().map(|field| field.name())),
+        })
+}
+
+fn is_index_field(name: &str, model: &Model) -> bool {
+    model
+        .unique_indexes()
+        .any(|index| index_field_name_matches(name, index))
+}
+
+fn index_field_name_matches(name: &str, index: psl::parser_database::walkers::IndexWalker<'_>) -> bool {
+    match index.name() {
+        Some(index_name) => name == index_name,
+        None => compound_field_name_matches(name, index.fields().map(|field| field.name())),
+    }
+}
+
+fn compound_field_name_matches<'a>(mut name: &str, fields: impl Iterator<Item = &'a str>) -> bool {
+    let mut first = true;
+
+    for field in fields {
+        if first {
+            first = false;
+        } else if let Some(rest) = name.strip_prefix('_') {
+            name = rest;
+        } else {
+            return false;
+        }
+
+        let Some(rest) = name.strip_prefix(field) else {
+            return false;
+        };
+        name = rest;
+    }
+
+    !first && name.is_empty()
+}

query-compiler/core/src/query_graph_builder/write/upsert.rs

@@ -232,7 +232,7 @@ fn can_use_connector_native_upsert<'a>(
 fn is_unique_field(field_name: &str, model: &Model) -> bool {
     match model.fields().find_from_scalar(field_name) {
         Ok(field) => field.unique(),
-        Err(_) => resolve_compound_field(field_name, model).is_some(),
+        Err(_) => is_compound_field(field_name, model),
     }
 }
 

query-compiler/query-compiler/tests/data/query-compound-id.json

@@ -0,0 +1,19 @@
+{
+  "modelName": "ParentModelWithCompositeId",
+  "action": "findUnique",
+  "query": {
+    "arguments": {
+      "relationLoadStrategy": "query",
+      "where": {
+        "a_b": {
+          "a": 1,
+          "b": 1
+        }
+      }
+    },
+    "selection": {
+      "$composites": true,
+      "$scalars": true
+    }
+  }
+}

query-compiler/query-compiler/tests/snapshots/queries__queries@query-compound-id.json.snap

@@ -0,0 +1,18 @@
+---
+source: query-compiler/query-compiler/tests/queries.rs
+expression: pretty
+input_file: query-compiler/query-compiler/tests/data/query-compound-id.json
+snapshot_kind: text
+---
+dataMap {
+    a: Int (a)
+    b: Int (b)
+}
+unique (query «SELECT "public"."ParentModelWithCompositeId"."a",
+               "public"."ParentModelWithCompositeId"."b" FROM
+               "public"."ParentModelWithCompositeId" WHERE
+               ("public"."ParentModelWithCompositeId"."a" = $1 AND
+               "public"."ParentModelWithCompositeId"."b" = $2) LIMIT $3 OFFSET
+               $4»
+        params [const(BigInt(1)), const(BigInt(1)), const(BigInt(1)),
+                const(BigInt(0))])