Skip to main content

uni_query/query/
function_props.rs

1// M-CANONICAL-DOCS: This module defines property requirements for Cypher functions
2// M-CANONICAL-DOCS: to enable pushdown hydration optimization
3//
4// Pushdown hydration analyzes which properties a query needs and loads them during
5// the initial scan, transforming property loading from O(N*M) to O(N) complexity.
6
7use std::sync::LazyLock;
8
9/// Specification of property requirements for a Cypher function.
10///
11/// This helps the query planner understand which properties need to be loaded
12/// for entity arguments to a function, enabling pushdown hydration.
13#[derive(Debug, Clone, Copy)]
14pub struct FunctionPropertySpec {
15    /// Argument positions containing entity references (0-indexed).
16    /// For example, in `validAt(entity, start, end, ts)`, position 0 is the entity.
17    pub entity_args: &'static [usize],
18
19    /// (arg_index, entity_arg_index) pairs for property name arguments.
20    /// For example, in `validAt(entity, 'start', 'end', ts)`:
21    /// - (1, 0) means argument 1 is a property name for entity at position 0
22    /// - (2, 0) means argument 2 is a property name for entity at position 0
23    pub property_name_args: &'static [(usize, usize)],
24
25    /// If true, requires all properties of entity (e.g., keys(), properties()).
26    pub needs_full_entity: bool,
27}
28
29/// Static registry of function property specifications.
30/// Function names are uppercase for case-insensitive lookup.
31static FUNCTION_SPECS: LazyLock<std::collections::HashMap<&'static str, FunctionPropertySpec>> =
32    LazyLock::new(|| {
33        // Helper specs for common patterns
34        let full_entity = FunctionPropertySpec {
35            entity_args: &[0],
36            property_name_args: &[],
37            needs_full_entity: true,
38        };
39        let entity_arg_only = FunctionPropertySpec {
40            entity_args: &[0],
41            property_name_args: &[],
42            needs_full_entity: false,
43        };
44        let no_entity = FunctionPropertySpec {
45            entity_args: &[],
46            property_name_args: &[],
47            needs_full_entity: false,
48        };
49
50        std::collections::HashMap::from([
51            // uni.temporal.validAt(entity, start_prop, end_prop, timestamp)
52            (
53                "UNI.TEMPORAL.VALIDAT",
54                FunctionPropertySpec {
55                    entity_args: &[0],
56                    property_name_args: &[(1, 0), (2, 0)],
57                    needs_full_entity: false,
58                },
59            ),
60            // Functions that need full entity materialization
61            ("KEYS", full_entity),
62            ("PROPERTIES", full_entity),
63            ("LABELS", full_entity),
64            ("NODES", full_entity),
65            ("RELATIONSHIPS", full_entity),
66            // Functions that take entity arg but don't need full entity
67            ("COUNT", entity_arg_only),
68            // Functions where properties are extracted from PropertyAccess
69            ("COALESCE", no_entity),
70            ("SUM", no_entity),
71            ("AVG", no_entity),
72            ("MIN", no_entity),
73            ("MAX", no_entity),
74            ("COLLECT", no_entity),
75            ("PERCENTILEDISC", no_entity),
76            ("PERCENTILECONT", no_entity),
77        ])
78    });
79
80/// Look up the property specification for a function by name (case-insensitive).
81pub fn get_function_spec(name: &str) -> Option<&'static FunctionPropertySpec> {
82    let name_upper = name.to_uppercase();
83    FUNCTION_SPECS.get(name_upper.as_str())
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89
90    #[test]
91    fn test_validat_spec() {
92        let spec = get_function_spec("uni.temporal.validAt").unwrap();
93        assert_eq!(spec.entity_args, &[0]);
94        assert_eq!(spec.property_name_args, &[(1, 0), (2, 0)]);
95        assert!(!spec.needs_full_entity);
96    }
97
98    #[test]
99    fn test_keys_spec() {
100        let spec = get_function_spec("keys").unwrap();
101        assert_eq!(spec.entity_args, &[0]);
102        assert!(spec.needs_full_entity);
103    }
104
105    #[test]
106    fn test_properties_spec() {
107        let spec = get_function_spec("PROPERTIES").unwrap();
108        assert_eq!(spec.entity_args, &[0]);
109        assert!(spec.needs_full_entity);
110    }
111
112    #[test]
113    fn test_unknown_function_returns_none() {
114        assert!(get_function_spec("unknownFunction").is_none());
115    }
116
117    #[test]
118    fn test_count_spec_exists() {
119        let spec = get_function_spec("COUNT").unwrap();
120        assert!(!spec.needs_full_entity);
121        assert_eq!(spec.entity_args, &[0]);
122    }
123
124    #[test]
125    fn test_all_aggregates_registered() {
126        for func in ["COUNT", "SUM", "AVG", "MIN", "MAX", "COLLECT"] {
127            let spec = get_function_spec(func);
128            assert!(
129                spec.is_some(),
130                "Aggregate function {} should be registered",
131                func
132            );
133            assert!(
134                !spec.unwrap().needs_full_entity,
135                "Aggregate function {} should not need full entity",
136                func
137            );
138        }
139    }
140
141    #[test]
142    fn test_aggregate_case_insensitive() {
143        // Test that aggregate functions work with different case
144        assert!(get_function_spec("count").is_some());
145        assert!(get_function_spec("Count").is_some());
146        assert!(get_function_spec("COUNT").is_some());
147    }
148}