|
5 | 5 | use crate::wrappers::hll::CardinalityValue; |
6 | 6 | use ::polars::prelude::*; |
7 | 7 | use std::collections::HashMap; |
| 8 | +use std::fmt::Write as _; |
8 | 9 |
|
9 | 10 | /// the resulting key→count table for O(1) per-key queries. |
10 | 11 | /// Shared by `cms/polars` and `countsketch/polars`. |
@@ -262,28 +263,39 @@ impl<T: PolarsColumnItem> PolarsQuantileCore<T> { |
262 | 263 | } |
263 | 264 | } |
264 | 265 |
|
265 | | -/// The `;`-joined key of one label subset, in column order: the subset key |
266 | | -/// asap_sketchlib 0.2.2's `Hydra::update` built, so the baseline matches a |
267 | | -/// value in any label column. 0.3.0's Hydra keys a subset by column, so the |
268 | | -/// two differ when label columns share a value domain. |
| 266 | +/// The key of the label subset `mask` picks out of `parts`, in column order, |
| 267 | +/// each value tagged with its column (`label0:a;label2:c`, with `\`, `:` and `;` |
| 268 | +/// escaped): asap_sketchlib 0.3.0's subkey over the `label{i}` schema |
| 269 | +/// `hydra_shared` declares. So a value in one column never matches the same |
| 270 | +/// value in another (#74), in the baseline as in the sketch. |
269 | 271 | pub fn subset_key(parts: &[&str], mask: usize) -> String { |
270 | 272 | let mut out = String::new(); |
271 | 273 | for (j, part) in parts.iter().enumerate() { |
272 | 274 | if (mask >> j) & 1 == 1 { |
273 | 275 | if !out.is_empty() { |
274 | 276 | out.push(';'); |
275 | 277 | } |
276 | | - out.push_str(part); |
| 278 | + let _ = write!(out, "label{j}:"); |
| 279 | + for ch in part.chars() { |
| 280 | + if matches!(ch, '\\' | ':' | ';') { |
| 281 | + out.push('\\'); |
| 282 | + } |
| 283 | + out.push(ch); |
| 284 | + } |
277 | 285 | } |
278 | 286 | } |
279 | 287 | out |
280 | 288 | } |
281 | 289 |
|
| 290 | +/// The key a probe of the leading label columns `labels` reads. |
| 291 | +pub fn prefix_key(labels: &[&str]) -> String { |
| 292 | + subset_key(labels, (1 << labels.len()) - 1) |
| 293 | +} |
| 294 | + |
282 | 295 | /// Expand one record into `(subset_key, value)` rows, one per non-empty subset |
283 | | -/// of its labels. Empty label parts are dropped, as asap_sketchlib 0.2.2's |
284 | | -/// Hydra did (0.3.0's keeps them in their column). |
| 296 | +/// of its labels. An empty label keeps its column, as 0.3.0's Hydra does. |
285 | 297 | pub fn fan_out<V: Clone>(key: &str, value: &V, keys: &mut Vec<String>, values: &mut Vec<V>) { |
286 | | - let parts: Vec<&str> = key.split(';').filter(|s| !s.is_empty()).collect(); |
| 298 | + let parts: Vec<&str> = key.split(';').collect(); |
287 | 299 | for mask in 1..(1usize << parts.len()) { |
288 | 300 | keys.push(subset_key(&parts, mask)); |
289 | 301 | values.push(value.clone()); |
@@ -350,7 +362,7 @@ impl<T: PolarsFrequencyItem> PolarsSubpopFrequencyCore<T> { |
350 | 362 |
|
351 | 363 | pub fn query(&self, labels: &[&str], value: &T) -> f64 { |
352 | 364 | self.counts |
353 | | - .get(&(labels.join(";"), value.count_key())) |
| 365 | + .get(&(prefix_key(labels), value.count_key())) |
354 | 366 | .copied() |
355 | 367 | .unwrap_or(0) as f64 |
356 | 368 | } |
@@ -415,7 +427,7 @@ impl<T: CardinalityValue> PolarsSubpopCardinalityCore<T> { |
415 | 427 | } |
416 | 428 |
|
417 | 429 | pub fn query(&self, labels: &[&str]) -> f64 { |
418 | | - self.distinct.get(&labels.join(";")).copied().unwrap_or(0) as f64 |
| 430 | + self.distinct.get(&prefix_key(labels)).copied().unwrap_or(0) as f64 |
419 | 431 | } |
420 | 432 |
|
421 | 433 | pub fn memory_bytes(&self) -> usize { |
@@ -484,7 +496,7 @@ impl<T: CardinalityValue> PolarsSubpopVectorCore<T> { |
484 | 496 | } |
485 | 497 |
|
486 | 498 | pub fn query(&self, labels: &[&str]) -> f64 { |
487 | | - self.exact.get(&labels.join(";")).copied().unwrap_or(0.0) |
| 499 | + self.exact.get(&prefix_key(labels)).copied().unwrap_or(0.0) |
488 | 500 | } |
489 | 501 |
|
490 | 502 | pub fn memory_bytes(&self) -> usize { |
@@ -521,3 +533,28 @@ fn grouped_values<T: CardinalityValue>( |
521 | 533 | .expect("polars group_by collect"), |
522 | 534 | ) |
523 | 535 | } |
| 536 | + |
| 537 | +#[cfg(test)] |
| 538 | +mod tests { |
| 539 | + use super::*; |
| 540 | + |
| 541 | + #[test] |
| 542 | + fn subset_keys_tag_each_value_with_its_column() { |
| 543 | + assert_eq!(subset_key(&["a", "b", "c"], 0b101), "label0:a;label2:c"); |
| 544 | + assert_eq!(subset_key(&["x:y;z", ""], 0b11), r"label0:x\:y\;z;label1:"); |
| 545 | + assert_eq!(prefix_key(&["a", "b"]), "label0:a;label1:b"); |
| 546 | + } |
| 547 | + |
| 548 | + /// #74: the column-0 group `a` holds 2 records (`a;a`, `a;b`). The old |
| 549 | + /// untagged key counted 3: `a;a`'s `a` once per column. |
| 550 | + #[test] |
| 551 | + fn a_value_shared_by_two_columns_is_counted_once() { |
| 552 | + let mut core = PolarsSubpopFrequencyCore::<i64>::default(); |
| 553 | + core.update(&("a;a".to_string(), 7)); |
| 554 | + core.update(&("a;b".to_string(), 7)); |
| 555 | + core.finalize(); |
| 556 | + assert_eq!(core.query(&["a"], &7), 2.0); |
| 557 | + assert_eq!(core.query(&["a", "a"], &7), 1.0); |
| 558 | + assert_eq!(core.query(&["b"], &7), 0.0); |
| 559 | + } |
| 560 | +} |
0 commit comments