perf(python): add pre-eval cache and context key filtering to PyO3 bindings

aepfli · claude · aepfli · commit 41888548888c · 2026-02-10T21:57:14.000+01:00
Add host-side optimizations to the PyO3 FlagEvaluator matching the WASM
and Go implementations:
- Pre-evaluated cache for static/disabled flags (skip Rust evaluation)
- Context key filtering: extract only required keys from PyDict before
  converting to serde_json::Value, avoiding full dict serialization
- evaluate_flag_pre_enriched path for pre-enriched filtered contexts

Co-Authored-By: Claude Opus 4.6 &lt;noreply@anthropic.com&gt;
diff --git a/python/src/lib.rs b/python/src/lib.rs
@@ -1,7 +1,8 @@
-use ::flagd_evaluator::ValidationMode;
+use ::flagd_evaluator::{EvaluationResult, ValidationMode};
 use pyo3::prelude::*;
 use pyo3::types::PyDict;
-use serde_json::Value;
+use serde_json::{Map, Value};
+use std::collections::{HashMap, HashSet};
 
 /// FlagEvaluator - Stateful feature flag evaluator
 ///
@@ -26,6 +27,84 @@ use serde_json::Value;
 struct FlagEvaluator {
     /// Wrap the Rust FlagEvaluator directly - no duplication!
     inner: ::flagd_evaluator::FlagEvaluator,
+    /// Cache of pre-evaluated results for static and disabled flags
+    pre_evaluated: HashMap<String, EvaluationResult>,
+    /// Required context keys per flag for context filtering
+    required_context_keys: HashMap<String, HashSet<String>>,
+}
+
+/// Build a filtered, pre-enriched context containing only the required keys.
+///
+/// Extracts only the keys in `required_keys` from the PyDict, adds `targetingKey`
+/// (defaulting to `""`) and `$flagd` enrichment with `flagKey` and `timestamp`.
+fn build_filtered_context(
+    flag_key: &str,
+    context: &Bound<'_, PyDict>,
+    required_keys: &HashSet<String>,
+) -> PyResult<Value> {
+    let mut map = Map::new();
+
+    // Extract only required keys from PyDict
+    for key in required_keys {
+        if let Some(py_val) = context.get_item(key)? {
+            let val: Value = pythonize::depythonize(&py_val)?;
+            map.insert(key.clone(), val);
+        }
+    }
+
+    // Ensure targetingKey is present
+    if !map.contains_key("targetingKey") {
+        if let Some(py_val) = context.get_item("targetingKey")? {
+            let val: Value = pythonize::depythonize(&py_val)?;
+            map.insert("targetingKey".to_string(), val);
+        } else {
+            map.insert("targetingKey".to_string(), Value::String(String::new()));
+        }
+    }
+
+    // Add $flagd enrichment
+    let timestamp = std::time::SystemTime::now()
+        .duration_since(std::time::UNIX_EPOCH)
+        .unwrap_or_default()
+        .as_secs();
+
+    let mut flagd_props = Map::new();
+    flagd_props.insert("flagKey".to_string(), Value::String(flag_key.to_string()));
+    flagd_props.insert("timestamp".to_string(), Value::Number(timestamp.into()));
+    map.insert("$flagd".to_string(), Value::Object(flagd_props));
+
+    Ok(Value::Object(map))
+}
+
+/// Resolve the context for evaluation: use filtered context if required keys are
+/// known, otherwise fall back to full depythonize.
+fn resolve_context(
+    flag_key: &str,
+    context: &Bound<'_, PyDict>,
+    required_context_keys: &HashMap<String, HashSet<String>>,
+) -> PyResult<(Value, bool)> {
+    if let Some(required_keys) = required_context_keys.get(flag_key) {
+        let filtered = build_filtered_context(flag_key, context, required_keys)?;
+        Ok((filtered, true))
+    } else {
+        let full: Value = pythonize::depythonize(context.as_any())?;
+        Ok((full, false))
+    }
+}
+
+/// Call the appropriate Rust evaluate method depending on whether context
+/// is pre-enriched.
+fn do_evaluate(
+    inner: &::flagd_evaluator::FlagEvaluator,
+    flag_key: &str,
+    context: &Value,
+    pre_enriched: bool,
+) -> EvaluationResult {
+    if pre_enriched {
+        inner.evaluate_flag_pre_enriched(flag_key, context)
+    } else {
+        inner.evaluate_flag(flag_key, context)
+    }
 }
 
 #[pymethods]
@@ -47,6 +126,8 @@ impl FlagEvaluator {
 
         FlagEvaluator {
             inner: ::flagd_evaluator::FlagEvaluator::new(mode),
+            pre_evaluated: HashMap::new(),
+            required_context_keys: HashMap::new(),
         }
     }
 
@@ -77,6 +158,20 @@ impl FlagEvaluator {
             ))
         })?;
 
+        // Extract pre-evaluated cache
+        self.pre_evaluated = response.pre_evaluated.clone().unwrap_or_default();
+
+        // Extract required context keys (convert Vec<String> to HashSet<String>)
+        self.required_context_keys = response
+            .required_context_keys
+            .as_ref()
+            .map(|m| {
+                m.iter()
+                    .map(|(k, v)| (k.clone(), v.iter().cloned().collect()))
+                    .collect()
+            })
+            .unwrap_or_default();
+
         // Convert response to Python dict
         pythonize::pythonize(py, &response)
             .map(|bound| bound.unbind())
