maxframe 0.1.0b5__cp39-cp39-macosx_11_0_arm64.whl

maxframe/lib/filesystem/arrow.py

@@ -0,0 +1,236 @@
+# 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 os
+import subprocess
+import weakref
+from typing import BinaryIO, Dict, Iterator, List, TextIO, Tuple, Union
+from urllib.parse import urlparse
+
+import pyarrow as pa
+from pyarrow.fs import FileInfo, FileSelector
+from pyarrow.fs import FileSystem as ArrowFileSystem
+from pyarrow.fs import FileType
+from pyarrow.fs import HadoopFileSystem as ArrowHadoopFileSystem
+from pyarrow.fs import LocalFileSystem as ArrowLocalFileSystem
+
+from ...utils import implements, stringify_path
+from .core import FileSystem, path_type
+
+__all__ = ("ArrowBasedLocalFileSystem", "HadoopFileSystem")
+
+
+# When pyarrow.fs.FileSystem gc collected,
+# the underlying connection will be closed,
+# so we hold the reference to make sure
+# FileSystem will not be gc collected before file object
+_file_to_filesystems = weakref.WeakKeyDictionary()
+
+
+class ArrowBasedFileSystem(FileSystem):
+    """
+    FileSystem implemented with arrow fs API (>=2.0.0).
+    """
+
+    def __init__(self, arrow_fs: ArrowFileSystem, sequential_read=False):
+        self._arrow_fs = arrow_fs
+        # for open('rb'), open a sequential reading only or not
+        self._sequential_read = sequential_read
+
+    @staticmethod
+    def _process_path(path):
+        return stringify_path(path)
+
+    @implements(FileSystem.cat)
+    def cat(self, path: path_type) -> bytes:
+        path = self._process_path(path)
+        file: pa.NativeFile = self._arrow_fs.open_input_stream(path)
+        return file.read()
+
+    @implements(FileSystem.ls)
+    def ls(self, path: path_type) -> List[path_type]:
+        path = self._process_path(path)
+        file_selector: FileSelector = FileSelector(path)
+        paths = []
+        for file_info in self._arrow_fs.get_file_info(file_selector):
+            paths.append(file_info.path)
+        return paths
+
+    def _get_file_info(self, path: path_type) -> FileInfo:
+        path = self._process_path(path)
+        file_info: FileInfo = self._arrow_fs.get_file_info([path])[0]
+        return file_info
+
+    @implements(FileSystem.delete)
+    def delete(self, path: path_type, recursive: bool = False):
+        path = self._process_path(path)
+        info = self._get_file_info(path)
+        if info.is_file:
+            self._arrow_fs.delete_file(path)
+        elif info.type == FileType.Directory:
+            if not recursive and len(self.ls(path)) > 0:
+                raise OSError(f"[Errno 66] Directory not empty: '{path}'")
+            self._arrow_fs.delete_dir(path)
+        else:  # pragma: no cover
+            raise TypeError(f"path({path}) to delete must be a file or directory")
+
+    @implements(FileSystem.rename)
+    def rename(self, path: path_type, new_path: path_type):
+        path = self._process_path(path)
+        new_path = self._process_path(new_path)
+        self._arrow_fs.move(path, new_path)
+
+    @implements(FileSystem.stat)
+    def stat(self, path: path_type) -> Dict:
+        path = self._process_path(path)
+        info = self._get_file_info(path)
+        stat = dict(name=path, size=info.size, modified_time=info.mtime_ns / 1e9)
+        if info.type == FileType.File:
+            stat["type"] = "file"
+        elif info.type == FileType.Directory:
+            stat["type"] = "directory"
+        else:  # pragma: no cover
+            stat["type"] = "other"
+        return stat
+
+    @implements(FileSystem.mkdir)
+    def mkdir(self, path: path_type, create_parents: bool = True):
+        path = self._process_path(path)
+        self._arrow_fs.create_dir(path, recursive=create_parents)
+
+    @implements(FileSystem.isdir)
+    def isdir(self, path: path_type) -> bool:
+        path = self._process_path(path)
+        info = self._get_file_info(path)
+        return info.type == FileType.Directory
+
+    @implements(FileSystem.isfile)
+    def isfile(self, path: path_type) -> bool:
+        path = self._process_path(path)
+        info = self._get_file_info(path)
+        return info.is_file
+
+    @implements(FileSystem._isfilestore)
+    def _isfilestore(self) -> bool:
+        return True
+
+    @implements(FileSystem.exists)
+    def exists(self, path: path_type):
+        path = self._process_path(path)
+        info = self._get_file_info(path)
+        return info.type != FileType.NotFound
+
+    @implements(FileSystem.open)
+    def open(self, path: path_type, mode: str = "rb") -> Union[BinaryIO, TextIO]:
+        path = self._process_path(path)
+        is_binary = mode.endswith("b")
+        if not is_binary:  # pragma: no cover
+            raise ValueError(
+                f"mode can only be binary for arrow based filesystem, got {mode}"
+            )
+        mode = mode.rstrip("b")
+        if mode == "w":
+            file = self._arrow_fs.open_output_stream(path)
+        elif mode == "r":
+            if self._sequential_read:  # pragma: no cover
+                file = self._arrow_fs.open_input_stream(path)
+            else:
+                file = self._arrow_fs.open_input_file(path)
+        elif mode == "a":
+            file = self._arrow_fs.open_append_stream(path)
+        else:  # pragma: no cover
+            raise ValueError(
+                f'mode can only be "wb", "rb" and "ab" for '
+                f"arrow based filesystem, got {mode}"
+            )
+
+        _file_to_filesystems[file] = self._arrow_fs
+        return file
+
+    @implements(FileSystem.walk)
+    def walk(self, path: path_type) -> Iterator[Tuple[str, List[str], List[str]]]:
+        path = self._process_path(path)
+        q = [path]
+        while q:
+            curr = q.pop(0)
+            file_selector: FileSelector = FileSelector(curr)
+            dirs, files = [], []
+            for info in self._arrow_fs.get_file_info(file_selector):
+                if info.type == FileType.File:
+                    files.append(info.base_name)
+                elif info.type == FileType.Directory:
+                    dirs.append(info.base_name)
+                    q.append(info.path)
+                else:  # pragma: no cover
+                    continue
+            yield curr, dirs, files
+
+    @implements(FileSystem.glob)
+    def glob(self, path: path_type, recursive: bool = False) -> List[path_type]:
+        from ._glob import FileSystemGlob
+
+        path = self._process_path(path)
+        return FileSystemGlob(self).glob(path, recursive=recursive)
+
+
+class ArrowBasedLocalFileSystem(ArrowBasedFileSystem):
+    def __init__(self):
+        super().__init__(ArrowLocalFileSystem())
+
+    _instance = None
+
+    @classmethod
+    def get_instance(cls):
+        if cls._instance is None:
+            cls._instance = ArrowBasedLocalFileSystem()
+        return cls._instance
+
+
+class HadoopFileSystem(ArrowBasedFileSystem):
+    def __init__(
+        self,
+        host="default",
+        port=0,
+        user=None,
+        kerb_ticket=None,
+        driver="libhdfs",
+        extra_conf=None,
+    ):
+        assert driver == "libhdfs"
+        if "HADOOP_HOME" in os.environ and "CLASSPATH" not in os.environ:
+            classpath_proc = subprocess.run(
+                [os.environ["HADOOP_HOME"] + "/bin/hdfs", "classpath", "--glob"],
+                stdout=subprocess.PIPE,
+            )
+            os.environ["CLASSPATH"] = classpath_proc.stdout.decode().strip()
+        arrow_fs = ArrowHadoopFileSystem(
+            host=host,
+            port=port,
+            user=user,
+            kerb_ticket=kerb_ticket,
+            extra_conf=extra_conf,
+        )
+        super().__init__(arrow_fs)
+
+    @staticmethod
+    def _process_path(path):
+        path = ArrowBasedFileSystem._process_path(path)
+        # use urlparse to extract path from like:
+        # hdfs://localhost:8020/tmp/test/simple_test.csv,
+        # due to the reason that pa.fs.HadoopFileSystem cannot accept
+        # path with hdfs:// prefix
+        if path.startswith("hdfs://"):
+            return urlparse(path).path
+        else:
+            return path

