polars: tag Hydra baseline subset keys by column (#74)

configs/datagen/hydra_columns.yaml

@@ -30,12 +30,10 @@
 # comparable as `d` changes.
 #
 # The columns here share the default alphabet, so their value domains overlap.
-# The sketch is not affected: asap_sketchlib 0.3.0 keys each subset by column
-# (`key1:abc`), so a key1 value never collides with the same key2 value. The
-# `polars` exact baselines still build untagged subset keys (the 0.2.2
-# encoding, `polars_shared.rs`), so on this file they merge the key1 and key2
-# groups and do not score zero (#74; docs/rqe_optimizer_hydra.md §5.2). Give
-# each column its own `string.alphabet` and the baselines drop to exactly zero.
+# Neither side aliases them: asap_sketchlib 0.3.0 keys each subset by column
+# (`label0:abc`), and the `polars` exact baselines build the same tagged keys
+# (`polars_shared::subset_key`), so a key1 value never matches the same key2
+# value and the baselines score exactly zero (#74).
 
 column_num: 3
 column_label: [key1, key2, value]

configs/datagen/hydra_kll_disjoint_labels.yaml

@@ -1,10 +1,9 @@
 # `hydra_columns_f64.yaml` with a separate alphabet per label column, so a key1
-# value can never render the same as a key2 value (#74). The sketch keys each
-# subset by column (asap_sketchlib 0.3.0), so it doesn't need this; the
-# `polars` baseline still builds untagged subset keys, and on the shared
-# alphabet it merges a key1 group with the matching key2 group. Here the
-# baseline scores zero and any rank error is the sketch's. `export_rqe_optimizer_costs.sh` priced `hydra-kll` on it
-# until #174; HydraKLL has no saturation curve, so nothing measures it now.
+# value can never render the same as a key2 value. Since #74 the sketch and the
+# `polars` baseline both key subsets by column, so the shared alphabet no longer
+# aliases either; this file stays as the disjoint-domain case.
+# `export_rqe_optimizer_costs.sh` priced `hydra-kll` on it until #174; HydraKLL
+# has no saturation curve, so nothing measures it now.
 
 column_num: 3
 column_label: [key1, key2, value]

sketch-bench/src/wrappers/hydra_kll/polars.rs

@@ -47,7 +47,7 @@ impl<T: QuantileValue + PolarsColumnItem> PolarsSubpopQuantile<T> {
     /// rank interval contains `phi * n`, which is what the rank-error
     /// comparator scores against, so an exact answer scores zero.
     pub fn estimate_subpop_quantile(&self, labels: &[&str], phi: f64) -> f64 {
-        let Some(values) = self.sorted.get(&labels.join(";")) else {
+        let Some(values) = self.sorted.get(&prefix_key(labels)) else {
             return f64::NAN;
         };
         if values.is_empty() {

sketch-bench/src/wrappers/polars_shared.rs

@@ -5,6 +5,7 @@
 use crate::wrappers::hll::CardinalityValue;
 use ::polars::prelude::*;
 use std::collections::HashMap;
+use std::fmt::Write as _;
 
 /// the resulting key→count table for O(1) per-key queries.
 /// Shared by `cms/polars` and `countsketch/polars`.
@@ -262,28 +263,39 @@ impl<T: PolarsColumnItem> PolarsQuantileCore<T> {
     }
 }
 
-/// The `;`-joined key of one label subset, in column order: the subset key
-/// asap_sketchlib 0.2.2's `Hydra::update` built, so the baseline matches a
-/// value in any label column. 0.3.0's Hydra keys a subset by column, so the
-/// two differ when label columns share a value domain.
+/// The key of the label subset `mask` picks out of `parts`, in column order,
+/// each value tagged with its column (`label0:a;label2:c`, with `\`, `:` and `;`
+/// escaped): asap_sketchlib 0.3.0's subkey over the `label{i}` schema
+/// `hydra_shared` declares. So a value in one column never matches the same
+/// value in another (#74), in the baseline as in the sketch.
 pub fn subset_key(parts: &[&str], mask: usize) -> String {
     let mut out = String::new();
     for (j, part) in parts.iter().enumerate() {
         if (mask >> j) & 1 == 1 {
             if !out.is_empty() {
                 out.push(';');
             }
-            out.push_str(part);
+            let _ = write!(out, "label{j}:");
+            for ch in part.chars() {
+                if matches!(ch, '\\' | ':' | ';') {
+                    out.push('\\');
+                }
+                out.push(ch);
+            }
         }
     }
     out
 }
 
+/// The key a probe of the leading label columns `labels` reads.
+pub fn prefix_key(labels: &[&str]) -> String {
+    subset_key(labels, (1 << labels.len()) - 1)
+}
+
 /// Expand one record into `(subset_key, value)` rows, one per non-empty subset
-/// of its labels. Empty label parts are dropped, as asap_sketchlib 0.2.2's
-/// Hydra did (0.3.0's keeps them in their column).
+/// of its labels. An empty label keeps its column, as 0.3.0's Hydra does.
 pub fn fan_out<V: Clone>(key: &str, value: &V, keys: &mut Vec<String>, values: &mut Vec<V>) {
-    let parts: Vec<&str> = key.split(';').filter(|s| !s.is_empty()).collect();
+    let parts: Vec<&str> = key.split(';').collect();
     for mask in 1..(1usize << parts.len()) {
         keys.push(subset_key(&parts, mask));
         values.push(value.clone());
@@ -350,7 +362,7 @@ impl<T: PolarsFrequencyItem> PolarsSubpopFrequencyCore<T> {
 
     pub fn query(&self, labels: &[&str], value: &T) -> f64 {
         self.counts
-            .get(&(labels.join(";"), value.count_key()))
+            .get(&(prefix_key(labels), value.count_key()))
             .copied()
             .unwrap_or(0) as f64
     }
@@ -415,7 +427,7 @@ impl<T: CardinalityValue> PolarsSubpopCardinalityCore<T> {
     }
 
     pub fn query(&self, labels: &[&str]) -> f64 {
-        self.distinct.get(&labels.join(";")).copied().unwrap_or(0) as f64
+        self.distinct.get(&prefix_key(labels)).copied().unwrap_or(0) as f64
     }
 
     pub fn memory_bytes(&self) -> usize {
@@ -484,7 +496,7 @@ impl<T: CardinalityValue> PolarsSubpopVectorCore<T> {
     }
 
     pub fn query(&self, labels: &[&str]) -> f64 {
-        self.exact.get(&labels.join(";")).copied().unwrap_or(0.0)
+        self.exact.get(&prefix_key(labels)).copied().unwrap_or(0.0)
     }
 
     pub fn memory_bytes(&self) -> usize {
@@ -521,3 +533,28 @@ fn grouped_values<T: CardinalityValue>(
             .expect("polars group_by collect"),
     )
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn subset_keys_tag_each_value_with_its_column() {
+        assert_eq!(subset_key(&["a", "b", "c"], 0b101), "label0:a;label2:c");
+        assert_eq!(subset_key(&["x:y;z", ""], 0b11), r"label0:x\:y\;z;label1:");
+        assert_eq!(prefix_key(&["a", "b"]), "label0:a;label1:b");
+    }
+
+    /// #74: the column-0 group `a` holds 2 records (`a;a`, `a;b`). The old
+    /// untagged key counted 3: `a;a`'s `a` once per column.
+    #[test]
+    fn a_value_shared_by_two_columns_is_counted_once() {
+        let mut core = PolarsSubpopFrequencyCore::<i64>::default();
+        core.update(&("a;a".to_string(), 7));
+        core.update(&("a;b".to_string(), 7));
+        core.finalize();
+        assert_eq!(core.query(&["a"], &7), 2.0);
+        assert_eq!(core.query(&["a", "a"], &7), 1.0);
+        assert_eq!(core.query(&["b"], &7), 0.0);
+    }
+}