From ad41d19ffadbd6bc1387524b6ba1f3ac9e0c1741 Mon Sep 17 00:00:00 2001 From: Montana Low Date: Tue, 12 Sep 2023 11:18:38 -0700 Subject: [PATCH 1/2] handle errors cleanly rather than crashing --- pgml-extension/src/metrics.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/pgml-extension/src/metrics.rs b/pgml-extension/src/metrics.rs index 4526d971e..1e726efe3 100644 --- a/pgml-extension/src/metrics.rs +++ b/pgml-extension/src/metrics.rs @@ -51,13 +51,17 @@ impl ConfusionMatrix { y_hat: &ArrayView1, num_classes: usize, ) -> ConfusionMatrix { - assert_eq!(ground_truth.len(), y_hat.len()); + if ground_truth.len() != y_hat.len() { + error!("Can't compute metrics when the ground truth labels are a different size than the predicted labels. {} != {}", ground_truth.len(), y_hat.len()) + }; // Distinct classes. let mut classes = ground_truth.iter().collect::>(); classes.extend(&mut y_hat.iter().collect::>().into_iter()); - assert_eq!(num_classes, classes.len()); + if num_classes != classes.len() { + error!("Can't compute metrics when the number of classes in the test set is different than the number of classes in the training set. {} != {}", num_classes, classes.len()) + }; // Class value = index in the confusion matrix // e.g. class value 5 will be index 4 if there are classes 1, 2, 3 and 4 present. From bf7e1acd10ff1b61034ace964cf1f5bcd573f391 Mon Sep 17 00:00:00 2001 From: Montana Low Date: Tue, 12 Sep 2023 11:22:00 -0700 Subject: [PATCH 2/2] use --- pgml-extension/src/metrics.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/pgml-extension/src/metrics.rs b/pgml-extension/src/metrics.rs index 1e726efe3..c23f27b64 100644 --- a/pgml-extension/src/metrics.rs +++ b/pgml-extension/src/metrics.rs @@ -1,4 +1,5 @@ /// Module providing various metrics used to rank the algorithms. +use pgrx::*; use std::collections::{BTreeSet, HashMap}; use ndarray::{Array2, ArrayView1};