maxframe 0.1.0b5__cp310-cp310-macosx_11_0_arm64.whl

maxframe/tests/utils.py

@@ -0,0 +1,164 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import functools
+import os
+import queue
+import socket
+import types
+from threading import Thread
+from typing import Dict, List, Optional, Set, Tuple
+
+import pytest
+from tornado import netutil
+
+from ..core import Tileable, TileableGraph
+from ..utils import lazy_import
+
+try:
+    from flaky import flaky
+except ImportError:
+
+    def flaky(func=None, **_):
+        if func is not None:
+            return func
+
+        def decorator(f):
+            return f
+
+        return decorator
+
+
+cupy = lazy_import("cupy")
+cudf = lazy_import("cudf")
+ray = lazy_import("ray")
+
+_test_tables_to_drop = set()
+
+
+def bind_unused_port(
+    reuse_port: bool = False, address: str = "127.0.0.1"
+) -> Tuple[socket.socket, int]:
+    sock = netutil.bind_sockets(
+        0, address, family=socket.AF_INET, reuse_port=reuse_port
+    )[0]
+    port = sock.getsockname()[1]
+    return sock, port
+
+
+def tn(s, limit=128):
+    if os.environ.get("TEST_NAME_SUFFIX") is not None:
+        suffix = "_" + os.environ.get("TEST_NAME_SUFFIX").lower()
+        if len(s) + len(suffix) > limit:
+            s = s[: limit - len(suffix)]
+        table_name = s + suffix
+        if table_name.count(".") <= 2:
+            _test_tables_to_drop.add(table_name)
+        return table_name
+    else:
+        if len(s) > limit:
+            s = s[:limit]
+        return s
+
+
+def run_app_in_thread(app_func):
+    def app_thread_func(
+        loop: asyncio.AbstractEventLoop,
+        q: queue.Queue,
+        exit_event: asyncio.Event,
+        args,
+        kwargs,
+    ):
+        asyncio.set_event_loop(loop)
+
+        async def app_main():
+            iterator = app_func(*args, **kwargs)
+            is_gen = isinstance(iterator, types.GeneratorType)
+            try:
+                q.put(next(iterator) if is_gen else iterator)
+                await exit_event.wait()
+            finally:
+                try:
+                    if is_gen:
+                        list(iterator)
+                except StopIteration:
+                    pass
+
+        loop.run_until_complete(app_main())
+
+    @functools.wraps(app_func)
+    def fixture_func(*args, **kwargs):
+        app_loop = asyncio.new_event_loop()
+        q = queue.Queue()
+        exit_event = asyncio.Event(loop=app_loop)
+        app_thread = Thread(
+            name="TestAppThread",
+            target=app_thread_func,
+            args=(app_loop, q, exit_event, args, kwargs),
+        )
+        app_thread.start()
+
+        try:
+            yield q.get()
+        finally:
+
+            async def set_exit_event():
+                exit_event.set()
+
+            asyncio.run_coroutine_threadsafe(set_exit_event(), app_loop)
+            app_thread.join()
+
+    return fixture_func
+
+
+def assert_graph_equal(
+    graph: TileableGraph,
+    expected_node_count: int,
+    expected_nodes: Set[str],
+    expected_edges: Dict[str, List[str]],
+    expected_results: Optional[List[Tileable]],
+):
+    assert len(graph) == expected_node_count
+    key_to_node = {node.key: node for node in graph.iter_nodes()}
+    assert set(graph) == {key_to_node[key] for key in expected_nodes}
+    for pred, successors in expected_edges.items():
+        pred_node = key_to_node[pred]
+        assert graph.count_successors(pred_node) == len(successors)
+        for successor in successors:
+            assert graph.has_successor(pred_node, key_to_node[successor])
+    assert graph.results == expected_results
+
+
+def require_cupy(func):
+    if pytest:
+        func = pytest.mark.cuda(func)
+    func = pytest.mark.skipif(cupy is None, reason="cupy not installed")(func)
+    return func
+
+
+def require_cudf(func):
+    if pytest:
+        func = pytest.mark.cuda(func)
+    func = pytest.mark.skipif(cudf is None, reason="cudf not installed")(func)
+    return func
+
+
+def require_hadoop(func):
+    if pytest:
+        func = pytest.mark.hadoop(func)
+    func = pytest.mark.skipif(
+        not os.environ.get("WITH_HADOOP"), reason="Only run when hadoop is installed"
+    )(func)
+    return func

maxframe/typing_.py