maxframe/lib/filesystem/base.py

@@ -0,0 +1,263 @@
+# 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 os
+from abc import ABC, abstractmethod
+from typing import BinaryIO, Dict, Iterator, List, TextIO, Tuple, Union
+from urllib.parse import urlparse
+
+from ...utils import stringify_path
+
+path_type = Union[str, os.PathLike]
+
+
+class FileSystem(ABC):
+    """
+    Abstract filesystem interface
+    """
+
+    @abstractmethod
+    def cat(self, path: path_type) -> bytes:
+        """
+        Return contents of file as a bytes object
+
+        Parameters
+        ----------
+        path : str or path-like
+            File path to read content from.
+
+        Returns
+        -------
+        contents : bytes
+        """
+
+    @abstractmethod
+    def ls(self, path: path_type) -> List[path_type]:
+        """
+        Return list of file paths
+
+        Returns
+        -------
+        paths : list
+        """
+
+    @abstractmethod
+    def delete(self, path: path_type, recursive: bool = False):
+        """
+        Delete the indicated file or directory
+
+        Parameters
+        ----------
+        path : str
+        recursive : bool, default False
+            If True, also delete child paths for directories
+        """
+
+    def disk_usage(self, path: path_type) -> int:
+        """
+        Compute bytes used by all contents under indicated path in file tree
+
+        Parameters
+        ----------
+        path : string
+            Can be a file path or directory
+
+        Returns
+        -------
+        usage : int
+        """
+        path = stringify_path(path)
+        path_info = self.stat(path)
+        if path_info["type"] == "file":
+            return path_info["size"]
+
+        total = 0
+        for root, directories, files in self.walk(path):
+            for child_path in files:
+                abspath = self.path_join(root, child_path)
+                total += self.stat(abspath)["size"]
+
+        return total
+
+    def path_join(self, *args):
+        return self.pathsep.join(args)
+
+    def path_split(self, path):
+        """
+        Split a pathname. Returns tuple "(head, tail)" where "tail" is everything after the final slash. Either part
+        may be empty.
+
+        Parameters
+        ----------
+        path : string
+            Can be a file path or directory
+
+        Returns
+        -------
+        usage : int
+        """
+        splits = path.rsplit(self.pathsep, 1)
+        if len(splits) == 1:
+            return "", splits[0]
+        else:
+            return splits
+
+    @abstractmethod
+    def stat(self, path: path_type) -> Dict:
+        """
+        Information about a filesystem entry.
+
+        Returns
+        -------
+        stat : dict
+        """
+
+    def rm(self, path: path_type, recursive: bool = False):
+        """
+        Alias for FileSystem.delete
+        """
+        return self.delete(path, recursive=recursive)
+
+    def mv(self, path, new_path):
+        """
+        Alias for FileSystem.rename
+        """
+        return self.rename(path, new_path)
+
+    @abstractmethod
+    def rename(self, path: path_type, new_path: path_type):
+        """
+        Rename file, like UNIX mv command
+
+        Parameters
+        ----------
+        path : string
+            Path to alter
+        new_path : string
+            Path to move to
+        """
+
+    @abstractmethod
+    def mkdir(self, path: path_type, create_parents: bool = True):
+        """
+        Create a directory.
+
+        Parameters
+        ----------
+        path : str
+            Path to the directory.
+        create_parents : bool, default True
+            If the parent directories don't exists create them as well.
+        """
+
+    @abstractmethod
+    def exists(self, path: path_type):
+        """
+        Return True if path exists.
+
+        Parameters
+        ----------
+        path : str
+            Path to check.
+        """
+
+    @abstractmethod
+    def isdir(self, path: path_type) -> bool:
+        """
+        Return True if path is a directory.
+
+        Parameters
+        ----------
+        path : str
+            Path to check.
+        """
+
+    @abstractmethod
+    def isfile(self, path: path_type) -> bool:
+        """
+        Return True if path is a file.
+
+        Parameters
+        ----------
+        path : str
+            Path to check.
+        """
+
+    @abstractmethod
+    def _isfilestore(self) -> bool:
+        """
+        Returns True if this FileSystem is a unix-style file store with
+        directories.
+        """
+
+    @abstractmethod
+    def open(self, path: path_type, mode: str = "rb") -> Union[BinaryIO, TextIO]:
+        """
+        Open file for reading or writing.
+        """
+
+    @abstractmethod
+    def walk(self, path: path_type) -> Iterator[Tuple[str, List[str], List[str]]]:
+        """
+        Directory tree generator.
+
+        Parameters
+        ----------
+        path : str
+
+        Returns
+        -------
+        generator
+        """
+
+    @abstractmethod
+    def glob(self, path: path_type, recursive: bool = False) -> List[path_type]:
+        """
+        Return a list of paths matching a pathname pattern.
+
+        Parameters
+        ----------
+        path : str
+            Pattern may contain simple shell-style wildcards
+        recursive : bool
+            If recursive is true, the pattern '**' will match any files and
+            zero or more directories and subdirectories.
+
+        Returns
+        -------
+        paths : List
+        """
+
+    @property
+    def pathsep(self) -> str:
+        return "/"
+
+    @staticmethod
+    def parse_from_path(uri: str):
+        parsed_uri = urlparse(uri)
+        options = dict()
+        options["host"] = parsed_uri.netloc.rsplit("@", 1)[-1].rsplit(":", 1)[0]
+        if parsed_uri.port:
+            options["port"] = parsed_uri.port
+        if parsed_uri.username:
+            options["user"] = parsed_uri.username
+        if parsed_uri.password:
+            options["password"] = parsed_uri.password
+        return options
+
+    @classmethod
+    def get_storage_options(cls, storage_options: Dict, uri: str) -> Dict:
+        options = cls.parse_from_path(uri)
+        storage_options.update(options)
+        return storage_options

