aryn-ai · bsowell · Oct 11, 2024 · Oct 11, 2024
diff --git a/lib/sycamore/sycamore/tests/integration/transforms/test_random_sample.py b/lib/sycamore/sycamore/tests/integration/transforms/test_random_sample.py
@@ -0,0 +1,6 @@
+import sycamore
+from sycamore.tests.unit.transforms.test_random_sample import TestRandomSample
+
+
+class TestRandomSampleRay(TestRandomSample):
+    exec_mode = sycamore.EXEC_RAY
diff --git a/lib/sycamore/sycamore/tests/unit/transforms/test_random_sample.py b/lib/sycamore/sycamore/tests/unit/transforms/test_random_sample.py
@@ -8,6 +8,8 @@
 
 
 class TestRandomSample:
+    exec_mode = sycamore.EXEC_LOCAL
+
     @pytest.fixture()
     def docs(self) -> list[Document]:
         print("Generating docs")
@@ -18,7 +20,7 @@ def docs(self) -> list[Document]:
 
     @pytest.fixture()
     def docset(self, docs: list[Document]) -> DocSet:
-        context = sycamore.init()
+        context = sycamore.init(exec_mode=self.exec_mode)
         return context.read.document(docs)
 
     def test_empty_sample(self, docset: DocSet):

diff --git a/lib/sycamore/sycamore/transforms/random_sample.py b/lib/sycamore/sycamore/transforms/random_sample.py
@@ -1,6 +1,8 @@
+import random
 from typing import Optional, TYPE_CHECKING
 
 from sycamore.plan_nodes import Node, Transform
+from sycamore.data import Document
 
 if TYPE_CHECKING:
     from ray.data import Dataset
@@ -25,3 +27,8 @@ def __init__(self, child: Node, fraction: float, seed: Optional[int] = None, **r
     def execute(self, **kwargs) -> "Dataset":
         dataset = self.child().execute()
         return dataset.random_sample(self.fraction, seed=self.seed)
+
+    def local_execute(self, all_docs: list[Document]) -> list[Document]:
+        if self.seed is not None:
+            random.seed(self.seed)
+        return random.sample(all_docs, int(len(all_docs) * self.fraction))