postgresml
diff --git a/‎pgml-extension/pgml_rust/src/api.rs
Lines changed: 1 addition & 1 deletion b/‎pgml-extension/pgml_rust/src/api.rs
Lines changed: 1 addition & 1 deletion
diff --git a/‎pgml-extension/pgml_rust/src/engines/mod.rs
Lines changed: 1 addition & 0 deletions b/‎pgml-extension/pgml_rust/src/engines/mod.rs
Lines changed: 1 addition & 0 deletions
diff --git a/‎pgml-extension/pgml_rust/src/engines/sklearn.rs
Lines changed: 16 additions & 5 deletions b/‎pgml-extension/pgml_rust/src/engines/sklearn.rs
Lines changed: 16 additions & 5 deletions
@@ -40,7 +40,7 @@ fn train(
     search_args: default!(JsonB, "'{}'"),
     test_size: default!(f32, 0.25),
     test_sampling: default!(Sampling, "'last'"),
-    engine: default!(Engine, "'sklearn'"),
+    engine: Option<default!(Engine, "NULL")>,
 ) -> impl std::iter::Iterator<
     Item = (
         name!(project, String),
 
@@ -1,3 +1,4 @@
 pub mod engine;
 pub mod sklearn;
 pub mod smartcore;
+pub mod xgboost;
@@ -1,14 +1,22 @@
-use pgx::*;
+/// Scikit-Learn implementation.
+///
+/// Scikit needs no introduction. It implements dozens of industry-standard
+/// algorithms used in data science and machine learning.
+///
+/// It uses numpy as its dense matrix.
+///
+/// Our implementation below calls into Python wrappers
+/// defined in `src/engines/wrappers.py`.
 use pyo3::prelude::*;
 use pyo3::types::PyTuple;
 
-use std::collections::HashMap;
-
 use crate::orm::algorithm::Algorithm;
 use crate::orm::dataset::Dataset;
 use crate::orm::estimator::SklearnBox;
 use crate::orm::task::Task;
 
+use pgx::*;
+
 #[pg_extern]
 pub fn sklearn_version() -> String {
     let mut version = String::new();
@@ -25,7 +33,7 @@ pub fn sklearn_train(
     task: Task,
     algorithm: Algorithm,
     dataset: &Dataset,
-    hyperparams: &JsonB,
+    hyperparams: &serde_json::Map<std::string::String, serde_json::Value>,
 ) -> SklearnBox {
     let module = include_str!(concat!(
         env!("CARGO_MANIFEST_DIR"),
@@ -75,12 +83,15 @@ pub fn sklearn_train(
     SklearnBox::new(estimator)
 }
 
-pub fn sklearn_test(estimator: &SklearnBox, x_test: &[f32], num_features: usize) -> Vec<f32> {
+pub fn sklearn_test(estimator: &SklearnBox, dataset: &Dataset) -> Vec<f32> {
     let module = include_str!(concat!(
         env!("CARGO_MANIFEST_DIR"),
         "/src/engines/wrappers.py"
     ));
 
+    let x_test = dataset.x_test();
+    let num_features = dataset.num_features;
+
     let y_hat: Vec<f32> = Python::with_gil(|py| -> Vec<f32> {
         let module = PyModule::from_code(py, module, "", "").unwrap();
         let predictor = module.getattr("predictor").unwrap();