fix: restore KLL UDF impl mode rendering (#297)

Author: GnaneshGnani

asap-summary-ingest/run_arroyosketch.py

@@ -952,6 +952,10 @@ def main(args):
             parameters["impl_mode"] = getattr(
                 args, "sketch_cmwh_impl", "legacy"
             ).capitalize()
+        elif agg_function in ("datasketcheskll_", "hydrakll_"):
+            parameters["impl_mode"] = getattr(
+                args, "sketch_kll_impl", "sketchlib"
+            ).capitalize()
 
         sql_queries.append(sql_query)
         # if not is_labels_accumulator:

asap-summary-ingest/templates/udfs/datasketcheskll_.rs.j2

@@ -1,17 +1,31 @@
 /*
 [dependencies]
+dsrs = { git = "https://github.com/ProjectASAP/datasketches-rs", rev = "d748ec75c80fff21f7b24897244dd1c895df2e9a" }
 asap_sketchlib = { git = "https://github.com/ProjectASAP/asap_sketchlib" }
 arroyo-udf-plugin = "0.1"
 rmp-serde = "1.1"
 serde = { version = "1.0", features = ["derive"] }
 */
 
 use arroyo_udf_plugin::udf;
+use dsrs::KllDoubleSketch;
 use serde::{Deserialize, Serialize};
 use asap_sketchlib::KLL;
 
 const DEFAULT_K: u16 = {{ k }};
 
+enum ImplMode {
+    Legacy,
+    Sketchlib,
+}
+
+{% set _impl_mode = impl_mode | default("Sketchlib") %}
+const IMPL_MODE: ImplMode = ImplMode::{% if _impl_mode == "Legacy" or _impl_mode == "Sketchlib" %}{{ _impl_mode }}{% else %}Sketchlib{% endif %};
+
+fn use_sketchlib_for_kll() -> bool {
+    matches!(IMPL_MODE, ImplMode::Sketchlib)
+}
+
 #[derive(Serialize, Deserialize)]
 struct KllSketchData {
     k: u16,
@@ -20,11 +34,19 @@ struct KllSketchData {
 
 #[udf]
 fn datasketcheskll_(values: Vec<f64>) -> Option<Vec<u8>> {
-    let mut sketch = KLL::init_kll(DEFAULT_K as i32);
-    for &value in &values {
-        sketch.update(&value);
-    }
-    let sketch_bytes = sketch.serialize_to_bytes().ok()?;
+    let sketch_bytes = if use_sketchlib_for_kll() {
+        let mut sketch = KLL::init_kll(DEFAULT_K as i32);
+        for &value in &values {
+            sketch.update(&value);
+        }
+        sketch.serialize_to_bytes().ok()?
+    } else {
+        let mut sketch = KllDoubleSketch::with_k(DEFAULT_K);
+        for &value in &values {
+            sketch.update(value);
+        }
+        sketch.serialize().as_ref().to_vec()
+    };
     let serialized = KllSketchData {
         k: DEFAULT_K,
         sketch_bytes,

asap-summary-ingest/templates/udfs/hydrakll_.rs.j2

@@ -1,5 +1,6 @@
 /*
 [dependencies]
+dsrs = { git = "https://github.com/ProjectASAP/datasketches-rs", rev = "d748ec75c80fff21f7b24897244dd1c895df2e9a" }
 asap_sketchlib = { git = "https://github.com/ProjectASAP/asap_sketchlib" }
 arroyo-udf-plugin = "0.1"
 rmp-serde = "1.1"
@@ -8,11 +9,24 @@ xxhash-rust = { version = "0.8", features = ["xxh32"] }
 */
 
 use arroyo_udf_plugin::udf;
+use dsrs::KllDoubleSketch;
 use rmp_serde::Serializer;
 use serde::{Deserialize, Serialize};
 use asap_sketchlib::KLL;
 use xxhash_rust::xxh32::xxh32;
 
+enum ImplMode {
+    Legacy,
+    Sketchlib,
+}
+
+{% set _impl_mode = impl_mode | default("Sketchlib") %}
+const IMPL_MODE: ImplMode = ImplMode::{% if _impl_mode == "Legacy" or _impl_mode == "Sketchlib" %}{{ _impl_mode }}{% else %}Sketchlib{% endif %};
+
+fn use_sketchlib_for_kll() -> bool {
+    matches!(IMPL_MODE, ImplMode::Sketchlib)
+}
+
 const ROW_NUM: usize = {{ row_num }};
 const COL_NUM: usize = {{ col_num }};
 const DEFAULT_K: u16 = {{ k }};
@@ -32,40 +46,76 @@ struct HydraKllSketchData {
 
 #[udf]
 fn hydrakll_(keys: Vec<&str>, values: Vec<f64>) -> Option<Vec<u8>> {
-    let mut sketches: Vec<Vec<KLL>> = (0..ROW_NUM)
-        .map(|_| {
-            (0..COL_NUM)
-                .map(|_| KLL::init_kll(DEFAULT_K as i32))
-                .collect()
-        })
-        .collect();
+    let sketch_data: Vec<Vec<KllSketchData>> = if use_sketchlib_for_kll() {
+        let mut sketches: Vec<Vec<KLL>> = (0..ROW_NUM)
+            .map(|_| {
+                (0..COL_NUM)
+                    .map(|_| KLL::init_kll(DEFAULT_K as i32))
+                    .collect()
+            })
+            .collect();
 
-    for (i, &key) in keys.iter().enumerate() {
-        if i >= values.len() {
-            break;
+        for (i, &key) in keys.iter().enumerate() {
+            if i >= values.len() {
+                break;
+            }
+            let key_bytes = key.as_bytes();
+            for row in 0..ROW_NUM {
+                let hash_value = xxh32(key_bytes, row as u32);
+                let col_index = (hash_value as usize) % COL_NUM;
+                sketches[row][col_index].update(&values[i]);
+            }
         }
-        let key_bytes = key.as_bytes();
-        for row in 0..ROW_NUM {
-            let hash_value = xxh32(key_bytes, row as u32);
-            let col_index = (hash_value as usize) % COL_NUM;
-            sketches[row][col_index].update(&values[i]);
+
+        sketches
+            .iter()
+            .map(|row| {
+                row.iter()
+                    .map(|sketch| {
+                        let sketch_bytes = sketch.serialize_to_bytes().ok()?;
+                        Some(KllSketchData {
+                            k: DEFAULT_K,
+                            sketch_bytes,
+                        })
+                    })
+                    .collect::<Option<Vec<_>>>()
+            })
+            .collect::<Option<Vec<_>>>()?
+    } else {
+        let mut sketches: Vec<Vec<KllDoubleSketch>> = (0..ROW_NUM)
+            .map(|_| {
+                (0..COL_NUM)
+                    .map(|_| KllDoubleSketch::with_k(DEFAULT_K))
+                    .collect()
+            })
+            .collect();
+
+        for (i, &key) in keys.iter().enumerate() {
+            if i >= values.len() {
+                break;
+            }
+            let key_bytes = key.as_bytes();
+            for row in 0..ROW_NUM {
+                let hash_value = xxh32(key_bytes, row as u32);
+                let col_index = (hash_value as usize) % COL_NUM;
+                sketches[row][col_index].update(values[i]);
+            }
         }
-    }
 
-    let sketch_data: Vec<Vec<KllSketchData>> = sketches
-        .iter()
-        .map(|row| {
-            row.iter()
-                .map(|sketch| {
-                    let sketch_bytes = sketch.serialize_to_bytes().ok()?;
-                    Some(KllSketchData {
-                        k: DEFAULT_K,
-                        sketch_bytes,
+        sketches
+            .iter()
+            .map(|row| {
+                row.iter()
+                    .map(|sketch| {
+                        Some(KllSketchData {
+                            k: DEFAULT_K,
+                            sketch_bytes: sketch.serialize().as_ref().to_vec(),
+                        })
                     })
-                })
-                .collect::<Option<Vec<_>>>()
-        })
-        .collect::<Option<Vec<_>>>()?;
+                    .collect::<Option<Vec<_>>>()
+            })
+            .collect::<Option<Vec<_>>>()?
+    };
 
     let hydra_data = HydraKllSketchData {
         row_num: ROW_NUM,

asap-summary-ingest/tests/test_integration.py

@@ -430,5 +430,48 @@ def test_parse_promql_config_file(self):
         agg_configs[0].validate(metric_config, query_language="promql")
 
 
+class TestKllUdfTemplates:
+    """Tests for KLL UDF template rendering."""
+
+    @pytest.fixture
+    def udf_template_dir(self):
+        """Load the UDF template directory."""
+        return os.path.join(
+            os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+            "templates",
+            "udfs",
+        )
+
+    def test_template_variables_ignore_local_jinja_assignments(self):
+        """Test that local Jinja variables are not treated as required params."""
+        template_source = """
+{% set _impl_mode = impl_mode | default("Sketchlib") %}
+{{ _impl_mode }} {{ k }}
+"""
+        template_vars = jinja_utils.get_template_variables(template_source)
+
+        assert template_vars == {"impl_mode", "k"}
+
+    def test_datasketches_kll_template_renders_legacy_mode(self, udf_template_dir):
+        """Test rendering the single KLL UDF with the legacy backend."""
+        template = jinja_utils.load_template(udf_template_dir, "datasketcheskll_.rs.j2")
+        rendered = template.render(k=200, impl_mode="Legacy")
+
+        assert "const IMPL_MODE: ImplMode = ImplMode::Legacy;" in rendered
+        assert "KllDoubleSketch::with_k(DEFAULT_K)" in rendered
+        assert "asap_sketchlib" in rendered
+        assert "{{" not in rendered
+
+    def test_hydra_kll_template_renders_sketchlib_mode(self, udf_template_dir):
+        """Test rendering the Hydra KLL UDF with the sketchlib backend."""
+        template = jinja_utils.load_template(udf_template_dir, "hydrakll_.rs.j2")
+        rendered = template.render(row_num=3, col_num=128, k=200, impl_mode="Sketchlib")
+
+        assert "const IMPL_MODE: ImplMode = ImplMode::Sketchlib;" in rendered
+        assert "KLL::init_kll(DEFAULT_K as i32)" in rendered
+        assert "KllDoubleSketch::with_k(DEFAULT_K)" in rendered
+        assert "{{" not in rendered
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])

asap-summary-ingest/utils/jinja_utils.py

@@ -1,4 +1,4 @@
-from jinja2 import Environment, FileSystemLoader, nodes
+from jinja2 import Environment, FileSystemLoader, meta
 
 
 def load_template(template_dir, template_name):
@@ -23,5 +23,4 @@ def get_template_variables(template_source, environment=None):
         environment = Environment()
 
     ast = environment.parse(template_source)
-    template_vars = ast.find_all(nodes.Name)
-    return {var.name for var in template_vars if var.ctx == "load"}
+    return meta.find_undeclared_variables(ast)