disambiguate api (#566)

montanalow · web-flow · commit 5a54db42aa72 · 2023-03-15T11:30:33.000-07:00
diff --git a/pgml-extension/examples/image_classification.sql b/pgml-extension/examples/image_classification.sql
@@ -22,7 +22,7 @@ SELECT left(image::text, 40) || ',...}', target FROM pgml.digits LIMIT 10;
 SELECT * FROM pgml.train('Handwritten Digits', 'classification', 'pgml.digits', 'target');
 
 -- check out the predictions
-SELECT target, pgml.predict('Handwritten Digits', image::FLOAT4[]) AS prediction
+SELECT target, pgml.predict('Handwritten Digits', image) AS prediction
 FROM pgml.digits 
 LIMIT 10;
 
@@ -95,6 +95,6 @@ SELECT * FROM pgml.deploy('Handwritten Digits', 'rollback');
 SELECT * FROM pgml.deploy('Handwritten Digits', 'best_score', 'svm');
 
 -- check out the improved predictions
-SELECT target, pgml.predict('Handwritten Digits', image::FLOAT4[]) AS prediction
+SELECT target, pgml.predict('Handwritten Digits', image) AS prediction
 FROM pgml.digits 
 LIMIT 10;
diff --git a/pgml-extension/src/api.rs b/pgml-extension/src/api.rs
@@ -414,10 +414,35 @@ fn deploy(
 }
 
 #[pg_extern(strict, name = "predict")]
-fn predict(project_name: &str, features: Vec<f32>) -> f32 {
+fn predict_f32(project_name: &str, features: Vec<f32>) -> f32 {
     predict_model(Project::get_deployed_model_id(project_name), features)
 }
 
+#[pg_extern(strict, name = "predict")]
+fn predict_f64(project_name: &str, features: Vec<f64>) -> f32 {
+    predict_f32(project_name, features.iter().map(|&i| i as f32).collect())
+}
+
+#[pg_extern(strict, name = "predict")]
+fn predict_i16(project_name: &str, features: Vec<i16>) -> f32 {
+    predict_f32(project_name, features.iter().map(|&i| i as f32).collect())
+}
+
+#[pg_extern(strict, name = "predict")]
+fn predict_i32(project_name: &str, features: Vec<i32>) -> f32 {
+    predict_f32(project_name, features.iter().map(|&i| i as f32).collect())
+}
+
+#[pg_extern(strict, name = "predict")]
+fn predict_i64(project_name: &str, features: Vec<i64>) -> f32 {
+    predict_f32(project_name, features.iter().map(|&i| i as f32).collect())
+}
+
+#[pg_extern(strict, name = "predict")]
+fn predict_bool(project_name: &str, features: Vec<bool>) -> f32 {
+    predict_f32(project_name, features.iter().map(|&i| i as u8 as f32).collect())
+}
+
 #[pg_extern(strict, name = "predict_proba")]
 fn predict_proba(project_name: &str, features: Vec<f32>) -> Vec<f32> {
     predict_model_proba(Project::get_deployed_model_id(project_name), features)
