cocoindex 0.2.1__cp313-cp313t-manylinux_2_28_aarch64.whl → 0.2.2__cp313-cp313t-manylinux_2_28_aarch64.whl

cocoindex/flow.py

@@ -120,9 +120,10 @@ def _transform_helper(
     else:
         raise ValueError("transform() can only be called on a CocoIndex function")
 
-    return _create_data_slice(
-        flow_builder_state,
-        lambda target_scope, name: flow_builder_state.engine_flow_builder.transform(
+    def _create_data_slice_inner(
+        target_scope: _engine.DataScopeRef | None, name: str | None
+    ) -> _engine.DataSlice:
+        result = flow_builder_state.engine_flow_builder.transform(
             kind,
             dump_engine_object(spec),
             transform_args,
@@ -130,7 +131,12 @@ def _transform_helper(
             flow_builder_state.field_name_builder.build_name(
                 name, prefix=_to_snake_case(_spec_kind(fn_spec)) + "_"
             ),
-        ),
+        )
+        return result
+
+    return _create_data_slice(
+        flow_builder_state,
+        _create_data_slice_inner,
         name,
     )
 
@@ -166,6 +172,7 @@ class _DataSliceState:
     def engine_data_slice(self) -> _engine.DataSlice:
         """
         Get the internal DataSlice.
+        This can be blocking.
         """
         if self._lazy_lock is None:
             if self._data_slice is None:
@@ -179,6 +186,13 @@ class _DataSliceState:
                     self._data_slice = self._data_slice_creator(None)
                 return self._data_slice
 
+    async def engine_data_slice_async(self) -> _engine.DataSlice:
+        """
+        Get the internal DataSlice.
+        This can be blocking.
+        """
+        return await asyncio.to_thread(lambda: self.engine_data_slice)
+
     def attach_to_scope(self, scope: _engine.DataScopeRef, field_name: str) -> None:
         """
         Attach the current data slice (if not yet attached) to the given scope.
@@ -795,9 +809,8 @@ class Flow:
         """
         Setup persistent backends of the flow. The async version.
         """
-        await make_setup_bundle([self]).describe_and_apply_async(
-            report_to_stdout=report_to_stdout
-        )
+        bundle = await make_setup_bundle_async([self])
+        await bundle.describe_and_apply_async(report_to_stdout=report_to_stdout)
 
     def drop(self, report_to_stdout: bool = False) -> None:
         """
@@ -814,9 +827,8 @@ class Flow:
         """
         Drop persistent backends of the flow. The async version.
         """
-        await make_drop_bundle([self]).describe_and_apply_async(
-            report_to_stdout=report_to_stdout
-        )
+        bundle = await make_drop_bundle_async([self])
+        await bundle.describe_and_apply_async(report_to_stdout=report_to_stdout)
 
     def close(self) -> None:
         """
@@ -1071,19 +1083,16 @@ class TransformFlow(Generic[T]):
                 _DataSliceState(flow_builder_state, engine_ds)
             )
 
-        output = self._flow_fn(**kwargs)
-        flow_builder_state.engine_flow_builder.set_direct_output(
-            _data_slice_state(output).engine_data_slice
-        )
+        output = await asyncio.to_thread(lambda: self._flow_fn(**kwargs))
+        output_data_slice = await _data_slice_state(output).engine_data_slice_async()
+
+        flow_builder_state.engine_flow_builder.set_direct_output(output_data_slice)
         engine_flow = (
             await flow_builder_state.engine_flow_builder.build_transient_flow_async(
                 execution_context.event_loop
             )
         )
-
-        engine_return_type = (
-            _data_slice_state(output).engine_data_slice.data_type().schema()
-        )
+        engine_return_type = output_data_slice.data_type().schema()
         python_return_type: type[T] | None = _get_data_slice_annotation_type(
             inspect.signature(self._flow_fn).return_annotation
         )
@@ -1142,28 +1151,42 @@ def transform_flow() -> Callable[[Callable[..., DataSlice[T]]], TransformFlow[T]
     return _transform_flow_wrapper
 
 
-def make_setup_bundle(flow_iter: Iterable[Flow]) -> SetupChangeBundle:
+async def make_setup_bundle_async(flow_iter: Iterable[Flow]) -> SetupChangeBundle:
     """
     Make a bundle to setup flows with the given names.
     """
     full_names = []
     for fl in flow_iter:
-        fl.internal_flow()
+        await fl.internal_flow_async()
         full_names.append(fl.full_name)
     return SetupChangeBundle(_engine.make_setup_bundle(full_names))
 
 
-def make_drop_bundle(flow_iter: Iterable[Flow]) -> SetupChangeBundle:
+def make_setup_bundle(flow_iter: Iterable[Flow]) -> SetupChangeBundle:
+    """
+    Make a bundle to setup flows with the given names.
+    """
+    return execution_context.run(make_setup_bundle_async(flow_iter))
+
+
+async def make_drop_bundle_async(flow_iter: Iterable[Flow]) -> SetupChangeBundle:
     """
     Make a bundle to drop flows with the given names.
     """
     full_names = []
     for fl in flow_iter:
-        fl.internal_flow()
+        await fl.internal_flow_async()
         full_names.append(fl.full_name)
     return SetupChangeBundle(_engine.make_drop_bundle(full_names))
 
 
+def make_drop_bundle(flow_iter: Iterable[Flow]) -> SetupChangeBundle:
+    """
+    Make a bundle to drop flows with the given names.
+    """
+    return execution_context.run(make_drop_bundle_async(flow_iter))
+
+
 def setup_all_flows(report_to_stdout: bool = False) -> None:
     """
     Setup all flows registered in the current process.

cocoindex/tests/test_transform_flow.py

@@ -166,6 +166,25 @@ class GpuAppendSuffixExecutor:
         return f"{text}{self.spec.suffix}"
 
 
+class GpuAppendSuffixWithAnalyzePrepare(cocoindex.op.FunctionSpec):
+    suffix: str
+
+
+@cocoindex.op.executor_class(gpu=True)
+class GpuAppendSuffixWithAnalyzePrepareExecutor:
+    spec: GpuAppendSuffixWithAnalyzePrepare
+    suffix: str
+
+    def analyze(self) -> Any:
+        return str
+
+    def prepare(self) -> None:
+        self.suffix = self.spec.suffix
+
+    def __call__(self, text: str) -> str:
+        return f"{text}{self.suffix}"
+
+
 def test_gpu_function() -> None:
     @cocoindex.transform_flow()
     def transform_flow(text: cocoindex.DataSlice[str]) -> cocoindex.DataSlice[str]:
@@ -174,3 +193,15 @@ def test_gpu_function() -> None:
     result = transform_flow.eval("Hello")
     expected = "Hello world!"
     assert result == expected, f"Expected {expected}, got {result}"
+
+    @cocoindex.transform_flow()
+    def transform_flow_with_analyze_prepare(
+        text: cocoindex.DataSlice[str],
+    ) -> cocoindex.DataSlice[str]:
+        return text.transform(gpu_append_world).transform(
+            GpuAppendSuffixWithAnalyzePrepare(suffix="!!")
+        )
+
+    result = transform_flow_with_analyze_prepare.eval("Hello")
+    expected = "Hello world!!"
+    assert result == expected, f"Expected {expected}, got {result}"

{cocoindex-0.2.1.dist-info → cocoindex-0.2.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cocoindex
-Version: 0.2.1
+Version: 0.2.2
 Requires-Dist: click>=8.1.8
 Requires-Dist: rich>=14.0.0
 Requires-Dist: python-dotenv>=1.1.0
@@ -84,7 +84,7 @@ CocoIndex makes it effortless to transform data with AI, and keep source data an
 </br>
 
 <p align="center">
-<img width="4187" height="1883" alt="CocoIndex Features" src="https://github.com/user-attachments/assets/6147673f-0daf-4313-a0c7-7c4205ba2e31" />
+<img alt="CocoIndex Features" src="https://cocoindex.io/images/venn2.svg" />
 </p>
 
 </br>

{cocoindex-0.2.1.dist-info → cocoindex-0.2.2.dist-info}/RECORD

@@ -1,13 +1,13 @@
-cocoindex-0.2.1.dist-info/METADATA,sha256=YjOsHGGJpjD5ioS_7bQkhkVBCmgEFH2NwFxu24PAY24,12187
-cocoindex-0.2.1.dist-info/WHEEL,sha256=SQ8BH9TyCgXz3Ds3g6o4Uyw-YxShg-r1RAudNXd9Wds,110
-cocoindex-0.2.1.dist-info/entry_points.txt,sha256=_NretjYVzBdNTn7dK-zgwr7YfG2afz1u1uSE-5bZXF8,46
-cocoindex-0.2.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+cocoindex-0.2.2.dist-info/METADATA,sha256=INC54JkKwVqZVSQHVM5mtJJ9dcFRMegCtCRwMwtK2mg,12118
+cocoindex-0.2.2.dist-info/WHEEL,sha256=SQ8BH9TyCgXz3Ds3g6o4Uyw-YxShg-r1RAudNXd9Wds,110
+cocoindex-0.2.2.dist-info/entry_points.txt,sha256=_NretjYVzBdNTn7dK-zgwr7YfG2afz1u1uSE-5bZXF8,46
+cocoindex-0.2.2.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
 cocoindex/__init__.py,sha256=sLpSVO5Cotgn_82lawxvXnaqfa-qj33rytWBAe2MTtU,2201
-cocoindex/_engine.cpython-313t-aarch64-linux-gnu.so,sha256=rQWSRJw7RjNuTw_eMA5LhSX36QPIPZMMeqrTSUX-UOg,69818216
+cocoindex/_engine.cpython-313t-aarch64-linux-gnu.so,sha256=6Hm1txg9Y0dqjThfydNyJ7HTnHPPFdFNJWMQkR0ELB4,69832168
 cocoindex/auth_registry.py,sha256=PE1-kVkcyC1G2C_V7b1kvYzeq73OFQehWKQP7ln7fJ8,1478
 cocoindex/cli.py,sha256=8Q2D0gI_yHorxfFC7tzV96oJAR7PtTn_T4qWLorPPoM,20675
 cocoindex/convert.py,sha256=x3eYx-NQD-z9h-jmpVgUMiSTqll6Baa_SfZJF3mZSUc,22071
-cocoindex/flow.py,sha256=7133Z6QU4LF_GMV81AmkO_aGhfAQCStwOnMW7BABymU,36258
+cocoindex/flow.py,sha256=AoWt8i05MfEnnoAQR0EWzZU-9k36UmCFuNLb3itlGO0,37204
 cocoindex/functions.py,sha256=09erNt3WbzY9l1KER-akBF2O5-6xEahV2ORBECaL6yk,12260
 cocoindex/index.py,sha256=j93B9jEvvLXHtpzKWL88SY6wCGEoPgpsQhEGHlyYGFg,540
 cocoindex/lib.py,sha256=f--9dAYd84CZosbDZqNW0oGbBLsY3dXiUTR1VrfQ_QY,817
@@ -23,11 +23,11 @@ cocoindex/targets.py,sha256=Nfh_tpFd1goTnS_cxBjIs4j9zl3Z4Z1JomAQ1dl3Sic,2796
 cocoindex/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cocoindex/tests/test_convert.py,sha256=UsQDuz7qukT9UAU5LBv3atNTqdUVc6uW8lrmLA7ys_w,50207
 cocoindex/tests/test_optional_database.py,sha256=snAmkNa6wtOSaxoZE1HgjvL5v_ylitt3Jt_9df4Cgdc,8506
-cocoindex/tests/test_transform_flow.py,sha256=MdIzLdIMBPlIfKRFAB7n8gIZixqChL6xb4vScvcyiNU,5148
+cocoindex/tests/test_transform_flow.py,sha256=G69w-n-vnCTo3r9hVIk2lJNAQEkGUA7PZfHsXna3oS0,6030
 cocoindex/tests/test_typing.py,sha256=9OF3lO2uSpZBefkEJx7WRbnkXjwQtvlQIeeARYQID68,12391
 cocoindex/tests/test_validation.py,sha256=X6AQzVs-hVKIXcrHMEMQnhfUE8at7iXQnPq8nHNhZ2Q,4543
 cocoindex/typing.py,sha256=lEQYIzAGVKQp6RnhyeopY9Q7xEED7yQj3ZMxvTPblV8,14200
 cocoindex/user_app_loader.py,sha256=o2n25y7CGw9Qrb-RN6tiNlv2YJ9YLM8NSMhJZjLhhKQ,1914
 cocoindex/utils.py,sha256=hUhX-XV6XGCtJSEIpBOuDv6VvqImwPlgBxztBTw7u0U,598
 cocoindex/validation.py,sha256=PZnJoby4sLbsmPv9fOjOQXuefjfZ7gmtsiTGU8SH-tc,3090
-cocoindex-0.2.1.dist-info/RECORD,,
+cocoindex-0.2.2.dist-info/RECORD,,