maxframe/lib/filesystem/core.py

@@ -0,0 +1,95 @@
+# 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 glob as glob_
+import os
+from typing import Dict, List
+from urllib.parse import urlparse
+
+from ..compression import compress
+from .base import FileSystem, path_type
+from .local import LocalFileSystem
+from .oss import OSSFileSystem
+
+_filesystems = {"file": LocalFileSystem, "oss": OSSFileSystem}
+_scheme_to_dependencies = {
+    "hdfs": ["pyarrow"],
+    "az": ["fsspec", "adlfs"],
+    "abfs": ["fsspec", "adlfs"],
+    "s3": ["fsspec", "s3fs"],
+}
+
+
+def register_filesystem(name: str, fs):
+    _filesystems[name] = fs
+
+
+def get_fs(path: path_type, storage_options: Dict = None) -> FileSystem:
+    if storage_options is None:
+        storage_options = dict()
+
+    # detect scheme
+    if os.path.exists(path) or glob_.glob(path):
+        scheme = "file"
+    else:
+        scheme = urlparse(path).scheme
+    if scheme == "" or len(scheme) == 1:  # len == 1 for windows
+        scheme = "file"
+
+    if scheme in _filesystems:
+        file_system_type = _filesystems[scheme]
+        if scheme == "file" or scheme == "oss":
+            # local file systems are singletons.
+            return file_system_type.get_instance()
+        else:
+            storage_options = file_system_type.get_storage_options(
+                storage_options, path
+            )
+            return file_system_type(**storage_options)
+    elif scheme in _scheme_to_dependencies:  # pragma: no cover
+        dependencies = ", ".join(_scheme_to_dependencies[scheme])
+        raise ImportError(f"Need to install {dependencies} to access {scheme}.")
+    else:
+        raise ValueError(
+            f"Unknown file system type: {scheme}, "
+            f'available include: {", ".join(_scheme_to_dependencies.keys())}'
+        )
+
+
+def glob(path: path_type, storage_options: Dict = None) -> List[path_type]:
+    if "*" in path:
+        fs = get_fs(path, storage_options)
+        return fs.glob(path)
+    else:
+        return [path]
+
+
+def file_size(path: path_type, storage_options: Dict = None) -> int:
+    fs = get_fs(path, storage_options)
+    return fs.stat(path)["size"]
+
+
+def open_file(
+    path: path_type,
+    mode: str = "rb",
+    compression: str = None,
+    storage_options: Dict = None,
+):
+    fs = get_fs(path, storage_options)
+    file = fs.open(path, mode=mode)
+
+    if compression is not None:
+        file = compress(file, compression)
+
+    return file