Hot fixing 2.0 (#409)

Author: levkk

pgml-extension/sql/schema.sql

@@ -2,6 +2,7 @@
 --- Validate we have the necessary Python dependencies.
 ---
 SELECT pgml.validate_python_dependencies();
+SELECT pgml.validate_shared_library();
 
 --- 
 --- Track of updates to data

pgml-extension/src/api.rs

@@ -60,6 +60,21 @@ pub fn validate_python_dependencies() {
 #[pg_extern]
 pub fn validate_python_dependencies() {}
 
+#[pg_extern]
+pub fn validate_shared_library() {
+    let shared_preload_libraries: String = Spi::get_one(
+        "SELECT setting
+         FROM pg_settings
+         WHERE name = 'shared_preload_libraries'
+         LIMIT 1",
+    )
+    .unwrap();
+
+    if !shared_preload_libraries.contains("pgml") {
+        error!("`pgml` must be added to `shared_preload_libraries` setting or models cannot be deployed");
+    }
+}
+
 #[cfg(feature = "python")]
 #[pg_extern]
 pub fn sklearn_version() -> String {
@@ -102,6 +117,7 @@ pub fn python_version() -> String {
 fn version() -> String {
     crate::VERSION.to_string()
 }
+
 #[allow(clippy::too_many_arguments)]
 #[pg_extern]
 fn train(

pgml-extension/src/bindings/xgboost.rs

@@ -9,6 +9,7 @@ use xgboost::parameters::*;
 use xgboost::{Booster, DMatrix};
 
 use crate::orm::dataset::Dataset;
+use crate::orm::task::Task;
 use crate::orm::Hyperparams;
 
 use crate::bindings::Bindings;
@@ -133,20 +134,27 @@ fn get_tree_params(hyperparams: &Hyperparams) -> tree::TreeBoosterParameters {
 }
 
 pub fn fit_regression(dataset: &Dataset, hyperparams: &Hyperparams) -> Box<dyn Bindings> {
-    fit(dataset, hyperparams, learning::Objective::RegLinear)
+    fit(
+        dataset,
+        hyperparams,
+        Task::regression,
+        learning::Objective::RegLinear,
+    )
 }
 
 pub fn fit_classification(dataset: &Dataset, hyperparams: &Hyperparams) -> Box<dyn Bindings> {
     fit(
         dataset,
         hyperparams,
+        Task::classification,
         learning::Objective::MultiSoftprob(dataset.num_distinct_labels.try_into().unwrap()),
     )
 }
 
 fn fit(
     dataset: &Dataset,
     hyperparams: &Hyperparams,
+    task: Task,
     objective: learning::Objective,
 ) -> Box<dyn Bindings> {
     // split the train/test data into DMatrix
@@ -205,7 +213,11 @@ fn fit(
     Box::new(Estimator {
         estimator: booster,
         num_features: dataset.num_features,
-        num_classes: dataset.num_distinct_labels,
+        num_classes: if task == Task::regression {
+            1
+        } else {
+            dataset.num_distinct_labels
+        },
     })
 }
 

pgml-extension/src/orm/dataset.rs

@@ -22,8 +22,8 @@ impl Display for Dataset {
     fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
         write!(
             f,
-            "Dataset {{ num_features: {}, num_labels: {}, num_rows: {}, num_train_rows: {}, num_test_rows: {} }}",
-            self.num_features, self.num_labels, self.num_rows, self.num_train_rows, self.num_test_rows,
+            "Dataset {{ num_features: {}, num_labels: {}, num_distinct_labels: {}, num_rows: {}, num_train_rows: {}, num_test_rows: {} }}",
+            self.num_features, self.num_labels, self.num_distinct_labels, self.num_rows, self.num_train_rows, self.num_test_rows,
         )
     }
 }

pgml-extension/src/orm/snapshot.rs

@@ -196,6 +196,11 @@ impl Snapshot {
                 let nullable = row[3].value::<bool>().unwrap();
                 let position = row[4].value::<i32>().unwrap() as usize;
                 let label = self.y_column_name.contains(&name);
+
+                if nullable {
+                    warning!("Column \"{}\" can contain nulls which can cause errors", name);
+                }
+
                 columns.push(
                     Column {
                         name,