postgresml · SilasMarvin · Nov 9, 2023 · Nov 2, 2023 · Nov 2, 2023 · Nov 3, 2023
diff --git a/pgml-extension/src/api.rs b/pgml-extension/src/api.rs
@@ -4,6 +4,8 @@ use std::str::FromStr;
 use ndarray::Zip;
 use pgrx::iter::{SetOfIterator, TableIterator};
 use pgrx::*;
+use pyo3::prelude::*;
+use pyo3::types::{IntoPyDict, PyDict};
 
 #[cfg(feature = "python")]
 use serde_json::json;
@@ -632,6 +634,75 @@ pub fn transform_string(
     }
 }
 
+struct TransformStreamIterator {
+    locals: Py<PyDict>,
+}
+
+impl TransformStreamIterator {
+    fn new(python_iter: Py<PyAny>) -> Self {
+        let locals = Python::with_gil(|py| -> Result<Py<PyDict>, PyErr> {
+            Ok([("python_iter", python_iter)].into_py_dict(py).into())
+        })
+        .map_err(|e| error!("{e}"))
+        .unwrap();
+        Self { locals }
+    }
+}
+
+impl Iterator for TransformStreamIterator {
+    type Item = String;
+    fn next(&mut self) -> Option<Self::Item> {
+        // We can unwrap this becuase if there is an error the current transaction is aborted in the map_err call
+        Python::with_gil(|py| -> Result<Option<String>, PyErr> {
+            let code = "next(python_iter)";
+            let res: &PyAny = py.eval(code, Some(self.locals.as_ref(py)), None)?;
+            if res.is_none() {
+                Ok(None)
+            } else {
+                let res: String = res.extract()?;
+                Ok(Some(res))
+            }
+        })
+        .map_err(|e| error!("{e}"))
+        .unwrap()
+    }
+}
+
+#[cfg(all(feature = "python", not(feature = "use_as_lib")))]
+#[pg_extern(immutable, parallel_safe, name = "transform_stream")]
+#[allow(unused_variables)] // cache is maintained for api compatibility
+pub fn transform_stream_json(
+    task: JsonB,
+    args: default!(JsonB, "'{}'"),
+    input: default!(&str, "''"),
+    cache: default!(bool, false),
+) -> SetOfIterator<'static, String> {
+    // We can unwrap this becuase if there is an error the current transaction is aborted in the map_err call
+    let python_iter = crate::bindings::transformers::transform_stream(&task.0, &args.0, input)
+        .map_err(|e| error!("{e}"))
+        .unwrap();
+    let res = TransformStreamIterator::new(python_iter);
+    SetOfIterator::new(res)
+}
+
+#[cfg(all(feature = "python", not(feature = "use_as_lib")))]
+#[pg_extern(immutable, parallel_safe, name = "transform_stream")]
+#[allow(unused_variables)] // cache is maintained for api compatibility
+pub fn transform_stream_string(
+    task: String,
+    args: default!(JsonB, "'{}'"),
+    input: default!(&str, "''"),
+    cache: default!(bool, false),
+) -> SetOfIterator<'static, String> {
+    let task_json = json!({ "task": task });
+    // We can unwrap this becuase if there is an error the current transaction is aborted in the map_err call
+    let python_iter = crate::bindings::transformers::transform_stream(&task_json, &args.0, input)
+        .map_err(|e| error!("{e}"))
+        .unwrap();
+    let res = TransformStreamIterator::new(python_iter);
+    SetOfIterator::new(res)
+}
+
 #[cfg(feature = "python")]
 #[pg_extern(immutable, parallel_safe, name = "generate")]
 fn generate(project_name: &str, inputs: &str, config: default!(JsonB, "'{}'")) -> String {

diff --git a/pgml-extension/src/bindings/transformers/mod.rs b/pgml-extension/src/bindings/transformers/mod.rs
@@ -16,41 +16,10 @@ use super::TracebackError;
 
 pub mod whitelist;
 
-create_pymodule!("/src/bindings/transformers/transformers.py");
-
-pub fn transform(
-    task: &serde_json::Value,
-    args: &serde_json::Value,
-    inputs: Vec<&str>,
-) -> Result<serde_json::Value> {
-    crate::bindings::python::activate()?;
-
-    whitelist::verify_task(task)?;
-
-    let task = serde_json::to_string(task)?;
-    let args = serde_json::to_string(args)?;
-    let inputs = serde_json::to_string(&inputs)?;
+mod transformers;
+pub use transformers::*;
 
-    let results = Python::with_gil(|py| -> Result<String> {
-        let transform: Py<PyAny> = get_module!(PY_MODULE)
-            .getattr(py, "transform")
-            .format_traceback(py)?;
-
-        let output = transform
-            .call1(
-                py,
-                PyTuple::new(
-                    py,
-                    &[task.into_py(py), args.into_py(py), inputs.into_py(py)],
-                ),
-            )
-            .format_traceback(py)?;
-
-        output.extract(py).format_traceback(py)
-    })?;
-
-    Ok(serde_json::from_str(&results)?)
-}
+create_pymodule!("/src/bindings/transformers/transformers.py");
 
 pub fn get_model_from(task: &Value) -> Result<String> {
     Python::with_gil(|py| -> Result<String> {