Added pgml.transform to queries

santiadavani · santiadavani · commit 0f3bf87917f5 · 2023-06-08T15:05:23.000-07:00
diff --git a/pgml-sdks/python/pgml/pgml/__init__.py b/pgml-sdks/python/pgml/pgml/__init__.py
@@ -5,4 +5,4 @@
     run_select_statement,
     run_drop_or_delete_statement,
 )
-from .queries import Embed, CosineDistance
+from .queries import Embed, CosineDistance, Transform
diff --git a/pgml-sdks/python/pgml/pgml/queries.py b/pgml-sdks/python/pgml/pgml/queries.py
@@ -1,7 +1,8 @@
 from typing import Any
-from pypika.functions import Function
-from typing import Dict
-from pypika import JSON, Array
+from pypika.functions import Function, Cast
+from typing import Dict, List
+from pypika import JSON, Array, Field
+import json
 
 
 class Embed(Function):
@@ -20,3 +21,29 @@ def __init__(
 class CosineDistance(Function):
     def __init__(self, lhs: Array, rhs: Array, alias: str = "cosine") -> None:
         super(CosineDistance, self).__init__("cosine_distance", lhs, rhs, alias=alias)
+
+
+class Transform(Function):
+    def __init__(
+        self,
+        task: str | Dict[str, Any],
+        inputs: List,
+        args: Dict[str, Any] = {},
+        alias: str = "transform",
+    ) -> None:
+        super(Transform, self).__init__(
+            "pgml.transform", task=task, inputs=inputs, args=args, alias=alias
+        )
+        self.task = task
+        self.inputs = inputs
+        self.args = args
+
+    def get_function_sql(self, **kwargs: Any) -> str:
+        if isinstance(self.task, str):
+            return "pgml.transform(task => '{}', inputs => ARRAY{}, args => '{}'::JSONB)".format(
+                self.task, self.inputs, self.args
+            )
+        elif isinstance(self.task, dict):
+            return "pgml.transform(task => '{}'::JSONB, inputs => ARRAY{}, args => '{}'::JSONB)".format(
+                json.dumps(self.task), self.inputs, self.args
+            )

Original file line number	Diff line number	Diff line change
`@@ -5,4 +5,4 @@`
`5`	`5`	`run_select_statement,`
`6`	`6`	`run_drop_or_delete_statement,`
`7`	`7`	`)`
`8`		`-from .queries import Embed, CosineDistance`
	`8`	`+from .queries import Embed, CosineDistance, Transform`