impl stratified

postgresml · montanalow · Feb 29, 2024 · Jan 17, 2024 · Jan 17, 2024 · Jan 17, 2024
commit accaab001f02130048871bf2ab6b407660683452
diff --git a/pgml-extension/src/orm/sampling.rs b/pgml-extension/src/orm/sampling.rs
@@ -1,11 +1,14 @@
 use pgrx::*;
 use serde::Deserialize;
 
+use super::snapshot::Column;
+
 #[derive(PostgresEnum, Copy, Clone, Eq, PartialEq, Debug, Deserialize)]
 #[allow(non_camel_case_types)]
 pub enum Sampling {
     random,
     last,
+    stratified_random,
 }
 
 impl std::str::FromStr for Sampling {
@@ -15,6 +18,7 @@ impl std::str::FromStr for Sampling {
         match input {
             "random" => Ok(Sampling::random),
             "last" => Ok(Sampling::last),
+            "stratified_random" => Ok(Sampling::stratified_random),
             _ => Err(()),
         }
     }
@@ -25,6 +29,111 @@ impl std::string::ToString for Sampling {
         match *self {
             Sampling::random => "random".to_string(),
             Sampling::last => "last".to_string(),
+            Sampling::stratified_random => "stratified_random".to_string(),
         }
     }
 }
+
+impl Sampling {
+    // Implementing the sampling strategy in SQL
+    // Effectively orders the table according to the train/test split
+    // e.g. first N rows are train, last M rows are test
+    // where M is configured by the user
+    pub fn get_sql(&self, relation_name: &str, y_column_names: Vec<Column>) -> String {
+        let col_string = y_column_names
+            .iter()
+            .map(|c| c.quoted_name())
+            .collect::<Vec<String>>()
+            .join(", ");
+        match *self {
+            Sampling::random => {
+                format!("SELECT {col_string} FROM {relation_name} ORDER BY RANDOM()")
+            }
+            Sampling::last => {
+                format!("SELECT {col_string} FROM {relation_name}")
+            }
+            Sampling::stratified_random => {
+                format!(
+                    "
+                    SELECT *
+                    FROM (
+                        SELECT 
+                            *,
+                        ROW_NUMBER() OVER(PARTITION BY {col_string} ORDER BY RANDOM()) AS rn
+                        FROM {relation_name}
+                    ) AS subquery
+                    ORDER BY rn, RANDOM();
+                "
+                )
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::orm::snapshot::{Preprocessor, Statistics};
+
+    use super::*;
+
+    fn get_column_fixtures() -> Vec<Column> {
+        vec![
+            Column {
+                name: "col1".to_string(),
+                pg_type: "text".to_string(),
+                nullable: false,
+                label: true,
+                position: 0,
+                size: 0,
+                array: false,
+                preprocessor: Preprocessor::default(),
+                statistics: Statistics::default(),
+            },
+            Column {
+                name: "col2".to_string(),
+                pg_type: "text".to_string(),
+                nullable: false,
+                label: true,
+                position: 0,
+                size: 0,
+                array: false,
+                preprocessor: Preprocessor::default(),
+                statistics: Statistics::default(),
+            },
+        ]
+    }
+
+    #[test]
+    fn test_get_sql_random_sampling() {
+        let sampling = Sampling::random;
+        let columns = get_column_fixtures();
+        let sql = sampling.get_sql("my_table", columns);
+        assert_eq!(sql, "SELECT \"col1\", \"col2\" FROM my_table ORDER BY RANDOM()");
+    }
+
+    #[test]
+    fn test_get_sql_last_sampling() {
+        let sampling = Sampling::last;
+        let columns = get_column_fixtures();
+        let sql = sampling.get_sql("my_table", columns);
+        assert_eq!(sql, "SELECT \"col1\", \"col2\" FROM my_table");
+    }
+
+    #[test]
+    fn test_get_sql_stratified_random_sampling() {
+        let sampling = Sampling::stratified_random;
+        let columns = get_column_fixtures();
+        let sql = sampling.get_sql("my_table", columns);
+        let expected_sql = "
+                    SELECT *
+                    FROM (
+                        SELECT 
+                            *,
+                        ROW_NUMBER() OVER(PARTITION BY \"col1\", \"col2\" ORDER BY RANDOM()) AS rn
+                        FROM my_table
+                    ) AS subquery
+                    ORDER BY rn, RANDOM();
+                ";
+        assert_eq!(sql, expected_sql);
+    }
+}
diff --git a/pgml-extension/src/orm/snapshot.rs b/pgml-extension/src/orm/snapshot.rs
@@ -119,7 +119,7 @@ pub(crate) struct Preprocessor {
 }
 
 #[derive(Debug, PartialEq, Serialize, Deserialize, Clone)]
-pub(crate) struct Column {
+pub struct Column {
     pub(crate) name: String,
     pub(crate) pg_type: String,
     pub(crate) nullable: bool,
@@ -147,7 +147,7 @@ impl Column {
         )
     }
 
-    fn quoted_name(&self) -> String {
+    pub(crate) fn quoted_name(&self) -> String {
         format!(r#""{}""#, self.name)
     }
 
@@ -608,13 +608,8 @@ impl Snapshot {
             };
 
             if materialized {
-                let mut sql = format!(
-                    r#"CREATE TABLE "pgml"."snapshot_{}" AS SELECT * FROM {}"#,
-                    s.id, s.relation_name
-                );
-                if s.test_sampling == Sampling::random {
-                    sql += " ORDER BY random()";
-                }
+                let sampled_query = s.test_sampling.get_sql(&s.relation_name, s.columns.clone());
+                let sql = format!(r#"CREATE TABLE "pgml"."snapshot_{}" AS {}"#, s.id, sampled_query);
                 client.update(&sql, None, None).unwrap();
             }
             snapshot = Some(s);
@@ -742,26 +737,7 @@ impl Snapshot {
     }
 
     fn select_sql(&self) -> String {
-        format!(
-            "SELECT {} FROM {} {}",
-            self.columns
-                .iter()
-                .map(|c| c.quoted_name())
-                .collect::<Vec<String>>()
-                .join(", "),
-            self.relation_name(),
-            match self.materialized {
-                // If the snapshot is materialized, we already randomized it.
-                true => "",
-                false => {
-                    if self.test_sampling == Sampling::random {
-                        "ORDER BY random()"
-                    } else {
-                        ""
-                    }
-                }
-            },
-        )
+        self.test_sampling.get_sql(&self.relation_name(), self.columns.clone())
     }
 
     fn train_test_split(&self, num_rows: usize) -> (usize, usize) {