@@ -102,11 +197,23 @@ impl FlagEvaluator {
         flag_key: String,
         context: &Bound<'_, PyDict>,
     ) -> PyResult<PyObject> {
-        // Convert context to JSON Value
-        let context_value: Value = pythonize::depythonize(context.as_any())?;
+        // Check pre-evaluation cache first
+        if let Some(cached) = self.pre_evaluated.get(&flag_key) {
+            return pythonize::pythonize(py, cached)
+                .map(|bound| bound.unbind())
+                .map_err(|e| {
+                    PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
+                        "Failed to convert result: {}",
+                        e
+                    ))
+                });
+        }
 
-        // Delegate to the Rust FlagEvaluator
-        let result = self.inner.evaluate_flag(&flag_key, &context_value);
+        // Resolve context (filtered or full)
+        let (context_value, pre_enriched) =
+            resolve_context(&flag_key, context, &self.required_context_keys)?;
+
+        let result = do_evaluate(&self.inner, &flag_key, &context_value, pre_enriched);
 
         // Convert result to Python dict
         pythonize::pythonize(py, &result)
@@ -134,10 +241,21 @@ impl FlagEvaluator {
         context: &Bound<'_, PyDict>,
         default_value: bool,
     ) -> PyResult<bool> {
-        let context_value: Value = pythonize::depythonize(context.as_any())?;
-        let result = self.inner.evaluate_bool(&flag_key, &context_value);
+        // Check pre-evaluation cache first
+        if let Some(cached) = self.pre_evaluated.get(&flag_key) {
+            if cached.error_code.is_some() {
+                return Ok(default_value);
+            }
+            return match &cached.value {
+                Value::Bool(b) => Ok(*b),
+                _ => Ok(default_value),
+            };
+        }
+
+        let (context_value, pre_enriched) =
+            resolve_context(&flag_key, context, &self.required_context_keys)?;
+        let result = do_evaluate(&self.inner, &flag_key, &context_value, pre_enriched);
 
-        // If there's an error, return the default value
         if result.error_code.is_some() {
             return Ok(default_value);
         }
@@ -163,10 +281,21 @@ impl FlagEvaluator {
         context: &Bound<'_, PyDict>,
         default_value: String,
     ) -> PyResult<String> {
-        let context_value: Value = pythonize::depythonize(context.as_any())?;
-        let result = self.inner.evaluate_string(&flag_key, &context_value);
+        // Check pre-evaluation cache first
+        if let Some(cached) = self.pre_evaluated.get(&flag_key) {
+            if cached.error_code.is_some() {
+                return Ok(default_value);
+            }
+            return match &cached.value {
+                Value::String(s) => Ok(s.clone()),
+                _ => Ok(default_value),
+            };
+        }
+
+        let (context_value, pre_enriched) =
+            resolve_context(&flag_key, context, &self.required_context_keys)?;
+        let result = do_evaluate(&self.inner, &flag_key, &context_value, pre_enriched);
 
-        // If there's an error, return the default value
         if result.error_code.is_some() {
             return Ok(default_value);
         }
@@ -192,10 +321,21 @@ impl FlagEvaluator {
         context: &Bound<'_, PyDict>,
         default_value: i64,
     ) -> PyResult<i64> {
-        let context_value: Value = pythonize::depythonize(context.as_any())?;
-        let result = self.inner.evaluate_int(&flag_key, &context_value);
+        // Check pre-evaluation cache first
+        if let Some(cached) = self.pre_evaluated.get(&flag_key) {
+            if cached.error_code.is_some() {
+                return Ok(default_value);
+            }
+            return match &cached.value {
+                Value::Number(n) => Ok(n.as_i64().unwrap_or(default_value)),
+                _ => Ok(default_value),
+            };
+        }
+
+        let (context_value, pre_enriched) =
+            resolve_context(&flag_key, context, &self.required_context_keys)?;
+        let result = do_evaluate(&self.inner, &flag_key, &context_value, pre_enriched);
 
-        // If there's an error, return the default value
         if result.error_code.is_some() {
             return Ok(default_value);
         }
@@ -221,10 +361,21 @@ impl FlagEvaluator {
         context: &Bound<'_, PyDict>,
         default_value: f64,
     ) -> PyResult<f64> {
-        let context_value: Value = pythonize::depythonize(context.as_any())?;
-        let result = self.inner.evaluate_float(&flag_key, &context_value);
+        // Check pre-evaluation cache first
+        if let Some(cached) = self.pre_evaluated.get(&flag_key) {
+            if cached.error_code.is_some() {
+                return Ok(default_value);
+            }
+            return match &cached.value {
+                Value::Number(n) => Ok(n.as_f64().unwrap_or(default_value)),
+                _ => Ok(default_value),
+            };
+        }
+
+        let (context_value, pre_enriched) =
+            resolve_context(&flag_key, context, &self.required_context_keys)?;
+        let result = do_evaluate(&self.inner, &flag_key, &context_value, pre_enriched);
 
-        // If there's an error, return the default value
         if result.error_code.is_some() {
             return Ok(default_value);
         }