@@ -0,0 +1,37 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import TypeVar, Union
+
+import pandas as pd
+import pyarrow as pa
+
+TimeoutType = Union[int, float, None]
+
+
+ArrowTableType = Union[pa.Table, pa.RecordBatch]
+PandasObjectTypes = Union[
+    pd.DataFrame,
+    pd.Series,
+    pd.Index,
+]
+
+OperatorType = TypeVar("OperatorType")
+TileableType = TypeVar("TileableType")
+ChunkType = TypeVar("ChunkType")
+EntityType = TypeVar("EntityType")
+SessionType = TypeVar("SessionType")
+
+ClusterType = TypeVar("ClusterType")
+ClientType = TypeVar("ClientType")

maxframe/udf.py

@@ -0,0 +1,134 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import shlex
+from typing import Callable, List, Optional, Union
+
+from odps.models import Resource
+
+from .serialization.serializables import (
+    BoolField,
+    FieldTypes,
+    FunctionField,
+    ListField,
+    Serializable,
+    StringField,
+)
+from .utils import tokenize
+
+
+class PythonPackOptions(Serializable):
+    key = StringField("key")
+    requirements = ListField("requirements", FieldTypes.string, default_factory=list)
+    force_rebuild = BoolField("force_rebuild", default=False)
+    prefer_binary = BoolField("prefer_binary", default=False)
+    pre_release = BoolField("pre_release", default=False)
+    pack_instance_id = StringField("pack_instance_id", default=None)
+
+    def __init__(self, key: str = None, **kw):
+        super().__init__(key=key, **kw)
+        if self.key is None:
+            args = {
+                "force_rebuild": self.force_rebuild,
+                "prefer_binary": self.prefer_binary,
+                "pre_release": self.pre_release,
+            }
+            self.key = tokenize(set(self.requirements), args)
+
+    def __repr__(self):
+        return (
+            f"<PythonPackOptions {self.requirements} force_rebuild={self.force_rebuild} "
+            f"prefer_binary={self.prefer_binary} pre_release={self.pre_release}>"
+        )
+
+
+class MarkedFunction(Serializable):
+    func = FunctionField("func")
+    resources = ListField("resources", FieldTypes.string, default_factory=list)
+    pythonpacks = ListField("pythonpacks", FieldTypes.reference, default_factory=list)
+
+    def __init__(self, func: Optional[Callable] = None, **kw):
+        super().__init__(func=func, **kw)
+
+    def __getattr__(self, item):
+        return getattr(self.func, item)
+
+    def __call__(self, *args, **kw):
+        return self.func(*args, **kw)
+
+    def __repr__(self):
+        return f"<MarkedFunction {self.func!r}>"
+
+
+def with_resources(*resources: Union[str, Resource], use_wrapper_class: bool = True):
+    def res_to_str(res: Union[str, Resource]) -> str:
+        if isinstance(res, str):
+            return res
+        res_parts = [res.project.name]
+        if res.schema:
+            res_parts.extend(["schemas", res.schema])
+        res_parts.extend(["resources", res.name])
+        return "/".join(res_parts)
+
+    def func_wrapper(func):
+        str_resources = [res_to_str(r) for r in resources]
+        if not use_wrapper_class:
+            existing = getattr(func, "resources") or []
+            func.resources = existing + str_resources
+            return func
+
+        if isinstance(func, MarkedFunction):
+            func.resources = func.resources + str_resources
+            return func
+        return MarkedFunction(func, resources=str_resources)
+
+    return func_wrapper
+
+
+def with_python_requirements(
+    *requirements: str,
+    force_rebuild: bool = False,
+    prefer_binary: bool = False,
+    pre_release: bool = False,
+):
+    result_req = []
+    for req in requirements:
+        result_req.extend(shlex.split(req))
+
+    def func_wrapper(func):
+        pack_item = PythonPackOptions(
+            requirements=requirements,
+            force_rebuild=force_rebuild,
+            prefer_binary=prefer_binary,
+            pre_release=pre_release,
+        )
+        if isinstance(func, MarkedFunction):
+            func.pythonpacks.append(pack_item)
+            return func
+        return MarkedFunction(func, pythonpacks=[pack_item])
+
+    return func_wrapper
+
+
+with_resource_libraries = with_resources
+
+
+def get_udf_resources(
+    func: Callable,
+) -> List[Union[Resource, str]]:
+    return getattr(func, "resources", None) or []
+
+
+def get_udf_pythonpacks(func: Callable) -> List[PythonPackOptions]:
+    return getattr(func, "pythonpacks", None) or []